Spaces:
Running
on
Zero
Running
on
Zero
add gpu for spatracker
Browse files- app.py +14 -14
- 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 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
42 |
-
|
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=
|
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=
|
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=
|
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=
|
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=
|
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 |
|