baondi commited on
Commit
fef413e
verified
1 Parent(s): 1314a7f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +265 -0
app.py ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ import pdb
4
+ import random
5
+ import numpy as np
6
+ from PIL import Image
7
+ import base64
8
+ from io import BytesIO
9
+
10
+ import torch
11
+ from torchvision import transforms
12
+ import torchvision.transforms.functional as TF
13
+ import gradio as gr
14
+
15
+ from src.model import make_1step_sched
16
+ from src.pix2pix_turbo import Pix2Pix_Turbo
17
+
18
+ model = Pix2Pix_Turbo("sketch_to_image_stochastic")
19
+
20
+ style_list = [
21
+ {
22
+ "name": "No Style",
23
+ "prompt": "{prompt}",
24
+ },
25
+ {
26
+ "name": "Cinematic",
27
+ "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
28
+ },
29
+ {
30
+ "name": "3D Model",
31
+ "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
32
+ },
33
+ {
34
+ "name": "Anime",
35
+ "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime, highly detailed",
36
+ },
37
+ {
38
+ "name": "Digital Art",
39
+ "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
40
+ },
41
+ {
42
+ "name": "Photographic",
43
+ "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
44
+ },
45
+ {
46
+ "name": "Pixel art",
47
+ "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
48
+ },
49
+ {
50
+ "name": "Fantasy art",
51
+ "prompt": "ethereal fantasy concept art of {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
52
+ },
53
+ {
54
+ "name": "Neonpunk",
55
+ "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
56
+ },
57
+ {
58
+ "name": "Manga",
59
+ "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
60
+ },
61
+ ]
62
+
63
+ styles = {k["name"]: k["prompt"] for k in style_list}
64
+ STYLE_NAMES = list(styles.keys())
65
+ DEFAULT_STYLE_NAME = "Fantasy art"
66
+ MAX_SEED = np.iinfo(np.int32).max
67
+
68
+
69
+ def pil_image_to_data_uri(img, format='PNG'):
70
+ buffered = BytesIO()
71
+ img.save(buffered, format=format)
72
+ img_str = base64.b64encode(buffered.getvalue()).decode()
73
+ return f"data:image/{format.lower()};base64,{img_str}"
74
+
75
+
76
+ def run(image, prompt, prompt_template, style_name, seed, val_r, brush_size, brush_color):
77
+ print(f"seed: {seed}, r_val: {val_r}")
78
+ print("sketch updated")
79
+ if image is None:
80
+ ones = Image.new("L", (512, 512), 255)
81
+ temp_uri = pil_image_to_data_uri(ones)
82
+ return ones, gr.update(link=temp_uri), gr.update(link=temp_uri)
83
+ prompt = prompt_template.replace("{prompt}", prompt)
84
+ image = image.convert("RGB")
85
+ image_t = TF.to_tensor(image) > 0.5
86
+ image_pil = TF.to_pil_image(image_t.to(torch.float32))
87
+ print(f"r_val={val_r}, seed={seed}")
88
+ with torch.no_grad():
89
+ c_t = image_t.unsqueeze(0).cuda().float()
90
+ torch.manual_seed(seed)
91
+ B,C,H,W = c_t.shape
92
+ noise = torch.randn((1,4,H//8, W//8), device=c_t.device)
93
+ output_image = model(c_t, prompt, deterministic=False, r=val_r, noise_map=noise)
94
+ output_pil = TF.to_pil_image(output_image[0].cpu()*0.5+0.5)
95
+ input_sketch_uri = pil_image_to_data_uri(Image.fromarray(255-np.array(image)))
96
+ output_image_uri = pil_image_to_data_uri(output_pil)
97
+ return output_pil, gr.update(link=input_sketch_uri), gr.update(link=output_image_uri)
98
+
99
+
100
+ def update_canvas(brush_size, brush_color):
101
+ return gr.update(brush_radius=brush_size, brush_color=brush_color, interactive=True)
102
+
103
+
104
+ def upload_sketch(file):
105
+ _img = Image.open(file.name)
106
+ _img = _img.convert("L")
107
+ return gr.update(value=_img, source="upload", interactive=True)
108
+
109
+
110
+ style_css = """
111
+ /* Colores del bot贸n */
112
+ .gradio .input_color input[type="color"]::-webkit-color-swatch {
113
+ border: 1px solid black;
114
+ }
115
+
116
+ /* Tama帽o del bot贸n */
117
+ .gradio .input_number input[type="number"] {
118
+ width: 120px;
119
+ }
120
+
121
+ /* Tama帽o del deslizador */
122
+ .gradio .input_slider input[type="range"] {
123
+ width: 120px;
124
+ }
125
+ """
126
+
127
+ scripts = """
128
+ async () => {
129
+ globalThis.theSketchDownloadFunction = () => {
130
+ console.log("test")
131
+ var link = document.createElement("a");
132
+ dataUri = document.getElementById('download_sketch').href
133
+ link.setAttribute("href", dataUri)
134
+ link.setAttribute("download", "sketch.png")
135
+ document.body.appendChild(link); // Required for Firefox
136
+ link.click();
137
+ document.body.removeChild(link); // Clean up
138
+
139
+ // also call the output download function
140
+ theOutputDownloadFunction();
141
+ return false
142
+ }
143
+ globalThis.theOutputDownloadFunction = () => {
144
+ console.log("test output download function")
145
+ var link = document.createElement("a");
146
+ dataUri = document.getElementById('download_output').href
147
+ link.setAttribute("href", dataUri);
148
+ link.setAttribute("download", "output.png");
149
+ document.body.appendChild(link); // Required for Firefox
150
+ link.click();
151
+ document.body.removeChild(link); // Clean up
152
+ return false
153
+ }
154
+ globalThis.UNDO_SKETCH_FUNCTION = () => {
155
+ console.log("undo sketch function")
156
+ var button_undo = document.querySelector('#input_image > div.image-container.svelte-p3y7hu > div.svelte-s6ybro > button:nth-child(1)');
157
+ // Create a new 'click' event
158
+ var event = new MouseEvent('click', {
159
+ 'view': window,
160
+ 'bubbles': true,
161
+ 'cancelable': true
162
+ });
163
+ button_undo.dispatchEvent(event);
164
+ }
165
+ globalThis.DELETE_SKETCH_FUNCTION = () => {
166
+ console.log("delete sketch function")
167
+ var button_del = document.querySelector('#input_image > div.image-container.svelte-p3y7hu > div.svelte-s6ybro > button:nth-child(2)');
168
+ // Create a new 'click' event
169
+ var event = new MouseEvent('click', {
170
+ 'view': window,
171
+ 'bubbles': true,
172
+ 'cancelable': true
173
+ });
174
+ button_del.dispatchEvent(event);
175
+ }
176
+ globalThis.togglePencil = () => {
177
+ el_pencil = document.getElementById('my-toggle-pencil');
178
+ el_pencil.classList.toggle('clicked');
179
+ // simulate a click on the gradio button
180
+ btn_gradio = document.querySelector("#cb-line > label > input");
181
+ var event = new MouseEvent('click', {
182
+ 'view': window,
183
+ 'bubbles': true,
184
+ 'cancelable': true
185
+ });
186
+ btn_gradio.dispatchEvent(event);
187
+ if (el_pencil.classList.contains('clicked')) {
188
+ document.getElementById('my-toggle-eraser').classList.remove('clicked');
189
+ document.getElementById('my-div-pencil').style.backgroundColor = "gray";
190
+ document.getElementById('my-div-eraser').style.backgroundColor = "white";
191
+ }
192
+ else {
193
+ document.getElementById('my-toggle-eraser').classList.add('clicked');
194
+ document.getElementById('my-div-pencil').style.backgroundColor = "white";
195
+ document.getElementById('my-div-eraser').style.backgroundColor = "gray";
196
+ }
197
+
198
+ }
199
+ globalThis.toggleEraser = () => {
200
+ element = document.getElementById('my-toggle-eraser');
201
+ element.classList.toggle('clicked');
202
+ // simulate a click on the gradio button
203
+ btn_gradio = document.querySelector("#cb-eraser > label > input");
204
+ var event = new MouseEvent('click', {
205
+ 'view': window,
206
+ 'bubbles': true,
207
+ 'cancelable': true
208
+ });
209
+ btn_gradio.dispatchEvent(event);
210
+ if (element.classList.contains('clicked')) {
211
+ document.getElementById('my-toggle-pencil').classList.remove('clicked');
212
+ document.getElementById('my-div-pencil').style.backgroundColor = "white";
213
+ document.getElementById('my-div-eraser').style.backgroundColor = "gray";
214
+ }
215
+ else {
216
+ document.getElementById('my-toggle-pencil').classList.add('clicked');
217
+ document.getElementById('my-div-pencil').style.backgroundColor = "gray";
218
+ document.getElementById('my-div-eraser').style.backgroundColor = "white";
219
+ }
220
+ }
221
+ }
222
+ """
223
+
224
+
225
+ inputs = [
226
+ gr.inputs.Image(label="Input Sketch", type="pil"),
227
+ gr.inputs.Text(label="Prompt", default=""),
228
+ gr.inputs.Dropdown(label="Style", choices=STYLE_NAMES, default=DEFAULT_STYLE_NAME),
229
+ gr.inputs.Slider(label="Seed", min_value=0, max_value=MAX_SEED, default=42, step=1,),
230
+ gr.inputs.Slider(label="R", min_value=0, max_value=1, default=0.6, step=0.01),
231
+ gr.inputs.Slider(label="Brush Size", min_value=1, max_value=50, default=4, step=1),
232
+ gr.inputs.ColorPicker(label="Brush Color", default="#000000")
233
+ ]
234
+
235
+ outputs = [
236
+ gr.outputs.Image(label="Result"),
237
+ gr.outputs.Image(label="Input Sketch", type="pil"),
238
+ gr.outputs.Image(label="Output Image", type="pil"),
239
+ ]
240
+
241
+ title = "pix2pix-Turbo: Sketch"
242
+ description = "One-Step Image Translation with Text-to-Image Models. Paper: [One-Step Image Translation with Text-to-Image Models](https://arxiv.org/abs/2403.12036). GitHub: [pix2pix-Turbo](https://github.com/GaParmar/img2img-turbo)"
243
+ examples = [["example.jpg", "A cat", "Fantasy art", 42, 0.6, 4, "#000000"]]
244
+ server_port = 7860
245
+
246
+
247
+ if __name__ == "__main__":
248
+ gr.Interface(
249
+ run,
250
+ inputs,
251
+ outputs,
252
+ title=title,
253
+ description=description,
254
+ examples=examples,
255
+ theme="huggingface",
256
+ allow_flagging=False,
257
+ layout="unaligned",
258
+ live=True,
259
+ capture_session=True,
260
+ server_port=server_port,
261
+ css=style_css,
262
+ scripts=scripts,
263
+ update_canvas=update_canvas,
264
+ upload_sketch=upload_sketch,
265
+ ).launch(share=True)