1"""OpenAI-style /v1/images/generations server for SANA 1.5 1.6B (m64 capsule).23Ported from an earlier standalone SANA server. Differences from the original:4- weights load from /weights/sana (an m64 artifact mount), port from M64_PORT;5- size, steps and guidance are baked constants (the capsule's [params]); a request carries6  only `prompt` and `seed` (the m64 front rejects anything else with a 400);7- progress is reported as JSON lines appended to M64_STATUS.8"""910import base6411import io12import json13import os14import secrets15import threading16import time17import traceback18from typing import Optional1920import torch21import uvicorn22from fastapi import FastAPI, HTTPException23from fastapi.responses import JSONResponse24from pydantic import BaseModel, Field2526SERVED_MODEL_ID = "sana-1.5-1.6b"27WEIGHTS = "/weights/sana"28PORT = int(os.environ["M64_PORT"])29STATUS = os.environ["M64_STATUS"]3031# Baked generation settings. Keep in sync with [params] in capsule.toml.32SIZE = "1024x1024"33NUM_INFERENCE_STEPS = 2034GUIDANCE_SCALE = 4.535WIDTH, HEIGHT = (int(x) for x in SIZE.split("x"))36MAX_SEED = 2**63 - 13738app = FastAPI()39pipe = None40lock = threading.Lock()  # one generation at a time on the single pipeline414243def status(phase: str, **extra) -> None:44    with open(STATUS, "a") as f:45        f.write(json.dumps({"phase": phase, **extra}) + "\n")464748def load_model() -> None:49    global pipe50    from diffusers import SanaPipeline5152    status("loading", progress=0.0, detail="loading SANA 1.5 1.6B (bf16)")53    print("Loading SANA 1.5 1.6B...", flush=True)54    t0 = time.time()55    pipe = SanaPipeline.from_pretrained(WEIGHTS, torch_dtype=torch.bfloat16)56    pipe.to("cuda")57    pipe.text_encoder.to(torch.bfloat16)58    print(f"Model loaded in {time.time() - t0:.1f}s.", flush=True)5960    status("warming", detail="one short warm-up generation")61    t0 = time.time()62    with lock:63        pipe(prompt="warm-up", width=WIDTH, height=HEIGHT, num_inference_steps=2,64             guidance_scale=GUIDANCE_SCALE, generator=torch.Generator("cuda").manual_seed(0))65    print(f"Warm-up done in {time.time() - t0:.1f}s.", flush=True)666768class ImageRequest(BaseModel):69    prompt: str = Field(..., min_length=1, description="Text prompt")70    seed: Optional[int] = Field(None, ge=0, le=MAX_SEED, description="RNG seed. null = random; the one used is returned.")717273def baked() -> dict:74    return {"size": SIZE, "num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE}757677@app.get("/health")78def health():79    if pipe is None:80        raise HTTPException(503, "Model not loaded")81    return {"status": "ok", "model": SERVED_MODEL_ID}828384@app.get("/v1/config")85def config():86    return {87        "model": SERVED_MODEL_ID,88        "params": baked(),89        "request_fields": {90            "prompt": {"required": True, "description": "Text prompt"},91            "seed": {"required": False, "default": None, "description": "RNG seed; null = random"},92        },93    }949596@app.post("/v1/images/generations")97def generate(req: ImageRequest):98    if pipe is None:99        raise HTTPException(503, "Model not loaded")100    seed = req.seed if req.seed is not None else secrets.randbelow(2**31)101    t0 = time.time()102    with lock:103        gen = torch.Generator("cuda").manual_seed(seed)104        images = pipe(105            prompt=req.prompt,106            width=WIDTH, height=HEIGHT,107            num_inference_steps=NUM_INFERENCE_STEPS,108            guidance_scale=GUIDANCE_SCALE,109            num_images_per_prompt=1,110            generator=gen,111        )[0]112    elapsed = time.time() - t0113    print(f"Generated {len(images)} image(s) in {elapsed:.1f}s (seed {seed})", flush=True)114    data = []115    for img in images:116        buf = io.BytesIO()117        img.save(buf, format="PNG")118        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})119    return JSONResponse({120        "created": int(time.time()),121        "model": SERVED_MODEL_ID,122        "seed": seed,123        "elapsed_seconds": round(elapsed, 2),124        "data": data,125        "params_used": {**baked(), "seed": seed},126    })127128129def main() -> None:130    try:131        load_model()132    except Exception:133        traceback.print_exc()134        status("error", detail="model load failed; see container logs")135        raise136    status("ready")137    uvicorn.run(app, host="0.0.0.0", port=PORT)  # the gate reaches us over the capsule network138139140if __name__ == "__main__":141    main()