104 lines
4.2 KiB
Python
104 lines
4.2 KiB
Python
"""图片 4K 高清放大 — upscale reference image to 4K.
|
||
|
||
tripocli img-upscale -a <asset_id> --wait -o ./upscaled.png
|
||
tripocli img-upscale (自动取最新一张参考图放大)
|
||
|
||
实测流程（页面上 4K 图标抓到的请求）：
|
||
POST /v2/studio/image/upscale
|
||
{"asset_id": "...", "resource_key": "generate-image/..."}
|
||
→ {"task_id": "...", "asset_id": "..."}
|
||
|
||
轮询：
|
||
POST /v2/studio/image/get_image_asset {"asset_id": "..."}
|
||
→ status == "success", type == "upscale_image", output.data[0].url
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
|
||
from ..errors import TripoAPIError, TripoTimeout
|
||
|
||
NAME = "img-upscale"
|
||
HELP = "图片 4K 高清放大：提升参考图分辨率（对应网页端 4K 按钮）"
|
||
|
||
|
||
def add_arguments(parser):
|
||
parser.add_argument("-a", "--asset", help="图片资产 ID（默认自动取最近生成的一张）")
|
||
parser.add_argument("--key", help="指定图片 resource_key（默认自动解析）")
|
||
parser.add_argument("-w", "--wait", action="store_true", default=True, help="等待放大完成（默认等待）")
|
||
parser.add_argument("--no-wait", dest="wait", action="store_false", help="不等待，提交后直接退出")
|
||
parser.add_argument("-o", "--output", help="放大后的图片保存路径")
|
||
parser.add_argument("--timeout", type=int, default=180, help="等待超时秒数（默认 180）")
|
||
|
||
|
||
def resolve_asset_and_key(ctx, asset_id, resource_key):
|
||
if not asset_id:
|
||
# 取最新的一张图片资产
|
||
data = ctx.client.get_data(
|
||
"POST", "/v2/studio/image/image_assets",
|
||
body={"page_num": 1, "page_size": 1},
|
||
) or {}
|
||
assets = data.get("assets") or []
|
||
if not assets:
|
||
raise TripoAPIError("未找到任何图片资产，请用 -a 指定 asset_id")
|
||
asset_id = assets[0].get("asset_id")
|
||
|
||
if not resource_key:
|
||
info = ctx.client.get_data("POST", "/v2/studio/image/get_image_asset", body={"asset_id": asset_id}) or {}
|
||
out_data = (info.get("output") or {}).get("data") or []
|
||
if out_data and out_data[0].get("key"):
|
||
resource_key = out_data[0].get("key")
|
||
else:
|
||
in_imgs = (info.get("input") or {}).get("prompt_image") or []
|
||
if in_imgs and in_imgs[0].get("key"):
|
||
resource_key = in_imgs[0].get("key")
|
||
|
||
if not resource_key:
|
||
raise TripoAPIError(f"无法从图片资产 {asset_id} 解析 resource_key，请用 --key 显式传入")
|
||
|
||
return asset_id, resource_key
|
||
|
||
|
||
def run(ctx):
|
||
args = ctx.args
|
||
asset_id, key = resolve_asset_and_key(ctx, args.asset, args.key)
|
||
ctx.out.line(f"🔍 提交 4K 放大任务: asset={asset_id}")
|
||
|
||
resp = ctx.client.get_data("POST", "/v2/studio/image/upscale", body={
|
||
"asset_id": asset_id,
|
||
"resource_key": key,
|
||
}) or {}
|
||
|
||
new_asset_id = resp.get("asset_id") or resp.get("task_id")
|
||
if not new_asset_id:
|
||
raise TripoAPIError(f"upscale 未返回 asset_id: {resp}")
|
||
|
||
if not args.wait:
|
||
ctx.out.data(resp, human=lambda: ctx.out.line(f"✅ 任务已提交，asset_id: {new_asset_id}"))
|
||
return resp
|
||
|
||
# 轮询直到完成
|
||
ctx.out.line(f"⏳ 等待 4K 放大完成: {new_asset_id}")
|
||
start = time.time()
|
||
while time.time() - start < args.timeout:
|
||
info = ctx.client.get_data("POST", "/v2/studio/image/get_image_asset", body={"asset_id": new_asset_id}) or {}
|
||
status = info.get("status")
|
||
if status == "success":
|
||
ctx.out.line("🎉 4K 放大完成！")
|
||
out_imgs = (info.get("output") or {}).get("data") or []
|
||
url = out_imgs[0].get("url") if out_imgs else None
|
||
if url:
|
||
ctx.out.line(f"图片地址: {url}")
|
||
if args.output:
|
||
save_path = args.output
|
||
if os.path.isdir(save_path):
|
||
save_path = os.path.join(save_path, f"{new_asset_id}_4k.png")
|
||
ctx.client.download_file(url, save_path)
|
||
ctx.out.line(f"💾 已保存到: {save_path}")
|
||
return info
|
||
if status in ("failed", "expired"):
|
||
raise TripoAPIError(f"4K 放大失败: status={status}")
|
||
time.sleep(2)
|
||
|
||
raise TripoTimeout(f"4K 放大超时（超过 {args.timeout} 秒）")
|