Update app.py

app.py

@@ -1,8 +1,7 @@
+import spaces
 import random
 import numpy as np
 from PIL import Image
-import base64
-from io import BytesIO
 
 import torch
 import torchvision.transforms.functional as F
@@ -69,45 +68,40 @@ DEFAULT_STYLE_NAME = "No Style"
 MAX_SEED = np.iinfo(np.int32).max
 
 
-def pil_image_to_data_url(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 randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
 
+@spaces.GPU
 def run(
-    image, 
-    prompt, 
-    prompt_template, 
-    style_name, 
+    image,
+    prompt,
+    prompt_template,
+    style_name,
     controlnet_conditioning_scale,
     device_type="GPU",
-    param_dtype='torch.float16',
+    param_dtype="torch.float16",
 ):
     if device_type == "CPU":
-        device = "cpu" 
-        param_dtype = 'torch.float32'
+        device = "cpu"
+        param_dtype = "torch.float32"
     else:
         device = "cuda"
-    
-    pipe.to(torch_device=device, torch_dtype=torch.float16 if param_dtype == 'torch.float16' else torch.float32)
+
+    pipe.to(
+        torch_device=device,
+        torch_dtype=torch.float16 if param_dtype == "torch.float16" else torch.float32,
+    )
 
     print(f"prompt: {prompt}")
     print("sketch updated")
     if image is None:
         ones = Image.new("L", (512, 512), 255)
-        temp_url = pil_image_to_data_url(ones)
-        return ones, gr.update(link=temp_url), gr.update(link=temp_url)
+        return ones
     prompt = prompt_template.replace("{prompt}", prompt)
-    control_image = image.convert("RGB")
-    control_image = Image.fromarray(255 - np.array(control_image))
+    control_image = Image.fromarray(255 - np.array(image["composite"])[:, :, -1])
 
     output_pil = pipe(
         prompt=prompt,
@@ -121,205 +115,84 @@ def run(
         controlnet_conditioning_scale=float(controlnet_conditioning_scale),
     ).images[0]
 
-    input_sketch_url = pil_image_to_data_url(control_image)
-    output_image_url = pil_image_to_data_url(output_pil)
-    return (
-        output_pil,
-        gr.update(link=input_sketch_url),
-        gr.update(link=output_image_url),
-    )
-
-
-def update_canvas(use_line, use_eraser):
-    if use_eraser:
-        _color = "#ffffff"
-        brush_size = 20
-    if use_line:
-        _color = "#000000"
-        brush_size = 8
-    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)
-
+    return output_pil
 
-scripts = """
-async () => {
-    globalThis.theSketchDownloadFunction = () => {
-        console.log("test")
-        var link = document.createElement("a");
-        dataUrl = document.getElementById('download_sketch').href
-        link.setAttribute("href", dataUrl)
-        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");
-        dataUrl = document.getElementById('download_output').href
-        link.setAttribute("href", dataUrl);
-        link.setAttribute("download", "output.png");
-        document.body.appendChild(link); // Required for Firefox
-        link.click();
-        document.body.removeChild(link); // Clean up
-      return false
-    }
-
-    globalThis.UNDO_SKETCH_FUNCTION = () => {
-        console.log("undo sketch function")
-        var button_undo = document.querySelector('#input_image > div.image-container.svelte-p3y7hu > div.svelte-s6ybro > button:nth-child(1)');
-        // Create a new 'click' event
-        var event = new MouseEvent('click', {
-            'view': window,
-            'bubbles': true,
-            'cancelable': true
-        });
-        button_undo.dispatchEvent(event);
-    }
-
-    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:
+with gr.Blocks() as demo:
     gr.Markdown("# SDXS-512-DreamShaper-Sketch")
     gr.Markdown("[SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions](https://arxiv.org/abs/2403.16627) | [GitHub](https://github.com/IDKiro/sdxs)")
-    # these are hidden buttons that are used to trigger the canvas changes
-    line = gr.Checkbox(label="line", value=False, elem_id="cb-line")
-    eraser = gr.Checkbox(label="eraser", value=False, elem_id="cb-eraser")
     with gr.Row(elem_id="main_row"):
         with gr.Column(elem_id="column_input"):
             gr.Markdown("## INPUT", elem_id="input_header")
-            image = gr.Image(
-                source="canvas", tool="color-sketch", type="pil", image_mode="L",
-                invert_colors=True, shape=(512, 512), brush_radius=8, height=440, width=440,
-                brush_color="#000000", interactive=True, show_download_button=True, elem_id="input_image", show_label=False)
-            download_sketch = gr.Button("Download sketch", scale=1, elem_id="download_sketch")
-            
-            gr.HTML("""
-            <div class="button-row">
-                <div id="my-div-pencil" class="pad2"> <button id="my-toggle-pencil" onclick="return togglePencil(this)"></button> </div>
-                <div id="my-div-eraser" class="pad2"> <button id="my-toggle-eraser" onclick="return toggleEraser(this)"></button> </div>
-                <div class="pad2"> <button id="my-button-undo" onclick="return UNDO_SKETCH_FUNCTION(this)"></button> </div>
-                <div class="pad2"> <button id="my-button-clear" onclick="return DELETE_SKETCH_FUNCTION(this)"></button> </div>
-                <div class="pad2"> <button href="TODO" download="image" id="my-button-down" onclick='return theSketchDownloadFunction()'></button> </div>
-            </div>
-            """)
+            image = gr.Sketchpad(
+                type="pil",
+                image_mode="RGBA",
+                brush=gr.Brush(colors=["#000000"], color_mode="fixed", default_size=8),
+                crop_size=(512, 512),
+            )
+
             # gr.Markdown("## Prompt", elem_id="tools_header")
             prompt = gr.Textbox(label="Prompt", value="", show_label=True)
             with gr.Row():
-                style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME, scale=1)
-                prompt_temp = gr.Textbox(label="Prompt Style Template", value=styles[DEFAULT_STYLE_NAME], scale=2, max_lines=1)
-            
-            controlnet_conditioning_scale = gr.Slider(label="Control Strength", minimum=0, maximum=1, step=0.01, value=0.8)
-
-                 
-            device_choices = ['GPU','CPU']
-            device_type = gr.Radio(device_choices, label='Device',  
-                                        value=device_choices[0],
-                                        interactive=True,
-                                        info='Many thanks to the community for the GPU!')
-            
-            dtype_choices = ['torch.float16','torch.float32']
-            param_dtype = gr.Radio(dtype_choices,label='torch.weight_type',  
-                                        value=dtype_choices[0],
-                                        interactive=True,
-                                        info='To save GPU memory, use torch.float16. For better quality, use torch.float32.')
-                
-
-        with gr.Column(elem_id="column_process", min_width=50, scale=0.4):
-            gr.Markdown("## SDXS-Sketch", elem_id="description")
-            run_button = gr.Button("Run", min_width=50)
+                style = gr.Dropdown(
+                    label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME, scale=1
+                )
+                prompt_temp = gr.Textbox(
+                    label="Prompt Style Template",
+                    value=styles[DEFAULT_STYLE_NAME],
+                    scale=2,
+                    max_lines=1,
+                )
+
+            controlnet_conditioning_scale = gr.Slider(
+                label="Control Strength", minimum=0, maximum=1, step=0.01, value=0.8
+            )
+
+            device_choices = ["GPU", "CPU"]
+            device_type = gr.Radio(
+                device_choices,
+                label="Device",
+                value=device_choices[0],
+                interactive=True,
+                info="Many thanks to the community for the GPU!",
+            )
+
+            dtype_choices = ["torch.float16", "torch.float32"]
+            param_dtype = gr.Radio(
+                dtype_choices,
+                label="torch.weight_type",
+                value=dtype_choices[0],
+                interactive=True,
+                info="To save GPU memory, use torch.float16. For better quality, use torch.float32.",
+            )
 
         with gr.Column(elem_id="column_output"):
             gr.Markdown("## OUTPUT", elem_id="output_header")
-            result = gr.Image(label="Result", height=440, width=440, elem_id="output_image", show_label=False, show_download_button=True)
-            download_output = gr.Button("Download output", elem_id="download_output")
-            gr.Markdown("### Instructions")
-            gr.Markdown("**1**. Enter a text prompt (e.g. cat)")
-            gr.Markdown("**2**. Start sketching")
-            gr.Markdown("**3**. Change the image style using a style template")
-            gr.Markdown("**4**. Adjust the effect of sketch guidance using the slider")
-
-    
-    eraser.change(fn=lambda x: gr.update(value=not x), inputs=[eraser], outputs=[line]).then(update_canvas, [line, eraser], [image])
-    line.change(fn=lambda x: gr.update(value=not x), inputs=[line], outputs=[eraser]).then(update_canvas, [line, eraser], [image])
-
-    demo.load(None,None,None,_js=scripts)
-    inputs = [image, prompt, prompt_temp, style, controlnet_conditioning_scale, device_type, param_dtype]
-    outputs = [result, download_sketch, download_output]
-    prompt.submit(fn=run, inputs=inputs, outputs=outputs)
+            result = gr.Image(
+                label="Result",
+                height=512,
+                width=512,
+                elem_id="output_image",
+                show_label=False,
+                show_download_button=True,
+            )
+
+    inputs = [
+        image,
+        prompt,
+        prompt_temp,
+        style,
+        controlnet_conditioning_scale,
+        device_type,
+        param_dtype,
+    ]
+    outputs = [result]
+
+    prompt.change(fn=run, inputs=inputs, outputs=outputs)
     style.change(lambda x: styles[x], inputs=[style], outputs=[prompt_temp]).then(
-        fn=run, inputs=inputs, outputs=outputs,)
-    run_button.click(fn=run, inputs=inputs, outputs=outputs)
+            fn=run, inputs=inputs, outputs=outputs,)
     image.change(run, inputs=inputs, outputs=outputs,)
+    controlnet_conditioning_scale.change(run, inputs=inputs, outputs=outputs,)
 
 if __name__ == "__main__":
-    demo.queue().launch(debug=True)
+    demo.queue().launch()