Spaces:
Running on Zero
Running on Zero
| """Text-to-image demo for the MOSS-Image weights published at Comfy-Org/MOSS-Image. | |
| Comfy-Org/MOSS-Image is a ComfyUI repack of OpenMOSS MOSS-Image. This Space | |
| runs the matching diffusers checkpoint (MOSS-Image-SFT): bf16 transformer and | |
| Qwen3.5-9B text encoder, float32 FLUX.2 VAE, flow-matching Euler, CFG 3, and | |
| negative-branch skip of layer 22. | |
| """ | |
| import secrets | |
| import spaces | |
| import torch | |
| import gradio as gr | |
| from diffusers import DiffusionPipeline | |
| MODEL_ID = "OpenMOSS-Team/MOSS-Image-SFT" | |
| WEIGHT_PACK = "Comfy-Org/MOSS-Image" | |
| RESOLUTIONS = { | |
| "1024 × 1024": (1024, 1024), | |
| "1152 × 896": (1152, 896), | |
| "896 × 1152": (896, 1152), | |
| "1280 × 720": (1280, 720), | |
| "720 × 1280": (720, 1280), | |
| "1344 × 768": (1344, 768), | |
| "768 × 1344": (768, 1344), | |
| } | |
| pipe = DiffusionPipeline.from_pretrained( | |
| MODEL_ID, | |
| custom_pipeline="pipeline_moss_image", | |
| trust_remote_code=True, | |
| torch_dtype={"default": torch.bfloat16, "vae": torch.float32}, | |
| ) | |
| pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config, shift=5.0) | |
| pipe.to("cuda") | |
| def _duration(*args, **kwargs): | |
| steps = kwargs.get("steps", args[3] if len(args) > 3 else 50) | |
| try: | |
| steps_i = int(steps) | |
| except (TypeError, ValueError): | |
| steps_i = 50 | |
| # Two transformer forwards per step when guidance is on, plus encode/decode. | |
| return min(180, 40 + max(1, steps_i) * 2) | |
| def generate( | |
| prompt: str, | |
| negative_prompt: str = "", | |
| resolution: str = "1024 × 1024", | |
| steps: float = 50, | |
| guidance: float = 3.0, | |
| shift: float = 5.0, | |
| seed: float = -1, | |
| progress: gr.Progress = gr.Progress(track_tqdm=True), | |
| ): | |
| """Generate an image with MOSS-Image SFT. | |
| The checkpoint is the diffusers release of the weights packaged for ComfyUI | |
| at Comfy-Org/MOSS-Image (moss_image_sft, Qwen3.5-9B, FLUX.2 VAE). | |
| Pass seed -1 for a random seed. Width and height must be multiples of 16. | |
| """ | |
| if not (prompt or "").strip(): | |
| raise gr.Error("Enter a prompt.") | |
| width, height = RESOLUTIONS.get(resolution, (1024, 1024)) | |
| steps_i = int(steps) | |
| guidance_f = float(guidance) | |
| shift_f = float(shift) | |
| seed_i = int(seed) | |
| if steps_i < 1 or steps_i > 80: | |
| raise gr.Error("Steps must be between 1 and 80.") | |
| if guidance_f < 0: | |
| raise gr.Error("Guidance scale must be nonnegative.") | |
| if shift_f <= 0: | |
| raise gr.Error("Shift must be positive.") | |
| if seed_i == -1: | |
| seed_i = secrets.randbelow(2**31) | |
| elif not 0 <= seed_i < 2**31: | |
| raise gr.Error("Seed must be -1 (random) or an integer from 0 to 2^31 - 1.") | |
| pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config, shift=shift_f) | |
| try: | |
| image = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt or "", | |
| height=height, | |
| width=width, | |
| num_inference_steps=steps_i, | |
| guidance_scale=guidance_f, | |
| uncond_skip_layers=[22], | |
| generator=torch.Generator("cpu").manual_seed(seed_i), | |
| ).images[0] | |
| except torch.cuda.OutOfMemoryError as exc: | |
| torch.cuda.empty_cache() | |
| raise gr.Error("GPU memory ran out. Try a smaller resolution or fewer steps.") from exc | |
| return image, seed_i | |
| CSS = """ | |
| #col-container { max-width: 980px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="MOSS-Image") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| f""" | |
| # MOSS-Image | |
| Text-to-image demo for [{WEIGHT_PACK}](https://huggingface.co/{WEIGHT_PACK}), | |
| the ComfyUI weight pack of [MOSS-Image-SFT](https://huggingface.co/{MODEL_ID}). | |
| Default sampling matches the model card: 50 steps, CFG 3, shift 5, skip layer 22 on the negative branch. | |
| """.strip() | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| lines=5, | |
| placeholder="Describe the image. Layout JSON and rendered text are supported.", | |
| ) | |
| with gr.Accordion("Advanced settings", open=False): | |
| negative = gr.Textbox(label="Negative prompt", lines=2) | |
| resolution = gr.Dropdown( | |
| label="Resolution", | |
| choices=list(RESOLUTIONS), | |
| value="1024 × 1024", | |
| ) | |
| with gr.Row(): | |
| steps = gr.Slider(1, 80, value=50, step=1, label="Steps") | |
| guidance = gr.Slider(0, 10, value=3.0, step=0.1, label="Guidance") | |
| with gr.Row(): | |
| shift = gr.Slider(0.5, 12, value=5.0, step=0.1, label="Shift") | |
| seed = gr.Number(label="Seed (−1 = random)", value=-1, precision=0) | |
| run = gr.Button("Generate", variant="primary") | |
| with gr.Column(scale=2): | |
| image = gr.Image(label="Result", type="pil", height=512) | |
| used_seed = gr.Number(label="Seed used", precision=0) | |
| gr.Examples( | |
| examples=[ | |
| [ | |
| "A wide view of a refrigerator beside a kitchen counter with a fruit bowl. " | |
| "A small sticky note on the refrigerator door is handwritten in Chinese:\n" | |
| "\"记得买牛奶和鸡蛋！\n牛奶两盒，鸡蛋一打。\n顺便带一袋面包。\"" | |
| ], | |
| [ | |
| "A sunlit ceramic mug on a wooden table, the words \"MOSS Image\" painted in blue glaze, " | |
| "soft morning light, photorealistic" | |
| ], | |
| [ | |
| "A small bookstore window at dusk, rain on the glass, a poster inside reading " | |
| "\"OPEN LATE\", warm interior light" | |
| ], | |
| ], | |
| inputs=[prompt], | |
| outputs=[image, used_seed], | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| run.click( | |
| generate, | |
| inputs=[prompt, negative, resolution, steps, guidance, shift, seed], | |
| outputs=[image, used_seed], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(css=CSS, mcp_server=True) | |