Update app.py

app.py

@@ -4,13 +4,28 @@ import numpy as np
 from transformers import DPTForDepthEstimation, DPTImageProcessor
 import gradio as gr
 import torch.quantization
+import torch.nn.utils.prune as prune
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 model = DPTForDepthEstimation.from_pretrained("Intel/dpt-swinv2-tiny-256", torch_dtype=torch.float32)
 model.eval()
+
+# Apply global unstructured pruning
+parameters_to_prune = [
+    (module, "weight") for module in filter(lambda m: isinstance(m, (torch.nn.Conv2d, torch.nn.Linear)), model.modules())
+]
+prune.global_unstructured(
+    parameters_to_prune,
+    pruning_method=prune.L1Unstructured,
+    amount=0.4,  # Prune 40% of weights
+)
+
+# Apply quantization after pruning
 model = torch.quantization.quantize_dynamic(
     model, {torch.nn.Linear, torch.nn.Conv2d}, dtype=torch.qint8
 ).to(device)
+
 processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
 color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)