Function sample_diffusion

src/sample.py:6–48

(args)

def sample_diffusion(args):
    from utils.triplane_util import load_triplane_data, decompose_featmaps, save_triplane_data
    from diffusion.script_util import create_model_and_diffusion_from_args
    
    # dist_util.setup_dist(args.gpu_id)

    src_data, sizes = load_triplane_data(encoding_feat_path(args.tag), device=dist_util.dev())

    model, diffusion = create_model_and_diffusion_from_args(args)
    model_path = diffusion_model_path(args.tag, args.ema_rate, args.diff_n_iters)
    model.load_state_dict(dist_util.load_state_dict(model_path, map_location="cpu"))
    model.to(dist_util.dev()).eval()

    sample_fn = (
        diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
    )

    result_dir = os.path.join(args.tag, args.output)
    os.makedirs(result_dir, exist_ok=True)

    C = src_data.shape[0]
    H, W, D = sizes
    batch_size = args.diff_batch_size
    H, W, D = int(H * args.resize[0]), int(W * args.resize[1]), int(D * args.resize[2])
    print("H, W, D:", H, W, D)

    result_paths = []
    for i in range(0, args.n_samples, batch_size):
        bs = min(batch_size, args.n_samples - i)
        out_shape = [bs, C, H + D, W + D]

        cond = {'H': H, 'W': W, 'D': D}
        samples = sample_fn(model, out_shape, progress=True, model_kwargs=cond)
        samples_xy, samples_xz, samples_yz = decompose_featmaps(samples, (H, W, D))
        samples_xy = samples_xy.detach().cpu().numpy()
        samples_xz = samples_xz.detach().cpu().numpy()
        samples_yz = samples_yz.detach().cpu().numpy()

        for j in range(bs):
            save_path = os.path.join(result_dir, f"{i+j:03d}", "feat.npz")
            save_triplane_data(save_path, samples_xy[j], samples_xz[j], samples_yz[j])
            result_paths.append(save_path)
    return result_paths


def decode(args, paths):