File size: 1,975 Bytes
bf5a03f ae221cd 51d2760 ae221cd 95aa254 51d2760 95aa254 bf5a03f ae221cd bf5a03f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
from diffusers import StableDiffusionPipeline
import torch
import requests
import base64
import gradio as gr
import os
from gradio.components import Textbox
HF_TOKEN = os.getenv('HF_TOKEN')
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "movie-img-generator")
auth_token = os.environ.get("auth_token")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=auth_token)
pipe = pipe.to(device)
generator = torch.Generator(device=device)
seed = generator.seed()
def improve_image(img):
url = "https://hf.space/embed/abidlabs/GFPGAN/+/api/predict"
img_to_str = gr.processing_utils.encode_pil_to_base64(img)
inputs = {'data': [img_to_str, 2]}
resp = requests.post(url, json = inputs)
resp_img = gr.processing_utils.decode_base64_to_image(resp.json()['data'][0])
return resp_img
def generate(celebrity, setting):
prompt = "A movie poster with title and caption of {} in {}".format(celebrity, setting)
image = pipe(prompt,
latents=torch.randn((1, 4, 64, 64), generator=torch.manual_seed(12)),
guidance_scale=7).images[0]
enhanced_img = improve_image(image)
return enhanced_img
demo = gr.Interface(
title="🎬 Movie Poster Image Generator",
description="Generate an image of a celebrity in a cinematic setting by entering in any celebrity name and selecting a TV show/movie to see what image is generated!🍿",
inputs=[gr.Textbox(lines=1, label="Enter a celebrity name:"),
gr.Dropdown(["Stranger Things",
"Mean Girls",
"Titanic",
"The Big Bang Theory",
"Game of Thrones"], label="Select a setting:"),],
outputs="image",
fn=generate,
allow_flagging="manual",
flagging_options=["Wrong Celebrity", "Wrong Setting", "Image Not Clear"],
flagging_callback=hf_writer
)
demo.launch(debug=True)
|