import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import requests
import base64
import os
import gradio as gr

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
auth_token = os.getenv("auth_token")
hf_writer = gr.HuggingFaceDatasetSaver(auth_token, "poster-generator-demo")

def improve_image(img, factor):
  encoded_img = gr.processing_utils.encode_pil_to_base64(img)
  URL = "https://hf.space/embed/abidlabs/GFPGAN/+/api/predict"
  PARAMS = {"data": [encoded_img, factor]}
  improved_b64 = requests.post(url = URL, json = PARAMS).json()["data"][0]
  improved_img = gr.processing_utils.decode_base64_to_image(improved_b64)
  return improved_img

def generate(celebrity, setting):
  pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=auth_token)
  pipe = pipe.to(device)
  prompt = f"{celebrity} in {setting} adventure movie poster"
  image = (pipe(prompt,num_inference_steps = 5,guidance_scale=5).images[0])
  return improve_image(image,2)
  

poster_demo = gr.Interface(
fn=generate, 
inputs=[gr.Textbox(lines=2, placeholder="Enter Celebrity Name"),gr.Textbox(lines=2, placeholder="Enter Movie Setting")],
outputs="image", 
allow_flagging="manual",
flagging_options=["Poor Face Quality", "Poor Setting Representation"],
flagging_callback=hf_writer,
)  
poster_demo.launch(debug = True)