Spaces:
Runtime error
Runtime error
File size: 3,725 Bytes
ec0fdfd |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
"""
Generate a large batch of samples from a super resolution model, given a batch
of samples from a regular model from image_sample.py.
"""
import argparse
import os
import blobfile as bf
import numpy as np
import torch as th
import torch.distributed as dist
from guided_diffusion import dist_util, logger
from guided_diffusion.script_util import (
sr_model_and_diffusion_defaults,
sr_create_model_and_diffusion,
args_to_dict,
add_dict_to_argparser,
)
def main():
args = create_argparser().parse_args()
dist_util.setup_dist()
logger.configure()
logger.log("creating model...")
model, diffusion = sr_create_model_and_diffusion(
**args_to_dict(args, sr_model_and_diffusion_defaults().keys())
)
model.load_state_dict(
dist_util.load_state_dict(args.model_path, map_location="cpu")
)
model.to(dist_util.dev())
if args.use_fp16:
model.convert_to_fp16()
model.eval()
logger.log("loading data...")
data = load_data_for_worker(args.base_samples, args.batch_size, args.class_cond)
logger.log("creating samples...")
all_images = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = next(data)
model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()}
sample = diffusion.p_sample_loop(
model,
(args.batch_size, 3, args.large_size, args.large_size),
clip_denoised=args.clip_denoised,
model_kwargs=model_kwargs,
)
sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
sample = sample.permute(0, 2, 3, 1)
sample = sample.contiguous()
all_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
dist.all_gather(all_samples, sample) # gather not supported with NCCL
for sample in all_samples:
all_images.append(sample.cpu().numpy())
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[: args.num_samples]
if dist.get_rank() == 0:
shape_str = "x".join([str(x) for x in arr.shape])
out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz")
logger.log(f"saving to {out_path}")
np.savez(out_path, arr)
dist.barrier()
logger.log("sampling complete")
def load_data_for_worker(base_samples, batch_size, class_cond):
with bf.BlobFile(base_samples, "rb") as f:
obj = np.load(f)
image_arr = obj["arr_0"]
if class_cond:
label_arr = obj["arr_1"]
rank = dist.get_rank()
num_ranks = dist.get_world_size()
buffer = []
label_buffer = []
while True:
for i in range(rank, len(image_arr), num_ranks):
buffer.append(image_arr[i])
if class_cond:
label_buffer.append(label_arr[i])
if len(buffer) == batch_size:
batch = th.from_numpy(np.stack(buffer)).float()
batch = batch / 127.5 - 1.0
batch = batch.permute(0, 3, 1, 2)
res = dict(low_res=batch)
if class_cond:
res["y"] = th.from_numpy(np.stack(label_buffer))
yield res
buffer, label_buffer = [], []
def create_argparser():
defaults = dict(
clip_denoised=True,
num_samples=10000,
batch_size=16,
use_ddim=False,
base_samples="",
model_path="",
)
defaults.update(sr_model_and_diffusion_defaults())
parser = argparse.ArgumentParser()
add_dict_to_argparser(parser, defaults)
return parser
if __name__ == "__main__":
main()
|