shaoan xie commited on
Commit
d124f41
1 Parent(s): bc2c9f6
Files changed (2) hide show
  1. .idea/.gitignore +8 -0
  2. app.py +180 -0
.idea/.gitignore ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ # Default ignored files
2
+ /shelf/
3
+ /workspace.xml
4
+ # Editor-based HTTP Client requests
5
+ /httpRequests/
6
+ # Datasource local storage ignored files
7
+ /dataSources/
8
+ /dataSources.local.xml
app.py ADDED
@@ -0,0 +1,180 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import pickle
4
+ from torchvision.utils import save_image
5
+ import numpy as np
6
+ from diffusers import StableDiffusionUpscalePipeline
7
+ from huggingface_hub import hf_hub_download
8
+ import torch
9
+
10
+ # Load the model from Hugging Face Hub
11
+ model_path = hf_hub_download(repo_id="Shaoan/ConceptGAN", filename="augceleba_6451.pkl")
12
+
13
+ with open(model_path, 'rb') as f:
14
+ G = pickle.load(f)['G_ema'].cpu().float() # torch.nn.Module
15
+
16
+
17
+ cchoices = ['Bald',
18
+ 'Black Hair',
19
+ 'Blond Hair',
20
+ 'Smiling',
21
+ 'NoSmile',
22
+ 'Male',
23
+ 'Female'
24
+ ]
25
+
26
+ model_choices = [
27
+ 'Change Dim = 8',
28
+ 'Change Dim = 15',
29
+ 'Change Dim = 30',
30
+ 'Change Dim = 60'
31
+ ]
32
+
33
+
34
+ cchoices = [
35
+ 'Big Nose',
36
+ 'Black Hair',
37
+ 'Blond Hair',
38
+ 'Chubby',
39
+ 'Eyeglasses',
40
+ 'Male',
41
+ 'Pale Skin',
42
+ 'Smiling',
43
+ 'Straight Hair',
44
+ 'Wavy Hair',
45
+ 'Wearing Hat',
46
+ 'Young'
47
+ ]
48
+
49
+
50
+ import requests
51
+ from PIL import Image
52
+ from io import BytesIO
53
+ from diffusers import LDMSuperResolutionPipeline
54
+ import torch
55
+
56
+ device = "cuda" if torch.cuda.is_available() else "cpu"
57
+ model_id = "CompVis/ldm-super-resolution-4x-openimages"
58
+
59
+ # load model and scheduler
60
+ pipeline = LDMSuperResolutionPipeline.from_pretrained(model_id)
61
+ pipeline = pipeline.to(device)
62
+ model_id = "stabilityai/stable-diffusion-x4-upscaler"
63
+ text_pipeline = StableDiffusionUpscalePipeline.from_pretrained(
64
+ model_id, variant="fp32", torch_dtype=torch.float32
65
+ )
66
+ # let's download an image
67
+
68
+
69
+ def super_res(low_res_img, num_steps):
70
+ # run pipeline in inference (sample random noise and denoise)
71
+ upscaled_image = pipeline(low_res_img, num_inference_steps=num_steps, eta=1).images[0]
72
+ #upscaled_image = text_pipeline(prompt="a sharp image of human face", image=low_res_img, num_inference_steps=75).images[0]
73
+ return upscaled_image
74
+
75
+
76
+ @torch.no_grad()
77
+ def generate(seed, upscale, upscale_steps,*checkboxes):
78
+ z = torch.randn([1, G.z_dim], generator=torch.Generator().manual_seed(seed))
79
+ #m = torch.tensor([[1, 0, 0, 0, 1, 1, 0.]]).repeat(1, 1)
80
+ checkboxes_vector = torch.zeros([20])
81
+ for i in range(len(checkboxes)):
82
+ if i == 1:
83
+ checkboxes_vector[cchoices.index('Black Hair')] = checkboxes[i]
84
+ elif i == 2:
85
+ checkboxes_vector[cchoices.index('Blond Hair')] = checkboxes[i]
86
+ elif i == 3:
87
+ checkboxes_vector[cchoices.index('Straight Hair')] = checkboxes[i]
88
+ elif i == 4:
89
+ checkboxes_vector[cchoices.index('Wavy Hair')] = checkboxes[i]
90
+ elif i == 5:
91
+ checkboxes_vector[cchoices.index('Young')] = checkboxes[i] * 2
92
+ elif i == 6:
93
+ checkboxes_vector[cchoices.index('Male')] = checkboxes[i]
94
+ elif i == 9:
95
+ checkboxes_vector[cchoices.index('Big Nose')] = checkboxes[i]
96
+ elif i == 10:
97
+ checkboxes_vector[cchoices.index('Chubby')] = checkboxes[i]
98
+ elif i == 11:
99
+ checkboxes_vector[cchoices.index('Eyeglasses')] = checkboxes[i] * 2
100
+ elif i == 12:
101
+ checkboxes_vector[cchoices.index('Pale Skin')] = checkboxes[i]
102
+ elif i == 13:
103
+ checkboxes_vector[cchoices.index('Smiling')] = checkboxes[i]
104
+ elif i == 14:
105
+ checkboxes_vector[cchoices.index('Wearing Hat')] = checkboxes[i] * 2
106
+
107
+
108
+ is_young = checkboxes[5]
109
+ is_male = checkboxes[6]
110
+ is_bald = checkboxes[0]
111
+ is_goatee = checkboxes[7]
112
+ is_mustache = checkboxes[8]
113
+
114
+ checkboxes_vector[12] = is_mustache * 2
115
+ checkboxes_vector[13] = is_mustache * 2
116
+ checkboxes_vector[14] = is_goatee *2
117
+ checkboxes_vector[15] = is_goatee*2
118
+
119
+ checkboxes_vector[16] = is_bald
120
+ checkboxes_vector[17] = is_bald
121
+ checkboxes_vector[18] = is_bald
122
+ checkboxes_vector[19] = is_bald
123
+
124
+
125
+
126
+ print(checkboxes_vector)
127
+
128
+ m = checkboxes_vector.view(1, 20)
129
+ ws = G.mapping(z, m, truncation_psi=0.5)
130
+ img = (G.synthesis(ws, force_fp32=True).clip(-1,1)+1)/2
131
+ if upscale:
132
+ up_img = np.array(super_res(img*2-1, upscale_steps))
133
+ return up_img
134
+ else:
135
+ return img[0].permute(1, 2, 0).numpy()
136
+
137
+
138
+ # Create the interface using gr.Blocks
139
+ with gr.Blocks() as demo:
140
+ with gr.Row():
141
+ sliders = [
142
+ gr.Slider(label='Bald', minimum=0, maximum=1, step=0.01),
143
+ gr.Slider(label='Black Hair', minimum=0, maximum=1, step=0.01),
144
+ gr.Slider(label='Blond Hair', minimum=0, maximum=1, step=0.01),
145
+ gr.Slider(label='Straight Hair', minimum=0, maximum=1, step=0.01),
146
+ gr.Slider(label='Wavy Hair', minimum=0, maximum=1, step=0.01),
147
+ ]
148
+
149
+ with gr.Row():
150
+ sliders += [gr.Slider(label='Young', minimum=0, maximum=1, step=0.01)]
151
+ sliders += [gr.Slider(label='Male', minimum=0, maximum=1, step=0.01)]
152
+
153
+ with gr.Row():
154
+ sliders += [gr.Slider(label='Goatee', minimum=0, maximum=1, step=0.01)]
155
+ sliders += [gr.Slider(label='Mustache', minimum=0, maximum=1, step=0.01)]
156
+
157
+ with gr.Row():
158
+ sliders += [
159
+ gr.Slider(label='Big Nose', minimum=0, maximum=1, step=0.01),
160
+ gr.Slider(label='Chubby', minimum=0, maximum=1, step=0.01),
161
+ gr.Slider(label='Eyeglasses', minimum=0, maximum=1, step=0.01),
162
+ gr.Slider(label='Pale Skin', minimum=0, maximum=1, step=0.01),
163
+ gr.Slider(label='Smiling', minimum=0, maximum=1, step=0.01),
164
+ gr.Slider(label='Wearing Hat', minimum=0, maximum=1, step=0.01),
165
+ ]
166
+
167
+ seed_input = gr.Number(label="Seed", value=6)
168
+ upscale_funcs = []
169
+ with gr.Row():
170
+ upscale_funcs = [gr.Checkbox(label="Upscale 4x")]
171
+ upscale_funcs += [gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=10)]
172
+ generate_button = gr.Button("Generate")
173
+
174
+ output_image = gr.Image(label="Generated Image")
175
+
176
+ # Set the action for the button
177
+ generate_button.click(fn=generate, inputs=[seed_input] + upscale_funcs +sliders, outputs=output_image)
178
+
179
+ # Launch the demo
180
+ demo.launch()