Update app.py

app.py

@@ -11,13 +11,14 @@ class ObjectDetector:
     def __init__(self):
         self.api_url = "https://api-inference.huggingface.co/models/facebook/detr-resnet-50"
         
-    def detect(self, image):
-        # Use HF_TOKEN environment variable if available, otherwise require user token
-        hf_token = os.getenv("HF_TOKEN")
-        if not hf_token:
-            raise Exception("HF_TOKEN not found in environment variables")
+    def detect(self, image, hf_token=None):
+        # Try multiple ways to get HF token
+        token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+        
+        if not token:
+            raise Exception("HF Token required. Please set HF_TOKEN in Space secrets or environment variables")
             
-        headers = {"Authorization": f"Bearer {hf_token}"}
+        headers = {"Authorization": f"Bearer {token}"}
         
         # Convert PIL image to bytes
         img_buffer = io.BytesIO()
@@ -48,13 +49,13 @@ COCO_CLASSES = [
     'toothbrush'
 ]
 
-def detect_objects(image, target_class, confidence_threshold):
+def detect_objects(image, target_class, confidence_threshold, hf_token=None):
     """
     Detect objects in the image and return bounding boxes for the target class
     """
     try:
         # Use Hugging Face Inference API for object detection
-        results = object_detector.detect(image)
+        results = object_detector.detect(image, hf_token)
         
         # Filter results for target class
         target_detections = []
@@ -89,7 +90,7 @@ def create_mask_from_detections(image, detections, mask_expansion=10):
     return mask
 
 @spaces.GPU
-def remove_objects(image, object_class, confidence_threshold, mask_expansion, inpaint_prompt):
+def remove_objects(image, object_class, confidence_threshold, mask_expansion, inpaint_prompt, hf_token):
     """
     Main function to remove objects from image using SDXL inpainting
     """
@@ -97,8 +98,13 @@ def remove_objects(image, object_class, confidence_threshold, mask_expansion, in
         if image is None:
             raise gr.Error("Please upload an image")
         
+        # Try to get token from multiple sources
+        token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+        if not token:
+            raise gr.Error("Please provide your Hugging Face token or set HF_TOKEN in Space secrets")
+        
         # Step 1: Detect objects
-        detections = detect_objects(image, object_class, confidence_threshold)
+        detections = detect_objects(image, object_class, confidence_threshold, token)
         
         if not detections:
             return image, None, f"No {object_class} objects detected with confidence > {confidence_threshold}"
@@ -109,12 +115,7 @@ def remove_objects(image, object_class, confidence_threshold, mask_expansion, in
         # Step 3: Use SDXL for inpainting via Hugging Face Inference API
         inpaint_api_url = "https://api-inference.huggingface.co/models/diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
         
-        # Get HF token from environment
-        hf_token = os.getenv("HF_TOKEN")
-        if not hf_token:
-            raise gr.Error("HF_TOKEN not found in environment variables")
-            
-        headers = {"Authorization": f"Bearer {hf_token}"}
+        headers = {"Authorization": f"Bearer {token}"}
         
         # Convert images to bytes
         img_buffer = io.BytesIO()
@@ -234,6 +235,13 @@ with gr.Blocks(
                     info="Be specific about the desired background/replacement"
                 )
             
+            hf_token = gr.Textbox(
+                label="🔑 Hugging Face Token (Optional)",
+                type="password",
+                placeholder="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+                info="Get token from https://huggingface.co/settings/tokens (or set HF_TOKEN in Space secrets)"
+            )
+            
             remove_btn = gr.Button("🚀 Remove Objects", variant="primary", size="lg")
         
         with gr.Column(scale=2):
@@ -266,7 +274,8 @@ with gr.Blocks(
             object_class,
             confidence_threshold,
             mask_expansion,
-            inpaint_prompt
+            inpaint_prompt,
+            hf_token
         ],
         outputs=[output_image, mask_image, status_text]
     )
@@ -307,7 +316,8 @@ with gr.Blocks(
             **And 60+ more COCO classes!**
             
             ### ⚠️ Important Notes:
-            - This Space requires a Hugging Face token to be set in the environment
+            - **Token Required**: Either enter your HF token above OR set `HF_TOKEN` in Space secrets
+            - **Get Token**: Visit https://huggingface.co/settings/tokens to create one
             - Processing may take 30-60 seconds depending on image size
             - Results depend on object detection accuracy and image complexity
             - Red overlay indicates detected areas when inpainting fails