File size: 7,520 Bytes
98322fb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5f7150e
98322fb
6eb8bc3
98322fb
 
 
 
 
 
 
 
 
de5fa74
98322fb
 
 
55d75d7
5fadb62
98322fb
 
 
 
ddaa35b
98322fb
 
f872b37
98322fb
5164e6e
98322fb
 
 
 
 
 
 
6eb8bc3
98322fb
 
 
 
 
 
 
 
 
 
 
 
2d27002
84aac8a
b5da494
1e561c3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98322fb
a3c3ebf
98322fb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2ec3d1b
5164e6e
 
 
 
ade53d7
 
 
 
09be82c
5f7150e
ade53d7
 
 
 
a98ece3
ade53d7
09be82c
ade53d7
5f7150e
 
 
 
98322fb
 
fb8291f
ade53d7
ed65a98
1e561c3
98322fb
 
4bf0d68
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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)