import gradio as gr
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe = pipe.to("cuda")
client = Client('abidlabs/GFPGAN')

# setting=Tom Hanks nightwing in real life camera, high-def, high-res, 4k, dark night, raining, thunder
def predict(celebrity, movie,setting):
  prompt=f" Create a movie poster featuring {celebrity} along with the movie title. The poster should include the title of the movie as {movie},and in a given settings as {setting}"
  image = pipe(prompt).images[0]
  tmpFileName = 'tmp_image.png'
  image.save(tmpFileName, 'PNG')
  gfpGanFileName = client.predict(tmpFileName, 2)

  return Image.open(gfpGanFileName)

import gradio as gr
import os

HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer =  gr.HuggingFaceDatasetSaver(HF_TOKEN, "movie-poster")

gr.Interface(
  fn = predict,  
    inputs=[
        gr.inputs.Textbox(label="celebrity"),
        gr.inputs.Dropdown(label="Movies", choices=["Batman", "Gladiator", "Mission Impossible"]),
        gr.inputs.Textbox(label="Setting")               
    ],
    outputs= gr.outputs.Image(label="Generated movie poster",type='pil'),
    description="Generate a movie poster with the given information",
    title = "Movie Poster Generator",
    allow_flagging="manual",
    flagging_options=["Good", "Better", "Worst"],
    flagging_callback=hf_writer
).launch(share=True)