push app

app.py

@@ -177,18 +177,22 @@ else:
     # Prompt input for inpainting
     prompt = st.text_input("Prompt for inpainting (describe what should replace the selected area):")
     
-    # Only proceed when a point is selected and prompt is provided
-    # if coords and prompt:
-    #     cx, cy = int(coords['x']), int(coords['y'])
-    #     st.write("Generating mask with FastSAM...")
-    #     results = fastsam_model(img, points=[[cx, cy]], labels=[1])
-    #     mask_data = results[0].masks.data[0]
-    #     mask_array = mask_data.cpu().numpy()
-    #     mask_image = Image.fromarray((mask_array * 255).astype(np.uint8))
-        
-    #     st.write("Running Stable Diffusion Inpainting...")
-    #     result = sd_pipe(prompt=prompt, image=img, mask_image=mask_image).images[0]
-        
-    #     # Finally, update the same canvas with the inpainted image
-    #     with canvas.container():
-    #         canvas.image(result, caption="Inpainted Image", use_container_width=True)
+   # --- 4) INPAINTING LOGIC ---
+    prompt = st.text_input("Prompt for inpainting (describe what should replace the selected area):")
+
+    # If there's a prompt and we have at least one mask from the combined points, do inpainting
+    if prompt and combined_mask_img is not None:
+        if st.button("Run Inpainting"):
+            with st.spinner("Inpainting..."):
+                # Run Stable Diffusion Inpainting on the entire combined mask
+                inpainted_img = sd_pipe(
+                    prompt=prompt,
+                    image=img,
+                    mask_image=combined_mask_img
+                ).images[0]
+
+                # Update the session image to the newly inpainted result
+                st.session_state.img = inpainted_img
+                # Optionally reset the points or keep them
+                st.session_state.coords_list = []
+                st.rerun()