55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
"""Upscale a saved direct-runtime H3 latent with the learned 3D model."""
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.latent_upscaler import load_h3_latent_upscaler, upscale_h3_latent
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--latent", type=Path, required=True)
|
|
parser.add_argument("--model", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--scale", type=float, default=2.0)
|
|
parser.add_argument("--precision", choices=("float16", "bfloat16", "float32"), default="float16")
|
|
args = parser.parse_args()
|
|
|
|
dtype = getattr(torch, args.precision)
|
|
state = torch.load(args.latent, map_location="cpu", weights_only=False)
|
|
if not isinstance(state, dict) or "latent" not in state:
|
|
state = {"latent": state}
|
|
source = state["latent"]
|
|
|
|
torch.cuda.synchronize()
|
|
started = time.perf_counter()
|
|
model = load_h3_latent_upscaler(args.model, dtype=dtype)
|
|
torch.cuda.synchronize()
|
|
loaded = time.perf_counter()
|
|
upscaled = upscale_h3_latent(model, source, scale=args.scale).cpu()
|
|
torch.cuda.synchronize()
|
|
finished = time.perf_counter()
|
|
|
|
result = dict(state)
|
|
result["latent"] = upscaled
|
|
result["width"] = upscaled.shape[-1] * 16
|
|
result["height"] = upscaled.shape[-2] * 16
|
|
result["upscale"] = {
|
|
"model": args.model.name,
|
|
"scale": args.scale,
|
|
"precision": args.precision,
|
|
"source_shape": tuple(source.shape),
|
|
"output_shape": tuple(upscaled.shape),
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save(result, args.output)
|
|
print(json.dumps({
|
|
"output": str(args.output),
|
|
"source_shape": tuple(source.shape),
|
|
"output_shape": tuple(upscaled.shape),
|
|
"model_load_seconds": loaded - started,
|
|
"upscale_seconds": finished - loaded,
|
|
}, indent=2))
|