Spaces:
Sleeping
Sleeping
File size: 9,124 Bytes
2b7b059 89321ee 2b7b059 |
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 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 |
from base64 import b64encode
from utils import *
import numpy
import torch
from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel
from huggingface_hub import notebook_login
# For video display:
from IPython.display import HTML
from matplotlib import pyplot as plt
from pathlib import Path
from PIL import Image
from torch import autocast
from torchvision import transforms as tfms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer, logging
import os
import shutil
from device import torch_device,vae,text_encoder,unet,tokenizer,scheduler,token_emb_layer,pos_emb_layer,position_embeddings
torch.manual_seed(1)
if not (Path.home()/'.cache/huggingface'/'token').exists(): notebook_login()
# Supress some unnecessary warnings when loading the CLIPTextModel
logging.set_verbosity_error()
# Set device
def generate_distorted_image(pil_image,vae):
# View a noised version
encoded = pil_to_latent(pil_image)
noise = torch.randn_like(encoded) # Random noise
sampling_step = 5 # Equivalent to step 10 out of 15 in the schedule above
# encoded_and_noised = scheduler.add_noise(encoded, noise, timestep) # Diffusers 0.3 and below
encoded_and_noised = scheduler.add_noise(encoded, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]]))
return latents_to_pil(encoded_and_noised)[0] # Display
def set_timesteps(scheduler, num_inference_steps):
scheduler.set_timesteps(num_inference_steps)
scheduler.timesteps = scheduler.timesteps.to(torch.float32)
# Some settings
def generate_image(prompt,concept_embed,num_inference_steps=50,color_postprocessing=False,noised_image=False,loss_scale=10,seed=42):
height = 512 # default height of Stable Diffusion
width = 512 # default width of Stable Diffusion
num_inference_steps = num_inference_steps # Number of denoising steps
guidance_scale = 7.5 # Scale for classifier-free guidance
generator = torch.manual_seed(seed) # Seed generator to create the inital latent noise
batch_size = 1
# Define the directory name
directory_name = "steps"
# Check if the directory exists, and if so, delete it
if os.path.exists(directory_name):
shutil.rmtree(directory_name)
#Create the directory
os.makedirs(directory_name)
# Prep text
#text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
# with torch.no_grad():
# text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
input_ids = text_input.input_ids.to(torch_device)
custom_style_token=tokenizer.encode("cs",add_special_token=False)[0]
# Get token embeddings
token_embeddings = token_emb_layer(input_ids)
embed_key=list(concept_embed.keys())[0]
# The new embedding. In this case just the input embedding of token 2368...
replacement_token_embedding = concept_embed[embed_key]
token_embeddings[0,torch.where(input_ids[0]==custom_style_token)]=replacement_token_embedding.to(torch_device)
# Combine with pos embs
input_embeddings = token_embeddings + position_embeddings
# Feed through to get final output embs
modified_output_embeddings = get_output_embeds(input_embeddings)
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
[""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
with torch.no_grad():
uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
text_embeddings = torch.cat([uncond_embeddings, modified_output_embeddings])
# minor fix to ensure MPS compatibility, fixed in diffusers PR 3925
set_timesteps(scheduler,num_inference_steps)
# Prep latents
latents = torch.randn(
(batch_size, unet.in_channels, height // 8, width // 8),
generator=generator,
)
latents = latents.to(torch_device)
latents = latents * scheduler.init_noise_sigma # Scaling (previous versions did latents = latents * self.scheduler.sigmas[0]
# Loop
with autocast("cpu"): # will fallback to CPU if no CUDA; no autocast for MPS
for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):
# expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
latent_model_input = torch.cat([latents] * 2)
sigma = scheduler.sigmas[i]
# Scale the latents (preconditioning):
# latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) # Diffusers 0.3 and below
latent_model_input = scheduler.scale_model_input(latent_model_input, t)
# predict the noise residual
with torch.no_grad():
noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
# perform guidance
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
# latents = scheduler.step(noise_pred, i, latents)["prev_sample"] # Diffusers 0.3 and below
#latents = torch.tensor(initial_latents, requires_grad=True)
### ADDITIONAL GUIDANCE ###
# Requires grad on the latents
if color_postprocessing:
latents = latents.detach().requires_grad_()
# Get the predicted x0:
latents_x0 = latents - sigma * noise_pred
# Decode to image space
denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
#denoised_images = vae.decode((1 / 0.18215) * latents_x0) / 2 + 0.5 # (0, 1)
# Calculate loss
loss = orange_loss(denoised_images) * loss_scale
#loss = color_loss(denoised_images,postporcessing_color) * color_loss_scale
if i%10==0:
print(i, 'loss:', loss.item())
# Get gradient
cond_grad = -torch.autograd.grad(loss, latents)[0]
# Modify the latents based on this gradient
latents = latents.detach() + cond_grad * sigma**2
### And saving as before ###
# Get the predicted x0:
latents_x0 = latents - sigma * noise_pred
im_t0 = latents_to_pil(latents_x0)[0]
# And the previous noisy sample x_t -> x_t-1
latents = scheduler.step(noise_pred, t, latents)["prev_sample"]
im_next = latents_to_pil(latents)[0]
# Combine the two images and save for later viewing
im = Image.new('RGB', (1024, 512))
im.paste(im_next, (0, 0))
im.paste(im_t0, (512, 0))
im.save(f'steps/{i:04}.jpeg')
else:
latents = scheduler.step(noise_pred, t, latents).prev_sample
if noised_image:
output = generate_distorted_image(latents_to_pil(latents)[0],vae)
else:
output = latents_to_pil(latents)[0]
return output
def get_output_embeds(input_embeddings):
# CLIP's text model uses causal mask, so we prepare it here:
bsz, seq_len = input_embeddings.shape[:2]
causal_attention_mask = text_encoder.text_model._build_causal_attention_mask(bsz, seq_len, dtype=input_embeddings.dtype)
# Getting the output embeddings involves calling the model with passing output_hidden_states=True
# so that it doesn't just return the pooled final predictions:
encoder_outputs = text_encoder.text_model.encoder(
inputs_embeds=input_embeddings,
attention_mask=None, # We aren't using an attention mask so that can be None
causal_attention_mask=causal_attention_mask.to(torch_device),
output_attentions=None,
output_hidden_states=True, # We want the output embs not the final output
return_dict=None,
)
# We're interested in the output hidden state only
output = encoder_outputs[0]
# There is a final layer norm we need to pass these through
output = text_encoder.text_model.final_layer_norm(output)
# And now they're ready!
return output
|