File size: 2,270 Bytes
cfaf85b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Sample the exported custom TraceDiT inference bundle without loading pickle."""
import argparse
import json
from pathlib import Path
import numpy as np
from PIL import Image
from safetensors.torch import load_file
import torch
from models import build_backbone
from objectives import build_objective
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", required=True)
parser.add_argument("--output", default="samples.png")
parser.add_argument("--n-per-class", type=int, default=4)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--weights", default="model.safetensors")
parser.add_argument("--nfe", type=int, default=1)
args = parser.parse_args()
if args.n_per_class <= 0:
parser.error("n-per-class must be positive")
folder = Path(args.model_dir)
cfg = json.loads((folder / "config.json").read_text())
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(args.seed)
model = build_backbone(cfg).to(device)
model.load_state_dict(load_file(str(folder / args.weights), device=device), strict=True)
objective = build_objective(cfg)
if args.nfe != 1 and not hasattr(objective, "max_nfe"):
parser.error("This model supports only one NFE")
if objective.num_classes is None:
labels, columns = None, 10
else:
columns = objective.num_classes
labels = torch.arange(columns, device=device).repeat(args.n_per_class)
images = objective.sample(model, n_samples=columns * args.n_per_class,
labels=labels, device=device,
**({"nfe": args.nfe} if hasattr(objective, "max_nfe") else {}))
pixels = (images.permute(0, 2, 3, 1).cpu().numpy() * 255).round().astype(np.uint8)
rows = args.n_per_class
if pixels.shape[-1] == 1:
pixels = pixels[..., 0]
grid = np.concatenate([np.concatenate(pixels[row*columns:(row+1)*columns], axis=1)
for row in range(rows)], axis=0)
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
Image.fromarray(grid).save(output)
print(f"{args.nfe}-NFE samples saved: {output}")
if __name__ == "__main__":
main()
|