app.py

"""UniAR: Unified Multimodal Autoregressive Modeling — image understanding & generation.

This Space demos the UniAR model (ShareLab-SII/UniAR-RL), a single Transformer
that handles both visual understanding (VQA) and text-to-image generation via
a shared discrete visual tokenizer (BSQ).
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import sys

# Make the bundled uniar/ and inference/ packages importable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
from transformers import AutoProcessor

from uniar import (
    UniARForConditionalGeneration,
    UniARVisualDecoder,
    CHAT_TEMPLATE,
)
from inference.visual_inputs import prepare_visual_inputs

MODEL_ID = "ShareLab-SII/UniAR-RL"

# ---------------------------------------------------------------------------
# Load all model components at module scope (ZeroGPU packs to disk on startup)
# ---------------------------------------------------------------------------

print(f"[UniAR] Loading AR model from {MODEL_ID} …")
ar_processor = AutoProcessor.from_pretrained(MODEL_ID, padding_side="left")
ar_model = UniARForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
ar_model.eval()

print("[UniAR] Loading visual decoder …")
visual_decoder = UniARVisualDecoder.from_pretrained(MODEL_ID, device="cuda")

print("[UniAR] All components loaded.")


# ---------------------------------------------------------------------------
# Inference functions
# ---------------------------------------------------------------------------

@spaces.GPU(duration=180)
def understand_image(image, prompt: str, max_new_tokens: int = 1024) -> str:
    """Answer a question about an input image (visual question answering).

    Args:
        image: Input PIL image.
        prompt: Text question about the image.
        max_new_tokens: Maximum number of text tokens to generate.
    """
    if image is None:
        return "Please upload an image first."
    if not prompt.strip():
        return "Please enter a question."

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]

    inputs = ar_processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")
    inputs.pop("mm_token_type_ids", None)

    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        generated_ids = ar_model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
        )

    generated_ids_trimmed = [
        output_ids[len(input_ids):]
        for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = ar_processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    return output_text[0] if output_text else ""


@spaces.GPU(duration=300)
def generate_image(
    prompt: str,
    ar_height: int = 960,
    ar_width: int = 960,
    upsampling_ratio: float = 1.0,
    temperature: float = 0.5,
    cfg: float = 1.5,
    decoder_steps: int = 28,
    decoder_cfg: float = 1.5,
):
    """Generate an image from a text prompt.

    Args:
        prompt: Text description of the image to generate.
        ar_height: AR visual-token rollout height in pixels.
        ar_width: AR visual-token rollout width in pixels.
        upsampling_ratio: SD3 decoder output scale (1.0–2.0).
        temperature: Sampling temperature for visual tokens.
        cfg: Classifier-free guidance scale for AR rollout.
        decoder_steps: SD3 visual decoder denoising steps.
        decoder_cfg: SD3 visual decoder CFG scale.
    """
    if not prompt.strip():
        return None, "Please enter a prompt."

    visual_inputs = prepare_visual_inputs(
        [prompt], ar_model, ar_processor,
        ar_height, ar_width,
    )
    # Cast attention mask to bool — SDPA expects bool/float, not long
    visual_inputs["attention_mask"] = visual_inputs["attention_mask"].bool()
    generated_visual_ids = ar_model.generate_visual(
        prefix_input_ids=visual_inputs["prefix_input_ids"],
        attention_mask=visual_inputs["attention_mask"],
        pos_ids_image=visual_inputs["pos_ids_image"],
        pos_ids_all=visual_inputs["pos_ids_all"],
        image_token_num=visual_inputs["image_token_num"],
        temperature=temperature,
        cfg=cfg,
        show_progress=False,
    )

    images = visual_decoder.decode(
        generated_visual_ids,
        ar_height=ar_height,
        ar_width=ar_width,
        upsampling_ratio=upsampling_ratio,
        num_inference_steps=decoder_steps,
        guidance_scale=decoder_cfg,
    )
    return images[0], "Image generated successfully."


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        """
        # UniAR: Unified Multimodal Autoregressive Modeling
        A single Transformer for **image understanding** and **image generation**
        using a shared discrete visual tokenizer (BSQ).

        [[Paper](https://arxiv.org/abs/2606.18249)] ·
        [[Model](https://huggingface.co/ShareLab-SII/UniAR-RL)] ·
        [[Code](https://github.com/ShareLab-SII/UniAR)]
        """
    )

    with gr.Tabs():
        # ----- Tab 1: Image Understanding -----
        with gr.Tab("🖼️ Image Understanding"):
            with gr.Row():
                with gr.Column(scale=1):
                    understand_image_in = gr.Image(
                        label="Input Image", type="pil", height=350
                    )
                    understand_prompt = gr.Textbox(
                        label="Question / Prompt",
                        placeholder="Describe this image in detail.",
                        lines=2,
                    )
                    understand_btn = gr.Button("Analyze", variant="primary")
                with gr.Column(scale=1):
                    understand_out = gr.Textbox(
                        label="Response", lines=12, interactive=False
                    )

            gr.Examples(
                examples=[
                    ["example_input.jpg", "Describe this image in detail."],
                    ["example_input.jpg", "What objects can you see in this image?"],
                    ["example_input.jpg", "What is the color palette of this image?"],
                ],
                inputs=[understand_image_in, understand_prompt],
                outputs=understand_out,
                fn=understand_image,
                cache_examples=False,
                run_on_click=True,
            )

        # ----- Tab 2: Image Generation -----
        with gr.Tab("🎨 Image Generation"):
            with gr.Row():
                with gr.Column(scale=4):
                    gen_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="A cute anime girl.",
                        lines=2,
                        show_label=False,
                        container=False,
                    )
                with gr.Column(scale=1):
                    gen_btn = gr.Button("Generate", variant="primary")
            gen_out = gr.Image(label="Generated Image", height=450)
            gen_status = gr.Textbox(label="Status", interactive=False, show_label=False)

            with gr.Accordion("Advanced settings", open=False):
                with gr.Row():
                    gen_height = gr.Slider(
                        256, 1280, value=960, step=32,
                        label="AR Height", info="Higher = more detail, slower"
                    )
                    gen_width = gr.Slider(
                        256, 1280, value=960, step=32,
                        label="AR Width", info="Higher = more detail, slower"
                    )
                with gr.Row():
                    gen_upsample = gr.Slider(
                        1.0, 2.0, value=2.0, step=0.1,
                        label="Upsampling Ratio",
                        info="Final output = AR size × ratio"
                    )
                    gen_temp = gr.Slider(
                        0.5, 1.0, value=0.5, step=0.1,
                        label="Temperature"
                    )
                with gr.Row():
                    gen_cfg = gr.Slider(
                        1.0, 5.0, value=1.5, step=0.1,
                        label="AR CFG Scale"
                    )
                    gen_decoder_steps = gr.Slider(
                        10, 50, value=28, step=1,
                        label="Decoder Steps"
                    )
                    gen_decoder_cfg = gr.Slider(
                        0.5, 5.0, value=1.5, step=0.1,
                        label="Decoder CFG Scale"
                    )

            gen_btn.click(
                generate_image,
                inputs=[
                    gen_prompt, gen_height, gen_width,
                    gen_upsample, gen_temp, gen_cfg,
                    gen_decoder_steps, gen_decoder_cfg,
                ],
                outputs=[gen_out, gen_status],
                api_name="generate",
            )

            gr.Examples(
                examples=[
                    [
                        "The image features a young individual sitting outdoors, likely at a "
                        "playground or amusement area. The person is wearing a maroon beanie, "
                        "round glasses, and a white long-sleeve shirt layered over a dark blue "
                        "sweater. They have a warm, friendly expression, with a slight smile "
                        "showing their teeth. Their hair is dark and falls naturally around "
                        "their face. The individual is leaning slightly forward, resting their "
                        "chin on their hand, which gives a relaxed and casual pose. The "
                        "background is blurred but shows greenery, trees, and some colorful "
                        "structures, suggesting a lively outdoor setting. There are hints of "
                        "other people and objects in the background, adding to the sense of a "
                        "public space. The lighting is soft and natural, indicating daytime, "
                        "possibly during autumn given the muted tones and leafless branches "
                        "visible in the background.",
                        960, 960, 1.067, 0.5, 1.5, 28, 1.5,
                    ],
                    [
                        "The image is a black-and-white portrait of a woman standing outdoors near a body of water, likely a beach or seaside. She has long, wet hair pulled back, suggesting she has been in the water recently. Her facial features include prominent eyes with defined eyelashes, a straight nose, and full lips. She appears to be wearing minimal makeup, with noticeable mascara enhancing her eyelashes. The woman's skin is smooth, and there are small droplets of water on her face and neck, indicating recent contact with water. She is wearing a dark, possibly one-piece swimsuit, which is visible from the shoulders up. The background is blurred but shows a soft, misty horizon with hints of land or hills in the distance, creating a serene and natural setting. The lighting is soft, giving the image a calm and reflective mood.",
                        960, 960, 1.067, 0.5, 1.5, 28, 1.5,
                    ],
                    [
                        "The image depicts a young woman standing outdoors in a park-like setting. She has short black hair styled into two braids, which she is holding with her right hand. She is wearing a light beige bucket hat and a white sleeveless dress with thin straps. Her skin appears smooth, and she has a gentle, warm smile with rosy cheeks. The background is softly blurred, featuring green grass, trees, and dappled sunlight filtering through the foliage, creating a serene and natural atmosphere. The tree trunk next to her is prominent, showing its rough bark texture, and she leans slightly against it. The overall mood of the image is calm and peaceful.",
                        1280, 704, 1, 0.5, 1.5, 28, 1.5,
                    ],
                    [
                        "The image features a person standing indoors against a clean, minimalist white background. The individual has long, wavy brown hair that cascades over their shoulders. They are wearing a fitted, off-the-shoulder dress in a muted purple color. The dress has long sleeves and a draped design at the front, accentuating the figure. The person is also wearing transparent high-heeled sandals with a nude tone. Their posture is confident, with one hand resting on their hip and the other hanging naturally. The setting includes a white couch with fluffy white cushions in the background, and the floor is covered with a soft, white shaggy rug. The overall ambiance is bright and elegant, suggesting a stylish and sophisticated environment.",
                        1280, 704, 1, 0.5, 1.5, 28, 1.5,
                    ],
                    [
                        "a breathtaking photograph capturing the vibrant hues of a sunset with streaks of pink and orange painting the sky behind the majestic Grand Canyon. The canyon's intricate rock formations are silhouetted against the illuminated backdrop, showcasing the deep crevices and towering spires. In the foreground, the Colorado River can be glimpsed winding its way through the ancient geological marvel.",
                        704, 1280, 1, 0.5, 1.5, 28, 1.5,
                    ],
                    [
                        "A striking statue of the Egyptian god Anubis, depicted with a jackal's head, is dressed unconventionally in modern attire consisting of a crisp white t-shirt, a black leather jacket, and a pair of aviator goggles resting atop its head. The statue stands in sharp contrast against the backdrop of a full moon that illuminates the night sky over the sprawling cityscape of Los Angeles. The city lights twinkle in the distance, creating a mosaic of urban life behind the ancient deity's contemporary ensemble.",
                        704, 1280, 1, 0.5, 1.5, 28, 1.5,
                    ]
                ],
                inputs=[
                    gen_prompt, gen_height, gen_width,
                    gen_upsample, gen_temp, gen_cfg,
                    gen_decoder_steps, gen_decoder_cfg,
                ],
                outputs=[gen_out, gen_status],
                fn=generate_image,
                cache_examples=False,
                run_on_click=True,
            )

    understand_btn.click(
        understand_image,
        inputs=[understand_image_in, understand_prompt],
        outputs=understand_out,
        api_name="understand",
    )


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