interface update

.gitignore

@@ -1,2 +1,4 @@
 **/__pycache__/
-*.pyc
+*.pyc
+.vscode
+*.mp4

app.py

@@ -1,97 +1,278 @@
+import functools
+from dataclasses import dataclass
+import random
+
 import gradio as gr
 import spaces
-import torch
-import os
 
-# ==========================================
-# 1. 核心处理函数 (骨架)
-# ==========================================
+
+@dataclass(frozen=True)
+class SliderConfig:
+    label: str
+    minimum: float
+    maximum: float
+    step: float
+    value: float
+    info: str
+
+
+@dataclass(frozen=True)
+class PresetConfig:
+    shift_gamma: float
+    steps: int
+    guidance_scale: float
+
+
+GAMMA_SLIDER = SliderConfig(
+    label="Shift Gamma",
+    minimum=1.0,
+    maximum=10.0,
+    step=0.5,
+    value=5.0,
+    info="Scheduler adjustment parameter."
+)
+
+STEP_SLIDER = SliderConfig(
+    label="Inference Steps",
+    minimum=10,
+    maximum=50,
+    step=1,
+    value=28,
+    info="More steps improve quality but take longer."
+)
+
+GUIDANCE_SLIDER = SliderConfig(
+    label="Guidance Scale (CFG)",
+    minimum=1.0,
+    maximum=20.0,
+    step=0.5,
+    value=1.5,
+    info="Controls adherence to the text prompt."
+)
+
+
+STYLE_CHOICES = [
+    "Oil painting style, vivid colors",
+    "Neon cyberpunk, futuristic city",
+    "Minimalist sketch, soft shading",
+    "Anime aesthetic, bold lines",
+]
+
+
+EXAMPLE_INPUTS = [
+    ["assets/video_00000000.mp4", "Oil painting style, vivid colors"],
+    ["assets/video_00000007.mp4", "Neon cyberpunk, futuristic city"],
+    ["assets/video_00000107.mp4", "Minimalist sketch, soft shading"],
+]
+
+
+PRESET_MODES = {
+    "Fast": PresetConfig(shift_gamma=5.0, steps=10, guidance_scale=1.),
+    "Balanced": PresetConfig(shift_gamma=5.0, steps=10, guidance_scale=1.5),
+    "Quality": PresetConfig(shift_gamma=5.0, steps=20, guidance_scale=1.5),
+}
+
+
+def _create_slider(config: SliderConfig) -> gr.Slider:
+    """Helper to keep slider creation consistent."""
+    return gr.Slider(
+        label=config.label,
+        minimum=config.minimum,
+        maximum=config.maximum,
+        step=config.step,
+        value=config.value,
+        info=config.info,
+    )
+
+
 @spaces.GPU
-def run_stylization(input_video_path, prompt, noise_scale, shift_gamma, steps, guidance_scale, seed):
-    """
-    这里是实际推理逻辑的占位符。
-    """
+def run_stylization(
+    input_video_path,
+    prompt,
+    shift_gamma,
+    steps,
+    guidance_scale,
+    seed,
+    randomize_seed,
+):
+    """Placeholder stylization pipeline that echoes the input video."""
     if not input_video_path:
         return None
-    
+
+    resolved_seed = _resolve_seed(seed, randomize_seed)
+
     print("========== Inference Start ==========")
     print(f"Video Path: {input_video_path}")
     print(f"Prompt: {prompt}")
-    print(f"Params: Noise={noise_scale}, Gamma={shift_gamma}, Steps={steps}, CFG={guidance_scale}, Seed={seed}")
+    print(
+        "Params: "
+        f"Gamma={shift_gamma}, "
+        f"Steps={steps}, "
+        f"CFG={guidance_scale}, "
+        f"Seed={resolved_seed}"
+    )
 
-    # [Prototype Logic] 直接返回输入视频用于演示
     return input_video_path
 
