Update app.py

app.py

@@ -1,6 +1,92 @@
 # Import app from main.py for actual background removal
 try:
-    from main import app
+    from main import app as main_app
+    import os
+    import secrets
+    import json
+    import uuid
+    from fastapi import FastAPI, Depends, HTTPException, Header, status, File, UploadFile
+    
+    # Create the FastAPI app
+    app = FastAPI(title="Background Removal API with API Key Support")
+    
+    # API key storage file
+    API_KEYS_FILE = "api_keys.json"
+    
+    # Initialize API keys from file or create default
+    def init_api_keys():
+        if os.path.exists(API_KEYS_FILE):
+            try:
+                with open(API_KEYS_FILE, "r") as f:
+                    return json.load(f)
+            except:
+                # If file is corrupted, create new
+                pass
+                
+        # Generate a default API key if none exists
+        default_key = secrets.token_urlsafe(16)
+        keys = {"default": {"key": default_key, "active": True}}
+        with open(API_KEYS_FILE, "w") as f:
+            json.dump(keys, f)
+        print(f"Created default API key: {default_key}")
+        return keys
+    
+    # API keys dictionary
+    API_KEYS = init_api_keys()
+    
+    # Function to verify API key
+    def verify_api_key(api_key: str = Header(None, alias="X-API-Key")):
+        if api_key is None:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="API key is missing"
+            )
+            
+        for key_info in API_KEYS.values():
+            if key_info["key"] == api_key and key_info["active"]:
+                return True
+                
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API key"
+        )
+    
+    # Root endpoint
+    @app.get("/")
+    async def root():
+        return {
+            "message": "Background Removal API with API Key Support",
+            "docs": "/docs",
+            "endpoints": ["/remove-bg", "/api-key"]
+        }
+    
+    # Generate API key endpoint
+    @app.get("/api-key")
+    async def generate_api_key():
+        new_key = secrets.token_urlsafe(16)
+        key_id = str(uuid.uuid4())[:8]
+        API_KEYS[key_id] = {"key": new_key, "active": True}
+        
+        # Save to file
+        with open(API_KEYS_FILE, "w") as f:
+            json.dump(API_KEYS, f)
+            
+        return {"api_key": new_key, "message": "Keep this key secure. You'll need it for API calls."}
+    
+    # Import the background removal function from main
+    from main import remove_background
+    
+    # Protect the remove-bg endpoint with API key
+    @app.post("/remove-bg")
+    async def protected_remove_background(api_key_valid: bool = Depends(verify_api_key), file: UploadFile = File(...)):
+        # Forward to the original function
+        return await remove_background(file=file)
+    
+    # Print the default API key for testing
+    if "default" in API_KEYS:
+        print(f"Default API key for testing: {API_KEYS['default']['key']}")
+        print(f"To use: Add 'X-API-Key: {API_KEYS['default']['key']}' to your request headers")
+    
 except ImportError:
     # Fallback to the simple version if the main version fails
     from fallback import app