examples/client_render.py

#!/usr/bin/env python3
"""上传工程、检查资源、逐批渲染并下载结果；仅依赖 gradio_client。

新任务：python examples/client_render.py --space USER/SPACE --project scene.zip --config examples/video.json
续渲：python examples/client_render.py --space USER/SPACE --resume-job --output outputs
"""

from __future__ import annotations

import argparse
import inspect
import json
import os
from pathlib import Path
import shutil
import sys
import time
from typing import Any

from gradio_client import Client, handle_file


def emit(event: str, **values: Any) -> None:
    """使用 JSON 行输出，方便 agent 读取进度和最终结果。"""
    print(json.dumps({"event": event, **values}, ensure_ascii=False), flush=True)


def save_json(path: Path, value: dict[str, Any]) -> None:
    """原子保存任务 ID，避免进程被中断后重新创建重复任务。"""
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + ".tmp")
    temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    try:
        temporary.chmod(0o600)
    except OSError:
        pass
    temporary.replace(path)


def make_client(space: str, cache: Path) -> Client:
    # Gradio Client 新版参数为 token；兼容使用 hf_token 的旧版本。
    token = os.environ.get("HF_TOKEN") or None
    parameters = inspect.signature(Client).parameters
    # 云端初次连接、跨区网络和大文件下载可能超过 HTTPX 默认等待时间。
    kwargs: dict[str, Any] = {"verbose": False, "download_files": str(cache),
                              "httpx_kwargs": {"timeout": 120}}
    kwargs["token" if "token" in parameters else "hf_token"] = token
    return Client(space, **kwargs)


def local_file(value: Any) -> Path:
    """Client 已负责鉴权下载；只复制本机文件，不向任意 URL 转发令牌。"""
    if isinstance(value, str):
        path = Path(value)
    elif isinstance(value, dict) and value.get("path"):
        path = Path(value["path"])
    else:
        raise RuntimeError("API 返回了无法识别的文件结构，请更新 gradio_client。")
    if not path.is_file():
        raise RuntimeError(f"Client 未能下载输出文件：{path.name}")
    return path


def export_results(client: Client, job_id: str, output: Path) -> dict[str, Any]:
    summary, files, _image, _video = client.predict(job_id, api_name="/finalize_job")
    downloaded: list[str] = []
    for value in files or []:
        source = local_file(value)
        destination = output / source.name
        if source.resolve() != destination.resolve():
            shutil.copy2(source, destination)
        downloaded.append(str(destination.resolve()))
    save_json(output / "client-summary.json", summary)
    emit("export", job_id=job_id, export_complete=summary.get("export_complete", False),
         files=downloaded, error=summary.get("export_error"))
    return summary


def ensure_healthy(status: dict[str, Any]) -> None:
    # 部分帧成功后出错时 status 仍可能是 partial，必须同时检查 error。
    if status.get("error"):
        raise RuntimeError(str(status["error"]))
    if status.get("status") in {"failed", "cancelled", "cancelling"}:
        raise RuntimeError(f"任务已停止：{status.get('status')}。确认原因后使用 --resume-job。")


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    result.add_argument("--space", required=True, help="HF Space ID（USER/SPACE），或 Gradio 根 URL")
    result.add_argument("--project", nargs="+", type=Path, help="一个项目 ZIP，或 .blend 与资源文件")
    result.add_argument("--config", type=Path, help="UTF-8 JSON 渲染配置；新任务必需")
    result.add_argument("--output", type=Path, default=Path("outputs"), help="下载目录，默认 outputs")
    result.add_argument("--job-file", type=Path, help="任务凭据文件，默认 OUTPUT/render-job.json")
    result.add_argument("--resume-job", nargs="?", const="@job-file", metavar="JOB_ID",
                        help="恢复指定任务；不传 ID 时从 --job-file 读取")
    result.add_argument("--inspect", "--inspect-only", dest="inspect_only", action="store_true",
                        help="准备并检查任务后退出，不渲染；正常渲染也会先检查")
    result.add_argument("--max-batches", type=int, default=0, help="本次最多渲染几批；0 表示直到完成")
    result.add_argument("--self-test", action="store_true", help="渲染前运行真实 GPU 小场景自检")
    return result


