app.py

from __future__ import annotations

import os

# Long image-conditioned prefixes cause large transient allocations; expandable
# segments keep the caching allocator from fragmenting into a hard failure
# (NVML_SUCCESS == r INTERNAL ASSERT FAILED in CUDACachingAllocator).
# MUST be set before torch is imported.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST be imported before torch / transformers / sensenova_u1

import random
import threading
import time

import gradio as gr
import numpy as np
import torch
from PIL import Image
from transformers import AutoConfig, AutoModel, AutoTokenizer

import sensenova_u1
from sensenova_u1.models.neo_unify.utils import smart_resize

MODEL_ID = "sensenova/SenseNova-U1.5-8B-MoT"

NORM_MEAN = (0.5, 0.5, 0.5)
NORM_STD = (0.5, 0.5, 0.5)

# U1.5 trained T2I aspect-ratio buckets (from the upstream examples/t2i/inference.py
# SUPPORTED_RESOLUTIONS table).
T2I_RESOLUTIONS: dict[str, tuple[int, int]] = {
    "1:1": (2048, 2048),
    "16:9": (2720, 1536),
    "9:16": (1536, 2720),
    "3:2": (2496, 1664),
    "2:3": (1664, 2496),
    "4:3": (2368, 1760),
    "3:4": (1760, 2368),
    "1:2": (1440, 2880),
    "2:1": (2880, 1440),
    "1:3": (1152, 3456),
    "3:1": (3456, 1152),
}

# Reference config for SenseNova-U1.5 (from the model card Quick Start):
#   cfg_scale=4.0, timestep_shift=3.0, num_steps=50
DEFAULT_CFG_SCALE = 4.0
DEFAULT_TIMESTEP_SHIFT = 3.0
DEFAULT_IMG_CFG_SCALE = 1.0
REFERENCE_NUM_STEPS = 50

# A fixed-seed A/B against the 50-step reference showed 28 steps keeps
# composition, prompt adherence and text rendering intact — the only cost is
# some micro-texture in landscape/skin — while running ~1.8x faster. Editing is
# near step-invariant (0.97 SSIM vs 50 steps). Below ~24 steps structure starts
# to smear, so the slider floor stays well above that.
DEFAULT_NUM_STEPS = 28
MIN_NUM_STEPS = 8
MAX_NUM_STEPS = REFERENCE_NUM_STEPS

# --- NCII prompt guard ------------------------------------------------------
# hfmlsoc/ncii-guard-v02 flags image-editing prompts that ask to strip, undress
# or otherwise sexualize a person in a photograph. It runs on CPU in a separate
# subprocess (see ncii_guard.py), and is consulted before any GPU worker is
# allocated, so a refused prompt costs no GPU time.
#
# Threshold: the model deliberately ships no default. On its own held-out sweep
# (eval/threshold-sweep.json, 980 prompts / 70 positives) precision climbs from
# 0.870 at 0.5 to 0.923 at 0.80 while recall stays flat at 0.857 — so anything
# below ~0.7 just donates false positives. Scores on real prompts are strongly
# bimodal (0.000 or 1.000), so the threshold only decides the obfuscated
# mid-band: character-separated evasions ("r e m o v e  h e r  d r e s s") land
# at 0.73. 0.70 catches those for the cost of one extra false positive per 980
# clean prompts vs 0.80.
GUARD_ID = "hfmlsoc/ncii-guard-v02"
GUARD_THRESHOLD = 0.70
# The classifier is trained on *edit* prompts, and the NCII harm needs a real
# photo of a real person, so only the editing path is screened. Screening plain
# text-to-image costs real false positives (e.g. "classical marble statue of a
# nude figure in a museum" scores 0.9999) for no gain in NCII coverage.
GUARD_SCREEN_TEXT_TO_IMAGE = False
# The classifier runs in a subprocess (see ncii_guard.py): loaded in this
# process it leaves CUDA driver state behind that kills every subsequent
# ZeroGPU worker at worker_init with "No CUDA GPUs are available".

# Editing output grid factor (= patch_size * merge_size = 32).
EDIT_GRID_FACTOR = 32
EDIT_TARGET_PIXELS = 2048 * 2048
# Total input-image pixel budget for the editing prefix, shared across all
# reference images. The it2i prefix forward uses eager attention, whose memory
# grows quadratically with the number of input image tokens; capping the total
# keeps a multi-image edit from blowing past the ZeroGPU slot.
EDIT_INPUT_TOTAL_MAX_PIXELS = 2048 * 2048
EDIT_INPUT_MIN_PIXELS = 512 * 512
MAX_INPUT_IMAGES = 4

