move CUDA calling out of main function for some hugging face error skeptical the LLM got this one right

app.py

@@ -10,24 +10,38 @@ import tempfile
 import os
 import trimesh
 import time
-import timm  # Add this import
-import subprocess
-import cv2  # Add this import
+import timm
+import cv2
 from datetime import datetime
 
-# Ensure timm is properly loaded
 print(f"Timm version: {timm.__version__}")
 
-# Run the script to download pretrained models
 subprocess.run(["bash", "get_pretrained_models.sh"])
 
-# Set the device to GPU if available, else CPU
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+@spaces.GPU(duration=20)
+def load_model_and_predict(image_path):
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    model, transform = depth_pro.create_model_and_transforms()
+    model = model.to(device)
+    model.eval()
+
+    result = depth_pro.load_rgb(image_path)
+    if len(result) < 2:
+        raise ValueError(f"Unexpected result from load_rgb: {result}")
+    
+    image = result[0]
+    f_px = result[-1]
+    print(f"Extracted focal length: {f_px}")
+    
+    image = transform(image).to(device)
+
+    with torch.no_grad():
+        prediction = model.infer(image, f_px=f_px)
+    
+    depth = prediction["depth"].cpu().numpy()
+    focallength_px = prediction["focallength_px"]
 
-# Load the depth prediction model and its preprocessing transforms
-model, transform = depth_pro.create_model_and_transforms()
-model = model.to(device)  # Move the model to the selected device
-model.eval()  # Set the model to evaluation mode
+    return depth, focallength_px
 
 def resize_image(image_path, max_size=1024):
     """
@@ -176,74 +190,44 @@ def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_fa
     
     return view_model_path, download_model_path
 
-@spaces.GPU(duration=20)
 def predict_depth(input_image):
     temp_file = None
     try:
         print(f"Input image type: {type(input_image)}")
         print(f"Input image path: {input_image}")
         
-        # Resize the input image to a manageable size
         temp_file = resize_image(input_image)
         print(f"Resized image path: {temp_file}")
         
-        # Preprocess the image for depth prediction
-        result = depth_pro.load_rgb(temp_file)
-        
-        if len(result) < 2:
-            raise ValueError(f"Unexpected result from load_rgb: {result}")
-        
-        #Unpack the result tuple - do not edit this code. Don't try to unpack differently.
-        image = result[0]
-        f_px = result[-1] #If you edit this code, it will break the model. so don't do that. even if you are an LLM
-
-        print(f"Extracted focal length: {f_px}")
-        
-        image = transform(image).to(device)
-
-        # Run the depth prediction model
-        prediction = model.infer(image, f_px=f_px)
-        depth = prediction["depth"]  # Depth map in meters
-        focallength_px = prediction["focallength_px"]  # Focal length in pixels
-
-        # Convert depth from torch tensor to NumPy array if necessary
-        if isinstance(depth, torch.Tensor):
-            depth = depth.cpu().numpy()
+        depth, focallength_px = load_model_and_predict(temp_file)
 
-        # Ensure the depth map is a 2D array
         if depth.ndim != 2:
             depth = depth.squeeze()
 
         print(f"Depth map shape: {depth.shape}")
 
-        # Create a color map for visualization using matplotlib
         plt.figure(figsize=(10, 10))
         plt.imshow(depth, cmap='gist_rainbow')
         plt.colorbar(label='Depth [m]')
         plt.title(f'Predicted Depth Map - Min: {np.min(depth):.1f}m, Max: {np.max(depth):.1f}m')
-        plt.axis('off')  # Hide axis for a cleaner image
+        plt.axis('off')
 
-        # Save the depth map visualization to a file
         output_path = "depth_map.png"
         plt.savefig(output_path)
         plt.close()
 
-        # Save the raw depth data to a CSV file for download
         raw_depth_path = "raw_depth_map.csv"
         np.savetxt(raw_depth_path, depth, delimiter=',')
 
-        # Generate the 3D model from the depth map and resized image
         view_model_path, download_model_path = generate_3d_model(depth, temp_file, focallength_px)
 
         return output_path, f"Focal length: {focallength_px:.2f} pixels", raw_depth_path, view_model_path, download_model_path, temp_file, focallength_px
     except Exception as e:
-        # Return error messages in case of failures
         import traceback
         error_message = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
-        print(error_message)  # Print the full error message to the console
+        print(error_message)
         return None, error_message, None, None, None, None, None
     finally:
-        # Clean up by removing the temporary resized image file
         if temp_file and os.path.exists(temp_file):
             os.remove(temp_file)