"""
Genie-Envisioner (GE-Base) 视频生成 —— Modal 推理脚本
=====================================================

在 Modal 云 GPU 上跑 AgibotTech/Genie-Envisioner 的 GE-Base 视频生成，
输入来自我们自己的数据集 HuggingFriends/mllm-as-embodied-world-judge。

整体流程（三个 Modal 函数 + 一个本地编排入口）：
  1. download_weights()  —— 把 GE-Base checkpoint 和 LTX-Video 组件下到 Volume（只需跑一次）
  2. prepare_samples()   —— 从 HF 数据集取 N 条样本，转成 GE 要求的 image_root 格式，写到 Volume
  3. run_inference()     —— GPU 函数：改配置 + 跑 video_gen_examples/infer.py，逐样本出 video.mp4

用法见同目录 README_RUN.md。先跑通 1-2 条，再放量。

数据桥接说明（重要）：
  GE-Base 训练时用的是 3 相机视角(head/hand_left/hand_right) × 每视角 4 张历史帧。
  我们的数据集每条只有「单张 init_frame + 单条 prompt」。
  这里默认用 BRIDGE_MODE="replicate"：把 init_frame 复制成 3 视角 × 4 帧，
  让模型留在它训练时的 3 视角/4 帧 regime（内容是复制的，侧视角是假的，质量需实测）。
  另一条路是改成单视角单帧推理(见 README)，属于 out-of-distribution，二选一可对比。
"""

import os
import modal

APP_NAME = "genie-envisioner-infer"
GE_REPO = "https://github.com/AgibotTech/Genie-Envisioner.git"

# ---- 可调参数 -------------------------------------------------------------
GE_CKPT_REPO = "agibot-world/Genie-Envisioner"
GE_CKPT_FILE = "ge_base_slow_v0.1.safetensors"   # 或 "GE_base_fast_v0.1.safetensors"
LTX_REPO = "Lightricks/LTX-Video"
DATASET_REPO = "HuggingFriends/mllm-as-embodied-world-judge"
# GPU：默认 L40S(实测性价比最高)；可临时切换，例如 GE_GPU=H100 modal run ...
GPU_TYPE = os.environ.get("GE_GPU", "L40S")
# --------------------------------------------------------------------------

app = modal.App(APP_NAME)

# 三个持久化 Volume：权重、转换后的样本、输出视频
weights_vol = modal.Volume.from_name("ge-weights", create_if_missing=True)
data_vol = modal.Volume.from_name("ge-data", create_if_missing=True)
output_vol = modal.Volume.from_name("ge-outputs", create_if_missing=True)

WEIGHTS_DIR = "/weights"   # /weights/ge/<ckpt>.safetensors , /weights/ltx/<LTX组件>
DATA_DIR = "/data"         # /data/samples/<sample_id>/{head_color,hand_left_color,hand_right_color}/*.png + prompt.txt
OUT_DIR = "/outputs"

# 容器镜像：python3.10 + GE 依赖 + 把 GE 仓库 clone 进镜像
image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install("git", "ffmpeg", "libgl1", "libglib2.0-0")
    .pip_install("huggingface_hub", "hf_transfer")
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
    .run_commands(
        f"git clone --depth 1 {GE_REPO} /root/Genie-Envisioner",
        "pip install -r /root/Genie-Envisioner/requirements.txt",
        # deepspeed 仅训练用；推理不需要，且它在 import 时会找 CUDA_HOME 编译算子而报错，
        # GE 代码本身没有 import deepspeed，卸掉可让 transformers 跳过它的导入。
        "pip uninstall -y deepspeed",
    )
)


# ==========================================================================
# 1) 下载权重到 Volume（只需跑一次，之后命中缓存）
# ==========================================================================
@app.function(image=image, volumes={WEIGHTS_DIR: weights_vol}, timeout=60 * 60)
def download_weights():
    import os
    from huggingface_hub import hf_hub_download, snapshot_download

    ge_dir = os.path.join(WEIGHTS_DIR, "ge")
    ltx_dir = os.path.join(WEIGHTS_DIR, "ltx")
    os.makedirs(ge_dir, exist_ok=True)
    os.makedirs(ltx_dir, exist_ok=True)

    # GE-Base checkpoint（diffusion transformer 权重）
    print(f"下载 GE checkpoint: {GE_CKPT_FILE} ...")
    hf_hub_download(repo_id=GE_CKPT_REPO, filename=GE_CKPT_FILE, local_dir=ge_dir)

    # LTX-Video：推理只需 vae / tokenizer / text_encoder / scheduler / model_index.json
    print("下载 LTX-Video 组件 (vae/tokenizer/text_encoder/scheduler/model_index.json) ...")
    snapshot_download(
        repo_id=LTX_REPO,
        local_dir=ltx_dir,
        allow_patterns=[
            "vae/*",
            "tokenizer/*",
            "text_encoder/*",
            "scheduler/*",
            "model_index.json",
        ],
    )
    weights_vol.commit()
    print("权重下载完成。")


