import os os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" import spaces import sys import time import tempfile import numpy as np import torch import torch.nn.functional as F import gradio as gr from PIL import Image from einops import rearrange from tqdm import tqdm # ------------------------------------------------------------------ # # Model + config setup (module scope, eagerly on GPU) # ------------------------------------------------------------------ # MODEL_ID = "gangweix/next-forcing-base" from huggingface_hub import snapshot_download _model_path = snapshot_download( MODEL_ID, repo_type="model", allow_patterns=[ "transformer/*", "vae/*", "text_encoder/*", "tokenizer/*", ], ) from wan_va.modules.utils import ( WanVAEStreamingWrapper, load_text_encoder, load_tokenizer, load_transformer, load_vae, ) from wan_va.utils.scheduler import FlowMatchScheduler from wan_va.utils.utils import get_mesh_id, data_seq_to_patch DTYPE = torch.bfloat16 DEVICE = "cuda" # ---- Demo config (matches va_demo_cfg.py) ---- CONFIG = dict( attn_window=30, frame_chunk_size=4, env_type="none", height=256, width=256, action_dim=30, action_per_frame=8, obs_cam_keys=["observation.images.top", "observation.images.wrist"], guidance_scale=5, action_guidance_scale=1, num_inference_steps=5, video_exec_step=-1, action_num_inference_steps=10, snr_shift=5.0, action_snr_shift=1.0, patch_size=(1, 2, 2), used_action_channel_ids=list(range(0, 5)) + list(range(28, 29)), action_norm_method="quantiles", norm_stat={ "q01": [ -90.60303497314453, -98.73043060302734, -79.9008560180664, 48.95470428466797, -32.794578552246094, ] + [0.0] * 23 + [0.8250824809074402, 0], "q99": [ 71.735107421875, 65.89081573486328, 92.87967681884766, 100.0, 22.784151077270508, ] + [0.0] * 23 + [100.0, 0], }, ) # Inverse action channel mapping inverse_used_action_channel_ids = [len(CONFIG["used_action_channel_ids"])] * CONFIG[ "action_dim" ] for i, j in enumerate(CONFIG["used_action_channel_ids"]): inverse_used_action_channel_ids[j] = i CONFIG["inverse_used_action_channel_ids"] = inverse_used_action_channel_ids # ---- Load model components ---- vae = load_vae(os.path.join(_model_path, "vae"), torch_dtype=DTYPE, torch_device=DEVICE) streaming_vae = WanVAEStreamingWrapper(vae) tokenizer = load_tokenizer(os.path.join(_model_path, "tokenizer")) text_encoder = load_text_encoder( os.path.join(_model_path, "text_encoder"), torch_dtype=DTYPE, torch_device=DEVICE ) transformer = load_transformer( os.path.join(_model_path, "transformer"), torch_dtype=DTYPE, torch_device=DEVICE, attn_mode="torch", disable_mcp=True, ) transformer.eval().requires_grad_(False) scheduler = FlowMatchScheduler(shift=CONFIG["snr_shift"], sigma_min=0.0, extra_one_step=True) action_scheduler = FlowMatchScheduler( shift=CONFIG["action_snr_shift"], sigma_min=0.0, extra_one_step=True ) scheduler.set_timesteps(1000, training=True) action_scheduler.set_timesteps(1000, training=True) action_mask = torch.zeros([CONFIG["action_dim"]]).bool() action_mask[CONFIG["used_action_channel_ids"]] = True actions_q01 = torch.tensor(CONFIG["norm_stat"]["q01"], dtype=torch.float32).reshape(-1, 1, 1) actions_q99 = torch.tensor(CONFIG["norm_stat"]["q99"], dtype=torch.float32).reshape(-1, 1, 1) from diffusers.video_processor import VideoProcessor video_processor = VideoProcessor(vae_scale_factor=1) # ------------------------------------------------------------------ # # Inference helpers # ------------------------------------------------------------------ # def _get_t5_prompt_embeds(prompt, max_sequence_length=512): from diffusers.pipelines.wan.pipeline_wan import prompt_clean prompt_list = [prompt] if isinstance(prompt, str) else prompt prompt_list = [prompt_clean(u) for u in prompt_list] batch_size = len(prompt_list) text_inputs = tokenizer( prompt_list, padding="max_length", max_length=max_sequence_length, truncation=True, add_special_tokens=True, return_attention_mask=True, return_tensors="pt", ) text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() text_encoder_device = next(text_encoder.parameters()).device prompt_embeds = text_encoder( text_input_ids.to(text_encoder_device), mask.to(text_encoder_device) ).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=DTYPE, device=DEVICE) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( [ torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds ], dim=0, ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, 1, 1) prompt_embeds = prompt_embeds.view(batch_size, seq_len, -1) return prompt_embeds.to(DEVICE) def encode_prompt(prompt): prompt_embeds = _get_t5_prompt_embeds(prompt) neg_prompt_embeds = _get_t5_prompt_embeds("") return prompt_embeds, neg_prompt_embeds def normalize_latents(latents, latents_mean, latents_std): latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(device=latents.device) latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device) latents = ((latents.float() - latents_mean) * latents_std).to(latents) return latents def encode_obs(images_dict_list): """Encode observation images into latent space. Args: images_dict_list: list of dicts, each mapping cam_key -> np.ndarray(H,W,3) uint8 """ images = images_dict_list if not isinstance(images, list): images = [images] if len(images) < 1: return None videos = [] for k_i, k in enumerate(CONFIG["obs_cam_keys"]): height_i, width_i = CONFIG["height"], CONFIG["width"] history_video_k = ( torch.from_numpy(np.stack([each[k] for each in images])) .float() .permute(3, 0, 1, 2) ) history_video_k = F.interpolate( history_video_k, size=(height_i, width_i), mode="bilinear", align_corners=False, ).unsqueeze(0) videos.append(history_video_k) videos = torch.cat(videos, dim=0) / 255.0 * 2.0 - 1.0 vae_device = next(streaming_vae.vae.parameters()).device videos_chunk = videos.to(vae_device).to(DTYPE) enc_out = streaming_vae.encode_chunk(videos_chunk) mu, logvar = torch.chunk(enc_out, 2, dim=1) latents_mean = torch.tensor(vae.config.latents_mean).to(mu.device) latents_std = torch.tensor(vae.config.latents_std).to(mu.device) mu_norm = normalize_latents(mu, latents_mean, 1.0 / latents_std) video_latent = torch.cat(mu_norm.split(1, dim=0), dim=-1) return video_latent.to(DEVICE) def _repeat_input_for_cfg(input_dict, use_cfg, prompt_embeds, negative_prompt_embeds): if use_cfg: input_dict["noisy_latents"] = input_dict["noisy_latents"].repeat(2, 1, 1, 1, 1) input_dict["text_emb"] = torch.cat( [prompt_embeds.to(DTYPE).clone(), negative_prompt_embeds.to(DTYPE).clone()], dim=0, ) input_dict["grid_id"] = input_dict["grid_id"][None].repeat(2, 1, 1) input_dict["timesteps"] = input_dict["timesteps"][None].repeat(2, 1) else: input_dict["grid_id"] = input_dict["grid_id"][None] input_dict["timesteps"] = input_dict["timesteps"][None] return input_dict def _prepare_latent_input( latent_model_input, action_model_input, latent_t=0, action_t=0, latent_cond=None, action_cond=None, frame_st_id=0, patch_size=(1, 2, 2), prompt_embeds=None, use_cfg=False, negative_prompt_embeds=None, ): input_dict = dict() if latent_model_input is not None: input_dict["latent_res_lst"] = { "noisy_latents": latent_model_input, "timesteps": torch.ones( [latent_model_input.shape[2]], dtype=torch.float32, device=DEVICE ) * latent_t, "grid_id": get_mesh_id( latent_model_input.shape[-3] // patch_size[0], latent_model_input.shape[-2] // patch_size[1], latent_model_input.shape[-1] // patch_size[2], 0, 1, frame_st_id, ).to(DEVICE), "text_emb": prompt_embeds.to(DTYPE).clone(), } if latent_cond is not None: input_dict["latent_res_lst"]["noisy_latents"][:, :, 0:1] = latent_cond[:, :, 0:1] input_dict["latent_res_lst"]["timesteps"][0:1] *= 0 if action_model_input is not None: input_dict["action_res_lst"] = { "noisy_latents": action_model_input, "timesteps": torch.ones( [action_model_input.shape[2]], dtype=torch.float32, device=DEVICE ) * action_t, "grid_id": get_mesh_id( action_model_input.shape[-3], action_model_input.shape[-2], action_model_input.shape[-1], 1, 1, frame_st_id, action=True, ).to(DEVICE), "text_emb": prompt_embeds.to(DTYPE).clone(), } if action_cond is not None: input_dict["action_res_lst"]["noisy_latents"][:, :, 0:1] = action_cond[:, :, 0:1] input_dict["action_res_lst"]["timesteps"][0:1] *= 0 input_dict["action_res_lst"]["noisy_latents"][:, ~action_mask] *= 0 return input_dict def infer_chunk( init_latent, frame_st_id, prompt_embeds, negative_prompt_embeds, use_cfg, guidance_scale, action_guidance_scale, num_chunks_to_infer, ): """Generate one video chunk (video latents + action latents).""" frame_chunk_size = CONFIG["frame_chunk_size"] latent_height = CONFIG["height"] // 16 latent_width = (CONFIG["width"] // 16) * len(CONFIG["obs_cam_keys"]) latents = torch.randn( 1, 48, frame_chunk_size, latent_height, latent_width, device=DEVICE, dtype=DTYPE ) actions = torch.randn( 1, CONFIG["action_dim"], frame_chunk_size, CONFIG["action_per_frame"], 1, device=DEVICE, dtype=DTYPE, ) video_inference_step = CONFIG["num_inference_steps"] action_inference_step = CONFIG["action_num_inference_steps"] video_step = CONFIG["video_exec_step"] scheduler.set_timesteps(video_inference_step) action_scheduler.set_timesteps(action_inference_step) timesteps = scheduler.timesteps action_timesteps = action_scheduler.timesteps timesteps = F.pad(timesteps, (0, 1), mode="constant", value=0) if video_step != -1: timesteps = timesteps[:video_step] action_timesteps = F.pad(action_timesteps, (0, 1), mode="constant", value=0) with torch.no_grad(): # 1. Video generation loop for i, t in enumerate(timesteps): last_step = i == len(timesteps) - 1 latent_cond = init_latent[:, :, 0:1].to(DTYPE) if frame_st_id == 0 else None input_dict = _prepare_latent_input( latents, None, t, t, latent_cond, None, frame_st_id=frame_st_id, patch_size=CONFIG["patch_size"], prompt_embeds=prompt_embeds, use_cfg=use_cfg, negative_prompt_embeds=negative_prompt_embeds, ) video_noise_pred = transformer( _repeat_input_for_cfg( input_dict["latent_res_lst"], use_cfg, prompt_embeds, negative_prompt_embeds, ), update_cache=1 if last_step else 0, cache_name="pos", action_mode=False, ) if not last_step or video_step != -1: video_noise_pred = data_seq_to_patch( CONFIG["patch_size"], video_noise_pred, frame_chunk_size, latent_height, latent_width, batch_size=2 if use_cfg else 1, ) if guidance_scale > 1: video_noise_pred = video_noise_pred[1:] + guidance_scale * ( video_noise_pred[:1] - video_noise_pred[1:] ) else: video_noise_pred = video_noise_pred[:1] latents = scheduler.step(video_noise_pred, t, latents, return_dict=False) latents[:, :, 0:1] = ( latent_cond if frame_st_id == 0 else latents[:, :, 0:1] ) # 2. Action generation loop for i, t in enumerate(action_timesteps): last_step = i == len(action_timesteps) - 1 action_cond = ( torch.zeros( [1, CONFIG["action_dim"], 1, CONFIG["action_per_frame"], 1], device=DEVICE, dtype=DTYPE, ) if frame_st_id == 0 else None ) input_dict = _prepare_latent_input( None, actions, t, t, None, action_cond, frame_st_id=frame_st_id, patch_size=CONFIG["patch_size"], prompt_embeds=prompt_embeds, use_cfg=use_cfg, negative_prompt_embeds=negative_prompt_embeds, ) action_noise_pred = transformer( _repeat_input_for_cfg( input_dict["action_res_lst"], use_cfg, prompt_embeds, negative_prompt_embeds, ), update_cache=1 if last_step else 0, cache_name="pos", action_mode=True, ) if not last_step: action_noise_pred = rearrange( action_noise_pred, "b (f n) c -> b c f n 1", f=frame_chunk_size ) if action_guidance_scale > 1: action_noise_pred = action_noise_pred[1:] + action_guidance_scale * ( action_noise_pred[:1] - action_noise_pred[1:] ) else: action_noise_pred = action_noise_pred[:1] actions = action_scheduler.step( action_noise_pred, t, actions, return_dict=False ) actions[:, :, 0:1] = ( action_cond if frame_st_id == 0 else actions[:, :, 0:1] ) actions[:, ~action_mask] *= 0 return actions, latents def decode_video(pred_latent): """Decode latent tensor to video frames.""" vae_device = next(vae.parameters()).device latents = pred_latent.to(vae_device).to(vae.dtype) latents_mean = ( torch.tensor(vae.config.latents_mean) .view(1, vae.config.z_dim, 1, 1, 1) .to(latents.device, latents.dtype) ) latents_std = 1.0 / torch.tensor(vae.config.latents_std).view( 1, vae.config.z_dim, 1, 1, 1 ).to(latents.device, latents.dtype) latents = latents / latents_std + latents_mean with torch.no_grad(): video = vae.decode(latents, return_dict=False)[0] video = video_processor.postprocess_video(video, output_type="np")[0] return video # ------------------------------------------------------------------ # # Gradio inference function # ------------------------------------------------------------------ # @spaces.GPU(duration=60, size="xlarge") def generate( top_img: "np.ndarray", wrist_img: "np.ndarray", prompt: str, num_chunks: int = 5, seed: int = 0, progress=gr.Progress(track_tqdm=True), ): """Generate a robot manipulation video from initial observations and a text prompt. Args: top_img: Top-down camera observation image. wrist_img: Wrist camera observation image. prompt: Natural language instruction for the robot task. num_chunks: Number of video chunks to generate autoregressively (each chunk = 4 frames). seed: Random seed for reproducibility. """ torch.manual_seed(seed) torch.cuda.manual_seed(seed) use_cfg = CONFIG["guidance_scale"] > 1 or CONFIG["action_guidance_scale"] > 1 # Prepare observations obs = [ { CONFIG["obs_cam_keys"][0]: top_img, CONFIG["obs_cam_keys"][1]: wrist_img, } ] # Reset KV cache transformer.clear_cache("pos") streaming_vae.clear_cache() # Encode initial observation init_latent = encode_obs(obs) # Encode prompt prompt_embeds, negative_prompt_embeds = encode_prompt(prompt) # Latent dimensions latent_height = CONFIG["height"] // 16 latent_width = (CONFIG["width"] // 16) * len(CONFIG["obs_cam_keys"]) patch_size = CONFIG["patch_size"] latent_token_per_chunk = ( CONFIG["frame_chunk_size"] * latent_height * latent_width ) // (patch_size[0] * patch_size[1] * patch_size[2]) action_token_per_chunk = CONFIG["frame_chunk_size"] * CONFIG["action_per_frame"] # Create KV cache transformer.create_empty_cache( "pos", CONFIG["attn_window"], latent_token_per_chunk, action_token_per_chunk, device=DEVICE, dtype=DTYPE, batch_size=2 if use_cfg else 1, ) # Autoregressive chunk generation pred_latent_lst = [] for chunk_id in range(num_chunks): frame_st_id = chunk_id * CONFIG["frame_chunk_size"] actions, latents = infer_chunk( init_latent, frame_st_id, prompt_embeds, negative_prompt_embeds, use_cfg, CONFIG["guidance_scale"], CONFIG["action_guidance_scale"], num_chunks, ) pred_latent_lst.append(latents) pred_latent = torch.cat(pred_latent_lst, dim=2) # Free VRAM before VAE decode: move transformer and text encoder to CPU transformer.clear_cache("pos") streaming_vae.clear_cache() transformer.to("cpu") text_encoder.to("cpu") torch.cuda.empty_cache() # Decode video pred_latent_cpu = pred_latent.cpu() del pred_latent torch.cuda.empty_cache() video = decode_video(pred_latent_cpu) # Save to temp file from diffusers.utils import export_to_video tmp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) tmp_file.close() export_to_video(video, tmp_file.name, fps=10) return tmp_file.name # ------------------------------------------------------------------ # # Gradio UI # ------------------------------------------------------------------ # CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks() as demo: gr.Markdown( """ # Next Forcing: Causal World Modeling with Multi-Chunk Prediction Generate robot manipulation video from initial observations and a text instruction. Upload top-down and wrist camera images, describe the task, and the model autoregressively predicts future video frames. [Paper](https://arxiv.org/abs/2606.11187) | [Code](https://github.com/gangweix/next-forcing) | [Model](https://huggingface.co/gangweix/next-forcing-base) """ ) with gr.Row(): with gr.Column(): top_img = gr.Image( label="Top Camera", type="numpy", height=256, ) wrist_img = gr.Image( label="Wrist Camera", type="numpy", height=256, ) prompt = gr.Textbox( label="Task Instruction", placeholder="e.g. Pick the green cube and place it inside the blue box", lines=2, ) with gr.Accordion("Advanced Settings", open=False): num_chunks = gr.Slider( label="Number of chunks (4 frames each)", minimum=1, maximum=10, value=5, step=1, ) seed = gr.Number(label="Seed", value=0, precision=0) run_btn = gr.Button("Generate Video", variant="primary") with gr.Column(): video_out = gr.Video(label="Generated Video") gr.Examples( examples=[ [ "examples/observation.images.top.png", "examples/observation.images.wrist.png", "Pick the green cube and place it inside the blue box", 5, 0, ], [ "examples/observation.images.top.png", "examples/observation.images.wrist.png", "Move the red block to the left side of the table", 5, 42, ], ], inputs=[top_img, wrist_img, prompt, num_chunks, seed], outputs=video_out, fn=generate, cache_examples=True, cache_mode="lazy", ) run_btn.click( fn=generate, inputs=[top_img, wrist_img, prompt, num_chunks, seed], outputs=video_out, api_name="generate", ) demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)