Upload folder using huggingface_hub

README.md

@@ -4,25 +4,22 @@ emoji: 🌊
 colorFrom: pink
 colorTo: indigo
 sdk: gradio
-sdk_version: 6.15.1
+sdk_version: 6.20.0
 app_file: app.py
-short_description: Flow-based world model for robot video prediction
-python_version: "3.12"
+short_description: Predict future RGB video + optical flow from an image
+python_version: "3.10"
 startup_duration_timeout: 1h
 ---
 
 # FlowWAM
 
-**FlowWAM** generates future RGB and optical-flow video from a scene image and a robot task instruction.
+Interactive demo of **FlowWAM: Optical Flow as a Unified Action Representation
+for World Action Models**.
 
-Given an initial scene observation and a natural-language instruction, the dual-stream video diffusion model (based on Wan2.2-TI2V-5B) jointly predicts:
+A single dual-stream video diffusion model (built on Wan2.2-TI2V-5B) jointly
+predicts a **future RGB video** and its **optical-flow field** from one input
+image and a short text instruction.
 
-- **Future RGB frames** — the predicted visual trajectory
-- **Optical flow frames** — per-pixel motion representing the action signal
-
-## Model
-
-- **Backbone**: Wan2.2-TI2V-5B (dual-stream video DiT)
-- **Checkpoint**: `YixiangChen/FlowWAM` (flowwam_robotwin)
-- **Paper**: [FlowWAM: Optical Flow as a Unified Action Representation for World Action Models](https://arxiv.org/abs/2607.13017)
-- **Project Page**: [flow-wam.github.io](https://flow-wam.github.io/)
+- Paper: https://huggingface.co/papers/2607.13017
+- Code: https://github.com/YixiangChen515/FlowWAM
+- Weights: https://huggingface.co/YixiangChen/FlowWAM

app.py

@@ -1,84 +1,81 @@
 import os
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-import spaces
 import sys
-import tempfile
 import time
-import numpy as np
+import tempfile
+
+import spaces  # noqa: E402  (must precede torch / CUDA imports)
 import torch
+import numpy as np
 import gradio as gr
 from PIL import Image
-from tqdm import tqdm
+from huggingface_hub import hf_hub_download
 
-# ── Make the diffsynth library importable ──────────────────────────────
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+# Make the vendored diffsynth package importable.
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+if SCRIPT_DIR not in sys.path:
+    sys.path.insert(0, SCRIPT_DIR)
 
-from diffsynth.models.wan_video_dit_dual_stream import FlowStreamModule, init_flow_stream
+from diffsynth.models.utils import load_state_dict
+from diffsynth.models.wan_video_dit_dual_stream import init_flow_stream
 from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig
 from diffsynth.pipelines.wan_video_dual_stream import model_fn_wan_video_dual_stream
-from diffsynth.schedulers.flow_match import FlowMatchScheduler
 from diffsynth.data.video import save_video
-from diffsynth.models.utils import load_state_dict
 
-# ── Constants ───────────────────────────────────────────────────────────
-CHECKPOINT_ID = "YixiangChen/FlowWAM"
-CHECKPOINT_FILE = "flowwam_robotwin.safetensors"
-CAMERA_PREFIX = (
-    "A multi-view video of an aloha robot in T-shape layout: "
-    "the top row shows the full-size rear camera view, "
-    "the bottom-left shows the half-size left arm camera view, "
-    "and the bottom-right shows the half-size right arm camera view. "
-    "The robot is performing the following task: "
-)
-SIZE_W, SIZE_H = 320, 256
-NUM_VIDEO_FRAMES = 9
-VIDEO_INFERENCE_STEPS = 25
-SIGMA_SHIFT = 5.0
-
-# ── Model loading (module scope) ────────────────────────────────────────
-log_msgs = []
-def _log(msg):
-    log_msgs.append(msg)
-    print(msg, flush=True)
-
-_log("Loading Wan2.2-TI2V-5B pipeline (VAE + T5 + DiT) ...")
+# ----------------------------------------------------------------------------
+# Model setup (module scope — ZeroGPU packs weights to disk at startup).
+# ----------------------------------------------------------------------------
+BASE_MODEL = "Wan-AI/Wan2.2-TI2V-5B"
+TOKENIZER_MODEL = "Wan-AI/Wan2.1-T2V-1.3B"
+FLOWWAM_REPO = "YixiangChen/FlowWAM"
+FLOWWAM_CKPT = "flowwam_worldarena_stage1.safetensors"
+
+MODELS_DIR = os.path.join(SCRIPT_DIR, "models")
+os.makedirs(MODELS_DIR, exist_ok=True)
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+
+def _mc(pattern, offload="cpu"):
+    return ModelConfig(
+        model_id=BASE_MODEL,
+        origin_file_pattern=pattern,
+        offload_device=offload,
+        local_model_path=MODELS_DIR,
+        download_resource="huggingface",
+    )
+
+
+print("Loading Wan2.2-TI2V-5B dual-stream pipeline (VAE + T5 + DiT) ...", flush=True)
 pipe = WanVideoPipeline.from_pretrained(
-    torch_dtype=torch.bfloat16,
-    device="cuda",
+    torch_dtype=DTYPE,
+    device=DEVICE,
     model_configs=[
-        ModelConfig(
-            model_id="Wan-AI/Wan2.2-TI2V-5B",
-            origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth",
-            offload_device="cpu",
-        ),
-        ModelConfig(
-            model_id="Wan-AI/Wan2.2-TI2V-5B",
-            origin_file_pattern="diffusion_pytorch_model*.safetensors",
-            offload_device="cpu",
-        ),
-        ModelConfig(
-            model_id="Wan-AI/Wan2.2-TI2V-5B",
-            origin_file_pattern="Wan2.2_VAE.pth",
-            offload_device="cpu",
-        ),
+        _mc("models_t5_umt5-xxl-enc-bf16.pth"),
+        _mc("diffusion_pytorch_model*.safetensors"),
+        _mc("Wan2.2_VAE.pth"),
     ],
     tokenizer_config=ModelConfig(
-        model_id="Wan-AI/Wan2.1-T2V-1.3B",
+        model_id=TOKENIZER_MODEL,
         origin_file_pattern="google/*",
+        local_model_path=MODELS_DIR,
+        download_resource="huggingface",
     ),
+    redirect_common_files=False,
 )
 
-_log("Initialising FlowStreamModule ...")
+# Flow stream: deep-copied patch-embed + head from the DiT.
 flow_stream = init_flow_stream(pipe.dit)
 
-_log("Downloading and loading FlowWAM checkpoint ...")
-from huggingface_hub import hf_hub_download
-ckpt_path = hf_hub_download(CHECKPOINT_ID, CHECKPOINT_FILE, repo_type="model")
+# Load the FlowWAM checkpoint: DiT + flow_stream keys (no action_expert in
+# the world-model stage-1 checkpoint).
+print(f"Downloading FlowWAM checkpoint {FLOWWAM_CKPT} ...", flush=True)
+ckpt_path = hf_hub_download(FLOWWAM_REPO, FLOWWAM_CKPT)
 state_dict = load_state_dict(ckpt_path)
 
-dit_keys = {}
-flow_keys = {}
+dit_keys, flow_keys = {}, {}
 for k, v in state_dict.items():
     if k.startswith("action_expert."):
         continue
@@ -87,160 +84,137 @@ for k, v in state_dict.items():
     else:
         dit_keys[k] = v
 
+# Params trained in fp32 (modulation / time-MLP / LayerNorm) — restore later.
 fp32_dit_values = {k: v.clone() for k, v in dit_keys.items() if v.dtype == torch.float32}
 
 missing, unexpected = pipe.dit.load_state_dict(dit_keys, strict=False)
-_log(f"DiT (full): loaded {len(dit_keys) - len(unexpected)} keys, {len(missing)} missing, {len(unexpected)} unexpected")
-
-if flow_keys:
-    missing_f, unexpected_f = flow_stream.load_state_dict(flow_keys, strict=False)
-    _log(f"FlowStream (full): loaded {len(flow_keys) - len(unexpected_f)} keys, {len(missing_f)} missing, {len(unexpected_f)} unexpected")
-
-# Apply fp32 modulation restoration
-from diffsynth.vram_management.layers import AutoWrappedLinear, WanAutoCastLayerNorm
-param_map = dict(pipe.dit.named_parameters())
-restored = 0
-for key, fp32_value in fp32_dit_values.items():
-    if key in param_map:
-        param_map[key].data = fp32_value.to(device=param_map[key].device)
-        restored += param_map[key].numel()
-
-for seq_module in [pipe.dit.time_embedding, pipe.dit.time_projection]:
-    for sub in seq_module.modules():
-        if isinstance(sub, AutoWrappedLinear):
-            sub.offload_dtype = torch.float32
-            sub.onload_dtype = torch.float32
-            sub.computation_dtype = torch.float32
-
-def _pre_hook(_mod, args):
-    return tuple(a.float() if isinstance(a, torch.Tensor) else a for a in args)
-
-def _post_hook(_mod, _args, output):
-    return output.bfloat16() if isinstance(output, torch.Tensor) else output
-
-for seq_module in [pipe.dit.time_embedding, pipe.dit.time_projection]:
-    seq_module.register_forward_pre_hook(_pre_hook)
-    seq_module.register_forward_hook(_post_hook)
-
-for module in pipe.dit.modules():
-    if isinstance(module, WanAutoCastLayerNorm):
-        module.offload_dtype = torch.float32
-        module.onload_dtype = torch.float32
-
-_log(f"[FP32Modulation] Restored {restored:,} fp32 params")
-
-# Enable VRAM management (needed for 48 GB ZeroGPU large tier)
+print(f"DiT (full): loaded {len(dit_keys) - len(unexpected)} keys, "
+      f"{len(missing)} missing, {len(unexpected)} unexpected", flush=True)
+missing, unexpected = flow_stream.load_state_dict(flow_keys, strict=False)
+print(f"FlowStream (full): loaded {len(flow_keys) - len(unexpected)} keys, "
+      f"{len(missing)} missing, {len(unexpected)} unexpected", flush=True)
+
 pipe.enable_vram_management()
-flow_stream = flow_stream.to(device="cuda", dtype=torch.bfloat16).eval()
-_log("Model loaded successfully.")
-
-
-# ── T-shape tiling helper ───────────────────────────────────────────────
-def tshape_tile(head_img, left_img, right_img):
-    """T-shape spatial concatenation: head on top, left+right half-res below."""
-    orig_h, orig_w = head_img.shape[:2]
-    half_h, half_w = orig_h // 2, orig_w // 2
-    import cv2
-    left_half = cv2.resize(left_img, (half_w, half_h))
-    right_half = cv2.resize(right_img, (half_w, half_h))
-    bottom = np.hstack([left_half, right_half])
-    return np.vstack([head_img, bottom])
-
-
-# ── Inference ──────────────────────────────────────────────────────────
-@spaces.GPU(duration=300)
-def generate_video(
-    input_image,
-    instruction: str,
-    seed: int = 1,
-    num_video_frames: int = NUM_VIDEO_FRAMES,
-    video_inference_steps: int = VIDEO_INFERENCE_STEPS,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """Generate future RGB and optical-flow video frames from a single input image.
-
-    Given a scene image and a language instruction describing the robot task,
-    FlowWAM's dual-stream video diffusion jointly predicts the future RGB
-    trajectory and the corresponding optical-flow video.
+
+
+def _apply_fp32_modulation(dit, fp32_state_values):
+    """Restore fp32 precision for modulation / time-MLP / LayerNorm params."""
+    from diffsynth.vram_management.layers import AutoWrappedLinear, WanAutoCastLayerNorm
+    param_map = dict(dit.named_parameters())
+    for key, fp32_value in fp32_state_values.items():
+        if key in param_map:
+            param_map[key].data = fp32_value.to(device=param_map[key].device)
+    for seq_module in [dit.time_embedding, dit.time_projection]:
+        for sub in seq_module.modules():
+            if isinstance(sub, AutoWrappedLinear):
+                sub.offload_dtype = torch.float32
+                sub.onload_dtype = torch.float32
+                sub.computation_dtype = torch.float32
+
+    def _pre_hook(_mod, args):
+        return tuple(a.float() if isinstance(a, torch.Tensor) else a for a in args)
+
+    def _post_hook(_mod, _args, output):
+        return output.bfloat16() if isinstance(output, torch.Tensor) else output
+
+    for seq_module in [dit.time_embedding, dit.time_projection]:
+        seq_module.register_forward_pre_hook(_pre_hook)
+        seq_module.register_forward_hook(_post_hook)
+    for module in dit.modules():
+        if isinstance(module, WanAutoCastLayerNorm):
+            module.offload_dtype = torch.float32
+            module.onload_dtype = torch.float32
+
+
+if fp32_dit_values:
+    _apply_fp32_modulation(pipe.dit, fp32_dit_values)
+
+flow_stream = flow_stream.to(device=DEVICE, dtype=DTYPE).eval()
+print("FlowWAM pipeline ready.", flush=True)
+
+
+# ----------------------------------------------------------------------------
+# Inference — dual-stream world-model rollout (stage 1 only).
+# ----------------------------------------------------------------------------
+def _estimate(image, instruction, num_frames=49, num_inference_steps=25,
+              sigma_shift=5.0, seed=1, *args, **kwargs):
+    steps = int(num_inference_steps)
+    return min(160, 45 + int(steps * 4.0))
+
+
+@spaces.GPU(duration=_estimate)
+@torch.no_grad()
+def generate(image, instruction, num_frames=49, num_inference_steps=25,
+             sigma_shift=5.0, seed=1,
+             progress=gr.Progress(track_tqdm=True)):
+    """Predict a future RGB video and its optical-flow field from one image + instruction.
 
     Args:
-        input_image: The initial scene image (any resolution; will be resized).
-        instruction: Natural-language description of the robot task.
-        seed: Random seed for reproducibility.
-        num_video_frames: Number of future video frames to generate (default 9).
-        video_inference_steps: Number of diffusion denoising steps (default 25).
+        image: the conditioning first frame (PIL image).
+        instruction: text describing the action / motion to imagine.
+        num_frames: number of frames to generate (4k+1).
+        num_inference_steps: dual-stream denoising steps.
+        sigma_shift: flow-match scheduler sigma shift.
+        seed: RNG seed.
+
+    Returns:
+        (rgb_video_path, flow_video_path): mp4 files for the predicted future
+        RGB frames and the predicted optical-flow field.
     """
-    t0 = time.perf_counter()
-
-    if input_image is None:
+    if image is None:
         raise gr.Error("Please provide an input image.")
+    instruction = (instruction or "").strip()
 
-    w, h = SIZE_W, SIZE_H
     device = pipe.device
     dtype = pipe.torch_dtype
     vae_z_dim = getattr(pipe.vae, "z_dim", 16)
-
-    # ── Prepare the T-shape tiled image ────────────────────────────────
-    img_np = np.array(input_image.convert("RGB"))
-    head = np.array(Image.fromarray(img_np).resize((w, h), Image.BICUBIC))
-    # For single-image input, create left/right views by cropping the image
-    half_h = h // 2
-    half_w = w // 2
-    left_crop = img_np[:, :img_np.shape[1]//2]
-    right_crop = img_np[:, img_np.shape[1]//2:]
-    if left_crop.size == 0 or right_crop.size == 0:
-        left_crop = img_np
-        right_crop = img_np
-    left = np.array(Image.fromarray(left_crop).resize((w, h), Image.BICUBIC))
-    right = np.array(Image.fromarray(right_crop).resize((w, h), Image.BICUBIC))
-    tiled = tshape_tile(head, left, right)
-
-    tiled_h, tiled_w = tiled.shape[:2]
+    seed = int(seed)
+    num_frames = int(num_frames)
+
+    # ---- Resize conditioning frame to a valid grid ----
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    image = image.convert("RGB")
+    w, h = image.size
+    # Keep a compact aspect-preserving size (~320x256 like the reference).
+    target_w = 320
+    target_h = max(1, round(h * target_w / w))
     tiled_h, tiled_w, video_frames = pipe.check_resize_height_width(
-        tiled_h, tiled_w, num_video_frames)
-    tiled_pil = Image.fromarray(tiled).resize((tiled_w, tiled_h), Image.BICUBIC)
-    flow_h, flow_w = tiled_h, tiled_w
+        target_h, target_w, num_frames)
+    cond_pil = image.resize((tiled_w, tiled_h), Image.BICUBIC)
 
-    # ── Text encoding ──────────────────────────────────────────────────
-    video_prompt = CAMERA_PREFIX + instruction
+    # ---- Text encoding ----
     pipe.load_models_to_device(["text_encoder"])
-    context = pipe.prompter.encode_prompt(video_prompt, positive=True, device=device)
+    context = pipe.prompter.encode_prompt(instruction, positive=True, device=device)
 
-    # ── VAE encode the single conditioning frame ───────────────────────
+    # ---- VAE encode the single conditioning frame (RGB + zero-flow prefix) ----
     pipe.load_models_to_device(["vae"])
     upscale = pipe.vae.upsampling_factor
     T_lat = (video_frames - 1) // 4 + 1
     rgb_H_lat = tiled_h // upscale
     rgb_W_lat = tiled_w // upscale
-    flow_H_lat = flow_h // upscale
-    flow_W_lat = flow_w // upscale
 
-    rgb_vid = pipe.preprocess_video([tiled_pil])
+    rgb_vid = pipe.preprocess_video([cond_pil])
     rgb_prefix = pipe.vae.encode(rgb_vid, device=device).to(dtype=dtype, device=device)
 
-    zero_flow_pil = Image.new("RGB", (flow_w, flow_h), (255, 255, 255))
+    zero_flow_pil = Image.new("RGB", (tiled_w, tiled_h), (255, 255, 255))
     flow_vid = pipe.preprocess_video([zero_flow_pil])
     flow_prefix = pipe.vae.encode(flow_vid, device=device).to(dtype=dtype, device=device)
 
     rgb_noise_shape = (1, vae_z_dim, T_lat, rgb_H_lat, rgb_W_lat)
-    flow_noise_shape = (1, vae_z_dim, T_lat, flow_H_lat, flow_W_lat)
+    flow_noise_shape = (1, vae_z_dim, T_lat, rgb_H_lat, rgb_W_lat)
     rgb_noise = pipe.generate_noise(rgb_noise_shape, seed=seed, rand_device="cpu").to(dtype=dtype, device=device)
     rgb_noise[:, :, :1] = rgb_prefix
-    flow_noise = pipe.generate_noise(
-        flow_noise_shape, seed=(seed + 1) if seed is not None else None, rand_device="cpu"
-    ).to(dtype=dtype, device=device)
+    flow_noise = pipe.generate_noise(flow_noise_shape, seed=seed + 1, rand_device="cpu").to(dtype=dtype, device=device)
     flow_noise[:, :, :1] = flow_prefix
 
     rgb_latents = rgb_noise.clone()
     flow_latents = flow_noise.clone()
 
-    # ── Stage 1: dual-stream video denoising ────────────────────────────
-    pipe.scheduler.set_timesteps(video_inference_steps, shift=SIGMA_SHIFT)
+    # ---- Dual-stream video denoising ----
+    pipe.scheduler.set_timesteps(int(num_inference_steps), shift=float(sigma_shift))
     pipe.load_models_to_device(pipe.in_iteration_models)
-
-    for progress_id, timestep in enumerate(
-        tqdm(pipe.scheduler.timesteps, desc="Video DiT Denoising")
-    ):
+    for progress_id, timestep in enumerate(pipe.scheduler.timesteps):
         t_tensor = timestep.unsqueeze(0).to(dtype=dtype, device=device)
         rgb_pred, flow_pred = model_fn_wan_video_dual_stream(
             dit=pipe.dit,
@@ -257,86 +231,76 @@ def generate_video(
         rgb_latents[:, :, :1] = rgb_prefix
         flow_latents[:, :, :1] = flow_prefix
 
-    # ── Decode both streams ─────────────────────────────────────────────
+    # ---- Decode both streams ----
     pipe.load_models_to_device(["vae"])
-    rgb_video = pipe.vae_output_to_video(pipe.vae.decode(rgb_latents, device=device))
-    flow_video = pipe.vae_output_to_video(pipe.vae.decode(flow_latents, device=device))
+    rgb_frames = pipe.vae_output_to_video(pipe.vae.decode(rgb_latents, device=device))
+    flow_frames = pipe.vae_output_to_video(pipe.vae.decode(flow_latents, device=device))
     pipe.load_models_to_device([])
 
-    # ── Save videos ─────────────────────────────────────────────────────
-    rgb_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    flow_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    save_video(rgb_video, rgb_path, fps=12)
-    save_video(flow_video, flow_path, fps=12)
-
-    elapsed = time.perf_counter() - t0
-    _log(f"Generation complete in {elapsed:.1f}s")
+    rgb_path = tempfile.NamedTemporaryFile(suffix="_rgb.mp4", delete=False).name
+    flow_path = tempfile.NamedTemporaryFile(suffix="_flow.mp4", delete=False).name
+    save_video(rgb_frames, rgb_path, fps=12)
+    save_video(flow_frames, flow_path, fps=12)
+    return rgb_path, flow_path
 
-    return rgb_path, flow_path, f"{elapsed:.1f}s"
 
-
-# ── Gradio UI ───────────────────────────────────────────────────────────
+# ----------------------------------------------------------------------------
+# UI
+# ----------------------------------------------------------------------------
 CSS = """
 #col-container { max-width: 1100px; margin: 0 auto; }
 .dark .gradio-container { color: var(--body-text-color); }
 """
 
+DESCRIPTION = """
+# FlowWAM — Optical Flow as a Unified Action Representation
+
+A single dual-stream video diffusion model (built on **Wan2.2-TI2V-5B**) that
+jointly predicts a **future RGB video** and its **optical-flow field** from one
+image and a short text instruction. From the paper
+*FlowWAM: Optical Flow as a Unified Action Representation for World Action Models*.
+
+Give it a starting frame and describe the motion — it imagines how the scene
+evolves and the dense per-pixel motion (flow) that drives it.
+
+[Paper](https://huggingface.co/papers/2607.13017) · [Code](https://github.com/YixiangChen515/FlowWAM) · [Weights](https://huggingface.co/YixiangChen/FlowWAM)
+"""
+
 with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(
-            "# 🌊 FlowWAM: Optical Flow as a Unified Action Representation for World Action Models\n"
-            "Generate future RGB and optical-flow video from a scene image and a robot task instruction.\n\n"
-            "[Paper](https://arxiv.org/abs/2607.13017) | "
-            "[Project Page](https://flow-wam.github.io/) | "
-            "[Model](https://huggingface.co/YixiangChen/FlowWAM)"
-        )
-
+        gr.Markdown(DESCRIPTION)
         with gr.Row():
-            with gr.Column(scale=1):
-                input_image = gr.Image(
-                    label="Input Scene Image",
-                    type="pil",
-                    height=256,
-                )
+            with gr.Column():
+                image = gr.Image(label="Input image (first frame)", type="pil")
                 instruction = gr.Textbox(
-                    label="Task Instruction",
-                    placeholder="e.g., pick up the block and place it on the right",
-                    lines=2,
+                    label="Instruction",
+                    placeholder="describe the motion, e.g. 'the astronaut walks forward'",
                 )
-                run_btn = gr.Button("Generate Video", variant="primary")
-
-            with gr.Column(scale=1):
-                rgb_video_out = gr.Video(label="Predicted Future RGB")
-                flow_video_out = gr.Video(label="Predicted Optical Flow")
-                time_out = gr.Textbox(label="Inference Time", interactive=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            seed_input = gr.Number(label="Seed", value=1, precision=0)
-            num_frames_input = gr.Number(
-                label="Number of Video Frames", value=NUM_VIDEO_FRAMES,
-                precision=0, minimum=5, maximum=49,
-            )
-            num_steps_input = gr.Slider(
-                label="Denoising Steps", value=VIDEO_INFERENCE_STEPS,
-                minimum=5, maximum=50, step=1,
-            )
+                run = gr.Button("Generate", variant="primary")
+            with gr.Column():
+                rgb_out = gr.Video(label="Predicted future RGB")
+                flow_out = gr.Video(label="Predicted optical flow")
+        with gr.Accordion("Advanced settings", open=False):
+            num_frames = gr.Slider(13, 49, value=49, step=4, label="Frames (4k+1)")
+            num_inference_steps = gr.Slider(10, 40, value=25, step=1, label="Denoising steps")
+            sigma_shift = gr.Slider(1.0, 8.0, value=5.0, step=0.5, label="Sigma shift")
+            seed = gr.Number(value=1, precision=0, label="Seed")
+
+        inputs = [image, instruction, num_frames, num_inference_steps, sigma_shift, seed]
+        run.click(generate, inputs=inputs, outputs=[rgb_out, flow_out], api_name="generate")
 
         gr.Examples(
             examples=[
-                ["example_input.jpg", "pick up the block and place it in the basket"],
-                ["example_input.jpg", "grab the bottle and move it to the right side"],
+                ["astronaut.jpg", "the astronaut walks across the surface"],
+                ["husky_dog.jpg", "the dog runs forward"],
+                ["hot_air_balloon.jpg", "the balloon drifts upward across the sky"],
             ],
-            inputs=[input_image, instruction],
-            outputs=[rgb_video_out, flow_video_out, time_out],
-            fn=generate_video,
+            inputs=[image, instruction],
+            outputs=[rgb_out, flow_out],
+            fn=generate,
             cache_examples=True,
             cache_mode="lazy",
         )
 
-    run_btn.click(
-        fn=generate_video,
-        inputs=[input_image, instruction, seed_input, num_frames_input, num_steps_input],
-        outputs=[rgb_video_out, flow_video_out, time_out],
-    )
-
-demo.launch(mcp_server=True)
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

requirements.txt

@@ -1,18 +1,16 @@
-transformers
+torch
 torchvision
-imageio
-imageio-ffmpeg
-pillow
-regex
+transformers
+accelerate
 safetensors
 einops
+imageio
+imageio[ffmpeg]
+opencv-python-headless
+pillow
 sentencepiece
 protobuf
-modelscope
 ftfy
-pynvml
-accelerate
-opencv-python
-pyyaml
-tqdm
-numpy<2
+regex
+modelscope
+numpy<2