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)