Update app.py

app.py

@@ -1,8 +1,9 @@
-import spaces  # Import spaces first
+import os
 import cv2
 import torch
 import numpy as np
 import gradio as gr
+from huggingface_hub import hf_hub_download
 from depth_anything_v2.dpt import DepthAnythingV2
 
 # Model initialization
@@ -175,15 +176,22 @@ def process_normal_map(image):
     result = (converted_normal.squeeze(0).numpy() * 255).astype(np.uint8)
     return result
 
-# Initialize model globally
+# Download and initialize model
 def initialize_model():
     encoder = 'vitl'
     max_depth = 1
     
     model = DepthAnythingV2(**{**model_configs[encoder], 'max_depth': max_depth})
     
+    # Download model from private repo
+    model_path = hf_hub_download(
+        "NightRaven109/DepthAnythingv2custom",
+        "model2.pth",
+        use_auth_token=os.environ['Read']
+    )
+    
     # Load checkpoint
-    checkpoint = torch.load('checkpoints/model2.pth', map_location='cpu')
+    checkpoint = torch.load(model_path, map_location='cpu')
     
     # Get state dict
     state_dict = {}
@@ -200,15 +208,15 @@ def initialize_model():
     model.load_state_dict(my_state_dict)
     return model
 
+# Initialize model at startup
 MODEL = initialize_model()
 
-@spaces.GPU
 def process_image(input_image):
     """
-    Process the input image and return depth maps, intensity map, blended result, and normal map
+    Process the input image and return depth map and normal map
     """
     if input_image is None:
-        return None, None, None, None, None
+        return None, None
     
     # Move model to GPU for processing
     MODEL.to('cuda')
@@ -223,10 +231,6 @@ def process_image(input_image):
         
         # Normalize depth for visualization (0-255)
         depth_normalized = ((depth - depth.min()) / (depth.max() - depth.min()) * 255).astype(np.uint8)
-        
-        # Apply colormap for better visualization
-        depth_colormap = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_INFERNO)
-        depth_colormap = cv2.cvtColor(depth_colormap, cv2.COLOR_BGR2RGB)
     
     # Move model back to CPU after processing
     MODEL.to('cpu')
@@ -245,16 +249,15 @@ def process_image(input_image):
     # Generate normal map from blended result
     normal_map = process_normal_map(blended_result)
     
-    return depth_normalized, depth_colormap, intensity_map, blended_result, normal_map
+    return depth_normalized, normal_map
 
-@spaces.GPU
 def gradio_interface(input_img):
     try:
-        depth_raw, depth_colored, intensity, blended, normal = process_image(input_img)
-        return [input_img, depth_raw, depth_colored, intensity, blended, normal]
+        depth_raw, normal = process_image(input_img)
+        return [depth_raw, normal]
     except Exception as e:
         print(f"Error processing image: {str(e)}")
-        return [input_img, None, None, None, None, None]
+        return [None, None]
 
 # Define interface
 iface = gr.Interface(
@@ -264,8 +267,8 @@ iface = gr.Interface(
         gr.Image(label="Raw Depth Map"),
         gr.Image(label="Normal Map")
     ],
-    title="Depth, Intensity, and Normal Map Estimation",
-    description="Upload an image to generate its depth map, intensity map, blended result, and normal map.",
+    title="Depth and Normal Map Generation",
+    description="Upload an image to generate its depth map and normal map.",
     examples=["image.jpg"]
 )