731 lines
22 KiB
Python
731 lines
22 KiB
Python
"""
|
|
clone the following repo if haven't
|
|
- git clone 'https://github.com/openai/CLIP'
|
|
- git clone 'https://github.com/CompVis/taming-transformers'
|
|
"""
|
|
|
|
import sys
|
|
import tempfile
|
|
import warnings
|
|
import numpy as np
|
|
from pathlib import Path
|
|
import argparse
|
|
import torch
|
|
from torch import nn, optim
|
|
from torch.nn import functional as F
|
|
from torchvision import transforms
|
|
from torchvision.transforms import functional as TF
|
|
from torch.cuda import get_device_properties
|
|
from omegaconf import OmegaConf
|
|
from torch_optimizer import DiffGrad, AdamP, RAdam
|
|
import kornia.augmentation as K
|
|
import imageio
|
|
from tqdm import tqdm
|
|
import cog
|
|
from CLIP import clip
|
|
from PIL import ImageFile, Image, PngImagePlugin, ImageChops
|
|
|
|
sys.path.append("taming-transformers")
|
|
from taming.models import cond_transformer, vqgan
|
|
|
|
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
|
torch.backends.cudnn.benchmark = False
|
|
warnings.filterwarnings("ignore")
|
|
|
|
|
|
class ReplaceGrad(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, x_forward, x_backward):
|
|
ctx.shape = x_backward.shape
|
|
return x_forward
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_in):
|
|
return None, grad_in.sum_to_size(ctx.shape)
|
|
|
|
|
|
class ClampWithGrad(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, input, min, max):
|
|
ctx.min = min
|
|
ctx.max = max
|
|
ctx.save_for_backward(input)
|
|
return input.clamp(min, max)
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_in):
|
|
(input,) = ctx.saved_tensors
|
|
return (
|
|
grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0),
|
|
None,
|
|
None,
|
|
)
|
|
|
|
|
|
replace_grad = ReplaceGrad.apply
|
|
clamp_with_grad = ClampWithGrad.apply
|
|
|
|
|
|
class Predictor(cog.Predictor):
|
|
def setup(self):
|
|
self.device = torch.device("cuda:0")
|
|
# Check for GPU and reduce the default image size if low VRAM
|
|
default_image_size = 512 # >8GB VRAM
|
|
if not torch.cuda.is_available():
|
|
default_image_size = 256 # no GPU found
|
|
elif (
|
|
get_device_properties(0).total_memory <= 2 ** 33
|
|
): # 2 ** 33 = 8,589,934,592 bytes = 8 GB
|
|
default_image_size = 318 # <8GB VRAM
|
|
|
|
self.args = get_args()
|
|
self.args.size = [default_image_size, default_image_size]
|
|
self.model = load_vqgan_model(
|
|
self.args.vqgan_config, self.args.vqgan_checkpoint
|
|
).to(self.device)
|
|
print("Model loaded!")
|
|
jit = True if float(torch.__version__[:3]) < 1.8 else False
|
|
self.perceptor = (
|
|
clip.load(self.args.clip_model, jit=jit)[0]
|
|
.eval()
|
|
.requires_grad_(False)
|
|
.to(self.device)
|
|
)
|
|
cut_size = self.perceptor.visual.input_resolution
|
|
# choose latest Cutout class as default
|
|
self.make_cutouts = MakeCutouts(
|
|
cut_size, self.args.cutn, self.args, cut_pow=self.args.cut_pow
|
|
)
|
|
|
|
self.z_min = self.model.quantize.embedding.weight.min(dim=0).values[
|
|
None, :, None, None
|
|
]
|
|
self.z_max = self.model.quantize.embedding.weight.max(dim=0).values[
|
|
None, :, None, None
|
|
]
|
|
|
|
print("Using device:", self.device)
|
|
print("Optimising using:", self.args.optimiser)
|
|
|
|
@cog.input(
|
|
"image",
|
|
type=Path,
|
|
default=None,
|
|
help="Initial Image, optional. When the image is provided, the prompts will be used to create some 'style transfer' effect",
|
|
)
|
|
@cog.input(
|
|
"prompts",
|
|
type=str,
|
|
default="A cute, smiling, Nerdy Rodent",
|
|
help="Prompts for generating images. Supports multiple prompts separated by pipe | ",
|
|
)
|
|
@cog.input(
|
|
"iterations",
|
|
type=int,
|
|
default=300,
|
|
help="total iterations for generating images. Set to lower iterations when initial image is uploaded",
|
|
)
|
|
@cog.input(
|
|
"display_frequency",
|
|
type=int,
|
|
default=20,
|
|
help="display frequency for intermediate generated images",
|
|
)
|
|
def predict(self, image, prompts, iterations, display_frequency):
|
|
# gumbel is False
|
|
e_dim = self.model.quantize.e_dim
|
|
n_toks = self.model.quantize.n_e
|
|
f = 2 ** (self.model.decoder.num_resolutions - 1)
|
|
toksX, toksY = self.args.size[0] // f, self.args.size[1] // f
|
|
sideX, sideY = toksX * f, toksY * f
|
|
|
|
if image is not None:
|
|
self.args.init_image = str(image)
|
|
self.args.step_size = 0.25
|
|
if "http" in self.args.init_image:
|
|
img = Image.open(urlopen(self.args.init_image))
|
|
else:
|
|
img = Image.open(self.args.init_image)
|
|
pil_image = img.convert("RGB")
|
|
pil_image = pil_image.resize((sideX, sideY), Image.LANCZOS)
|
|
pil_tensor = TF.to_tensor(pil_image)
|
|
z, *_ = self.model.encode(pil_tensor.to(self.device).unsqueeze(0) * 2 - 1)
|
|
else:
|
|
one_hot = F.one_hot(
|
|
torch.randint(n_toks, [toksY * toksX], device=self.device), n_toks
|
|
).float()
|
|
# gumbel is False
|
|
z = one_hot @ self.model.quantize.embedding.weight
|
|
z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2)
|
|
|
|
z_orig = z.clone()
|
|
z.requires_grad_(True)
|
|
|
|
self.opt = get_opt(self.args.optimiser, self.args.step_size, z)
|
|
|
|
self.args.display_freq = display_frequency
|
|
self.args.max_iterations = iterations
|
|
|
|
story_phrases = [phrase.strip() for phrase in prompts.split("^")]
|
|
|
|
# Make a list of all phrases
|
|
all_phrases = []
|
|
for phrase in story_phrases:
|
|
all_phrases.append(phrase.split("|"))
|
|
|
|
# First phrase
|
|
prompts = all_phrases[0]
|
|
|
|
pMs = []
|
|
for prompt in prompts:
|
|
txt, weight, stop = split_prompt(prompt)
|
|
embed = self.perceptor.encode_text(
|
|
clip.tokenize(txt).to(self.device)
|
|
).float()
|
|
pMs.append(Prompt(embed, weight, stop).to(self.device))
|
|
# args.image_prompts is None for now
|
|
# args.noise_prompt_seeds, args.noise_prompt_weights None for now
|
|
print(f"Using text prompts: {prompts}")
|
|
if self.args.init_image:
|
|
print(f"Using initial image: {self.args.init_image}")
|
|
|
|
if self.args.seed is None:
|
|
seed = torch.seed()
|
|
else:
|
|
seed = self.args.seed
|
|
torch.manual_seed(seed)
|
|
print(f"Using seed: {seed}")
|
|
i = 0 # Iteration counter
|
|
# j = 0 # Zoom video frame counter
|
|
# p = 1 # Phrase counter
|
|
# smoother = 0 # Smoother counter
|
|
# this_video_frame = 0 # for video styling
|
|
|
|
out_path = Path(tempfile.mkdtemp()) / "out.png"
|
|
# Do it
|
|
for i in range(1, self.args.max_iterations + 1):
|
|
self.opt.zero_grad(set_to_none=True)
|
|
lossAll = ascend_txt(
|
|
i, z, self.perceptor, self.args, self.model, self.make_cutouts, pMs
|
|
)
|
|
|
|
if i % self.args.display_freq == 0 and not i == self.args.max_iterations:
|
|
yield checkin(i, lossAll, prompts, self.model, z, out_path)
|
|
|
|
loss = sum(lossAll)
|
|
loss.backward()
|
|
self.opt.step()
|
|
|
|
# with torch.no_grad():
|
|
with torch.inference_mode():
|
|
z.copy_(z.maximum(self.z_min).minimum(self.z_max))
|
|
|
|
# Ready to stop yet?
|
|
if i == self.args.max_iterations:
|
|
yield checkin(i, lossAll, prompts, self.model, z, out_path)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def checkin(i, losses, prompts, model, z, outpath):
|
|
losses_str = ", ".join(f"{loss.item():g}" for loss in losses)
|
|
tqdm.write(f"i: {i}, loss: {sum(losses).item():g}, losses: {losses_str}")
|
|
out = synth(z, model)
|
|
info = PngImagePlugin.PngInfo()
|
|
info.add_text("comment", f"{prompts}")
|
|
TF.to_pil_image(out[0].cpu()).save(str(outpath), pnginfo=info)
|
|
return outpath
|
|
|
|
|
|
def get_args():
|
|
vq_parser = argparse.ArgumentParser(description="Image generation using VQGAN+CLIP")
|
|
|
|
# Add the arguments
|
|
vq_parser.add_argument(
|
|
"-p", "--prompts", type=str, help="Text prompts", default=None, dest="prompts"
|
|
)
|
|
vq_parser.add_argument(
|
|
"-ip",
|
|
"--image_prompts",
|
|
type=str,
|
|
help="Image prompts / target image",
|
|
default=[],
|
|
dest="image_prompts",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-i",
|
|
"--iterations",
|
|
type=int,
|
|
help="Number of iterations",
|
|
default=500,
|
|
dest="max_iterations",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-se",
|
|
"--save_every",
|
|
type=int,
|
|
help="Save image iterations",
|
|
default=50,
|
|
dest="display_freq",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-s",
|
|
"--size",
|
|
nargs=2,
|
|
type=int,
|
|
help="Image size (width height) (default: %(default)s)",
|
|
dest="size",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-ii",
|
|
"--init_image",
|
|
type=str,
|
|
help="Initial image",
|
|
default=None,
|
|
dest="init_image",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-in",
|
|
"--init_noise",
|
|
type=str,
|
|
help="Initial noise image (pixels or gradient)",
|
|
default=None,
|
|
dest="init_noise",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-iw",
|
|
"--init_weight",
|
|
type=float,
|
|
help="Initial weight",
|
|
default=0.0,
|
|
dest="init_weight",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-m",
|
|
"--clip_model",
|
|
type=str,
|
|
help="CLIP model (e.g. ViT-B/32, ViT-B/16)",
|
|
default="ViT-B/32",
|
|
dest="clip_model",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-conf",
|
|
"--vqgan_config",
|
|
type=str,
|
|
help="VQGAN config",
|
|
default=f"checkpoints/vqgan_imagenet_f16_16384.yaml",
|
|
dest="vqgan_config",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-ckpt",
|
|
"--vqgan_checkpoint",
|
|
type=str,
|
|
help="VQGAN checkpoint",
|
|
default=f"checkpoints/vqgan_imagenet_f16_16384.ckpt",
|
|
dest="vqgan_checkpoint",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-nps",
|
|
"--noise_prompt_seeds",
|
|
nargs="*",
|
|
type=int,
|
|
help="Noise prompt seeds",
|
|
default=[],
|
|
dest="noise_prompt_seeds",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-npw",
|
|
"--noise_prompt_weights",
|
|
nargs="*",
|
|
type=float,
|
|
help="Noise prompt weights",
|
|
default=[],
|
|
dest="noise_prompt_weights",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-lr",
|
|
"--learning_rate",
|
|
type=float,
|
|
help="Learning rate",
|
|
default=0.1,
|
|
dest="step_size",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-cutm",
|
|
"--cut_method",
|
|
type=str,
|
|
help="Cut method",
|
|
choices=["original", "updated", "nrupdated", "updatedpooling", "latest"],
|
|
default="latest",
|
|
dest="cut_method",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-cuts", "--num_cuts", type=int, help="Number of cuts", default=32, dest="cutn"
|
|
)
|
|
vq_parser.add_argument(
|
|
"-cutp",
|
|
"--cut_power",
|
|
type=float,
|
|
help="Cut power",
|
|
default=1.0,
|
|
dest="cut_pow",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-sd", "--seed", type=int, help="Seed", default=None, dest="seed"
|
|
)
|
|
vq_parser.add_argument(
|
|
"-opt",
|
|
"--optimiser",
|
|
type=str,
|
|
help="Optimiser",
|
|
choices=[
|
|
"Adam",
|
|
"AdamW",
|
|
"Adagrad",
|
|
"Adamax",
|
|
"DiffGrad",
|
|
"AdamP",
|
|
"RAdam",
|
|
"RMSprop",
|
|
],
|
|
default="Adam",
|
|
dest="optimiser",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-o",
|
|
"--output",
|
|
type=str,
|
|
help="Output filename",
|
|
default="output.png",
|
|
dest="output",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-vid",
|
|
"--video",
|
|
action="store_true",
|
|
help="Create video frames?",
|
|
dest="make_video",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-zvid",
|
|
"--zoom_video",
|
|
action="store_true",
|
|
help="Create zoom video?",
|
|
dest="make_zoom_video",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-zs",
|
|
"--zoom_start",
|
|
type=int,
|
|
help="Zoom start iteration",
|
|
default=0,
|
|
dest="zoom_start",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-zse",
|
|
"--zoom_save_every",
|
|
type=int,
|
|
help="Save zoom image iterations",
|
|
default=10,
|
|
dest="zoom_frequency",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-zsc",
|
|
"--zoom_scale",
|
|
type=float,
|
|
help="Zoom scale %",
|
|
default=0.99,
|
|
dest="zoom_scale",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-zsx",
|
|
"--zoom_shift_x",
|
|
type=int,
|
|
help="Zoom shift x (left/right) amount in pixels",
|
|
default=0,
|
|
dest="zoom_shift_x",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-zsy",
|
|
"--zoom_shift_y",
|
|
type=int,
|
|
help="Zoom shift y (up/down) amount in pixels",
|
|
default=0,
|
|
dest="zoom_shift_y",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-cpe",
|
|
"--change_prompt_every",
|
|
type=int,
|
|
help="Prompt change frequency",
|
|
default=0,
|
|
dest="prompt_frequency",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-vl",
|
|
"--video_length",
|
|
type=float,
|
|
help="Video length in seconds (not interpolated)",
|
|
default=10,
|
|
dest="video_length",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-ofps",
|
|
"--output_video_fps",
|
|
type=float,
|
|
help="Create an interpolated video (Nvidia GPU only) with this fps (min 10. best set to 30 or 60)",
|
|
default=30,
|
|
dest="output_video_fps",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-ifps",
|
|
"--input_video_fps",
|
|
type=float,
|
|
help="When creating an interpolated video, use this as the input fps to interpolate from (>0 & <ofps)",
|
|
default=15,
|
|
dest="input_video_fps",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-d",
|
|
"--deterministic",
|
|
action="store_true",
|
|
help="Enable cudnn.deterministic?",
|
|
dest="cudnn_determinism",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-aug",
|
|
"--augments",
|
|
nargs="+",
|
|
action="append",
|
|
type=str,
|
|
choices=["Ji", "Sh", "Gn", "Pe", "Ro", "Af", "Et", "Ts", "Cr", "Er", "Re"],
|
|
help="Enabled augments (latest vut method only)",
|
|
default=[["Af", "Pe", "Ji", "Er"]],
|
|
dest="augments",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-vsd",
|
|
"--video_style_dir",
|
|
type=str,
|
|
help="Directory with video frames to style",
|
|
default=None,
|
|
dest="video_style_dir",
|
|
)
|
|
vq_parser.add_argument(
|
|
"-cd",
|
|
"--cuda_device",
|
|
type=str,
|
|
help="Cuda device to use",
|
|
default="cuda:0",
|
|
dest="cuda_device",
|
|
)
|
|
|
|
# Execute the parse_args() method
|
|
args = vq_parser.parse_args("")
|
|
return args
|
|
|
|
|
|
def load_vqgan_model(config_path, checkpoint_path):
|
|
config = OmegaConf.load(config_path)
|
|
# config.model.target == 'taming.models.vqgan.VQModel':
|
|
model = vqgan.VQModel(**config.model.params)
|
|
model.eval().requires_grad_(False)
|
|
model.init_from_ckpt(checkpoint_path)
|
|
del model.loss
|
|
return model
|
|
|
|
|
|
class MakeCutouts(nn.Module):
|
|
def __init__(self, cut_size, cutn, args, cut_pow=1.0):
|
|
super().__init__()
|
|
self.cut_size = cut_size
|
|
self.cutn = cutn
|
|
self.cut_pow = cut_pow # not used with pooling
|
|
|
|
# Pick your own augments & their order
|
|
augment_list = []
|
|
for item in args.augments[0]:
|
|
if item == "Ji":
|
|
augment_list.append(
|
|
K.ColorJitter(
|
|
brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1, p=0.7
|
|
)
|
|
)
|
|
elif item == "Sh":
|
|
augment_list.append(K.RandomSharpness(sharpness=0.3, p=0.5))
|
|
elif item == "Gn":
|
|
augment_list.append(K.RandomGaussianNoise(mean=0.0, std=1.0, p=0.5))
|
|
elif item == "Pe":
|
|
augment_list.append(K.RandomPerspective(distortion_scale=0.7, p=0.7))
|
|
elif item == "Ro":
|
|
augment_list.append(K.RandomRotation(degrees=15, p=0.7))
|
|
elif item == "Af":
|
|
augment_list.append(
|
|
K.RandomAffine(
|
|
degrees=15,
|
|
translate=0.1,
|
|
shear=5,
|
|
p=0.7,
|
|
padding_mode="zeros",
|
|
keepdim=True,
|
|
)
|
|
) # border, reflection, zeros
|
|
elif item == "Et":
|
|
augment_list.append(K.RandomElasticTransform(p=0.7))
|
|
elif item == "Ts":
|
|
augment_list.append(
|
|
K.RandomThinPlateSpline(scale=0.8, same_on_batch=True, p=0.7)
|
|
)
|
|
elif item == "Cr":
|
|
augment_list.append(
|
|
K.RandomCrop(
|
|
size=(self.cut_size, self.cut_size),
|
|
pad_if_needed=True,
|
|
padding_mode="reflect",
|
|
p=0.5,
|
|
)
|
|
)
|
|
elif item == "Er":
|
|
augment_list.append(
|
|
K.RandomErasing(
|
|
scale=(0.1, 0.4),
|
|
ratio=(0.3, 1 / 0.3),
|
|
same_on_batch=True,
|
|
p=0.7,
|
|
)
|
|
)
|
|
elif item == "Re":
|
|
augment_list.append(
|
|
K.RandomResizedCrop(
|
|
size=(self.cut_size, self.cut_size),
|
|
scale=(0.1, 1),
|
|
ratio=(0.75, 1.333),
|
|
cropping_mode="resample",
|
|
p=0.5,
|
|
)
|
|
)
|
|
|
|
self.augs = nn.Sequential(*augment_list)
|
|
self.noise_fac = 0.1
|
|
# self.noise_fac = False
|
|
|
|
# Uncomment if you like seeing the list ;)
|
|
# print(augment_list)
|
|
|
|
# Pooling
|
|
self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))
|
|
self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))
|
|
|
|
def forward(self, input):
|
|
cutouts = []
|
|
|
|
for _ in range(self.cutn):
|
|
# Use Pooling
|
|
cutout = (self.av_pool(input) + self.max_pool(input)) / 2
|
|
cutouts.append(cutout)
|
|
|
|
batch = self.augs(torch.cat(cutouts, dim=0))
|
|
|
|
if self.noise_fac:
|
|
facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)
|
|
batch = batch + facs * torch.randn_like(batch)
|
|
return batch
|
|
|
|
|
|
def get_opt(opt_name, opt_lr, z):
|
|
if opt_name == "Adam":
|
|
opt = optim.Adam([z], lr=opt_lr) # LR=0.1 (Default)
|
|
elif opt_name == "AdamW":
|
|
opt = optim.AdamW([z], lr=opt_lr)
|
|
elif opt_name == "Adagrad":
|
|
opt = optim.Adagrad([z], lr=opt_lr)
|
|
elif opt_name == "Adamax":
|
|
opt = optim.Adamax([z], lr=opt_lr)
|
|
elif opt_name == "DiffGrad":
|
|
opt = DiffGrad(
|
|
[z], lr=opt_lr, eps=1e-9, weight_decay=1e-9
|
|
) # NR: Playing for reasons
|
|
elif opt_name == "AdamP":
|
|
opt = AdamP([z], lr=opt_lr)
|
|
elif opt_name == "RAdam":
|
|
opt = RAdam([z], lr=opt_lr)
|
|
elif opt_name == "RMSprop":
|
|
opt = optim.RMSprop([z], lr=opt_lr)
|
|
else:
|
|
print("Unknown optimiser. Are choices broken?")
|
|
opt = optim.Adam([z], lr=opt_lr)
|
|
return opt
|
|
|
|
|
|
def ascend_txt(i, z, perceptor, args, model, make_cutouts, pMs):
|
|
normalize = transforms.Normalize(
|
|
mean=[0.48145466, 0.4578275, 0.40821073],
|
|
std=[0.26862954, 0.26130258, 0.27577711],
|
|
)
|
|
out = synth(z, model)
|
|
iii = perceptor.encode_image(normalize(make_cutouts(out))).float()
|
|
|
|
result = []
|
|
|
|
if args.init_weight:
|
|
# result.append(F.mse_loss(z, z_orig) * args.init_weight / 2)
|
|
result.append(
|
|
F.mse_loss(z, torch.zeros_like(z_orig))
|
|
* ((1 / torch.tensor(i * 2 + 1)) * args.init_weight)
|
|
/ 2
|
|
)
|
|
|
|
for prompt in pMs:
|
|
result.append(prompt(iii))
|
|
|
|
if args.make_video:
|
|
img = np.array(
|
|
out.mul(255).clamp(0, 255)[0].cpu().detach().numpy().astype(np.uint8)
|
|
)[:, :, :]
|
|
img = np.transpose(img, (1, 2, 0))
|
|
imageio.imwrite("steps/" + str(i) + ".png", np.array(img))
|
|
|
|
return result
|
|
|
|
|
|
def synth(z, model):
|
|
# gumbel is False
|
|
z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(
|
|
3, 1
|
|
)
|
|
return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)
|
|
|
|
|
|
def vector_quantize(x, codebook):
|
|
d = (
|
|
x.pow(2).sum(dim=-1, keepdim=True)
|
|
+ codebook.pow(2).sum(dim=1)
|
|
- 2 * x @ codebook.T
|
|
)
|
|
indices = d.argmin(-1)
|
|
x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook
|
|
return replace_grad(x_q, x)
|
|
|
|
|
|
def split_prompt(prompt):
|
|
vals = prompt.rsplit(":", 2)
|
|
vals = vals + ["", "1", "-inf"][len(vals) :]
|
|
return vals[0], float(vals[1]), float(vals[2])
|
|
|
|
|
|
class Prompt(nn.Module):
|
|
def __init__(self, embed, weight=1.0, stop=float("-inf")):
|
|
super().__init__()
|
|
self.register_buffer("embed", embed)
|
|
self.register_buffer("weight", torch.as_tensor(weight))
|
|
self.register_buffer("stop", torch.as_tensor(stop))
|
|
|
|
def forward(self, input):
|
|
input_normed = F.normalize(input.unsqueeze(1), dim=2)
|
|
embed_normed = F.normalize(self.embed.unsqueeze(0), dim=2)
|
|
dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)
|
|
dists = dists * self.weight.sign()
|
|
return (
|
|
self.weight.abs()
|
|
* replace_grad(dists, torch.maximum(dists, self.stop)).mean()
|
|
)
|