Update app.py

app.py

@@ -38,7 +38,7 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            canvas = gr.Sketchpad(label="Draw Here", tool="brush")
+            canvas = gr.Sketchpad(label="Draw Here", shape=(400, 400))  # Updated: Removed 'tool' argument
             clear_btn = gr.Button("Clear")
             generate_btn = gr.Button("Generate Image")
         
@@ -52,9 +52,13 @@ with gr.Blocks() as demo:
         return np.zeros((400, 400, 3), dtype=np.uint8)
 
     clear_btn.click(fn=clear_canvas, inputs=None, outputs=canvas)
-    
+
+    # Assign color to brush for each color button
+    def change_color(color):
+        return color
+
     for color_name, color in color_mapping.items():
-        color_btns[color_name].click(fn=lambda color=color: canvas.update(value=color), inputs=None, outputs=None)
+        color_btns[color_name].click(fn=change_color, inputs=None, outputs=canvas)
 
     # Click to generate an image
     generate_btn.click(fn=predict, inputs=canvas, outputs=output_image)