Update app.py

app.py

@@ -29,9 +29,6 @@ model = torch.quantization.quantize_dynamic(
 
 model = model.to(device)
 
-# Convert model to TorchScript
-model = torch.jit.script(model)
-
 processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
 color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)
@@ -39,23 +36,29 @@ color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFER
 input_tensor = torch.zeros((1, 3, 128, 128), dtype=torch.float32, device=device)
 
 def preprocess_image(image):
-    return cv2.resize(image, (128, 128), interpolation=cv2.INTER_AREA).transpose(2, 0, 1).astype(np.float32) / 255.0
+    image = torch.from_numpy(image).to(device)
+    image = torch.nn.functional.interpolate(image.permute(2, 0, 1).unsqueeze(0), size=(128, 128), mode='bilinear', align_corners=False)
+    return (image.squeeze(0) / 255.0)
+
+static_input = torch.zeros((1, 3, 128, 128), device=device, dtype=torch.float16)  # Use float16 if using FP16
+g = torch.cuda.CUDAGraph()
+with torch.cuda.graph(g):
+    static_output = model(static_input)
 
 @torch.inference_mode()
 def process_frame(image):
     if image is None:
         return None
     preprocessed = preprocess_image(image)
-    input_tensor = torch.from_numpy(preprocessed).unsqueeze(0).to(device)
-    
-    predicted_depth = model(input_tensor).predicted_depth
-    depth_map = predicted_depth.squeeze().cpu().numpy()
+    static_input.copy_(preprocessed)
+    g.replay()
+    depth_map = static_output.predicted_depth.squeeze().cpu().numpy()
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
     depth_map = (depth_map * 255).astype(np.uint8)
     depth_map_colored = cv2.applyColorMap(depth_map, cv2.COLORMAP_INFERNO)
     
     return cv2.cvtColor(depth_map_colored, cv2.COLOR_BGR2RGB)
-
+    
 interface = gr.Interface(
     fn=process_frame,
     inputs=gr.Image(sources="webcam", streaming=True),