Update app.py

app.py

@@ -1,65 +1,36 @@
-# app.py
-
 import gradio as gr
+import numpy as np
 from PIL import Image, ImageDraw
 import torch
-import numpy as np
 from transformers import SamModel, SamProcessor
 from diffusers import StableDiffusionInpaintPipeline
 
 # Constants
 IMG_SIZE = 512
 
-# 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 on CPU with a compatible model
-inpaint_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-inpainting",
-    torch_dtype=torch.float32
-).to("cpu")
-# No need for model_cpu_offload on CPU
-
 # Global variables to store points and the original image
 input_points = []
 input_image = None
 
-def mask_to_rgba(mask):
+def generate_mask(image, points):
     """
-    Converts a binary mask to an RGBA image for visualization.
+    Generates a mask using SAM based on input points.
     """
-    bg_transparent = np.zeros(mask.shape + (4,), dtype=np.uint8)
-    bg_transparent[mask == 1] = [0, 255, 0, 127]  # Green with transparency
-    return bg_transparent
-
-def generate_mask(image, input_points):
-    """
-    Generates a binary mask using SAM based on input points.
-    
-    Args:
-        image (PIL.Image): The input image.
-        input_points (list of lists): List of points selected by the user.
-        
-    Returns:
-        np.ndarray: Binary mask where the object is marked with 1s.
-    """
-    if not input_points:
+    if not points:
         return None
-    
-    # Convert image to RGB if not already
+
     image = image.convert("RGB")
+    points = [tuple(point) for point in points]
     
-    # Flatten the list of points
-    points = [tuple(point) for point in input_points]
+    # 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")
     
-    # Prepare inputs for SAM
     inputs = sam_processor(image, points=points, return_tensors="pt").to("cpu")
     
     with torch.no_grad():
         outputs = sam_model(**inputs)
     
-    # Post-process masks
     masks = sam_processor.image_processor.post_process_masks(
         outputs.pred_masks.cpu(),
         inputs["original_sizes"].cpu(),
@@ -69,32 +40,24 @@ def generate_mask(image, input_points):
     if len(masks) == 0:
         return None
     
-    # Select the mask with the highest IoU score
     best_mask = masks[0][0][outputs.iou_scores.argmax()]
-    
-    # Invert mask: object=1, background=0
     binary_mask = ~best_mask.numpy().astype(bool).astype(int)
     
     return binary_mask
 
 def replace_object(image, mask, prompt, negative_prompt, seed, guidance_scale):
     """
-    Replaces the selected object in the image based on the prompt.
-    
-    Args:
-        image (PIL.Image): The original image.
-        mask (np.ndarray): Binary mask of the selected object.
-        prompt (str): Text prompt describing the replacement.
-        negative_prompt (str): Negative text prompt to refine generation.
-        seed (int): Random seed for reproducibility.
-        guidance_scale (float): Guidance scale for the inpainting model.
-        
-    Returns:
-        PIL.Image: The augmented image with the object replaced.
+    Replaces the object in the image based on the mask and prompt.
     """
     if mask is None:
         return image
     
+    # Initialize Inpainting pipeline on CPU with a compatible model
+    inpaint_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-inpainting",
+        torch_dtype=torch.float32
+    ).to("cpu")
+    
     mask_image = Image.fromarray((mask * 255).astype(np.uint8))
     
     generator = torch.Generator("cpu").manual_seed(seed)
@@ -116,46 +79,33 @@ def replace_object(image, mask, prompt, negative_prompt, seed, guidance_scale):
 def visualize_mask(image, mask):
     """
     Overlays the mask on the image for visualization.
-    
-    Args:
-        image (PIL.Image): The original image.
-        mask (np.ndarray): Binary mask of the selected object.
-        
-    Returns:
-        PIL.Image: Image with mask overlay.
     """
     if mask is None:
         return image
     
-    mask_rgba = mask_to_rgba(mask)
-    mask_pil = Image.fromarray(mask_rgba)
-    overlay = Image.alpha_composite(image.convert("RGBA"), mask_pil)
+    bg_transparent = np.zeros(mask.shape + (4,), dtype=np.uint8)
+    bg_transparent[mask == 1] = [0, 255, 0, 127]  # Green with transparency
+    mask_rgba = Image.fromarray(bg_transparent)
+    
+    overlay = Image.alpha_composite(image.convert("RGBA"), mask_rgba)
     return overlay.convert("RGB")
 
 def get_points(img, evt: gr.SelectData):
     """
     Captures points selected by the user on the image.
-    
-    Args:
-        img (PIL.Image): The uploaded image.
-        evt (gr.SelectData): Event data containing the point coordinates.
-        
-    Returns:
-        Tuple: (Updated mask visualization, Updated image with crossmarks)
     """
     global input_points
     global input_image
     
