Update app.py

app.py

@@ -32,13 +32,13 @@ def pil_image_to_data_uri(img, format='PNG'):
 def run(image, item_name):
     print("sketch updated")
     print(image)
-    if image["layers"][0] is None:
+    if image["composite"] is None:
         ones = Image.new("L", (512, 512), 255)
         return ones
     prompt = item_name + " professional 3d model. octane render, highly detailed, volumetric, dramatic lighting"
-    inverted_image = ImageOps.invert(image["layers"][0])
-    image = inverted_image.convert("RGB")
-    image_t = TF.to_tensor(image) > 0.5
+    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)
@@ -177,7 +177,7 @@ with gr.Blocks(css="style.css") as demo:
             with gr.Group(elem_id="input_image_container", elem_classes="image_container" ):
                 image = gr.Sketchpad(type="pil", image_mode="L",container=False, height="100%", width="100%",
                     # invert_colors=True, shape=(512, 512), brush_radius=4, 
-                    interactive=True, show_download_button=False, elem_id="input_image", show_label=False)
+                    interactive=True, show_download_button=True, elem_id="input_image", 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>