def main():
args = create_argparser().parse_args()
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(
**args_to_dict(args,
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())
model.eval()
logger.log("creating data loader...")
data = load_data(
data_dir=args.data_dir,
batch_size=args.batch_size,
image_size=args.image_size,
class_cond=args.class_cond,
deterministic=True,
)
logger.log("evaluating...")
run_bpd_evaluation(model, diffusion, data, args.num_samples,
args.clip_denoised)
def main():
args = create_argparser().parse_args()
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(
**args_to_dict(args, model_and_diffusion_defaults().keys())
)
model.to(dist_util.dev())
schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion)
logger.log("creating data loader...")
data = load_data(
data_dir=args.data_dir,
batch_size=args.batch_size,
image_size=args.image_size,
class_cond=args.class_cond,
)
logger.log("training...")
TrainLoop(
model=model,
diffusion=diffusion,
data=data,
batch_size=args.batch_size,
microbatch=args.microbatch,
lr=args.lr,
ema_rate=args.ema_rate,
log_interval=args.log_interval,
save_interval=args.save_interval,
resume_checkpoint=args.resume_checkpoint,
use_fp16=args.use_fp16,
fp16_scale_growth=args.fp16_scale_growth,
schedule_sampler=schedule_sampler,
weight_decay=args.weight_decay,
lr_anneal_steps=args.lr_anneal_steps,
).run_loop()
def main():
args = create_argparser().parse_args()
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(
**args_to_dict(args, 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("sampling...")
all_images = []
all_labels = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = {}
if args.class_cond:
classes = th.randint(
low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev()
)
model_kwargs["y"] = classes
sample_fn = (
diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
)
sample = sample_fn(
model,
(args.batch_size, 3, args.image_size, args.image_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()
gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
dist.all_gather(gathered_samples, sample) # gather not supported with NCCL
all_images.extend([sample.cpu().numpy() for sample in gathered_samples])
if args.class_cond:
gathered_labels = [
th.zeros_like(classes) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_labels, classes)
all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[: args.num_samples]
if args.class_cond:
label_arr = np.concatenate(all_labels, axis=0)
label_arr = label_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}")
if args.class_cond:
np.savez(out_path, arr, label_arr)
else:
np.savez(out_path, arr)
dist.barrier()
logger.log("sampling complete")
# timesteps.
'image_size': 512,
'learn_sigma': True,
'noise_schedule': 'linear',
'num_channels': 256,
'num_head_channels': 64,
'num_res_blocks': 2,
'resblock_updown': True,
'use_fp16': True,
'use_scale_shift_norm': True,
})
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
model, diffusion = create_model_and_diffusion(**model_config)
model.load_state_dict(
torch.load(
'./guided_diffusion_hd/512x512_diffusion_uncond_finetune_008100.pt',
map_location='cpu'))
model.requires_grad_(False).eval().to(device)
for name, param in model.named_parameters():
if 'qkv' in name or 'norm' in name or 'proj' in name:
param.requires_grad_()
if model_config['use_fp16']:
model.convert_to_fp16()
clip_model = clip.load('ViT-B/16',
jit=False)[0].eval().requires_grad_(False).to(device)
clip_size = clip_model.visual.input_resolution
normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
def main(**kwargs):
args = create_argparser().parse_args()
kw = {k: v for k, v in kwargs.items() if k in args}
args.__dict__.update(**kw)
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
pprint(args_to_dict(args, model_and_diffusion_defaults().keys()))
model, diffusion = create_model_and_diffusion(
**args_to_dict(args,
model_and_diffusion_defaults().keys()))
print(f"loading state dict {args.model_path} -> model")
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 classifier...")
classifier = create_classifier(
**args_to_dict(args,
classifier_defaults().keys()))
print(f"loading state dict {args.classifier_path}")
classifier.load_state_dict(
dist_util.load_state_dict(args.classifier_path, map_location="cpu"))
classifier.to(dist_util.dev())
if args.classifier_use_fp16:
classifier.convert_to_fp16()
classifier.eval()
def cond_fn(x, t, y=None):
assert y is not None
with th.enable_grad():
# print(f" cond_fn(x,t,y) t {t}")
# t, batch_size time sampling 999 -> 0
x_in = x.detach().requires_grad_(True)
logits = classifier(x_in, t)
# print(f" .. x:{x_in.shape}, t:{t.shape}, logits:{logits.shape}, y: {y.shape}")
# x:(batch_size, channels, image_size, image_size), t:(batch_size), (batch_size, numclasses), y: (batch_size)
log_probs = F.log_softmax(logits, dim=-1)
selected = log_probs[range(len(logits)), y.view(-1)]
# print(f" .. selected, softmax(logits)[range(), y] {selected}") # (batch_size) floats
cond = th.autograd.grad(selected.sum(),
x_in)[0] * args.classifier_scale
# print(f" .. cond: {tuple(cond.shape)}, args.classifier_scale {args.classifier_scale}")
# cond: (batch_size, 3, image_size, image_size), args.classifier_scale 0.5
return cond
def logt(x):
if isinstance(x, th.Tensor):
out = f" {tuple(x.shape)}"
if x.ndim == 1:
out += f"{x}"
elif isinstance(x, (int, float)):
out = f" {x}"
return out
def model_fn(x, t, y=None):
assert y is not None
print(f"timestep {t.tolist()} conditional y {y.tolist()}")
#print(f"model_fn, x {logt(x)}, t {logt(t)} y {logt(y)}")
return model(x, t, y if args.class_cond else None)
logger.log("sampling...")
all_images = []
all_labels = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = {}
classes = th.randint(low=0,
high=NUM_CLASSES,
size=(args.batch_size, ),
device=dist_util.dev())
model_kwargs["y"] = classes
if args.use_ddim:
print("sample_fn = diffusion.ddim_sample_loop: args.use_ddim")
else:
print("sample_fn = diffusion.p_sample_loop: not args.use_ddim")
sample_fn = (diffusion.p_sample_loop
if not args.use_ddim else diffusion.ddim_sample_loop)
# print(f"sample_fn args.batch_size {args.batch_size}, args.image_size {args.image_size} args.clip_denoised {args.clip_denoised} model_kwargs {model_kwargs}")
# model_kwargs['y']: class conditioner e.g [ 53, 37, 609, 498, 679, 38, 242, 705, 253, 822, 721, 762, 64, 42, 337, 483]
sample = sample_fn(
model_fn,
(args.batch_size, 3, args.image_size, args.image_size),
clip_denoised=args.clip_denoised,
model_kwargs=model_kwargs,
cond_fn=cond_fn,
device=dist_util.dev(),
)
sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
sample = sample.permute(0, 2, 3, 1)
sample = sample.contiguous()
gathered_samples = [
th.zeros_like(sample) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_samples,
sample) # gather not supported with NCCL
all_images.extend(
[sample.cpu().numpy() for sample in gathered_samples])
gathered_labels = [
th.zeros_like(classes) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_labels, classes)
all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[:args.num_samples]
label_arr = np.concatenate(all_labels, axis=0)
label_arr = label_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, label_arr)
dist.barrier()
logger.log("sampling complete")
def main():
args = create_argparser().parse_args()
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(
**args_to_dict(args,
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 classifier...")
classifier = create_classifier(
**args_to_dict(args,
classifier_defaults().keys()))
classifier.load_state_dict(
dist_util.load_state_dict(args.classifier_path, map_location="cpu"))
classifier.to(dist_util.dev())
if args.classifier_use_fp16:
classifier.convert_to_fp16()
classifier.eval()
def cond_fn(x, t, y=None):
assert y is not None
with th.enable_grad():
x_in = x.detach().requires_grad_(True)
logits = classifier(x_in, t)
log_probs = F.log_softmax(logits, dim=-1)
selected = log_probs[range(len(logits)), y.view(-1)]
return th.autograd.grad(selected.sum(),
x_in)[0] * args.classifier_scale
def model_fn(x, t, y=None):
assert y is not None
return model(x, t, y if args.class_cond else None)
logger.log("sampling...")
all_images = []
all_labels = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = {}
classes = th.randint(low=0,
high=NUM_CLASSES,
size=(args.batch_size, ),
device=dist_util.dev())
model_kwargs["y"] = classes
sample_fn = (diffusion.p_sample_loop
if not args.use_ddim else diffusion.ddim_sample_loop)
sample = sample_fn(
model_fn,
(args.batch_size, 3, args.image_size, args.image_size),
clip_denoised=args.clip_denoised,
model_kwargs=model_kwargs,
cond_fn=cond_fn,
device=dist_util.dev(),
)
sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
sample = sample.permute(0, 2, 3, 1)
sample = sample.contiguous()
gathered_samples = [
th.zeros_like(sample) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_samples,
sample) # gather not supported with NCCL
all_images.extend(
[sample.cpu().numpy() for sample in gathered_samples])
gathered_labels = [
th.zeros_like(classes) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_labels, classes)
all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[:args.num_samples]
label_arr = np.concatenate(all_labels, axis=0)
label_arr = label_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, label_arr)
dist.barrier()
logger.log("sampling complete")