Update app.py

app.py

@@ -66,7 +66,7 @@ def download_models():
 download_models()
 
 # DepthAnythingV2
-if 'dam2' not in globals():
+def get_dam2_model():
     model_configs = {
         'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
         'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
@@ -83,22 +83,25 @@ if 'dam2' not in globals():
     dam2_checkpoint = f'checkpoints/depth_anything_v2_{encoder}.pth'
     dam2.load_state_dict(torch.load(dam2_checkpoint, map_location='cpu'))
     dam2 = dam2.to(DEVICE).eval()
+    return dam2
 
 # GenStereo
-if 'genstereo' not in globals():
+def get_genstereo_model():
     genwarp_cfg = dict(
         pretrained_model_path='checkpoints',
         checkpoint_name=CHECKPOINT_NAME,
         half_precision_weights=True
     )
     genstereo = GenStereo(cfg=genwarp_cfg, device=DEVICE)
+    return genstereo
 
 # Adaptive Fusion
-if 'fusion_model' not in globals():
+def get_fusion_model():
     fusion_model = AdaptiveFusionLayer()
     fusion_checkpoint = join('checkpoints', CHECKPOINT_NAME, 'fusion_layer.pth')
     fusion_model.load_state_dict(torch.load(fusion_checkpoint, map_location='cpu'))
     fusion_model = fusion_model.to(DEVICE).eval()
+    return fusion_model
 
 # Crop the image to the shorter side.
 def crop(img: Image) -> Image:
@@ -190,6 +193,7 @@ with tempfile.TemporaryDirectory() as tmpdir:
 
             image_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
+            dam2 = get_dam2_model()
             depth_dam2 = dam2.infer_image(image_bgr)
             depth = torch.tensor(depth_dam2).unsqueeze(0).unsqueeze(0).float().cuda()
 
@@ -202,6 +206,9 @@ with tempfile.TemporaryDirectory() as tmpdir:
             norm_disp = normalize_disp(depth)
             disp = norm_disp * scale_factor / 100 * IMAGE_SIZE
 
+            genstereo = get_genstereo_model()
+            fusion_model = get_fusion_model()
+
             renders = genstereo(
                 src_image=image,
                 src_disparity=disp,
@@ -231,4 +238,4 @@ with tempfile.TemporaryDirectory() as tmpdir:
         )
 
     if __name__ == '__main__':
-        demo.launch()
+        demo.launch()