Align output layout to partner (task/episode/1 + prompt/ + video frames); --view all|head

ge_modal_app.py

@@ -121,6 +121,34 @@ def stable_uniform_0_1(*parts) -> float:
     return int(h[:16], 16) / float(0xFFFFFFFFFFFFFFFF)
 
 
+def _resolve_task_episode(e):
+    """与搭档 run_local.py 的 resolve_task_episode 一致：优先字段，否则从 gt_path 推。"""
+    if "task_name" in e and "episode_name" in e:
+        return str(e["task_name"]), str(e["episode_name"])
+    gt = e.get("gt_path")
+    if gt:
+        parts = str(gt).split("/")
+        if len(parts) >= 3:
+            return parts[-3], parts[-2]
+    return "unknown_task", str(e.get("episode_name", "episode_00000"))
+
+
+def _extract_frames(video_path, frames_dir):
+    """与搭档一致：逐帧抽成 video/frame_00000.jpg ..."""
+    import os, cv2
+    os.makedirs(frames_dir, exist_ok=True)
+    cap = cv2.VideoCapture(video_path)
+    idx = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        cv2.imwrite(os.path.join(frames_dir, f"frame_{idx:05d}.jpg"), frame)
+        idx += 1
+    cap.release()
+    return idx
+
+
 @app.function(image=image, volumes={DATA_DIR: data_vol}, timeout=60 * 30)
 def prepare_samples(
     source: str = "agibot_world",
@@ -159,9 +187,10 @@ def prepare_samples(
     made = []
     n_skip_hash = 0
     for i, e in enumerate(entries[:limit]):
-        base_id = f"{source}_{e.get('task_name','t')}_{e.get('episode_name','e')}"
+        task_name, episode_name = _resolve_task_episode(e)
+        base_id = f"{source}_{task_name}_{episode_name}"
         # 哈希采样的 sample_id：优先用 gt_path 原始字符串（与搭档 run_local.py 完全一致）
-        hash_id = e.get("gt_path") or f"{source}/{e.get('task_name','t')}/{e.get('episode_name','e')}"
+        hash_id = e.get("gt_path") or f"{source}/{task_name}/{episode_name}"
 
         # 先按 (gt_path, prompt_key) 逐 variant 做哈希采样，决定哪些要跑
         variants_to_run = []
@@ -198,7 +227,14 @@ def prepare_samples(
 
             with open(os.path.join(sdir, "prompt.txt"), "w") as f:
                 f.write(prompt + "\n")
-            made.append(sample_id)
+            shutil.copy(init_png, os.path.join(sdir, "init_frame.png"))   # 供输出 prompt bundle
+            made.append({
+                "sample_id": sample_id,        # data volume 上的输入目录名
+                "source": source,
+                "task_name": task_name,
+                "episode_name": episode_name,
+                "prompt_key": field,           # prompt_prefix / prompt_rewrite
+            })
 
     data_vol.commit()
     print(f"已准备 {len(made)} 个样本目录（哈希按 variant 跳过 {n_skip_hash} 个）到 {samples_root}")
@@ -215,13 +251,14 @@ def prepare_samples(
     timeout=60 * 60 * 12,   # 一个 shard 可能跑几小时
 )
 def run_inference(
-    sample_ids: list,
+    samples: list,                       # 每项是 prepare_samples 返回的 dict
     config_name: str = "video_model_infer_slow.yaml",
     n_chunk: int = 4,
     seed: int = 42,
+    view: str = "all",                   # "all"=三视角横拼(默认,信息全); "head"=只存head单视角
 ):
-    """处理一个 shard 的样本：模型只加载一次，循环推理；
-    断点续跑(已有 video.mp4 则跳过) + 逐样本容错(单条失败不影响整批)。"""
+    """处理一个 shard：模型只加载一次，循环推理；
+    断点续跑 + 逐样本容错。输出对齐搭档 run_local.py 的目录结构，视频只保留 head 单视角。"""
     import os, sys, yaml, shutil, torch
     from einops import rearrange
 
@@ -250,21 +287,28 @@ def run_inference(
         video_fps = 30
 
     device = "cuda"
-    print(f"加载模型中 (本 shard 共 {len(sample_ids)} 个样本) ...")
+    print(f"加载模型中 (本 shard 共 {len(samples)} 个样本) ...")
     tokenizer, text_encoder, vae, diffusion_model, scheduler, pipe = prepare_model(args, device=device)
     valid_cams = [c + "_color" for c in args.data["train"]["valid_cam"]]
+    head_view_idx = args.data["train"]["valid_cam"].index("head")   # 只保留 head 视角
     TEMPORAL_DOWN_RATIO = vae.temporal_compression_ratio
     sample_size = (args.data["train"]["sample_size"][1], args.data["train"]["sample_size"][0])
 
     samples_root = os.path.join(DATA_DIR, "samples")
     done, skipped, failed = [], [], []
 
-    for i, sid in enumerate(sample_ids):
+    for i, s in enumerate(samples):
+        sid = s["sample_id"]
         image_root = os.path.join(samples_root, sid)
-        out_path = os.path.join(OUT_DIR, sid)
-        video_file = os.path.join(out_path, "video.mp4")
-
-        if os.path.exists(video_file):          # 断点续跑：已生成则跳过
+        # 搭档风格输出路径： <source>/<prompt_key>/<task>/<episode>/1/
+        out_dir = os.path.join(OUT_DIR, s["source"], s["prompt_key"],
+                               s["task_name"], s["episode_name"], "1")
+        final_mp4 = os.path.join(out_dir, f"{s['task_name']}_{s['episode_name']}.mp4")
+        frames_dir = os.path.join(out_dir, "video")
+        prompt_dir = os.path.join(out_dir, "prompt")
+
+        # 断点续跑：视频和抽帧都在则跳过
+        if os.path.exists(final_mp4) and os.path.isdir(frames_dir) and os.listdir(frames_dir):
             skipped.append(sid)
             continue
         try:
@@ -283,15 +327,32 @@ def run_inference(
                 history_action_state=None, pixel_wise_timestep=args.pixel_wise_timestep,
                 n_chunk=n_chunk,
             )[0]
-            os.makedirs(out_path, exist_ok=True)
+            os.makedirs(out_dir, exist_ok=True)
+
+            # 保存视频：默认 all=三视角横拼(信息全)，可选 head=只存 head 单视角
             video = preds["video"].data.cpu()
-            save_video(rearrange(video, "(b v) c t h w -> b c t h (v w)", v=v)[0], video_file, fps=video_fps)
-            shutil.copy(os.path.join(image_root, "prompt.txt"), os.path.join(out_path, "prompt.txt"))
+            if view == "head":
+                video = rearrange(video, "(b v) c t h w -> b v c t h w", v=v)
+                out_video = video[0, head_view_idx]                     # c t h w (单视角)
+            else:
+                out_video = rearrange(video, "(b v) c t h w -> b c t h (v w)", v=v)[0]  # 三视角横拼
+            save_video(out_video, final_mp4, fps=video_fps)
+
+            # prompt bundle: prompt/prompt.txt + prompt/init_frame.png
+            os.makedirs(prompt_dir, exist_ok=True)
+            shutil.copy(os.path.join(image_root, "prompt.txt"), os.path.join(prompt_dir, "prompt.txt"))
+            init_src = os.path.join(image_root, "init_frame.png")
+            if os.path.exists(init_src):
+                shutil.copy(init_src, os.path.join(prompt_dir, "init_frame.png"))
+
+            # 抽帧: video/frame_XXXXX.jpg
+            _extract_frames(final_mp4, frames_dir)
+
             done.append(sid)
-            print(f"  ✅ [{i+1}/{len(sample_ids)}] {sid}")
+            print(f"  ✅ [{i+1}/{len(samples)}] {s['task_name']}/{s['episode_name']} [{s['prompt_key']}]")
         except Exception as e:
             failed.append((sid, str(e)[:200]))
-            print(f"  ❌ [{i+1}/{len(sample_ids)}] {sid}: {e}")
+            print(f"  ❌ [{i+1}/{len(samples)}] {sid}: {e}")
 
         if (i + 1) % 20 == 0:                    # 每 20 条提交一次，崩了也保得住进度
             output_vol.commit()
@@ -316,6 +377,7 @@ def main(
     keep_prob: float = 1.0,             # 哈希采样：0.3 = 约跑 30% 子集(确定性、可复现、可断点续跑)
     sample_seed: int = 20260609,        # 与搭档对齐用
     sample_model_key: str = "",         # 与搭档对齐的关键：两边相同才抽到同一子集；空则用 "ge"
+    view: str = "all",                  # "all"=三视角(默认); "head"=只存head单视角
 ):
     if not skip_weights:
         print("== 步骤 1: 下载权重 ==")
@@ -332,7 +394,7 @@ def main(
     print("== 步骤 3: GPU 并行推理 (断点续跑，已生成的会跳过) ==")
     chunks = [sample_ids[i::shards] for i in range(shards)]   # 轮转切片，负载均衡
     chunks = [c for c in chunks if c]
-    handles = [run_inference.spawn(c, config_name=config_name, n_chunk=n_chunk) for c in chunks]
+    handles = [run_inference.spawn(c, config_name=config_name, n_chunk=n_chunk, view=view) for c in chunks]
     results = [h.get() for h in handles]
 
     tot_done = sum(len(r["done"]) for r in results)