app.py

# Import libraries
from transformers import pipeline
from numpy import random
import gradio as gr
import re
import torch
from torch import autocast
import os

# Array with song cover art styles
image_input_styles = ["Random", "Pencil sketch", "Oil painting", "Pop art", "Piet Mondriaan"]

# Get image type for image input
"""
The default setting for the art style dropdown is "Random". The below function determines which style is chosen
If set to "Random", copy the art style array and remove "Random" to prevent "Random" from being a chosen art style
"""
def get_image_input(title, given_input_style):
    if given_input_style == 'Random':
      image_input_styles_new = image_input_styles.copy()
      image_input_styles_new.pop(0)
      random_choice = random.randint(len(image_input_styles_new)-1)
      final_style = image_input_styles_new[random_choice]
    else:
      final_style = given_input_style
    image_input = 'Cover for ' + title.lower() + ' in style of ' + final_style 
    return image_input, final_style

# Available models for generate lyrics pipeline
# checkpoint = 'wvangils/GPT-Medium-Beatles-Lyrics-finetuned-newlyrics'
checkpoint = 'wvangils/GPT-Neo-125m-Beatles-Lyrics-finetuned-newlyrics'
# checkpoint = 'wvangils/BLOOM-560m-Beatles-Lyrics-finetuned'

# Setup all the pipelines we need
title_generator = pipeline('summarization', model='czearing/story-to-title')
lyrics_generator = pipeline("text-generation", model=checkpoint)

# For the image generator we use stable diffusion from an existing HuggingFace space, Gradio accelerated backend
stable_diffusion = gr.Blocks.load(name="spaces/stabilityai/stable-diffusion-1")

# Create 4 images for the given prompt and receive the first one
# This function uses an existing HuggingFace space where the number of created images cannot be modified
def get_image(prompt):
    gallery_dir = stable_diffusion(prompt, fn_index=2)
    images = [os.path.join(gallery_dir, img) for img in os.listdir(gallery_dir)]
    return [images[0]]

# Lyrics generation
def generate_beatles(input_prompt, temperature, top_p, given_input_style):
    # Create generator for different models
    generated_lyrics = lyrics_generator(input_prompt
                              , max_length = 100
                              , num_return_sequences = 1
                              , return_full_text = True
                              , temperature = temperature    
                              , top_p = top_p    # Default 1.0
                              , no_repeat_ngram_size = 3   # Default = 0
                              , repetition_penalty = 1.0   # Default = 1.0
                              )[0]["generated_text"]

    # Put lyrics in the right form
    lyrics_sentences = re.sub('\n', '. ', generated_lyrics)

    # Create a title based on the generated lyrics
    title = title_generator(lyrics_sentences, min_length=1, max_length=10, repetition_penalty=2.5)[0]['summary_text']

    # Create an image based on the generated title
    image_input, image_style = get_image_input(title, given_input_style)

    # Generate the image
    image = get_image(image_input)
    return (title, generated_lyrics, image, image_style)
    
# Create textboxes for input and output
input_box = gr.Textbox(label="Write the start of a song here", placeholder="Write the start of a new song here", value="It's been three days ", lines=2, max_lines=5)
gen_lyrics = gr.Textbox(label="Song lyrics", lines=15)
gen_title = gr.Textbox(label="Proposed songtitle", lines=1)
gen_image = gr.Gallery(label="Proposed song cover").style(grid=1, height="auto")
gen_image_style = gr.Textbox(label="Image style", lines=1)

# Layout and text around the app
title='Beatles lyrics generator'
description="<p style='text-align: center'>We've fine-tuned multiple language models on lyrics from The Beatles to generate Beatles-like text. Below are the results we obtained fine-tuning a GPT Neo model. After generation a title is generated using <a href='https://huggingface.co/czearing/story-to-title' target='_blank'>this model</a>. On top we use the generated title to suggest an album cover using <a href='https://huggingface.co/CompVis/stable-diffusion-v1-4' target='_blank'>Stable Diffusion 1.4</a>. Give it a try!</p>"
article="""<p style='text-align: left'>These text generation models that output Beatles-like text were created by data
    scientists working for <a href='https://cmotions.nl/' target="_blank">Cmotions.</a>
    We tried several text generation models that we were able to load in Colab: a general <a
        href='https://huggingface.co/gpt2-medium' target='_blank'>GPT2-medium</a> model, the Eleuther AI small-sized GPT
    model <a href='https://huggingface.co/EleutherAI/gpt-neo-125M' target='_blank'>GPT-Neo</a> and the new kid on the
    block build by the <a
        href='https://bigscience.notion.site/BLOOM-BigScience-176B-Model-ad073ca07cdf479398d5f95d88e218c4'
        target='_blank'>Bigscience</a> initiative <a href='https://huggingface.co/bigscience/bloom-560m'
        target='_blank'>BLOOM 560m</a>.
    Further we've put together a <a href='https://huggingface.co/datasets/cmotions/Beatles_lyrics' target='_blank'>
        Huggingface dataset</a> containing all known lyrics created by The Beatles. A general blog on building this text generator can be found <a
        href='https://www.theanalyticslab.nl/building-a-beatles-lyrics-generator/' target='_blank'>here. </a> If you are interested in how we evaluated the results of these models <a
        href='https://www.theanalyticslab.nl/how-to-evaluate-a-text-generation-model-strengths-and-limitations-of-popular-evaluation-metrics/' target='_blank'>this blog</a> will enlighten you. Finally we also wrote on our experience to showcase the results of our model
        by using HF spaces, details can be found <a href='https://www.theanalyticslab.nl/how-to-showcase-your-demo-on-a-hugging-face-space/' target='_blank'>here. </a>
    The default output contains 100 tokens and has a repetition penalty of 1.0.
</p>"""
css = """
        .gr-button-primary {
          text-indent: -9999px;
          line-height: 0;
        }  
        .gr-button-primary:after {
          content: "Beatlify!";
          text-indent: 0;
          display: block;
          line-height: initial;
        }
""" 

# Let users select their own temperature and top-p
temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, label="Change the temperature \r\n (higher temperature = more creative in lyrics generation, but posibbly less Beatly)", value=0.7, show_label=True)  #high = sensitive for low probability tokens
top_p = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, label="Change top probability of the next word \n (higher top probability = more words to choose from for the next word, but possibly less Beatly)", value=0.5, show_label=True)
given_input_style = gr.Dropdown(choices=image_input_styles, value="Random", label="Choose the art style for the lyrics cover", show_label=True)

# Use generate Beatles function in demo-app Gradio
gr.Interface(fn=generate_beatles
             , inputs=[input_box, temperature, top_p, given_input_style]
             , outputs=[gen_title, gen_lyrics, gen_image, gen_image_style]
             , title=title
             , css=css
             , description=description
             , article=article
             , allow_flagging='never'
             ).launch()