Update app.py

app.py

@@ -1,39 +1,83 @@
 import gradio as gr
-import requests
 import io
+import torch
+import numpy as np
 from PIL import Image
+import os
+import sys
 
-def remove_background(image):
-    try:
-        # Convert image to bytes
-        img_bytes = io.BytesIO()
-        image.save(img_bytes, format="PNG")
-        img_bytes = img_bytes.getvalue()
-        
-        # Call your local API (will change to HF URL later)
-        response = requests.post(
-            "http://localhost:8001/api/v1/background-removal/remove-background",
-            files={"file": ("image.png", img_bytes, "image/png")}
-        )
+# Add current directory to path for model files
+sys.path.append("/app")
+
+# Import model components
+from briarmbg import BriaRMBG
+from utilities import preprocess_image, postprocess_image
+
+class BackgroundRemover:
+    def __init__(self):
+        self.model = None
+        self.device = None
+        self.load_model()
+    
+    def load_model(self):
+        """Load the RMBG-1.4 model"""
+        try:
+            print("🔄 Loading background removal model...")
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.model = BriaRMBG.from_pretrained("/app")
+            self.model.to(self.device)
+            self.model.eval()
+            print("✅ Model loaded successfully!")
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            self.model = None
+    
+    def remove_background(self, image):
+        """Remove background from image"""
+        if self.model is None:
+            raise Exception("Model not loaded")
         
-        if response.status_code == 200:
-            # Get the response and return image
-            result_data = response.json()
-            download_url = result_data["download_url"]
+        try:
+            # Convert to RGB if needed
+            input_image = image.convert("RGB")
+            
+            # Preprocess
+            model_input_size = [1024, 1024]
+            orig_im = np.array(input_image)
+            orig_im_size = orig_im.shape[0:2]
+            processed_image = preprocess_image(orig_im, model_input_size).to(self.device)
+            
+            # Inference
+            with torch.no_grad():
+                result = self.model(processed_image)
             
-            # Download the processed image
-            result_response = requests.get(f"http://localhost:8001{download_url}")
-            result_bytes = io.BytesIO(result_response.content)
-            return Image.open(result_bytes)
-        else:
-            raise Exception(f"API call failed: {response.status_code}")
+            # Postprocess
+            result_image = postprocess_image(result[0][0], orig_im_size)
             
+            # Create transparent image
+            pil_mask = Image.fromarray(result_image)
+            no_bg_image = input_image.copy()
+            no_bg_image.putalpha(pil_mask)
+            
+            return no_bg_image
+            
+        except Exception as e:
+            raise Exception(f"Background removal failed: {str(e)}")
+
+# Initialize the remover
+remover = BackgroundRemover()
+
+def process_image(image):
+    """Gradio interface function"""
+    try:
+        result = remover.remove_background(image)
+        return result
     except Exception as e:
-        raise gr.Error(f"Background removal failed: {str(e)}")
+        raise gr.Error(str(e))
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=remove_background,
+    fn=process_image,
     inputs=gr.Image(type="pil", label="📷 Upload Image"),
     outputs=gr.Image(type="pil", label="🎨 Background Removed"),
     title="🎨 Professional Background Remover",