"""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)


@spaces.GPU(duration=_duration)
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)