sample_diffusion.py

importargparse, os, sys, glob, datetime, yaml
importtorch
importtime
importnumpyasnp
fromtqdmimporttrange

fromomegaconfimportOmegaConf
fromPILimportImage

fromldm.models.diffusion.ddimimportDDIMSampler
fromldm.utilimportinstantiate_from_config

rescale=lambdax: (x+1.) /2.

defcustom_to_pil(x):
x=x.detach().cpu()
x=torch.clamp(x, -1., 1.)
x= (x+1.) /2.
x=x.permute(1, 2, 0).numpy()
x= (255*x).astype(np.uint8)
x=Image.fromarray(x)
ifnotx.mode=="RGB":
x=x.convert("RGB")
returnx


defcustom_to_np(x):
# saves the batch in adm style as in https://github.com/openai/guided-diffusion/blob/main/scripts/image_sample.py
sample=x.detach().cpu()
sample= ((sample+1) *127.5).clamp(0, 255).to(torch.uint8)
sample=sample.permute(0, 2, 3, 1)
sample=sample.contiguous()
returnsample


deflogs2pil(logs, keys=["sample"]):
imgs=dict()
forkinlogs:
try:
iflen(logs[k].shape) ==4:
img=custom_to_pil(logs[k][0, ...])
eliflen(logs[k].shape) ==3:
img=custom_to_pil(logs[k])
else:
print(f"Unknown format for key {k}. ")
img=None
except:
img=None
imgs[k] =img
returnimgs


@torch.no_grad()
defconvsample(model, shape, return_intermediates=True,
verbose=True,
make_prog_row=False):


ifnotmake_prog_row:
returnmodel.p_sample_loop(None, shape,
return_intermediates=return_intermediates, verbose=verbose)
else:
returnmodel.progressive_denoising(
None, shape, verbose=True
 )


@torch.no_grad()
defconvsample_ddim(model, steps, shape, eta=1.0
 ):
ddim=DDIMSampler(model)
bs=shape[0]
shape=shape[1:]
samples, intermediates=ddim.sample(steps, batch_size=bs, shape=shape, eta=eta, verbose=False,)
returnsamples, intermediates


@torch.no_grad()
defmake_convolutional_sample(model, batch_size, vanilla=False, custom_steps=None, eta=1.0,):


log=dict()

shape= [batch_size,
model.model.diffusion_model.in_channels,
model.model.diffusion_model.image_size,
model.model.diffusion_model.image_size]

withmodel.ema_scope("Plotting"):
t0=time.time()
ifvanilla:
sample, progrow=convsample(model, shape,
make_prog_row=True)
else:
sample, intermediates=convsample_ddim(model, steps=custom_steps, shape=shape,
eta=eta)

t1=time.time()

x_sample=model.decode_first_stage(sample)

log["sample"] =x_sample
log["time"] =t1-t0
log['throughput'] =sample.shape[0] / (t1-t0)
print(f'Throughput for this batch: {log["throughput"]}')
returnlog

defrun(model, logdir, batch_size=50, vanilla=False, custom_steps=None, eta=None, n_samples=50000, nplog=None):
ifvanilla:
print(f'Using Vanilla DDPM sampling with {model.num_timesteps} sampling steps.')
else:
print(f'Using DDIM sampling with {custom_steps} sampling steps and eta={eta}')


tstart=time.time()
n_saved=len(glob.glob(os.path.join(logdir,'*.png')))-1
# path = logdir
ifmodel.cond_stage_modelisNone:
all_images= []

