Update app.py

app.py

@@ -7,11 +7,30 @@ import os
 import spaces
 import json
 import re
+import torch
+from diffusers import FluxKontextPipeline
 
-# Initialize object detection using proven working YOLO models
-class AdvancedYOLODetector:
+# Initialize FLUX model for advanced inpainting
+@spaces.GPU
+def load_flux_model():
+    """Load FLUX.1 Kontext model for high-quality object removal"""
+    try:
+        pipe = FluxKontextPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-Kontext-dev", 
+            torch_dtype=torch.bfloat16
+        ).to("cuda")
+        return pipe
+    except Exception as e:
+        print(f"Failed to load FLUX model: {e}")
+        return None
+
+# Global variable to store the model (loaded once)
+flux_pipe = None
+
+# Initialize object detection using proven working models
+class AdvancedObjectDetector:
     def __init__(self):
-        # Using proven working YOLO models on Hugging Face Inference API
+        # Using proven working object detection models on Hugging Face Inference API
         self.api_url = "https://api-inference.huggingface.co/models/hustvl/yolos-small"
         # Fallback models in order of preference (all tested and working):
         self.fallback_models = [
@@ -93,7 +112,7 @@ class AdvancedYOLODetector:
                        "2. High API traffic - try again in a few minutes\n" +
                        "3. Check your HF token is valid and has sufficient quota")
 