MAX_SEED = 2**31 - 1


def _denorm(x: torch.Tensor) -> torch.Tensor:
    mean = torch.tensor(NORM_MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
    std = torch.tensor(NORM_STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
    return (x * std + mean).clamp(0, 1)


def _to_pil(batch: torch.Tensor) -> list[Image.Image]:
    arr = _denorm(batch.float()).permute(0, 2, 3, 1).cpu().numpy()
    arr = (arr * 255.0).round().astype(np.uint8)
    return [Image.fromarray(a) for a in arr]


def _coerce_pil(img) -> Image.Image | None:
    """Best-effort conversion of one Gradio gallery entry into a PIL image.

    Returns ``None`` for anything that isn't actually an image (stray strings
    such as example/placeholder values, ``None`` slots, unreadable paths...)
    so callers can drop it instead of crashing inside the GPU worker.
    """
    if img is None:
        return None
    # gr.Gallery hands over (path_or_image, caption) tuples.
    if isinstance(img, (tuple, list)):
        if not img:
            return None
        img = img[0]
        if isinstance(img, (tuple, list)):  # nested, give up
            return None
    if isinstance(img, dict):
        img = img.get("image") or img.get("path") or img.get("name") or img.get("url")
        if img is None:
            return None
    if isinstance(img, Image.Image):
        return img
    if isinstance(img, np.ndarray):
        try:
            return Image.fromarray(img.astype(np.uint8))
        except Exception:
            return None
    if isinstance(img, os.PathLike):
        img = os.fspath(img)
    if isinstance(img, str):
        # Only treat it as an image if it really is a readable image file.
        if not img or not os.path.isfile(img):
            return None
        try:
            loaded = Image.open(img)
            loaded.load()
            return loaded
        except Exception:
            return None
    return None


def _normalize_images(images) -> list[Image.Image]:
    """Turn whatever Gradio handed us into a clean list of PIL images.

    Runs on CPU, before the ZeroGPU worker is entered, so that a bad input is
    rejected up front instead of aborting a GPU call mid-flight.
    """
    if images is None:
        return []
    # A bare string / PIL image / array (not a list) is a single input, never
    # something to iterate over character by character.
    if isinstance(images, (str, bytes, os.PathLike, Image.Image, np.ndarray, dict)):
        items = [images]
    elif isinstance(images, (list, tuple)):
        items = list(images)
    else:
        items = [images]

    pils: list[Image.Image] = []
    dropped = 0
    for item in items:
        pil = _coerce_pil(item)
        if pil is None:
            dropped += 1
            continue
        pils.append(pil)

    if dropped and not pils:
        raise gr.Error(
            "The uploaded input could not be read as an image. "
            "Please upload an image file, or clear the gallery for text-to-image."
        )
    if dropped:
        print(f"[generate] ignored {dropped} non-image input(s).")
    if len(pils) > MAX_INPUT_IMAGES:
        raise gr.Error(f"Please use at most {MAX_INPUT_IMAGES} input images.")
    return pils


def _input_pixel_budget(num_images: int) -> int:
    """Per-image pixel budget so all inputs together fit the prefix budget."""
    return max(EDIT_INPUT_MIN_PIXELS, EDIT_INPUT_TOTAL_MAX_PIXELS // max(1, num_images))


def _prep_input_image(img: Image.Image, max_pixels: int) -> Image.Image:
    if img.mode == "RGBA":
        bg = Image.new("RGB", img.size, (255, 255, 255))
        bg.paste(img, mask=img.split()[3])
        img = bg
    img = img.convert("RGB")
    h, w = smart_resize(
        height=img.height,
        width=img.width,
        factor=EDIT_GRID_FACTOR,
        min_pixels=max_pixels,
        max_pixels=max_pixels,
    )
    if (w, h) != img.size:
        img = img.resize((w, h), Image.LANCZOS)
    return img


def _editing_output_size(input_img: Image.Image, target_pixels: int) -> tuple[int, int]:
    h, w = smart_resize(
        height=input_img.height,
        width=input_img.width,
        factor=EDIT_GRID_FACTOR,
        min_pixels=target_pixels,
        max_pixels=target_pixels,
    )
    return w, h


# Start the guard subprocess before the main model is loaded: Popen forks this
# process, and forking it after 35GB of weights are resident is far costlier.


print("[startup] loading SenseNova-U1.5-8B-MoT (this may take a few minutes)...")
sensenova_u1.set_attn_backend("auto")
print(f"[startup] attn backend: {sensenova_u1.effective_attn_backend()!r}")

config = AutoConfig.from_pretrained(MODEL_ID)
sensenova_u1.check_checkpoint_compatibility(config)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModel.from_pretrained(MODEL_ID, config=config, dtype=torch.bfloat16).to("cuda").eval()
print("[startup] model ready.")

import sn_aoti

# Swap the 42 Qwen3DecoderLayer gen paths for one AOTI-compiled package.
# No-op unless SN_AOTI=1 AND the published key matches this card exactly
# (torch version, sm arch, attention backend, layer count, hidden size); any
# mismatch prints one line and leaves the model eager. Nothing here touches the
# GPU: the download is CPU work and the .pt2 is not opened until the first
# denoising step inside the GPU worker.
print(f"[startup] {sn_aoti.status()}")
sn_aoti.maybe_load(model)




def _estimate_duration(
    pil_inputs, prompt, aspect_ratio, seed, num_steps=DEFAULT_NUM_STEPS, *args, **kwargs
):
    # Sampling cost is ~linear in step count, so scale the 50-step baselines
    # instead of always reserving the worst case.
    try:
        steps = max(MIN_NUM_STEPS, min(int(num_steps), MAX_NUM_STEPS))
    except (TypeError, ValueError):
        steps = DEFAULT_NUM_STEPS
    # Editing is heavier (image conditioning); give it more headroom.
    if pil_inputs:
        base = 180
    else:
        # T2I at 2048x2048 with 50 steps is the heaviest t2i case
        w, h = T2I_RESOLUTIONS.get(aspect_ratio, (2048, 2048))
        base = 180 if w * h > 2048 * 2048 else 120
    # The prefix pass and pixel decode don't shrink with the step count, so keep
    # a fixed floor and scale only the sampling part.
    return int(30 + (base - 30) * steps / REFERENCE_NUM_STEPS)


def _register_step_hook(counter: dict) -> object | None:
    """Count denoising steps by hooking the per-step timestep embedder.

    The sampling loops in ``t2i_generate`` / ``it2i_generate`` call
    ``fm_modules['timestep_embedder']`` exactly once per step, so a forward
    hook on it is a cheap, non-invasive progress signal.
    """
    try:
        embedder = model.fm_modules["timestep_embedder"]
    except Exception:  # pragma: no cover - defensive
        return None

    def _tick(_module, _args, _output):
        counter["step"] += 1

    try:
        return embedder.register_forward_hook(_tick)
    except Exception:  # pragma: no cover - defensive
        return None


@spaces.GPU(duration=_estimate_duration)
def _generate_on_gpu(
    pil_inputs: list,
    prompt: str,
    aspect_ratio: str,
    seed: int,
    num_steps: int = DEFAULT_NUM_STEPS,
    cfg_scale: float = DEFAULT_CFG_SCALE,
    timestep_shift: float = DEFAULT_TIMESTEP_SHIFT,
    img_cfg_scale: float = DEFAULT_IMG_CFG_SCALE,
):
    """GPU worker. Inputs are already validated / preprocessed on CPU.

    Runs sampling on a background thread inside the GPU fork and yields
    ``("progress", steps_done, total_steps)`` ticks while it works, then a
    final ``("result", image)``. The caller turns those ticks into a
    ``gr.Progress`` bar on the output component.
    """
    torch.cuda.reset_peak_memory_stats()
    counter = {"step": 0}
    total_steps = int(num_steps)
    hook = _register_step_hook(counter)
    result: dict = {}

    def _run():
        try:
            with torch.inference_mode():
                if not pil_inputs:
                    width, height = T2I_RESOLUTIONS[aspect_ratio]
                    tensor = model.t2i_generate(
                        tokenizer,
                        prompt,
                        image_size=(width, height),
                        cfg_scale=float(cfg_scale),
                        cfg_norm="none",
                        timestep_shift=float(timestep_shift),
                        cfg_interval=(0.0, 1.0),
                        num_steps=total_steps,
                        batch_size=1,
                        seed=int(seed),
                        think_mode=False,
                    )
                else:
                    out_w, out_h = _editing_output_size(pil_inputs[0], EDIT_TARGET_PIXELS)
                    tensor = model.it2i_generate(
                        tokenizer,
                        prompt,
                        pil_inputs,
                        image_size=(out_w, out_h),
                        cfg_scale=float(cfg_scale),
                        img_cfg_scale=float(img_cfg_scale),
                        cfg_norm="none",
                        timestep_shift=float(timestep_shift),
                        cfg_interval=(0.0, 1.0),
                        num_steps=total_steps,
                        batch_size=1,
                        think_mode=False,
                        seed=int(seed),
                    )
            result["image"] = _to_pil(tensor)[0]
            del tensor
        except BaseException as exc:  # re-raised on the generator thread below
            result["error"] = exc

    worker = threading.Thread(target=_run, daemon=True)
    started = time.time()
    try:
        worker.start()
        while True:
            worker.join(0.4)
            done = min(counter["step"], total_steps)
            if not worker.is_alive():
                break
            yield ("progress", done, total_steps, time.time() - started)
        print(f"[generate] peak GPU memory: {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB")
        exc = result.get("error")
        if isinstance(exc, torch.cuda.OutOfMemoryError):  # pragma: no cover
            raise gr.Error(
                "Ran out of GPU memory for this request. Try a smaller input image "
                "or a smaller aspect ratio."
            ) from exc
        if exc is not None:
            raise exc
        yield ("result", result["image"], total_steps, time.time() - started)
    finally:
        if hook is not None:
            hook.remove()
        # Never leave a failed call's tensors behind for the next request.
        torch.cuda.empty_cache()


def generate(
    images: list | None,
    prompt: str,
    aspect_ratio: str = "1:1",
    seed: int = 42,
    randomize_seed: bool = True,
    num_steps: int = DEFAULT_NUM_STEPS,
    cfg_scale: float = DEFAULT_CFG_SCALE,
    timestep_shift: float = DEFAULT_TIMESTEP_SHIFT,
    img_cfg_scale: float = DEFAULT_IMG_CFG_SCALE,
    progress=gr.Progress(),
):
    """Generate an image from a text prompt, or edit an uploaded image.

    Args:
        images: optional uploaded image(s) to edit; leave empty for text-to-image.
        prompt: what to generate, or the edit instruction to apply to the input image.
        aspect_ratio: output aspect ratio for text-to-image (ignored when editing).
        seed: RNG seed for reproducible sampling.
        randomize_seed: if True, pick a fresh random seed each run.
        num_steps: denoising steps; 50 is the model card reference, 28 is ~1.8x faster.
        cfg_scale: prompt guidance strength (reference: 4.0).
        timestep_shift: flow-matching schedule shift (reference: 3.0).
        img_cfg_scale: input-image guidance strength, editing only (reference: 1.0).
    """
    # Everything below runs on CPU: reject bad inputs *before* a GPU worker is
    # allocated, so a malformed request can never abort a GPU call mid-flight.
    if not isinstance(prompt, str) or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    prompt = prompt.strip()
    if aspect_ratio not in T2I_RESOLUTIONS:
        aspect_ratio = "1:1"
    try:
        seed = int(seed)
    except (TypeError, ValueError):
        seed = 42
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = max(0, min(int(seed), MAX_SEED))

    def _clamp(value, low, high, fallback):
        try:
            value = float(value)
        except (TypeError, ValueError):
            return fallback
        if value != value:  # NaN
            return fallback
        return max(low, min(value, high))

    num_steps = int(_clamp(num_steps, MIN_NUM_STEPS, MAX_NUM_STEPS, DEFAULT_NUM_STEPS))
    cfg_scale = _clamp(cfg_scale, 1.0, 10.0, DEFAULT_CFG_SCALE)
    timestep_shift = _clamp(timestep_shift, 0.5, 6.0, DEFAULT_TIMESTEP_SHIFT)
    img_cfg_scale = _clamp(img_cfg_scale, 1.0, 4.0, DEFAULT_IMG_CFG_SCALE)

    pil_images = _normalize_images(images)
    # Screen on CPU, before any GPU worker is allocated
    budget = _input_pixel_budget(len(pil_images))
    pil_inputs = [_prep_input_image(img, budget) for img in pil_images]
    if pil_inputs:
        sizes = ", ".join(f"{im.width}x{im.height}" for im in pil_inputs)
        print(f"[generate] editing with {len(pil_inputs)} input image(s) at {sizes}")

    # Drive a progress bar on the output component from the GPU worker's ticks.
    estimated = _estimate_duration(pil_inputs, prompt, aspect_ratio, seed, num_steps)
    progress(0.0, desc="Starting generation…")
    image_out = None
    for kind, payload, total_steps, elapsed in _generate_on_gpu(
        pil_inputs, prompt, aspect_ratio, seed, num_steps,
        cfg_scale, timestep_shift, img_cfg_scale,
    ):
        if kind == "progress":
            steps_done = int(payload)
            if steps_done > 0:
                progress(
                    min(steps_done / max(total_steps, 1), 1.0),
                    desc=f"Denoising step {steps_done}/{total_steps}",
                )
            else:
                # Prefix / conditioning pass, before the first denoising step.
                progress(
                    min(elapsed / max(estimated, 1), 0.05),
                    desc="Encoding prompt…",
                )
        else:
            image_out = payload
            progress(1.0, desc="Done")
    if image_out is None:  # pragma: no cover - worker always yields a result
        raise gr.Error("Generation produced no image. Please try again.")
    return image_out, seed


def generate_t2i(
    prompt: str,
    aspect_ratio: str = "1:1",
    seed: int = 42,
    randomize_seed: bool = True,
    num_steps: int = DEFAULT_NUM_STEPS,
    cfg_scale: float = DEFAULT_CFG_SCALE,
    timestep_shift: float = DEFAULT_TIMESTEP_SHIFT,
    progress=gr.Progress(),
):
    """Text-to-image only: generate an image from a prompt (no input image).

    Args:
        prompt: what to generate.
        aspect_ratio: output aspect ratio.
        seed: RNG seed for reproducible sampling.
        randomize_seed: if True, pick a fresh random seed each run.
        num_steps: denoising steps (default 28; 50 is the reference config).
        cfg_scale: prompt guidance strength.
        timestep_shift: flow-matching schedule shift.
    """
    # Wired to the text-to-image gr.Examples: those rows carry (prompt,
    # aspect_ratio) only, so they must NOT be bound straight to `generate`,
    # whose first argument is the input-image list.
    return generate(
        None, prompt, aspect_ratio, seed, randomize_seed,
        num_steps, cfg_scale, timestep_shift, progress=progress,
    )


def generate_edit(
    images: list | None,
    prompt: str,
    seed: int = 42,
    randomize_seed: bool = True,
    num_steps: int = DEFAULT_NUM_STEPS,
    cfg_scale: float = DEFAULT_CFG_SCALE,
    timestep_shift: float = DEFAULT_TIMESTEP_SHIFT,
    img_cfg_scale: float = DEFAULT_IMG_CFG_SCALE,
    progress=gr.Progress(),
):
    """Image editing: apply an edit instruction to the uploaded image(s).

    Args:
        images: input image(s) to edit.
        prompt: the edit instruction.
        seed: RNG seed for reproducible sampling.
        randomize_seed: if True, pick a fresh random seed each run.
        num_steps: denoising steps (default 28; 50 is the reference config).
        cfg_scale: prompt guidance strength.
        timestep_shift: flow-matching schedule shift.
        img_cfg_scale: input-image guidance strength.
    """
    return generate(
        images, prompt, "1:1", seed, randomize_seed,
        num_steps, cfg_scale, timestep_shift, img_cfg_scale, progress=progress,
    )


# T2I examples: prompt + aspect ratio
T2I_EXAMPLES = [
    [
        "A cinematic mountain lake at sunrise, realistic photography, golden mist over still water, snow-capped peaks reflected in the lake, ultra-detailed.",
        "1:1",
    ],
    [
        'A neon bar sign that clearly reads "OPEN LATE", dark interior, moody reflections, easy text rendering.',
        "16:9",
    ],
    [
        "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look.",
        "2:3",
    ],
    [
        "A cute fluffy corgi puppy wearing a tiny chef's hat, sitting at a wooden table with fresh-baked cookies, warm kitchen lighting, photorealistic.",
        "1:1",
    ],
    [
        "Lavender fields stretching to the horizon under a pastel sunset, a small stone farmhouse, highly detailed flowers, romantic countryside scene.",
        "4:3",
    ],
]

# Editing examples: gallery input (list of paths) + prompt
EDIT_EXAMPLES = [
    [["examples/edit_1.webp"], "Change the jacket of the person on the left to bright yellow."],
    [["examples/edit_2.webp"], "Make the person in the image smile."],
    [["examples/edit_3.webp"], "Add a bouquet of flowers."],
    [["examples/edit_4.webp"], "Turn the image into an American comic style."],
]


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="SenseNova-U1.5-8B-MoT") as demo:
    gr.Markdown(
        """
        # SenseNova-U1.5-8B-MoT

        Unified text-to-image **and** image editing with
        [**SenseNova-U1.5-8B-MoT**](https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT),
        a natively unified multimodal model built on the
        [NEO-unify](https://huggingface.co/blog/sensenova/neo-unify) architecture.
        Leave the image upload empty for text-to-image, or upload an image and
        write an edit instruction.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input_gallery = gr.Gallery(
                label="Upload image(s) to edit (leave empty for text-to-image)",
                file_types=["image"],
                columns=4,
            )
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image to generate, or how to edit your input.",
                lines=3,
            )
            aspect_ratio = gr.Dropdown(
                label="Aspect ratio (text-to-image only — editing keeps input ratio)",
                choices=list(T2I_RESOLUTIONS.keys()),
                value="1:1",
            )
            generate_button = gr.Button("Generate", variant="primary")

            with gr.Accordion("Advanced options", open=False):
                num_steps = gr.Slider(
                    label="Denoising steps",
                    info=(
                        f"{DEFAULT_NUM_STEPS} keeps composition and prompt adherence while "
                        f"running ~1.8x faster than the {REFERENCE_NUM_STEPS}-step reference; "
                        "raise it for fine texture in landscapes and skin."
                    ),
                    minimum=MIN_NUM_STEPS,
                    maximum=MAX_NUM_STEPS,
                    step=1,
                    value=DEFAULT_NUM_STEPS,
                )
                cfg_scale = gr.Slider(
                    label="Guidance scale (CFG)",
                    info="How strictly to follow the prompt. Reference: 4.0.",
                    minimum=1.0,
                    maximum=10.0,
                    step=0.1,
                    value=DEFAULT_CFG_SCALE,
                )
                timestep_shift = gr.Slider(
                    label="Timestep shift",
                    info=(
                        "Shifts sampling toward high-noise steps. Reference: 3.0 — lowering it "
                        "changes the image rather than sharpening it."
                    ),
                    minimum=0.5,
                    maximum=6.0,
                    step=0.1,
                    value=DEFAULT_TIMESTEP_SHIFT,
                )
                img_cfg_scale = gr.Slider(
                    label="Image guidance (editing only)",
                    info="How closely an edit sticks to the input image. Reference: 1.0.",
                    minimum=1.0,
                    maximum=4.0,
                    step=0.1,
                    value=DEFAULT_IMG_CFG_SCALE,
                )
                with gr.Row():
                    seed = gr.Slider(
                        label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42
                    )
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Output",
                type="pil",
                format="png",
                interactive=False,
            )
            used_seed = gr.Number(label="Seed used", interactive=False)

    with gr.Accordion("Image Editing Examples", open=True):
        gr.Examples(
            examples=EDIT_EXAMPLES,
            inputs=[image_input_gallery, prompt_input],
            outputs=[output_image, used_seed],
            fn=generate_edit,
            cache_examples=False,
            run_on_click=True,
        )
        
    with gr.Accordion("Text-to-Image Examples", open=True):
        gr.Examples(
            examples=T2I_EXAMPLES,
            inputs=[prompt_input, aspect_ratio],
            outputs=[output_image, used_seed],
            fn=generate_t2i,
            cache_examples=False,
            run_on_click=True,
        )

    

    generate_button.click(
        fn=generate,
        inputs=[
            image_input_gallery, prompt_input, aspect_ratio, seed,
            randomize_seed, num_steps, cfg_scale, timestep_shift, img_cfg_scale,
        ],
        outputs=[output_image, used_seed],
        api_name="generate",
    )
    prompt_input.submit(
        fn=generate,
        inputs=[
            image_input_gallery, prompt_input, aspect_ratio, seed,
            randomize_seed, num_steps, cfg_scale, timestep_shift, img_cfg_scale,
        ],
        outputs=[output_image, used_seed],
        api_name="generate_submit",
    )


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)