import transformers
import gradio as gr
import datasets
import torch
import os
from diffusers import StableDiffusionPipeline

def generate(movie, celebrity, scale):
  prompt = f'Poster from movie {movie}, {celebrity} wearing lion head opal opal opal opal opal opal + high-tech ornate diamond bejeweled seductress with sapphire, ruby, emerald, gold, live-action fairy tale character, hyper realistic, live-action epic dreams, cinematic concept art, beautiful dark-skinned, beautiful lighting, intricate gold-and-silver-plated armor'
  image = pipe(prompt, generator=generator, guidance_scale=scale).images[0]
  return image
  

# Initialise the token
HF_TOKEN = os.getenv('HF_TOKEN')

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

hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, 'crowdsourced-movieposter-demo')

gr.Interface(
  fn=generate,
  inputs=[gr.Textbox(label='Celebrity Name'), gr.Dropdown(['Mission Impossible', 'Avatar', 'Abyss', 'TOO OLD TO DIE YOUNG', 'The Last Dual','The beginning of everything'], label='Movie'), gr.Slider(label='Image Accuracy', minimum=7.5, maximum=20, value=12)],
  outputs=gr.Image(type='pil') ,
  allow_flagging="manual",
  flagging_options=["No Image", "good", "bad"],
  flagging_callback=hf_writer
).launch()