from diffusers import StableDiffusionPipeline
import gradio as gr
import requests
import base64
from PIL import Image, PngImagePlugin
from io import BytesIO
import torch

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

def encode_pil_to_base64(pil_image):
  ''' From: https://github.com/gradio-app/gradio/blob/main/gradio/processing_utils.py'''
  with BytesIO() as output_bytes:

      # Copy any text-only metadata
      use_metadata = False
      metadata = PngImagePlugin.PngInfo()
      for key, value in pil_image.info.items():
          if isinstance(key, str) and isinstance(value, str):
              metadata.add_text(key, value)
              use_metadata = True

      pil_image.save(
          output_bytes, "PNG", pnginfo=(metadata if use_metadata else None)
      )
      bytes_data = output_bytes.getvalue()
  base64_str = str(base64.b64encode(bytes_data), "utf-8")
  return "data:image/png;base64," + base64_str

def decode_base64_to_image(encoding):
  ''' From: https://github.com/gradio-app/gradio/blob/main/gradio/processing_utils.py'''
  content = encoding.split(";")[1]
  image_encoded = content.split(",")[1]
  return Image.open(BytesIO(base64.b64decode(image_encoded)))

def improve_image(img, scale=2):
  ''' Improves an input image using GFP-GAN
  Inputs
  img (PIL): image to improve
  scale (int): scale factor for new image
  Output
  Improved image. If the request to GFPGAN is unsuccesful, it returns
  a black image.
  '''
  url = "https://hf.space/embed/NotFungibleIO/GFPGAN/+/api/predict"
  request_objt = {"data":[encode_pil_to_base64(img),'v1.3', scale]}
  try:
    imp_img = decode_base64_to_image(requests.post(url, json=request_objt).json()['data'][0])
  except AttributeError:
    return Image.new('RGB', size=(512, 512))
  return imp_img

def generate(celebrity, movie, guidance, improve_flag, scale): # add scale as var
  prompt = f"A movie poster of {celebrity} in {movie}."
  image = pipe(prompt, guidance=guidance).images[0]
  if improve_flag:
    image = improve_image(image, scale=scale)
  return image

movie_options = ["James Bond", "Snatch", "Saving Private Ryan", "Scarface", "Avatar", "Top Gun"]
title = "Movie Poster Celebrity Swap"
description = "Write the name of a celebrity, and pick a movie from the dropdown menu.\
               It will generate a new movie poster (inspired by your chosen movie)\
               with the chosen celebrity in it. See below for explanation of the\
               input variables."
article= "Inputs explained: \n Guidance: the lower, the more random the output\
          image. Improve and scale: if selected, the image will be refined\
          using GFP-GAN (google it), and scaled (if scale is >1)."

demo = gr.Interface(
    fn=generate,
    inputs=[gr.Textbox(value="Daniel Craig"), 
            gr.Dropdown(movie_options, value="Saving Private Ryan"), 
            gr.Slider(1, 20, value=7.5, step=0.5),
            gr.Checkbox(label="Improve and scale? (extra time)"),
            gr.Slider(1, 3, value=1, step=0.5)
            ],
    outputs='image',
    title=title,
    description=description,
    article=article
) 

demo.launch()