Simplify for HF Spaces deployment

- Remove SAM mask refinement functionality
- Remove output_path file saving features
- Simplify UI to focus on core drag inpainting
- Remove complex dependencies for better HF Spaces compatibility

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import tempfile
+import os
 from utils.ui_utils import *
 
 CANVAS_SIZE = 400
@@ -22,7 +24,6 @@ def create_interface():
                     canvas = gr.Image(type="numpy", tool="sketch", label=" ", height=CANVAS_SIZE, width=CANVAS_SIZE)
                     with gr.Row():
                         fit_btn = gr.Button("Resize Image")
-                        if_sam_box = gr.Checkbox(label='Refine mask (SAM)')
 
                 # Control Points Column
                 with gr.Column():
@@ -40,20 +41,15 @@ def create_interface():
                         run_btn = gr.Button("Inpaint")
                         reset_btn = gr.Button("Reset All")
 
-        # Output Settings
-        with gr.Row("Generation Parameters"):
-            sam_ks = gr.Slider(minimum=11, maximum=51, value=21, step=2, label='How much to refine mask with SAM', interactive=True)
+        # Generation Parameters
+        with gr.Row():
             inpaint_ks = gr.Slider(minimum=0, maximum=25, value=5, step=1, label='How much to expand inpainting mask', interactive=True)
-            output_path = gr.Textbox(value='output/app', label="Output path")
 
         setup_events(
             components={
                 'canvas': canvas,
                 'input_img': input_img,
                 'output_img': output_img,
-                'output_path': output_path,
-                'if_sam_box': if_sam_box,
-                'sam_ks': sam_ks,
                 'inpaint_ks': inpaint_ks,
             },
             state=state,
@@ -75,21 +71,21 @@ def setup_events(components, state, buttons):
             clear_all,
             [state['canvas_size']],
             [components['canvas'], components['input_img'], components['output_img'], 
-             state['points_list'], components['sam_ks'], components['inpaint_ks'], components['output_path'], state['inpaint_mask']]
+             state['points_list'], components['inpaint_ks'], state['inpaint_mask']]
         )
         
         components['canvas'].clear(
             clear_all,
             [state['canvas_size']],
             [components['canvas'], components['input_img'], components['output_img'], 
-             state['points_list'], components['sam_ks'], components['inpaint_ks'], components['output_path'], state['inpaint_mask']]
+             state['points_list'], components['inpaint_ks'], state['inpaint_mask']]
         )
 
     # Image manipulation events
     def setup_image_events():
         buttons['fit'].click(
             clear_point,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['input_img']]
         ).then(
             resize,
@@ -101,41 +97,21 @@ def setup_events(components, state, buttons):
     def setup_canvas_events():
         components['canvas'].edit(
             visualize_user_drag,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list']],
             [components['input_img']]
         ).then(
             preview_out_image,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['output_img'], state['inpaint_mask']]
         )
 
-        components['if_sam_box'].change(
-            visualize_user_drag,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box']],
-            [components['input_img']]
-        ).then(
-            preview_out_image,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
-            [components['output_img'], state['inpaint_mask']]
-        )
-
-        components['sam_ks'].change(
-            visualize_user_drag,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box']],
-            [components['input_img']]
-        ).then(
-            preview_out_image,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
-            [components['output_img'], state['inpaint_mask']]
-        )
-        
         components['inpaint_ks'].change(
             visualize_user_drag,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box']],
+            [components['canvas'], state['points_list']],
             [components['input_img']]
         ).then(
             preview_out_image,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['output_img'], state['inpaint_mask']]
         )
 
@@ -143,11 +119,11 @@ def setup_events(components, state, buttons):
     def setup_input_events():
         components['input_img'].select(
             add_point,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['input_img']]
         ).then(
             preview_out_image,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['output_img'], state['inpaint_mask']]
         )
 
@@ -155,21 +131,21 @@ def setup_events(components, state, buttons):
     def setup_point_events():
         buttons['undo'].click(
             undo_point,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['input_img']]
         ).then(
             preview_out_image,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['output_img'], state['inpaint_mask']]
         )
         
         buttons['clear'].click(
             clear_point,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['input_img']]
         ).then(
             preview_out_image,
-            [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['output_img'], state['inpaint_mask']]
         )
 
