import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") import spaces # noqa: E402 (must come before torch) import hashlib import json import shutil import struct import subprocess import sys import tempfile import time from types import SimpleNamespace import gradio as gr import imageio.v2 as imageio import numpy as np import requests import torch from huggingface_hub import hf_hub_url, snapshot_download from huggingface_hub.utils import build_hf_headers from PIL import Image # --------------------------------------------------------------------------- # 1. Source code (official WorldGuide repo, pinned commit) + checkpoints # --------------------------------------------------------------------------- REPO_URL = "https://github.com/mbzuai-oryx/WorldGuide.git" REPO_COMMIT = "3d40b193c2f69e835714d5b839b6cbf3e744018e" REPO_DIR = os.path.abspath("WorldGuide") CKPT_DIR = os.path.abspath("ckpts/HunyuanVideo-1.5") WG_REPO = "ankanmbz/WorldGuide-Ckpt" if not os.path.isdir(os.path.join(REPO_DIR, "hyvideo")): subprocess.run(["git", "clone", "--depth", "1", REPO_URL, REPO_DIR], check=True) subprocess.run(["git", "-C", REPO_DIR, "fetch", "--depth", "1", "origin", REPO_COMMIT], check=False) subprocess.run(["git", "-C", REPO_DIR, "checkout", REPO_COMMIT], check=False) sys.path.insert(0, REPO_DIR) t0 = time.time() # WorldGuide ContextPlanner == DiT text encoder (fine-tuned Qwen2.5-VL-7B), one shared copy. wg_local = snapshot_download(WG_REPO, allow_patterns=["text_encoder/llm/*", "transformer/config.json"], ignore_patterns=["*.pt"]) LLM_DIR = os.path.join(wg_local, "text_encoder", "llm") DIT_CONFIG_DIR = os.path.join(wg_local, "transformer") # HunyuanVideo-1.5 VAE + scheduler (the WorldGuide DiT checkpoint holds all DiT weights). snapshot_download("tencent/HunyuanVideo-1.5", allow_patterns=["vae/*", "scheduler/*"], local_dir=CKPT_DIR) # SigLIP vision encoder (byte-identical to FLUX.1-Redux-dev/image_encoder, ungated mirror). snapshot_download( "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v", allow_patterns=["image_encoder/*", "feature_extractor/*"], local_dir=os.path.join(CKPT_DIR, "vision_encoder", "siglip"), ) # ByT5 glyph encoder. snapshot_download( "google/byt5-small", local_dir=os.path.join(CKPT_DIR, "text_encoder", "byt5-small"), ignore_patterns=["*.h5", "*.msgpack"], ) glyph_root = os.path.join(CKPT_DIR, "text_encoder", "Glyph-SDXL-v2") glyph_ckpt = os.path.join(glyph_root, "checkpoints", "byt5_model.pt") if not os.path.exists(glyph_ckpt): tmp = snapshot_download("multimodalart/glyph-sdxl-v2-byt5-small", allow_patterns=["assets/*", "pytorch_model.bin"]) os.makedirs(os.path.join(glyph_root, "assets"), exist_ok=True) os.makedirs(os.path.join(glyph_root, "checkpoints"), exist_ok=True) for f in os.listdir(os.path.join(tmp, "assets")): shutil.copy(os.path.join(tmp, "assets", f), os.path.join(glyph_root, "assets", f)) shutil.copy(os.path.join(tmp, "pytorch_model.bin"), glyph_ckpt) os.makedirs(os.path.join(CKPT_DIR, "text_encoder"), exist_ok=True) llm_link = os.path.join(CKPT_DIR, "text_encoder", "llm") if not os.path.exists(llm_link): os.symlink(LLM_DIR, llm_link) print(f"[setup] downloads done in {time.time() - t0:.0f}s", flush=True) # --------------------------------------------------------------------------- # 2. hyvideo imports + ZeroGPU-friendly patches # --------------------------------------------------------------------------- import hyvideo # noqa: E402,F401 (sets single-process dist env vars) import hyvideo.commons as hy_commons # noqa: E402 import hyvideo.commons.parallel_states as hy_parallel # noqa: E402 from hyvideo.commons.infer_state import initialize_infer_state # noqa: E402 hy_commons.get_gpu_memory = lambda device=None: 96 * 1024**3 class _SingleProcessParallel: sp = 1 world_size = 1 sp_enabled = False dp_enabled = False sp_rank = 0 # Single GPU, no sequence parallelism: avoid creating a torch.distributed device mesh. setattr(hy_parallel, "__parallel_dims", _SingleProcessParallel()) import hyvideo.models.text_encoders as hy_te # noqa: E402 import hyvideo.pipelines.worldplay_video_pipeline as hy_pipe_mod # noqa: E402 from hyvideo.models.transformers.worldplay_1_5_transformer import HunyuanVideo_1_5_DiffusionTransformer # noqa: E402 from hyvideo.models.vision_encoder import VisionEncoder # noqa: E402 from hyvideo.models.autoencoders import hunyuanvideo_15_vae_w_cache # noqa: E402 from hyvideo.schedulers.scheduling_flow_match_discrete import FlowMatchDiscreteScheduler # noqa: E402 from hyvideo.commons import PIPELINE_CONFIGS # noqa: E402 from hyvideo.qwen_planner_utils import ( # noqa: E402 PLANNER_SYSTEM_PROMPT, TASK_COMPLETED_TOKEN, build_planner_prompt, clean_planner_prediction, split_completion, ) from hyvideo.qwen_video_reader import configure_planner_video_backend # noqa: E402 hy_pipe_mod.get_gpu_memory = hy_commons.get_gpu_memory HunyuanVideo_1_5_Pipeline = hy_pipe_mod.HunyuanVideo_1_5_Pipeline initialize_infer_state( SimpleNamespace( sage_blocks_range="0-53", use_sageattn=False, enable_torch_compile=False, use_fp8_gemm=False, quant_type="fp8-per-block", include_patterns="double_blocks", use_vae_parallel=False, ) ) configure_planner_video_backend("ffmpeg") DTYPE = torch.bfloat16 # --------------------------------------------------------------------------- # 3. WorldGuide video DiT: build on meta, stream the fp32 checkpoint from the # Hub straight into bf16 tensors (no 35 GB file on disk, no fp32 copy in RAM) # --------------------------------------------------------------------------- _ST_DTYPES = { "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16, "I64": torch.int64, "I32": torch.int32, "U8": torch.uint8, "BOOL": torch.bool, } class _ResumableStream: def __init__(self, url): self.url = url self.pos = 0 self._open() def _open(self): headers = build_hf_headers() if self.pos: headers["Range"] = f"bytes={self.pos}-" self.resp = requests.get(self.url, headers=headers, stream=True, timeout=120, allow_redirects=True) self.resp.raise_for_status() self.raw = self.resp.raw self.raw.decode_content = True def read(self, n): buf = bytearray(n) mv = memoryview(buf) got, retries = 0, 0 while got < n: try: k = self.raw.readinto(mv[got:]) except Exception as e: # network hiccup: resume with a Range request k = 0 print(f"[dit-stream] read error {e!r}; resuming at {self.pos}", flush=True) if not k: retries += 1 if retries > 20: raise RuntimeError("DiT checkpoint stream failed repeatedly") time.sleep(2) self.resp.close() self._open() continue got += k self.pos += k return buf def stream_safetensors(url, dtype): stream = _ResumableStream(url) header_len = struct.unpack(" data_pos: stream.read(start - data_pos) nbytes = end - start if nbytes: buf = stream.read(nbytes) t = torch.frombuffer(buf, dtype=_ST_DTYPES[info["dtype"]]).reshape(info["shape"]) else: t = torch.empty(info["shape"], dtype=_ST_DTYPES[info["dtype"]]) state_dict[name] = t.to(dtype) if t.is_floating_point() else t.clone() data_pos = end if i % 200 == 0: print(f"[dit-stream] {i}/{len(items)} tensors, {stream.pos / 1e9:.1f} GB", flush=True) stream.resp.close() return state_dict t0 = time.time() from accelerate import init_empty_weights # noqa: E402 with init_empty_weights(): transformer = HunyuanVideo_1_5_DiffusionTransformer.from_config( HunyuanVideo_1_5_DiffusionTransformer.load_config(DIT_CONFIG_DIR) ) transformer.add_action_parameters() dit_sd = stream_safetensors(hf_hub_url(WG_REPO, "transformer/diffusion_pytorch_model.safetensors"), DTYPE) incompatible = transformer.load_state_dict(dit_sd, strict=False, assign=True) del dit_sd if incompatible.unexpected_keys: raise RuntimeError(f"Unexpected DiT keys: {incompatible.unexpected_keys[:10]}") if incompatible.missing_keys: raise RuntimeError(f"Missing DiT keys: {incompatible.missing_keys[:10]}") meta_left = [n for n, p in list(transformer.named_parameters()) + list(transformer.named_buffers()) if p.is_meta] if meta_left: raise RuntimeError(f"DiT tensors left on meta: {meta_left[:10]}") transformer.memory_compression_available = True transformer = transformer.to(DTYPE).eval().requires_grad_(False) transformer.to("cuda") print(f"[setup] WorldGuide DiT loaded in {time.time() - t0:.0f}s", flush=True) # --------------------------------------------------------------------------- # 4. Planner (Qwen2.5-VL) shared with the DiT text encoder # --------------------------------------------------------------------------- from transformers import AutoModelForImageTextToText, AutoProcessor # noqa: E402 t0 = time.time() planner_processor = AutoProcessor.from_pretrained(LLM_DIR) if TASK_COMPLETED_TOKEN not in planner_processor.tokenizer.get_vocab(): planner_processor.tokenizer.add_special_tokens({"additional_special_tokens": [TASK_COMPLETED_TOKEN]}) planner_model = AutoModelForImageTextToText.from_pretrained(LLM_DIR, dtype=DTYPE).eval().requires_grad_(False) planner_model.to("cuda") def _shared_text_encoder(text_encoder_type, text_encoder_precision=None, text_encoder_path=None, logger=None, device=None): lm = planner_model.model.language_model lm.final_layer_norm = lm.norm return lm, text_encoder_path hy_te.load_text_encoder = _shared_text_encoder text_encoder = hy_te.TextEncoder( text_encoder_type="llm", tokenizer_type="llm", text_encoder_path=LLM_DIR, max_length=1000, text_encoder_precision="bf16", prompt_template=hy_te.PROMPT_TEMPLATE["li-dit-encode-image-json"], prompt_template_video=hy_te.PROMPT_TEMPLATE["li-dit-encode-video-json"], hidden_state_skip_layer=2, apply_final_norm=False, reproduce=False, logger=hy_pipe_mod.loguru.logger, device=None, ) print(f"[setup] planner / text encoder loaded in {time.time() - t0:.0f}s", flush=True) # --------------------------------------------------------------------------- # 5. VAE, scheduler, ByT5, SigLIP -> pipeline # --------------------------------------------------------------------------- vae_cfg = HunyuanVideo_1_5_Pipeline.get_vae_inference_config() vae = hunyuanvideo_15_vae_w_cache.AutoencoderKLConv3D.from_pretrained( os.path.join(CKPT_DIR, "vae"), torch_dtype=vae_cfg["dtype"] ).to("cuda") scheduler = FlowMatchDiscreteScheduler.from_pretrained(os.path.join(CKPT_DIR, "scheduler")) byt5_kwargs, prompt_format = HunyuanVideo_1_5_Pipeline._load_byt5(CKPT_DIR, True, 256, device=torch.device("cpu")) byt5_kwargs["byt5_model"].to("cuda") vision_encoder = VisionEncoder( vision_encoder_type="siglip", vision_encoder_precision="fp16", vision_encoder_path=os.path.join(CKPT_DIR, "vision_encoder", "siglip"), processor_type=None, processor_path=None, output_key=None, logger=hy_pipe_mod.logger, device=torch.device("cuda"), ) pipe = HunyuanVideo_1_5_Pipeline( vae=vae, text_encoder=text_encoder, transformer=transformer, scheduler=scheduler, text_encoder_2=None, progress_bar_config=None, byt5_model=byt5_kwargs["byt5_model"], byt5_tokenizer=byt5_kwargs["byt5_tokenizer"], byt5_max_length=byt5_kwargs["byt5_max_length"], prompt_format=prompt_format, execution_device="cuda", vision_encoder=vision_encoder, enable_offloading=False, **PIPELINE_CONFIGS["480p_i2v"], ) print("[setup] pipeline ready", flush=True) # --------------------------------------------------------------------------- # 6. Closed-loop step (mirrors hyvideo/generate_continuous_step_promptsV5_qwen_planner.py # with the authors' run_v5_qwen_closed_loop_memory.sh settings) # --------------------------------------------------------------------------- HEIGHT, WIDTH, FPS = 480, 832, 16 PLANNER_NFRAMES, PLANNER_MAX_PIXELS, PLANNER_MIN_PIXELS = 8, 65536, 16384 PLANNER_HISTORY_K, PLANNER_MAX_NEW_TOKENS = 3, 128 MEMORY_STEPS, MEMORY_BLEND = 2, 1.0 CATEGORIES = [ "origami", "paper engineering", "building blocks", "cooking", "crafts", "assembly", "knotting", "cloth folding", ] def build_task_goal(category: str, subject: str) -> str: subject = (subject or "").strip().rstrip(".") category = (category or "").strip() tutorial = f"{category} tutorial" if category else "tutorial" return f"A step-by-step {tutorial} showing how to make {subject.lower()}." def latent_length_from_video_length(video_length: int) -> int: return ((video_length - 1) // 4) + 1 def pad_or_trim_memory_latents(m: torch.Tensor, target_frames: int) -> torch.Tensor: if m.shape[2] > target_frames: return m[:, :, -target_frames:] if m.shape[2] < target_frames: pad = m.new_zeros(m.shape[0], m.shape[1], target_frames - m.shape[2], m.shape[3], m.shape[4]) return torch.cat([pad, m], dim=2) return m def normalize_step_prompt(raw_step: str, step_number: int): import re match = re.match(r"^\s*Step\s+(\d+)\s*:\s*(.+?)\s*$", raw_step, flags=re.IGNORECASE) text = match.group(2).strip() if match else raw_step.strip() return f"Step {step_number}: {text}", text def save_video(video: torch.Tensor, path: str, fps: int = FPS): if video.ndim == 5: video = video[0] video = (video * 255).clamp(0, 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy() imageio.mimwrite(path, video, fps=fps) def last_frame(path: str) -> Image.Image: reader = imageio.get_reader(path) frame = None try: for frame in reader: pass finally: reader.close() return Image.fromarray(frame[:, :, :3]).convert("RGB") def make_preview_clip(image: Image.Image, path: str): """1-second static 'environment preview' clip for the planner's first step.""" im = image.convert("RGB") w, h = im.size scale = max(WIDTH / w, HEIGHT / h) im = im.resize((max(WIDTH, round(w * scale)), max(HEIGHT, round(h * scale))), Image.BICUBIC) w, h = im.size left, top = (w - WIDTH) // 2, (h - HEIGHT) // 2 frame = np.array(im.crop((left, top, left + WIDTH, top + HEIGHT))) imageio.mimwrite(path, [frame] * FPS, fps=FPS) @torch.no_grad() def encode_preview_memory(preview_path: str, latent_length: int): """Authors' build_global_clip_memory_latents: VAE-encode the environment preview clip as step-1 mode2 memory (the executor was trained with this memory present).""" reader = imageio.get_reader(preview_path) try: frames = np.stack([f[:, :, :3] for f in reader]) finally: reader.close() video = torch.from_numpy(frames).permute(3, 0, 1, 2).float() / 127.5 - 1.0 # C,T,H,W video = video.unsqueeze(0).to("cuda") with torch.autocast(device_type="cuda", dtype=torch.float16): mem = pipe.vae.encode(video).latent_dist.mode() * pipe.vae.config.scaling_factor return pad_or_trim_memory_latents(mem, max(latent_length - 1, 0)) * MEMORY_BLEND @torch.no_grad() def planner_generate(messages, suppress_initial_special_tokens=False) -> str: from qwen_vl_utils import process_vision_info text = planner_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) image_inputs, video_inputs = process_vision_info(messages) inputs = planner_processor(text=[text], images=image_inputs, videos=video_inputs, return_tensors="pt", padding=True) inputs = {k: v.to("cuda") if torch.is_tensor(v) else v for k, v in inputs.items()} gen_kwargs = {} if suppress_initial_special_tokens: tok = planner_processor.tokenizer completion_id = tok.get_vocab().get(TASK_COMPLETED_TOKEN) gen_kwargs["begin_suppress_tokens"] = sorted(set(tok.all_special_ids) - {completion_id}) out = planner_model.generate( **inputs, max_new_tokens=PLANNER_MAX_NEW_TOKENS, do_sample=False, temperature=None, top_p=None, top_k=None, **gen_kwargs ) trimmed = out[0][inputs["input_ids"].shape[1]:] return planner_processor.decode(trimmed, skip_special_tokens=False).strip() def plan_next(task_goal: str, step_number: int, visual_path: str, history: list, is_env_preview: bool): prompt = build_planner_prompt( task_goal=task_goal, history=history, step_number=step_number, is_env_preview=is_env_preview, max_history_text=PLANNER_HISTORY_K, visual_feedback=True, ) messages = [ {"role": "system", "content": PLANNER_SYSTEM_PROMPT}, { "role": "user", "content": [ {"type": "video", "video": visual_path, "nframes": PLANNER_NFRAMES, "max_pixels": PLANNER_MAX_PIXELS, "min_pixels": PLANNER_MIN_PIXELS}, {"type": "text", "text": prompt}, ], }, ] raw = planner_generate(messages) pred = clean_planner_prediction(raw, planner_processor) if not pred: # authors' --planner_retry_empty raw = planner_generate(messages, suppress_initial_special_tokens=True) pred = clean_planner_prediction(raw, planner_processor) return split_completion(pred) def _image_key(img: Image.Image, goal: str) -> str: h = hashlib.md5(np.asarray(img.convert("RGB").resize((64, 64))).tobytes()).hexdigest() return f"{h}|{goal}" def stitch(clips: list, out_path: str): """Concatenate clips for display, dropping each clip's first frame (latent frame 0 of the mode2 window is a gray placeholder; the authors already drop it for clips after step 1).""" writer = imageio.get_writer(out_path, fps=FPS) try: for ci, clip in enumerate(clips): reader = imageio.get_reader(clip) try: for fi, frame in enumerate(reader): if fi == 0: continue writer.append_data(frame) finally: reader.close() finally: writer.close() def render_log(state) -> str: if not state or not state.get("steps"): return "_No steps yet. Upload a start image, describe the goal and press **Plan & execute next step**._" lines = [f"**Goal:** {state['goal']}", ""] for s in state["steps"]: src = "you" if s["source"] == "user" else "planner" lines.append(f"{s['n']}. {s['text']} _(by {src}, {s['secs']:.0f}s)_") if state.get("done"): lines += ["", "✅ **The planner reports the task as completed.**"] return "\n".join(lines) def _duration(start_image=None, category=None, subject=None, state=None, override_action="", video_length=45, num_inference_steps=20, *args, **kwargs): # measured on ZeroGPU xlarge: ~7.7 s per denoising step at 45 frames + ~10 s fixed (+15%) return int(15 + int(num_inference_steps) * 8.8 * (int(video_length) / 45.0)) @spaces.GPU(duration=_duration, size="xlarge") def next_step( start_image: Image.Image, category: str, subject: str, state: dict | None = None, override_action: str = "", video_length: int = 45, num_inference_steps: int = 20, guidance_scale: float = 7.5, flow_shift: float = 5.0, seed: int = 3208, progress=gr.Progress(track_tqdm=True), ): """Run one closed-loop WorldGuide step: the ContextPlanner looks at the current state and predicts the next atomic action, then the video DiT executes it as a short clip. Args: start_image: Initial observation (starting state of the task). category: Task category, e.g. "origami" or "cooking". subject: What to make, e.g. "a paper boat". state: Rollout state returned by the previous call (None starts a new rollout). override_action: Optional action text to execute instead of the planner's prediction. video_length: Frames per step clip (4k+1). num_inference_steps: Denoising steps per clip. guidance_scale: Classifier-free guidance scale. flow_shift: Flow-matching scheduler shift. seed: Base seed (step k uses seed + k - 1). Returns: Latest step clip, full stitched rollout video, step log (markdown) and the updated state. """ if start_image is None: raise gr.Error("Please provide a start image.") if not (subject or "").strip(): raise gr.Error("Please describe what to make (e.g. 'a paper boat').") t_start = time.time() goal = build_task_goal(category, subject) key = _image_key(start_image, goal) if not state or state.get("key") != key: state = {"key": key, "goal": goal, "dir": tempfile.mkdtemp(prefix="worldguide_"), "clips": [], "history": [], "latents": [], "refs": [], "steps": [], "done": False} if state.get("done") and not (override_action or "").strip(): gr.Info("The planner already marked this task as completed. Press Reset to start over, or type an action to continue.") return state.get("shown"), state.get("full"), render_log(state), state video_length = int(video_length) step_idx = len(state["clips"]) step_number = step_idx + 1 out_dir = state["dir"] # 1) Plan the next atomic action from the current (generated) observation. if step_idx == 0: preview = os.path.join(out_dir, "preview.mp4") make_preview_clip(start_image, preview) visual, is_preview = preview, True else: visual, is_preview = state["clips"][-1], False if (override_action or "").strip(): action, completed, source = override_action.strip(), False, "user" else: action, completed = plan_next(goal, step_number, visual, state["history"], is_preview) source = "planner" print(f"[step {step_number}] action={action!r} completed={completed} plan_time={time.time() - t_start:.1f}s", flush=True) if not action: state["done"] = True gr.Info("The planner reports the task is complete.") return state.get("shown"), state.get("full"), render_log(state), state # 2) Execute it with the video DiT (reference = last frame of previous clip; mode2 memory). reference_image = start_image.convert("RGB") if step_idx == 0 else last_frame(state["clips"][-1]) step_prompt, step_text = normalize_step_prompt(action, step_number) latent_length = latent_length_from_video_length(video_length) memory = None if step_idx == 0: memory = encode_preview_memory(os.path.join(out_dir, "preview.mp4"), latent_length) elif state["latents"]: memory = pad_or_trim_memory_latents(torch.cat(state["latents"], dim=2), max(latent_length - 1, 0)) memory = memory.to("cuda") * MEMORY_BLEND viewmats = torch.eye(4).unsqueeze(0).repeat(latent_length, 1, 1).unsqueeze(0) Ks = torch.tensor([[1.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.0, 0.0, 1.0]]).unsqueeze(0).repeat(latent_length, 1, 1).unsqueeze(0) act = torch.zeros((1, latent_length), dtype=torch.long) out = pipe( prompt=step_prompt, aspect_ratio="16:9", video_length=video_length, prompt_rewrite=False, num_inference_steps=int(num_inference_steps), negative_prompt="", seed=int(seed) + step_idx, flow_shift=float(flow_shift), embedded_guidance_scale=None, reference_image=reference_image, user_height=HEIGHT, user_width=WIDTH, chunk_latent_frames=latent_length, guidance_scale=float(guidance_scale), viewmats=viewmats, Ks=Ks, action=act, enable_sr=False, output_type="pt", return_dict=True, few_step=False, model_type="bi", memory_frames=20, temporal_context_size=12, memory_frame_policy="recent", origami_memory_mode="mode2", origami_memory_steps=MEMORY_STEPS, origami_memory_policy="latest_k", origami_memory_blend=MEMORY_BLEND, memory_history_latents=memory, memory_reference_images=list(state["refs"]), ) print(f"[step {step_number}] generation done at {time.time() - t_start:.1f}s", flush=True) clip_path = os.path.join(out_dir, f"step_{step_number:02d}.mp4") save_video(out.videos, clip_path) state["latents"] = (state["latents"] + [out.latents.detach().to("cpu", copy=True)])[-MEMORY_STEPS:] state["refs"] = (state["refs"] + [reference_image.copy()])[-MEMORY_STEPS:] del out, memory torch.cuda.empty_cache() state["clips"].append(clip_path) state["history"].append(step_text) if completed: state["done"] = True full_path = os.path.join(out_dir, f"rollout_{step_number:02d}.mp4") stitch(state["clips"], full_path) shown_clip = os.path.join(out_dir, f"show_step_{step_number:02d}.mp4") stitch([clip_path], shown_clip) state["shown"] = shown_clip state["full"] = full_path state["steps"].append({"n": step_number, "text": step_text, "source": source, "secs": time.time() - t_start}) print(f"[step {step_number}] done in {time.time() - t_start:.1f}s", flush=True) return shown_clip, full_path, render_log(state), state def reset(): return None, None, render_log(None), None, "" # --------------------------------------------------------------------------- # 7. UI # --------------------------------------------------------------------------- CSS = """ #col-container { max-width: 1200px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(title="WorldGuide") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# WorldGuide: goal-directed video world model\n" "Give it a starting scene and a goal. Each step, the **ContextPlanner** (fine-tuned Qwen2.5-VL) " "looks at the latest frames and predicts the next atomic action, then the **video DiT executor** " "(HunyuanVideo-1.5 based, with step-level latent memory) generates that action as a short clip. " "Press the button again to keep going until the planner reports the task is done.\n\n" "[Paper](https://huggingface.co/papers/2610.12459) · [Code](https://github.com/mbzuai-oryx/WorldGuide) · " "[Model](https://huggingface.co/ankanmbz/WorldGuide-Ckpt) · [Project page](https://mbzuai-oryx.github.io/WorldGuide/)" ) state = gr.State(None) with gr.Row(): with gr.Column(scale=1): start_image = gr.Image(label="Start state", type="pil", height=300) with gr.Row(): category = gr.Dropdown(CATEGORIES, value="cooking", label="Task category", allow_custom_value=True) subject = gr.Textbox(label="What to make", placeholder="e.g. a paper boat", value="") override = gr.Textbox( label="Next action (optional)", placeholder="Leave empty to let the planner decide, or type an action to steer the next step", ) with gr.Row(): run_btn = gr.Button("Plan & execute next step", variant="primary", scale=3) reset_btn = gr.Button("Reset", scale=1) with gr.Accordion("Advanced settings", open=False): video_length = gr.Slider(33, 61, value=45, step=4, label="Frames per step (16 fps)") num_steps = gr.Slider(10, 50, value=20, step=1, label="Denoising steps (authors used 30)") guidance = gr.Slider(1.0, 10.0, value=7.5, step=0.5, label="Guidance scale") shift = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Flow shift") seed = gr.Number(value=3208, precision=0, label="Seed (step k uses seed + k - 1)") with gr.Column(scale=1): latest = gr.Video(label="Latest step", autoplay=True) full = gr.Video(label="Full rollout so far") log = gr.Markdown(render_log(None)) inputs = [start_image, category, subject, state, override, video_length, num_steps, guidance, shift, seed] outputs = [latest, full, log, state] run_btn.click(next_step, inputs=inputs, outputs=outputs, api_name="next_step").then( lambda: "", outputs=override ) reset_btn.click(reset, outputs=[latest, full, log, state, override], api_name=False) gr.Examples( examples=[ ["examples/seafood_pancake_ingredients.jpg", "cooking", "a korean seafood pancake"], ["examples/strawberries_cutting_board.jpg", "cooking", "strawberry cheesecake cups"], ["examples/lego_bricks_desk.jpg", "building blocks", "a lego car"], ], inputs=[start_image, category, subject], outputs=outputs, fn=next_step, cache_examples=False, run_on_click=True, ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)