Step-1 mode2 memory from environment preview clip (authors' training-aligned path); timing logs

app.py

@@ -373,6 +373,22 @@ def make_preview_clip(image: Image.Image, path: str):
     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
@@ -456,7 +472,7 @@ def render_log(state) -> str:
 
 def _duration(start_image=None, category=None, subject=None, state=None, override_action="",
               video_length=45, num_inference_steps=30, *args, **kwargs):
-    return int(40 + int(num_inference_steps) * 3.2 * (int(video_length) / 45.0))
+    return 300
 
 
 @spaces.GPU(duration=_duration, size="xlarge")
@@ -522,7 +538,7 @@ def next_step(
     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}", flush=True)
+    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.")
@@ -534,7 +550,9 @@ def next_step(
     step_prompt, step_text = normalize_step_prompt(action, step_number)
     latent_length = latent_length_from_video_length(video_length)
     memory = None
-    if state["latents"]:
+    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)
@@ -573,6 +591,7 @@ def next_step(
         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:]