RKoops commited on
Commit
d812b13
·
1 Parent(s): 3c01271

Create new file

Browse files
Files changed (1) hide show
  1. app.py +52 -0
app.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import StableDiffusionPipeline
2
+ import requests
3
+ import os
4
+ import gradio as gr
5
+ import torch
6
+
7
+ SEED = 42
8
+ AUTH_TOKEN = os.getenv("auth_token")
9
+ HF_TOKEN = os.getenv('HF_TOKEN')
10
+ DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
11
+
12
+ pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=AUTH_TOKEN)
13
+ pip = pipe.to(DEVICE)
14
+ hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "celebrity-set-dataset")
15
+
16
+ # Ensure consistently generated images
17
+ generator = torch.Generator(device=DEVICE).manual_seed(SEED)
18
+ latent = torch.randn(
19
+ (1, 4, 64, 64),
20
+ generator = generator,
21
+ device = DEVICE
22
+ )
23
+
24
+ def generate(celebrity, setting):
25
+ prompt = "A movie poster with {} in {}.".format(celebrity, setting)
26
+ return improve_image(pipe(prompt, latents=latent).images[0], 2)
27
+
28
+ # Use the GANS model of Abubakar
29
+ def improve_image(img, rescaling_factor = 1):
30
+ return gr.processing_utils.decode_base64_to_image(
31
+ requests.post(
32
+ url = 'https://hf.space/embed/abidlabs/GFPGAN/+/api/predict',
33
+ json = {
34
+ "data": [
35
+ gr.processing_utils.encode_pil_to_base64(img),
36
+ rescaling_factor
37
+ ]}
38
+ ).json()['data'][0])
39
+
40
+ gr.Interface(
41
+ inputs = [
42
+ gr.Textbox(label = 'Celebrity'),
43
+ gr.Dropdown(
44
+ choices = ['Star Trek', 'Star Wars', 'The Wire', 'Breaking Bad', 'a rainforest', 'a skyscraper.'],
45
+ label = 'Movie / Show / Setting')
46
+ ],
47
+ fn = generate,
48
+ outputs = "image"
49
+ allow_flagging = "manual",
50
+ flagging_options = ["Looks good", "Looks bad"],
51
+ flagging_callback = hf_writer
52
+ ).launch()