File size: 5,190 Bytes
d643072 |
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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
"""Adapted from TODO"""
import argparse
import json
import os
import numpy as np
import torch
from diffusers import DiffusionPipeline, StableDiffusionPipeline
from einops import rearrange
from PIL import Image
from pytorch_lightning import seed_everything
from torchvision.transforms import ToTensor
from torchvision.utils import make_grid
from tqdm import tqdm, trange
torch.set_grad_enabled(False)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("metadata_file", type=str, help="JSONL file containing lines of metadata for each prompt")
parser.add_argument("--model", type=str, default="runwayml/stable-diffusion-v1-5", help="Huggingface model name")
parser.add_argument("--outdir", type=str, nargs="?", help="dir to write results to", default="outputs")
parser.add_argument(
"--n_samples",
type=int,
default=4,
help="number of samples",
)
parser.add_argument(
"--steps",
type=int,
default=50,
help="number of ddim sampling steps",
)
parser.add_argument(
"--negative-prompt",
type=str,
nargs="?",
const="ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, bad anatomy, watermark, signature, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face",
default=None,
help="negative prompt for guidance",
)
parser.add_argument(
"--H",
type=int,
default=None,
help="image height, in pixel space",
)
parser.add_argument(
"--W",
type=int,
default=None,
help="image width, in pixel space",
)
parser.add_argument(
"--scale",
type=float,
default=9.0,
help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))",
)
parser.add_argument(
"--seed",
type=int,
default=42,
help="the seed (for reproducible sampling)",
)
parser.add_argument(
"--batch_size",
type=int,
default=1,
help="how many samples can be produced simultaneously",
)
parser.add_argument(
"--skip_grid",
action="store_true",
help="skip saving grid",
)
opt = parser.parse_args()
return opt
def main(opt):
# Load prompts
with open(opt.metadata_file) as fp:
metadatas = [json.loads(line) for line in fp]
# Load model
if opt.model == "stabilityai/stable-diffusion-xl-base-1.0":
model = DiffusionPipeline.from_pretrained(
opt.model, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
)
model.enable_xformers_memory_efficient_attention()
else:
model = StableDiffusionPipeline.from_pretrained(opt.model, torch_dtype=torch.float16)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = model.to(device)
model.enable_attention_slicing()
for index, metadata in enumerate(metadatas):
seed_everything(opt.seed)
outpath = os.path.join(opt.outdir, f"{index:0>5}")
os.makedirs(outpath, exist_ok=True)
prompt = metadata["prompt"]
n_rows = batch_size = opt.batch_size
print(f"Prompt ({index: >3}/{len(metadatas)}): '{prompt}'")
sample_path = os.path.join(outpath, "samples")
os.makedirs(sample_path, exist_ok=True)
with open(os.path.join(outpath, "metadata.jsonl"), "w") as fp:
json.dump(metadata, fp)
sample_count = 0
with torch.no_grad():
all_samples = list()
for n in trange((opt.n_samples + batch_size - 1) // batch_size, desc="Sampling"):
# Generate images
samples = model(
prompt,
height=opt.H,
width=opt.W,
num_inference_steps=opt.steps,
guidance_scale=opt.scale,
num_images_per_prompt=min(batch_size, opt.n_samples - sample_count),
negative_prompt=opt.negative_prompt or None,
).images
for sample in samples:
sample.save(os.path.join(sample_path, f"{sample_count:05}.png"))
sample_count += 1
if not opt.skip_grid:
all_samples.append(torch.stack([ToTensor()(sample) for sample in samples], 0))
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=n_rows)
# to image
grid = 255.0 * rearrange(grid, "c h w -> h w c").cpu().numpy()
grid = Image.fromarray(grid.astype(np.uint8))
grid.save(os.path.join(outpath, f"grid.png"))
del grid
del all_samples
print("Done.")
if __name__ == "__main__":
opt = parse_args()
main(opt)
|