"""OpenAI-compatible /v1/images/generations front end for the MOSS-Image family. ONE SERVER, SEVERAL CHECKPOINTS. MOSS-Image-Base, MOSS-Image-SFT (OpenMOSS-Team) and Kromoss (lodestones) share an architecture, an inference pipeline, a text encoder, a VAE and a repo layout, so which one is served is purely MODEL_DIR. Kromoss is a full fine-tune of MOSS-Image-Base at 1024 px — same architecture, pipeline, text encoder (Qwen3.5-9B) and VAE (the FLUX.2 SMALL decoder) as MOSS-Image-SFT, published as the EMA weights at update 11000 over 704k training images. Author-measured on PartiPrompts (100 prompts, CLIP ViT-L/14): 27.90 against MOSS-Image-Base's 27.19 at 50 steps, and 28.13 with the DMD2 LoRA. THREE SAMPLER CONTRACTS OVER ONE SET OF WEIGHTS, selected by MOSS_LORA, straight from the README's settings table: "" full model 50 steps, CFG 3, negative branch skipping layer 22 dmd2 + LoRA 8 steps, CFG 1, NO negative branch <- the author's pick dmd + LoRA 8 steps, CFG 1, NO negative branch base-dmd2 + LoRA 8 steps, CFG 1, NO negative branch, OWN SAMPLER (below) Flow-matching Euler with shift 5 for the first three; shift is not a request field because it already lives in the scheduler config, and the pipeline reads it from there. base-dmd2 is data-archetype/moss_image_base_dmd2: a separate 8-step DMD2 LoRA (rank 128, 256 on ffn_up_gate) distilled on MOSS-Image-Base, shipped as a PEFT adapter. Its card is explicit that it was trained for ONE sampler and others give worse results: at each step jump to the clean estimate and re-noise it with FRESH noise to the next level (ComfyUI's lcm), over 8 evenly spaced times under flow shift e^1.5 = 4.4817 (not 5). FreshNoiseScheduler below is the card's class; the re-noise draws from a per-request CUDA generator seeded from the request seed, so a seed still reproduces an image. uncond_skip_layers IS NOT A NEGATIVE PROMPT KNOB. It defaults to (22,) and REPLACES CFG's negative-prompt prediction with a skipped-block prediction of the positive prompt, which is why the full model still costs two forwards per step while ignoring what you put in negative_prompt. Pass null or [] to get an ordinary negative branch back. Under the LoRAs, guidance_scale is 1.0 so there is no second branch at all and the field is moot — the server sets it to null for them, matching the README. THE LoRAs ARE APPLIED UNMERGED, VIA FORWARD HOOKS, using the author's own apply_lora from the model card. That is not a style choice: the card is explicit that the delta is close to bf16's rounding step, so merging it into the bf16 weights would quietly lose most of it. Both are rank 128 / alpha 128 fp32 with kohya keys, and only one is used at a time. The text encoder is Qwen3.5-9B loaded as a transformers Qwen3_5Model — a PIPELINE COMPONENT, not a served LLM, so this rig's vLLM gap on the qwen3_5 architecture (the wall q-judger and the FLUX 3 prompt upsampler both hit) never applies here. Dtypes follow the card exactly: transformer and text encoder in bfloat16, VAE in float32, and the initial noise generated on CPU in float32. LICENSE: Apache-2.0 for both the weights and the bundled inference code (pipeline_moss_image.py, from MOSS-Image-SFT), and not gated — one of the genuinely unencumbered image models on this rig, unlike flux3-klein (non-commercial) or lens (MIT but research-only). Defaults are hardcoded below and echoed per call under "params_used"; query them via GET /v1/config. """ import base64, io, os, threading, time from typing import List, Optional import torch from fastapi import FastAPI, HTTPException from fastapi.responses import JSONResponse from pydantic import BaseModel, Field from safetensors.torch import load_file import uvicorn from diffusers import DiffusionPipeline # lora file, default steps, default CFG, default uncond_skip_layers VARIANTS = { "": {"lora": None, "steps": 50, "cfg": 3.0, "skip": [22], "about": "full model, negative branch skipping layer 22"}, "dmd2": {"lora": "8-steps-dmd2-lora.safetensors", "steps": 8, "cfg": 1.0, "skip": None, "about": "8-step DMD2 turbo LoRA (DMD plus a GAN stage against real images)"}, "dmd": {"lora": "8-steps-dmd-lora.safetensors", "steps": 8, "cfg": 1.0, "skip": None, "about": "8-step data-space DMD turbo LoRA"}, "base-dmd2": {"lora": "transformer/adapter_model.safetensors", "format": "peft", "sampler": "fresh-noise", "steps": 8, "cfg": 1.0, "skip": None, "about": "data-archetype's 8-step DMD2 LoRA for MOSS-Image-Base, fresh-noise sampler"}, } MOSS_LORA = os.environ.get("MOSS_LORA", "") # The turbo LoRAs ship in the Kromoss repo but were distilled from # MOSS-Image-Base, so they transfer across the family; point this at whichever # directory actually holds them. MOSS_LORA_DIR = os.environ.get("MOSS_LORA_DIR", os.environ.get("MODEL_DIR", "/model")) if MOSS_LORA not in VARIANTS: raise SystemExit(f"MOSS_LORA must be one of {list(VARIANTS)}; got {MOSS_LORA!r}") SPEC = VARIANTS[MOSS_LORA] MODEL_DIR = os.environ.get("MODEL_DIR", "/model") PORT = int(os.environ.get("PORT", "9218")) SERVED_MODEL_ID = os.environ.get("SERVED_MODEL_ID", "moss-image" + (f"-{MOSS_LORA}" if MOSS_LORA else "")) VAE_TILE = 16 # the pipeline rejects height/width that are not multiples of 16 app = FastAPI() _lock = threading.Lock() _state = {} def apply_lora(transformer, path, strength=1.0): """Kromoss's own apply_lora, verbatim in behaviour. Registers a forward hook per adapted module instead of merging, because the card states the delta is close to bf16's rounding step — merging would lose it. """ state = load_file(path, device="cuda") modules = dict(transformer.named_modules()) n = 0 for key in [k for k in state if k.endswith(".lora_down.weight")]: name = key[len("diffusion_model."):-len(".lora_down.weight")] down, up = state[key].float(), state[key.replace("lora_down", "lora_up")].float() scale = strength * state[f"diffusion_model.{name}.alpha"].item() / down.shape[0] def hook(module, args, out, down=down, up=up, scale=scale): return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype) modules[name].register_forward_hook(hook) n += 1 return n def apply_peft_lora(transformer, adapter_dir, strength=1.0): """A PEFT LoRA adapter, applied unmerged through forward hooks like apply_lora. Same reason as above: the delta is below bf16's rounding step, so it is never merged. Scale per module is lora_alpha / r with the adapter's rank/alpha patterns (rsLoRA and DoRA are refused rather than silently mis-applied). """ import json, re cfg = json.load(open(os.path.join(adapter_dir, "adapter_config.json"))) if cfg.get("use_dora") or cfg.get("use_rslora"): raise SystemExit("adapter uses DoRA/rsLoRA, which apply_peft_lora does not implement") state = load_file(os.path.join(adapter_dir, "adapter_model.safetensors"), device="cuda") modules = dict(transformer.named_modules()) def pattern(table, name, default): return next((v for k, v in table.items() if re.search(rf"(^|\.){re.escape(k)}$", name)), default) n = 0 for key in [k for k in state if k.endswith(".lora_A.weight")]: name = key[len("base_model.model."):-len(".lora_A.weight")] if name not in modules: raise SystemExit(f"adapter targets {name!r}, which the transformer does not have") down, up = state[key].float(), state[key.replace("lora_A", "lora_B")].float() r = pattern(cfg.get("rank_pattern") or {}, name, cfg["r"]) alpha = pattern(cfg.get("alpha_pattern") or {}, name, cfg["lora_alpha"]) if down.shape[0] != r: raise SystemExit(f"{name}: rank {down.shape[0]} disagrees with adapter_config ({r})") scale = strength * alpha / r def hook(module, args, out, down=down, up=up, scale=scale): return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype) modules[name].register_forward_hook(hook) n += 1 return n def _fresh_noise_scheduler(): """The base-dmd2 card's FreshNoiseScheduler, with a seedable re-noise generator.""" import math from diffusers import FlowMatchEulerDiscreteScheduler shift = math.exp(1.5) class FreshNoiseScheduler(FlowMatchEulerDiscreteScheduler): generator = None # set per request def set_timesteps(self, num_inference_steps=8, device=None): t = torch.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) sigmas = shift * t / (1 + (shift - 1) * t) super().set_timesteps(device=device, sigmas=sigmas.tolist()) def step(self, model_output, timestep, sample, return_dict=True, **kwargs): if self.step_index is None: self._init_step_index(timestep) sigma, sigma_next = self.sigmas[self.step_index], self.sigmas[self.step_index + 1] clean = sample.float() - sigma * model_output.float() if sigma_next > 0: noise = torch.randn(clean.shape, generator=self.generator, device=clean.device, dtype=clean.dtype) clean = (1 - sigma_next) * clean + sigma_next * noise self._step_index += 1 return (clean.to(sample.dtype),) return FreshNoiseScheduler(shift=1.0), shift def _load(): t0 = time.time() pipe = DiffusionPipeline.from_pretrained( MODEL_DIR, custom_pipeline="pipeline_moss_image", trust_remote_code=True, torch_dtype={"default": torch.bfloat16, "vae": torch.float32}, ) # The card rebuilds the scheduler from its own config before use. pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config) shift = (getattr(pipe.scheduler.config, "shift", None) if not isinstance(pipe.scheduler.config, dict) else pipe.scheduler.config.get("shift")) if SPEC.get("sampler") == "fresh-noise": pipe.scheduler, shift = _fresh_noise_scheduler() pipe.to("cuda") adapted = 0 if SPEC.get("format") == "peft": adapted = apply_peft_lora(pipe.transformer, os.path.dirname(os.path.join(MOSS_LORA_DIR, SPEC["lora"]))) elif SPEC["lora"]: adapted = apply_lora(pipe.transformer, os.path.join(MOSS_LORA_DIR, SPEC["lora"])) _state.update(pipe=pipe, adapted=adapted, shift=round(shift, 4) if shift else shift) print(f"moss-image loaded in {time.time()-t0:.1f}s | variant={MOSS_LORA or 'full'} " f"({SPEC['about']}) | defaults {SPEC['steps']} steps / CFG {SPEC['cfg']} / " f"skip {SPEC['skip']} | lora modules hooked: {adapted} | shift {_state['shift']}", flush=True) class ImageRequest(BaseModel): prompt: str = Field(..., description="Text prompt, encoded by Qwen3.5-9B (token budget 4096).") n: int = Field(1, description="Images per prompt (num_images_per_prompt).") 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.") 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.") num_inference_steps: Optional[int] = Field(None, description=f"Flow-matching Euler steps. Default for this variant: {SPEC['steps']}.") 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 "")) 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.") 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.") slg_scale: float = Field(0.0, description="Skip-Layer Guidance strength. 0 disables the extra forward.") max_sequence_length: int = Field(4096, description="Text encoder token budget.") seed: Optional[int] = Field(None, description="Random seed. The initial noise is generated on CPU in float32, as the card specifies.") response_format: str = Field("b64_json", description="Only b64_json supported.") def _params_doc(): out = {} for name, info in ImageRequest.model_fields.items(): d = info.default out[name] = {"default": None if d is ... else d, "description": info.description} out["num_inference_steps"]["resolved_default"] = SPEC["steps"] out["guidance_scale"]["resolved_default"] = SPEC["cfg"] out["uncond_skip_layers"]["resolved_default"] = SPEC["skip"] return out @app.get("/health") def health(): return {"status": "ok" if _state else "degraded", "model": SERVED_MODEL_ID, "loaded": bool(_state), "variant": MOSS_LORA or "full", "lora_modules": _state.get("adapted")} @app.get("/v1/models") def models(): return {"object": "list", "data": [{"id": SERVED_MODEL_ID, "object": "model"}]} @app.get("/v1/config") def config(): return {"model": SERVED_MODEL_ID, "variant": MOSS_LORA or "full", "lora": SPEC["lora"], "params": _params_doc(), "scheduler": {"type": "fresh-noise (clean estimate, re-noise; ComfyUI lcm)" if SPEC.get("sampler") == "fresh-noise" else "flow-matching Euler", "shift": _state.get("shift")}, "notes": [ f"This variant is the {SPEC['about']}: {SPEC['steps']} steps at CFG " f"{SPEC['cfg']}. The three Kromoss variants do NOT share a sampler " "contract — send only a prompt to get the right one.", "uncond_skip_layers is NOT a negative-prompt control: by default the CFG " "uncond branch is a skipped-block prediction of the POSITIVE prompt, so " "negative_prompt is ignored unless you pass uncond_skip_layers: [].", "The LoRA is applied UNMERGED through forward hooks, as the card requires " "— its delta is close to bf16's rounding step, so merging would lose it." if SPEC["lora"] else "No LoRA: this is the full 50-step model. kromoss-dmd2 is the author's " "recommended 8-step variant and scores higher on PartiPrompts.", "Text encoder is Qwen3.5-9B as a pipeline component, not a served LLM, so " "this rig's vLLM gap on the qwen3_5 architecture does not apply.", "VAE is the FLUX.2 small decoder, run in float32 while the transformer and " "text encoder run in bfloat16, exactly as the card specifies.", "LICENSE: Apache-2.0, weights and bundled inference code, not gated.", ]} @app.post("/v1/images/generations") def generate_images(req: ImageRequest): if not _state: raise HTTPException(503, "model not loaded") if req.response_format != "b64_json": raise HTTPException(400, "only response_format=b64_json is supported") try: w, h = (int(x) for x in req.size.lower().split("x")) except ValueError: raise HTTPException(400, f"size must look like 1024x1024; got {req.size}") if w <= 0 or h <= 0 or w % VAE_TILE or h % VAE_TILE: raise HTTPException(400, f"size must be positive multiples of {VAE_TILE} in both " f"dimensions; got {req.size}") steps = req.num_inference_steps if req.num_inference_steps is not None else SPEC["steps"] cfg = req.guidance_scale if req.guidance_scale is not None else SPEC["cfg"] skip = req.uncond_skip_layers if req.uncond_skip_layers is not None else SPEC["skip"] seed = req.seed if req.seed is not None else int(time.time() * 1000) % (2**31) t0 = time.time() with _lock: queued = time.time() - t0 # CPU generator in fp32: the card is explicit that the initial noise is fp32. generator = torch.Generator("cpu").manual_seed(seed) if SPEC.get("sampler") == "fresh-noise": _state["pipe"].scheduler.generator = torch.Generator("cuda").manual_seed(seed) try: out = _state["pipe"]( req.prompt, height=h, width=w, num_inference_steps=steps, guidance_scale=cfg, negative_prompt=req.negative_prompt, num_images_per_prompt=req.n, generator=generator, max_sequence_length=req.max_sequence_length, uncond_skip_layers=skip, slg_layers=req.slg_layers, slg_scale=req.slg_scale) except ValueError as exc: # The pipeline validates skip-layer indices, slg bounds and the size rule. raise HTTPException(400, str(exc)) images = out.images data = [] for image in images: buf = io.BytesIO(); image.save(buf, format="PNG") data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()}) elapsed = time.time() - t0 print(f"Generated {len(data)} image(s) {req.size} in {elapsed:.1f}s " f"({MOSS_LORA or 'full'}, {steps} steps, CFG {cfg}, skip {skip})", flush=True) return JSONResponse({ "created": int(time.time()), "data": data, "params_used": { "n": req.n, "variant": MOSS_LORA or "full", "lora": SPEC["lora"], "size": req.size, "num_inference_steps": steps, "guidance_scale": cfg, "uncond_skip_layers": skip, "slg_layers": req.slg_layers, "slg_scale": req.slg_scale, "shift": _state.get("shift"), "negative_prompt": req.negative_prompt, "negative_prompt_effective": bool(req.negative_prompt) and cfg > 1 and not skip, "max_sequence_length": req.max_sequence_length, "seed": seed, "queued_s": round(queued, 3), "elapsed_s": round(elapsed, 3), }, }) if __name__ == "__main__": print(f"{SERVED_MODEL_ID} adapter on :{PORT}", flush=True) _load() uvicorn.run(app, host="0.0.0.0", port=PORT)