qiny17's picture
Create app.py
7229e67
import torch
import gradio as gr
from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
# generate lyrics
lyrics_generator = pipeline("text-generation", "ECE1786-AG/lyrics-generator")
# summarize lyrics
model = T5ForConditionalGeneration.from_pretrained("Michau/t5-base-en-generate-headline")
tokenizer = T5Tokenizer.from_pretrained("Michau/t5-base-en-generate-headline")
# generate single cover
scheduler = EulerDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-2", subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)
def generate_lyrics(genre, prompt):
complete_prompt = "<BOS> <{0}>\n{1}".format(genre, prompt)
lyrics = lyrics_generator(complete_prompt, max_length=1024)
lyrics = lyrics[0]['generated_text']
lyrics = lyrics.split('\n', 1)[1] # remove first line from the generated lyrics
return lyrics
def summarize_lyrics(lyrics):
text = "headline: " + lyrics
encoding = tokenizer.encode_plus(text, return_tensors = "pt")
input_ids = encoding["input_ids"]
attention_masks = encoding["attention_mask"]
beam_outputs = model.generate(
input_ids = input_ids,
attention_mask = attention_masks,
max_length = 100,
num_beams = 5,
early_stopping = True,
)
result = tokenizer.decode(beam_outputs[0])
result = result.replace('<pad>', '')
result = result.replace('</s>', '')
return result
def generate_cover(prompt, style, effect):
prompt = summarize_lyrics(prompt) # call function summarize_lyrics to summarize lyrics
prompt = prompt + ", " + effect + ", album cover, artistic, " + style
print(prompt)
image = pipe(prompt).images[0]
return image
demo = gr.Blocks()
with demo:
gr.HTML(
"""
<div style="text-align: center; max-width: 700px; margin: 0 auto;">
<div style="display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;">
<h1 style="font-weight: 900; margin-bottom: 7px;">ArtIstic GENREator</h1>
</div>
<p style="margin-bottom: 10px; font-size: 94%">Generate Inspirational Lyrics and Single Cover</p>
</div>
"""
)
with gr.Row():
# Left column (lyrics generation)
with gr.Column():
with gr.Accordion("Step 1. Generate Lyrics"):
gr.Markdown("Enter the starting text and select genre to generate lyrics")
with gr.Row():
input_start_text = gr.Textbox(placeholder='I am', label="Starting Text")
input_lyrics_type = gr.Radio(choices=['pop', 'rap', 'country', 'rock', 'r&b'], value='pop', label="Lyrics Genre")
button_gen_lyrics = gr.Button("Generate Lyrics", variant="primary")
output_generated_lyrics = gr.Textbox(label="Generated Lyrics", lines=8)
# Right column (single cover generation)
with gr.Column():
with gr.Accordion("Step 2. Generate Single Cover"):
gr.Markdown("Cover will be generated based on style, effect and generated lyrics")
with gr.Row():
input_cover_style = gr.Dropdown(choices=['painted', 'abstract', 'minimalist', 'illustrated', 'photographic', 'vintage'], value='painted', label="Track Cover Style")
input_cover_effect = gr.Radio(choices=['black and white', 'highly detailed', 'blurred'], value='highly detailed', label="Track Cover Effect")
button_gen_cover = gr.Button("Generate Cover", variant="primary")
output_generated_cover = gr.Image(label="Generated Cover")
# Bind functions to buttons
button_gen_lyrics.click(fn=generate_lyrics, inputs=[input_lyrics_type , input_start_text], outputs=output_generated_lyrics)
button_gen_cover.click(fn=generate_cover, inputs=[output_generated_lyrics, input_cover_style, input_cover_effect], outputs=output_generated_cover)
demo.launch(debug=True)