A19grey commited on
Commit
47669a5
1 Parent(s): 39f1439

more CUDA initialization fixes via perplexity

Browse files
Files changed (1) hide show
  1. app.py +10 -5
app.py CHANGED
@@ -1,18 +1,20 @@
 
1
  import gradio as gr
2
  from PIL import Image
3
- import src.depth_pro as depth_pro
4
  import numpy as np
5
  import matplotlib.pyplot as plt
6
  import subprocess
7
- import spaces
8
- import torch
9
  import tempfile
10
  import os
11
  import trimesh
12
  import time
 
 
 
 
 
13
  import timm
14
  import cv2
15
- from datetime import datetime
16
 
17
  print(f"Timm version: {timm.__version__}")
18
 
@@ -20,7 +22,7 @@ subprocess.run(["bash", "get_pretrained_models.sh"])
20
 
21
  @spaces.GPU(duration=20)
22
  def load_model_and_predict(image_path):
23
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
24
  model, transform = depth_pro.create_model_and_transforms()
25
  model = model.to(device)
26
  model.eval()
@@ -68,6 +70,7 @@ def resize_image(image_path, max_size=1024):
68
  img.save(temp_file, format="PNG")
69
  return temp_file.name
70
 
 
71
  def generate_3d_model(depth, image_path, focallength_px, simplification_factor=0.8, smoothing_iterations=1, thin_threshold=0.01):
72
  """
73
  Generate a textured 3D mesh from the depth map and the original image.
@@ -178,6 +181,7 @@ def remove_thin_features(mesh, thickness_threshold=0.01):
178
 
179
  return mesh
180
 
 
181
  def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_factor, smoothing_iterations, thin_threshold):
182
  # Load depth from CSV
183
  depth = np.loadtxt(depth_csv, delimiter=',')
@@ -190,6 +194,7 @@ def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_fa
190
 
191
  return view_model_path, download_model_path
192
 
 
193
  def predict_depth(input_image):
194
  temp_file = None
195
  try:
 
1
+ import spaces
2
  import gradio as gr
3
  from PIL import Image
 
4
  import numpy as np
5
  import matplotlib.pyplot as plt
6
  import subprocess
 
 
7
  import tempfile
8
  import os
9
  import trimesh
10
  import time
11
+ from datetime import datetime
12
+
13
+ # Import potentially CUDA-initializing modules after 'spaces'
14
+ import torch
15
+ import src.depth_pro as depth_pro
16
  import timm
17
  import cv2
 
18
 
19
  print(f"Timm version: {timm.__version__}")
20
 
 
22
 
23
  @spaces.GPU(duration=20)
24
  def load_model_and_predict(image_path):
25
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
26
  model, transform = depth_pro.create_model_and_transforms()
27
  model = model.to(device)
28
  model.eval()
 
70
  img.save(temp_file, format="PNG")
71
  return temp_file.name
72
 
73
+ @spaces.GPU(duration=20)
74
  def generate_3d_model(depth, image_path, focallength_px, simplification_factor=0.8, smoothing_iterations=1, thin_threshold=0.01):
75
  """
76
  Generate a textured 3D mesh from the depth map and the original image.
 
181
 
182
  return mesh
183
 
184
+ @spaces.GPU(duration=20)
185
  def regenerate_3d_model(depth_csv, image_path, focallength_px, simplification_factor, smoothing_iterations, thin_threshold):
186
  # Load depth from CSV
187
  depth = np.loadtxt(depth_csv, delimiter=',')
 
194
 
195
  return view_model_path, download_model_path
196
 
197
+ @spaces.GPU(duration=20)
198
  def predict_depth(input_image):
199
  temp_file = None
200
  try: