Update app.py

app.py

@@ -1,280 +1,280 @@
-import gradio as gr
-import torch
-from PIL import Image
-import numpy as np
-import time
-import os
-
-# --- Import Custom Pipelines ---
-# Ensure these files are in the same directory or accessible in PYTHONPATH
-try:
-    from gen2seg_sd_pipeline import gen2segSDPipeline
-    from gen2seg_mae_pipeline import gen2segMAEInstancePipeline
-except ImportError as e:
-    print(f"Error importing pipeline modules: {e}")
-    print("Please ensure gen2seg_sd_pipeline.py and gen2seg_mae_pipeline.py are in the same directory.")
-    # Optionally, raise an error or exit if pipelines are critical at startup
-    # raise ImportError("Could not import custom pipeline modules. Check file paths.") from e
-
-from transformers import ViTMAEForPreTraining, AutoImageProcessor
-
-# --- Configuration ---
-MODEL_IDS = {
-    "SD": "reachomk/gen2seg-sd",
-    "MAE-H": "reachomk/gen2seg-mae-h"
-}
-
-# Check if a GPU is available and set the device accordingly
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {DEVICE}")
-
-# --- Global Variables for Caching Pipelines ---
-sd_pipe_global = None
-mae_pipe_global = None
-
-# --- Model Loading Functions ---
-def get_sd_pipeline():
-    """Loads and caches the gen2seg Stable Diffusion pipeline."""
-    global sd_pipe_global
-    if sd_pipe_global is None:
-        model_id_sd = MODEL_IDS["SD"]
-        print(f"Attempting to load SD pipeline from Hugging Face Hub: {model_id_sd}")
-        try:
-            sd_pipe_global = gen2segSDPipeline.from_pretrained(
-                model_id_sd,
-                use_safetensors=True,
-                # torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, # Optional: use float16 on GPU
-            ).to(DEVICE)
-            print(f"SD Pipeline loaded successfully from {model_id_sd} on {DEVICE}.")
-        except Exception as e:
-            print(f"Error loading SD pipeline from Hugging Face Hub ({model_id_sd}): {e}")
-            sd_pipe_global = None # Ensure it remains None on failure
-            # Do not raise gr.Error here; let the main function handle it.
-    return sd_pipe_global
-
-def get_mae_pipeline():
-    """Loads and caches the gen2seg MAE-H pipeline."""
-    global mae_pipe_global
-    if mae_pipe_global is None:
-        model_id_mae = MODEL_IDS["MAE-H"]
-        print(f"Loading MAE-H pipeline with model {model_id_mae} on {DEVICE}...")
-        try:
-            model = ViTMAEForPreTraining.from_pretrained(model_id_mae)
-            model.to(DEVICE)
-            model.eval() # Set to evaluation mode
-
-            # Load the official MAE-H image processor
-            # Using "facebook/vit-mae-huge" as per the original app_mae.py
-            image_processor = AutoImageProcessor.from_pretrained("facebook/vit-mae-huge")
-            
-            mae_pipe_global = gen2segMAEInstancePipeline(model=model, image_processor=image_processor)
-            # The custom MAE pipeline's model is already on the DEVICE.
-            print(f"MAE-H Pipeline with model {model_id_mae} loaded successfully on {DEVICE}.")
-        except Exception as e:
-            print(f"Error loading MAE-H model or pipeline from Hugging Face Hub ({model_id_mae}): {e}")
-            mae_pipe_global = None # Ensure it remains None on failure
-            # Do not raise gr.Error here; let the main function handle it.
-    return mae_pipe_global
-
-# --- Unified Prediction Function ---
-def segment_image(input_image: Image.Image, model_choice: str) -> Image.Image:
-    """
-    Takes a PIL Image and model choice, performs segmentation, and returns the segmented image.
-    """
-    if input_image is None:
-        raise gr.Error("No image provided. Please upload an image.")
-
-    print(f"Model selected: {model_choice}")
-    # Ensure image is in RGB format
-    image_rgb = input_image.convert("RGB")
-    original_resolution = image_rgb.size # (width, height)
-    seg_array = None
-
-    try:
-        if model_choice == "SD":
-            pipe_sd = get_sd_pipeline()
-            if pipe_sd is None:
-                raise gr.Error("The SD segmentation pipeline could not be loaded. "
-                               "Please check the Space logs for more details, or try again later.")
-
-            print(f"Running SD inference with image size: {image_rgb.size}")
-            start_time = time.time()
-            with torch.no_grad():
-                # The gen2segSDPipeline expects a single image or a list
-                # The pipeline's __call__ method handles preprocessing internally
-                seg_output = pipe_sd(image_rgb, match_input_resolution=False).prediction # Output is before resize
-                
-                # seg_output is expected to be a numpy array (N,H,W,1) or (N,1,H,W) or tensor
-                # Based on gen2seg_sd_pipeline.py, if output_type="np" (default), it's [N,H,W,1]
-                # If output_type="pt", it's [N,1,H,W]
-                # The original app_sd.py converted tensor to numpy and squeezed.
-                if isinstance(seg_output, torch.Tensor):
-                    seg_output = seg_output.cpu().numpy()
-
-                if seg_output.ndim == 4 and seg_output.shape[0] == 1: # Batch size 1
-                    if seg_output.shape[1] == 1: # Grayscale, (1, 1, H, W)
-                        seg_array = seg_output.squeeze(0).squeeze(0).astype(np.uint8)
-                    elif seg_output.shape[-1] == 1: # Grayscale, (1, H, W, 1)
-                         seg_array = seg_output.squeeze(0).squeeze(-1).astype(np.uint8)
-                    elif seg_output.shape[1] == 3: # RGB, (1, 3, H, W) -> (H, W, 3)
-                        seg_array = np.transpose(seg_output.squeeze(0), (1, 2, 0)).astype(np.uint8)
-                    elif seg_output.shape[-1] == 3: # RGB, (1, H, W, 3)
-                         seg_array = seg_output.squeeze(0).astype(np.uint8)
-                    else: # Fallback for unexpected shapes
-                        seg_array = seg_output.squeeze().astype(np.uint8)
-
-                elif seg_output.ndim == 3: # (H, W, C) or (C, H, W)
-                    seg_array = seg_output.astype(np.uint8)
-                elif seg_output.ndim == 2: # (H,W)
-                    seg_array = seg_output.astype(np.uint8)
-                else:
-                    raise TypeError(f"Unexpected SD segmentation output type/shape: {type(seg_output)}, {seg_output.shape}")
-            end_time = time.time()
-            print(f"SD Inference completed in {end_time - start_time:.2f} seconds.")
-
-
-        elif model_choice == "MAE-H":
-            pipe_mae = get_mae_pipeline()
-            if pipe_mae is None:
-                raise gr.Error("The MAE-H segmentation pipeline could not be loaded. "
-                               "Please check the Space logs for more details, or try again later.")
-
-            print(f"Running MAE-H inference with image size: {image_rgb.size}")
-            start_time = time.time()
-            with torch.no_grad():
-                # The gen2segMAEInstancePipeline expects a list of images
-                # output_type="np" returns a NumPy array
-                pipe_output = pipe_mae([image_rgb], output_type="np")
-                # Prediction is (batch_size, height, width, 3) for MAE
-                prediction_np = pipe_output.prediction[0] # Get the first (and only) image prediction
-            
-            end_time = time.time()
-            print(f"MAE-H Inference completed in {end_time - start_time:.2f} seconds.")
-            
-            if not isinstance(prediction_np, np.ndarray):
-                # This case should ideally not be reached if output_type="np"
-                prediction_np = prediction_np.cpu().numpy()
-
-            # Ensure it's in the expected (H, W, C) format and uint8
-            if prediction_np.ndim == 3 and prediction_np.shape[-1] == 3: # Expected (H, W, 3)
-                seg_array = prediction_np.astype(np.uint8)
-            else:
-                # Attempt to handle other shapes if necessary, or raise error
-                raise gr.Error(f"Unexpected MAE-H prediction shape: {prediction_np.shape}. Expected (H, W, 3).")
-            
-            # The MAE pipeline already does gamma correction and scaling to 0-255.
-            # It also ensures 3 channels.
-
-        else:
-            raise gr.Error(f"Invalid model choice: {model_choice}. Please select a valid model.")
-
-        if seg_array is None:
-             raise gr.Error("Segmentation array was not generated. An unknown error occurred.")
-
-        print(f"Segmentation array generated with shape: {seg_array.shape}, dtype: {seg_array.dtype}")
-
-        # Convert numpy array to PIL Image
-        # Handle grayscale or RGB based on seg_array channels
-        if seg_array.ndim == 2: # Grayscale
-            segmented_image_pil = Image.fromarray(seg_array, mode='L')
-        elif seg_array.ndim == 3 and seg_array.shape[-1] == 3: # RGB
-            segmented_image_pil = Image.fromarray(seg_array, mode='RGB')
-        elif seg_array.ndim == 3 and seg_array.shape[-1] == 1: # Grayscale with channel dim
-            segmented_image_pil = Image.fromarray(seg_array.squeeze(-1), mode='L')
-        else:
-            raise gr.Error(f"Cannot convert seg_array with shape {seg_array.shape} to PIL Image.")
-
-        # Resize back to original image resolution using LANCZOS for high quality
-        segmented_image_pil = segmented_image_pil.resize(original_resolution, Image.Resampling.LANCZOS)
-        
-        print(f"Segmented image processed. Output size: {segmented_image_pil.size}, mode: {segmented_image_pil.mode}")
-        return segmented_image_pil
-
-    except Exception as e:
-        print(f"Error during segmentation with {model_choice}: {e}")
-        # Re-raise as gr.Error for Gradio to display, if not already one
-        if not isinstance(e, gr.Error):
-            # It's often helpful to include the type of the original exception
-            error_type = type(e).__name__
-            raise gr.Error(f"An error occurred during segmentation: {error_type} - {str(e)}")
-        else:
-            raise e # Re-raise if it's already a gr.Error
-
-# --- Gradio Interface ---
-title = "gen2seg: Generative Models Enable Generalizable Instance Segmentation Demo (SD & MAE-H)"
-description = f"""
-<div style="text-align: center; font-family: 'Arial', sans-serif;">
-    <p>Upload an image and choose a model architecture to see the instance segmentation result generated by the respective model. </p>
-    <p>
-        Currently, inference is running on CPU.
-        Performance will be significantly better on GPU.
-    </p>
-    <ul>
-        <li><strong>SD</strong>: Based on Stable Diffusion 2.
-            <a href="https://huggingface.co/{MODEL_IDS['SD']}" target="_blank">Model Link</a>.
-            <em>Approx. CPU inference time: ~1-2 minutes per image.</em>
-        </li>
-        <li><strong>MAE-H</strong>: Based on Masked Autoencoder (Huge).
-            <a href="https://huggingface.co/{MODEL_IDS['MAE-H']}" target="_blank">Model Link</a>.
-            <em>Approx. CPU inference time: ~15-45 seconds per image.</em>
-            If you experience tokenizer artifacts or very dark images, you can use gamma correction to handle this. 
-        </li>
-    </ul>
-    <p>
-        For faster inference, please check out our GitHub to run the models locally on a GPU:
-        <a href="https://github.com/UCDvision/gen2seg" target="_blank">https://github.com/UCDvision/gen2seg</a>
-    </p>
-    <p>If the demo experiences issues, please open an issue on our GitHub.</p>
-    <p> If you have not already, please see our webpage at <a href="https://reachomk.github.io/gen2seg" target="_blank">https://reachomk.github.io/gen2seg</a>
-
-</div>
-"""
-
-article = """
-"""
-
-# Define Gradio inputs
-input_image_component = gr.Image(type="pil", label="Input Image")
-model_choice_component = gr.Dropdown(
-    choices=list(MODEL_IDS.keys()),
-    value="SD",  # Default model
-    label="Choose Segmentation Model Architecture"
-)
-
-# Define Gradio output
-output_image_component = gr.Image(type="pil", label="Segmented Image")
-
-# Example images (ensure these paths are correct if you upload examples to your Space)
-# For example, if you create an "examples" folder in your Space repo:
-# example_paths = [
-#     os.path.join("examples", "example1.jpg"),
-#     os.path.join("examples", "example2.png")
-# ]
-# Filter out non-existent example files to prevent errors
-# example_paths = [ex for ex in example_paths if os.path.exists(ex)]
-example_paths = [] # Add paths to example images here if you have them
-
-iface = gr.Interface(
-    fn=segment_image,
-    inputs=[input_image_component, model_choice_component],
-    outputs=output_image_component,
-    title=title,
-    description=description,
-    article=article,
-    examples=example_paths if example_paths else None, # Pass None if no examples
-    allow_flagging="never",
-    theme=gr.themes.Soft() # Using a soft theme for a slightly modern look
-)
-
-if __name__ == "__main__":
-    # Optional: Pre-load a default model on startup if desired.
-    # This can make the first inference faster but increases startup time.
-    # print("Attempting to pre-load default SD model on startup...")
-    # try:
-    #    get_sd_pipeline() # Pre-load the default SD model
-    #    print("Default SD model pre-loaded successfully or was already cached.")
-    # except Exception as e:
-    #    print(f"Could not pre-load default SD model: {e}")
-
-    print("Launching Gradio interface...")
-    iface.launch()
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+import time
+import os
+
+# --- Import Custom Pipelines ---
+# Ensure these files are in the same directory or accessible in PYTHONPATH
+try:
+    from gen2seg_sd_pipeline import gen2segSDPipeline
+    from gen2seg_mae_pipeline import gen2segMAEInstancePipeline
+except ImportError as e:
+    print(f"Error importing pipeline modules: {e}")
+    print("Please ensure gen2seg_sd_pipeline.py and gen2seg_mae_pipeline.py are in the same directory.")
+    # Optionally, raise an error or exit if pipelines are critical at startup
+    # raise ImportError("Could not import custom pipeline modules. Check file paths.") from e
+
+from transformers import ViTMAEForPreTraining, AutoImageProcessor
+
+# --- Configuration ---
+MODEL_IDS = {
+    "SD": "reachomk/gen2seg-sd",
+    "MAE-H": "reachomk/gen2seg-mae-h"
+}
+
+# Check if a GPU is available and set the device accordingly
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {DEVICE}")
+
+# --- Global Variables for Caching Pipelines ---
+sd_pipe_global = None
+mae_pipe_global = None
+
+# --- Model Loading Functions ---
+def get_sd_pipeline():
+    """Loads and caches the gen2seg Stable Diffusion pipeline."""
+    global sd_pipe_global
+    if sd_pipe_global is None:
+        model_id_sd = MODEL_IDS["SD"]
+        print(f"Attempting to load SD pipeline from Hugging Face Hub: {model_id_sd}")
+        try:
+            sd_pipe_global = gen2segSDPipeline.from_pretrained(
+                model_id_sd,
+                use_safetensors=True,
+                # torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, # Optional: use float16 on GPU
+            ).to(DEVICE)
+            print(f"SD Pipeline loaded successfully from {model_id_sd} on {DEVICE}.")
+        except Exception as e:
+            print(f"Error loading SD pipeline from Hugging Face Hub ({model_id_sd}): {e}")
+            sd_pipe_global = None # Ensure it remains None on failure
+            # Do not raise gr.Error here; let the main function handle it.
+    return sd_pipe_global
+
+def get_mae_pipeline():
+    """Loads and caches the gen2seg MAE-H pipeline."""
+    global mae_pipe_global
+    if mae_pipe_global is None:
+        model_id_mae = MODEL_IDS["MAE-H"]
+        print(f"Loading MAE-H pipeline with model {model_id_mae} on {DEVICE}...")
+        try:
+            model = ViTMAEForPreTraining.from_pretrained(model_id_mae)
+            model.to(DEVICE)
+            model.eval() # Set to evaluation mode
+
+            # Load the official MAE-H image processor
+            # Using "facebook/vit-mae-huge" as per the original app_mae.py
+            image_processor = AutoImageProcessor.from_pretrained("facebook/vit-mae-huge")
+            
+            mae_pipe_global = gen2segMAEInstancePipeline(model=model, image_processor=image_processor)
+            # The custom MAE pipeline's model is already on the DEVICE.
+            print(f"MAE-H Pipeline with model {model_id_mae} loaded successfully on {DEVICE}.")
+        except Exception as e:
+            print(f"Error loading MAE-H model or pipeline from Hugging Face Hub ({model_id_mae}): {e}")
+            mae_pipe_global = None # Ensure it remains None on failure
+            # Do not raise gr.Error here; let the main function handle it.
+    return mae_pipe_global
+
+# --- Unified Prediction Function ---
+def segment_image(input_image: Image.Image, model_choice: str) -> Image.Image:
+    """
+    Takes a PIL Image and model choice, performs segmentation, and returns the segmented image.
+    """
+    if input_image is None:
+        raise gr.Error("No image provided. Please upload an image.")
+
+    print(f"Model selected: {model_choice}")
+    # Ensure image is in RGB format
+    image_rgb = input_image.convert("RGB")
+    original_resolution = image_rgb.size # (width, height)
+    seg_array = None
+
+    try:
+        if model_choice == "SD":
+            pipe_sd = get_sd_pipeline()
+            if pipe_sd is None:
+                raise gr.Error("The SD segmentation pipeline could not be loaded. "
+                               "Please check the Space logs for more details, or try again later.")
+
+            print(f"Running SD inference with image size: {image_rgb.size}")
+            start_time = time.time()
+            with torch.no_grad():
+                # The gen2segSDPipeline expects a single image or a list
+                # The pipeline's __call__ method handles preprocessing internally
+                seg_output = pipe_sd(image_rgb, match_input_resolution=False).prediction # Output is before resize
+                
+                # seg_output is expected to be a numpy array (N,H,W,1) or (N,1,H,W) or tensor
+                # Based on gen2seg_sd_pipeline.py, if output_type="np" (default), it's [N,H,W,1]
+                # If output_type="pt", it's [N,1,H,W]
+                # The original app_sd.py converted tensor to numpy and squeezed.
+                if isinstance(seg_output, torch.Tensor):
+                    seg_output = seg_output.cpu().numpy()
+
+                if seg_output.ndim == 4 and seg_output.shape[0] == 1: # Batch size 1
+                    if seg_output.shape[1] == 1: # Grayscale, (1, 1, H, W)
+                        seg_array = seg_output.squeeze(0).squeeze(0).astype(np.uint8)
+                    elif seg_output.shape[-1] == 1: # Grayscale, (1, H, W, 1)
+                         seg_array = seg_output.squeeze(0).squeeze(-1).astype(np.uint8)
+                    elif seg_output.shape[1] == 3: # RGB, (1, 3, H, W) -> (H, W, 3)
+                        seg_array = np.transpose(seg_output.squeeze(0), (1, 2, 0)).astype(np.uint8)
+                    elif seg_output.shape[-1] == 3: # RGB, (1, H, W, 3)
+                         seg_array = seg_output.squeeze(0).astype(np.uint8)
+                    else: # Fallback for unexpected shapes
+                        seg_array = seg_output.squeeze().astype(np.uint8)
+
+                elif seg_output.ndim == 3: # (H, W, C) or (C, H, W)
+                    seg_array = seg_output.astype(np.uint8)
+                elif seg_output.ndim == 2: # (H,W)
+                    seg_array = seg_output.astype(np.uint8)
+                else:
+                    raise TypeError(f"Unexpected SD segmentation output type/shape: {type(seg_output)}, {seg_output.shape}")
+            end_time = time.time()
+            print(f"SD Inference completed in {end_time - start_time:.2f} seconds.")
+
+
+        elif model_choice == "MAE-H":
+            pipe_mae = get_mae_pipeline()
+            if pipe_mae is None:
+                raise gr.Error("The MAE-H segmentation pipeline could not be loaded. "
+                               "Please check the Space logs for more details, or try again later.")
+
+            print(f"Running MAE-H inference with image size: {image_rgb.size}")
+            start_time = time.time()
+            with torch.no_grad():
+                # The gen2segMAEInstancePipeline expects a list of images
+                # output_type="np" returns a NumPy array
+                pipe_output = pipe_mae([image_rgb], output_type="np")
+                # Prediction is (batch_size, height, width, 3) for MAE
+                prediction_np = pipe_output.prediction[0] # Get the first (and only) image prediction
+            
+            end_time = time.time()
+            print(f"MAE-H Inference completed in {end_time - start_time:.2f} seconds.")
+            
+            if not isinstance(prediction_np, np.ndarray):
+                # This case should ideally not be reached if output_type="np"
+                prediction_np = prediction_np.cpu().numpy()
+
+            # Ensure it's in the expected (H, W, C) format and uint8
+            if prediction_np.ndim == 3 and prediction_np.shape[-1] == 3: # Expected (H, W, 3)
+                seg_array = prediction_np.astype(np.uint8)
+            else:
+                # Attempt to handle other shapes if necessary, or raise error
+                raise gr.Error(f"Unexpected MAE-H prediction shape: {prediction_np.shape}. Expected (H, W, 3).")
+            
+            # The MAE pipeline already does gamma correction and scaling to 0-255.
+            # It also ensures 3 channels.
+
+        else:
+            raise gr.Error(f"Invalid model choice: {model_choice}. Please select a valid model.")
+
+        if seg_array is None:
+             raise gr.Error("Segmentation array was not generated. An unknown error occurred.")
+
+        print(f"Segmentation array generated with shape: {seg_array.shape}, dtype: {seg_array.dtype}")
+
+        # Convert numpy array to PIL Image
+        # Handle grayscale or RGB based on seg_array channels
+        if seg_array.ndim == 2: # Grayscale
+            segmented_image_pil = Image.fromarray(seg_array, mode='L')
+        elif seg_array.ndim == 3 and seg_array.shape[-1] == 3: # RGB
+            segmented_image_pil = Image.fromarray(seg_array, mode='RGB')
+        elif seg_array.ndim == 3 and seg_array.shape[-1] == 1: # Grayscale with channel dim
+            segmented_image_pil = Image.fromarray(seg_array.squeeze(-1), mode='L')
+        else:
+            raise gr.Error(f"Cannot convert seg_array with shape {seg_array.shape} to PIL Image.")
+
+        # Resize back to original image resolution using LANCZOS for high quality
+        segmented_image_pil = segmented_image_pil.resize(original_resolution, Image.Resampling.LANCZOS)
+        
+        print(f"Segmented image processed. Output size: {segmented_image_pil.size}, mode: {segmented_image_pil.mode}")
+        return segmented_image_pil
+
+    except Exception as e:
+        print(f"Error during segmentation with {model_choice}: {e}")
+        # Re-raise as gr.Error for Gradio to display, if not already one
+        if not isinstance(e, gr.Error):
+            # It's often helpful to include the type of the original exception
+            error_type = type(e).__name__
+            raise gr.Error(f"An error occurred during segmentation: {error_type} - {str(e)}")
+        else:
+            raise e # Re-raise if it's already a gr.Error
+
+# --- Gradio Interface ---
+title = "gen2seg: Generative Models Enable Generalizable Instance Segmentation Demo (SD & MAE-H)"
+description = f"""
+<div style="text-align: center; font-family: 'Arial', sans-serif;">
+    <p>Upload an image and choose a model architecture to see the instance segmentation result generated by the respective model. </p>
+    <p>
+        Currently, inference is running on CPU.
+        Performance will be significantly better on GPU.
+    </p>
+    <ul>
+        <li><strong>SD</strong>: Based on Stable Diffusion 2.
+            <a href="https://huggingface.co/{MODEL_IDS['SD']}" target="_blank">Model Link</a>.
+            <em>Approx. CPU inference time: ~1-2 minutes per image.</em>
+        </li>
+        <li><strong>MAE-H</strong>: Based on Masked Autoencoder (Huge).
+            <a href="https://huggingface.co/{MODEL_IDS['MAE-H']}" target="_blank">Model Link</a>.
+            <em>Approx. CPU inference time: ~15-45 seconds per image.</em>
+            If you experience tokenizer artifacts or very dark images, you can use gamma correction to handle this. 
+        </li>
+    </ul>
+    <p>
+        For faster inference, please check out our GitHub to run the models locally on a GPU:
+        <a href="https://github.com/UCDvision/gen2seg" target="_blank">https://github.com/UCDvision/gen2seg</a>
+    </p>
+    <p>If the demo experiences issues, please open an issue on our GitHub.</p>
+    <p> If you have not already, please see our webpage at <a href="https://reachomk.github.io/gen2seg" target="_blank">https://reachomk.github.io/gen2seg</a>
+
+</div>
+"""
+
+article = """
+"""
+
+# Define Gradio inputs
+input_image_component = gr.Image(type="pil", label="Input Image")
+model_choice_component = gr.Dropdown(
+    choices=list(MODEL_IDS.keys()),
+    value="SD",  # Default model
+    label="Choose Segmentation Model Architecture"
+)
+
+# Define Gradio output
+output_image_component = gr.Image(type="pil", label="Segmented Image")
+
+# Example images (ensure these paths are correct if you upload examples to your Space)
+# For example, if you create an "examples" folder in your Space repo:
+# example_paths = [
+#     os.path.join("examples", "example1.jpg"),
+#     os.path.join("examples", "example2.png")
+# ]
+# Filter out non-existent example files to prevent errors
+# example_paths = [ex for ex in example_paths if os.path.exists(ex)]
+example_paths = ["cats-on-rock-1948.jpg", "dogs.png", "000000484893.jpg", "https://reachomk.github.io/gen2seg/images/comparison/vertical/7.png", "https://reachomk.github.io/gen2seg/images/comparison/horizontal/11.png", "https://reachomk.github.io/gen2seg/images/comparison/vertical/2.jpg"] # Add paths to example images here if you have them
+
+iface = gr.Interface(
+    fn=segment_image,
+    inputs=[input_image_component, model_choice_component],
+    outputs=output_image_component,
+    title=title,
+    description=description,
+    article=article,
+    examples=example_paths if example_paths else None, # Pass None if no examples
+    allow_flagging="never",
+    theme="shivi/calm_seafoam"
+)
+
+if __name__ == "__main__":
+    # Optional: Pre-load a default model on startup if desired.
+    # This can make the first inference faster but increases startup time.
+    # print("Attempting to pre-load default SD model on startup...")
+    try:
+       get_sd_pipeline() # Pre-load the default SD model
+       print("Default SD model pre-loaded successfully or was already cached.")
+    except Exception as e:
+       print(f"Could not pre-load default SD model: {e}")
+
+    print("Launching Gradio interface...")
+    iface.launch()