1"""OpenAI-compatible /v1/images/generations front end for the MOSS-Image family.23ONE SERVER, SEVERAL CHECKPOINTS. MOSS-Image-Base, MOSS-Image-SFT (OpenMOSS-Team) and4Kromoss (lodestones) share an architecture, an inference pipeline, a text encoder, a5VAE and a repo layout, so which one is served is purely MODEL_DIR. Kromoss is a full6fine-tune of MOSS-Image-Base at 1024 px — same architecture, pipeline, text7encoder (Qwen3.5-9B) and VAE (the FLUX.2 SMALL decoder) as MOSS-Image-SFT, published as8the EMA weights at update 11000 over 704k training images. Author-measured on9PartiPrompts (100 prompts, CLIP ViT-L/14): 27.90 against MOSS-Image-Base's 27.19 at 5010steps, and 28.13 with the DMD2 LoRA.1112THREE SAMPLER CONTRACTS OVER ONE SET OF WEIGHTS, selected by MOSS_LORA, straight from13the README's settings table:1415    ""      full model   50 steps, CFG 3, negative branch skipping layer 2216    dmd2     + LoRA        8 steps, CFG 1, NO negative branch      <- the author's pick17    dmd      + LoRA        8 steps, CFG 1, NO negative branch18    base-dmd2  + LoRA      8 steps, CFG 1, NO negative branch, OWN SAMPLER (below)1920Flow-matching Euler with shift 5 for the first three; shift is not a request field because21it already lives in the scheduler config, and the pipeline reads it from there.2223base-dmd2 is data-archetype/moss_image_base_dmd2: a separate 8-step DMD2 LoRA (rank 128,24256 on ffn_up_gate) distilled on MOSS-Image-Base, shipped as a PEFT adapter. Its card is25explicit that it was trained for ONE sampler and others give worse results: at each step26jump to the clean estimate and re-noise it with FRESH noise to the next level (ComfyUI's27lcm), over 8 evenly spaced times under flow shift e^1.5 = 4.4817 (not 5). FreshNoiseScheduler28below is the card's class; the re-noise draws from a per-request CUDA generator seeded29from the request seed, so a seed still reproduces an image.3031uncond_skip_layers IS NOT A NEGATIVE PROMPT KNOB. It defaults to (22,) and REPLACES CFG's32negative-prompt prediction with a skipped-block prediction of the positive prompt, which33is why the full model still costs two forwards per step while ignoring what you put in34negative_prompt. Pass null or [] to get an ordinary negative branch back. Under the35LoRAs, guidance_scale is 1.0 so there is no second branch at all and the field is moot —36the server sets it to null for them, matching the README.3738THE LoRAs ARE APPLIED UNMERGED, VIA FORWARD HOOKS, using the author's own apply_lora from39the model card. That is not a style choice: the card is explicit that the delta is close40to bf16's rounding step, so merging it into the bf16 weights would quietly lose most of41it. Both are rank 128 / alpha 128 fp32 with kohya keys, and only one is used at a time.4243The text encoder is Qwen3.5-9B loaded as a transformers Qwen3_5Model — a PIPELINE44COMPONENT, not a served LLM, so this rig's vLLM gap on the qwen3_5 architecture (the wall45q-judger and the FLUX 3 prompt upsampler both hit) never applies here.4647Dtypes follow the card exactly: transformer and text encoder in bfloat16, VAE in float32,48and the initial noise generated on CPU in float32.4950LICENSE: Apache-2.0 for both the weights and the bundled inference code51(pipeline_moss_image.py, from MOSS-Image-SFT), and not gated — one of the genuinely52unencumbered image models on this rig, unlike flux3-klein (non-commercial) or lens (MIT53but research-only).5455Defaults are hardcoded below and echoed per call under "params_used"; query them via56GET /v1/config.57"""58import base64, io, os, threading, time59from typing import List, Optional6061import torch62from fastapi import FastAPI, HTTPException63from fastapi.responses import JSONResponse64from pydantic import BaseModel, Field65from safetensors.torch import load_file66import uvicorn67from diffusers import DiffusionPipeline6869# lora file, default steps, default CFG, default uncond_skip_layers70VARIANTS = {71    "":     {"lora": None, "steps": 50, "cfg": 3.0, "skip": [22],72             "about": "full model, negative branch skipping layer 22"},73    "dmd2": {"lora": "8-steps-dmd2-lora.safetensors", "steps": 8, "cfg": 1.0, "skip": None,74             "about": "8-step DMD2 turbo LoRA (DMD plus a GAN stage against real images)"},75    "dmd":  {"lora": "8-steps-dmd-lora.safetensors", "steps": 8, "cfg": 1.0, "skip": None,76             "about": "8-step data-space DMD turbo LoRA"},77    "base-dmd2": {"lora": "transformer/adapter_model.safetensors", "format": "peft",78                  "sampler": "fresh-noise", "steps": 8, "cfg": 1.0, "skip": None,79                  "about": "data-archetype's 8-step DMD2 LoRA for MOSS-Image-Base, fresh-noise sampler"},80}8182MOSS_LORA = os.environ.get("MOSS_LORA", "")83# The turbo LoRAs ship in the Kromoss repo but were distilled from84# MOSS-Image-Base, so they transfer across the family; point this at whichever85# directory actually holds them.86MOSS_LORA_DIR = os.environ.get("MOSS_LORA_DIR", os.environ.get("MODEL_DIR", "/model"))87if MOSS_LORA not in VARIANTS:88    raise SystemExit(f"MOSS_LORA must be one of {list(VARIANTS)}; got {MOSS_LORA!r}")89SPEC = VARIANTS[MOSS_LORA]90MODEL_DIR = os.environ.get("MODEL_DIR", "/model")91PORT = int(os.environ.get("PORT", "9218"))92SERVED_MODEL_ID = os.environ.get("SERVED_MODEL_ID",93                                 "moss-image" + (f"-{MOSS_LORA}" if MOSS_LORA else ""))94VAE_TILE = 16   # the pipeline rejects height/width that are not multiples of 169596app = FastAPI()97_lock = threading.Lock()98_state = {}99100101def apply_lora(transformer, path, strength=1.0):102    """Kromoss's own apply_lora, verbatim in behaviour.103104    Registers a forward hook per adapted module instead of merging, because the card105    states the delta is close to bf16's rounding step — merging would lose it.106    """107    state = load_file(path, device="cuda")108    modules = dict(transformer.named_modules())109    n = 0110    for key in [k for k in state if k.endswith(".lora_down.weight")]:111        name = key[len("diffusion_model."):-len(".lora_down.weight")]112        down, up = state[key].float(), state[key.replace("lora_down", "lora_up")].float()113        scale = strength * state[f"diffusion_model.{name}.alpha"].item() / down.shape[0]114115        def hook(module, args, out, down=down, up=up, scale=scale):116            return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype)117118        modules[name].register_forward_hook(hook)119        n += 1120    return n121122123def apply_peft_lora(transformer, adapter_dir, strength=1.0):124    """A PEFT LoRA adapter, applied unmerged through forward hooks like apply_lora.125126    Same reason as above: the delta is below bf16's rounding step, so it is never merged.127    Scale per module is lora_alpha / r with the adapter's rank/alpha patterns (rsLoRA and128    DoRA are refused rather than silently mis-applied).129    """130    import json, re131    cfg = json.load(open(os.path.join(adapter_dir, "adapter_config.json")))132    if cfg.get("use_dora") or cfg.get("use_rslora"):133        raise SystemExit("adapter uses DoRA/rsLoRA, which apply_peft_lora does not implement")134    state = load_file(os.path.join(adapter_dir, "adapter_model.safetensors"), device="cuda")135    modules = dict(transformer.named_modules())136137    def pattern(table, name, default):138        return next((v for k, v in table.items() if re.search(rf"(^|\.){re.escape(k)}$", name)), default)139140    n = 0141    for key in [k for k in state if k.endswith(".lora_A.weight")]:142        name = key[len("base_model.model."):-len(".lora_A.weight")]143        if name not in modules:144            raise SystemExit(f"adapter targets {name!r}, which the transformer does not have")145        down, up = state[key].float(), state[key.replace("lora_A", "lora_B")].float()146        r = pattern(cfg.get("rank_pattern") or {}, name, cfg["r"])147        alpha = pattern(cfg.get("alpha_pattern") or {}, name, cfg["lora_alpha"])148        if down.shape[0] != r:149            raise SystemExit(f"{name}: rank {down.shape[0]} disagrees with adapter_config ({r})")150        scale = strength * alpha / r151152        def hook(module, args, out, down=down, up=up, scale=scale):153            return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype)154155        modules[name].register_forward_hook(hook)156        n += 1157    return n158159160def _fresh_noise_scheduler():161    """The base-dmd2 card's FreshNoiseScheduler, with a seedable re-noise generator."""162    import math163    from diffusers import FlowMatchEulerDiscreteScheduler164    shift = math.exp(1.5)165166    class FreshNoiseScheduler(FlowMatchEulerDiscreteScheduler):167        generator = None  # set per request168169        def set_timesteps(self, num_inference_steps=8, device=None):170            t = torch.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)171            sigmas = shift * t / (1 + (shift - 1) * t)172            super().set_timesteps(device=device, sigmas=sigmas.tolist())173174        def step(self, model_output, timestep, sample, return_dict=True, **kwargs):175            if self.step_index is None:176                self._init_step_index(timestep)177            sigma, sigma_next = self.sigmas[self.step_index], self.sigmas[self.step_index + 1]178            clean = sample.float() - sigma * model_output.float()179            if sigma_next > 0:180                noise = torch.randn(clean.shape, generator=self.generator,181                                    device=clean.device, dtype=clean.dtype)182                clean = (1 - sigma_next) * clean + sigma_next * noise183            self._step_index += 1184            return (clean.to(sample.dtype),)185186    return FreshNoiseScheduler(shift=1.0), shift187188189def _load():190    t0 = time.time()191    pipe = DiffusionPipeline.from_pretrained(192        MODEL_DIR, custom_pipeline="pipeline_moss_image", trust_remote_code=True,193        torch_dtype={"default": torch.bfloat16, "vae": torch.float32},194    )195    # The card rebuilds the scheduler from its own config before use.196    pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config)197    shift = (getattr(pipe.scheduler.config, "shift", None)198             if not isinstance(pipe.scheduler.config, dict) else pipe.scheduler.config.get("shift"))199    if SPEC.get("sampler") == "fresh-noise":200        pipe.scheduler, shift = _fresh_noise_scheduler()201    pipe.to("cuda")202    adapted = 0203    if SPEC.get("format") == "peft":204        adapted = apply_peft_lora(pipe.transformer,205                                  os.path.dirname(os.path.join(MOSS_LORA_DIR, SPEC["lora"])))206    elif SPEC["lora"]:207        adapted = apply_lora(pipe.transformer, os.path.join(MOSS_LORA_DIR, SPEC["lora"]))208    _state.update(pipe=pipe, adapted=adapted, shift=round(shift, 4) if shift else shift)209    print(f"moss-image loaded in {time.time()-t0:.1f}s | variant={MOSS_LORA or 'full'} "210          f"({SPEC['about']}) | defaults {SPEC['steps']} steps / CFG {SPEC['cfg']} / "211          f"skip {SPEC['skip']} | lora modules hooked: {adapted} | shift {_state['shift']}",212          flush=True)213214215class ImageRequest(BaseModel):216    prompt: str = Field(..., description="Text prompt, encoded by Qwen3.5-9B (token budget 4096).")217    n: int = Field(1, description="Images per prompt (num_images_per_prompt).")218    size: str = Field("1024x1024", description=f"WxH in pixels, both positive multiples of {VAE_TILE}; the pipeline rejects anything else. The fine-tune was done at 1024 px.")219    negative_prompt: str = Field("", description="Negative prompt. ONLY TAKES EFFECT when guidance_scale > 1 AND uncond_skip_layers is null/empty — by default the uncond branch is a skipped-block prediction of the POSITIVE prompt, not this text.")220    num_inference_steps: Optional[int] = Field(None, description=f"Flow-matching Euler steps. Default for this variant: {SPEC['steps']}.")221    guidance_scale: Optional[float] = Field(None, description=f"CFG scale. Default for this variant: {SPEC['cfg']}." + (" The turbo LoRAs are distilled for CFG 1; raising it is off-contract." if MOSS_LORA else ""))222    uncond_skip_layers: Optional[List[int]] = Field(None, description=f"Zero-based block indices whose skipped prediction replaces the CFG negative branch. Default for this variant: {SPEC['skip']}. Pass [] for a conventional negative branch driven by negative_prompt.")223    slg_layers: Optional[List[int]] = Field(None, description="Skip-Layer Guidance block indices. Adds a third forward per step and the correction slg_scale * (cond - cond_skip). Off by default.")224    slg_scale: float = Field(0.0, description="Skip-Layer Guidance strength. 0 disables the extra forward.")225    max_sequence_length: int = Field(4096, description="Text encoder token budget.")226    seed: Optional[int] = Field(None, description="Random seed. The initial noise is generated on CPU in float32, as the card specifies.")227    response_format: str = Field("b64_json", description="Only b64_json supported.")228229230def _params_doc():231    out = {}232    for name, info in ImageRequest.model_fields.items():233        d = info.default234        out[name] = {"default": None if d is ... else d, "description": info.description}235    out["num_inference_steps"]["resolved_default"] = SPEC["steps"]236    out["guidance_scale"]["resolved_default"] = SPEC["cfg"]237    out["uncond_skip_layers"]["resolved_default"] = SPEC["skip"]238    return out239240241@app.get("/health")242def health():243    return {"status": "ok" if _state else "degraded", "model": SERVED_MODEL_ID,244            "loaded": bool(_state), "variant": MOSS_LORA or "full",245            "lora_modules": _state.get("adapted")}246247248@app.get("/v1/models")249def models():250    return {"object": "list", "data": [{"id": SERVED_MODEL_ID, "object": "model"}]}251252253@app.get("/v1/config")254def config():255    return {"model": SERVED_MODEL_ID, "variant": MOSS_LORA or "full",256            "lora": SPEC["lora"], "params": _params_doc(),257            "scheduler": {"type": "fresh-noise (clean estimate, re-noise; ComfyUI lcm)" if SPEC.get("sampler") == "fresh-noise" else "flow-matching Euler", "shift": _state.get("shift")},258            "notes": [259                f"This variant is the {SPEC['about']}: {SPEC['steps']} steps at CFG "260                f"{SPEC['cfg']}. The three Kromoss variants do NOT share a sampler "261                "contract — send only a prompt to get the right one.",262                "uncond_skip_layers is NOT a negative-prompt control: by default the CFG "263                "uncond branch is a skipped-block prediction of the POSITIVE prompt, so "264                "negative_prompt is ignored unless you pass uncond_skip_layers: [].",265                "The LoRA is applied UNMERGED through forward hooks, as the card requires "266                "— its delta is close to bf16's rounding step, so merging would lose it."267                if SPEC["lora"] else268                "No LoRA: this is the full 50-step model. kromoss-dmd2 is the author's "269                "recommended 8-step variant and scores higher on PartiPrompts.",270                "Text encoder is Qwen3.5-9B as a pipeline component, not a served LLM, so "271                "this rig's vLLM gap on the qwen3_5 architecture does not apply.",272                "VAE is the FLUX.2 small decoder, run in float32 while the transformer and "273                "text encoder run in bfloat16, exactly as the card specifies.",274                "LICENSE: Apache-2.0, weights and bundled inference code, not gated.",275            ]}276277278@app.post("/v1/images/generations")279def generate_images(req: ImageRequest):280    if not _state:281        raise HTTPException(503, "model not loaded")282    if req.response_format != "b64_json":283        raise HTTPException(400, "only response_format=b64_json is supported")284    try:285        w, h = (int(x) for x in req.size.lower().split("x"))286    except ValueError:287        raise HTTPException(400, f"size must look like 1024x1024; got {req.size}")288    if w <= 0 or h <= 0 or w % VAE_TILE or h % VAE_TILE:289        raise HTTPException(400, f"size must be positive multiples of {VAE_TILE} in both "290                                 f"dimensions; got {req.size}")291292    steps = req.num_inference_steps if req.num_inference_steps is not None else SPEC["steps"]293    cfg = req.guidance_scale if req.guidance_scale is not None else SPEC["cfg"]294    skip = req.uncond_skip_layers if req.uncond_skip_layers is not None else SPEC["skip"]295    seed = req.seed if req.seed is not None else int(time.time() * 1000) % (2**31)296    t0 = time.time()297298    with _lock:299        queued = time.time() - t0300        # CPU generator in fp32: the card is explicit that the initial noise is fp32.301        generator = torch.Generator("cpu").manual_seed(seed)302        if SPEC.get("sampler") == "fresh-noise":303            _state["pipe"].scheduler.generator = torch.Generator("cuda").manual_seed(seed)304        try:305            out = _state["pipe"](306                req.prompt, height=h, width=w, num_inference_steps=steps,307                guidance_scale=cfg, negative_prompt=req.negative_prompt,308                num_images_per_prompt=req.n, generator=generator,309                max_sequence_length=req.max_sequence_length,310                uncond_skip_layers=skip, slg_layers=req.slg_layers,311                slg_scale=req.slg_scale)312        except ValueError as exc:313            # The pipeline validates skip-layer indices, slg bounds and the size rule.314            raise HTTPException(400, str(exc))315        images = out.images316317    data = []318    for image in images:319        buf = io.BytesIO(); image.save(buf, format="PNG")320        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})321322    elapsed = time.time() - t0323    print(f"Generated {len(data)} image(s) {req.size} in {elapsed:.1f}s "324          f"({MOSS_LORA or 'full'}, {steps} steps, CFG {cfg}, skip {skip})", flush=True)325    return JSONResponse({326        "created": int(time.time()),327        "data": data,328        "params_used": {329            "n": req.n, "variant": MOSS_LORA or "full", "lora": SPEC["lora"],330            "size": req.size, "num_inference_steps": steps, "guidance_scale": cfg,331            "uncond_skip_layers": skip, "slg_layers": req.slg_layers,332            "slg_scale": req.slg_scale, "shift": _state.get("shift"),333            "negative_prompt": req.negative_prompt,334            "negative_prompt_effective": bool(req.negative_prompt) and cfg > 1 and not skip,335            "max_sequence_length": req.max_sequence_length, "seed": seed,336            "queued_s": round(queued, 3), "elapsed_s": round(elapsed, 3),337        },338    })339340341if __name__ == "__main__":342    print(f"{SERVED_MODEL_ID} adapter on :{PORT}", flush=True)343    _load()344    uvicorn.run(app, host="0.0.0.0", port=PORT)