# ==========================================================================
# 2) 准备样本：从 HF 数据集取 N 条，转成 GE image_root 格式
# ==========================================================================
# summary.json 里 prompt 的字段名映射；prefix/rewrite 各自分开跑
PROMPT_FIELD_MAP = {
    "prompt": "prompt",
    "prefix": "prompt_prefix",
    "rewrite": "prompt_rewrite",
}


def _extract_text(value):
    """字段可能是 str（如 agibot_world 的 prefix）也可能是 list（droid/egodex），统一取文本。"""
    if isinstance(value, list):
        return value[0]
    return value


def stable_uniform_0_1(*parts) -> float:
    """与搭档 run_local.py 完全一致的确定性哈希打分，返回 [0,1)。
    同样的 (seed, model_key, prompt_key, sample_id) -> 同样的分数 -> 可复现、可断点续跑。"""
    import hashlib
    text = "||".join(str(x) for x in parts)
    h = hashlib.sha256(text.encode("utf-8")).hexdigest()
    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",
    limit: int = 2,
    bridge_mode: str = "replicate",   # "replicate" = 3视角×4帧复制 init_frame
    prompt_variants=("prefix", "rewrite"),  # 每条样本对这些 prompt 版本各建一个目录、分开跑
    n_previous: int = 4,
    valid_cams=("head", "hand_left", "hand_right"),
    keep_prob: float = 0.4,           # 哈希采样：约保留 40%(确定性、可复现)；1.0=全跑
    sample_seed: int = 20260609,      # 与搭档对齐用；同 seed/model_key/prompt_key/gt_path -> 同决定
    sample_model_key: str = "genie",  # GE 的固定标识，决定抽哪个子集；从一而终别改
):
    """从 data/<source>/summary.json 取前 limit 条，下 init_frame，
    对每个 prompt 版本(prefix/rewrite)各铺一个 GE 格式目录。"""
    import os, json, shutil
    from huggingface_hub import hf_hub_download

    samples_root = os.path.join(DATA_DIR, "samples")
    # 每次重建，避免脏数据
    if os.path.exists(samples_root):
        shutil.rmtree(samples_root)
    os.makedirs(samples_root, exist_ok=True)

    # 下 summary.json
    summary_path = hf_hub_download(
        repo_id=DATASET_REPO,
        filename=f"data/{source}/summary.json",
        repo_type="dataset",
    )
    with open(summary_path) as f:
        entries = json.load(f)

    model_key = (sample_model_key or "").strip() or "genie"
    print(f"{source}: 共 {len(entries)} 条 × {prompt_variants}; "
          f"keep_prob={keep_prob} seed={sample_seed} model_key={model_key}")
    made = []
    n_skip_hash = 0
    for i, e in enumerate(entries[:limit]):
        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}/{task_name}/{episode_name}"

        # 先按 (gt_path, prompt_key) 逐 variant 做哈希采样，决定哪些要跑
        variants_to_run = []
        for variant in prompt_variants:
            field = PROMPT_FIELD_MAP[variant]   # "prompt_prefix" / "prompt_rewrite"，与搭档 prompt_key 一致
            if field not in e:
                continue
            if keep_prob < 1.0:
                score = stable_uniform_0_1(sample_seed, model_key, field, hash_id)
                if not (score < keep_prob):
                    n_skip_hash += 1
                    continue
            variants_to_run.append((variant, field))

        if not variants_to_run:
            continue

        # 有要跑的 variant 才下 init_frame（省带宽）
        init_png = hf_hub_download(repo_id=DATASET_REPO, filename=e["image"], repo_type="dataset")

        for variant, field in variants_to_run:
            prompt = _extract_text(e[field]).strip()
            sample_id = f"{base_id}__{variant}"      # 输出目录名
            sdir = os.path.join(samples_root, sample_id)

            if bridge_mode == "replicate":
                for cam in valid_cams:
                    camdir = os.path.join(sdir, f"{cam}_color")
                    os.makedirs(camdir, exist_ok=True)
                    for k in range(n_previous):
                        shutil.copy(init_png, os.path.join(camdir, f"{k}.png"))
            else:
                raise ValueError(f"暂只实现 replicate；single-view 见 README。bridge_mode={bridge_mode}")

            with open(os.path.join(sdir, "prompt.txt"), "w") as f:
                f.write(prompt + "\n")
            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}")
    return made


