PushkarA07 commited on
Commit
93e417b
·
1 Parent(s): 9a10134

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +145 -0
app.py ADDED
@@ -0,0 +1,145 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
3
+ from diffusers import LDMTextToImagePipeline
4
+ import gradio as gr
5
+ import PIL.Image
6
+ import numpy as np
7
+ import random
8
+ import torch
9
+ import subprocess
10
+ from transformers import AutoModelWithLMHead, AutoModelForCausalLM, AutoTokenizer
11
+ from transformers import WhisperForConditionalGeneration, WhisperConfig, WhisperProcessor
12
+ import torchaudio
13
+ import nltk
14
+ from pydub import AudioSegment
15
+ import re
16
+ from datasets import load_dataset
17
+ from transformers import AutoModelWithLMHead, AutoTokenizer, set_seed, pipeline
18
+ import torch
19
+ from transformers import GPT2Tokenizer, GPT2LMHeadModel
20
+ import torch
21
+ from diffusers import StableDiffusionPipeline, AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DPMSolverMultistepScheduler, LMSDiscreteScheduler
22
+ from transformers import CLIPTextModel, CLIPTokenizer
23
+ from tqdm.auto import tqdm
24
+ from torch import autocast
25
+ from PIL import Image
26
+ torch_device = "cuda" if torch.cuda.is_available() else "cpu"
27
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
28
+
29
+ def generate_lyrics(sample):
30
+ model_name = "openai/whisper-tiny.en"
31
+ model_config = WhisperConfig.from_pretrained(model_name)
32
+ processor = WhisperProcessor.from_pretrained(model_name)
33
+ asr_model = WhisperForConditionalGeneration.from_pretrained(model_name, config=model_config)
34
+ asr_model.eval()
35
+ input_features = processor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features
36
+ transcript = asr_model.generate(input_features)
37
+ predicted_ids = asr_model.generate(input_features)
38
+ transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
39
+ lyrics = transcription[0]
40
+ return lyrics
41
+
42
+ def generate_summary(lyrics):
43
+ summarizer = pipeline("summarization", model="philschmid/bart-large-cnn-samsum")
44
+ summary = summarizer(lyrics)
45
+ return summary
46
+
47
+ def generate_prompt(summary):
48
+ # model_name = "gpt2"
49
+ # tokenizer = AutoTokenizer.from_pretrained(model_name)
50
+ # model = AutoModelWithLMHead.from_pretrained(model_name)
51
+ # Set up GPT-2 model and tokenizer
52
+ model_name = 'gpt2'
53
+ tokenizer = GPT2Tokenizer.from_pretrained(model_name)
54
+ model = GPT2LMHeadModel.from_pretrained(model_name)
55
+ # Set the device to GPU if available
56
+ model = model.to(device)
57
+ # Generate prompt text using GPT-2
58
+ prompt = f"Create an image that represents the feeling of '{summary}'"
59
+ # Generate the image prompt
60
+ input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
61
+ output = model.generate(input_ids, do_sample=True, max_length=100, temperature=0.7)
62
+ prompt_text = tokenizer.decode(output[0], skip_special_tokens=True)
63
+ return prompt_text
64
+
65
+ def generate_image(prompt,
66
+ height = 512, # default height of Stable Diffusion
67
+ width = 512 , # default width of Stable Diffusion
68
+ num_inference_steps = 50 , # Number of denoising steps
69
+ guidance_scale = 7.5 , # Scale for classifier-free guidance
70
+ generator = torch.manual_seed(32), # Seed generator to create the inital latent noise
71
+ batch_size = 1,):
72
+ pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
73
+ pipe = pipe.to(torch_device)
74
+ # 1. Load the autoencoder model which will be used to decode the latents into image space.
75
+ vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
76
+ # 2. Load the tokenizer and text encoder to tokenize and encode the text.
77
+ tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
78
+ text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
79
+ # 3. The UNet model for generating the latents.
80
+ unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
81
+ scheduler = DPMSolverMultistepScheduler.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="scheduler")
82
+ vae = vae.to(torch_device)
83
+ text_encoder = text_encoder.to(torch_device)
84
+ unet = unet.to(torch_device)
85
+ text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
86
+ with torch.no_grad():
87
+ text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]
88
+ max_length = text_input.input_ids.shape[-1]
89
+ uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
90
+ with torch.no_grad():
91
+ uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]
92
+ text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
93
+ latents = torch.randn((batch_size, unet.in_channels, height // 8, width // 8), generator=generator,)
94
+ latents = latents.to(torch_device)
95
+ scheduler.set_timesteps(num_inference_steps)
96
+ latents = latents * scheduler.init_noise_sigma
97
+ for t in tqdm(scheduler.timesteps):
98
+ # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
99
+ latent_model_input = torch.cat([latents] * 2)
100
+ latent_model_input = scheduler.scale_model_input(latent_model_input, t)
101
+ # predict the noise residual
102
+ with torch.no_grad():
103
+ noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
104
+ # perform guidance
105
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
106
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
107
+ # compute the previous noisy sample x_t -> x_t-1
108
+ latents = scheduler.step(noise_pred, t, latents).prev_sample
109
+ # scale and decode the image latents with vae
110
+ latents = 1 / 0.18215 * latents
111
+ with torch.no_grad():
112
+ image = vae.decode(latents).sample
113
+ # DPM Solver Multistep scheduler
114
+ image = (image / 2 + 0.5).clamp(0, 1)
115
+ image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
116
+ images = (image * 255).round().astype("uint8")
117
+ pil_images = [Image.fromarray(image) for image in images]
118
+ f_images = pil_images
119
+ return f_images
120
+
121
+ def predict(lyrics, steps=100, seed=42, guidance_scale=6.0):
122
+
123
+ print(subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT).decode("utf8"))
124
+ generator = torch.manual_seed(seed)
125
+ summary_1 = generate_summary(lyrics)
126
+ prompt_text_1 = generate_prompt(summary_1[0]['summary_text'])
127
+ images = generate_image(prompt= prompt_text_1)
128
+ # images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=guidance_scale)["images"]
129
+ print(subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT).decode("utf8"))
130
+ return images[0]
131
+
132
+ random_seed = random.randint(0, 2147483647)
133
+ gr.Interface(
134
+ predict,
135
+ inputs=[
136
+ gr.inputs.Textbox(label='Text', default='a chalk pastel drawing of a llama wearing a wizard hat'),
137
+ gr.inputs.Slider(1, 100, label='Inference Steps', default=50, step=1),
138
+ gr.inputs.Slider(0, 2147483647, label='Seed', default=random_seed, step=1),
139
+ gr.inputs.Slider(1.0, 20.0, label='Guidance Scale - how much the prompt will influence the results', default=6.0, step=0.1),
140
+ ],
141
+ outputs=gr.Image(shape=[256,256], type="pil", elem_id="output_image"),
142
+ css="#output_image{width: 256px}",
143
+ title="Cover Generator (text-to-image)",
144
+ description="Application of OpenAI tools such as Whisper, ChatGPT, and DALL-E to produce covers for the given text",
145
+ ).launch()