import os # Video DiT + full-resolution VAE decode on a partitioned GPU: keep the caching # allocator from tripping over one huge contiguous allocation. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # MUST come before torch / any CUDA-touching import import contextlib import logging import random import re import sys import tempfile import time from pathlib import Path import torch import gradio as gr from PIL import Image, ImageOps from huggingface_hub import snapshot_download logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", stream=sys.stderr, ) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Weights # --------------------------------------------------------------------------- WAN_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B" REMIND_MODEL_ID = "AppliedIntuitionResearch/ReMind" logger.info("Downloading %s ...", WAN_MODEL_ID) WAN_MODEL_DIR = snapshot_download( repo_id=WAN_MODEL_ID, repo_type="model", local_dir="/tmp/checkpoints/Wan2.2-TI2V-5B", ) logger.info("Downloading %s ...", REMIND_MODEL_ID) REMIND_DIR = snapshot_download( repo_id=REMIND_MODEL_ID, repo_type="model", local_dir="/tmp/checkpoints/ReMind-5B", ) BASE_CHECKPOINT = os.path.join(REMIND_DIR, "ReMind-5B.safetensors") EMA_CHECKPOINT = os.path.join(REMIND_DIR, "ReMind-5b-dmd-ema.safetensors") # --------------------------------------------------------------------------- # The ReMind inference code ships alongside this app # --------------------------------------------------------------------------- REPO_ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(REPO_ROOT)) from omegaconf import OmegaConf from safetensors.torch import load_file as safe_load_file from pipeline.cache_adapter import InferencePipelineAdapter from pipeline.causal_inference import CausalInferencePipeline from pipeline.checkpoints import ( STUDENT_ADAPTER, configure_inference_lora, load_adapter_state_dict, load_generator_checkpoint, ) from pipeline.dmd_rollout import DMDInferenceRollout from pipeline.known_context import build_known_context from pipeline.remind_inference import ( build_pixel_cameras, build_prompt, image_to_repeated_video, pixel_to_latent_cameras, ) from utils.video_io import write_mp4 DEVICE = "cuda" DTYPE = torch.bfloat16 # Fixed rollout geometry — the released DMD-EMA checkpoint is trained for # exactly this shape (81 pixel frames -> 21 latent frames -> 7 chunks of 3). FRAMES = 81 HEIGHT = 480 WIDTH = 832 FPS = 16 CHUNK_SIZE = 3 SEED_LATENT_FRAMES = 1 # I2V: one clean seed latent frame DEFAULT_SEED = 20261700 MAX_SEED = 2**31 - 1 def _load_model_config(): cfg = OmegaConf.merge( OmegaConf.load(REPO_ROOT / "configs" / "default_config.yaml"), OmegaConf.load(REPO_ROOT / "configs" / "model_5b.yaml"), ) folder = str(WAN_MODEL_DIR) cfg.wan_model_folder = folder cfg.text_encoder_model_folder = folder cfg.vae_model_folder = folder if cfg.get("generator", {}).get("weight_list"): for weight in cfg.generator.weight_list: weight.path = folder return cfg # Build the pipeline exactly like the project-page runner # (pipeline/remind_inference.py::load_dmd_ema_pipeline): Wan base -> ReMind-5B # generator overlay -> cast to bf16 -> merge the rank-128 DMD-EMA student LoRA. # The cast has to happen before the merge; merging in fp32 and casting after # gives measurably different videos from the released EMA. logger.info("Building pipeline ...") pipeline = CausalInferencePipeline(_load_model_config(), device=DEVICE) # Move the finished sub-modules to "cuda" as we go: ZeroGPU intercepts the call # at module scope, so RAM is released component by component instead of holding # the generator, the 11 GB umt5-xxl encoder and the VAE all at once. pipeline.text_encoder.to(device=DEVICE, dtype=DTYPE).requires_grad_(False).eval() pipeline.vae.to(device=DEVICE, dtype=DTYPE).requires_grad_(False).eval() logger.info("Loading ReMind-5B generator ...") load_generator_checkpoint(pipeline.generator, BASE_CHECKPOINT, "generator") pipeline.generator.to(dtype=DTYPE) pipeline.generator.to(device=DEVICE).requires_grad_(False).eval() @contextlib.contextmanager def _cpu_rng_only(): """Hide CUDA from ``torch.random.fork_rng`` while the LoRA is attached. ``configure_inference_lora`` forks the CUDA RNG so its LoRA init is reproducible. On ZeroGPU no GPU is attached to this process, and ``torch.cuda.get_rng_state`` is one of the few calls the CUDA emulation layer does not intercept, so it reaches ``torch._C._cuda_init`` and dies. Answering "no CUDA" makes the helper take its own CPU-only branch instead. This cannot change the weights: ``load_adapter_state_dict`` overwrites every LoRA parameter from the EMA checkpoint and raises unless key coverage is exact, so nothing the init RNG produced survives the next two lines. """ real_is_available = torch.cuda.is_available torch.cuda.is_available = lambda: False try: yield finally: torch.cuda.is_available = real_is_available logger.info("Merging the DMD-EMA student LoRA ...") with _cpu_rng_only(): lora_model, lora_targets = configure_inference_lora( pipeline.generator.model, rank=128, alpha=128, dropout=0.0 ) load_adapter_state_dict(lora_model, safe_load_file(EMA_CHECKPOINT), STUDENT_ADAPTER) pipeline.generator.model = lora_model.merge_and_unload() pipeline.generator.requires_grad_(False).eval() pipeline.scheduler.set_timesteps(1000, training=True) DENOISING_STEPS = [ int(value) for value in pipeline.denoising_step_list.tolist() if int(value) > 0 ] logger.info( "Model ready — %d LoRA targets merged, %d-step schedule %s", len(lora_targets), len(DENOISING_STEPS), DENOISING_STEPS, ) # --------------------------------------------------------------------------- # Disturbance schedules # # ReMind's contribution is remembering what happened while it could not be # seen. Every disturbance below is expressed as CHUNK-LOCAL TEXT (and, for the # camera mode, a real camera trajectory through PM-RoPE) — never as a pixel # overlay. Chunks 3-5 of 7 carry the "hidden" phase, the last chunks ask for # recovery, and the scene state has to have advanced continuously. # --------------------------------------------------------------------------- MODE_CLEAN = "None — plain image-to-video" MODE_OCCLUDER = "Occluder — an object covers the view, then leaves" MODE_LIGHTS = "Lights out — the scene goes dark, then bright again" MODE_CAMERA = "Camera pan — the camera turns away, then comes back" MODES = [MODE_CLEAN, MODE_OCCLUDER, MODE_LIGHTS, MODE_CAMERA] MODE_TO_KIND = { MODE_CLEAN: "clean", MODE_OCCLUDER: "occluder", MODE_LIGHTS: "light", MODE_CAMERA: "camera", } SIDES = ["left", "right", "top", "bottom"] def _build_control( disturbance, occluder_description, enter_side, exit_side, camera_axis, camera_degrees, camera_trajectory, ): kind = MODE_TO_KIND.get(str(disturbance), "clean") if kind == "clean": return {"kind": "clean"} if kind == "light": return {"kind": "light", "recovery_chunks": 2} if kind == "occluder": label = (occluder_description or "brown cardboard box").strip() slug = re.sub(r"[^a-z0-9]+", "_", label.lower()).strip("_") or "occluder" return { "kind": "occluder", "label": label, "slug": slug, "enter_side": str(enter_side or "right"), "exit_side": str(exit_side or "left"), "recovery_chunks": 2, } control = { "kind": "camera", "axis": str(camera_axis or "yaw"), "degrees": float(camera_degrees), "return": True, } if camera_trajectory: # Paired InSpatio-style trajectory: extrinsics/intrinsics drive PM-RoPE # directly and the caption stays content-only. control["_trajectory_path"] = str(camera_trajectory) return control def _schedule_text(prompts, seed, disturbance): prompts = list(prompts) header = ( f"seed {seed} · {len(prompts)} chunks × {CHUNK_SIZE} latent frames · " f"{len(DENOISING_STEPS)}-step DMD rollout · {disturbance}" ) if len(set(prompts)) == 1: return f"{header}\n\nOne caption for every chunk:\n{prompts[0]}" prefix = os.path.commonprefix(prompts) lines = [header, "", f"Shared caption: {prefix.strip()}", ""] for index, text in enumerate(prompts): suffix = text[len(prefix) :].strip() lines.append(f"Chunk {index + 1}: {suffix or '(shared caption only)'}") return "\n".join(lines) # Measured on zero-a10g: 13.8-14.0 s for every mode (the rollout geometry is # fixed, so cost does not vary with the input). Kept tight so the demo does not # waste visitors' ZeroGPU quota. @spaces.GPU(duration=18) def generate( input_image: str, prompt: str, disturbance: str = MODE_CLEAN, seed: int = DEFAULT_SEED, camera_trajectory: str = None, occluder_description: str = "brown cardboard box", enter_side: str = "right", exit_side: str = "left", camera_axis: str = "yaw", camera_degrees: float = 20.0, randomize_seed: bool = False, progress=gr.Progress(track_tqdm=True), ): """Generate an 81-frame 832x480 video from one image with ReMind. Args: input_image: Path to the first frame of the video. prompt: Content-only temporal caption: what happens in the scene, start to end. Do not describe the camera or the disturbance. disturbance: Out-of-sight event to schedule mid-clip. One of "None — plain image-to-video", "Occluder — an object covers the view, then leaves", "Lights out — the scene goes dark, then bright again", "Camera pan — the camera turns away, then comes back". seed: RNG seed. camera_trajectory: Optional path to a camera .npz (keys `extrinsics` [T,3,4] or [T,4,4] and `intrinsics` [T,3,3]) used by the camera mode instead of the synthetic pan. occluder_description: Appearance of the occluding object. enter_side: Side the occluder enters from. exit_side: Side the occluder leaves towards. camera_axis: "yaw" (horizontal) or "pitch" (vertical) synthetic pan. camera_degrees: Peak rotation of the synthetic pan, in degrees. randomize_seed: Draw a fresh random seed for this run. Returns: The generated MP4 and the per-chunk prompt schedule that produced it. """ if input_image is None: raise gr.Error("Please provide an input image.") if not prompt or not prompt.strip(): raise gr.Error("Please describe what happens in the scene.") seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed) started = time.perf_counter() if isinstance(input_image, Image.Image): image = ImageOps.exif_transpose(input_image).convert("RGB") else: image = ImageOps.exif_transpose(Image.open(str(input_image))).convert("RGB") image = image.resize((WIDTH, HEIGHT), Image.Resampling.LANCZOS) control = _build_control( disturbance, occluder_description, enter_side, exit_side, camera_axis, camera_degrees, camera_trajectory, ) progress(0.05, desc="Encoding the input image") source_video = image_to_repeated_video( image, frames=FRAMES, device=DEVICE, dtype=DTYPE ) source_latent = pipeline.vae.encode_to_latent(source_video).to( device=DEVICE, dtype=DTYPE ) latent_frames = int(source_latent.shape[1]) latent_frames -= latent_frames % CHUNK_SIZE if latent_frames <= 1: raise gr.Error(f"Invalid latent length {latent_frames}.") source_latent = source_latent[:, :latent_frames] num_chunks = latent_frames // CHUNK_SIZE pixel_extrinsics, pixel_intrinsics = build_pixel_cameras( control, frames=FRAMES, height=HEIGHT, width=WIDTH ) viewmats, intrinsics = pixel_to_latent_cameras( pixel_extrinsics.to(DEVICE), pixel_intrinsics.to(DEVICE), latent_frames ) prompt_schedule = build_prompt( str(prompt).strip(), control, pixel_frames=FRAMES, latent_frames=latent_frames, chunk_size=CHUNK_SIZE, ) conditional = pipeline.text_encoder(text_prompts=[prompt_schedule]) if isinstance(prompt_schedule, list): conditional["prompt_chunk_size"] = CHUNK_SIZE adapter = InferencePipelineAdapter( pipeline, chunk_size=CHUNK_SIZE, context_timestep=0 ) rollout = DMDInferenceRollout( adapter, pipeline.scheduler, DENOISING_STEPS, chunk_size=CHUNK_SIZE, context_timestep=0, ) known_latents, known_mask = build_known_context(source_latent, SEED_LATENT_FRAMES) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) noise = torch.randn_like(source_latent) progress(0.15, desc=f"Rolling out {num_chunks} chunks") generated = rollout.rollout( noise, conditional, {"viewmats": viewmats, "Ks": intrinsics}, known_latents=known_latents, known_mask=known_mask, exit_indices=[len(DENOISING_STEPS) - 1] * num_chunks, ) adapter.clear_cache() progress(0.85, desc="Decoding frames") decoded = pipeline.vae.decode_to_pixel(generated, use_cache=False) decoded = decoded.float().mul(0.5).add(0.5).clamp(0, 1) if decoded.shape[1] in (1, 3, 4): frames = decoded[0].permute(1, 0, 2, 3) elif decoded.shape[2] in (1, 3, 4): frames = decoded[0] else: raise gr.Error(f"Cannot identify decoded layout {tuple(decoded.shape)}.") output_path = tempfile.NamedTemporaryFile( suffix=".mp4", delete=False, dir="/tmp" ).name write_mp4(output_path, frames, fps=FPS) logger.info( "generated %s in %.1fs (%s, seed=%d)", output_path, time.perf_counter() - started, MODE_TO_KIND.get(str(disturbance), "clean"), seed, ) schedule = prompt_schedule if isinstance(prompt_schedule, list) else [prompt_schedule] return output_path, _schedule_text(schedule, seed, str(disturbance)) # --------------------------------------------------------------------------- # Examples — the seven image-to-video presets from the ReMind repo # (examples/presets/*.yaml), with their original captions and seeds. # --------------------------------------------------------------------------- LATTE_PROMPT = ( "An espresso machine steadily pours into a short glass cup: at the start the " "cup shows a dark coffee base topped by a thin crema disc with latte-art. As " "the pour continues a narrow stream and occasional droplet fall from the " "spout into the center, causing concentric ripples and a swirling pattern in " "the crema. Midway the light-colored foam layer grows and spreads outward, " "covering more of the darker coffee underneath. By the end the foam has built " "up into a slightly domed, glossy surface that nearly reaches the rim while " "the darker liquid is largely hidden beneath; the falling stream momentarily " "forms a thin column that contacts and feeds the expanding foam. The scene " "shows pouring, surface deformation, spreading foam, and gentle mixing of layers." ) PANCAKE_PROMPT = ( "A thick, pale batter is poured from a scoop into a hot frying pan, forming a " "high, coiled mound at the start. As more batter flows, the stream lands on " "the mound and the mixture spreads outward, flattening against the pan and " "creating concentric folds and swirls that ripple away from center. Midway the " "pouring narrows and the batter continues to flow into the center while the " "existing layers slump and smooth under their own weight. By the end the batter " "has spread into a broad, low disc filling most of the pan surface, the central " "swirl partly blended into a smoother pancake-like layer, and the pouring tool " "has been withdrawn. The motion shows viscous flow, impact, and lateral " "spreading of the liquid batter." ) WHISKING_PROMPT = ( "A metal whisk stirs a thick pale batter sitting on a wooden board. At the " "start a rounded mound of batter rests with the whisk entering near its top; " "the whisk repeatedly presses into the center and drags material outward, " "carving concentric swirls and a shallow trough. Midway the batter spreads and " "the whisk tilts and shifts rightward while its wires create radial grooves; " "the mixture behaves viscously, flowing slowly back into the stirred channel " "without splashing. By the end the batter is noticeably more spread with a " "distinct central depression and spiral ripples, the whisk positioned to the " "right and a single hand steadying it at the board edge. The scene shows " "continuous stirring-driven deformation and outward spreading of the batter." ) CEREAL_PROMPT = ( "A steady stream of small brown seeds is poured onto a shallow plate lined with " "a white paper layer, building a growing conical heap at center. At the start " "there is a thin, low layer of seeds spread across the paper; as pouring " "continues the falling grains strike the pile, bounce and slide outward, and " "accumulate into a taller, denser mound. Midway the central peak rises while " "impacts send individual seeds hopping and rolling to form a shoulder and a low " "rim at the plate edge. By the end the flow produces a pronounced central cone " "with scattered grains around its base; repeated collisions compact the pile and " "increase lateral scatter, while the falling stream remains above the apex " "feeding the mound." ) DOG_PROMPT = ( "A small dog paces and shifts its weight across a living-room floor in front of " "a shelving unit. At the start it stands with a fairly level back, head turned " "slightly to the right, then begins a series of low, quick steps: hindquarters " "lower and tail lifting as it crouches and pushes off. Midway the dog bends its " "legs and leans forward, front paws skidding slightly as it executes a rapid " "lateral motion and turns its head. By the end the animal is more hunched, tail " "up and blurred from wagging, with its body rotated and moved toward the right " "side of the scene. The shelves and objects behind remain static throughout." ) DRY_ICE_PROMPT = ( "A small clear glass initially holds a pale yellow liquid while thick white " "vapor and frothy foam curl and billow above the rim. The foam rapidly grows, " "rolling into larger lobes that spill and drip over the glass edge; wisps of " "vapor detach and fall. Midway the foam forms a dense cap, pressing down and " "streaking along the glass sides while the interior liquid becomes visually " "darker and more amber as the foam collapses and mixes. By the end the cup is " "largely enveloped by a soft, blobby white foam that has overflowed into rounded " "puddles on the table; amber liquid fills most of the vessel and streams of foam " "run down and pool at the base. The sequence shows expansion, overflow, " "deformation, and runoff of bubbly, wet material." ) FLOUR_PROMPT = ( "A fine white granular material is being poured onto a shallow wooden tray, " "forming a conical pile that steadily grows. At the start a modest mound sits on " "the tray while a diffuse stream of grains falls into its center. As pouring " "continues the apex rises and the slope steepens; individual grains bounce " "outward on impact, creating a scattered ring of particles across the tray. " "Intermittently a larger compact clump falls through the stream and strikes near " "the summit, causing short-lived sprays of grains that roll down the pile. By " "the end the heap is noticeably taller and broader, with more grains spread " "along the tray edges and a concentrated falling mass suspended briefly above " "the peak. The motion is dominated by continuous inflow, impacts, and small " "granular avalanches down the slopes." ) IMAGE_EXAMPLES = [ ["assets/latte_art_espresso.png", LATTE_PROMPT, MODE_OCCLUDER, 20261406], ["assets/pancake_batter_pour.png", PANCAKE_PROMPT, MODE_OCCLUDER, 20261418], ["assets/dog_tail_wag_bark.png", DOG_PROMPT, MODE_CLEAN, 20261700], ["assets/dry_ice_fog_bowl.png", DRY_ICE_PROMPT, MODE_CLEAN, 20261708], ["assets/flour_sugar_pour.png", FLOUR_PROMPT, MODE_CLEAN, 20261714], ] CAMERA_EXAMPLES = [ [ "assets/whisking_batter.png", WHISKING_PROMPT, MODE_CAMERA, 20261310, "assets/whisking_batter_camera.npz", ], [ "assets/cereal_bowl_pour.png", CEREAL_PROMPT, MODE_CAMERA, 20261317, "assets/cereal_bowl_pour_camera.npz", ], ] CSS = """ #col-container { max-width: 1200px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ # Gradio 6 moved `theme` and `css` from the Blocks constructor to launch(). with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # 🫥 ReMind — video generation with dynamic memory [ReMind](https://huggingface.co/AppliedIntuitionResearch/ReMind) is a Wan2.2-TI2V-5B world model taught to **remember what it can no longer see**. Schedule a disturbance mid-clip — an object covers the lens, the lights go out, the camera turns away — and the scene keeps evolving out of sight, so when the view comes back the state has moved on instead of resetting. 81 frames · 832×480 · 16 fps · four-step DMD rollout over seven 3-frame chunks. """ ) with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(label="First frame", type="filepath", height=260) prompt = gr.Textbox( label="What happens in the scene", placeholder=( "A content-only caption from start to end — the material, the " "motion, how the state advances. No camera or occluder talk." ), lines=5, ) disturbance = gr.Dropdown( MODES, value=MODE_CLEAN, label="Out-of-sight event", info="Scheduled over chunks 3–5 of 7, with two recovery chunks after it.", ) run_btn = gr.Button("Generate video", variant="primary") with gr.Column(scale=1): output_video = gr.Video(label="Generated video", autoplay=True) schedule_box = gr.Textbox( label="Chunk prompt schedule", info="What the model was actually told, chunk by chunk.", lines=9, max_lines=9, buttons=["copy"], interactive=False, ) with gr.Accordion("Advanced settings", open=False): with gr.Row(): seed = gr.Number( label="Seed", value=DEFAULT_SEED, precision=0, minimum=0 ) randomize_seed = gr.Checkbox(label="Randomize seed", value=False) with gr.Row(): occluder_description = gr.Textbox( label="Occluder appearance", value="brown cardboard box", info="Used by the occluder schedule only.", ) enter_side = gr.Dropdown( SIDES, value="right", label="Occluder enters from" ) exit_side = gr.Dropdown(SIDES, value="left", label="Occluder exits to") with gr.Row(): camera_axis = gr.Dropdown( ["yaw", "pitch"], value="yaw", label="Synthetic pan axis" ) camera_degrees = gr.Slider( -40.0, 40.0, value=20.0, step=1.0, label="Synthetic pan angle (degrees)", ) camera_trajectory = gr.File( label="Camera trajectory (.npz, optional)", file_types=[".npz"], type="filepath", ) gr.Markdown("### Examples from the ReMind project page") gr.Examples( examples=IMAGE_EXAMPLES, inputs=[input_image, prompt, disturbance, seed], outputs=[output_video, schedule_box], fn=generate, cache_examples=True, cache_mode="lazy", label="Image-to-video (two occluder-recovery cases, three clean)", examples_per_page=5, ) gr.Examples( examples=CAMERA_EXAMPLES, inputs=[input_image, prompt, disturbance, seed, camera_trajectory], outputs=[output_video, schedule_box], fn=generate, cache_examples=True, cache_mode="lazy", label="Camera-controlled (paired trajectories, PM-RoPE)", ) gr.Markdown( """ The occluder and the lights-out event are **never painted into the pixels** — they are described to the generator through chunk-local text, exactly as in the paper. The camera mode instead feeds a real trajectory (extrinsics + intrinsics) into ReMind's camera-phase RoPE; leave the trajectory file empty to use a synthetic pan-away-and-return. 📄 [Paper](https://huggingface.co/papers/2605.25333) · 💻 [Code](https://github.com/Applied-Intuition-Open-Source/ReMind) · 🤗 [Weights](https://huggingface.co/AppliedIntuitionResearch/ReMind) (CC-BY-NC-4.0, research use) · base model [Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) """ ) run_btn.click( fn=generate, inputs=[ input_image, prompt, disturbance, seed, camera_trajectory, occluder_description, enter_side, exit_side, camera_axis, camera_degrees, randomize_seed, ], outputs=[output_video, schedule_box], api_name="generate", ) demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)