app.py

# Copyright (c) 2024-present, BAAI. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ------------------------------------------------------------------------
"""URSA TI2V application."""

import argparse
import os
import sys

import gradio as gr
import numpy as np
import PIL.Image
import torch

from diffnext.pipelines import URSAPipeline
from diffnext.utils import export_to_image, export_to_video

# =========================
# 🔧 修复 Qwen3Config 缺失 rope_theta
# =========================
# =========================

# =========================
def _patch_diffnext_qwen3():
    try:
        from diffnext.models.text_encoders import qwen3
    except Exception:
        return

    orig_from_config = qwen3.Qwen3RotaryEmbedding.from_config

    @classmethod
    def patched_from_config(cls, config):
        # rope_theta
        if not hasattr(config, "rope_theta"):
            rp = getattr(config, "rope_parameters", None)
            if isinstance(rp, dict) and "rope_theta" in rp:
                config.rope_theta = rp["rope_theta"]
            else:
                config.rope_theta = float(1e6)

        # max_position_embeddings
        if not hasattr(config, "max_position_embeddings"):
            config.max_position_embeddings = int(
                getattr(config, "max_seq_len", 2048)
            )

        # ✅ 正确：只传 config，不要传 cls
        return orig_from_config(config)

    qwen3.Qwen3RotaryEmbedding.from_config = patched_from_config


_patch_diffnext_qwen3()
# =========================


# ✅ HF Token（从环境变量读取，不写死）
os.environ["HF_TOKEN"] = os.environ.get("HF_TOKEN", "")
os.environ["HUGGINGFACE_HUB_TOKEN"] = os.environ.get("HF_TOKEN", "")
os.environ["HF_HUB_ENABLE_HUGGINGFACE_TRANSFER"] = "1"
os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "300"

os.environ["TOKENIZERS_PARALLELISM"] = "true"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"


def parse_args():
    parser = argparse.ArgumentParser(description="Serve URSA TI2V application")
    parser.add_argument("--model", default="BAAI/URSA-0.6B-FSQ320", help="model path")
    parser.add_argument("--device", type=int, default=0, help="device index")
    parser.add_argument("--precision", default="float16", help="compute precision")
    return parser.parse_args()


def crop_image(image, target_h, target_w):
    h, w = image.height, image.width
    aspect_ratio_target, aspect_ratio = target_w / target_h, w / h
    if aspect_ratio > aspect_ratio_target:
        new_w = int(h * aspect_ratio_target)
        x_start = (w - new_w) // 2
        image = image.crop((x_start, 0, x_start + new_w, h))
    else:
        new_h = int(w / aspect_ratio_target)
        y_start = (h - new_h) // 2
        image = image.crop((0, y_start, w, y_start + new_h))
    return np.array(image.resize((target_w, target_h), PIL.Image.Resampling.BILINEAR))


def generate_image(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    guidance_scale,
    num_inference_steps=25,
):
    args = {**locals(), **video_presets["t2i"]}
    seed = np.random.randint(2147483647) if randomize_seed else seed
    device = getattr(pipe, "_offload_device", pipe.device)
    generator = torch.Generator(device=device).manual_seed(seed)
    images = pipe(generator=generator, **args).frames
    return [export_to_image(image, quality=95) for image in images] + [seed]


def generate_video(
    prompt,
    negative_prompt,
    image,
    motion_score,
    seed,
    randomize_seed,
    guidance_scale,
    num_inference_steps,
    output_type="np",
):
    args = {**locals(), **video_presets["ti2v"]}
    args["prompt"] = f"motion={motion_score:.1f}, {prompt}"
    args["image"] = crop_image(image, args["height"], args["width"]) if image else None
    seed = np.random.randint(2147483647) if randomize_seed else seed
    device = getattr(pipe, "_offload_device", pipe.device)
    generator = torch.Generator(device=device).manual_seed(seed)
    frames = pipe(generator=generator, **args).frames[0]
    return export_to_video(frames, fps=12), seed


css = """#col-container {margin: 0 auto; max-width: 1366px}"""
header = (
    "<div align='center'>"
    "<h2>Uniform Discrete Diffusion with Metric Path for Video Generation</h2>"
    "<h3><a href='https://arxiv.org/abs/2510.24717' target='_blank' rel='noopener'>[paper]</a>"
    "<a href='https://github.com/baaivision/URSA' target='_blank' rel='noopener'>[code]</a></h3>"
    "</div>"
)

