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

SEED = 42
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)
hf_writer = gr.HuggingFaceDatasetSaver(AUTH_TOKEN, "celebrity-set-dataset")

# Ensure consistently generated images
generator = torch.Generator(device=DEVICE).manual_seed(SEED)
latent = torch.randn(
    (1, 4, 64, 64),
    generator = generator,
    device = DEVICE
)

def generate(celebrity, setting):
  prompt = "A movie poster with {} in {}.".format(celebrity, setting)
  return improve_image(pipe(prompt, latents=latent).images[0], 2)

# Use the GANS model of Abubakar
def improve_image(img, rescaling_factor = 1):
  return gr.processing_utils.decode_base64_to_image(
      requests.post(
          url = 'https://hf.space/embed/abidlabs/GFPGAN/+/api/predict', 
          json = { 
              "data": [
                  gr.processing_utils.encode_pil_to_base64(img), 
                  rescaling_factor
                  ]}
                  ).json()['data'][0])
                  
gr.Interface(
    inputs = [
        gr.Textbox(label = 'Celebrity'), 
        gr.Dropdown(
            choices = ['Star Trek', 'Star Wars', 'The Wire', 'Breaking Bad', 'a rainforest', 'a skyscraper.'], 
            label = 'Movie / Show / Setting')
        ],
    fn = generate,
    outputs = "image",
    allow_flagging = "manual",
    flagging_options = ["Looks good", "Looks bad"],
    flagging_callback = hf_writer
).launch()