def run(args: argparse.Namespace) -> int:
    if args.max_batches < 0:
        raise ValueError("--max-batches 不能小于零。")
    args.output = args.output.expanduser().resolve()
    args.output.mkdir(parents=True, exist_ok=True)
    job_file = (args.job_file or args.output / "render-job.json").expanduser().resolve()
    if args.resume_job is None and job_file.exists():
        raise ValueError(f"任务文件已存在：{job_file}。请使用 --resume-job 续渲，或改用新的 --output。")
    if args.resume_job is not None and (args.project or args.config):
        raise ValueError("续渲沿用服务端原配置；请勿同时传入 --project 或 --config。")
    if args.resume_job is None:
        if not args.project or args.config is None:
            raise ValueError("新任务必须提供 --project 和 --config。")
        for path in [*args.project, args.config]:
            if not path.is_file():
                raise ValueError(f"找不到文件：{path}")
        config_text = args.config.read_text(encoding="utf-8-sig")
        if not isinstance(json.loads(config_text), dict):
            raise ValueError("配置根节点必须是 JSON 对象。")
    else:
        config_text = "{}"

    client = make_client(args.space, args.output / ".client-downloads")
    job_id: str | None = None

    def remember(status: dict[str, Any]) -> None:
        save_json(job_file, {"space": args.space, "job_id": job_id, "updated_at": time.time(),
                             "status": status.get("status"), "completed_count": status.get("completed_count"),
                             "total_frames": status.get("total_frames")})

    try:
        environment = client.predict(api_name="/environment_info")
        emit("environment", info=environment)
        if args.self_test:
            probe, _files, _image = client.predict(api_name="/gpu_self_test")
            emit("self_test", report=probe)
            if not probe.get("ok"):
                raise RuntimeError("真实渲染自检失败，请先查看 Space 的自检日志。")
            if probe.get("deployment_backend") != "cpu" and not probe.get("gpu_verified"):
                raise RuntimeError("自检未确认 Blender GPU 渲染成功。")

        if args.resume_job:
            if args.resume_job == "@job-file":
                saved = json.loads(job_file.read_text(encoding="utf-8"))
                if saved.get("space") != args.space:
                    raise ValueError("任务文件中的 Space 与 --space 不一致；请使用原 Space 标识。")
                job_id = str(saved["job_id"])
            else:
                job_id = args.resume_job
            status = client.predict(job_id, api_name="/resume_job")
        else:
            uploads = [handle_file(str(path.resolve())) for path in args.project]
            job_id, status = client.predict(uploads, config_text, None, api_name="/prepare_job")
        remember(status)
        emit("job", job_id=job_id, job_file=str(job_file), status=status.get("status"))

        # 在申请 GPU 之前检查相机、场景和资源，避免浪费租约。
        status = client.predict(job_id, api_name="/inspect_job")
        remember(status)
        save_json(args.output / "inspection.json", status)
        ensure_healthy(status)
        inspection = status.get("inspection") or {}
        missing = inspection.get("missing_assets") or []
        emit("inspection", missing_assets=missing, warnings=inspection.get("warnings", []))
        if missing and status.get("config", {}).get("strict_assets", True):
            raise RuntimeError("工程存在缺失或被阻止的外部资源；检查 inspection.json 并修复后创建新任务。")
        if args.inspect_only:
            emit("inspection_complete", job_id=job_id, next_action="使用 --resume-job 开始渲染")
            return 0

        batches = 0
        while status.get("remaining_frames"):
            before = status.get("completed_count", 0)
            status = client.predict(job_id, api_name="/render_chunk")
            batches += 1
            remember(status)
            emit("progress", job_id=job_id, status=status.get("status"),
                 completed=status.get("completed_count"), total=status.get("total_frames"),
                 error=status.get("error"))
            ensure_healthy(status)
            if status.get("completed_count", 0) <= before:
                raise RuntimeError("本批没有完成新帧，停止重试；请检查日志和 GPU 时限。")
            if args.max_batches and batches >= args.max_batches and status.get("remaining_frames"):
                export_results(client, job_id, args.output)
                emit("paused", job_id=job_id, message="已达到本次批数上限，可使用 --resume-job 继续。")
                return 0

        summary = export_results(client, job_id, args.output)
        if not summary.get("export_complete"):
            raise RuntimeError(summary.get("export_error") or "输出尚未完整生成。")
        emit("complete", job_id=job_id, output=str(args.output))
        return 0
    except KeyboardInterrupt:
        if job_id:
            try:
                remember(client.predict(job_id, api_name="/cancel_job"))
            except Exception:
                pass
        emit("interrupted", job_id=job_id, message="已请求停止；保留任务文件后可续渲。")
        return 130
    except Exception:
        if job_id:
            # 即使某一帧失败，也尽量取回已完成帧和诊断日志。
            try:
                export_results(client, job_id, args.output)
            except Exception:
                pass
            emit("resume_available", job_id=job_id, job_file=str(job_file))
        raise


def main() -> int:
    try:
        return run(parser().parse_args())
    except Exception as exc:
        message = str(exc)
        token = os.environ.get("HF_TOKEN")
        if token:
            message = message.replace(token, "[REDACTED]")
        print(f"错误：{message}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())