import gradio as gr
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt
import os
import torch

token = os.environ.get('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(token, "crowdsourced-movie-poster-demo")

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=token)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
pipe = pipe.to(device)


def generate(celebrity, setting):
  prompt = f"A movie poster of {celebrity} in {setting}" 
  return (pipe(prompt).images[0]).resize((224,224))

gr.Interface(
  fn = generate,
  inputs=[gr.Textbox(label='Celebrity'),
          gr.Dropdown(['House of the Dragon',
                       'Good will Hunting',
                       'About Love',
                       'Friends',
                       "That '70s Show"], label="Movie/Tv Show")],
  outputs = gr.Image(type='pil'),
  allow_flagging = "manual",
  flagging_options=["Good Poster", "Not So Good Poster"],
  flagging_callback=hf_writer,
  description='Create a movie poster with whoever celebrity you like with Stable Diffusion'
).launch(debug=True,enable_queue=True)