@@ -177,7 +153,7 @@ def setup_events(components, state, buttons):
     def setup_processing_events():
         buttons['run'].click(
             preview_out_image,
-           [components['canvas'], state['points_list'], components['sam_ks'], components['inpaint_ks'], components['if_sam_box'], components['output_path']],
+            [components['canvas'], state['points_list'], components['inpaint_ks']],
             [components['output_img'], state['inpaint_mask']]
         ).then(
             inpaint,
@@ -195,7 +171,8 @@ def setup_events(components, state, buttons):
 
 def main():
     app = create_interface()
-    app.queue().launch(share=True, debug=True)
+    # HF Space compatible launch
+    app.queue().launch()
 
 if __name__ == '__main__':
     main()

utils/refine_mask.py

@@ -1,168 +0,0 @@
-import os
-import urllib.request
-from typing import Optional
-
-import cv2
-import numpy as np
-import torch
-import torch.nn as nn
-
-
-def download_model(checkpoint_path: str, model_name: str = "efficientvit_sam_l0.pt") -> str:
-    """
-    Download the model checkpoint if not found locally.
-
-    Args:
-        checkpoint_path: Local path where model should be saved
-        model_name: Name of the model file to download
-
-    Returns:
-        str: Path to the downloaded checkpoint
-    """
-    os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
-    
-    base_url = "https://huggingface.co/mit-han-lab/efficientvit-sam/resolve/main"
-    model_url = f"{base_url}/{model_name}"
-    
-    try:
-        print(f"Downloading model from {model_url}...")
-        urllib.request.urlretrieve(model_url, checkpoint_path)
-        print(f"Model successfully downloaded to {checkpoint_path}")
-        return checkpoint_path
-    except Exception as e:
-        raise RuntimeError(f"Failed to download model: {str(e)}")
-
-
-class SamMaskRefiner(nn.Module):
-    CHECKPOINT_DIR = 'checkpoints'
-    MODEL_CONFIGS = {
-        'l0': 'efficientvit_sam_l0.pt',
-        'l1': 'efficientvit_sam_l1.pt',
-        'l2': 'efficientvit_sam_l2.pt'
-    }
-    
-    def __init__(self, model_name: str = 'l0') -> None:
-        """
-        Initialize SAM predictor with specified model version.
-
-        Args:
-            model_name: Model version to use ('l0', 'l1', or 'l2'). Defaults to 'l0'.
-
-        Raises:
-            ValueError: If invalid model_name is provided
-            RuntimeError: If model loading fails after download attempt
-        """
-        super().__init__()
-        
-        if model_name not in self.MODEL_CONFIGS:
-            raise ValueError(f"Invalid model_name. Choose from: {list(self.MODEL_CONFIGS.keys())}")
-            
-        model_filename = self.MODEL_CONFIGS[model_name]
-        checkpoint_path = os.path.join(self.CHECKPOINT_DIR, model_filename)
-        
-        try:
-            from efficientvit.models.efficientvit.sam import EfficientViTSamPredictor
-            from efficientvit.sam_model_zoo import create_efficientvit_sam_model
-        except ImportError:
-            raise ImportError(
-                "Failed to import EfficientViT modules. Please ensure the package is installed:\n"
-                "pip install git+https://github.com/mit-han-lab/efficientvit.git"
-            )
-        
-        if not os.path.exists(checkpoint_path):
-            print(f"Checkpoint not found at {checkpoint_path}. Attempting to download...")
-            checkpoint_path = download_model(checkpoint_path, model_filename)
-        
-        try:
-            model_type = f'efficientvit-sam-{model_name}'
-            device = 'cuda' if torch.cuda.is_available() else 'cpu'
-            self.model = create_efficientvit_sam_model(model_type, True, checkpoint_path).eval()
-            self.model = self.model.requires_grad_(False).to(device)
-            self.predictor = EfficientViTSamPredictor(self.model)
-            print(f"\033[92mEfficientViT-SAM model loaded from: {checkpoint_path}\033[0m")
-        except Exception as e:
-            raise RuntimeError(f"Failed to load model: {str(e)}")
-
-    def sample_points_from_mask(self, mask: np.ndarray, max_points: int = 128) -> np.ndarray:
-        """
-        Sample points uniformly from masked regions.
-
-        Args:
-            mask: Binary mask array of shape (H, W) with 0-1 values.
-            max_points: Maximum number of points to sample.
-
-        Returns:
-            np.ndarray: Array of shape (N, 2) containing [x,y] coordinates.
-        """
-        y_indices, x_indices = np.where(mask > 0.5)
-        total_points = len(y_indices)
-        
-        if total_points <= max_points:
-            return np.stack([x_indices, y_indices], axis=1)
-        
-        y_min, y_max = y_indices.min(), y_indices.max()
-        x_min, x_max = x_indices.min(), x_indices.max()
-        
-        aspect_ratio = (x_max - x_min) / max(y_max - y_min, 1)
-        ny = int(np.sqrt(max_points / aspect_ratio))
-        nx = int(ny * aspect_ratio)
-        
-        x_bins = np.linspace(x_min, x_max + 1, nx + 1, dtype=np.int32)
-        y_bins = np.linspace(y_min, y_max + 1, ny + 1, dtype=np.int32)
-        
-        x_dig = np.digitize(x_indices, x_bins) - 1
-        y_dig = np.digitize(y_indices, y_bins) - 1
-        bin_indices = y_dig * nx + x_dig
-        unique_bins = np.unique(bin_indices)
-        
-        points = []
-        for idx in unique_bins:
-            bin_y = idx // nx
-            bin_x = idx % nx
-            mask = (y_dig == bin_y) & (x_dig == bin_x)
-            
-            if np.any(mask):
-                px = int(np.mean(x_indices[mask]))
-                py = int(np.mean(y_indices[mask]))
-                points.append([px, py])
-        
-        points = np.array(points)
-        
-        if len(points) > max_points:
-            indices = np.linspace(0, len(points) - 1, max_points, dtype=int)
-            points = points[indices]
-        
-        return points
-
-    def refine_mask(self, image: np.ndarray, input_mask: np.ndarray, kernel_size: int = 21) -> np.ndarray:
-        """
-        Refine an input mask using the SAM (Segment Anything Model) model.
-
-        Args:
-            image: RGB image, shape (H, W, 3), values in [0, 255]
-            input_mask: Binary mask, shape (H, W), values in {0, 1}
-            kernel_size: Size of morphological kernel (default: 21)
-
-        Returns:
-            Refined binary mask, shape (H, W), values in {0, 1}
-        """
-        points = self.sample_points_from_mask(input_mask, max_points=128)
-        if len(points) == 0:
-            return input_mask
-
-        self.predictor.set_image(image)
-        masks_pred, _, _ = self.predictor.predict(
-            point_coords=points,
-            point_labels=np.ones(len(points)),
-            multimask_output=False
-        )
-        sam_mask = masks_pred[0]
-
-        kernel = np.ones((kernel_size, kernel_size), np.uint8)
-        expanded_input = cv2.dilate(input_mask.astype(np.uint8), kernel)
-        preserved_input = cv2.erode(input_mask.astype(np.uint8), kernel)
-    
-        sam_mask = np.logical_and(expanded_input, sam_mask).astype(np.uint8)
-        sam_mask = np.logical_or(preserved_input, sam_mask).astype(np.uint8)
-        
-        return sam_mask

utils/ui_utils.py

@@ -1,6 +1,7 @@
 import os
 import pickle
 from time import perf_counter
+import tempfile
 
 import cv2
 import gradio as gr
@@ -10,7 +11,6 @@ from PIL import Image
 from diffusers import AutoPipelineForInpainting, AutoencoderTiny, LCMScheduler
 
 from utils.drag import bi_warp
-from utils.refine_mask import SamMaskRefiner
 
 
 __all__ = [
@@ -19,10 +19,13 @@ __all__ = [
     'add_point', 'undo_point', 'clear_point',
 ]
 
+# Global variables for lazy loading
+pipe = None
+
 # UI functions
 def clear_all(length):
     """Reset UI by clearing all input images and parameters."""
-    return (gr.Image(value=None, height=length, width=length),) * 3 + ([], 21, 2, "output/app", None)
+    return (gr.Image(value=None, height=length, width=length),) * 3 + ([], 2, None)
 
 def resize(canvas, gen_length, canvas_length):
     """Resize canvas while maintaining aspect ratio."""
@@ -51,57 +54,35 @@ def process_canvas(canvas):
     return image, mask
 
 # Point manipulation functions
-def add_point(canvas, points, sam_ks, if_sam, output_path, evt: gr.SelectData):
+def add_point(canvas, points, inpaint_ks, evt: gr.SelectData):
     """Add selected point to points list and update image."""
     if canvas is None:
         return None 
     points.append(evt.index)
-    return visualize_user_drag(canvas, points, sam_ks, if_sam, output_path)
+    return visualize_user_drag(canvas, points)
 
-def undo_point(canvas, points, sam_ks, if_sam, output_path):
+def undo_point(canvas, points, inpaint_ks):
     """Remove last point and update image."""
     if canvas is None:
         return None 
     if len(points) > 0:
         points.pop()
-    return visualize_user_drag(canvas, points, sam_ks, if_sam, output_path)
+    return visualize_user_drag(canvas, points)
 
-def clear_point(canvas, points, sam_ks, if_sam, output_path):
+def clear_point(canvas, points, inpaint_ks):
     """Clear all points and update image."""
     if canvas is None:
         return None 
     points.clear()
-    return visualize_user_drag(canvas, points, sam_ks, if_sam, output_path)
+    return visualize_user_drag(canvas, points)
 
 # Visualization tools
-def refine_mask(image, mask, kernel_size):
-    """Refine mask using SAM model if available."""
-    global sam_refiner
-    try:
-        if 'sam_refiner' not in globals():
-            sam_refiner = SamMaskRefiner()
-        return sam_refiner.refine_mask(image, mask, kernel_size)
-    except ImportError:
-        gr.Warning("EfficientVit not installed. Please install with: pip install git+https://github.com/mit-han-lab/efficientvit.git")
-        return mask
-    except Exception as e:
-        gr.Warning(f"Error refining mask: {str(e)}")
-        return mask
-
-def visualize_user_drag(canvas, points, sam_ks, if_sam=False, output_path=None):
-    """Visualize control points and motion vectors on the input image.
-    
-    Args:
-        canvas (dict): Gradio canvas containing image and mask
-        points (list): List of (x,y) coordinate pairs for control points
-        sam_ks (int): Kernel size for SAM mask refinement
-        if_sam (bool): Whether to use SAM refinement on mask
-    """
+def visualize_user_drag(canvas, points):
+    """Visualize control points and motion vectors on the input image."""
     if canvas is None:
         return None
     
     image, mask = process_canvas(canvas)
-    mask = refine_mask(image, mask, sam_ks) if if_sam and mask.sum() > 0 else mask
 
     # Apply colored mask overlay
     result = image.copy()
@@ -120,29 +101,11 @@ def visualize_user_drag(canvas, points, sam_ks, if_sam=False, output_path=None):
         else:
             cv2.circle(image, tuple(point), 10, (255, 0, 0), -1)  # Start point
             prev_point = point
-
-    if output_path:
-        os.makedirs(output_path, exist_ok=True)
-        Image.fromarray(image).save(os.path.join(output_path, 'user_drag_i4p.png'))
+    
     return image
 
-def preview_out_image(canvas, points, sam_ks, inpaint_ks, if_sam=False, output_path=None):
-    """Preview warped image result and generate inpainting mask.
-    
-    Args:
-        canvas (dict): Gradio canvas containing the input image and mask
-        points (list): List of (x,y) coordinate pairs defining source and target positions for warping
-        sam_ks (int): Kernel size parameter for SAM mask refinement
-        inpaint_ks (int): Kernel size parameter for inpainting mask generation
-        if_sam (bool): Whether to use SAM model for mask refinement
-        output_path (str, optional): Directory path to save original image and metadata
-        
-    Returns:
-        tuple: 
-            - ndarray: Warped image with grid pattern overlay on regions needing inpainting
-            - ndarray: Binary mask (255 for inpainting regions, 0 elsewhere)
-            - (None, None): If canvas is empty or fewer than 2 control points provided
-    """
+def preview_out_image(canvas, points, inpaint_ks):
+    """Preview warped image result and generate inpainting mask."""
     if canvas is None:
         return None, None
     
@@ -155,15 +118,7 @@ def preview_out_image(canvas, points, sam_ks, inpaint_ks, if_sam=False, output_p
     size_valid = all(max(x.shape[:2] if len(x.shape) > 2 else x.shape) == 512 for x in (image, mask))
     if not (shapes_valid and size_valid):
         gr.Warning('Click Resize Image Button first.')
-
-    mask = refine_mask(image, mask, sam_ks) if if_sam and mask.sum() > 0 else mask
-    
-    if output_path:
-        os.makedirs(output_path, exist_ok=True)
-        Image.fromarray(image).save(os.path.join(output_path, 'original_image.png'))
-        metadata = {'mask': mask, 'points': points}
-        with open(os.path.join(output_path, 'meta_data_i4p.pkl'), 'wb') as f:
-            pickle.dump(metadata, f)
+        return image, None
     
     handle_pts, target_pts, inpaint_mask = bi_warp(mask, points, inpaint_ks)
     image[target_pts[:, 1], target_pts[:, 0]] = image[handle_pts[:, 1], handle_pts[:, 0]]
@@ -172,9 +127,6 @@ def preview_out_image(canvas, points, sam_ks, inpaint_ks, if_sam=False, output_p
     background = np.ones_like(mask) * 255
     background[::10] = background[:, ::10] = 0
     image = np.where(inpaint_mask[..., np.newaxis]==1, background[..., np.newaxis], image)
-
-    if output_path:
-        Image.fromarray(image).save(os.path.join(output_path, 'preview_image.png'))
     
     return image, (inpaint_mask * 255).astype(np.uint8)
 
@@ -187,11 +139,26 @@ def setup_pipeline(device='cuda', model_version='v1-5'):
     }
     model_id, lora_id, vae_id = MODEL_CONFIGS[model_version]
 
-    pipe = AutoPipelineForInpainting.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", safety_checker=None)
+    # Check if CUDA is available, fallback to CPU
+    if not torch.cuda.is_available():
+        device = 'cpu'
+        torch_dtype = torch.float32
+        variant = None
+    else:
+        torch_dtype = torch.float16
+        variant = "fp16"
+
+    gr.Info('Loading inpainting pipeline...')
+    pipe = AutoPipelineForInpainting.from_pretrained(
+        model_id, 
+        torch_dtype=torch_dtype, 
+        variant=variant, 
+        safety_checker=None
+    )
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
     pipe.load_lora_weights(lora_id)
     pipe.fuse_lora()
-    pipe.vae = AutoencoderTiny.from_pretrained(vae_id, torch_dtype=torch.float16)
+    pipe.vae = AutoencoderTiny.from_pretrained(vae_id, torch_dtype=torch_dtype)
     pipe = pipe.to(device)
     
     # Pre-compute prompt embeddings during setup
@@ -206,19 +173,20 @@ def setup_pipeline(device='cuda', model_version='v1-5'):
             
     return pipe
 
-pipe = setup_pipeline(model_version='v1-5')
-pipe.cached_prompt_embeds = pipe.encode_prompt('', 'cuda', 1, False)[0]
+def get_pipeline():
+    """Lazy load pipeline only when needed."""
+    global pipe
+    if pipe is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        pipe = setup_pipeline(device=device, model_version='v1-5')
+        if device == 'cuda':
+            pipe.cached_prompt_embeds = pipe.encode_prompt('', 'cuda', 1, False)[0]
+        else:
+            pipe.cached_prompt_embeds = pipe.encode_prompt('', 'cpu', 1, False)[0]
+    return pipe
 
 def inpaint(image, inpaint_mask):
-    """Perform efficient inpainting on masked regions using Stable Diffusion.
-    
-    Args:
-        image (ndarray): Input RGB image array (warped preview image)
-        inpaint_mask (ndarray): Binary mask array where 255 indicates regions to inpaint
-        
-    Returns:
-        ndarray: Inpainted image with masked regions filled in
-    """
+    """Perform efficient inpainting on masked regions using Stable Diffusion."""
     if image is None:
         return None
 
@@ -226,6 +194,10 @@ def inpaint(image, inpaint_mask):
         return image
     
     start = perf_counter()
+    
+    # Get pipeline (lazy loading)
+    pipe = get_pipeline()
+    
     pipe_id = 'xl' if 'xl' in pipe.config._name_or_path else 'v1-5'
     inpaint_strength = 0.99 if pipe_id == 'xl' else 1.0
 
@@ -254,18 +226,22 @@ def inpaint(image, inpaint_mask):
     }
 
     # Run pipeline
-    if pipe_id == 'v1-5':
-        inpainted = pipe(
-            prompt_embeds=pipe.cached_prompt_embeds,
-            **common_params
-        ).images[0]
-    else:
-        inpainted = pipe(
-            prompt_embeds=pipe.cached_prompt_embeds,
-            pooled_prompt_embeds=pipe.cached_pooled_prompt_embeds,
-            **common_params
-        ).images[0]
+    try:
+        if pipe_id == 'v1-5':
+            inpainted = pipe(
+                prompt_embeds=pipe.cached_prompt_embeds,
+                **common_params
+            ).images[0]
+        else:
+            inpainted = pipe(
+                prompt_embeds=pipe.cached_prompt_embeds,
+                pooled_prompt_embeds=pipe.cached_pooled_prompt_embeds,
+                **common_params
+            ).images[0]
+    except Exception as e:
+        gr.Warning(f"Inpainting failed: {str(e)}")
+        return image
 
     # Post-process results
     inpaint_mask = (inpaint_mask[..., np.newaxis] / 255).astype(np.uint8)
-    return (inpainted * 255).astype(np.uint8) * inpaint_mask + image * (1 - inpaint_mask)
+    return (inpainted * 255).astype(np.uint8) * inpaint_mask + image * (1 - inpaint_mask)