video_presets = {
    "t2i": {"width": 512, "height": 320, "num_frames": 1},
    "ti2v": {"width": 512, "height": 320, "num_frames": 49},
}

prompts = [
    "a lone grizzly bear walks through a misty forest at dawn, sunlight catching its fur.",
    "Many spotted jellyfish pulsating under water. Their bodies are transparent and glowing in deep ocean.",
    "An intense close-up of a soldier’s face, covered in dirt and sweat, his eyes filled with determination as he surveys the battlefield.",
    "a close-up shot of a woman standing in a dimly lit room. she is wearing a traditional chinese outfit, which includes a red and gold dress with intricate designs and a matching headpiece. the woman has her hair styled in an updo, adorned with a gold accessory. her makeup is done in a way that accentuates her features, with red lipstick and dark eyeshadow. she is looking directly at the camera with a neutral expression. the room has a rustic feel, with wooden beams and a stone wall visible in the background. the lighting in the room is soft and warm, creating a contrast with the woman's vibrant attire. there are no texts or other objects in the video. the style of the video is a portrait, focusing on the woman and her attire.",
    "The camera slowly rotates around a massive stack of vintage televisions that are placed within a large New York museum gallery. Each of the televisions is showing a different program. There are 1950s sci-fi movies with their distinctive visuals, horror movies with their creepy scenes, news broadcasts with moving images and words, static on some screens, and a 1970s sitcom with its characteristic look. The televisions are of various sizes and designs, some with rounded edges and others with more angular shapes. The gallery is well-lit, with light falling on the stack of televisions and highlighting the different programs being shown. There are no people visible in the immediate vicinity, only the stack of televisions and the surrounding gallery space.",
]
motion_scores = [9, 9, 9, 9, 9]
examples = [list(x) for x in zip(prompts, motion_scores)]


if __name__ == "__main__":
    args = parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu", args.device)
    model_args = {
        "torch_dtype": getattr(torch, args.precision.lower()),
        "trust_remote_code": True,
    }

    pipe = URSAPipeline.from_pretrained(
        args.model,
        **model_args,
        resume_download=True,
    ).to(device)

    # Application
    app = gr.Blocks(css=css, theme="origin").__enter__()
    container = gr.Column(elem_id="col-container").__enter__()
    _, main_row = gr.Markdown(header), gr.Row().__enter__()

    # Input
    input_col = gr.Column().__enter__()
    prompt = gr.Text(
        label="Prompt",
        placeholder="Describe the video you want to generate",
        value="A lone grizzly bear walks through a misty forest at dawn, sunlight catching its fur.",
        lines=5,
    )
    negative_prompt = gr.Text(
        label="Negative Prompt",
        placeholder="Describe what you don't want in the video",
        value="worst quality, low quality, inconsistent motion, static, still, blurry, jittery, distorted, ugly",
        lines=1,
    )
    with gr.Row():
        generate_image_btn = gr.Button("Generate Image Prompt", variant="primary", size="lg")
        generate_video_btn = gr.Button("Generate Video", variant="primary", size="lg")
    image_prompt = gr.Image(label="Image Prompt", height=480, type="pil")

    options = gr.Accordion("Options", open=False).__enter__()
    seed = gr.Slider(label="Seed", maximum=2147483647, step=1, value=0)
    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
    guidance_scale = gr.Slider(label="Guidance scale", minimum=1, maximum=10.0, step=0.1, value=7.0)
    with gr.Row():
        num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=100, step=1, value=50)
    options.__exit__(), input_col.__exit__()

    # Results
    result_col = gr.Column().__enter__()
    motion = gr.Slider(label="Motion Score", minimum=1, maximum=10, step=1, value=9)
    result = gr.Video(label="Result", height=480, show_label=False, autoplay=True)
    result_col.__exit__(), main_row.__exit__()

    with gr.Row():
        gr.Examples(examples=examples, inputs=[prompt, motion])

    container.__exit__()
    gr.on(
        triggers=[generate_image_btn.click, prompt.submit, negative_prompt.submit],
        fn=generate_image,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            guidance_scale,
        ],
        outputs=[image_prompt, seed],
    )
    gr.on(
        triggers=[generate_video_btn.click, prompt.submit, negative_prompt.submit],
        fn=generate_video,
        inputs=[
            prompt,
            negative_prompt,
            image_prompt,
            motion,
            seed,
            randomize_seed,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )
    app.__exit__(), app.launch(share=False, ssr_mode=False, max_threads=1)