Spaces:
Running on Zero
Running on Zero
| """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)) | |
| 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) |