-    # The first time this is called, save the untouched input image
     if len(input_points) == 0:
         input_image = img.copy()
     
     x = evt.index[0]
     y = evt.index[1]
-
+    
     input_points.append([x, y])
     
-    # Run SAM to generate mask
+    # Generate mask based on selected points
     mask = generate_mask(input_image, input_points)
     
     # Mark selected points with a green crossmark
@@ -174,16 +124,6 @@ def get_points(img, evt: gr.SelectData):
 def run_inpaint(prompt, negative_prompt, cfg, seed, invert):
     """
     Runs the inpainting process based on user inputs.
-    
-    Args:
-        prompt (str): Prompt for infill.
-        negative_prompt (str): Negative prompt.
-        cfg (float): Classifier-Free Guidance Scale.
-        seed (int): Random seed.
-        invert (bool): Whether to infill the subject instead of the background.
-        
-    Returns:
-        PIL.Image: The inpainted image.
     """
     global input_image
     global input_points
@@ -209,9 +149,6 @@ def run_inpaint(prompt, negative_prompt, cfg, seed, invert):
 def reset_points_func():
     """
     Resets the selected points and the input image.
-    
-    Returns:
-        Tuple: (Reset mask visualization, Reset image, Empty inpainted image)
     """
     global input_points
     global input_image
@@ -222,19 +159,11 @@ def reset_points_func():
 def preprocess(input_img):
     """
     Preprocesses the uploaded image to ensure it is square and resized.
-    
-    Args:
-        input_img (PIL.Image): The uploaded image.
-        
-    Returns:
-        PIL.Image: The preprocessed image.
     """
     if input_img is None:
         return None
     
-    # Make sure the image is square
     width, height = input_img.size
-    
     if width != height:
         # Add white padding to make the image square
         new_size = max(width, height)
@@ -246,104 +175,142 @@ def preprocess(input_img):
     
     return input_img.resize((IMG_SIZE, IMG_SIZE))
 
-def build_app(get_processed_inputs, inpaint):
-    """
-    Builds and launches the Gradio app.
-    
-    Args:
-        get_processed_inputs (function): Function to process inputs for SAM.
-        inpaint (function): Function to perform inpainting.
-        
-    Returns:
-        None
-    """
-    with gr.Blocks() as demo:
+with gr.Blocks() as demo:
 
