bugfix

app.py

@@ -317,6 +317,12 @@ def generate_tracking_cotracker(video_tensor, density=30):
 @spaces.GPU(duration=240)
 def apply_tracking_unified(video_tensor, tracking_tensor, repaint_img_tensor, prompt, fps):
     """统一的应用跟踪函数"""
+    print("--- Entering apply_tracking_unified ---") # 添加打印语句
+    print(f"Prompt received: {prompt}")             # 添加打印语句
+    print(f"FPS received: {fps}")                   # 添加打印语句
+    print(f"Video tensor shape: {video_tensor.shape if video_tensor is not None else None}") # 添加打印语句
+    print(f"Tracking tensor shape: {tracking_tensor.shape if tracking_tensor is not None else None}") # 添加打印语句
+    print(f"Repaint tensor shape: {repaint_img_tensor.shape if repaint_img_tensor is not None else None}") # 添加打印语句
     try:
         if video_tensor is None or tracking_tensor is None:
             return None
@@ -328,7 +334,8 @@ def apply_tracking_unified(video_tensor, tracking_tensor, repaint_img_tensor, pr
             tracking_tensor=tracking_tensor,
             img_cond_tensor=repaint_img_tensor,
             prompt=prompt,
-            checkpoint_path=DEFAULT_MODEL_PATH
+            checkpoint_path=DEFAULT_MODEL_PATH,
+            num_inference_steps=15
         )
         
         print(f"生成的视频路径: {output_path}")
@@ -464,27 +471,17 @@ def load_examples():
 with gr.Blocks(title="Diffusion as Shader") as demo:
     gr.Markdown("# Diffusion as Shader Web UI")
     gr.Markdown("### [Project Page](https://igl-hkust.github.io/das/) | [GitHub](https://github.com/IGL-HKUST/DiffusionAsShader)")
-    
+
     # 创建隐藏状态变量来存储中间结果
     video_tensor_state = gr.State(None)
     tracking_tensor_state = gr.State(None)
     repaint_img_tensor_state = gr.State(None)
     fps_state = gr.State(None)
-    
+
     with gr.Row():
         left_column = gr.Column(scale=1)
         right_column = gr.Column(scale=1)
 
-    with right_column:
-        
-        gr.Markdown("### 4. Generate Tracking Video")
-        gr.Markdown("'Generate Tracking Video' is used to preserve all motion from the Source. You need to generate tracking video before producing the final result.")
-        mt_run_btn = gr.Button("Generate Tracking", variant="primary", size="lg")
-        tracking_video = gr.Video(label="Tracking Video")
-        
-        apply_tracking_btn = gr.Button("5. Generate Video", variant="primary", size="lg", interactive=False)
-        output_video = gr.Video(label="Generated Video")
-
     with left_column:
         gr.Markdown("### 1. Upload Source")
         gr.Markdown("Upload a video or image, We will extract the motion and space structure from it")
@@ -531,25 +528,49 @@ with gr.Blocks(title="Diffusion as Shader") as demo:
                     outputs=[mt_repaint_preview]
                 )
 
-                mt_run_btn.click(
-                    fn=process_motion_transfer,
-                    inputs=[
-                        source_upload, common_prompt,
-                        mt_repaint_option, mt_repaint_upload
-                    ],
-                    outputs=[tracking_video, video_tensor_state, tracking_tensor_state, repaint_img_tensor_state, fps_state]
-                ).then(
-                    fn=enable_apply_button,
-                    inputs=[tracking_video],
-                    outputs=[apply_tracking_btn]
-                )
-                
             with gr.TabItem("Camera Control"):
                  gr.Markdown("Camera Control is not available in Huggingface Space, please deploy our [GitHub project](https://github.com/IGL-HKUST/DiffusionAsShader) on your own machine")
 
             with gr.TabItem("Object Manipulation"):
                  gr.Markdown("Object Manipulation is not available in Huggingface Space, please deploy our [GitHub project](https://github.com/IGL-HKUST/DiffusionAsShader) on your own machine")
 
+    with right_column:
+        
+        gr.Markdown("### 4. Generate Tracking Video")
+        gr.Markdown("'Generate Tracking Video' is used to preserve all motion from the Source. You need to generate tracking video before producing the final result.")
+        mt_run_btn = gr.Button("Generate Tracking", variant="primary", size="lg")
+        tracking_video = gr.Video(label="Tracking Video")
+        
+        apply_tracking_btn = gr.Button("5. Generate Video", variant="primary", size="lg", interactive=False)
+        output_video = gr.Video(label="Generated Video")
+
+        # mt_run_btn 的 click 事件定义
+        mt_run_btn.click(
+            fn=process_motion_transfer,
+            inputs=[
+                source_upload, common_prompt,
+                mt_repaint_option, mt_repaint_upload
+            ],
+            outputs=[tracking_video, video_tensor_state, tracking_tensor_state, repaint_img_tensor_state, fps_state]
+        ).then(
+            fn=enable_apply_button,
+            inputs=[tracking_video],
+            outputs=[apply_tracking_btn]
+        )
+
+        # apply_tracking_btn 的 click 事件定义
+        apply_tracking_btn.click(
+            fn=apply_tracking_unified,
+            inputs=[
+                video_tensor_state,
+                tracking_tensor_state,
+                repaint_img_tensor_state,
+                common_prompt,  # common_prompt 现在可用
+                fps_state
+            ],
+            outputs=[output_video]
+        )
+
     examples_list = load_examples()
     gr.Markdown("### Examples (For Workflow Demo Only)")
     gr.Markdown("The following examples are only for demonstrating DaS's workflow and output quality. If you want to actually generate tracking or videos, the program will not run unless you manually upload files from your devices.")

models/pipelines.py

@@ -570,7 +570,7 @@ class DiffusionAsShaderPipeline:
         return tracking_path, tracking_video
 
     
-    def apply_tracking(self, video_tensor, fps=8, tracking_tensor=None, img_cond_tensor=None, prompt=None, checkpoint_path=None):
+    def apply_tracking(self, video_tensor, fps=8, tracking_tensor=None, img_cond_tensor=None, prompt=None, checkpoint_path=None, num_inference_steps=15):
         """Generate final video with motion transfer
         
         Args:
@@ -595,7 +595,7 @@ class DiffusionAsShaderPipeline:
             tracking_tensor=tracking_tensor,
             image_tensor=img_cond_tensor,
             output_path=final_output,
-            num_inference_steps=25,
+            num_inference_steps=num_inference_steps,
             guidance_scale=6.0,
             dtype=torch.bfloat16,
             fps=self.fps