add spaces decorator

app.py

@@ -6,11 +6,25 @@ import random
 import string
 import json
 from omegaconf import OmegaConf,ListConfig
+import spaces
 
 
 from train import main as train_main
 from inference import inference as inference_main
-# 模拟训练函数
+
+@spaces.GPU()
+def inference_app(
+        embedding_dir,
+        prompt, 
+        video_round,
+        save_dir,
+        motion_type,
+        seed,
+        inference_steps):
+    
+    return inference_main(embedding_dir, prompt, video_round, save_dir, motion_type, seed, inference_steps)
+
+
 def train_model(video, config):
     output_dir = 'results'
     os.makedirs(output_dir, exist_ok=True)
@@ -28,7 +42,7 @@ def train_model(video, config):
     # cur_save_dir = 'results/06'
     return cur_save_dir
 
-# 模拟推理函数
+
 def inference_model(text, checkpoint, inference_steps, video_type,seed):
     
     checkpoint = os.path.join('results',checkpoint)
@@ -36,7 +50,7 @@ def inference_model(text, checkpoint, inference_steps, video_type,seed):
     embedding_dir = '/'.join(checkpoint.split('/')[:-1])
     video_round = checkpoint.split('/')[-1]
 
-    video_path = inference_main(
+    video_path = inference_app(
         embedding_dir=embedding_dir,
         prompt=text, 
         video_round=video_round,
@@ -49,7 +63,6 @@ def inference_model(text, checkpoint, inference_steps, video_type,seed):
     return video_path
 
 
-# 获取checkpoint文件列表
 def get_checkpoints(checkpoint_dir):
     
     checkpoints = []
@@ -135,7 +148,6 @@ if __name__ == "__main__":
     ]
 
     gradio_theme = gr.themes.Default()
-    # 创建Gradio界面
     with gr.Blocks(
         theme=gradio_theme,
         title="Motion Inversion",
@@ -249,5 +261,5 @@ if __name__ == "__main__":
             checkpoint_output.change(update_checkpoints, inputs=checkpoint_output, outputs=checkpoint_dropdown)
             inference_button.click(inference_model, inputs=[text_input, checkpoint_dropdown,inference_steps,motion_type, seed], outputs=output_video)
             output_video.change(fn=update_generated_prompt, inputs=[text_input], outputs=generated_prompt)
-    # 启动Gradio界面
+
     demo.launch()