import logging from pathlib import Path import matplotlib.pyplot as plt import torch from diffusers import StableDiffusionPipeline from fastcore.all import concat from huggingface_hub import notebook_login from PIL import Image import numpy as np # from IPython.display import display from torchvision import transforms as tfms from transformers import CLIPTextModel, CLIPTokenizer from diffusers import AutoencoderKL, UNet2DConditionModel from diffusers import LMSDiscreteScheduler from tqdm.auto import tqdm logging.disable(logging.WARNING) class ImageGenerator(): def __init__(self, g:int=7.5, ): self.latent_images = [] self.g = g self.width = 512 self.height = 512 self.generator = torch.manual_seed(32) self.bs = 1 if torch.cuda.is_available(): self.device = torch.device("cuda") self.float_size = torch.float16 else: self.device = torch.device("cpu") self.float_size = torch.float32 def __repr__(self): return f"Image Generator with {self.g=}" def load_models(self): self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=self.float_size) self.text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=self.float_size).to( self.device) # vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema", torch_dtype=torch.float16 ).to(self.device) self.vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").to( self.device) self.unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet" ).to( self.device) #torch_dtype=torch.float16, def load_scheduler( self, beta_start : float=0.00085, beta_end : float=0.012, beta_schedule : str="scaled_linear", num_train_timesteps :int=1000): self.scheduler = LMSDiscreteScheduler( beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear", num_train_timesteps=num_train_timesteps) def load_image(self, filepath:str): return Image.open(filepath).resize(size=(self.width,self.height)) #.convert("RGB") # RGB = 3 dimensions, RGBA = 4 dimensions def pil_to_latent(self, image: Image) -> torch.Tensor: with torch.no_grad(): np_img = np.transpose( (( np.array(image) / 255)-0.5)*2, (2,0,1)) # turn pil image into np array with values between -1 and 1 # print(f"{np_img.shape=}") # 4, 64, 64 np_images = np.repeat(np_img[np.newaxis, :, :], self.bs, axis=0) # adding a new dimension and repeating the image for each prompt # print(f"{np_images.shape=}") decoded_latent = torch.from_numpy(np_images).to(self.device).float() #<-- stability-ai vae uses half(), compvis vae uses float? # print(f"{decoded_latent.shape=}") encoded_latent = 0.18215 * self.vae.encode(decoded_latent).latent_dist.sample() # print(f"{encoded_latent.shape=}") return encoded_latent def add_noise(self, latent: torch.Tensor, scheduler_steps: int = 10) -> torch.FloatTensor: # noise = torch.randn_like(latent) # missing generator parameter noise = torch.randn( size = (self.bs, self.unet.config.in_channels, self.height//8, self.width//8), generator = self.generator).to(self.device) timesteps = torch.tensor([self.scheduler.timesteps[scheduler_steps]]) noisy_latent = self.scheduler.add_noise(latent, noise, timesteps) # print(f"add_noise: {timesteps.shape=} {timesteps=} {noisy_latent.shape=}") return noisy_latent def latent_to_pil(self, latent:torch.Tensor) -> Image: # print(f"latent_to_pil {latent.dtype=}") with torch.no_grad(): decoded = self.vae.decode(1 / 0.18215 * latent).sample[0] # print(f"latent_to_pil {decoded.shape=}") image = (decoded/2+0.5).clamp(0,1).detach().cpu().permute(1, 2, 0).numpy() return Image.fromarray((image*255).round().astype("uint8")) def image_grid(self, imgs: [Image]) -> Image: w,h = imgs[0].size cols = len(imgs) grid = Image.new('RGB', size=(cols*w, h)) for i, img in enumerate(imgs): # print(f"{img.size=}") grid.paste(img, box=(i%cols*w, i//cols*h)) return grid def text_enc(self, prompt:str, maxlen=None) -> torch.Tensor: '''tokenize and encode a prompt''' if maxlen is None: maxlen = self.tokenizer.model_max_length inp = self.tokenizer([prompt], padding="max_length", max_length=maxlen, truncation=True, return_tensors="pt") return self.text_encoder(inp.input_ids.to(self.device))[0].float() def tensor_to_pil(self, t:torch.Tensor) -> Image: '''transforms a tensor decoded by the vae to a pil image''' # print(f"tensor_to_pil {t.shape=} {type(t)=}") image = (t/2+0.5).clamp(0,1).detach().cpu().permute(1, 2, 0).numpy() return Image.fromarray((image*255).round().astype("uint8")) def latent_callback(self, latent:torch.Tensor) -> None: '''store latents in an array so that we can inpect them later.''' with torch.no_grad(): # print(f"cb {latent.shape=}") decoded = self.vae.decode(1 / 0.18215 * latent).sample[0] self.latent_images.append(self.tensor_to_pil(decoded)) def generate(self, prompt : str, secondary_prompt: str=None, prompt_mix_ratio : float=0.5, negative_prompt="", seed : int=32, steps : int=30, start_step_ratio : float=1/5, init_image : Image=None, latent_callback_mod : int=10): self.latent_images = [] if not negative_prompt: negative_prompt = "" with torch.no_grad(): text = self.text_enc(prompt) if secondary_prompt: sec_prompt_text = self.text_enc(secondary_prompt) text = text * prompt_mix_ratio + sec_prompt_text * ( 1 - prompt_mix_ratio ) uncond = self.text_enc(negative_prompt * self.bs, text.shape[1]) emb = torch.cat([uncond, text]) if seed: torch.manual_seed(seed) self.scheduler.set_timesteps(steps) self.scheduler.timesteps = self.scheduler.timesteps.to(torch.float32) if (init_image == None): start_steps = 0 latents = torch.randn( size = (self.bs, self.unet.config.in_channels, self.height//8, self.width//8), generator = self.generator) latents = latents * self.scheduler.init_noise_sigma # print(f"{latents.shape=}") else: start_steps = int(steps * start_step_ratio) # 0%: too much noise, 100% no noise # print(f"{start_steps=}") latents =self. pil_to_latent(init_image) self.latent_callback(latents) latents = self.add_noise(latents, start_steps).to(self.device).float() self.latent_callback(latents) latents = latents.to(self.device).float() for i,ts in enumerate(tqdm(self.scheduler.timesteps, leave=False)): if i >= start_steps: inp = self.scheduler.scale_model_input(torch.cat([latents] * 2), ts) with torch.no_grad(): u,t = self.unet(inp, ts, encoder_hidden_states=emb).sample.chunk(2) #todo, grab those with callbacks pred = u + self.g*(t-u) # pred = u + self.g*(t-u)/torch.norm(t-u)*torch.norm(u) latents = self.scheduler.step(pred, ts, latents).prev_sample if latent_callback_mod and i % latent_callback_mod == 0: self.latent_callback(latents) return self.latent_to_pil(latents), self.latent_images