Spaces:
Running
Running
There
Browse files
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 |
-
|
231 |
-
<a href='https://taesung.me/'> Taesung Park,</a>
|
232 |
-
|
233 |
-
<a href='https://www.cs.cmu.edu/~srinivas/'>Srinivasa Narasimhan, </a>
|
234 |
-
|
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 |
-
|
244 |
-
<a href='https://github.com/GaParmar/img2img-turbo'>
|
245 |
-
<img src='https://img.shields.io/badge/github-%23121011.svg'>
|
246 |
-
</a>
|
247 |
-
|
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)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|