1"""OpenAI-style /v1/images/generations server for Kromoss + the 8-step DMD2 LoRA (m64 capsule).23Ported from an earlier standalone server, run there with MOSS_LORA=dmd2.4Differences from the original:5- weights (and the LoRA, which ships in the same repo) load from /weights/kromoss, an6  m64 artifact mount; port from M64_PORT;7- only the DMD2 variant: size, steps, guidance and shift are baked constants (the capsule's8  [params]); a request carries only `prompt` and `seed` (the m64 front rejects9  anything else with a 400);10- progress is reported as JSON lines appended to M64_STATUS.1112Kept exactly from the original (and the model card): the pipeline is the repo's own13pipeline_moss_image.py via trust_remote_code; transformer and text encoder bf16, VAE fp32;14the scheduler is rebuilt from its own config (flow-matching Euler, shift 5); the LoRA is15applied UNMERGED through forward hooks, because its delta is close to bf16's rounding step16and merging would lose most of it; the initial noise comes from a CPU generator; under17CFG 1 there is no negative branch, so uncond_skip_layers is None.18"""1920import base6421import io22import json23import os24import secrets25import threading26import time27import traceback28from typing import Optional2930import torch31import uvicorn32from fastapi import FastAPI, HTTPException33from fastapi.responses import JSONResponse34from pydantic import BaseModel, Field35from safetensors.torch import load_file3637SERVED_MODEL_ID = "kromoss-dmd2"38WEIGHTS = "/weights/kromoss"39LORA_NAME = "8-steps-dmd2-lora"40LORA_PATH = os.path.join(WEIGHTS, LORA_NAME + ".safetensors")41PORT = int(os.environ["M64_PORT"])42STATUS = os.environ["M64_STATUS"]4344# Baked generation settings: the README's DMD2 row. Keep in sync with [params] in capsule.toml.45SIZE = "1024x1024"46NUM_INFERENCE_STEPS = 847GUIDANCE_SCALE = 1.048SHIFT = 5.049WIDTH, HEIGHT = (int(x) for x in SIZE.split("x"))50MAX_SEED = 2**63 - 15152app = FastAPI()53pipe = None54lock = threading.Lock()  # one generation at a time on the single pipeline555657def status(phase: str, **extra) -> None:58    with open(STATUS, "a") as f:59        f.write(json.dumps({"phase": phase, **extra}) + "\n")606162def apply_lora(transformer, path: str, strength: float = 1.0) -> int:63    """Kromoss's own apply_lora: one forward hook per adapted module, never merged."""64    state = load_file(path, device="cuda")65    modules = dict(transformer.named_modules())66    n = 067    for key in [k for k in state if k.endswith(".lora_down.weight")]:68        name = key[len("diffusion_model."):-len(".lora_down.weight")]69        down, up = state[key].float(), state[key.replace("lora_down", "lora_up")].float()70        scale = strength * state[f"diffusion_model.{name}.alpha"].item() / down.shape[0]7172        def hook(module, args, out, down=down, up=up, scale=scale):73            return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype)7475        modules[name].register_forward_hook(hook)76        n += 177    return n787980def run(prompt: str, seed: int, steps: int = NUM_INFERENCE_STEPS):81    # CPU generator: the card generates the initial noise on CPU in float32.82    return pipe(83        prompt, height=HEIGHT, width=WIDTH, num_inference_steps=steps,84        guidance_scale=GUIDANCE_SCALE, negative_prompt="", num_images_per_prompt=1,85        generator=torch.Generator("cpu").manual_seed(seed),86        uncond_skip_layers=None,  # CFG 1: no negative branch at all87    ).images888990def load_model() -> None:91    global pipe92    from diffusers import DiffusionPipeline9394    status("loading", progress=0.0, detail="loading Kromoss (10 GB transformer + 17 GB Qwen3.5-9B encoder)")95    print("Loading Kromoss...", flush=True)96    t0 = time.time()97    p = DiffusionPipeline.from_pretrained(98        WEIGHTS, custom_pipeline="pipeline_moss_image", trust_remote_code=True,99        torch_dtype={"default": torch.bfloat16, "vae": torch.float32},100    )101    # The card rebuilds the scheduler from its own config before use.102    p.scheduler = type(p.scheduler).from_config(p.scheduler.config)103    shift = p.scheduler.config.get("shift") if isinstance(p.scheduler.config, dict) else getattr(p.scheduler.config, "shift", None)104    if shift != SHIFT:105        raise RuntimeError(f"scheduler shift is {shift}, expected {SHIFT}")106    status("loading", progress=0.8, detail="moving to GPU, hooking the DMD2 LoRA")107    p.to("cuda")108    hooked = apply_lora(p.transformer, LORA_PATH)109    if hooked == 0:110        raise RuntimeError("DMD2 LoRA hooked no modules")111    pipe = p112    print(f"Model loaded in {time.time() - t0:.1f}s | LoRA modules hooked: {hooked} | shift {shift}", flush=True)113114    status("warming", detail="one short warm-up generation")115    t0 = time.time()116    with lock:117        run("warm-up", 0, steps=2)118    print(f"Warm-up done in {time.time() - t0:.1f}s.", flush=True)119120121class ImageRequest(BaseModel):122    prompt: str = Field(..., min_length=1, description="Text prompt")123    seed: Optional[int] = Field(None, ge=0, le=MAX_SEED, description="RNG seed. null = random; the one used is returned.")124125126def baked() -> dict:127    return {"size": SIZE, "num_inference_steps": NUM_INFERENCE_STEPS,128            "guidance_scale": GUIDANCE_SCALE, "shift": SHIFT, "lora": LORA_NAME}129130131@app.get("/health")132def health():133    if pipe is None:134        raise HTTPException(503, "Model not loaded")135    return {"status": "ok", "model": SERVED_MODEL_ID}136137138@app.get("/v1/config")139def config():140    return {141        "model": SERVED_MODEL_ID,142        "params": baked(),143        "request_fields": {144            "prompt": {"required": True, "description": "Text prompt (Qwen3.5-9B encoder, up to 4096 tokens)"},145            "seed": {"required": False, "default": None, "description": "RNG seed; null = random"},146        },147    }148149150@app.post("/v1/images/generations")151def generate(req: ImageRequest):152    if pipe is None:153        raise HTTPException(503, "Model not loaded")154    seed = req.seed if req.seed is not None else secrets.randbelow(2**31)155    t0 = time.time()156    with lock:157        images = run(req.prompt, seed)158    elapsed = time.time() - t0159    print(f"Generated {len(images)} image(s) in {elapsed:.1f}s (seed {seed})", flush=True)160    data = []161    for img in images:162        buf = io.BytesIO()163        img.save(buf, format="PNG")164        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})165    return JSONResponse({166        "created": int(time.time()),167        "model": SERVED_MODEL_ID,168        "seed": seed,169        "elapsed_seconds": round(elapsed, 2),170        "data": data,171        "params_used": {**baked(), "seed": seed},172    })173174175def main() -> None:176    try:177        load_model()178    except Exception:179        traceback.print_exc()180        status("error", detail="model load failed; see container logs")181        raise182    status("ready")183    uvicorn.run(app, host="0.0.0.0", port=PORT)  # the gate reaches us over the capsule network184185186if __name__ == "__main__":187    main()