mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
add code
This commit is contained in:
ablattmann
41
scripts/download_first_stages.sh
Normal file
41
scripts/download_first_stages.sh
Normal file
@ -0,0 +1,41 @@
|
||||
#!/bin/bash
|
||||
wget -O models/first_stage_models/kl-f4/model.zip https://ommer-lab.com/files/latent-diffusion/kl-f4.zip
|
||||
wget -O models/first_stage_models/kl-f8/model.zip https://ommer-lab.com/files/latent-diffusion/kl-f8.zip
|
||||
wget -O models/first_stage_models/kl-f16/model.zip https://ommer-lab.com/files/latent-diffusion/kl-f16.zip
|
||||
wget -O models/first_stage_models/kl-f32/model.zip https://ommer-lab.com/files/latent-diffusion/kl-f32.zip
|
||||
wget -O models/first_stage_models/vq-f4/model.zip https://ommer-lab.com/files/latent-diffusion/vq-f4.zip
|
||||
wget -O models/first_stage_models/vq-f4-noattn/model.zip https://heibox.uni-heidelberg.de/f/9c6681f64bb94338a069/?dl=1
|
||||
wget -O models/first_stage_models/vq-f8/model.zip https://ommer-lab.com/files/latent-diffusion/vq-f8.zip
|
||||
wget -O models/first_stage_models/vq-f8-n256/model.zip https://ommer-lab.com/files/latent-diffusion/vq-f8-n256.zip
|
||||
wget -O models/first_stage_models/vq-f16/model.zip https://heibox.uni-heidelberg.de/f/0e42b04e2e904890a9b6/?dl=1
|
||||
|
||||
|
||||
|
||||
cd models/first_stage_models/kl-f4
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../kl-f8
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../kl-f16
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../kl-f32
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../vq-f4
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../vq-f4-noattn
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../vq-f8
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../vq-f8-n256
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../vq-f16
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../..
|
||||
45
scripts/download_models.sh
Normal file
45
scripts/download_models.sh
Normal file
@ -0,0 +1,45 @@
|
||||
#!/bin/bash
|
||||
wget -O models/ldm/celeba256/celeba-256.zip https://ommer-lab.com/files/latent-diffusion/celeba.zip
|
||||
wget -O models/ldm/ffhq256/ffhq-256.zip https://ommer-lab.com/files/latent-diffusion/ffhq.zip
|
||||
wget -O models/ldm/lsun_churches256/lsun_churches-256.zip https://ommer-lab.com/files/latent-diffusion/lsun_churches.zip
|
||||
wget -O models/ldm/lsun_beds256/lsun_beds-256.zip https://ommer-lab.com/files/latent-diffusion/lsun_bedrooms.zip
|
||||
wget -O models/ldm/text2img256/model.zip https://ommer-lab.com/files/latent-diffusion/text2img.zip
|
||||
wget -O models/ldm/cin256/model.zip https://ommer-lab.com/files/latent-diffusion/cin.zip
|
||||
wget -O models/ldm/semantic_synthesis512/model.zip https://ommer-lab.com/files/latent-diffusion/semantic_synthesis.zip
|
||||
wget -O models/ldm/bsr_sr/model.zip https://ommer-lab.com/files/latent-diffusion/sr_bsr.zip
|
||||
wget -O models/ldm/layout2img-openimages256/model.zip https://ommer-lab.com/files/latent-diffusion/layout2img_model.zip
|
||||
wget -O models/ldm/inpainting_big/model.zip #TODO @patrick add path
|
||||
|
||||
|
||||
|
||||
cd models/ldm/celeba256
|
||||
unzip -o celeba-256.zip
|
||||
|
||||
cd ../ffhq256
|
||||
unzip -o ffhq-256.zip
|
||||
|
||||
cd ../lsun_churches256
|
||||
unzip -o lsun_churches-256.zip
|
||||
|
||||
cd ../lsun_beds256
|
||||
unzip -o lsun_beds-256.zip
|
||||
|
||||
cd ../text2img256
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../cin256
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../semantic_synthesis512
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../bsr_sr
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../layout2img-openimages256
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../inpainting_big
|
||||
unzip -o model.zip
|
||||
|
||||
cd ../..
|
||||
98
scripts/inpaint.py
Normal file
98
scripts/inpaint.py
Normal file
@ -0,0 +1,98 @@
|
||||
import argparse, os, sys, glob
|
||||
from omegaconf import OmegaConf
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
import torch
|
||||
from main import instantiate_from_config
|
||||
from ldm.models.diffusion.ddim import DDIMSampler
|
||||
|
||||
|
||||
def make_batch(image, mask, device):
|
||||
image = np.array(Image.open(image).convert("RGB"))
|
||||
image = image.astype(np.float32)/255.0
|
||||
image = image[None].transpose(0,3,1,2)
|
||||
image = torch.from_numpy(image)
|
||||
|
||||
mask = np.array(Image.open(mask).convert("L"))
|
||||
mask = mask.astype(np.float32)/255.0
|
||||
mask = mask[None,None]
|
||||
mask[mask < 0.5] = 0
|
||||
mask[mask >= 0.5] = 1
|
||||
mask = torch.from_numpy(mask)
|
||||
|
||||
masked_image = (1-mask)*image
|
||||
|
||||
batch = {"image": image, "mask": mask, "masked_image": masked_image}
|
||||
for k in batch:
|
||||
batch[k] = batch[k].to(device=device)
|
||||
batch[k] = batch[k]*2.0-1.0
|
||||
return batch
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--indir",
|
||||
type=str,
|
||||
nargs="?",
|
||||
help="dir containing image-mask pairs (`example.png` and `example_mask.png`)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--outdir",
|
||||
type=str,
|
||||
nargs="?",
|
||||
help="dir to write results to",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--steps",
|
||||
type=int,
|
||||
default=50,
|
||||
help="number of ddim sampling steps",
|
||||
)
|
||||
opt = parser.parse_args()
|
||||
|
||||
masks = sorted(glob.glob(os.path.join(opt.indir, "*_mask.png")))
|
||||
images = [x.replace("_mask.png", ".png") for x in masks]
|
||||
print(f"Found {len(masks)} inputs.")
|
||||
|
||||
config = OmegaConf.load("models/ldm/inpainting_big/config.yaml")
|
||||
model = instantiate_from_config(config.model)
|
||||
model.load_state_dict(torch.load("models/ldm/inpainting_big/last.ckpt")["state_dict"],
|
||||
strict=False)
|
||||
|
||||
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
|
||||
model = model.to(device)
|
||||
sampler = DDIMSampler(model)
|
||||
|
||||
os.makedirs(opt.outdir, exist_ok=True)
|
||||
with torch.no_grad():
|
||||
with model.ema_scope():
|
||||
for image, mask in tqdm(zip(images, masks)):
|
||||
outpath = os.path.join(opt.outdir, os.path.split(image)[1])
|
||||
batch = make_batch(image, mask, device=device)
|
||||
|
||||
# encode masked image and concat downsampled mask
|
||||
c = model.cond_stage_model.encode(batch["masked_image"])
|
||||
cc = torch.nn.functional.interpolate(batch["mask"],
|
||||
size=c.shape[-2:])
|
||||
c = torch.cat((c, cc), dim=1)
|
||||
|
||||
shape = (c.shape[1]-1,)+c.shape[2:]
|
||||
samples_ddim, _ = sampler.sample(S=opt.steps,
|
||||
conditioning=c,
|
||||
batch_size=c.shape[0],
|
||||
shape=shape,
|
||||
verbose=False)
|
||||
x_samples_ddim = model.decode_first_stage(samples_ddim)
|
||||
|
||||
image = torch.clamp((batch["image"]+1.0)/2.0,
|
||||
min=0.0, max=1.0)
|
||||
mask = torch.clamp((batch["mask"]+1.0)/2.0,
|
||||
min=0.0, max=1.0)
|
||||
predicted_image = torch.clamp((x_samples_ddim+1.0)/2.0,
|
||||
min=0.0, max=1.0)
|
||||
|
||||
inpainted = (1-mask)*image+mask*predicted_image
|
||||
inpainted = inpainted.cpu().numpy().transpose(0,2,3,1)[0]*255
|
||||
Image.fromarray(inpainted.astype(np.uint8)).save(outpath)
|
||||
313
scripts/sample_diffusion.py
Normal file
313
scripts/sample_diffusion.py
Normal file
@ -0,0 +1,313 @@
|
||||
import argparse, os, sys, glob, datetime, yaml
|
||||
import torch
|
||||
import time
|
||||
import numpy as np
|
||||
from tqdm import trange
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
from PIL import Image
|
||||
|
||||
from ldm.models.diffusion.ddim import DDIMSampler
|
||||
from ldm.util import instantiate_from_config
|
||||
|
||||
rescale = lambda x: (x + 1.) / 2.
|
||||
|
||||
def custom_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)
|
||||
if not x.mode == "RGB":
|
||||
x = x.convert("RGB")
|
||||
return x
|
||||
|
||||
|
||||
def custom_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()
|
||||
return sample
|
||||
|
||||
|
||||
def logs2pil(logs, keys=["sample"]):
|
||||
imgs = dict()
|
||||
for k in logs:
|
||||
try:
|
||||
if len(logs[k].shape) == 4:
|
||||
img = custom_to_pil(logs[k][0, ...])
|
||||
elif len(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
|
||||
return imgs
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def convsample(model, shape, return_intermediates=True,
|
||||
verbose=True,
|
||||
make_prog_row=False):
|
||||
|
||||
|
||||
if not make_prog_row:
|
||||
return model.p_sample_loop(None, shape,
|
||||
return_intermediates=return_intermediates, verbose=verbose)
|
||||
else:
|
||||
return model.progressive_denoising(
|
||||
None, shape, verbose=True
|
||||
)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def convsample_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,)
|
||||
return samples, intermediates
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def make_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]
|
||||
|
||||
with model.ema_scope("Plotting"):
|
||||
t0 = time.time()
|
||||
if vanilla:
|
||||
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"]}')
|
||||
return log
|
||||
|
||||
def run(model, logdir, batch_size=50, vanilla=False, custom_steps=None, eta=None, n_samples=50000, nplog=None):
|
||||
if vanilla:
|
||||
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
|
||||
if model.cond_stage_model is None:
|
||||
all_images = []
|
||||
|
||||
print(f"Running unconditional sampling for {n_samples} samples")
|
||||
for _ in trange(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"])])
|
||||
if n_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) for x in all_img.shape])
|
||||
nppath = os.path.join(nplog, f"{shape_str}-samples.npz")
|
||||
np.savez(nppath, all_img)
|
||||
|
||||
else:
|
||||
raise NotImplementedError('Currently only sampling for unconditional models supported.')
|
||||
|
||||
print(f"sampling of {n_saved} images finished in {(time.time() - tstart) / 60.:.2f} minutes.")
|
||||
|
||||
|
||||
def save_logs(logs, path, n_saved=0, key="sample", np_path=None):
|
||||
for k in logs:
|
||||
if k == key:
|
||||
batch = logs[key]
|
||||
if np_path is None:
|
||||
for x in batch:
|
||||
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) for x in npbatch.shape])
|
||||
nppath = os.path.join(np_path, f"{n_saved}-{shape_str}-samples.npz")
|
||||
np.savez(nppath, npbatch)
|
||||
n_saved += npbatch.shape[0]
|
||||
return n_saved
|
||||
|
||||
|
||||
def get_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
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def load_model_from_config(config, sd):
|
||||
model = instantiate_from_config(config)
|
||||
model.load_state_dict(sd,strict=False)
|
||||
model.cuda()
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def load_model(config, ckpt, gpu, eval_mode):
|
||||
if ckpt:
|
||||
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"])
|
||||
|
||||
return model, 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
|
||||
|
||||
if not os.path.exists(opt.resume):
|
||||
raise ValueError("Cannot find {}".format(opt.resume))
|
||||
if os.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}')
|
||||
except ValueError:
|
||||
paths = opt.resume.split("/")
|
||||
idx = -2 # take a guess: path/to/logdir/checkpoints/model.ckpt
|
||||
logdir = "/".join(paths[:idx])
|
||||
ckpt = opt.resume
|
||||
else:
|
||||
assert os.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) for cfg in opt.base]
|
||||
cli = OmegaConf.from_dotlist(unknown)
|
||||
config = OmegaConf.merge(*configs, cli)
|
||||
|
||||
gpu = True
|
||||
eval_mode = True
|
||||
|
||||
if opt.logdir != "none":
|
||||
locallog = logdir.split(os.sep)[-1]
|
||||
if locallog == "": 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)
|
||||
|
||||
with open(sampling_file, 'w') as f:
|
||||
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.")
|
||||
Reference in New Issue
Block a user