app.py

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

HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "crowdsourced-build-a-movie-poster-demo")

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

generator = torch.Generator(device=device)
seed = 1016
generator = generator.manual_seed(seed)
latent_tensor = torch.randn(size=(1, 4, 64, 64),
                            generator = generator,
                            device=device)

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=HF_TOKEN)
pipe = pipe.to(device)

def generate(celebrity, movie):
  prompt = f"A movie poster of {celebrity} in {movie}"
  image = pipe(prompt, latents=latent_tensor, num_inference_steps=50, guidance_scale=7).images[0]
  return image

iface = gr.Interface(
  fn=generate,
  inputs=[gr.Textbox(label='Celebrity Name'),
          gr.Dropdown(['Terminator',
                       'Captain America',
                       'Karate Kid',
                       'Titanic',
                       'The Crown'], label="Movie")],
  outputs=gr.Image(type='pil'),
  allow_flagging="manual",
  flagging_options=["Celebrity Issues", "No Poster generated", "Good Poster", "Bad Poster"],
  flagging_callback=hf_writer
)

iface.launch(debug=True)