1"""OpenAI-style /v1/images/generations server for Kromoss + the 8-step DMD2 LoRA.23What makes it an m64 capsule server:4- the weights (and the LoRA, which ships in the same repo) load from /weights/kromoss,5  the read-only mount of the capsule's [[artifacts]]; nothing is downloaded here;6- it starts listening on $M64_PORT only once the model is loaded; m64 polls /health7  until it answers 200;8- size, steps, guidance and shift are baked constants, the same values as [params]9  in capsule.toml; a request carries only `prompt` and `seed` (m64's front refuses10  anything else with a 400), and every answer returns them as `params_used`;11- progress goes to $M64_STATUS as JSON lines (loading, warming, ready, error), which12  `m64 status` and the /status route show.1314Kept exactly from the model card: the pipeline is the repo's own15pipeline_moss_image.py via trust_remote_code; transformer and text encoder bf16, VAE fp32;16the scheduler is rebuilt from its own config (flow-matching Euler, shift 5); the LoRA is17applied UNMERGED through forward hooks, because its delta is close to bf16's rounding step18and merging would lose most of it; the initial noise comes from a CPU generator; under19CFG 1 there is no negative branch, so uncond_skip_layers is None.20"""2122import base6423import io24import json25import os26import secrets27import threading28import time29import traceback30from typing import Optional3132import torch33import uvicorn34from fastapi import FastAPI, HTTPException35from fastapi.responses import JSONResponse36from pydantic import BaseModel, Field37from safetensors.torch import load_file3839SERVED_MODEL_ID = "kromoss-dmd2"40WEIGHTS = "/weights/kromoss"41LORA_NAME = "8-steps-dmd2-lora"42LORA_PATH = os.path.join(WEIGHTS, LORA_NAME + ".safetensors")43PORT = int(os.environ["M64_PORT"])44STATUS = os.environ["M64_STATUS"]4546# Baked generation settings: the README's DMD2 row. Keep in sync with [params] in capsule.toml.47SIZE = "1024x1024"48NUM_INFERENCE_STEPS = 849GUIDANCE_SCALE = 1.050SHIFT = 5.051WIDTH, HEIGHT = (int(x) for x in SIZE.split("x"))52MAX_SEED = 2**63 - 15354app = FastAPI()55pipe = None56lock = threading.Lock()  # one generation at a time on the single pipeline575859def status(phase: str, **extra) -> None:60    with open(STATUS, "a") as f:61        f.write(json.dumps({"phase": phase, **extra}) + "\n")626364def apply_lora(transformer, path: str, strength: float = 1.0) -> int:65    """Kromoss's own apply_lora: one forward hook per adapted module, never merged."""66    state = load_file(path, device="cuda")67    modules = dict(transformer.named_modules())68    n = 069    for key in [k for k in state if k.endswith(".lora_down.weight")]:70        name = key[len("diffusion_model."):-len(".lora_down.weight")]71        down, up = state[key].float(), state[key.replace("lora_down", "lora_up")].float()72        scale = strength * state[f"diffusion_model.{name}.alpha"].item() / down.shape[0]7374        def hook(module, args, out, down=down, up=up, scale=scale):75            return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype)7677        modules[name].register_forward_hook(hook)78        n += 179    return n808182def run(prompt: str, seed: int, steps: int = NUM_INFERENCE_STEPS):83    # CPU generator: the card generates the initial noise on CPU in float32.84    return pipe(85        prompt, height=HEIGHT, width=WIDTH, num_inference_steps=steps,86        guidance_scale=GUIDANCE_SCALE, negative_prompt="", num_images_per_prompt=1,87        generator=torch.Generator("cpu").manual_seed(seed),88        uncond_skip_layers=None,  # CFG 1: no negative branch at all89    ).images909192def load_model() -> None:93    global pipe94    from diffusers import DiffusionPipeline9596    status("loading", progress=0.0, detail="loading Kromoss (10 GB transformer + 17 GB Qwen3.5-9B encoder)")97    print("Loading Kromoss...", flush=True)98    t0 = time.time()99    p = DiffusionPipeline.from_pretrained(100        WEIGHTS, custom_pipeline="pipeline_moss_image", trust_remote_code=True,101        torch_dtype={"default": torch.bfloat16, "vae": torch.float32},102    )103    # The card rebuilds the scheduler from its own config before use.104    p.scheduler = type(p.scheduler).from_config(p.scheduler.config)105    shift = p.scheduler.config.get("shift") if isinstance(p.scheduler.config, dict) else getattr(p.scheduler.config, "shift", None)106    if shift != SHIFT:107        raise RuntimeError(f"scheduler shift is {shift}, expected {SHIFT}")108    status("loading", progress=0.8, detail="moving to GPU, hooking the DMD2 LoRA")109    p.to("cuda")110    hooked = apply_lora(p.transformer, LORA_PATH)111    if hooked == 0:112        raise RuntimeError("DMD2 LoRA hooked no modules")113    pipe = p114    print(f"Model loaded in {time.time() - t0:.1f}s | LoRA modules hooked: {hooked} | shift {shift}", flush=True)115116    status("warming", detail="one short warm-up generation")117    t0 = time.time()118    with lock:119        run("warm-up", 0, steps=2)120    print(f"Warm-up done in {time.time() - t0:.1f}s.", flush=True)121122123class ImageRequest(BaseModel):124    prompt: str = Field(..., min_length=1, description="Text prompt")125    seed: Optional[int] = Field(None, ge=0, le=MAX_SEED, description="RNG seed. null = random; the one used is returned.")126127128def baked() -> dict:129    return {"size": SIZE, "num_inference_steps": NUM_INFERENCE_STEPS,130            "guidance_scale": GUIDANCE_SCALE, "shift": SHIFT, "lora": LORA_NAME}131132133@app.get("/health")134def health():135    if pipe is None:136        raise HTTPException(503, "Model not loaded")137    return {"status": "ok", "model": SERVED_MODEL_ID}138139140@app.get("/v1/config")141def config():142    return {143        "model": SERVED_MODEL_ID,144        "params": baked(),145        "request_fields": {146            "prompt": {"required": True, "description": "Text prompt (Qwen3.5-9B encoder, up to 4096 tokens)"},147            "seed": {"required": False, "default": None, "description": "RNG seed; null = random"},148        },149    }150151152@app.post("/v1/images/generations")153def generate(req: ImageRequest):154    if pipe is None:155        raise HTTPException(503, "Model not loaded")156    seed = req.seed if req.seed is not None else secrets.randbelow(2**31)157    t0 = time.time()158    with lock:159        images = run(req.prompt, seed)160    elapsed = time.time() - t0161    print(f"Generated {len(images)} image(s) in {elapsed:.1f}s (seed {seed})", flush=True)162    data = []163    for img in images:164        buf = io.BytesIO()165        img.save(buf, format="PNG")166        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})167    return JSONResponse({168        "created": int(time.time()),169        "model": SERVED_MODEL_ID,170        "seed": seed,171        "elapsed_seconds": round(elapsed, 2),172        "data": data,173        "params_used": {**baked(), "seed": seed},174    })175176177def main() -> None:178    try:179        load_model()180    except Exception:181        traceback.print_exc()182        status("error", detail="model load failed; see container logs")183        raise184    status("ready")185    uvicorn.run(app, host="0.0.0.0", port=PORT)  # the gate reaches us over the capsule network186187188if __name__ == "__main__":189    main()