Update parallax.py

parallax.py

@@ -13,7 +13,6 @@ load_dotenv(verbose=False)
 
 DEPTH_ANYTHING = DepthAnythingV2(**{'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]})
 DEPTH_ANYTHING.load_state_dict(torch.load(hf_hub_download(repo_id='depth-anything/Depth-Anything-V2-Large', filename='depth_anything_v2_vitl.pth', repo_type='model', token=os.environ['HF_TOKEN']), map_location='cpu'))
-DEPTH_ANYTHING = DEPTH_ANYTHING.to('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu').eval()
 HAND_YOLO = YOLO(hf_hub_download('Bingsu/adetailer', 'hand_yolov8n.pt', token=os.environ['HF_TOKEN']))
 PERSON_YOLO = YOLO(hf_hub_download('Bingsu/adetailer', 'person_yolov8n-seg.pt', token=os.environ['HF_TOKEN']))
 LAMA = None
@@ -233,6 +232,8 @@ def feather(image: Image.Image, gauss_radius=1, band_px=1, strength=1.0) -> Imag
 
 
 def generate_parallax_images(image, n_layers=5, maximum=2048):
+    global LAMA
+    
     rgb_image = resize_iamge(image.convert('RGB'), maximum)
     width, height = rgb_image.size
     rgb = np.asarray(rgb_image)
@@ -240,7 +241,8 @@ def generate_parallax_images(image, n_layers=5, maximum=2048):
     if LAMA is None:
         LAMA = SimpleLama(device='cuda' if torch.cuda.is_available() else 'cpu')
 
-    depth = DEPTH_ANYTHING.infer_image(rgb[:, :, ::-1])
+    depth_anything = DEPTH_ANYTHING.to('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu').eval()
+    depth = depth_anything.infer_image(rgb[:, :, ::-1])
 
     n_clusters = n_layers
     x = depth.reshape(-1, 1)
@@ -269,8 +271,8 @@ def generate_parallax_images(image, n_layers=5, maximum=2048):
 
     inpaint_mask = np.zeros_like(front_mask, dtype=bool)
 
-    person_results = PERSON_YOLO.predict(source=rgb, conf=0.5, iou=0.45, verbose=False)
-    hand_results = HAND_YOLO.predict(source=rgb, conf=0.5, iou=0.45, verbose=False)
+    person_results = PERSON_YOLO.predict(source=rgb, conf=0.5, iou=0.45, verbose=False, device='0' if torch.cuda.is_available() else 'cpu')
+    hand_results = HAND_YOLO.predict(source=rgb, conf=0.5, iou=0.45, verbose=False, device='0' if torch.cuda.is_available() else 'cpu')
     person_boxes = []
     hand_boxes = []