Donmill commited on
Commit
bab6abe
1 Parent(s): 0f981a7
Files changed (1) hide show
  1. app1.py +0 -324
app1.py DELETED
@@ -1,324 +0,0 @@
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):
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(use_line, use_eraser):
101
- if use_eraser:
102
- _color = "#ffffff"
103
- brush_size = 20
104
- if use_line:
105
- _color = "#000000"
106
- brush_size = 4
107
- return gr.update(brush_radius=brush_size, brush_color=_color, interactive=True)
108
-
109
-
110
- def upload_sketch(file):
111
- _img = Image.open(file.name)
112
- _img = _img.convert("L")
113
- return gr.update(value=_img, source="upload", interactive=True)
114
-
115
-
116
- scripts = """
117
- async () => {
118
- globalThis.theSketchDownloadFunction = () => {
119
- console.log("test")
120
- var link = document.createElement("a");
121
- dataUri = document.getElementById('download_sketch').href
122
- link.setAttribute("href", dataUri)
123
- link.setAttribute("download", "sketch.png")
124
- document.body.appendChild(link); // Required for Firefox
125
- link.click();
126
- document.body.removeChild(link); // Clean up
127
-
128
- // also call the output download function
129
- theOutputDownloadFunction();
130
- return false
131
- }
132
-
133
- globalThis.theOutputDownloadFunction = () => {
134
- console.log("test output download function")
135
- var link = document.createElement("a");
136
- dataUri = document.getElementById('download_output').href
137
- link.setAttribute("href", dataUri);
138
- link.setAttribute("download", "output.png");
139
- document.body.appendChild(link); // Required for Firefox
140
- link.click();
141
- document.body.removeChild(link); // Clean up
142
- return false
143
- }
144
-
145
- globalThis.UNDO_SKETCH_FUNCTION = () => {
146
- console.log("undo sketch function")
147
- var button_undo = document.querySelector('#input_image > div.image-container.svelte-p3y7hu > div.svelte-s6ybro > button:nth-child(1)');
148
- // Create a new 'click' event
149
- var event = new MouseEvent('click', {
150
- 'view': window,
151
- 'bubbles': true,
152
- 'cancelable': true
153
- });
154
- button_undo.dispatchEvent(event);
155
- }
156
-
157
- globalThis.DELETE_SKETCH_FUNCTION = () => {
158
- console.log("delete sketch function")
159
- var button_del = document.querySelector('#input_image > div.image-container.svelte-p3y7hu > div.svelte-s6ybro > button:nth-child(2)');
160
- // Create a new 'click' event
161
- var event = new MouseEvent('click', {
162
- 'view': window,
163
- 'bubbles': true,
164
- 'cancelable': true
165
- });
166
- button_del.dispatchEvent(event);
167
- }
168
-
169
- globalThis.togglePencil = () => {
170
- el_pencil = document.getElementById('my-toggle-pencil');
171
- el_pencil.classList.toggle('clicked');
172
- // simulate a click on the gradio button
173
- btn_gradio = document.querySelector("#cb-line > label > input");
174
- var event = new MouseEvent('click', {
175
- 'view': window,
176
- 'bubbles': true,
177
- 'cancelable': true
178
- });
179
- btn_gradio.dispatchEvent(event);
180
- if (el_pencil.classList.contains('clicked')) {
181
- document.getElementById('my-toggle-eraser').classList.remove('clicked');
182
- document.getElementById('my-div-pencil').style.backgroundColor = "gray";
183
- document.getElementById('my-div-eraser').style.backgroundColor = "white";
184
- }
185
- else {
186
- document.getElementById('my-toggle-eraser').classList.add('clicked');
187
- document.getElementById('my-div-pencil').style.backgroundColor = "white";
188
- document.getElementById('my-div-eraser').style.backgroundColor = "gray";
189
- }
190
-
191
- }
192
-
193
- globalThis.toggleEraser = () => {
194
- element = document.getElementById('my-toggle-eraser');
195
- element.classList.toggle('clicked');
196
- // simulate a click on the gradio button
197
- btn_gradio = document.querySelector("#cb-eraser > label > input");
198
- var event = new MouseEvent('click', {
199
- 'view': window,
200
- 'bubbles': true,
201
- 'cancelable': true
202
- });
203
- btn_gradio.dispatchEvent(event);
204
- if (element.classList.contains('clicked')) {
205
- document.getElementById('my-toggle-pencil').classList.remove('clicked');
206
- document.getElementById('my-div-pencil').style.backgroundColor = "white";
207
- document.getElementById('my-div-eraser').style.backgroundColor = "gray";
208
- }
209
- else {
210
- document.getElementById('my-toggle-pencil').classList.add('clicked');
211
- document.getElementById('my-div-pencil').style.backgroundColor = "gray";
212
- document.getElementById('my-div-eraser').style.backgroundColor = "white";
213
- }
214
- }
215
- }
216
- """
217
-
218
- with gr.Blocks(css="style.css") as demo:
219
- #gr.Markdown("# pix2pix-Turbo: **Sketch**", elem_id="description")
220
- #gr.Markdown("**Paper:** [One-Step Image Translation with Text-to-Image Models](https://arxiv.org/abs/2403.12036) ", elem_id="paper_name")
221
- #gr.Markdown("**GitHub:** https://github.com/GaParmar/img2img-turbo", elem_id="github")
222
- gr.HTML(
223
- """
224
- <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
225
- <div>
226
- <h2><a href="https://github.com/GaParmar/img2img-turbo">One-Step Image Translation with Text-to-Image Models</a></h2>
227
- <div>
228
- <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
229
- <a href='https://gauravparmar.com/'>Gaurav Parmar, </a>
230
- &nbsp;
231
- <a href='https://taesung.me/'> Taesung Park,</a>
232
- &nbsp;
233
- <a href='https://www.cs.cmu.edu/~srinivas/'>Srinivasa Narasimhan, </a>
234
- &nbsp;
235
- <a href='https://www.cs.cmu.edu/~junyanz/'> Jun-Yan Zhu </a>
236
- </div>
237
- </div>
238
- </br>
239
- <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
240
- <a href='https://arxiv.org/abs/2403.12036'>
241
- <img src="https://img.shields.io/badge/arXiv-2403.12036-red">
242
- </a>
243
- &nbsp;
244
- <a href='https://github.com/GaParmar/img2img-turbo'>
245
- <img src='https://img.shields.io/badge/github-%23121011.svg'>
246
- </a>
247
- &nbsp;
248
- <a href='https://github.com/GaParmar/img2img-turbo/blob/main/LICENSE'>
249
- <img src='https://img.shields.io/badge/license-MIT-lightgrey'>
250
- </a>
251
- </div>
252
- </div>
253
- </div>
254
- <div>
255
- </br>
256
- </div>
257
- """
258
- )
259
-
260
-
261
- # these are hidden buttons that are used to trigger the canvas changes
262
- line = gr.Checkbox(label="line", value=False, elem_id="cb-line")
263
- eraser = gr.Checkbox(label="eraser", value=False, elem_id="cb-eraser")
264
- with gr.Row(elem_id="main_row"):
265
- with gr.Column(elem_id="column_input"):
266
- gr.Markdown("## INPUT", elem_id="input_header")
267
- image = gr.Image(
268
- source="canvas", tool="color-sketch", type="pil", image_mode="L",
269
- invert_colors=True, shape=(512, 512), brush_radius=4, height=440, width=440,
270
- brush_color="#000000", interactive=True, show_download_button=True, elem_id="input_image", show_label=False)
271
- download_sketch = gr.Button("Download sketch", scale=1, elem_id="download_sketch")
272
-
273
- gr.HTML("""
274
- <div class="button-row">
275
- <div id="my-div-pencil" class="pad2"> <button id="my-toggle-pencil" onclick="return togglePencil(this)"></button> </div>
276
- <div id="my-div-eraser" class="pad2"> <button id="my-toggle-eraser" onclick="return toggleEraser(this)"></button> </div>
277
- <div class="pad2"> <button id="my-button-undo" onclick="return UNDO_SKETCH_FUNCTION(this)"></button> </div>
278
- <div class="pad2"> <button id="my-button-clear" onclick="return DELETE_SKETCH_FUNCTION(this)"></button> </div>
279
- <div class="pad2"> <button href="TODO" download="image" id="my-button-down" onclick='return theSketchDownloadFunction()'></button> </div>
280
- </div>
281
- """)
282
- # gr.Markdown("## Prompt", elem_id="tools_header")
283
- prompt = gr.Textbox(label="Prompt", value="", show_label=True)
284
- with gr.Row():
285
- style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME, scale=1)
286
- prompt_temp = gr.Textbox(label="Prompt Style Template", value=styles[DEFAULT_STYLE_NAME], scale=2, max_lines=1)
287
-
288
- with gr.Row():
289
- val_r = gr.Slider(label="Sketch guidance: ", show_label=True, minimum=0, maximum=1, value=0.4, step=0.01, scale=3)
290
- seed = gr.Textbox(label="Seed", value=42, scale=1, min_width=50)
291
- randomize_seed = gr.Button("Random", scale=1, min_width=50)
292
-
293
- with gr.Column(elem_id="column_process", min_width=50, scale=0.4):
294
- gr.Markdown("## pix2pix-turbo", elem_id="description")
295
- run_button = gr.Button("Run", min_width=50)
296
-
297
- with gr.Column(elem_id="column_output"):
298
- gr.Markdown("## OUTPUT", elem_id="output_header")
299
- result = gr.Image(label="Result", height=440, width=440, elem_id="output_image", show_label=False, show_download_button=True)
300
- download_output = gr.Button("Download output", elem_id="download_output")
301
- gr.Markdown("### Instructions")
302
- gr.Markdown("**1**. Enter a text prompt (e.g. cat)")
303
- gr.Markdown("**2**. Start sketching")
304
- gr.Markdown("**3**. Change the image style using a style template")
305
- gr.Markdown("**4**. Adjust the effect of sketch guidance using the slider")
306
- gr.Markdown("**5**. Try different seeds to generate different results")
307
-
308
-
309
- eraser.change(fn=lambda x: gr.update(value=not x), inputs=[eraser], outputs=[line]).then(update_canvas, [line, eraser], [image])
310
- line.change(fn=lambda x: gr.update(value=not x), inputs=[line], outputs=[eraser]).then(update_canvas, [line, eraser], [image])
311
-
312
- demo.load(None,None,None,_js=scripts)
313
- randomize_seed.click(lambda x: random.randint(0, MAX_SEED), inputs=[], outputs=seed)
314
- inputs = [image, prompt, prompt_temp, style, seed, val_r]
315
- outputs = [result, download_sketch, download_output]
316
- prompt.submit(fn=run, inputs=inputs, outputs=outputs)
317
- style.change(lambda x: styles[x], inputs=[style], outputs=[prompt_temp]).then(
318
- fn=run, inputs=inputs, outputs=outputs,)
319
- val_r.change(run, inputs=inputs, outputs=outputs)
320
- run_button.click(fn=run, inputs=inputs, outputs=outputs)
321
- image.change(run, inputs=inputs, outputs=outputs,)
322
-
323
- if __name__ == "__main__":
324
- demo.queue().launch(debug=True)