suresh-subramanian's picture
Update app.py
dbf5e52
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()