Update app.py

app.py

@@ -19,6 +19,11 @@ from transformers import AutoModelForImageSegmentation, AutoImageProcessor, Auto
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 torch.set_float32_matmul_precision('high')
 
+# Load models at startup
+rmbg_model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True).to(device).eval()
+depth_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
+depth_model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf").to(device)
+
 def load_image_from_link(url: str) -> Image.Image:
     """Downloads an image from a URL and returns a Pillow Image."""
     response = requests.get(url)
@@ -29,10 +34,6 @@ def load_image_from_link(url: str) -> Image.Image:
 # Gaussian Blur Functions
 def run_rmbg(image: Image.Image, threshold=0.5):
     """Runs the RMBG-2.0 model on the image and returns a binary mask."""
-    model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True)
-    model.to(device)
-    model.eval()
-
     image_size = (1024, 1024)
     transform_image = transforms.Compose([
         transforms.Resize(image_size),
@@ -43,7 +44,7 @@ def run_rmbg(image: Image.Image, threshold=0.5):
     input_images = transform_image(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        preds = model(input_images)
+        preds = rmbg_model(input_images)
         if isinstance(preds, list):
             mask_logits = preds[-1]
         else:
@@ -72,15 +73,11 @@ def apply_background_blur(image: Image.Image, mask: np.ndarray, sigma: float = 1
 # Lens Blur Functions
 def run_depth_estimation(image: Image.Image, target_size=(512, 512)):
     """Runs the Depth-Anything-V2-Small model and returns the depth map."""
-    image_processor = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
-    model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
-    model.to(device)
-
     image_resized = image.resize(target_size, resample=Image.BILINEAR)
-    inputs = image_processor(images=image_resized, return_tensors="pt").to(device)
+    inputs = depth_processor(images=image_resized, return_tensors="pt").to(device)
 
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = depth_model(**inputs)
         predicted_depth = outputs.predicted_depth
 
     prediction = torch.nn.functional.interpolate(