Update app.py

app.py

@@ -8,16 +8,16 @@ from transformers import SamModel, SamProcessor
 from diffusers import StableDiffusionInpaintPipeline
 import io
 
-# Initialize SAM model and processor
-sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to("cuda")
+# Initialize SAM model and processor on CPU
+sam_model = SamModel.from_pretrained("facebook/sam-vit-huge", torch_dtype=torch.float32).to("cpu")
 sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
 
-# Initialize Inpainting pipeline
+# Initialize Inpainting pipeline on CPU
 inpaint_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
     "runwayml/stable-diffusion-inpainting",
-    torch_dtype=torch.float16
-).to("cuda")
-inpaint_pipeline.enable_model_cpu_offload()
+    torch_dtype=torch.float32
+).to("cpu")
+# Removed model_cpu_offload as it's unnecessary for CPU
 
 def mask_to_rgba(mask):
     """
@@ -48,7 +48,7 @@ def generate_mask(image, input_points):
     points = [tuple(point) for point in input_points]
     
     # Prepare inputs for SAM
-    inputs = sam_processor(image, points=points, return_tensors="pt").to("cuda")
+    inputs = sam_processor(image, points=points, return_tensors="pt").to("cpu")
     
     with torch.no_grad():
         outputs = sam_model(**inputs)
@@ -91,7 +91,7 @@ def replace_object(image, mask, prompt, negative_prompt, seed, guidance_scale):
     
     mask_image = Image.fromarray((mask * 255).astype(np.uint8))
     
-    generator = torch.Generator("cuda").manual_seed(seed)
+    generator = torch.Generator("cpu").manual_seed(seed)
     
     try:
         result = inpaint_pipeline(
@@ -157,7 +157,7 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(label="Upload Image", type="pil")
+            image_input = gr.Image(label="Upload Image", type="pil", tool="point", interactive=True)
             prompt_input = gr.Textbox(label="Replacement Prompt", placeholder="e.g., a red sports car", lines=2)
             negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="e.g., blurry, low quality", lines=2)
             seed_input = gr.Number(label="Seed", value=42)
@@ -167,11 +167,17 @@ with gr.Blocks() as demo:
             masked_output = gr.Image(label="Selected Object Mask Overlay")
             augmented_output = gr.Image(label="Augmented Image")
     
-    image_input.change(fn=lambda img: img, inputs=image_input, outputs=masked_output)
+    # Capture points selected on the image
+    points = gr.State([])
+    
+    def update_points(selected_points):
+        return selected_points
+    
+    image_input.select(update_points, inputs=image_input, outputs=points)
     
     process_button.click(
         fn=process,
-        inputs=[image_input, gr.State(), prompt_input, negative_prompt_input, seed_input, guidance_scale_input],
+        inputs=[image_input, points, prompt_input, negative_prompt_input, seed_input, guidance_scale_input],
         outputs=[masked_output, augmented_output]
     )