add gpu for zerogpu

app.py

@@ -8,6 +8,7 @@ import numpy as np
 import torchvision.transforms as transforms
 from moviepy.editor import VideoFileClip
 from diffusers.utils import load_image, load_video
+import spaces
 
 project_root = os.path.dirname(os.path.abspath(__file__))
 os.environ["GRADIO_TEMP_DIR"] = os.path.join(project_root, "tmp", "gradio")
@@ -131,12 +132,14 @@ def save_uploaded_file(file):
 das_pipeline = None
 moge_model = None
 
+@spaces.GPU
 def get_das_pipeline():
     global das_pipeline
     if das_pipeline is None:
         das_pipeline = DiffusionAsShaderPipeline(gpu_id=GPU_ID, output_dir=OUTPUT_DIR)
     return das_pipeline
 
+@spaces.GPU
 def get_moge_model():
     global moge_model
     if moge_model is None:
@@ -144,6 +147,7 @@ def get_moge_model():
         moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(das.device)
     return moge_model
 
+
 def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image):
     """Process video motion transfer task"""
     try:
@@ -217,6 +221,7 @@ def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image)
         print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
         return None
 
+
 def process_camera_control(source, prompt, camera_motion, tracking_method):
     """Process camera control task"""
     try:
@@ -288,6 +293,7 @@ def process_camera_control(source, prompt, camera_motion, tracking_method):
         print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
         return None
 
+
 def process_object_manipulation(source, prompt, object_motion, object_mask, tracking_method):
     """Process object manipulation task"""
     try:
@@ -381,6 +387,7 @@ def process_object_manipulation(source, prompt, object_motion, object_mask, trac
         print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
         return None
 
+
 def process_mesh_animation(source, prompt, tracking_video, ma_repaint_option, ma_repaint_image):
     """Process mesh animation task"""
     try:

models/pipelines.py

@@ -24,6 +24,7 @@ from models.cogvideox_tracking import CogVideoXImageToVideoPipelineTracking
 from submodules.MoGe.moge.model import MoGeModel
 from image_gen_aux import DepthPreprocessor
 from moviepy.editor import ImageSequenceClip
+import spaces
 
 class DiffusionAsShaderPipeline:
     def __init__(self, gpu_id=0, output_dir='outputs'):
@@ -55,6 +56,7 @@ class DiffusionAsShaderPipeline:
             transforms.ToTensor()
         ])
 
+    @spaces.GPU(duration=240)
     @torch.no_grad()
     def _infer(
         self, 
@@ -181,6 +183,7 @@ class DiffusionAsShaderPipeline:
 
         return intr
 
+    @spaces.GPU
     def _apply_poses(self, pts, intr, poses):
         """
         Args:
@@ -217,6 +220,7 @@ class DiffusionAsShaderPipeline:
         return tracking_pts
     
     ##============= SpatialTracker =============##
+    @spaces.GPU
     def generate_tracking_spatracker(self, video_tensor, density=70):
         """Generate tracking video
         
@@ -271,6 +275,7 @@ class DiffusionAsShaderPipeline:
             del tracker, self.depth_preprocessor
             torch.cuda.empty_cache()
 
+    @spaces.GPU
     def visualize_tracking_spatracker(self, video, pred_tracks, pred_visibility, T_Firsts, save_tracking=True):
         video = video.unsqueeze(0).to(self.device)
         vis = Visualizer(save_dir=self.output_dir, grayscale=False, fps=24, pad_value=0)
@@ -360,6 +365,7 @@ class DiffusionAsShaderPipeline:
             outline=tuple(color),
         )
     
+    @spaces.GPU
     def visualize_tracking_moge(self, points, mask, save_tracking=True):
         """Visualize tracking results from MoGe model
         
@@ -446,6 +452,7 @@ class DiffusionAsShaderPipeline:
 
         return tracking_path, tracking_video
     
+    @spaces.GPU(duration=240)
     def apply_tracking(self, video_tensor, fps=8, tracking_tensor=None, img_cond_tensor=None, prompt=None, checkpoint_path=None):
         """Generate final video with motion transfer
         
@@ -486,6 +493,7 @@ class DiffusionAsShaderPipeline:
         """
         self.object_motion = motion_type
 
+@spaces.GPU(duration=120)
 class FirstFrameRepainter:
     def __init__(self, gpu_id=0, output_dir='outputs'):
         """Initialize FirstFrameRepainter
@@ -498,7 +506,8 @@ class FirstFrameRepainter:
         self.output_dir = output_dir
         self.max_depth = 65.0
         os.makedirs(output_dir, exist_ok=True)
-        
+    
+    @spaces.GPU(duration=120)
     def repaint(self, image_tensor, prompt, depth_path=None, method="dav"):
         """Repaint first frame using Flux