# ==========================================================================
# 3) GPU 推理：改配置 + 逐样本跑 infer.py
# ==========================================================================
@app.function(
    image=image,
    gpu=GPU_TYPE,
    volumes={WEIGHTS_DIR: weights_vol, DATA_DIR: data_vol, OUT_DIR: output_vol},
    timeout=60 * 60 * 12,   # 一个 shard 可能跑几小时
)
def run_inference(
    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单视角
    extract_frames: bool = False,        # 是否抽帧成 video/*.jpg；judge 读 mp4 故默认关(省时省空间)
):
    """处理一个 shard：模型只加载一次，循环推理；
    断点续跑 + 逐样本容错。输出对齐搭档 run_local.py 的目录结构。"""
    import os, sys, yaml, shutil, time, torch
    from einops import rearrange

    repo = "/root/Genie-Envisioner"

    # 改配置(权重路径)写到一份拷贝，不动原文件
    cfg_src = os.path.join(repo, "configs/ltx_model", config_name)
    cfg_dst = os.path.join(repo, "configs/ltx_model", "_infer_modal.yaml")
    with open(cfg_src) as f:
        cfg = yaml.safe_load(f)
    cfg["pretrained_model_name_or_path"] = os.path.join(WEIGHTS_DIR, "ltx")
    cfg["diffusion_model"]["model_path"] = os.path.join(WEIGHTS_DIR, "ge", GE_CKPT_FILE)
    with open(cfg_dst, "w") as f:
        yaml.safe_dump(cfg, f)

    # config 里的 vae_class_path 等是相对路径，GE 按 cwd 解析；切到 repo 复现原来 cwd=repo 的行为
    # (其余路径 image_root/out_dir/cfg_dst 都是绝对路径，不受影响)
    os.chdir(repo)

    # 复用 GE 自己的代码加载模型 —— 只加载一次！(关键提速)
    sys.path.insert(0, repo)
    sys.path.insert(0, os.path.join(repo, "video_gen_examples"))
    from infer import load_config, prepare_model, load_images
    from utils import save_video

    args = load_config(cfg_dst)
    if "action_chunk" in args.data["train"]:
        video_fps = 30 // (args.data["train"]["action_chunk"] // args.data["train"]["chunk"])
    else:
        video_fps = 30

    device = "cuda"
    print(f"⏱GPU: {torch.cuda.get_device_name(0)}")
    print(f"加载模型中 (本 shard 共 {len(samples)} 个样本) ...")
    _t_load = time.perf_counter()
    tokenizer, text_encoder, vae, diffusion_model, scheduler, pipe = prepare_model(args, device=device)
    print(f"⏱模型加载耗时: {time.perf_counter() - _t_load:.1f}s")
    gen_times = []
    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, s in enumerate(samples):
        sid = s["sample_id"]
        image_root = os.path.join(samples_root, sid)
        # 搭档风格输出路径： <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")

        # 断点续跑：mp4 在则跳过(开了抽帧则还要求 video/ 有内容)
        done_already = os.path.exists(final_mp4)
        if extract_frames:
            done_already = done_already and os.path.isdir(frames_dir) and bool(os.listdir(frames_dir))
        if done_already:
            skipped.append(sid)
            continue
        try:
            _tg = time.perf_counter()
            with open(os.path.join(image_root, "prompt.txt")) as f:
                prompt = f.readline().strip()
            obs = load_images(args, image_root, valid_cams, size=sample_size)
            v, c, t, h, w = obs.shape
            preds = pipe.infer(
                image=obs.to(device), prompt=[prompt], negative_prompt="",
                num_inference_steps=50, decode_timestep=0.03, decode_noise_scale=0.025,
                height=h, width=w, n_view=v, guidance_scale=1.0,
                return_action=args.return_action, n_prev=args.data["train"]["n_previous"],
                chunk=(args.data["train"]["chunk"] - 1) // TEMPORAL_DOWN_RATIO + 1,
                return_video=args.return_video, noise_seed=seed,
                action_chunk=args.data["train"]["action_chunk"],
                history_action_state=None, pixel_wise_timestep=args.pixel_wise_timestep,
                n_chunk=n_chunk,
            )[0]
            os.makedirs(out_dir, exist_ok=True)

            # 保存视频：默认 all=三视角横拼(信息全)，可选 head=只存 head 单视角
            video = preds["video"].data.cpu()
            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
            if extract_frames:
                _extract_frames(final_mp4, frames_dir)

            dt = time.perf_counter() - _tg
            gen_times.append(dt)
            done.append(sid)
            print(f"  ✅ [{i+1}/{len(samples)}] {s['task_name']}/{s['episode_name']} [{s['prompt_key']}]  ⏱{dt:.1f}s")
        except Exception as e:
            failed.append((sid, str(e)[:200]))
            print(f"  ❌ [{i+1}/{len(samples)}] {sid}: {e}")

        if (i + 1) % 20 == 0:                    # 每 20 条提交一次，崩了也保得住进度
            output_vol.commit()

    output_vol.commit()
    if gen_times:
        print(f"⏱每条生成: {[round(x, 1) for x in gen_times]}  平均 {sum(gen_times)/len(gen_times):.1f}s/条")
    print(f"shard 完成: done={len(done)} skipped={len(skipped)} failed={len(failed)}")
    return {"done": done, "skipped": skipped, "failed": failed,
            "gen_times": gen_times, "gpu": GPU_TYPE}


