# -*- coding: utf-8 -*- """ Created on Sun Apr 7 05:35:36 2024 @author: aditya.vipradas """ # import file system libraries import os from pathlib import Path #from huggingface_hub import notebook_login from tqdm.auto import tqdm import gradio as gr # import image and visualization libraries from PIL import Image # import modeling libraries import torch from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel from transformers import CLIPTextModel, CLIPTokenizer from torchvision import transforms as tfms # check huggingface token #if not (Path.home()/'.cache/huggingface/token').exists(): # notebook_login() # set torch device and suppress duplication warnings #torch_device = "cuda" if torch.cuda.is_available() else "cpu" torch_device = "cpu" #os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = "1" # load the autoencoder vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae") # load the tokenizer and text encoder tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") # load the unet unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet") # noise scheduler (linear multi-step) scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) def latents_to_pil(latents): # remove the scaling as mentioned in the documentation latents = (1 / 0.18215) * latents with torch.no_grad(): # this generates images with (1, 3, 512, 512) dimensions. # permute them to (1, 512, 512, 3) later image = vae.decode(latents).sample image = (image / 2 + 0.5).clamp(0, 1) image = image.detach().cpu().permute(0, 2, 3, 1).numpy() images = (image * 255).round().astype("uint8") pil_images = [Image.fromarray(image) for image in images] return pil_images[0] def pil_to_latent(input_im): with torch.no_grad(): latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0)*2-1) return 0.18215 * latent.latent_dist.sample() # fixed parameters height = 512 width = 512 batch_size = 1 uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pt") with torch.no_grad(): uncond_embeddings = text_encoder(uncond_input.input_ids)[0] def diffusion(image_conditioned, prompt_image, prompt, artist, sampling_step, guidance_scale, num_inference_steps, seed): generator = torch.manual_seed(seed) scheduler.set_timesteps(num_inference_steps+1) if artist != "": prompt = prompt + ", " + artist + " style" 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)[0] text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) noise = torch.randn((batch_size, unet.config.in_channels, height // 8, width // 8), generator=generator) if image_conditioned: encoded_sketch = pil_to_latent(Image.fromarray(prompt_image). resize((height, width))) encoded_noised = scheduler.add_noise(encoded_sketch, noise, timesteps=torch.tensor([scheduler.timesteps[sampling_step]])) else: encoded_noised = noise * scheduler.init_noise_sigma for i, t in enumerate(scheduler.timesteps): if i >= sampling_step: latent_model_input = torch.cat([encoded_noised] * 2) sigma = scheduler.sigmas[i] latent_model_input = scheduler.scale_model_input(latent_model_input, t) with torch.no_grad(): noise_pred = unet(latent_model_input, t, encoder_hidden_states = text_embeddings).sample noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) encoded_noised = scheduler.step(noise_pred, t, encoded_noised).prev_sample yield latents_to_pil(encoded_noised), f"Generating your image...step {i}/{num_inference_steps}" if i == num_inference_steps: yield latents_to_pil(encoded_noised), "Image generation complete" demo = gr.Interface( theme=gr.themes.Soft(), title="Image- and Text- Conditioned Stable Diffusion", description="Latent stable diffusion implementation with image and text conditioning. Additional artistic rendering options are also provided.\n \ Potential Applications: \n1. Implement style of one image on another \n2. Convert concepts (drawing, paintings) to realistic images \n3. Guide text prompts with images.", fn=diffusion, inputs=[gr.Checkbox(value=False, label="Condition on reference image", info="Recommendation: Select sampling step 7 if selected else 0"), gr.Image(height=height, width=width, label='Select Reference Image'), gr.Textbox(value="A highly realistic and majestic lion with wavy mane, high-definition", info="Enter the text prompt", max_lines=1, label='Prompt'), gr.Radio(["Van Gogh", "Johannes Vermeer", "Claude Monet", "Pablo Picasso", "Frida Kahlo", "None"], value="Van Gogh", label="Render", info="Choose your artist"), gr.Slider(0, 20, value=0, step=1, label='Sampling step', info="Higher step adds less noise (between 0 and 20)"), gr.Slider(0, 14, value=7.5, step=0.5, label='Guidance scale', info="Adherance to prompt (between 0 and 14)"), gr.Slider(0, 100, value=70, step=1, label='Number of inference steps', info="Choose between 0 and 100"), gr.Slider(0, 100, value=50, step=1, label='Random Seed', info="Change to generate a different image (between 0 and 100)")], outputs=[gr.Image(height=height, width=width, label='Stable Diffusion Progress'), gr.Textbox(max_lines=1, label='Progress')]) demo.launch()