movie-poster / app.py
Kavindu99's picture
Update app.py
d346d0b
import gradio as gr
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt
import os
import torch
token = os.environ.get('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(token, "crowdsourced-movie-poster-demo")
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=token)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
pipe = pipe.to(device)
def generate(celebrity, setting):
prompt = f"A movie poster of {celebrity} in {setting}"
return (pipe(prompt).images[0]).resize((224,224))
gr.Interface(
fn = generate,
inputs=[gr.Textbox(label='Celebrity'),
gr.Dropdown(['House of the Dragon',
'Good will Hunting',
'About Love',
'Friends',
"That '70s Show"], label="Movie/Tv Show")],
outputs = gr.Image(type='pil'),
allow_flagging = "manual",
flagging_options=["Good Poster", "Not So Good Poster"],
flagging_callback=hf_writer,
description='Create a movie poster with whoever celebrity you like with Stable Diffusion'
).launch(debug=True,enable_queue=True)