-# ==========================================
-# 2. 界面布局 (Gradio Blocks)
-# ==========================================
-
-# 移除 CSS 参数以修复 TypeError
-with gr.Blocks() as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("# 🎥 ViBT Video Stylization Interface")
-        gr.Markdown("上传视频并设置风格化参数。")
-
-        with gr.Row():
-            # --- 左侧：输入与设置 ---
-            with gr.Column():
-                # 视频输入
-                input_video = gr.Video(label="Source Video", sources=["upload"])
-                
-                # 提示词
-                prompt_input = gr.Textbox(
-                    label="Style Prompt", 
-                    placeholder="e.g., Van Gogh style, cyberpunk city...",
-                    value="Oil painting style, vivid colors"
-                )
-
-                # 高级参数折叠区
-                with gr.Accordion("Advanced Settings", open=True):
-                    with gr.Row():
-                        noise_scale = gr.Slider(
-                            label="Noise Scale", minimum=0.0, maximum=2.0, step=0.1, value=1.0, 
-                            info="Controls how much noise is added."
-                        )
-                        shift_gamma = gr.Slider(
-                            label="Shift Gamma", minimum=1.0, maximum=10.0, step=0.5, value=5.0,
-                            info="Scheduler parameter."
-                        )
-                    
-                    with gr.Row():
-                        num_steps = gr.Slider(
-                            label="Inference Steps", minimum=10, maximum=50, step=1, value=28,
-                            info="More steps = higher quality but slower."
-                        )
-                        guidance_scale = gr.Slider(
-                            label="Guidance Scale (CFG)", minimum=1.0, maximum=20.0, step=0.5, value=1.5,
-                            info="How closely to follow the prompt."
-                        )
-                    
-                    seed = gr.Number(label="Seed", value=42, precision=0)
-
-                # 提交按钮
-                run_btn = gr.Button("Generate Video", variant="primary")
-
-            # --- 右侧：结果输出 ---
-            with gr.Column():
-                output_video = gr.Video(label="Stylized Result", interactive=False)
-
-    # ==========================================
-    # 3. 事件绑定
-    # ==========================================
-    run_btn.click(
-        fn=run_stylization,
-        inputs=[
-            input_video, 
-            prompt_input, 
-            noise_scale, 
-            shift_gamma, 
-            num_steps, 
-            guidance_scale, 
-            seed
-        ],
-        outputs=[output_video]
+
+def _resolve_seed(seed_value, randomize):
+    """Return an integer seed, generating a random one when requested or missing."""
+    if randomize or seed_value in (None, ""):
+        return random.randint(0, 2**31 - 1)
+    return int(seed_value)
+
+
+def run_with_preset(input_video_path, prompt, seed, randomize_seed, preset_key):
+    """Wrap stylization with predefined presets for quick generation."""
+    preset = PRESET_MODES[preset_key]
+    return run_stylization(
+        input_video_path=input_video_path,
+        prompt=prompt,
+        shift_gamma=preset.shift_gamma,
+        steps=preset.steps,
+        guidance_scale=preset.guidance_scale,
+        seed=seed,
+        randomize_seed=randomize_seed,
+    )
+
+
+def _bind_preset_button(button, preset_key, inputs, output, extra_kwargs=None):
+    extra_kwargs = extra_kwargs or {}
+    button.click(
+        fn=functools.partial(run_with_preset, preset_key=preset_key, **extra_kwargs),
+        inputs=inputs,
+        outputs=[output],
     )
 
+
+def build_demo() -> gr.Blocks:
+    """Create the Gradio interface for video stylization."""
+    with gr.Blocks() as demo:
+        with gr.Column(elem_id="col-container"):
+            gr.HTML(
+                """
+                <style>
+                  #col-container { max-width: 1200px; margin: 0 auto; }
+                </style>
+                """
+            )
+            gr.Markdown(
+                """
+                # 🎥 ViBT: Vision Bridge Transformer at Scale  
+                <div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
+                <a href="https://arxiv.org/abs/2411.15098"><img src="https://img.shields.io/badge/ariXv-Paper-A42C25.svg" alt="arXiv"></a>
+                <a href="https://huggingface.co/spaces/Yuanshi/OminiControl"><img src="https://img.shields.io/badge/🤗OminiControl-Demo-ffbd45.svg" alt="HuggingFace"></a>
+                <a href="https://github.com/Yuanshi9815/OminiControl"><img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub"></a>
+                </div>
+                """
+            )
+
+            with gr.Row():
+                with gr.Column():
+                    input_video = gr.Video(label="Source Video", sources=["upload"])
+                with gr.Column():
+                    output_video = gr.Video(label="Stylized Result", interactive=False)
+
+            with gr.Row():
+                with gr.Column(scale=1) as control_col:
+                    with gr.Tabs():
+                        with gr.Tab("Quick Generate"):
+                            prompt_quick = gr.Dropdown(
+                                label="Style Instruction",
+                                choices=STYLE_CHOICES,
+                                value=STYLE_CHOICES[0],
+                                allow_custom_value=False,
+                            )
+
+                            with gr.Row():
+                                fast_btn = gr.Button("⚡ Fast Generate", variant="primary")
+                                balanced_btn = gr.Button("🎯 Balanced Generate", variant="primary")
+                                quality_btn = gr.Button("🌟 Quality Generate", variant="primary")
+
+                            _bind_preset_button(
+                                button=fast_btn,
+                                preset_key="Fast",
+                                inputs=[
+                                    input_video,
+                                    prompt_quick,
+                                ],
+                                output=output_video,
+                                extra_kwargs={"seed": None, "randomize_seed": True},
+                            )
+                            _bind_preset_button(
+                                button=balanced_btn,
+                                preset_key="Balanced",
+                                inputs=[
+                                    input_video,
+                                    prompt_quick,
+                                ],
+                                output=output_video,
+                                extra_kwargs={"seed": None, "randomize_seed": True},
+                            )
+                            _bind_preset_button(
+                                button=quality_btn,
+                                preset_key="Quality",
+                                inputs=[
+                                    input_video,
+                                    prompt_quick,
+                                ],
+                                output=output_video,
+                                extra_kwargs={"seed": None, "randomize_seed": True},
+                            )
+
+                        with gr.Tab("Advanced Settings"):
+                            prompt_adv = gr.Dropdown(
+                                label="Style Instruction",
+                                choices=STYLE_CHOICES,
+                                value=STYLE_CHOICES[0],
+                                allow_custom_value=True,
+                            )
+                            with gr.Row():
+                                shift_gamma = _create_slider(GAMMA_SLIDER)
+                                guidance_scale = _create_slider(GUIDANCE_SLIDER)
+
+                            with gr.Row():
+                                num_steps = _create_slider(STEP_SLIDER)
+                                randomize_seed_adv = gr.Checkbox(
+                                    label="Randomize Seed",
+                                    value=False,
+                                    info="Checked = new random seed each run. Uncheck to provide your own seed.",
+                                )
+
+                            seed_adv = gr.Number(
+                                label="Seed (used when Randomize is off)",
+                                value=42,
+                                precision=0,
+                            )
+
+                            run_btn = gr.Button("Generate", variant="primary")
+
+                            run_btn.click(
+                                fn=run_stylization,
+                                inputs=[
+                                    input_video,
+                                    prompt_adv,
+                                    shift_gamma,
+                                    num_steps,
+                                    guidance_scale,
+                                    seed_adv,
+                                    randomize_seed_adv,
+                                ],
+                                outputs=[output_video],
+                            )
+
+                with gr.Column(scale=1):
+                    gr.Examples(
+                        examples=EXAMPLE_INPUTS,
+                        inputs=[input_video, prompt_quick, prompt_adv],
+                        label="Example inputs",
+                    )
+
+    return demo
+
+
+demo = build_demo()
+
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()