-        gr.Markdown(
+    gr.Markdown(
+    """
+    # Object Replacement App
+    Upload an image, select points on the object you want to replace, provide a text prompt for the replacement, and view the augmented image.
+    
+    **Instructions:**
+    1. **Upload Image:** Click on the first image box to upload your image.
+    2. **Select Points:** Click on the image to select points on the object you wish to replace. Use multiple points for better mask accuracy.
+    3. **Enter Prompts:** Provide a replacement prompt and optionally a negative prompt to refine the output.
+    4. **Adjust Settings:** Set the seed for reproducibility and adjust the guidance scale as needed.
+    5. **Replace Object:** Click the "Replace Object" button to generate the augmented image.
+    6. **Reset:** Click the "Reset" button to clear selections and start over.
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            # Image upload and point selection
+            upload_image = gr.Image(
+                label="Upload Image", 
+                type="pil", 
+                interactive=True,
+                height=IMG_SIZE,
+                width=IMG_SIZE
+            )
+            mask_visualization = gr.Image(
+                label="Selected Object Mask Overlay", 
+                interactive=False,
+                height=IMG_SIZE,
+                width=IMG_SIZE
+            )
+            selected_image = gr.Image(
+                label="Image with Selected Points", 
+                type="pil", 
+                interactive=False,
+                height=IMG_SIZE,
+                width=IMG_SIZE,
+            )
+            
+            # Capture points using the select event
+            upload_image.select(get_points, inputs=[upload_image], outputs=[mask_visualization, selected_image])
+            
+            # Preprocess image on change
+            upload_image.change(preprocess, inputs=[upload_image], outputs=[upload_image])
+            
+            # Text inputs and settings
+            prompt = gr.Textbox(
+                label="Replacement Prompt", 
+                placeholder="e.g., a red sports car", 
+                lines=2
+            )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt", 
+                placeholder="e.g., blurry, low quality", 
+                lines=2
+            )
+            cfg = gr.Slider(
+                label="Classifier-Free Guidance Scale", 
+                minimum=1.0, 
+                maximum=20.0, 
+                value=7.5, 
+                step=0.5
+            )
+            seed = gr.Number(
+                label="Seed", 
+                value=42, 
+                precision=0
+            )
+            invert = gr.Checkbox(
+                label="Infill subject instead of background"
+            )
+            
+            # Buttons
+            replace_button = gr.Button("Replace Object")
+            reset_button = gr.Button("Reset")
+        with gr.Column():
+            # Output images
+            augmented_image = gr.Image(
+                label="Augmented Image", 
+                type="pil", 
+                interactive=False,
+                height=IMG_SIZE,
+                width=IMG_SIZE,
+            )
+    
+    # Define button actions
+    replace_button.click(
+        fn=run_inpaint,
+        inputs=[prompt, negative_prompt, cfg, seed, invert],
+        outputs=[augmented_image]
+    )
+    
+    reset_button.click(
+        fn=reset_points_func,
+        inputs=[],
+        outputs=[mask_visualization, selected_image, augmented_image]
+    )
+    
+    # Examples (optional)
+    gr.Markdown(
         """
-        # Object Replacement App
-        Upload an image, select points on the object you want to replace, provide a text prompt for the replacement, and view the augmented image.
-        
-        **Instructions:**
-        1. **Upload Image:** Click on the first image box to upload your image.
-        2. **Select Points:** Click on the image to select points on the object you wish to replace. Use multiple points for better mask accuracy.
-        3. **Enter Prompts:** Provide a replacement prompt and optionally a negative prompt to refine the output.
-        4. **Adjust Settings:** Set the seed for reproducibility and adjust the guidance scale as needed.
-        5. **Replace Object:** Click the "Replace Object" button to generate the augmented image.
-        6. **Reset:** Click the "Reset" button to clear selections and start over.
+        ## EXAMPLES
+        Click on an example to load it. Then, follow the instructions above.
         """)
-        
-        with gr.Row():
-            with gr.Column():
-                # Image upload and point selection
-                upload_image = gr.Image(label="Upload Image", type="pil", interactive=True)
-                mask_visualization = gr.Image(label="Selected Object Mask Overlay", interactive=False)
-                selected_image = gr.Image(label="Image with Selected Points", type="pil", interactive=False)
-                
-                # Capture points using the select event
-                upload_image.select(get_points, inputs=[upload_image], outputs=[mask_visualization, selected_image])
-                
-                # Preprocess image on change
-                upload_image.change(preprocess, inputs=[upload_image], outputs=[upload_image])
-                
-                # Text inputs and settings
-                prompt = gr.Textbox(label="Replacement Prompt", placeholder="e.g., a red sports car", lines=2)
-                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="e.g., blurry, low quality", lines=2)
-                cfg = gr.Slider(
-                    label="Classifier-Free Guidance Scale", 
-                    minimum=1.0, 
-                    maximum=20.0, 
-                    value=7.5, 
-                    step=0.5
-                )
-                seed = gr.Number(label="Seed", value=42, precision=0)
-                invert = gr.Checkbox(label="Infill subject instead of background")
-                
-                # Buttons
-                replace_button = gr.Button("Replace Object")
-                reset_button = gr.Button("Reset")
-            with gr.Column():
-                # Output images
-                augmented_image = gr.Image(label="Augmented Image", type="pil", interactive=False)
-        
-        # Define button actions
-        replace_button.click(
-            fn=run_inpaint,
-            inputs=[prompt, negative_prompt, cfg, seed, invert],
-            outputs=[augmented_image]
-        )
-        
-        reset_button.click(
-            fn=reset_points_func,
-            inputs=[],
-            outputs=[mask_visualization, selected_image, augmented_image]
-        )
-        
-        # Examples (optional)
-        gr.Markdown(
-            """
-            ## EXAMPLES
-            Click on an example to load it. Then, follow the instructions above.
-            """)
-        
-        with gr.Row():
-            examples = gr.Examples(
-                examples=[
-                    ["car.png", "a red sports car", "blurry, low quality", 42],
-                    ["house.jpg", "a modern villa", "dark, overexposed", 123],
-                    ["tree.png", "a blooming cherry tree", "underexposed, low contrast", 999]
+    
+    with gr.Row():
+        examples = gr.Examples(
+            examples=[
+                [
+                    "car.png", 
+                    "a red sports car", 
+                    "blurry, low quality", 
+                    42
                 ],
-                inputs=[
-                    upload_image,
-                    prompt,
-                    negative_prompt,
-                    seed
+                [
+                    "house.jpg", 
+                    "a modern villa", 
+                    "dark, overexposed", 
+                    123
                 ],
-                label="Click to load examples",
-                cache_examples=True
-            )
-        
-    demo.queue(max_size=10).launch()
-
-# Launch the app
-build_app(None, None)
+                [
+                    "tree.png", 
+                    "a blooming cherry tree", 
+                    "underexposed, low contrast", 
+                    999
+                ]
+            ],
+            inputs=[
+                upload_image,
+                prompt,
+                negative_prompt,
+                seed
+            ],
+            label="Click to load examples",
+            cache_examples=False  # Set to False to avoid the error
+        )
+    
+demo.queue(max_size=10).launch()