stable-diffusion-finetune/scripts/txt2img.py

218 lines
6.7 KiB
Python
Raw Normal View History

2022-04-04 16:17:48 +02:00
import argparse, os, sys, glob
import torch
import numpy as np
from omegaconf import OmegaConf
from PIL import Image
from tqdm import tqdm, trange
2022-06-05 19:22:43 +02:00
from itertools import islice
2022-04-04 16:17:48 +02:00
from einops import rearrange
from torchvision.utils import make_grid
from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler
2022-04-04 23:43:15 +02:00
from ldm.models.diffusion.plms import PLMSSampler
2022-04-04 16:17:48 +02:00
2022-06-05 19:22:43 +02:00
def chunk(it, size):
it = iter(it)
return iter(lambda: tuple(islice(it, size)), ())
2022-04-04 16:17:48 +02:00
def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
sd = pl_sd["state_dict"]
model = instantiate_from_config(config.model)
m, u = model.load_state_dict(sd, strict=False)
if len(m) > 0 and verbose:
print("missing keys:")
print(m)
if len(u) > 0 and verbose:
print("unexpected keys:")
print(u)
model.cuda()
model.eval()
return model
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--prompt",
type=str,
nargs="?",
default="a painting of a virus monster playing guitar",
help="the prompt to render"
)
parser.add_argument(
"--outdir",
type=str,
nargs="?",
help="dir to write results to",
default="outputs/txt2img-samples"
)
2022-06-10 10:54:14 +02:00
parser.add_argument(
"--skip_grid",
action='store_true',
help="do not save a grid, only individual samples. Helpful when evaluating lots of samples",
)
2022-04-04 16:17:48 +02:00
parser.add_argument(
"--ddim_steps",
type=int,
2022-06-05 19:22:43 +02:00
default=50,
2022-04-04 16:17:48 +02:00
help="number of ddim sampling steps",
)
2022-04-04 23:43:15 +02:00
parser.add_argument(
"--plms",
action='store_true',
help="use plms sampling",
)
2022-04-04 16:17:48 +02:00
parser.add_argument(
"--ddim_eta",
type=float,
default=0.0,
help="ddim eta (eta=0.0 corresponds to deterministic sampling",
)
parser.add_argument(
"--n_iter",
type=int,
default=1,
help="sample this often",
)
parser.add_argument(
"--H",
type=int,
default=256,
help="image height, in pixel space",
)
parser.add_argument(
"--W",
type=int,
default=256,
help="image width, in pixel space",
)
parser.add_argument(
"--n_samples",
type=int,
2022-06-05 19:22:43 +02:00
default=8,
help="how many samples to produce for each given prompt. A.k.a batch size",
2022-04-04 16:17:48 +02:00
)
parser.add_argument(
"--scale",
type=float,
default=5.0,
help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))",
)
2022-06-05 19:22:43 +02:00
parser.add_argument(
"--dyn",
type=float,
help="dynamic thresholding from Imagen, in latent space (TODO: try in pixel space with intermediate decode)",
)
parser.add_argument(
"--from-file",
type=str,
help="if specified, load prompts from this file",
)
parser.add_argument(
"--config",
type=str,
default="logs/f8-kl-clip-encoder-256x256-run1/configs/2022-06-01T22-11-40-project.yaml",
help="path to config which constructs model",
)
parser.add_argument(
"--ckpt",
type=str,
default="logs/f8-kl-clip-encoder-256x256-run1/checkpoints/last.ckpt",
help="path to checkpoint of model",
)
2022-04-04 16:17:48 +02:00
opt = parser.parse_args()
2022-06-05 19:22:43 +02:00
config = OmegaConf.load(f"{opt.config}")
model = load_model_from_config(config, f"{opt.ckpt}")
2022-04-04 16:17:48 +02:00
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = model.to(device)
2022-04-04 23:43:15 +02:00
if opt.plms:
sampler = PLMSSampler(model)
else:
sampler = DDIMSampler(model)
2022-04-04 16:17:48 +02:00
os.makedirs(opt.outdir, exist_ok=True)
outpath = opt.outdir
2022-06-05 19:22:43 +02:00
batch_size = opt.n_samples
if not opt.from_file:
prompt = opt.prompt
assert prompt is not None
data = [batch_size * [prompt]]
2022-04-04 16:17:48 +02:00
2022-06-05 19:22:43 +02:00
else:
print(f"reading prompts from {opt.from_file}")
with open(opt.from_file, "r") as f:
data = f.read().splitlines()
data = list(chunk(data, batch_size))
2022-04-04 16:17:48 +02:00
sample_path = os.path.join(outpath, "samples")
os.makedirs(sample_path, exist_ok=True)
base_count = len(os.listdir(sample_path))
2022-06-05 19:22:43 +02:00
grid_count = len(os.listdir(outpath)) - 1
2022-04-04 16:17:48 +02:00
with torch.no_grad():
with model.ema_scope():
for n in trange(opt.n_iter, desc="Sampling"):
2022-06-05 19:22:43 +02:00
all_samples = list()
for prompts in tqdm(data, desc="data"):
uc = None
if opt.scale != 1.0:
uc = model.get_learned_conditioning(batch_size * [""])
2022-06-10 10:54:14 +02:00
if isinstance(prompts, tuple):
prompts = list(prompts)
2022-06-05 19:22:43 +02:00
c = model.get_learned_conditioning(prompts)
shape = [4, opt.H//8, opt.W//8]
samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
conditioning=c,
batch_size=opt.n_samples,
shape=shape,
verbose=False,
unconditional_guidance_scale=opt.scale,
unconditional_conditioning=uc,
eta=opt.ddim_eta,
dynamic_threshold=opt.dyn)
x_samples_ddim = model.decode_first_stage(samples_ddim)
x_samples_ddim = torch.clamp((x_samples_ddim+1.0)/2.0, min=0.0, max=1.0)
for x_sample in x_samples_ddim:
x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
Image.fromarray(x_sample.astype(np.uint8)).save(os.path.join(sample_path, f"{base_count:05}.png"))
base_count += 1
all_samples.append(x_samples_ddim)
2022-06-10 10:54:14 +02:00
if not opt.skip_grid:
# additionally, save as grid
grid = torch.stack(all_samples, 0)
grid = rearrange(grid, 'n b c h w -> (n b) c h w')
grid = make_grid(grid, nrow=opt.n_samples)
2022-06-05 19:22:43 +02:00
2022-06-10 10:54:14 +02:00
# to image
grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
Image.fromarray(grid.astype(np.uint8)).save(os.path.join(outpath, f'grid-{grid_count:04}.png'))
grid_count += 1
2022-06-05 19:22:43 +02:00
print(f"Your samples are ready and waiting for you here: \n{outpath} \nEnjoy.")