print(f"Running unconditional sampling for {n_samples} samples")
for_intrange(n_samples//batch_size, desc="Sampling Batches (unconditional)"):
logs=make_convolutional_sample(model, batch_size=batch_size,
vanilla=vanilla, custom_steps=custom_steps,
eta=eta)
n_saved=save_logs(logs, logdir, n_saved=n_saved, key="sample")
all_images.extend([custom_to_np(logs["sample"])])
ifn_saved>=n_samples:
print(f'Finish after generating {n_saved} samples')
break
all_img=np.concatenate(all_images, axis=0)
all_img=all_img[:n_samples]
shape_str="x".join([str(x) forxinall_img.shape])
nppath=os.path.join(nplog, f"{shape_str}-samples.npz")
np.savez(nppath, all_img)

else:
raiseNotImplementedError('Currently only sampling for unconditional models supported.')

print(f"sampling of {n_saved} images finished in {(time.time() -tstart) /60.:.2f} minutes.")


defsave_logs(logs, path, n_saved=0, key="sample", np_path=None):
forkinlogs:
ifk==key:
batch=logs[key]
ifnp_pathisNone:
forxinbatch:
img=custom_to_pil(x)
imgpath=os.path.join(path, f"{key}_{n_saved:06}.png")
img.save(imgpath)
n_saved+=1
else:
npbatch=custom_to_np(batch)
shape_str="x".join([str(x) forxinnpbatch.shape])
nppath=os.path.join(np_path, f"{n_saved}-{shape_str}-samples.npz")
np.savez(nppath, npbatch)
n_saved+=npbatch.shape[0]
returnn_saved


defget_parser():
parser=argparse.ArgumentParser()
parser.add_argument(
"-r",
"--resume",
type=str,
nargs="?",
help="load from logdir or checkpoint in logdir",
 )
parser.add_argument(
"-n",
"--n_samples",
type=int,
nargs="?",
help="number of samples to draw",
default=50000
 )
parser.add_argument(
"-e",
"--eta",
type=float,
nargs="?",
help="eta for ddim sampling (0.0 yields deterministic sampling)",
default=1.0
 )
parser.add_argument(
"-v",
"--vanilla_sample",
default=False,
action='store_true',
help="vanilla sampling (default option is DDIM sampling)?",
 )
parser.add_argument(
"-l",
"--logdir",
type=str,
nargs="?",
help="extra logdir",
default="none"
 )
parser.add_argument(
"-c",
"--custom_steps",
type=int,
nargs="?",
help="number of steps for ddim and fastdpm sampling",
default=50
 )
parser.add_argument(
"--batch_size",
type=int,
nargs="?",
help="the bs",
default=10
 )
returnparser


defload_model_from_config(config, sd):
model=instantiate_from_config(config)
model.load_state_dict(sd,strict=False)
model.cuda()
model.eval()
returnmodel


defload_model(config, ckpt, gpu, eval_mode):
ifckpt:
print(f"Loading model from {ckpt}")
pl_sd=torch.load(ckpt, map_location="cpu")
global_step=pl_sd["global_step"]
else:
pl_sd= {"state_dict": None}
global_step=None
model=load_model_from_config(config.model,
pl_sd["state_dict"])

returnmodel, global_step


if__name__=="__main__":
now=datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
sys.path.append(os.getcwd())
command=" ".join(sys.argv)

parser=get_parser()
opt, unknown=parser.parse_known_args()
ckpt=None

ifnotos.path.exists(opt.resume):
raiseValueError("Cannot find {}".format(opt.resume))
ifos.path.isfile(opt.resume):
# paths = opt.resume.split("/")
try:
logdir='/'.join(opt.resume.split('/')[:-1])
# idx = len(paths)-paths[::-1].index("logs")+1
print(f'Logdir is {logdir}')
exceptValueError:
paths=opt.resume.split("/")
idx=-2# take a guess: path/to/logdir/checkpoints/model.ckpt
logdir="/".join(paths[:idx])
ckpt=opt.resume
else:
assertos.path.isdir(opt.resume), f"{opt.resume} is not a directory"
logdir=opt.resume.rstrip("/")
ckpt=os.path.join(logdir, "model.ckpt")

base_configs=sorted(glob.glob(os.path.join(logdir, "config.yaml")))
opt.base=base_configs

configs= [OmegaConf.load(cfg) forcfginopt.base]
cli=OmegaConf.from_dotlist(unknown)
config=OmegaConf.merge(*configs, cli)

gpu=True
eval_mode=True

ifopt.logdir!="none":
locallog=logdir.split(os.sep)[-1]
iflocallog=="": locallog=logdir.split(os.sep)[-2]
print(f"Switching logdir from '{logdir}' to '{os.path.join(opt.logdir, locallog)}'")
logdir=os.path.join(opt.logdir, locallog)

print(config)

model, global_step=load_model(config, ckpt, gpu, eval_mode)
print(f"global step: {global_step}")
print(75*"=")
print("logging to:")
logdir=os.path.join(logdir, "samples", f"{global_step:08}", now)
imglogdir=os.path.join(logdir, "img")
numpylogdir=os.path.join(logdir, "numpy")

os.makedirs(imglogdir)
os.makedirs(numpylogdir)
print(logdir)
print(75*"=")

# write config out
sampling_file=os.path.join(logdir, "sampling_config.yaml")
sampling_conf=vars(opt)

withopen(sampling_file, 'w') asf:
yaml.dump(sampling_conf, f, default_flow_style=False)
print(sampling_conf)


run(model, imglogdir, eta=opt.eta,
vanilla=opt.vanilla_sample, n_samples=opt.n_samples, custom_steps=opt.custom_steps,
batch_size=opt.batch_size, nplog=numpylogdir)

print("done.")