app.py

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("<Q", bytes(stream.read(8)))[0]
    header = json.loads(bytes(stream.read(header_len)))
    header.pop("__metadata__", None)
    items = sorted(header.items(), key=lambda kv: kv[1]["data_offsets"][0])
    data_pos = 0
    state_dict = {}
    for i, (name, info) in enumerate(items):
        start, end = info["data_offsets"]
        if start > 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)