tripocli/tripo/features/_common.py

"""Shared argument groups and the submit -> wait -> download flow."""

import os

from ..errors import TripoAPIError, TripoTimeout
from ..versions import default_for, resolve


def add_generation_options(parser):
    """Options shared by text / image / multiview / batch generation."""
    parser.add_argument("--smart-mesh", action="store_true", help="智能网格模式 (P2.0/P1.0，四边面低模)")
    parser.add_argument("-m", "--model", help="模型版本: P2.0 / P1.0 / H3.1 / v3.0 / v2.5（默认随模式选择）")
    parser.add_argument("--face-limit", type=int, help="面数上限（智能网格默认 25000，高精度默认 2000000；游戏场景建议 10000~25000）")
    parser.add_argument("--quad", action="store_true", help="四边面拓扑（智能网格默认开启）")
    parser.add_argument("--count", type=int, default=1, choices=[1, 2, 4], help="生成候选数量 1/2/4（默认 1）")
    parser.add_argument("--no-texture", action="store_true", help="不生成贴图")
    parser.add_argument("--no-pbr", action="store_true", help="不生成 PBR 材质")
    parser.add_argument("--no-delight", action="store_true", help="保留原图光影（不去光）")
    parser.add_argument("--texture-quality", default="detailed", choices=["standard", "detailed"], help="贴图质量")
    parser.add_argument("--geometry-quality", default="detailed", choices=["standard", "detailed"], help="几何精细度")
    parser.add_argument("--t-pose", action="store_true", help="人形角色输出 T-Pose")
    parser.add_argument("--visibility", default="shareable", choices=["shareable", "private"], help="模型可见性")


def add_wait_options(parser, output_help="生成完成后保存到该路径"):
    parser.add_argument("-w", "--wait", action="store_true", help="阻塞等待任务完成")
    parser.add_argument("-o", "--output", help=output_help)


def add_op_options(parser, output_help="完成后保存到该路径（可选）"):
    """Options shared by post-process operators, which act on an existing project."""
    parser.add_argument("-p", "--project", required=True, help="项目 ID (project_id)")
    parser.add_argument("-m", "--model", help="覆盖 model_version（默认取项目自身的模型版本）")
    add_wait_options(parser, output_help)


def generation_body(args):
    """Turn generation CLI args into the request body fields."""
    smart = args.smart_mesh
    model_version = resolve(args.model, smart_mesh=smart)
    face_limit = args.face_limit or (25000 if smart else 2000000)
    quad = True if smart else args.quad
    return {
        "model_version": model_version,
        "face_limit": face_limit,
        "quad": quad,
        "texture": not args.no_texture,
        "pbr": not args.no_pbr,
        "delight": not args.no_delight,
        "texture_quality": args.texture_quality,
        "geometry_quality": args.geometry_quality,
        "t_pose": args.t_pose,
        "visibility": args.visibility,
        "smart_poly": smart,
        "model_count": args.count,
    }


def wait_and_report(ctx, operator_id, timeout=900):
    """Block on a task, printing progress; returns the finished task."""

    def on_progress(status, progress, left_time, _task):
        ctx.out.progress(f"⏳ 进度 {progress}% (预估剩余 {left_time}s)")

    task = ctx.client.wait_task(operator_id, timeout=timeout, on_progress=on_progress)
    ctx.out.progress_done()
    return task


def fetch_model_url(ctx, project_id):
    detail = ctx.client.project_detail(project_id) or {}
    return detail.get("model_url"), detail


def finish_single(ctx, task_data, output=None, wait=False, timeout=900, label="模型"):
    """Standard flow for one generated asset: optionally wait and download."""
    project_id = task_data.get("project_id")
    operator_id = task_data.get("operator_id")

    if not wait:
        ctx.out.data(
            task_data,
            human=lambda: (
                ctx.out.line("✅ 任务已提交"),
                ctx.out.line(f"project_id  : {project_id}"),
                ctx.out.line(f"operator_id : {operator_id}"),
                ctx.out.line(f"查询进度    : tripocli task status {operator_id}"),
                ctx.out.line(f"等待并下载  : tripocli task wait {operator_id} -p {project_id} -o out.glb"),
            ),
        )
        return task_data

    task = wait_and_report(ctx, operator_id, timeout=timeout)
    model_url, detail = fetch_model_url(ctx, project_id)

    output_path = output or f"{project_id}.glb"
    if model_url:
        ctx.client.download(model_url, output_path)
        output_path = os.path.abspath(output_path)

    payload = {"task": task, "project": detail, "output_file": output_path if model_url else None}
    ctx.out.data(
        payload,
        human=lambda: (
            ctx.out.line(f"🎉 {label}完成"),
            ctx.out.line(f"模型地址: {model_url}"),
            ctx.out.line(f"已保存  : {output_path}") if model_url else ctx.out.line("⚠️ 未取到 model_url"),
        ),
    )
    return payload


def project_version(client, project_id, override=None):
    """Model version to send with post-process ops.

    The studio sends the project's own version for most operations, so we read it
    from the project rather than assuming a global default.
    """
    if override:
        return resolve(override)
    return client.project_operator(project_id).get("model_version")


def node_part_name(client, project_id):
    """The ``part_name_list`` the studio expects for a project's current mesh.

    Real names come from the mesh's nodes (only present inside the model file);
    when the model cannot be read we fall back to the single-node name
    ``tripo_node_<operator_id>`` (verified live).
    """
    names = client.part_names(project_id)
    if names:
        return names
    operator = client.project_operator(project_id)
    operator_id = operator.get("operator_id")
    if not operator_id:
        raise TripoAPIError(f"无法从项目 {project_id} 读取 operator_id")
    return [f"tripo_node_{operator_id}"]


def run_operator(ctx, result, output=None, wait=False, timeout=900, label="处理", multiple=False):
    """Shared flow for post-process operators (they return project_id/operator_id)."""
    if multiple:
        operator_ids = [item.get("operator_id") for item in result]
        project_id = result[0].get("project_id") if result else None
    else:
        project_id = result.get("project_id")
        operator_ids = [result.get("operator_id")]
    # Some endpoints answer without a task id; never wait on a missing one.
    operator_ids = [oid for oid in operator_ids if oid]

    if not wait:
        ctx.out.data(
            result,
            human=lambda: (
                ctx.out.line("✅ 任务已提交"),
                *[ctx.out.line(f"operator_id : {oid}") for oid in operator_ids],
            ),
        )
        return result

    for operator_id in operator_ids:
        wait_and_report(ctx, operator_id, timeout=timeout)

    model_url, detail = fetch_model_url(ctx, project_id) if project_id else (None, None)
    output_path = output
    if output_path and model_url:
        ctx.client.download(model_url, output_path)
        output_path = os.path.abspath(output_path)

    payload = {"result": result, "project": detail, "model_url": model_url, "output_file": output_path}
    ctx.out.data(
        payload,
        human=lambda: (
            ctx.out.line(f"🎉 {label}完成"),
            ctx.out.line(f"模型地址: {model_url}"),
        ),
    )
    return payload