Spaces:
Running
on
A10G
Running
on
A10G
import os | |
import sys | |
import pdb | |
import random | |
import numpy as np | |
from PIL import Image | |
import base64 | |
from io import BytesIO | |
import torch | |
from torchvision import transforms | |
import torchvision.transforms.functional as TF | |
import gradio as gr | |
from src.model import make_1step_sched | |
from src.pix2pix_turbo import Pix2Pix_Turbo | |
model = Pix2Pix_Turbo("sketch_to_image_stochastic") | |
ITEMS_NAMES = [ "π‘ Lamp","π Bag","ποΈ Sofa","πͺ Chair","ποΈ Car","ποΈ Motorbike"] | |
MAX_SEED = np.iinfo(np.int32).max | |
DEFAULT_ITEM_NAME = "π‘ Lamp" | |
def pil_image_to_data_uri(img, format='PNG'): | |
buffered = BytesIO() | |
img.save(buffered, format=format) | |
img_str = base64.b64encode(buffered.getvalue()).decode() | |
return f"data:image/{format.lower()};base64,{img_str}" | |
def run(image, item_name): | |
print("sketch updated") | |
if image is None: | |
ones = Image.new("L", (512, 512), 255) | |
return ones | |
prompt = item_name + " professional 3d model. octane render, highly detailed, volumetric, dramatic lighting" | |
image = image.convert("RGB") | |
image_t = TF.to_tensor(image) > 0.5 | |
with torch.no_grad(): | |
c_t = image_t.unsqueeze(0).cuda().float() | |
torch.manual_seed(42) | |
B,C,H,W = c_t.shape | |
noise = torch.randn((1,4,H//8, W//8), device=c_t.device) | |
output_image = model(c_t, prompt, deterministic=False, r=0.4, noise_map=noise) | |
output_pil = TF.to_pil_image(output_image[0].cpu()*0.5+0.5) | |
return output_pil | |
def update_canvas(use_line, use_eraser): | |
if use_eraser: | |
_color = "#ffffff" | |
brush_size = 20 | |
if use_line: | |
_color = "#000000" | |
brush_size = 4 | |
return gr.update(brush_radius=brush_size, brush_color=_color, interactive=True) | |
def upload_sketch(file): | |
_img = Image.open(file.name) | |
_img = _img.convert("L") | |
return gr.update(value=_img, source="upload", interactive=True) | |
scripts = """ | |
async () => { | |
globalThis.theSketchDownloadFunction = () => { | |
console.log("test") | |
var link = document.createElement("a"); | |
dataUri = document.getElementById('download_sketch').href | |
link.setAttribute("href", dataUri) | |
link.setAttribute("download", "sketch.png") | |
document.body.appendChild(link); // Required for Firefox | |
link.click(); | |
document.body.removeChild(link); // Clean up | |
// also call the output download function | |
theOutputDownloadFunction(); | |
return false | |
} | |
globalThis.theOutputDownloadFunction = () => { | |
console.log("test output download function") | |
var link = document.createElement("a"); | |
dataUri = document.getElementById('download_output').href | |
link.setAttribute("href", dataUri); | |
link.setAttribute("download", "output.png"); | |
document.body.appendChild(link); // Required for Firefox | |
link.click(); | |
document.body.removeChild(link); // Clean up | |
return false | |
} | |
globalThis.DELETE_SKETCH_FUNCTION = () => { | |
console.log("delete sketch function") | |
var button_del = document.querySelector('#input_image > div.image-container.svelte-p3y7hu > div.svelte-s6ybro > button:nth-child(2)'); | |
// Create a new 'click' event | |
var event = new MouseEvent('click', { | |
'view': window, | |
'bubbles': true, | |
'cancelable': true | |
}); | |
button_del.dispatchEvent(event); | |
} | |
globalThis.togglePencil = () => { | |
el_pencil = document.getElementById('my-toggle-pencil'); | |
el_pencil.classList.toggle('clicked'); | |
// simulate a click on the gradio button | |
btn_gradio = document.querySelector("#cb-line > label > input"); | |
var event = new MouseEvent('click', { | |
'view': window, | |
'bubbles': true, | |
'cancelable': true | |
}); | |
btn_gradio.dispatchEvent(event); | |
if (el_pencil.classList.contains('clicked')) { | |
document.getElementById('my-toggle-eraser').classList.remove('clicked'); | |
document.getElementById('my-div-pencil').style.backgroundColor = "gray"; | |
document.getElementById('my-div-eraser').style.backgroundColor = "white"; | |
} | |
else { | |
document.getElementById('my-toggle-eraser').classList.add('clicked'); | |
document.getElementById('my-div-pencil').style.backgroundColor = "white"; | |
document.getElementById('my-div-eraser').style.backgroundColor = "gray"; | |
} | |
} | |
globalThis.toggleEraser = () => { | |
element = document.getElementById('my-toggle-eraser'); | |
element.classList.toggle('clicked'); | |
// simulate a click on the gradio button | |
btn_gradio = document.querySelector("#cb-eraser > label > input"); | |
var event = new MouseEvent('click', { | |
'view': window, | |
'bubbles': true, | |
'cancelable': true | |
}); | |
btn_gradio.dispatchEvent(event); | |
if (element.classList.contains('clicked')) { | |
document.getElementById('my-toggle-pencil').classList.remove('clicked'); | |
document.getElementById('my-div-pencil').style.backgroundColor = "white"; | |
document.getElementById('my-div-eraser').style.backgroundColor = "gray"; | |
} | |
else { | |
document.getElementById('my-toggle-pencil').classList.add('clicked'); | |
document.getElementById('my-div-pencil').style.backgroundColor = "gray"; | |
document.getElementById('my-div-eraser').style.backgroundColor = "white"; | |
} | |
} | |
} | |
""" | |
with gr.Blocks(css="style.css") as demo: | |
gr.HTML("""<div id="header_block"> | |
<img src="file=assets/title.png" alt="title" /> | |
<img src="file=assets/logos.png" alt="title" /> | |
</div>""") | |
with gr.Column(elem_id="main_block"): | |
with gr.Row(elem_id="board_row"): | |
with gr.Group(elem_id="input_image_container", elem_classes="image_container" ): | |
image = gr.Image( | |
source="canvas", tool="color-sketch", type="pil", image_mode="L",container=False, | |
invert_colors=True, shape=(512, 512), brush_radius=4, height="100%", width="100%", | |
brush_color="#000000", interactive=True, show_download_button=False, elem_id="input_image", show_label=False) | |
gr.HTML("""<img src="assets/drawCta.png" id="draw_cta" alt="draw here image" /> | |
<button id="eraser" onclick="return DELETE_SKETCH_FUNCTION(this)"> | |
<span id="eraser_icon"></span> | |
</button>""",elem_id="eraser_container") | |
with gr.Group(elem_id="output_image_container", elem_classes="image_container"): | |
result = gr.Image(label="Result", height="100%", width="100%", elem_id="output_image", show_label=False, show_download_button=True,container=False,) | |
with gr.Row(elem_id="radio_row"): | |
item = gr.Radio(choices=ITEMS_NAMES, value=DEFAULT_ITEM_NAME, show_label=False, container=False) | |
# gr.HTML("""<div id="footer_block"> | |
# <button id="fullscreen" onclick="document.body.querySelector('.contain').requestFullscreen();" class="fullscreen" title="fullscreen"></button> | |
# </div>""") | |
demo.load(None,None,None,_js=scripts) | |
inputs = [image, item] | |
outputs = [result] | |
item.change(fn=run, inputs=inputs, outputs=outputs) | |
image.change(fn=run, inputs=inputs, outputs=outputs,) | |
if __name__ == "__main__": | |
demo.queue().launch(debug=True) | |