from diffusers import StableDiffusionPipeline
import gradio as gr
import os
import torch

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


def generate(celebrity, setting):
  prompt = f"A poster of {celebrity} in {setting}, 35 mm, ultra detailed, cinematic light, photorealistic"
  return pipe(prompt, num_inference_steps=50, num_images_per_prompt=1, guidance_scale=9).images[0]

HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "crowdsourced-text2img")

interface = gr.Interface(
              fn=generate, 
              title="Poster of celebrity X in setting Y",
              inputs=[gr.Textbox(placeholder="Elon Musk"),
                      gr.Dropdown(["Mad Max", "Game of Thrones", "Pulp Fiction", "Moneyball", "The Sopranos", "Jurassic Park", "Cinderella", "The Lion King"],  value="Mad Max")],                     
              outputs=gr.Image(),
              flagging_options=["great result", "satisfactory result", "needs improvement"],
              flagging_callback=hf_writer
            )

interface.launch(debug=True)