EXCAI commited on
Commit
0932790
·
1 Parent(s): f6c79a3

add gpu for spatracker

Browse files
Files changed (2) hide show
  1. app.py +14 -14
  2. models/pipelines.py +1 -1
app.py CHANGED
@@ -31,15 +31,15 @@ GPU_ID = args.gpu
31
  DEFAULT_MODEL_PATH = args.model_path
32
  OUTPUT_DIR = args.output_dir
33
 
34
- if 'CUDA_HOME' not in os.environ:
35
- for cuda_path in ['/usr/local/cuda', '/usr/cuda', '/opt/cuda']:
36
- if os.path.exists(cuda_path):
37
- os.environ['CUDA_HOME'] = cuda_path
38
- print(cuda_path)
39
- break
40
- if 'CUDA_HOME' not in os.environ:
41
- os.environ['CUDA_HOME'] = '/usr/local/cuda'
42
- print("set default cuda path in: /usr/local/cuda")
43
 
44
  # Create necessary directories
45
  os.makedirs("outputs", exist_ok=True)
@@ -110,7 +110,7 @@ def create_run_command(args):
110
 
111
  return cmd
112
 
113
- @spaces.GPU(duration=280)
114
  def run_process(cmd):
115
  """Run command and return output"""
116
  print(f"Running command: {' '.join(cmd)}")
@@ -138,7 +138,7 @@ def run_process(cmd):
138
 
139
  return "\n".join(output)
140
 
141
- @spaces.GPU(duration=280)
142
  def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image):
143
  """Process video motion transfer task"""
144
  try:
@@ -185,7 +185,7 @@ def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image)
185
  print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
186
  return None
187
 
188
- @spaces.GPU(duration=280)
189
  def process_camera_control(source, prompt, camera_motion, tracking_method):
190
  """Process camera control task"""
191
  try:
@@ -226,7 +226,7 @@ def process_camera_control(source, prompt, camera_motion, tracking_method):
226
  print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
227
  return None
228
 
229
- @spaces.GPU(duration=280)
230
  def process_object_manipulation(source, prompt, object_motion, object_mask, tracking_method):
231
  """Process object manipulation task"""
232
  try:
@@ -265,7 +265,7 @@ def process_object_manipulation(source, prompt, object_motion, object_mask, trac
265
  print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
266
  return None
267
 
268
- @spaces.GPU(duration=280)
269
  def process_mesh_animation(source, prompt, tracking_video, ma_repaint_option, ma_repaint_image):
270
  """Process mesh animation task"""
271
  try:
 
31
  DEFAULT_MODEL_PATH = args.model_path
32
  OUTPUT_DIR = args.output_dir
33
 
34
+ # if 'CUDA_HOME' not in os.environ:
35
+ # for cuda_path in ['/usr/local/cuda', '/usr/cuda', '/opt/cuda']:
36
+ # if os.path.exists(cuda_path):
37
+ # os.environ['CUDA_HOME'] = cuda_path
38
+ # print(cuda_path)
39
+ # break
40
+ # if 'CUDA_HOME' not in os.environ:
41
+ # os.environ['CUDA_HOME'] = '/usr/local/cuda'
42
+ # print("set default cuda path in: /usr/local/cuda")
43
 
44
  # Create necessary directories
45
  os.makedirs("outputs", exist_ok=True)
 
110
 
111
  return cmd
112
 
113
+ @spaces.GPU(duration=270)
114
  def run_process(cmd):
115
  """Run command and return output"""
116
  print(f"Running command: {' '.join(cmd)}")
 
138
 
139
  return "\n".join(output)
140
 
141
+ @spaces.GPU(duration=270)
142
  def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image):
143
  """Process video motion transfer task"""
144
  try:
 
185
  print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
186
  return None
187
 
188
+ @spaces.GPU(duration=270)
189
  def process_camera_control(source, prompt, camera_motion, tracking_method):
190
  """Process camera control task"""
191
  try:
 
226
  print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
227
  return None
228
 
229
+ @spaces.GPU(duration=270)
230
  def process_object_manipulation(source, prompt, object_motion, object_mask, tracking_method):
231
  """Process object manipulation task"""
232
  try:
 
265
  print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
266
  return None
267
 
268
+ @spaces.GPU(duration=270)
269
  def process_mesh_animation(source, prompt, tracking_video, ma_repaint_option, ma_repaint_image):
270
  """Process mesh animation task"""
271
  try:
models/pipelines.py CHANGED
@@ -217,7 +217,7 @@ class DiffusionAsShaderPipeline:
217
  return tracking_pts
218
 
219
  ##============= SpatialTracker =============##
220
-
221
  def generate_tracking_spatracker(self, video_tensor, density=70):
222
  """Generate tracking video
223
 
 
217
  return tracking_pts
218
 
219
  ##============= SpatialTracker =============##
220
+ @spaces.GPU
221
  def generate_tracking_spatracker(self, video_tensor, density=70):
222
  """Generate tracking video
223