-object_detector = AdvancedYOLODetector()
+object_detector = AdvancedObjectDetector()
 
 # Extended object class names including common variations and synonyms
 COMMON_OBJECTS = [
@@ -189,13 +208,13 @@ def fuzzy_match_object(user_input, detected_labels):
 
 def detect_objects(image, target_object, confidence_threshold, hf_token=None):
     """
-    Detect any object in the image using advanced YOLO and return bounding boxes
+    Detect any object in the image using advanced detection models and return bounding boxes
     """
     try:
         if not target_object or not target_object.strip():
             raise gr.Error("Please enter an object name to detect and remove")
             
-        # Use advanced YOLO for object detection
+        # Use advanced detection for object detection
         results = object_detector.detect(image, hf_token)
         
         if not results or not isinstance(results, list):
@@ -280,9 +299,56 @@ def create_mask_from_detections(image, detections, mask_expansion=10):
     return mask
 
 @spaces.GPU
-def remove_objects(image, object_name, confidence_threshold, mask_expansion, inpaint_prompt, hf_token):
+def flux_inpainting(image, object_name, guidance_scale=2.5, steps=28):
+    """
+    Use FLUX.1 Kontext for intelligent object removal
+    """
+    global flux_pipe
+    
+    try:
+        # Load FLUX model if not already loaded
+        if flux_pipe is None:
+            print("Loading FLUX.1 Kontext model...")
+            flux_pipe = load_flux_model()
+            
+        if flux_pipe is None:
+            raise Exception("Failed to load FLUX model")
+        
+        # Create intelligent removal prompt
+        removal_prompt = f"Remove the {object_name} from this image, fill with natural background that matches the surrounding environment, photorealistic, seamless, high quality"
+        
+        # Use FLUX for contextual editing
+        result = flux_pipe(
+            image=image.convert("RGB"),
+            prompt=removal_prompt,
+            guidance_scale=guidance_scale,
+            width=image.size[0],
+            height=image.size[1],
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(42),
+        ).images[0]
+        
+        return result, True
+        
+    except Exception as e:
+        print(f"FLUX inpainting error: {str(e)}")
+        return None, False
+
+def create_mask_overlay(image, mask):
+    """Create a visualization showing the mask overlay on the original image"""
+    result_image = image.copy()
+    overlay = Image.new('RGBA', image.size, (255, 0, 0, 100))
+    mask_rgba = mask.convert('RGBA')
+    result_image = Image.alpha_composite(
+        result_image.convert('RGBA'), 
+        Image.composite(overlay, Image.new('RGBA', image.size, (0,0,0,0)), mask)
+    )
+    return result_image.convert('RGB')
+
+@spaces.GPU
+def remove_objects(image, object_name, confidence_threshold, mask_expansion, guidance_scale, steps, hf_token):
     """
-    Main function to remove any specified object from image using advanced YOLO + SDXL
+    Main function to remove any specified object using advanced detection + FLUX inpainting
     """
     try:
         if image is None:
@@ -296,7 +362,7 @@ def remove_objects(image, object_name, confidence_threshold, mask_expansion, inp
         if not token:
             raise gr.Error("Please provide your Hugging Face token or set HF_TOKEN in Space secrets")
         
-        # Step 1: Detect objects using advanced YOLO
+        # Step 1: Detect objects
         detections = detect_objects(image, object_name, confidence_threshold, token)
         
         if not detections:
@@ -308,119 +374,49 @@ def remove_objects(image, object_name, confidence_threshold, mask_expansion, inp
             suggestion_msg += "• Checking if the object is clearly visible in the image"
             return image, None, suggestion_msg
         
-        # Step 2: Create mask with adaptive expansion
+        # Step 2: Create mask for debugging/visualization
         mask = create_mask_from_detections(image, detections, mask_expansion)
         
-        # Step 3: Use reliable inpainting model
-        # Try multiple inpainting models in order of preference
-        inpaint_models = [
-            "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-inpainting",
-            "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-inpainting",
-            "https://api-inference.huggingface.co/models/diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
-        ]
-        
-        headers = {"Authorization": f"Bearer {token}"}
-        
-        # Convert images to bytes
-        img_buffer = io.BytesIO()
-        image.save(img_buffer, format='PNG')
-        img_bytes = img_buffer.getvalue()
-        
-        mask_buffer = io.BytesIO()
-        mask.save(mask_buffer, format='PNG')
-        mask_bytes = mask_buffer.getvalue()
-        
-        # Prepare multipart form data
-        files = {
-            'image': ('image.png', img_bytes, 'image/png'),
-            'mask': ('mask.png', mask_bytes, 'image/png')
-        }
-        
-        # Enhanced inpainting prompt
-        enhanced_prompt = f"{inpaint_prompt}, photorealistic, high quality, detailed, natural lighting"
-        
-        data = {
-            'prompt': enhanced_prompt,
-            'negative_prompt': 'blurry, low quality, distorted, artifacts, unrealistic, pixelated, noise',
-            'num_inference_steps': 25,
-            'guidance_scale': 7.5,
-            'strength': 0.99
-        }
-        
-        # Try multiple inpainting models
-        inpainting_success = False
-        last_error = ""
-        
-        for i, inpaint_api_url in enumerate(inpaint_models):
-            try:
-                print(f"Trying inpainting model {i+1}/{len(inpaint_models)}: {inpaint_api_url.split('/')[-1]}")
-                
-                response = requests.post(inpaint_api_url, headers=headers, files=files, data=data, timeout=120)
-                
-                if response.status_code == 503:
-                    # Model is loading, wait and retry once
-                    import time
-                    time.sleep(10)
-                    response = requests.post(inpaint_api_url, headers=headers, files=files, data=data, timeout=120)
-                
-                if response.status_code == 200:
-                    result_image = Image.open(io.BytesIO(response.content))
-                    detected_labels = [d.get('label', 'unknown') for d in detections]
-                    status_msg = f"✅ Successfully removed {len(detections)} '{object_name}' object(s)\n"
-                    status_msg += f"🎯 Detected as: {', '.join(detected_labels)}\n"
-                    status_msg += f"🔧 Used: {inpaint_api_url.split('/')[-1]} for inpainting"
-                    inpainting_success = True
-                    break
-                else:
-                    last_error = f"HTTP {response.status_code}: {response.text[:200]}"
-                    print(f"Model {i+1} failed: {last_error}")
-                    continue
-                    
-            except Exception as e:
-                last_error = str(e)
-                print(f"Model {i+1} error: {last_error}")
-                continue
-        
-        if not inpainting_success:
-            # Fallback: return original with mask overlay for debugging
+        # Step 3: Use FLUX.1 Kontext for intelligent object removal
+        print("Using FLUX.1 Kontext for advanced object removal...")
+        result_image, flux_success = flux_inpainting(image, object_name, guidance_scale, steps)
+        
+        if flux_success and result_image:
+            detected_labels = [d.get('label', 'unknown') for d in detections]
+            status_msg = f"✅ Successfully removed {len(detections)} '{object_name}' object(s)\n"
+            status_msg += f"🎯 Detected as: {', '.join(detected_labels)}\n"
+            status_msg += f"🚀 Used: FLUX.1 Kontext for professional-quality removal\n"
+            status_msg += f"⚙️ Settings: Guidance={guidance_scale}, Steps={steps}"
+            return result_image, mask, status_msg
+        else:
+            # Fallback: show detection areas
             result_image = create_mask_overlay(image, mask)
-            status_msg = f"⚠️ All inpainting models failed. Last error: {last_error}\n"
-            status_msg += f"🎯 Found {len(detections)} '{object_name}' object(s) - detection was successful\n"
-            status_msg += f"📍 Showing detected areas in red overlay"
-        
-        return result_image, mask, status_msg
+            status_msg = f"⚠️ FLUX inpainting failed, but detection was successful\n"
+            status_msg += f"🎯 Found {len(detections)} '{object_name}' object(s)\n"
+            status_msg += f"📍 Showing detected areas in red overlay\n"
+            status_msg += f"💡 Try adjusting guidance scale or steps, or check GPU availability"
+            return result_image, mask, status_msg
         
     except Exception as e:
         return image, None, f"❌ Error: {str(e)}"
 
-def create_mask_overlay(image, mask):
-    """Create a visualization showing the mask overlay on the original image"""
-    result_image = image.copy()
-    overlay = Image.new('RGBA', image.size, (255, 0, 0, 100))
-    mask_rgba = mask.convert('RGBA')
-    result_image = Image.alpha_composite(
-        result_image.convert('RGBA'), 
-        Image.composite(overlay, Image.new('RGBA', image.size, (0,0,0,0)), mask)
-    )
-    return result_image.convert('RGB')
-
 # Create Gradio interface
 with gr.Blocks(
     fill_height=True, 
-    title="Advanced Object Removal with YOLOv8",
+    title="Professional Object Removal with FLUX",
     theme=gr.themes.Soft()
 ) as demo:
     
     gr.Markdown("""
-    # 🚀 Advanced Object Removal using Proven Detection Models + SDXL Inpainting
+    # 🚀 Professional Object Removal using Advanced AI
     
-    Upload an image and specify **ANY object** you want to remove - no limitations!
+    Upload an image and specify **ANY object** you want to remove with professional results!
     
-    **How it works:**
+    **Powered by cutting-edge AI:**
     1. 🔍 **Multi-Model Detection**: Uses multiple proven object detection models with auto-fallback
     2. 🧠 **Smart Matching**: Handles synonyms, plurals, and fuzzy object name matching
-    3. 🎭 **Adaptive Masking**: Creates intelligent removal masks  
-    4. 🎨 **SDXL Inpainting**: Uses state-of-the-art AI to fill removed areas seamlessly
+    3. 🎭 **Intelligent Masking**: Creates precise removal masks  
+    4. 🚀 **FLUX.1 Kontext**: Uses state-of-the-art contextual editing for professional results
     """)
     
     with gr.Row():
@@ -468,18 +464,29 @@ with gr.Blocks(
                     value=20,
                     step=5,
                     label="📏 Mask Expansion (pixels)",
-                    info="Expand mask around detected objects for better removal"
+                    info="Expand mask around detected objects for debugging"
+                )
+                
+                guidance_scale = gr.Slider(
+                    minimum=1.0,
+                    maximum=10.0,
+                    value=2.5,
+                    step=0.1,
+                    label="🎯 FLUX Guidance Scale",
+                    info="Higher = more faithful to prompt, lower = more creative"
                 )
                 
-                inpaint_prompt = gr.Textbox(
-                    label="✨ Inpainting Prompt",
-                    value="natural background, seamless, realistic environment",
-                    placeholder="Describe what should replace the removed object",
-                    info="Be specific about the desired background/replacement"
+                steps = gr.Slider(
+                    minimum=10,
+                    maximum=50,
+                    value=28,
+                    step=2,
+                    label="🔄 FLUX Steps",
+                    info="More steps = higher quality but slower processing"
                 )
             
             hf_token = gr.Textbox(
-                label="🔑 Hugging Face Token (Optional)",
+                label="🔑 Hugging Face Token (Required)",
                 type="password",
                 placeholder="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
                 info="Get token from https://huggingface.co/settings/tokens (or set HF_TOKEN in Space secrets)"
@@ -498,7 +505,7 @@ with gr.Blocks(
                     height=300
                 )
                 mask_image = gr.Image(
-                    label="🎭 Generated Mask",
+                    label="🎭 Detection Mask (Debug)",
                     type="pil", 
                     height=300
                 )
@@ -506,7 +513,7 @@ with gr.Blocks(
             status_text = gr.Textbox(
                 label="📊 Status & Detection Info",
                 interactive=False,
-                max_lines=4
+                max_lines=5
             )
     
     # Event handlers
@@ -517,7 +524,8 @@ with gr.Blocks(
             object_name,
             confidence_threshold,
             mask_expansion,
-            inpaint_prompt,
+            guidance_scale,
+            steps,
             hf_token
         ],
         outputs=[output_image, mask_image, status_text]
@@ -531,11 +539,13 @@ with gr.Blocks(
             
             1. **Upload an image** containing objects you want to remove
             2. **Enter ANY object name** in the text box - no restrictions!
-            3. **Adjust settings** if needed:
+            3. **Adjust detection settings** if needed:
                - **Confidence**: Start with 0.3, increase if too many false detections
-               - **Mask expansion**: Larger values ensure complete object removal
-               - **Inpainting prompt**: Describe the desired replacement scene
-            4. **Click "Remove Objects"** and wait for AI processing
+               - **Mask expansion**: For debugging - shows detection areas
+            4. **Fine-tune FLUX settings**:
+               - **Guidance Scale**: 2.5 is optimal for most cases
+               - **Steps**: 28 gives good quality/speed balance
+            5. **Click "Remove Objects"** and wait for professional AI processing
             
             ### 💡 Smart Object Recognition:
             - **Handles variations**: "car" = "vehicle" = "automobile"
@@ -558,17 +568,17 @@ with gr.Blocks(
             - **Nature**: tree, flower, rock, cloud, mountain
             - **And literally thousands more!**
             
-            ### ⚠️ System Info:
-            - **🔍 Detection**: YOLOS-Small + DETR + OWL-ViT (auto-fallback)
-            - **🎨 Inpainting**: Multiple SD models with auto-fallback
-            - **⏱️ Processing**: 30-120 seconds (first request may be slower)
-            - **🔧 Auto-retry**: Tries multiple models if one is busy
-            - **Token Required**: HF token needed for API access
+            ### ⚡ FLUX.1 Kontext Advantages:
+            - **🎨 Professional Quality**: State-of-the-art contextual editing
+            - **🧠 Intelligent Fill**: Understands scene context and lighting
+            - **⚡ GPU Accelerated**: Fast processing with high quality
+            - **🎯 Precise Control**: Fine-tunable guidance and steps
+            - **🔧 No API Limits**: Runs locally without external dependencies
             
-            **If you get "models unavailable" error:**
-            - Wait 2-3 minutes and try again (models loading)
-            - Check your HF token is valid and has API quota
-            - Try during off-peak hours for better performance
+            **System Requirements:**
+            - GPU-enabled environment (automatically handled in Spaces)
+            - HF token for object detection API access
+            - Processing time: 30-90 seconds depending on image size
             """)
 
 if __name__ == "__main__":