Update app.py

app.py

@@ -11,21 +11,50 @@ import warnings
 # Suppress specific warnings if desired
 warnings.filterwarnings("ignore", category=UserWarning)
 
-# Load the LangSAM model
-model = LangSAM()  # Use the default model or specify custom checkpoint if necessary
+# Device configuration: Use CUDA if available, else CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Load the LangSAM model and move it to the selected device
+model = LangSAM()
+model.to(device)
+model.eval()  # Set model to evaluation mode
 
 def extract_masks(image_pil, prompts):
+    """
+    Extracts masks for each prompt using the LangSAM model.
+
+    Args:
+        image_pil (PIL.Image): The input image.
+        prompts (str): Comma-separated prompts for segmentation.
+
+    Returns:
+        dict: A dictionary mapping each prompt to its corresponding binary mask.
+    """
     prompts_list = [p.strip() for p in prompts.split(',') if p.strip()]
     masks_dict = {}
-    for prompt in prompts_list:
-        masks, boxes, phrases, logits = model.predict(image_pil, prompt)
-        if masks is not None and len(masks) > 0:
-            masks_np = masks[0].cpu().numpy()
-            mask = (masks_np > 0).astype(np.uint8) * 255  # Binary mask
-            masks_dict[prompt] = mask
+    with torch.no_grad():  # Disable gradient computation for inference
+        for prompt in prompts_list:
+            # Ensure the model uses the correct device
+            masks, boxes, phrases, logits = model.predict(image_pil, prompt)
+            if masks is not None and len(masks) > 0:
+                # Move masks to CPU and convert to numpy
+                masks_np = masks[0].cpu().numpy()
+                mask = (masks_np > 0).astype(np.uint8) * 255  # Binary mask
+                masks_dict[prompt] = mask
     return masks_dict
 
 def apply_color_matching(source_img_np, ref_img_np):
+    """
+    Applies color matching from the reference image to the source image.
+
+    Args:
+        source_img_np (numpy.ndarray): Source image in NumPy array format.
+        ref_img_np (numpy.ndarray): Reference image in NumPy array format.
+
+    Returns:
+        numpy.ndarray: Color-matched image.
+    """
     # Initialize ColorMatcher
     cm = ColorMatcher()
 
@@ -38,6 +67,24 @@ def apply_color_matching(source_img_np, ref_img_np):
     return img_res
 
 def process_image(current_image_pil, selected_prompt, masks_dict, replacement_image_pil, color_ref_image_pil, apply_replacement, apply_color_grading, apply_color_to_full_image, blending_amount, image_history):
+    """
+    Processes the image by applying replacement and/or color grading based on user input.
+
+    Args:
+        current_image_pil (PIL.Image): The current image to be edited.
+        selected_prompt (str): The selected segment prompt.
+        masks_dict (dict): Dictionary of masks for each prompt.
+        replacement_image_pil (PIL.Image): Replacement image (optional).
+        color_ref_image_pil (PIL.Image): Color reference image (optional).
+        apply_replacement (bool): Flag to apply replacement.
+        apply_color_grading (bool): Flag to apply color grading.
+        apply_color_to_full_image (bool): Flag to apply color grading to the full image.
+        blending_amount (int): Amount for blending the mask.
+        image_history (list): History of images for undo functionality.
+
+    Returns:
+        tuple: Updated image, status message, updated history, and image display.
+    """
     # Check if current_image_pil is None
     if current_image_pil is None:
         return None, "No current image to edit.", image_history, None
@@ -119,6 +166,15 @@ def process_image(current_image_pil, selected_prompt, masks_dict, replacement_im
     return current_image_pil, f"Applied changes to '{selected_prompt}'", image_history, current_image_pil
 
 def undo(image_history):
+    """
+    Undoes the last image edit by reverting to the previous image in the history.
+
+    Args:
+        image_history (list): History of images.
+
+    Returns:
+        tuple: Reverted image, updated history, and image display.
+    """
     if image_history and len(image_history) > 1:
         # Pop the last image
         image_history.pop()
@@ -133,6 +189,9 @@ def undo(image_history):
         return None, [], None
 
 def gradio_interface():
+    """
+    Defines and launches the Gradio interface for continuous image editing.
+    """
     with gr.Blocks() as demo:
         # Define the state variables
         image_history = gr.State([])
@@ -160,7 +219,15 @@ def gradio_interface():
                 status = gr.Textbox(lines=2, interactive=False, label="Status")
 
         def initialize_image(initial_image_pil):
-            # Initialize image history with the initial image
+            """
+            Initializes the image history and sets up the initial image.
+
+            Args:
+                initial_image_pil (PIL.Image): The uploaded initial image.
+
+            Returns:
+                tuple: Updated states and status message.
+            """
             if initial_image_pil is not None:
                 image_history = [initial_image_pil]
                 current_image_pil = initial_image_pil
@@ -177,6 +244,16 @@ def gradio_interface():
 
         # Segment button click
         def segment_image_wrapper(current_image_pil, prompts):
+            """
+            Handles the segmentation of the image based on user prompts.
+
+            Args:
+                current_image_pil (PIL.Image): The current image.
+                prompts (str): Comma-separated prompts.
+
+            Returns:
+                tuple: Status message, updated masks, and dropdown updates.
+            """
             if current_image_pil is None:
                 return "No image uploaded.", {}, gr.update(choices=[], value=None)
             masks = extract_masks(current_image_pil, prompts)