Spaces:
Runtime error
Runtime error
Create app.py
Browse files
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)
|