import os import torch import gradio as gr from tqdm import tqdm from PIL import Image from torchvision import transforms as tfms from transformers import CLIPTextModel, CLIPTokenizer, logging from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel torch_device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" if "mps" == torch_device: os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = "1" # Load the autoencoder model which will be used to decode the latents into image space. vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae") # Load the tokenizer and text encoder to tokenize and encode the text. tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") # The UNet model for generating the latents. unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet") # The noise scheduler scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) style_token_dict = {'Concept':'', 'Realistic':'', 'Line':'', 'Ricky':'', 'Plane Scape':''} # To the GPU we go! vae = vae.to(torch_device) text_encoder = text_encoder.to(torch_device) unet = unet.to(torch_device) token_emb_layer = text_encoder.text_model.embeddings.token_embedding pos_emb_layer = text_encoder.text_model.embeddings.position_embedding position_ids = text_encoder.text_model.embeddings.position_ids[:, :77] position_embeddings = pos_emb_layer(position_ids) concept_art_embed = torch.load('concept-art.bin') doose_s_realistic_art_style_embed = torch.load('doose-s-realistic-art-style.bin') line_art_embed = torch.load('line-art.bin') rickyart_embed = torch.load('rickyart.bin') tony_diterlizzi_s_planescape_art_embed = torch.load('tony-diterlizzi-s-planescape-art.bin') tokenizer.add_tokens(['', '', '', '', '']) token_emb_layer_with_art = torch.nn.Embedding(49413, 768) token_emb_layer_with_art.load_state_dict({'weight': torch.cat((token_emb_layer.state_dict()['weight'], concept_art_embed[''].unsqueeze(0).to(torch_device), doose_s_realistic_art_style_embed[''].unsqueeze(0).to(torch_device), line_art_embed[''].unsqueeze(0).to(torch_device), rickyart_embed[''].unsqueeze(0).to(torch_device), tony_diterlizzi_s_planescape_art_embed[''].unsqueeze(0).to(torch_device)))}) token_emb_layer_with_art = token_emb_layer_with_art.to(torch_device) def set_timesteps(scheduler, num_inference_steps): scheduler.set_timesteps(num_inference_steps) scheduler.timesteps = scheduler.timesteps.to(torch.float32) def pil_to_latent(input_im): with torch.no_grad(): latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling return 0.18215 * latent.latent_dist.sample() def latents_to_pil(latents): latents = (1 / 0.18215) * latents with torch.no_grad(): 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 def build_causal_attention_mask(bsz, seq_len, dtype): mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype) mask.fill_(torch.tensor(torch.finfo(dtype).min)) # fill with large negative number (acts like -inf) mask = mask.triu_(1) # zero out the lower diagonal to enforce causality return mask.unsqueeze(1) # add a batch dimension 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 = 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 def generate_with_embs(num_inference_steps, guidance_scale, seed, text_input, text_embeddings): height = 512 # default height of Stable Diffusion width = 512 # default width of Stable Diffusion generator = torch.manual_seed(seed) # Seed generator to create the inital latent noise batch_size = 1 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, text_embeddings]) # Prep Scheduler 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 # Loop 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] 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, t, latents).prev_sample return latents_to_pil(latents)[0] def inference(text, style, inference_step, guidance_scale, seed): prompt = text + " the style of " + style_token_dict[style] # Tokenize 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) # Get token embeddings token_embeddings = token_emb_layer_with_art(input_ids) # 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) # And generate an image with this: image = generate_with_embs(inference_step, guidance_scale, seed, text_input, modified_output_embeddings) return image title = "Stable Diffusion with Textual Inversion" description = "A simple Gradio interface to infer Stable Diffusion and generate images with different art style" examples = [["A sweet potato farm", 'Concept', 30, 0.5, 1], ["Sky full of cotton candy", 'Realistic', 30, 1.5, 2], ["Coffin full of jello", 'Line', 30, 2.5, 3], ["Water skiing on a lake", 'Ricky', 30, 3.5, 4], ["Super slippery noodles", 'Plane Scape', 30, 4.5, 5], ["Beautiful sunset", 'Concept', 30, 5.5, 6], ["A glittering gem", 'Realistic', 30, 6.5, 7], ["River rafting", 'Line', 30, 7.5, 8], ["A green tea", 'Ricky', 30, 8.5, 9], ["Three sphered rocks", 'Plane Scape', 30, 9.5, 10]] demo = gr.Interface(inference, inputs = [gr.Textbox(label="Prompt", type="text"), gr.Dropdown(label="Style", choices=['Concept', 'Realistic', 'Line', 'Ricky', 'Plane Scape'], value="Concept"), gr.Slider(10, 50, 30, step = 10, label="Inference steps"), gr.Slider(1, 10, 7.5, step = 0.1, label="Guidance scale"), gr.Slider(0, 10000, 1, step = 1, label="Seed")], outputs= [gr.Image(width=320, height=320, label="Output SAM")], title=title, description=description, examples=examples) demo.launch()