# ==========================================================================
# 把某个 source 的结果(直接在云端)推到 HF dataset 仓库
# 需要先创建 Modal secret： modal secret create huggingface-secret HF_TOKEN=hf_xxx
# ==========================================================================
@app.function(
    image=image,
    volumes={OUT_DIR: output_vol},
    secrets=[modal.Secret.from_name("huggingface-secret")],
    timeout=60 * 60,
)
def upload_source_to_hf(source: str, repo_id: str = "Xinyi0214/ge_result",
                        upload_frames: bool = False):
    """把 ge-outputs 里某个 source 推到 HF dataset。
    默认只传 mp4 + prompt/(每条~3文件)，排除 video/ 抽帧(每条~232张,会撑爆文件数)。
    需要抽帧时设 upload_frames=True。"""
    import os
    from huggingface_hub import HfApi
    api = HfApi(token=os.environ["HF_TOKEN"])
    api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
    src_dir = os.path.join(OUT_DIR, source)
    if not os.path.isdir(src_dir):
        print(f"⚠️ {src_dir} 不存在，跳过上传")
        return
    ignore = None if upload_frames else ["**/video/**"]   # 默认不传抽帧
    api.upload_folder(
        folder_path=src_dir, path_in_repo=source,
        repo_id=repo_id, repo_type="dataset",
        ignore_patterns=ignore,
        commit_message=f"Add results: {source}",
    )
    print(f"✅ 已上传 {source}{'(含抽帧)' if upload_frames else '(仅mp4+prompt)'} "
          f"-> https://huggingface.co/datasets/{repo_id}/tree/main/{source}")


# ==========================================================================
# 清理某个 source 已生成的抽帧(只删 video/ 目录，保留 mp4 + prompt，不动 GPU)
# ==========================================================================
@app.function(image=image, volumes={OUT_DIR: output_vol}, timeout=60 * 30)
def clean_frames(source: str):
    import os, shutil
    root = os.path.join(OUT_DIR, source)
    if not os.path.isdir(root):
        print(f"⚠️ {root} 不存在")
        return 0
    to_del = []
    for dirpath, dirnames, _ in os.walk(root):
        if "video" in dirnames:
            to_del.append(os.path.join(dirpath, "video"))
    for d in to_del:
        shutil.rmtree(d, ignore_errors=True)
    output_vol.commit()
    print(f"✅ {source}: 删除了 {len(to_del)} 个 video/ 抽帧目录（mp4 + prompt 保留）")
    return len(to_del)


