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): self.latent_images = [] 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 elif torch.backends.mps.is_available(): self.device = torch.device("mps") self.float_size = torch.float32 else: if not torch.backends.mps.is_built(): print("MPS not available because the current PyTorch install was not " "built with MPS enabled.") else: print("MPS not available because the current MacOS version is not 12.3+ " "and/or you do not have an MPS-enabled device on this machine.") self.device = torch.device("cpu") self.float_size = torch.float32 print(f"pytorch device: {self.device}") def __repr__(self): return f"Image Generator with {self.width=} {self.height=}" 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, 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(): image = image.resize(size=(self.width,self.height)) 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).astype(np.float32) # adding a new dimension and repeating the image for each prompt, float32 required for mac # 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: print(len(imgs)) 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, guidance :float=7.5, steps : int=30, start_step_ratio : float=1/5, init_image : Image=None, latent_callback_mod : int=10, progress_tqdm: callable=tqdm): self.latent_images = [] if not negative_prompt: negative_prompt = "" print(f"ImageGenerator: {prompt=} {secondary_prompt=} {prompt_mix_ratio=} {negative_prompt=} {guidance=} {steps=} {init_image=} ") with torch.no_grad(): text = self.text_enc(prompt) if secondary_prompt: print("using 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: print("using base image") 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(progress_tqdm(self.scheduler.timesteps, desc="Latent Generation")): #leave=False, does not work with gradio 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 + guidance*(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