app.py

from __future__ import annotations

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

import os
import random

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_NUM_STEPS = 50

# Editing output grid factor (= patch_size * merge_size = 32).
EDIT_GRID_FACTOR = 32
EDIT_TARGET_PIXELS = 2048 * 2048
EDIT_INPUT_MAX_PIXELS = 2048 * 2048

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:
    if isinstance(img, Image.Image):
        return img
    if isinstance(img, tuple):
        img = img[0]
    if isinstance(img, str):
        return Image.open(img)
    return img


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


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


def _estimate_duration(images, prompt, aspect_ratio, seed, randomize_seed, *args, **kwargs):
    # Editing is heavier (image conditioning); give it more headroom.
    has_input = images is not None and len(images) > 0
    if has_input:
        return 180
    # T2I at 2048x2048 with 50 steps is the heaviest t2i case
    w, h = T2I_RESOLUTIONS.get(aspect_ratio, (2048, 2048))
    pixels = w * h
    if pixels > 2048 * 2048:
        return 180
    return 120


@spaces.GPU(duration=_estimate_duration)
def generate(
    images: list | None,
    prompt: str,
    aspect_ratio: str = "1:1",
    seed: int = 42,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """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.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    has_input = images is not None and len(images) > 0

    with torch.inference_mode():
        if not has_input:
            width, height = T2I_RESOLUTIONS[aspect_ratio]
            tensor = model.t2i_generate(
                tokenizer,
                prompt,
                image_size=(width, height),
                cfg_scale=DEFAULT_CFG_SCALE,
                cfg_norm="none",
                timestep_shift=DEFAULT_TIMESTEP_SHIFT,
                cfg_interval=(0.0, 1.0),
                num_steps=DEFAULT_NUM_STEPS,
                batch_size=1,
                seed=int(seed),
                think_mode=False,
            )
        else:
            pil_inputs = [_prep_input_image(_coerce_pil(item), EDIT_INPUT_MAX_PIXELS) for item in images]
            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=DEFAULT_CFG_SCALE,
                img_cfg_scale=1.0,
                cfg_norm="none",
                timestep_shift=DEFAULT_TIMESTEP_SHIFT,
                cfg_interval=(0.0, 1.0),
                num_steps=DEFAULT_NUM_STEPS,
                batch_size=1,
                think_mode=False,
                seed=int(seed),
            )

    images_out = _to_pil(tensor)
    return images_out[0], seed


# 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: image path + 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", theme=gr.themes.Citrus(), css=CSS) 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"],
                height=200,
                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",
            )
            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)
            generate_button = gr.Button("Generate", variant="primary")
        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("Text-to-Image Examples", open=True):
        gr.Examples(
            examples=T2I_EXAMPLES,
            inputs=[prompt_input, aspect_ratio],
            outputs=[output_image, used_seed],
            fn=generate,
            cache_examples=False,
            run_on_click=True,
        )

    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,
            cache_examples=False,
            run_on_click=True,
        )

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


if __name__ == "__main__":
    demo.launch(mcp_server=True)