File size: 8,842 Bytes
98322fb
 
 
 
 
6b0b623
98322fb
 
 
 
 
 
 
 
 
 
 
 
96a862d
4321da9
98322fb
6eb8bc3
1ffaa25
c6dc013
1ffaa25
f6f9541
365e727
f6f9541
98322fb
d142887
 
c114545
d142887
f6f9541
98322fb
 
 
 
 
 
 
c6dc013
98322fb
f9a9c3c
b7ab832
804bc5f
b7ab832
 
c6b02d8
55d75d7
478c67d
c6dc013
4f3b924
 
 
98322fb
 
ddaa35b
98322fb
 
bbd66f8
98322fb
5164e6e
98322fb
 
 
 
 
 
 
6eb8bc3
98322fb
 
 
 
 
 
 
 
 
 
 
 
2d27002
84aac8a
b5da494
1e561c3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98322fb
a3c3ebf
98322fb
 
 
9274913
98322fb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
015e494
98322fb
b35d3a4
5164e6e
94b2733
77282db
f26caf9
77282db
 
 
 
 
5164e6e
ade53d7
 
 
ef10226
 
 
039454e
5f5c0dd
ade53d7
a98ece3
ade53d7
09be82c
c6dc013
cd6961d
f4ab459
98322fb
9234bc4
c6dc013
ade53d7
ed65a98
c6dc013
5caec7a
e47f8b2
98322fb
9f0d446
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
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
import os
import sys
import pdb
import random
import numpy as np
from PIL import Image, ImageOps, ImageChops
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"
COLORS = ["blue", "red", "black", "yellow", "green", "pink", "grey"]
DEFAULT_COLOR = "yellow"
def empty_input_image():
    return { 'background': Image.new("L", (512, 512), 255),
             'layers': [Image.new("L", (512, 512), 255),Image.new("L", (512, 512), 255)],
             'composite': Image.new("L", (512, 512), 255)}

def clear_sketchbox(color_name):
    temp_color_list = COLORS.copy()
    temp_color_list.remove(color_name)
    return [empty_input_image(), random.choice(temp_color_list)]
    
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, color_name):
    print("sketch updated")
    print(image)
    empty_image = Image.new("L", (512, 512), 255)
    diff = ImageChops.difference(image["composite"], empty_image)
    # if image["composite"] is None:
    if not diff.getbbox():
        ones = empty_image
        return ones
    print(item_name.split()[1])
    prompt = color_name + " " + item_name.split()[1] + " professional 3d model. octane render, highly detailed, volumetric, dramatic lighting"
    inverted_image = ImageOps.invert(image["composite"])
    converted_image = inverted_image.convert("RGB")
    image_t = TF.to_tensor(converted_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.6, 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-1sbaaot > div.controls-wrap.svelte-4lttvb > div > button:nth-child(3)');
        // 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";
        }
    }
}
"""
head="""<meta name="theme-color" content="#000"><link href="https://fonts.cdnfonts.com/css/pp-neue-montreal" rel="stylesheet">"""

with gr.Blocks(css="style.css", head = head) as demo:
    gr.HTML("""<div id="header_block">
      <h1>Dai forma al nuovo<br />design Made in Italy</h1>
      <div id="logos_block">
          <img id="logos_row" src="file=assets/logos.png" alt="logo" />
        <div id="text_row">
          <span>krnl.ai</span><span>//</span
          ><span>eccellenza-italiana.com</span>
        </div>
      </div>
    </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.Sketchpad(type="pil", image_mode="L", elem_id="input_image",value = empty_input_image, 
                    container=False, height="100%", width="100%",  brush = gr.Brush(default_size=3, colors=["#000000"], color_mode="fixed"), layers = False,
                    interactive=True, show_download_button=True, show_label=False)
                gr.HTML("""<img src="file=assets/drawCta.png" id="draw_cta" alt="draw here image" />""",elem_id="draw_cta_container")
                gr.HTML("""<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,)
                color = gr.Radio(choices=COLORS, value=DEFAULT_COLOR, show_label=False, container=False, visible=False)
    with gr.Row(elem_id="radio_row"):
        item = gr.Radio(choices=ITEMS_NAMES, value=DEFAULT_ITEM_NAME, show_label=False, container=False)

    demo.load(None,None,None,js=scripts)
    inputs = [image, item, color]
    outputs = [result]
    item.change(fn=run, inputs=inputs, outputs=outputs)
    color.change(fn=run, inputs=inputs, outputs=outputs)
    image.change(fn=run, inputs=inputs, outputs=outputs, trigger_mode="always_last")
    image.clear(fn=clear_sketchbox, inputs=color, outputs=[image , color])
if __name__ == "__main__":
    demo.queue().launch(debug=True, allowed_paths=["."])