1"""OpenAI-style /v1/images/generations server for FLUX.2 [klein] 4B (m64 capsule).23Ported from gx10:~/code/vlmtest/flux2_server.py (run by machines/gx10/start_flux2_klein_4b.sh),4through the sana15-1.6b-gb10 capsule's server. Differences from the original:5- weights load from /weights/flux2 (an m64 artifact mount), port from M64_PORT; no CPU offload;6- size, steps and guidance are baked constants (the capsule's [params], the original's7  flux2-klein-4b defaults: 4 steps, guidance 1, as the distilled model wants); a request carries8  only `prompt` and `seed` (the m64 front rejects anything else with a 400); one image;9- progress is reported as JSON lines appended to M64_STATUS.10"""1112import base6413import io14import json15import os16import secrets17import threading18import time19import traceback20from typing import Optional2122import torch23import uvicorn24from fastapi import FastAPI, HTTPException25from fastapi.responses import JSONResponse26from pydantic import BaseModel, Field2728SERVED_MODEL_ID = "flux2-klein-4b"29WEIGHTS = "/weights/flux2"30PORT = int(os.environ["M64_PORT"])31STATUS = os.environ["M64_STATUS"]3233# Baked generation settings. Keep in sync with [params] in capsule.toml.34SIZE = "1024x1024"35NUM_INFERENCE_STEPS = 436GUIDANCE_SCALE = 1.0  # klein is distilled: 1.0, no CFG branch37WIDTH, HEIGHT = (int(x) for x in SIZE.split("x"))38MAX_SEED = 2**63 - 13940app = FastAPI()41pipe = None42lock = threading.Lock()  # one generation at a time on the single pipeline434445def status(phase: str, **extra) -> None:46    with open(STATUS, "a") as f:47        f.write(json.dumps({"phase": phase, **extra}) + "\n")484950def load_model() -> None:51    global pipe52    from diffusers import Flux2KleinPipeline5354    status("loading", progress=0.0, detail="loading FLUX.2 [klein] 4B (bf16)")55    print("Loading FLUX.2 [klein] 4B...", flush=True)56    t0 = time.time()57    pipe = Flux2KleinPipeline.from_pretrained(WEIGHTS, torch_dtype=torch.bfloat16)58    pipe.to("cuda")59    print(f"Model loaded in {time.time() - t0:.1f}s.", flush=True)6061    status("warming", detail="one short warm-up generation")62    t0 = time.time()63    with lock:64        pipe(prompt="warm-up", width=WIDTH, height=HEIGHT, num_inference_steps=2,65             guidance_scale=GUIDANCE_SCALE, generator=torch.Generator("cuda").manual_seed(0))66    print(f"Warm-up done in {time.time() - t0:.1f}s.", flush=True)676869class ImageRequest(BaseModel):70    prompt: str = Field(..., min_length=1, description="Text prompt")71    seed: Optional[int] = Field(None, ge=0, le=MAX_SEED, description="RNG seed. null = random; the one used is returned.")727374def baked() -> dict:75    return {"size": SIZE, "num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE}767778@app.get("/health")79def health():80    if pipe is None:81        raise HTTPException(503, "Model not loaded")82    return {"status": "ok", "model": SERVED_MODEL_ID}838485@app.get("/v1/config")86def config():87    return {88        "model": SERVED_MODEL_ID,89        "params": baked(),90        "request_fields": {91            "prompt": {"required": True, "description": "Text prompt"},92            "seed": {"required": False, "default": None, "description": "RNG seed; null = random"},93        },94    }959697@app.post("/v1/images/generations")98def generate(req: ImageRequest):99    if pipe is None:100        raise HTTPException(503, "Model not loaded")101    seed = req.seed if req.seed is not None else secrets.randbelow(2**31)102    t0 = time.time()103    with lock:104        gen = torch.Generator("cuda").manual_seed(seed)105        images = pipe(106            prompt=req.prompt,107            width=WIDTH, height=HEIGHT,108            num_inference_steps=NUM_INFERENCE_STEPS,109            guidance_scale=GUIDANCE_SCALE,110            num_images_per_prompt=1,111            generator=gen,112        ).images113    elapsed = time.time() - t0114    print(f"Generated {len(images)} image(s) in {elapsed:.1f}s (seed {seed})", flush=True)115    data = []116    for img in images:117        buf = io.BytesIO()118        img.save(buf, format="PNG")119        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})120    return JSONResponse({121        "created": int(time.time()),122        "model": SERVED_MODEL_ID,123        "seed": seed,124        "elapsed_seconds": round(elapsed, 2),125        "data": data,126        "params_used": {**baked(), "seed": seed},127    })128129130def main() -> None:131    try:132        load_model()133    except Exception:134        traceback.print_exc()135        status("error", detail="model load failed; see container logs")136        raise137    status("ready")138    uvicorn.run(app, host="0.0.0.0", port=PORT)  # the gate reaches us over the capsule network139140141if __name__ == "__main__":142    main()