import torch
from transformers import CLIPModel, CLIPTextModel, CLIPTokenizer
from omegaconf import OmegaConf
import math
import imageio
from PIL import Image
import torchvision
import torch.nn.functional as F
import torch
import numpy as np
from PIL import Image
import time
import datetime
import torch
import sys
import os
from torchvision import datasets
import pickle
# StableDiffusion P2P implementation originally from https://github.com/bloc97/CrossAttentionControl
use_half_prec = True
if use_half_prec:
from my_half_diffusers import AutoencoderKL, UNet2DConditionModel
from my_half_diffusers.schedulers.scheduling_utils import SchedulerOutput
from my_half_diffusers import LMSDiscreteScheduler, PNDMScheduler, DDPMScheduler, DDIMScheduler
else:
from my_diffusers import AutoencoderKL, UNet2DConditionModel
from my_diffusers.schedulers.scheduling_utils import SchedulerOutput
from my_diffusers import LMSDiscreteScheduler, PNDMScheduler, DDPMScheduler, DDIMScheduler
torch_dtype = torch.float16 if use_half_prec else torch.float64
np_dtype = np.float16 if use_half_prec else np.float64
import random
from tqdm.auto import tqdm
from torch import autocast
from difflib import SequenceMatcher
# Build our CLIP model
model_path_clip = "openai/clip-vit-large-patch14"
clip_tokenizer = CLIPTokenizer.from_pretrained(model_path_clip)
clip_model = CLIPModel.from_pretrained(model_path_clip, torch_dtype=torch_dtype)
clip = clip_model.text_model
# Getting our HF Auth token
auth_token = os.environ.get('auth_token')
if auth_token is None:
with open('hf_auth', 'r') as f:
auth_token = f.readlines()[0].strip()
model_path_diffusion = "CompVis/stable-diffusion-v1-4"
# Build our SD model
unet = UNet2DConditionModel.from_pretrained(model_path_diffusion, subfolder="unet", use_auth_token=auth_token, revision="fp16", torch_dtype=torch_dtype)
vae = AutoencoderKL.from_pretrained(model_path_diffusion, subfolder="vae", use_auth_token=auth_token, revision="fp16", torch_dtype=torch_dtype)
# Push to devices w/ double precision
device = 'cuda'
if use_half_prec:
unet.to(device)
vae.to(device)
clip.to(device)
else:
unet.double().to(device)
vae.double().to(device)
clip.double().to(device)
print("Loaded all models")
def EDICT_editing(im_path,
base_prompt,
edit_prompt,
use_p2p=False,
steps=50,
mix_weight=0.93,
init_image_strength=0.8,
guidance_scale=3,
run_baseline=False,
width=512, height=512):
"""
Main call of our research, performs editing with either EDICT or DDIM
Args:
im_path: path to image to run on
base_prompt: conditional prompt to deterministically noise with
edit_prompt: desired text conditoining
steps: ddim steps
mix_weight: Weight of mixing layers.
Higher means more consistent generations but divergence in inversion
Lower means opposite
This is fairly tuned and can get good results
init_image_strength: Editing strength. Higher = more dramatic edit.
Typically [0.6, 0.9] is good range.
Definitely tunable per-image/maybe best results are at a different value
guidance_scale: classifier-free guidance scale
3 I've found is the best for both our method and basic DDIM inversion
Higher can result in more distorted results
run_baseline:
VERY IMPORTANT
True is EDICT, False is DDIM
Output:
PAIR of Images (tuple)
If run_baseline=True then [0] will be edit and [1] will be original
If run_baseline=False then they will be two nearly identical edited versions
"""
# Resize/center crop to 512x512 (Can do higher res. if desired)
if isinstance(im_path, str):
orig_im = load_im_into_format_from_path(im_path)
elif Image.isImageType(im_path):
width, height = im_path.size
# add max dim for sake of memory
max_dim = max(width, height)
if max_dim > 1024:
factor = 1024 / max_dim
width *= factor
height *= factor
width = int(width)
height = int(height)
im_path = im_path.resize((width, height))
min_dim = min(width, height)
if min_dim < 512:
factor = 512 / min_dim
width *= factor
height *= factor
width = int(width)
height = int(height)
im_path = im_path.resize((width, height))
width = width - (width%64)
height = height - (height%64)
orig_im = im_path # general_crop(im_path, width, height)
else:
orig_im = im_path
# compute latent pair (second one will be original latent if run_baseline=True)
latents = coupled_stablediffusion(base_prompt,
reverse=True,
init_image=orig_im,
init_image_strength=init_image_strength,
steps=steps,
mix_weight=mix_weight,
guidance_scale=guidance_scale,
run_baseline=run_baseline,
width=width, height=height)
# Denoise intermediate state with new conditioning
gen = coupled_stablediffusion(edit_prompt if (not use_p2p) else base_prompt,
None if (not use_p2p) else edit_prompt,
fixed_starting_latent=latents,
init_image_strength=init_image_strength,
steps=steps,
mix_weight=mix_weight,
guidance_scale=guidance_scale,
run_baseline=run_baseline,
width=width, height=height)
return gen
def img2img_editing(im_path,
edit_prompt,
steps=50,
init_image_strength=0.7,
guidance_scale=3):
"""
Basic SDEdit/img2img, given an image add some noise and denoise with prompt
"""
orig_im = load_im_into_format_from_path(im_path)
return baseline_stablediffusion(edit_prompt,
init_image_strength=init_image_strength,
steps=steps,
init_image=orig_im,
guidance_scale=guidance_scale)
def center_crop(im):
width, height = im.size # Get dimensions
min_dim = min(width, height)
left = (width - min_dim)/2
top = (height - min_dim)/2
right = (width + min_dim)/2
bottom = (height + min_dim)/2
# Crop the center of the image
im = im.crop((left, top, right, bottom))
return im
def general_crop(im, target_w, target_h):
width, height = im.size # Get dimensions
min_dim = min(width, height)
left = target_w / 2 # (width - min_dim)/2
top = target_h / 2 # (height - min_dim)/2
right = width - (target_w / 2) # (width + min_dim)/2
bottom = height - (target_h / 2) # (height + min_dim)/2
# Crop the center of the image
im = im.crop((left, top, right, bottom))
return im
def load_im_into_format_from_path(im_path):
return center_crop(Image.open(im_path)).resize((512,512))
#### P2P STUFF ####
def init_attention_weights(weight_tuples):
tokens_length = clip_tokenizer.model_max_length
weights = torch.ones(tokens_length)
for i, w in weight_tuples:
if i < tokens_length and i >= 0:
weights[i] = w
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention" and "attn2" in name:
module.last_attn_slice_weights = weights.to(device)
if module_name == "CrossAttention" and "attn1" in name:
module.last_attn_slice_weights = None
def init_attention_edit(tokens, tokens_edit):
tokens_length = clip_tokenizer.model_max_length
mask = torch.zeros(tokens_length)
indices_target = torch.arange(tokens_length, dtype=torch.long)
indices = torch.zeros(tokens_length, dtype=torch.long)
tokens = tokens.input_ids.numpy()[0]
tokens_edit = tokens_edit.input_ids.numpy()[0]
for name, a0, a1, b0, b1 in SequenceMatcher(None, tokens, tokens_edit).get_opcodes():
if b0 < tokens_length:
if name == "equal" or (name == "replace" and a1-a0 == b1-b0):
mask[b0:b1] = 1
indices[b0:b1] = indices_target[a0:a1]
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention" and "attn2" in name:
module.last_attn_slice_mask = mask.to(device)
module.last_attn_slice_indices = indices.to(device)
if module_name == "CrossAttention" and "attn1" in name:
module.last_attn_slice_mask = None
module.last_attn_slice_indices = None
def init_attention_func():
def new_attention(self, query, key, value, sequence_length, dim):
batch_size_attention = query.shape[0]
hidden_states = torch.zeros(
(batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype
)
slice_size = self._slice_size if self._slice_size is not None else hidden_states.shape[0]
for i in range(hidden_states.shape[0] // slice_size):
start_idx = i * slice_size
end_idx = (i + 1) * slice_size
attn_slice = (
torch.einsum("b i d, b j d -> b i j", query[start_idx:end_idx], key[start_idx:end_idx]) * self.scale
)
attn_slice = attn_slice.softmax(dim=-1)
if self.use_last_attn_slice:
if self.last_attn_slice_mask is not None:
new_attn_slice = torch.index_select(self.last_attn_slice, -1, self.last_attn_slice_indices)
attn_slice = attn_slice * (1 - self.last_attn_slice_mask) + new_attn_slice * self.last_attn_slice_mask
else:
attn_slice = self.last_attn_slice
self.use_last_attn_slice = False
if self.save_last_attn_slice:
self.last_attn_slice = attn_slice
self.save_last_attn_slice = False
if self.use_last_attn_weights and self.last_attn_slice_weights is not None:
attn_slice = attn_slice * self.last_attn_slice_weights
self.use_last_attn_weights = False
attn_slice = torch.einsum("b i j, b j d -> b i d", attn_slice, value[start_idx:end_idx])
hidden_states[start_idx:end_idx] = attn_slice
# reshape hidden_states
hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
return hidden_states
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention":
module.last_attn_slice = None
module.use_last_attn_slice = False
module.use_last_attn_weights = False
module.save_last_attn_slice = False
module._attention = new_attention.__get__(module, type(module))
def use_last_tokens_attention(use=True):
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention" and "attn2" in name:
module.use_last_attn_slice = use
def use_last_tokens_attention_weights(use=True):
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention" and "attn2" in name:
module.use_last_attn_weights = use
def use_last_self_attention(use=True):
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention" and "attn1" in name:
module.use_last_attn_slice = use
def save_last_tokens_attention(save=True):
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention" and "attn2" in name:
module.save_last_attn_slice = save
def save_last_self_attention(save=True):
for name, module in unet.named_modules():
module_name = type(module).__name__
if module_name == "CrossAttention" and "attn1" in name:
module.save_last_attn_slice = save
####################################
##### BASELINE ALGORITHM, ONLY USED NOW FOR SDEDIT ####3
@torch.no_grad()
def baseline_stablediffusion(prompt="",
prompt_edit=None,
null_prompt='',
prompt_edit_token_weights=[],
prompt_edit_tokens_start=0.0,
prompt_edit_tokens_end=1.0,
prompt_edit_spatial_start=0.0,
prompt_edit_spatial_end=1.0,
clip_start=0.0,
clip_end=1.0,
guidance_scale=7,
steps=50,
seed=1,
width=512, height=512,
init_image=None, init_image_strength=0.5,
fixed_starting_latent = None,
prev_image= None,
grid=None,
clip_guidance=None,
clip_guidance_scale=1,
num_cutouts=4,
cut_power=1,
scheduler_str='lms',
return_latent=False,
one_pass=False,
normalize_noise_pred=False):
width = width - width % 64
height = height - height % 64
#If seed is None, randomly select seed from 0 to 2^32-1
if seed is None: seed = random.randrange(2**32 - 1)
generator = torch.cuda.manual_seed(seed)
#Set inference timesteps to scheduler
scheduler_dict = {'ddim':DDIMScheduler,
'lms':LMSDiscreteScheduler,
'pndm':PNDMScheduler,
'ddpm':DDPMScheduler}
scheduler_call = scheduler_dict[scheduler_str]
if scheduler_str == 'ddim':
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False, set_alpha_to_one=False)
else:
scheduler = scheduler_call(beta_schedule="scaled_linear",
num_train_timesteps=1000)
scheduler.set_timesteps(steps)
if prev_image is not None:
prev_scheduler = LMSDiscreteScheduler(beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
num_train_timesteps=1000)
prev_scheduler.set_timesteps(steps)
#Preprocess image if it exists (img2img)
if init_image is not None:
init_image = init_image.resize((width, height), resample=Image.Resampling.LANCZOS)
init_image = np.array(init_image).astype(np_dtype) / 255.0 * 2.0 - 1.0
init_image = torch.from_numpy(init_image[np.newaxis, ...].transpose(0, 3, 1, 2))
#If there is alpha channel, composite alpha for white, as the diffusion model does not support alpha channel
if init_image.shape[1] > 3:
init_image = init_image[:, :3] * init_image[:, 3:] + (1 - init_image[:, 3:])
#Move image to GPU
init_image = init_image.to(device)
#Encode image
with autocast(device):
init_latent = vae.encode(init_image).latent_dist.sample(generator=generator) * 0.18215
t_start = steps - int(steps * init_image_strength)
else:
init_latent = torch.zeros((1, unet.in_channels, height // 8, width // 8), device=device)
t_start = 0
#Generate random normal noise
if fixed_starting_latent is None:
noise = torch.randn(init_latent.shape, generator=generator, device=device, dtype=unet.dtype)
if scheduler_str == 'ddim':
if init_image is not None:
raise notImplementedError
latent = scheduler.add_noise(init_latent, noise,
1000 - int(1000 * init_image_strength)).to(device)
else:
latent = noise
else:
latent = scheduler.add_noise(init_latent, noise,
t_start).to(device)
else:
latent = fixed_starting_latent
t_start = steps - int(steps * init_image_strength)
if prev_image is not None:
#Resize and prev_image for numpy b h w c -> torch b c h w
prev_image = prev_image.resize((width, height), resample=Image.Resampling.LANCZOS)
prev_image = np.array(prev_image).astype(np_dtype) / 255.0 * 2.0 - 1.0
prev_image = torch.from_numpy(prev_image[np.newaxis, ...].transpose(0, 3, 1, 2))
#If there is alpha channel, composite alpha for white, as the diffusion model does not support alpha channel
if prev_image.shape[1] > 3:
prev_image = prev_image[:, :3] * prev_image[:, 3:] + (1 - prev_image[:, 3:])
#Move image to GPU
prev_image = prev_image.to(device)
#Encode image
with autocast(device):
prev_init_latent = vae.encode(prev_image).latent_dist.sample(generator=generator) * 0.18215
t_start = steps - int(steps * init_image_strength)
prev_latent = prev_scheduler.add_noise(prev_init_latent, noise, t_start).to(device)
else:
prev_latent = None
#Process clip
with autocast(device):
tokens_unconditional = clip_tokenizer(null_prompt, padding="max_length", max_length=clip_tokenizer.model_max_length, truncation=True, return_tensors="pt", return_overflowing_tokens=True)
embedding_unconditional = clip(tokens_unconditional.input_ids.to(device)).last_hidden_state
tokens_conditional = clip_tokenizer(prompt, padding="max_length", max_length=clip_tokenizer.model_max_length, truncation=True, return_tensors="pt", return_overflowing_tokens=True)
embedding_conditional = clip(tokens_conditional.input_ids.to(device)).last_hidden_state
#Process prompt editing
assert not ((prompt_edit is not None) and (prev_image is not None))
if prompt_edit is not None:
tokens_conditional_edit = clip_tokenizer(prompt_edit, padding="max_length", max_length=clip_tokenizer.model_max_length, truncation=True, return_tensors="pt", return_overflowing_tokens=True)
embedding_conditional_edit = clip(tokens_conditional_edit.input_ids.to(device)).last_hidden_state
init_attention_edit(tokens_conditional, tokens_conditional_edit)
elif prev_image is not None:
init_attention_edit(tokens_conditional, tokens_conditional)
init_attention_func()
init_attention_weights(prompt_edit_token_weights)
timesteps = scheduler.timesteps[t_start:]
# print(timesteps)
assert isinstance(guidance_scale, int)
num_cycles = 1 # guidance_scale + 1
last_noise_preds = None
for i, t in tqdm(enumerate(timesteps), total=len(timesteps)):
t_index = t_start + i
latent_model_input = latent
if scheduler_str=='lms':
sigma = scheduler.sigmas[t_index] # last is first and first is last
latent_model_input = (latent_model_input / ((sigma**2 + 1) ** 0.5)).to(unet.dtype)
else:
assert scheduler_str in ['ddim', 'pndm', 'ddpm']
#Predict the unconditional noise residual
if len(t.shape) == 0:
t = t[None].to(unet.device)
noise_pred_uncond = unet(latent_model_input, t, encoder_hidden_states=embedding_unconditional,
).sample
if prev_latent is not None:
prev_latent_model_input = prev_latent
prev_latent_model_input = (prev_latent_model_input / ((sigma**2 + 1) ** 0.5)).to(unet.dtype)
prev_noise_pred_uncond = unet(prev_latent_model_input, t,
encoder_hidden_states=embedding_unconditional,
).sample
# noise_pred_uncond = unet(latent_model_input, t,
# encoder_hidden_states=embedding_unconditional)['sample']
#Prepare the Cross-Attention layers
if prompt_edit is not None or prev_latent is not None:
save_last_tokens_attention()
save_last_self_attention()
else:
#Use weights on non-edited prompt when edit is None
use_last_tokens_attention_weights()
#Predict the conditional noise residual and save the cross-attention layer activations
if prev_latent is not None:
raise NotImplementedError # I totally lost track of what this is
prev_noise_pred_cond = unet(prev_latent_model_input, t, encoder_hidden_states=embedding_conditional,
).sample
else:
noise_pred_cond = unet(latent_model_input, t, encoder_hidden_states=embedding_conditional,
).sample
#Edit the Cross-Attention layer activations
t_scale = t / scheduler.num_train_timesteps
if prompt_edit is not None or prev_latent is not None:
if t_scale >= prompt_edit_tokens_start and t_scale <= prompt_edit_tokens_end:
use_last_tokens_attention()
if t_scale >= prompt_edit_spatial_start and t_scale <= prompt_edit_spatial_end:
use_last_self_attention()
#Use weights on edited prompt
use_last_tokens_attention_weights()
#Predict the edited conditional noise residual using the cross-attention masks
if prompt_edit is not None:
noise_pred_cond = unet(latent_model_input, t,
encoder_hidden_states=embedding_conditional_edit).sample
#Perform guidance
# if i%(num_cycles)==0: # cycle_i+1==num_cycles:
"""
if cycle_i+1==num_cycles:
noise_pred = noise_pred_uncond
else:
noise_pred = noise_pred_cond - noise_pred_uncond
"""
if last_noise_preds is not None:
# print( (last_noise_preds[0]*noise_pred_uncond).sum(), (last_noise_preds[1]*noise_pred_cond).sum())
# print(F.cosine_similarity(last_noise_preds[0].flatten(), noise_pred_uncond.flatten(), dim=0),
# F.cosine_similarity(last_noise_preds[1].flatten(), noise_pred_cond.flatten(), dim=0))
last_grad= last_noise_preds[1] - last_noise_preds[0]
new_grad = noise_pred_cond - noise_pred_uncond
# print( F.cosine_similarity(last_grad.flatten(), new_grad.flatten(), dim=0))
last_noise_preds = (noise_pred_uncond, noise_pred_cond)
use_cond_guidance = True
if use_cond_guidance:
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
else:
noise_pred = noise_pred_uncond
if clip_guidance is not None and t_scale >= clip_start and t_scale <= clip_end:
noise_pred, latent = new_cond_fn(latent, t, t_index,
embedding_conditional, noise_pred,clip_guidance,
clip_guidance_scale,
num_cutouts,
scheduler, unet,use_cutouts=True,
cut_power=cut_power)
if normalize_noise_pred:
noise_pred = noise_pred * noise_pred_uncond.norm() / noise_pred.norm()
if scheduler_str == 'ddim':
latent = forward_step(scheduler, noise_pred,
t,
latent).prev_sample
else:
latent = scheduler.step(noise_pred,
t_index,
latent).prev_sample
if prev_latent is not None:
prev_noise_pred = prev_noise_pred_uncond + guidance_scale * (prev_noise_pred_cond - prev_noise_pred_uncond)
prev_latent = prev_scheduler.step(prev_noise_pred, t_index, prev_latent).prev_sample
if one_pass: break
#scale and decode the image latents with vae
if return_latent: return latent
latent = latent / 0.18215
image = vae.decode(latent.to(vae.dtype)).sample
image = (image / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()
image = (image[0] * 255).round().astype("uint8")
return Image.fromarray(image)
####################################
#### HELPER FUNCTIONS FOR OUR METHOD #####
def get_alpha_and_beta(t, scheduler):
# want to run this for both current and previous timnestep
if t.dtype==torch.long:
alpha = scheduler.alphas_cumprod[t]
return alpha, 1-alpha
if t<0:
return scheduler.final_alpha_cumprod, 1 - scheduler.final_alpha_cumprod
low = t.floor().long()
high = t.ceil().long()
rem = t - low
low_alpha = scheduler.alphas_cumprod[low]
high_alpha = scheduler.alphas_cumprod[high]
interpolated_alpha = low_alpha * rem + high_alpha * (1-rem)
interpolated_beta = 1 - interpolated_alpha
return interpolated_alpha, interpolated_beta
# A DDIM forward step function
def forward_step(
self,
model_output,
timestep: int,
sample,
eta: float = 0.0,
use_clipped_model_output: bool = False,
generator=None,
return_dict: bool = True,
use_double=False,
) :
if self.num_inference_steps is None:
raise ValueError(
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
)
prev_timestep = timestep - self.config.num_train_timesteps / self.num_inference_steps
if timestep > self.timesteps.max():
raise NotImplementedError("Need to double check what the overflow is")
alpha_prod_t, beta_prod_t = get_alpha_and_beta(timestep, self)
alpha_prod_t_prev, _ = get_alpha_and_beta(prev_timestep, self)
alpha_quotient = ((alpha_prod_t / alpha_prod_t_prev)**0.5)
first_term = (1./alpha_quotient) * sample
second_term = (1./alpha_quotient) * (beta_prod_t ** 0.5) * model_output
third_term = ((1 - alpha_prod_t_prev)**0.5) * model_output
return first_term - second_term + third_term
# A DDIM reverse step function, the inverse of above
def reverse_step(
self,
model_output,
timestep: int,
sample,
eta: float = 0.0,
use_clipped_model_output: bool = False,
generator=None,
return_dict: bool = True,
use_double=False,
) :
if self.num_inference_steps is None:
raise ValueError(
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
)
prev_timestep = timestep - self.config.num_train_timesteps / self.num_inference_steps
if timestep > self.timesteps.max():
raise NotImplementedError
else:
alpha_prod_t = self.alphas_cumprod[timestep]
alpha_prod_t, beta_prod_t = get_alpha_and_beta(timestep, self)
alpha_prod_t_prev, _ = get_alpha_and_beta(prev_timestep, self)
alpha_quotient = ((alpha_prod_t / alpha_prod_t_prev)**0.5)
first_term = alpha_quotient * sample
second_term = ((beta_prod_t)**0.5) * model_output
third_term = alpha_quotient * ((1 - alpha_prod_t_prev)**0.5) * model_output
return first_term + second_term - third_term
@torch.no_grad()
def latent_to_image(latent):
image = vae.decode(latent.to(vae.dtype)/0.18215).sample
image = prep_image_for_return(image)
return image
def prep_image_for_return(image):
image = (image / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()
image = (image[0] * 255).round().astype("uint8")
image = Image.fromarray(image)
return image
#############################
##### MAIN EDICT FUNCTION #######
# Use EDICT_editing to perform calls
@torch.no_grad()
def coupled_stablediffusion(prompt="",
prompt_edit=None,
null_prompt='',
prompt_edit_token_weights=[],
prompt_edit_tokens_start=0.0,
prompt_edit_tokens_end=1.0,
prompt_edit_spatial_start=0.0,
prompt_edit_spatial_end=1.0,
guidance_scale=7.0, steps=50,
seed=1, width=512, height=512,
init_image=None, init_image_strength=1.0,
run_baseline=False,
use_lms=False,
leapfrog_steps=True,
reverse=False,
return_latents=False,
fixed_starting_latent=None,
beta_schedule='scaled_linear',
mix_weight=0.93):
#If seed is None, randomly select seed from 0 to 2^32-1
if seed is None: seed = random.randrange(2**32 - 1)
generator = torch.cuda.manual_seed(seed)
def image_to_latent(im):
if isinstance(im, torch.Tensor):
# assume it's the latent
# used to avoid clipping new generation before inversion
init_latent = im.to(device)
else:
#Resize and transpose for numpy b h w c -> torch b c h w
im = im.resize((width, height), resample=Image.Resampling.LANCZOS)
im = np.array(im).astype(np_dtype) / 255.0 * 2.0 - 1.0
# check if black and white
if len(im.shape) < 3:
im = np.stack([im for _ in range(3)], axis=2) # putting at end b/c channels
im = torch.from_numpy(im[np.newaxis, ...].transpose(0, 3, 1, 2))
#If there is alpha channel, composite alpha for white, as the diffusion model does not support alpha channel
if im.shape[1] > 3:
im = im[:, :3] * im[:, 3:] + (1 - im[:, 3:])
#Move image to GPU
im = im.to(device)
#Encode image
if use_half_prec:
init_latent = vae.encode(im).latent_dist.sample(generator=generator) * 0.18215
else:
with autocast(device):
init_latent = vae.encode(im).latent_dist.sample(generator=generator) * 0.18215
return init_latent
assert not use_lms, "Can't invert LMS the same as DDIM"
if run_baseline: leapfrog_steps=False
#Change size to multiple of 64 to prevent size mismatches inside model
width = width - width % 64
height = height - height % 64
#Preprocess image if it exists (img2img)
if init_image is not None:
assert reverse # want to be performing deterministic noising
# can take either pair (output of generative process) or single image
if isinstance(init_image, list):
if isinstance(init_image[0], torch.Tensor):
init_latent = [t.clone() for t in init_image]
else:
init_latent = [image_to_latent(im) for im in init_image]
else:
init_latent = image_to_latent(init_image)
# this is t_start for forward, t_end for reverse
t_limit = steps - int(steps * init_image_strength)
else:
assert not reverse, 'Need image to reverse from'
init_latent = torch.zeros((1, unet.in_channels, height // 8, width // 8), device=device)
t_limit = 0
if reverse:
latent = init_latent
else:
#Generate random normal noise
noise = torch.randn(init_latent.shape,
generator=generator,
device=device,
dtype=torch_dtype)
if fixed_starting_latent is None:
latent = noise
else:
if isinstance(fixed_starting_latent, list):
latent = [l.clone() for l in fixed_starting_latent]
else:
latent = fixed_starting_latent.clone()
t_limit = steps - int(steps * init_image_strength)
if isinstance(latent, list): # initializing from pair of images
latent_pair = latent
else: # initializing from noise
latent_pair = [latent.clone(), latent.clone()]
if steps==0:
if init_image is not None:
return image_to_latent(init_image)
else:
image = vae.decode(latent.to(vae.dtype) / 0.18215).sample
return prep_image_for_return(image)
#Set inference timesteps to scheduler
schedulers = []
for i in range(2):
# num_raw_timesteps = max(1000, steps)
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012,
beta_schedule=beta_schedule,
num_train_timesteps=1000,
clip_sample=False,
set_alpha_to_one=False)
scheduler.set_timesteps(steps)
schedulers.append(scheduler)
with autocast(device):
# CLIP Text Embeddings
tokens_unconditional = clip_tokenizer(null_prompt, padding="max_length",
max_length=clip_tokenizer.model_max_length,
truncation=True, return_tensors="pt",
return_overflowing_tokens=True)
embedding_unconditional = clip(tokens_unconditional.input_ids.to(device)).last_hidden_state
tokens_conditional = clip_tokenizer(prompt, padding="max_length",
max_length=clip_tokenizer.model_max_length,
truncation=True, return_tensors="pt",
return_overflowing_tokens=True)
embedding_conditional = clip(tokens_conditional.input_ids.to(device)).last_hidden_state
#Process prompt editing (if running Prompt-to-Prompt)
if prompt_edit is not None:
tokens_conditional_edit = clip_tokenizer(prompt_edit, padding="max_length",
max_length=clip_tokenizer.model_max_length,
truncation=True, return_tensors="pt",
return_overflowing_tokens=True)
embedding_conditional_edit = clip(tokens_conditional_edit.input_ids.to(device)).last_hidden_state
init_attention_edit(tokens_conditional, tokens_conditional_edit)
init_attention_func()
init_attention_weights(prompt_edit_token_weights)
timesteps = schedulers[0].timesteps[t_limit:]
if reverse: timesteps = timesteps.flip(0)
for i, t in tqdm(enumerate(timesteps), total=len(timesteps)):
t_scale = t / schedulers[0].num_train_timesteps
if (reverse) and (not run_baseline):
# Reverse mixing layer
new_latents = [l.clone() for l in latent_pair]
new_latents[1] = (new_latents[1].clone() - (1-mix_weight)*new_latents[0].clone()) / mix_weight
new_latents[0] = (new_latents[0].clone() - (1-mix_weight)*new_latents[1].clone()) / mix_weight
latent_pair = new_latents
# alternate EDICT steps
for latent_i in range(2):
if run_baseline and latent_i==1: continue # just have one sequence for baseline
# this modifies latent_pair[i] while using
# latent_pair[(i+1)%2]
if reverse and (not run_baseline):
if leapfrog_steps:
# what i would be from going other way
orig_i = len(timesteps) - (i+1)
offset = (orig_i+1) % 2
latent_i = (latent_i + offset) % 2
else:
# Do 1 then 0
latent_i = (latent_i+1)%2
else:
if leapfrog_steps:
offset = i%2
latent_i = (latent_i + offset) % 2
latent_j = ((latent_i+1) % 2) if not run_baseline else latent_i
latent_model_input = latent_pair[latent_j]
latent_base = latent_pair[latent_i]
#Predict the unconditional noise residual
noise_pred_uncond = unet(latent_model_input, t,
encoder_hidden_states=embedding_unconditional).sample
#Prepare the Cross-Attention layers
if prompt_edit is not None:
save_last_tokens_attention()
save_last_self_attention()
else:
#Use weights on non-edited prompt when edit is None
use_last_tokens_attention_weights()
#Predict the conditional noise residual and save the cross-attention layer activations
noise_pred_cond = unet(latent_model_input, t,
encoder_hidden_states=embedding_conditional).sample
#Edit the Cross-Attention layer activations
if prompt_edit is not None:
t_scale = t / schedulers[0].num_train_timesteps
if t_scale >= prompt_edit_tokens_start and t_scale <= prompt_edit_tokens_end:
use_last_tokens_attention()
if t_scale >= prompt_edit_spatial_start and t_scale <= prompt_edit_spatial_end:
use_last_self_attention()
#Use weights on edited prompt
use_last_tokens_attention_weights()
#Predict the edited conditional noise residual using the cross-attention masks
noise_pred_cond = unet(latent_model_input,
t,
encoder_hidden_states=embedding_conditional_edit).sample
#Perform guidance
grad = (noise_pred_cond - noise_pred_uncond)
noise_pred = noise_pred_uncond + guidance_scale * grad
step_call = reverse_step if reverse else forward_step
new_latent = step_call(schedulers[latent_i],
noise_pred,
t,
latent_base)# .prev_sample
new_latent = new_latent.to(latent_base.dtype)
latent_pair[latent_i] = new_latent
if (not reverse) and (not run_baseline):
# Mixing layer (contraction) during generative process
new_latents = [l.clone() for l in latent_pair]
new_latents[0] = (mix_weight*new_latents[0] + (1-mix_weight)*new_latents[1]).clone()
new_latents[1] = ((1-mix_weight)*new_latents[0] + (mix_weight)*new_latents[1]).clone()
latent_pair = new_latents
#scale and decode the image latents with vae, can return latents instead of images
if reverse or return_latents:
results = [latent_pair]
return results if len(results)>1 else results[0]
# decode latents to iamges
images = []
for latent_i in range(2):
latent = latent_pair[latent_i] / 0.18215
image = vae.decode(latent.to(vae.dtype)).sample
images.append(image)
# Return images
return_arr = []
for image in images:
image = prep_image_for_return(image)
return_arr.append(image)
results = [return_arr]
return results if len(results)>1 else results[0]