Update app.py

app.py

@@ -7,25 +7,24 @@ from io import BytesIO
 from transformers import AutoModelForImageSegmentation
 import uvicorn
 from fastapi.middleware.cors import CORSMiddleware
-import subprocess
-import threading
 
 app = FastAPI()
+
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # For production, specify your frontend domain
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Load model
+# Load model directly from Hugging Face with trust_remote_code=True
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModelForImageSegmentation.from_pretrained('ZhengPeng7/BiRefNet', trust_remote_code=True)
 model.to(device)
 model.eval()
 
-# Define transforms
+# Define transforms based on the other implementation
 transform = transforms.Compose([
     transforms.Resize((1024, 1024)),
     transforms.ToTensor(),
@@ -41,8 +40,9 @@ async def remove_bg(file: UploadFile = File(...)):
     # Transform the image
     input_tensor = transform(image).unsqueeze(0).to(device)
     
-    # Perform inference
+    # Perform inference - notice the key differences here
     with torch.no_grad():
+        # Take the last element [-1] from the model output and apply sigmoid
         pred = model(input_tensor)[-1].sigmoid().cpu()[0].squeeze()
     
     # Convert to PIL image and resize back to original size
@@ -60,11 +60,11 @@ async def remove_bg(file: UploadFile = File(...)):
     
     return Response(content=img_io.getvalue(), media_type="image/png")
 
+# Add a simple root route for health check
 @app.get("/")
 def read_root():
     return {"status": "ok", "message": "Background removal API is running"}
 
+# Make sure this is included for Hugging Face Spaces
 if __name__ == "__main__":
-    
-    # Start the FastAPI application
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=7860)