from diffusers import StableDiffusionPipeline
import torch
import requests
import base64
import gradio as gr
import os
from gradio.components import Textbox

HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "movie-img-generator")
auth_token = os.environ.get("auth_token")

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

generator = torch.Generator(device=device)
seed = generator.seed()

def improve_image(img):
  url = "https://hf.space/embed/abidlabs/GFPGAN/+/api/predict"
  img_to_str = gr.processing_utils.encode_pil_to_base64(img)
  inputs = {'data': [img_to_str, 2]}
  resp = requests.post(url, json = inputs)
  resp_img = gr.processing_utils.decode_base64_to_image(resp.json()['data'][0])
  return resp_img

def generate(celebrity, setting):
  prompt = "A movie poster with title and caption of {} in {}".format(celebrity, setting)
  image = pipe(prompt, 
               latents=torch.randn((1, 4, 64, 64), generator=torch.manual_seed(12)),
               guidance_scale=7).images[0]

  enhanced_img = improve_image(image)
  return enhanced_img

demo = gr.Interface(
  title="🎬 Movie Poster Image Generator",
  description="Generate an image of a celebrity in a cinematic setting by entering in any celebrity name and selecting a TV show/movie to see what image is generated!🍿",
  inputs=[gr.Textbox(lines=1, label="Enter a celebrity name:"), 
          gr.Dropdown(["Stranger Things", 
                       "Mean Girls", 
                       "Titanic", 
                       "The Big Bang Theory", 
                       "Game of Thrones"], label="Select a setting:"),],
  outputs="image",
  fn=generate,
  allow_flagging="manual",
  flagging_options=["Wrong Celebrity", "Wrong Setting", "Image Not Clear"],
  flagging_callback=hf_writer
)  
demo.launch(debug=True)