# ==========================================================================
# 把结果上传到「共享数据集」的 generated_data/<model>_<variant>/ 下
# 结构与搭档 wan26_flash 一致：data/<source>/generated_data/<model>_<variant>/<task>/<episode>/1/
# ⚠️ 目标是 HuggingFriends 共享仓库，secret 里的 HF_TOKEN 需对该 org 有写权限
# ==========================================================================
@app.function(
    image=image,
    volumes={OUT_DIR: output_vol},
    secrets=[modal.Secret.from_name("huggingface-secret")],
    timeout=60 * 60,
)
def upload_to_shared(
    source: str,
    variant: str,                      # "prefix" 或 "rewrite"
    model_name: str = "genie",         # 文件夹名 <model_name>_<variant>，对齐搭档 wan26_flash 风格
    repo_id: str = "HuggingFriends/mllm-as-embodied-world-judge",
):
    import os
    from huggingface_hub import HfApi
    api = HfApi(token=os.environ["HF_TOKEN"])
    # 我们 volume 里的目录名是 prompt_prefix / prompt_rewrite
    src_dir = os.path.join(OUT_DIR, source, f"prompt_{variant}")
    if not os.path.isdir(src_dir):
        print(f"⚠️ {src_dir} 不存在，跳过（该 source/variant 还没生成？）")
        return
    path_in_repo = f"data/{source}/generated_data/{model_name}_{variant}"
    api.upload_folder(
        folder_path=src_dir,
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type="dataset",
        ignore_patterns=["**/video/**"],     # 只传 mp4 + prompt，不传抽帧
        commit_message=f"Add {model_name} {variant} results: {source}",
    )
    print(f"✅ {source}/{variant} -> {repo_id}/{path_in_repo}")


# ==========================================================================
# 本地编排入口：modal run ge_modal_app.py
# ==========================================================================
@app.local_entrypoint()
def main(
    source: str = "agibot_world",
    limit: int = 0,                     # 0 = 该来源全部 episode
    skip_weights: bool = False,
    config_name: str = "video_model_infer_slow.yaml",
    variants: str = "prefix,rewrite",   # 用哪些 prompt 版本；prefix 和 rewrite 各跑各的
    n_chunk: int = 4,                   # 视频长度：4 ≈ 8s(232帧@30fps)
    shards: int = 4,                    # 并行 GPU 容器数(GPU 型号见 GPU_TYPE)：墙钟时间 ÷shards
    keep_prob: float = 0.4,             # 哈希采样：约跑 40% 子集(确定性、可复现、可断点续跑)；1.0=全跑
    sample_seed: int = 20260609,        # 与搭档对齐用
    sample_model_key: str = "genie",    # GE 的固定标识，决定抽哪个子集；从一而终别改
    view: str = "all",                  # "all"=三视角(默认); "head"=只存head单视角
    extract_frames: bool = False,       # 是否抽帧成 video/*.jpg；judge 读 mp4 故默认关
    upload: bool = False,               # 跑完直接把该 source 结果(云端)上传 HF dataset
    hf_repo: str = "Xinyi0214/ge_result",
):
    if not skip_weights:
        print("== 步骤 1: 下载权重 ==")
        download_weights.remote()

    print(f"== 步骤 2: 准备样本 (source={source}, limit={'全部' if not limit else limit}) ==")
    prompt_variants = tuple(v.strip() for v in variants.split(",") if v.strip())
    sample_ids = prepare_samples.remote(
        source=source, limit=(None if not limit else limit), prompt_variants=prompt_variants,
        keep_prob=keep_prob, sample_seed=sample_seed, sample_model_key=sample_model_key,
    )
    print(f"   共 {len(sample_ids)} 个样本待推理，分 {shards} 个 {GPU_TYPE} 并行")

    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,
                                   view=view, extract_frames=extract_frames) for c in chunks]
    results = [h.get() for h in handles]

    tot_done = sum(len(r["done"]) for r in results)
    tot_skip = sum(len(r["skipped"]) for r in results)
    fails = [f for r in results for f in r["failed"]]
    print(f"\n=== {source} 完成: 新生成 {tot_done} / 跳过 {tot_skip} / 失败 {len(fails)} ===")
    for sid, err in fails:
        print(f"  FAIL {sid}: {err}")

    if upload:
        print(f"== 上传 {source} 结果到 HF dataset {hf_repo} ==")
        upload_source_to_hf.remote(source, repo_id=hf_repo)
    else:
        print("下载结果： modal volume get ge-outputs / ./results/")