| """ViRDM causal4 demo with native Wan VAE streaming on ZeroGPU.""" |
|
|
| import spaces |
|
|
| import gc |
| import hashlib |
| import importlib.metadata |
| import os |
| from pathlib import Path |
| import queue |
| import secrets |
| import tempfile |
| import threading |
| import traceback |
|
|
| import av |
| import gradio as gr |
| from huggingface_hub import hf_hub_download, snapshot_download |
| import numpy as np |
| from omegaconf import OmegaConf |
| from packaging.version import Version |
| import torch |
|
|
| from pipeline import CausalInferencePipeline |
| from scripts.export_videos import load_generator_state |
| from utils.misc import set_seed |
| from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper |
| from wan.modules.attention import require_flash_attention_2 |
|
|
| DTYPE = torch.bfloat16 |
| FPS = 15 |
| INITIAL_BUFFER_BLOCKS = 4 |
| OUTPUTS = Path(tempfile.gettempdir()) / "virdm-streams" |
| OUTPUTS.mkdir(exist_ok=True) |
|
|
| WAN_REVISION = '37ec512624d61f7aa208f7ea8140a131f93afc9a' |
| WAN_FILES = {'config.json': 'ab37994c43740513f94b3ba6233a784035a67b43c8cde83c8f31aa90468c67ce', |
| 'diffusion_pytorch_model.safetensors': '96b6b242ca1c2f24e9d02cd6596066fab6d310e2d7538f33ae267cb18d957e8f', |
| 'models_t5_umt5-xxl-enc-bf16.pth': '7cace0da2b446bbbbc57d031ab6cf163a3d59b366da94e5afe36745b746fd81d', |
| 'Wan2.1_VAE.pth': '38071ab59bd94681c686fa51d75a1968f64e470262043be31f7a094e442fd981', |
| 'google/umt5-xxl/tokenizer.json': '6e197b4d3dbd71da14b4eb255f4fa91c9c1f2068b20a2de2472967ca3d22602b', |
| 'google/umt5-xxl/spiece.model': 'e3909a67b780650b35cf529ac782ad2b6b26e6d1f849d3fbb6a872905f452458', |
| 'google/umt5-xxl/special_tokens_map.json': '7b8a9f5040adb67b5805abdfd42c1f8d0f3d0e711f10726580eb3789cd0ad61d', |
| 'google/umt5-xxl/tokenizer_config.json': 'ed9a3a8b0faa71a70a32847e0435fe036e6e112d4df4edb7bb48a921e344dc05'} |
| CHECKPOINT_REVISION = 'b77aad226f97cf15e36caa9afac7db07713c96c7' |
| CHECKPOINT_FILE = 'checkpoints/virdm_causal4_dynamic_step20.pt' |
| CHECKPOINT_SHA256 = 'bc66de80b2c28ffdbb9f21b764436a35bd1d4643c393f2ebbfe9130b6d6320fd' |
| |
| EXAMPLES = [['A vibrant and lively vlog-style photo of a corgi in tropical Maui, showcasing the dog ' |
| 'energetically filming itself on a sandy beach. The corgi stands on the shore, one paw slightly ' |
| 'lifted, with a joyful and curious expression. It wears a colorful collar and a small backpack ' |
| 'camera slung over its neck. The background features a lush, palm-fringed beach with clear ' |
| 'turquoise waters and a bright blue sky. The photo has a warm, natural lighting effect, ' |
| 'capturing the corgi from a slightly elevated angle, emphasizing its playful and adventurous ' |
| 'spirit.'], |
| ['Early spring in a mountain tea valley, a young man in a cream knit sweater and navy vest stands ' |
| 'close to camera, ultra‑clean 4K. He plucks tender leaves into a woven basket, then wipes dew ' |
| 'from his fingers as sun breaks through thin fog. The camera floats low over bright rows of tea, ' |
| 'then rises to reveal layered ridges and a river ribbon shining below. Clothing stitches, leaf ' |
| 'veins, and mist particles remain razor sharp and natural.'], |
| ['A stylish woman strolls down a bustling Tokyo street, the warm glow of neon lights and animated ' |
| 'city signs casting vibrant reflections. She wears a sleek black leather jacket paired with a ' |
| 'flowing red dress and black boots, her black purse slung over her shoulder. Sunglasses perched ' |
| 'on her nose and a bold red lipstick add to her confident, casual demeanor. The street is damp ' |
| 'and reflective, creating a mirror-like effect that enhances the colorful lights and shadows. ' |
| 'Pedestrians move about, adding to the lively atmosphere. The scene is captured in a dynamic ' |
| 'medium shot with the woman walking slightly to one side, highlighting her graceful strides.'], |
| ['A cinematic closeup and detailed portrait of a reindeer standing in a snowy forest at sunset. ' |
| 'The lighting is gorgeous and soft, with a golden backlight creating a warm and dreamy effect. ' |
| 'Soft bokeh and lens flares add a magical touch, enhancing the cinematic quality of the image. ' |
| 'The reindeer has a gentle expression, its fur glistening in the fading light. The background ' |
| 'features a serene snowy landscape with tall trees silhouetted against the orange and pink hues ' |
| 'of the setting sun. The color grade is rich and magical, capturing the essence of a winter ' |
| 'wonderland at twilight. A close-up shot from a slightly elevated angle.'], |
| ['A detailed and heartwarming wildlife photograph capturing a mother bird tenderly feeding her ' |
| 'chicks in a cozy nest. The mother bird gently places food into the wide-open beaks of her ' |
| 'chirping chicks, who eagerly await their meal. Her feathers are soft and fluffy, and she has a ' |
| 'gentle, attentive expression. The chicks have small, round heads with wide-open beaks and big, ' |
| 'curious eyes. The nest is lined with soft grass and twigs, and the background features a ' |
| 'blurred forest scene with dappled sunlight filtering through the leaves. The photo has a ' |
| 'natural, documentary style. A close-up shot from a slightly elevated angle, focusing on the ' |
| 'interaction between the mother bird and her chicks.'], |
| ['A digital illustration in a whimsical cartoon style of a llama coding and typing on his laptop ' |
| 'in a cozy cafe. The llama has a friendly expression, with large, expressive eyes and a gentle ' |
| 'smile. It wears a colorful patterned scarf and a pair of round glasses perched on its nose. The ' |
| 'cafe setting includes a wooden table, a few chairs, and a window with a view of a bustling ' |
| 'street outside. The background is filled with the soft glow of ambient lighting and hints of ' |
| "other patrons. The llama's fingers dance over the keyboard, with a cup of steaming coffee " |
| "nearby. A close-up shot from a slightly elevated angle, capturing the llama's focused and " |
| 'engaged posture.'], |
| ['A high-tech, cartoon-style illustration of a computer mouse with legs running on a treadmill. ' |
| 'The mouse has a round body with a pair of tiny legs, one in front and one behind, and large, ' |
| 'round eyes with a determined expression. It is wearing a small, colorful running outfit with ' |
| 'stripes and a tail that wags as it runs. The treadmill is set up in a modern, minimalist room ' |
| "with sleek, metallic walls and a few scattered tech gadgets in the background. The mouse's " |
| 'movements are lively and energetic, with its paws gripping the treadmill belt tightly. A ' |
| "dynamic side view, capturing the mouse's mid-run position."], |
| ['A serene orchard scene in the style of a gentle watercolor painting, with trees heavily laden ' |
| 'with fragrant blossoms in soft pastel shades of pink and white. Bees buzz busily, darting from ' |
| 'flower to flower in a display of natural harmony. The sun filters through the branches, casting ' |
| 'dappled shadows on the ground. A gentle breeze rustles the leaves, adding a sense of movement ' |
| 'and life to the scene. The background features a soft blue sky with fluffy white clouds. A ' |
| 'medium shot with a slightly elevated perspective, capturing both the detailed flowers and the ' |
| 'vast expanse of the orchard.'], |
| ['In spring under a tunnel of wisteria, a dancer in a pale lavender dress and silver hairpin ' |
| 'fills the frame in cinematic 4K. She moves through slow turns, sleeves fluttering, then pauses ' |
| 'to let blossoms brush her shoulders as the camera pushes close to her serene eyes. Behind her, ' |
| 'a sunlit garden pond mirrors purple clusters and fresh green leaves. Natural tones, crisp petal ' |
| 'detail, and soft lens bloom create a dreamy realism that feels delicate and moving.'], |
| ['A surreal and haunting digital art piece in a dreamlike style, featuring mushrooms sprouting ' |
| 'from the base of a decaying bookshelf. The mushrooms have vibrant, colorful caps in shades of ' |
| 'orange, yellow, and green, contrasting sharply with the worn, weathered wood of the bookshelf. ' |
| 'The bookshelf is covered in dust and peeling paint, with several books lying open and pages ' |
| 'torn. The background is dimly lit, with flickering light casting shadows and highlighting the ' |
| 'decay. The mushrooms appear to be growing from cracks and crevices in the wood, giving the ' |
| 'scene a mysterious and eerie feel. A close-up shot from a slightly elevated angle, emphasizing ' |
| 'the textures and colors.'], |
| ['In late autumn in a candle lit chapel, a violinist in a black satin dress fills the foreground ' |
| 'in cinematic 4K. She plays a slow phrase, then accelerates into a shimmering run as the camera ' |
| 'moves close to bow hair and fingerwork. Outside stained glass, rain taps softly and leaves ' |
| 'swirl in the courtyard. Warm flame light paints her face and instrument, with true tones and ' |
| 'gentle film grain making the performance hauntingly beautiful.'], |
| ['In late autumn inside an art studio, a sculptor in a black turtleneck and clay dusted apron ' |
| 'fills the foreground in HDR 4K. She shapes a face in wet clay, then steps back to study it as ' |
| 'the camera moves from her hands to her focused eyes. Outside tall windows, rain taps softly and ' |
| 'leaves cling to glass. Warm sunlight turns the room amber, revealing tool marks, dust motes, ' |
| 'and rich textures with gentle film grain.']] |
|
|
| HEADER = """ |
| <div class="virdm-header"> |
| <h1>⚡ ViRDM <span class="divider">·</span> <span class="gradient-text">4-step</span> <span class="divider">·</span> Causal Video</h1> |
| <div class="virdm-badges"> |
| <span class="metric">VBench Total <b>84.87</b></span> |
| <span class="metric">Peak training memory <b>48.3 GB/GPU</b></span> |
| <span class="metric">Post-training <b>~2 h · 8 × A100</b></span> |
| <span class="metric"><b>20</b> generator updates</span> |
| </div> |
| <p class="single-gpu">Also supports training on a single 80GB GPU.</p> |
| <div class="virdm-links"> |
| <a href="https://arxiv.org/abs/2609.28923" target="_blank">📄 Paper</a> |
| <a href="https://neu-vi.github.io/ViRDM/" target="_blank">🌐 Project</a> |
| <a href="https://github.com/neu-vi/ViRDM" target="_blank">⌘ Code</a> |
| <a href="https://huggingface.co/cr8br0ze/ViRDM" target="_blank">🤗 Model</a> |
| </div> |
| </div> |
| """ |
|
|
| CSS = """ |
| .gradio-container { max-width: 1120px !important; margin: 0 auto !important; } |
| .virdm-header { text-align: center; padding: 28px 18px 30px; margin: 4px 0 18px; |
| border-radius: 24px; background: linear-gradient(130deg, rgba(234,88,12,.09), rgba(255,152,88,.07), rgba(255,237,213,.16)); } |
| .virdm-header h1 { font-size: clamp(26px, 3.5vw, 38px); line-height: 1.3; font-weight: 750; margin: 0 0 14px; } |
| .gradient-text { background: linear-gradient(90deg, #ff6733, #f59e0b); background-clip: text; -webkit-background-clip: text; color: transparent; } |
| .divider { opacity: .35; font-weight: 400; } |
| .virdm-badges, .virdm-links { display: flex; flex-wrap: wrap; justify-content: center; gap: 9px; } |
| .metric { border: 1px solid rgba(214,91,49,.24); background: rgba(255,152,88,.09); padding: 6px 12px; border-radius: 999px; font-size: 13px; } |
| .single-gpu { font-size: 12px; opacity: .72; margin: 10px 0 0; } |
| .virdm-links { margin-top: 19px; } |
| .virdm-links a { color: inherit; text-decoration: none; font-size: 13px; font-weight: 600; |
| border: 1px solid rgba(120,120,140,.23); border-radius: 9px; padding: 6px 13px; } |
| .virdm-links a:hover { border-color: #d65b31; color: #d65b31; } |
| #generate-btn { background: linear-gradient(100deg, #ff6733, #ff9858); border: none; |
| color: #13233e; font-size: 16px; font-weight: 700; min-height: 49px; border-radius: 12px; } |
| #generate-btn:hover { filter: brightness(1.07); } |
| #live-output { border-radius: 16px; overflow: hidden; min-height: 290px; } |
| #live-output video { border-radius: 14px; } |
| #prompt-examples > .gallery { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 8px; } |
| #prompt-examples > .gallery > .gallery-item { width: 100%; min-width: 0; margin: 0; text-align: left; } |
| footer { display: none !important; } |
| |
| .streaming-note { text-align: center; font-size: 12px; opacity: .7; margin: 14px 0 4px; } |
| """ |
|
|
|
|
| def verify_file(path, expected): |
| digest = hashlib.sha256() |
| with open(path, "rb") as handle: |
| for block in iter(lambda: handle.read(8 * 1024 * 1024), b""): |
| digest.update(block) |
| if digest.hexdigest() != expected: |
| raise RuntimeError(f"SHA-256 mismatch for {Path(path).name}") |
|
|
|
|
| def load_pipeline(): |
| installed = importlib.metadata.version("flash-attn") |
| if Version(installed).public != "2.8.3.post1": |
| raise RuntimeError(f"Unexpected FlashAttention version: {installed}") |
| require_flash_attention_2(installed) |
| wan_root = snapshot_download( |
| "Wan-AI/Wan2.1-T2V-1.3B", revision=WAN_REVISION, |
| allow_patterns=list(WAN_FILES), |
| ) |
| for filename, sha256 in WAN_FILES.items(): |
| verify_file(Path(wan_root) / filename, sha256) |
| os.environ["VIRDM_WAN_MODEL_ROOT"] = wan_root |
| checkpoint = hf_hub_download( |
| "cr8br0ze/ViRDM", CHECKPOINT_FILE, revision=CHECKPOINT_REVISION, |
| ) |
| verify_file(checkpoint, CHECKPOINT_SHA256) |
| config = OmegaConf.create({ |
| "model_kwargs": {"timestep_shift": 5.0}, |
| "denoising_step_list": [1000, 750, 500, 250], |
| "first_frame_denoising_step_list": [1000, 750, 500, 250], |
| "warp_denoising_step": True, |
| "num_frame_per_block": 3, |
| "independent_first_frame": False, |
| "context_noise": 0, |
| }) |
| generator = WanDiffusionWrapper(**config.model_kwargs) |
| generator.load_state_dict(load_generator_state(checkpoint, use_ema=False), |
| strict=True, assign=True) |
| generator.eval().requires_grad_(False).to(device="cuda", dtype=DTYPE) |
| gc.collect() |
| text_encoder = WanTextEncoder(dtype=DTYPE, assign_load=True) |
| text_encoder.eval().requires_grad_(False).to(device="cuda", dtype=DTYPE) |
| vae = WanVAEWrapper().eval().requires_grad_(False).to(device="cuda", dtype=DTYPE) |
| pipeline = CausalInferencePipeline(config, device="cuda", generator=generator, |
| text_encoder=text_encoder, vae=vae) |
| return pipeline.eval().requires_grad_(False) |
|
|
|
|
| def encode_chunk(frames): |
| path = OUTPUTS / f"{secrets.token_hex(12)}.ts" |
| with av.open(str(path), mode="w", format="mpegts") as container: |
| stream = container.add_stream("libx264", rate=FPS) |
| stream.width, stream.height = 832, 480 |
| stream.pix_fmt = "yuv420p" |
| stream.options = {"preset": "ultrafast", "tune": "zerolatency", "crf": "18"} |
| for index, pixels in enumerate(frames): |
| frame = av.VideoFrame.from_ndarray(pixels, format="rgb24") |
| frame.pts = index |
| for packet in stream.encode(frame): |
| container.mux(packet) |
| for packet in stream.encode(): |
| container.mux(packet) |
| return str(path) |
|
|
|
|
| def stream_video(pipeline, prompt, seed): |
| messages = queue.Queue(maxsize=2) |
| cancelled = threading.Event() |
| block_count = 0 |
|
|
| def send(item): |
| while not cancelled.is_set(): |
| try: |
| messages.put(item, timeout=0.1) |
| return |
| except queue.Full: |
| continue |
| raise RuntimeError("Streaming request cancelled") |
|
|
| def on_generator_call(module, args, kwargs, output): |
| nonlocal block_count |
| |
| if not bool(torch.all(kwargs["timestep"] == 0).item()): |
| return |
| if cancelled.is_set(): |
| raise RuntimeError("Streaming request cancelled") |
| decoded = pipeline.vae.decode_to_pixel(kwargs["noisy_image_or_video"], use_cache=True) |
| pixels = (decoded[0].mul(0.5).add(0.5).clamp(0, 1).mul(255) |
| .to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()) |
| expected = 9 if block_count == 0 else 12 |
| if pixels.shape != (expected, 480, 832, 3): |
| raise RuntimeError(f"Unexpected streaming block shape: {pixels.shape}") |
| block_count += 1 |
| send(pixels) |
|
|
| def produce(): |
| hook = None |
| try: |
| with torch.inference_mode(): |
| set_seed(seed) |
| pipeline.vae.model.clear_cache() |
| pipeline.kv_cache1 = None |
| pipeline.crossattn_cache = None |
| noise = torch.randn([1, 21, 16, 60, 104], device="cuda", dtype=DTYPE, |
| generator=torch.Generator(device="cuda").manual_seed(seed)) |
| hook = pipeline.generator.register_forward_hook(on_generator_call, with_kwargs=True) |
| pipeline.inference(noise, [prompt], return_video=False) |
| if block_count != 7: |
| raise RuntimeError("Streaming did not produce all seven blocks") |
| send(None) |
| except BaseException as exc: |
| if not cancelled.is_set(): |
| send(exc) |
| finally: |
| if hook is not None: |
| hook.remove() |
| pipeline.vae.model.clear_cache() |
| pipeline.kv_cache1 = None |
| pipeline.crossattn_cache = None |
|
|
| worker = threading.Thread(target=produce, daemon=True) |
| worker.start() |
| buffered = [] |
| started = False |
| try: |
| while True: |
| item = messages.get() |
| if isinstance(item, BaseException): |
| raise item |
| if item is None: |
| break |
| buffered.append(item) |
| if not started and len(buffered) < INITIAL_BUFFER_BLOCKS: |
| continue |
| frames = np.concatenate(buffered, axis=0) |
| buffered.clear() |
| started = True |
| yield encode_chunk(frames) |
| finally: |
| cancelled.set() |
| worker.join(timeout=10) |
|
|
|
|
| PIPELINE, STARTUP_ERROR = None, None |
| try: |
| print("Loading ViRDM...", flush=True) |
| PIPELINE = load_pipeline() |
| print("ViRDM ready.", flush=True) |
| except Exception as exc: |
| traceback.print_exc() |
| STARTUP_ERROR = f"{type(exc).__name__}: {exc}" |
|
|
|
|
| @spaces.GPU(duration=20, size="xlarge") |
| def generate(prompt, seed): |
| if PIPELINE is None: |
| raise gr.Error(f"Model startup failed: {STARTUP_ERROR}") |
| prompt = (prompt or "").strip() |
| if not prompt: |
| raise gr.Error("Enter a prompt or choose an example.") |
| if len(prompt) > 4000: |
| raise gr.Error("Please keep the prompt below 4,000 characters.") |
| requested_seed = int(seed) |
| seed = secrets.randbelow(2**31) if requested_seed == -1 else requested_seed |
| if not 0 <= seed < 2**32: |
| raise gr.Error("Use -1 for a random seed, or a seed between 0 and 4294967295.") |
| for path in stream_video(PIPELINE, prompt, seed): |
| yield path, requested_seed |
|
|
|
|
| theme = gr.themes.Soft(primary_hue="orange", secondary_hue="amber", neutral_hue="slate") |
| with gr.Blocks(title="ViRDM · Four-step Causal Video", theme=theme, css=CSS, |
| delete_cache=(3600, 3600)) as demo: |
| gr.HTML(HEADER) |
| if STARTUP_ERROR: |
| gr.Markdown(f"**Model initialization failed.**\n```text\n{STARTUP_ERROR}\n```") |
| with gr.Row(equal_height=False): |
| with gr.Column(scale=1): |
| prompt = gr.Textbox(label="Prompt", value=EXAMPLES[0][0], lines=5) |
| seed = gr.Number(value=-1, precision=0, label="Seed", info="Use -1 for a random seed.") |
| run = gr.Button("✨ Generate video", variant="primary", elem_id="generate-btn", |
| interactive=PIPELINE is not None) |
| gr.Examples(examples=EXAMPLES, inputs=prompt, examples_per_page=len(EXAMPLES), |
| elem_id="prompt-examples") |
| with gr.Column(scale=1): |
| output = gr.Video(label="Generated video", streaming=True, autoplay=True, |
| loop=False, elem_id="live-output") |
| gr.HTML('<p class="streaming-note">This demo uses a four-block startup buffer for smoother ' |
| 'streaming on Hugging Face ZeroGPU; local deployments can use one block when ' |
| 'generation keeps pace.</p>') |
| run.click(generate, [prompt, seed], [output, seed], concurrency_limit=1, |
| concurrency_id="gpu", api_name="generate") |
|
|
| demo.queue(max_size=8, default_concurrency_limit=1).launch(show_error=True) |
|
|
