"""Gradio demo for lodestones/Kromoss — full fine-tune of MOSS-Image-Base (text-to-image) with a bundled 8-step turbo LoRA. Kromoss keeps MOSS-Image's architecture, Qwen3.5 text encoder and FLUX.2 VAE, and mainly fixes the Base model's blurriness. Two sampling routes are exposed: - Turbo (default): the bundled `8-steps-dmd-lora.safetensors` applied UNMERGED via forward hooks on the bf16 transformer — 8 steps, CFG 1. - Full: 50 steps, CFG 3, negative branch with block 22 skipped (the model card's full-quality recipe). The fp32 LoRA delta is close to bf16's rounding step, so it must never be merged into the bf16 weights — hence the unmerged hooks at strength 1. """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import random import spaces # ZeroGPU: import before torch / diffusers / transformers. import gradio as gr import torch import torch.nn.functional as F from diffusers import DiffusionPipeline from safetensors.torch import load_file from huggingface_hub import hf_hub_download MODEL_REPO = "lodestones/Kromoss" LORA_FILE = "8-steps-dmd-lora.safetensors" DEFAULT_STRENGTH = 1.0 MAX_SEED = 2**31 - 1 # ---------------------------------------------------------------- pipeline print("[boot] loading Kromoss pipeline (bf16 transformer + text encoder, fp32 VAE) ...", flush=True) pipe = DiffusionPipeline.from_pretrained( MODEL_REPO, custom_pipeline="pipeline_moss_image", trust_remote_code=True, torch_dtype={"default": torch.bfloat16, "vae": torch.float32}, ) pipe = pipe.to("cuda") pipe.set_progress_bar_config(disable=True) # ---------------------------------------------------------------- LoRA # The turbo delta is fp32 and below bf16's rounding step, so it is applied # UNMERGED via forward hooks on every block's proj q/k/v/o/ffn_up_gate/ # ffn_down linears: y = W x + strength * (alpha/rank) * up(down(x)). # A global `LORA_STRENGTH` lets the UI dial it to 0 to fall back to the # plain fine-tune (then use the 50-step, CFG-3 route). print("[boot] applying 8-step turbo LoRA (fp32, unmerged hooks) ...", flush=True) _lora_path = hf_hub_download(MODEL_REPO, LORA_FILE) _lora = load_file(_lora_path) _modules = [] for layer_index, layer in enumerate(pipe.transformer.layers): for name in ("q", "k", "v", "o", "ffn_up_gate", "ffn_down"): module = layer.proj[name] prefix = f"diffusion_model.layers.{layer_index}.proj.{name}" down = _lora[f"{prefix}.lora_down.weight"] up = _lora[f"{prefix}.lora_up.weight"] alpha = float(_lora[f"{prefix}.alpha"].reshape(-1)[0]) scale = alpha / down.shape[0] # Register as module buffers and move to cuda so ZeroGPU's startup # packing ships them to the GPU worker together with the model. module.register_buffer("lora_down", down.float().to("cuda"), persistent=False) module.register_buffer("lora_up", up.float().to("cuda"), persistent=False) module.register_buffer( "lora_scale", torch.full((), scale, dtype=torch.float32).to("cuda"), persistent=False ) def _make_hook(): def hook(module, inputs, output): from kromoss_state import LORA_STRENGTH if LORA_STRENGTH["value"] == 0.0: return output x = inputs[0].float() delta = F.linear(F.linear(x, module.lora_down), module.lora_up) return output + (module.lora_scale * LORA_STRENGTH["value"] * delta).to(output.dtype) return hook module.register_forward_hook(_make_hook()) _modules.append(module) assert len(_modules) == 252, f"expected 252 LoRA modules, wired {len(_modules)}" del _modules _lora.clear() pipe.transformer.eval().requires_grad_(False) def _estimate(prompt, height, width, num_steps, *args, **kwargs): """ZeroGPU budget, calibrated against measured runs on this Space. Turbo route (CFG 1, one forward per step), warm worker: 1024x1024@8 = 7.5 s, 1360x768@8 = 7.6 s, 1536x1024@8 = 11.6 s, 2048x2048@8 ~ 33 s (sibling Space, same architecture). First call after a cold start adds ~6 s of weight streaming. The full route runs CFG 3 (two forwards per step): measured 768x768@50 = 28 s. Formula: 8 s fixed + 0.85 s per forward per megapixel, forwards doubled above 12 steps (the CFG regime), capped at 240 s. """ pixels = max(256, int(height)) * max(256, int(width)) steps = max(1, min(int(num_steps), 50)) passes = steps if steps <= 12 else 2 * steps # CFG 3 doubles the forwards return min(240, int(8 + passes * (pixels / 1_000_000) * 0.85)) @spaces.GPU(duration=_estimate, size="large") def generate( prompt: str, height: int = 1024, width: int = 1024, num_steps: int = 8, seed: int = 0, randomize_seed: bool = True, lora_strength: float = DEFAULT_STRENGTH, negative_prompt: str = "", use_turbo_cfg_fix: bool = True, progress=gr.Progress(track_tqdm=True), ): """Generate an image from a text prompt with Kromoss. Args: prompt: Text prompt. Plain natural language, danbooru-style tags and JSON scene-graph layouts are all supported (as in MOSS-Image). height: Image height in pixels, rounded down to a multiple of 16. width: Image width in pixels, rounded down to a multiple of 16. num_steps: Euler steps. 8 with the turbo LoRA at strength 1 (default); 50 with LoRA strength 0 for the full CFG-3 recipe. seed: RNG seed for the initial noise. randomize_seed: Draw a fresh random seed on every run. lora_strength: Turbo LoRA weight (1 = the distilled recipe with 8 steps, CFG 1; 0 = the plain Kromoss fine-tune, use 50 steps and CFG 3). negative_prompt: Only used when the LoRA is off (CFG != 1). use_turbo_cfg_fix: Fixes guidance_scale=1 when the LoRA is active and 3.0 with negative-branch layer 22 skipping when it is not. """ height = max(256, int(height) // 16 * 16) width = max(256, int(width) // 16 * 16) num_steps = max(1, min(int(num_steps), 50)) if randomize_seed or seed < 0: seed = random.randint(0, MAX_SEED) seed = int(seed) strength = float(lora_strength) import kromoss_state kromoss_state.LORA_STRENGTH["value"] = strength try: if use_turbo_cfg_fix and strength > 0: guidance_scale, uncond_skip_layers = 1.0, None else: guidance_scale, uncond_skip_layers = 3.0, [22] generator = torch.Generator("cpu").manual_seed(seed) images = pipe( prompt=prompt, height=height, width=width, num_inference_steps=num_steps, guidance_scale=guidance_scale, negative_prompt=negative_prompt or "", uncond_skip_layers=uncond_skip_layers, generator=generator, ).images finally: kromoss_state.LORA_STRENGTH["value"] = DEFAULT_STRENGTH return images[0], seed CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } main { max-width: 1100px !important; margin: 0 auto !important; } """ with gr.Blocks(title="Kromoss") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # Kromoss · text-to-image A full fine-tune of [MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) (5B DiT, Qwen3.5 text encoder, FLUX.2 VAE) that mainly fixes the Base model's blurriness. By default this demo uses the bundled **8-step turbo LoRA** (CFG 1); scroll to Advanced to switch to the full 50-step · CFG 3 recipe. [Model card](https://huggingface.co/lodestones/Kromoss) · [MOSS-Image](https://github.com/OpenMOSS/MOSS-Image) """ ) prompt = gr.Textbox( label="Prompt", lines=4, max_lines=12, placeholder="Describe the image — natural language, danbooru-style tags, or JSON scene-graph layouts all work…", scale=4, ) run = gr.Button("Generate", variant="primary", scale=1) output = gr.Image(label="Generated image", type="pil") seed_used = gr.Number(label="Seed used", precision=0, interactive=False) with gr.Accordion("Advanced settings", open=False): with gr.Row(): height = gr.Slider(256, 2048, value=1024, step=16, label="Height") width = gr.Slider(256, 2048, value=1024, step=16, label="Width") with gr.Row(): num_steps = gr.Slider(1, 50, value=8, step=1, label="Steps", info="8 with the turbo LoRA; 50 with LoRA strength 0") lora_strength = gr.Slider(0.0, 1.0, value=DEFAULT_STRENGTH, step=0.05, label="Turbo LoRA strength", info="1 = 8-step distilled recipe; 0 = plain fine-tune (50 steps, CFG 3)") with gr.Row(): seed = gr.Number(label="Seed", value=0, precision=0, minimum=-1, maximum=MAX_SEED) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) use_turbo_cfg_fix = gr.Checkbox( label="Auto CFG: 1 with LoRA, 3 (skip layer 22) without", value=True, info="Uncheck to force your own guidance via the negative prompt branch", ) negative_prompt = gr.Textbox( label="Negative prompt (only used when the LoRA is off)", value="", lines=2, ) gr.Markdown("## Examples") gr.Examples( examples=[ ["a tabby cat holding a sign that says KROMOSS", 1024, 1024], ["general, 1girl, best quality, most attractive face, cosplay, amiya (arknights), with \"Rhodes Island\" on her clothes, carring a sword", 1024, 1024], ["A cinematic, photorealistic nighttime scene of a small neighborhood convenience store on a quiet, rain-soaked street. The warmly illuminated storefront occupies the right half of the frame, while the street recedes into soft mist on the left. A lone young adult wearing a charcoal coat stands just outside the glow, looking back over their shoulder.", 1024, 1024], ["A candid classroom photograph, slightly off-center, with a worn dark green chalkboard filling most of the frame. Uneven window light, soft shadows, chalk dust, eraser streaks, subtle camera grain. Visible board frame and chalk tray, a blurred desk corner in the foreground.", 768, 1360], ["多层折叠式纸雕立体绘本，放在一张书桌上，背景纯净突出主题，比例 3:2 横版。翻开的跨页呈现「哪吒闹海」的场景：脚踏风火轮的少年、翻涌的层叠海浪与龙宫剪影，所有元素皆可精细折叠组合，呈现逼真细腻的纸张折叠质感与多层景深。构图统一采用正面视角，整体视觉梦幻唯美，色彩缤纷绚丽，充满奇幻生动的故事氛围。", 1536, 1024], ], inputs=[prompt, height, width], outputs=[output, seed_used], fn=generate, cache_examples=True, cache_mode="lazy", ) run.click( fn=generate, inputs=[prompt, height, width, num_steps, seed, randomize_seed, lora_strength, negative_prompt, use_turbo_cfg_fix], outputs=[output, seed_used], api_name="generate", concurrency_limit=1, ) demo.queue(max_size=8) demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)