Update gradio_app_with_frames.py

gradio_app_with_frames.py

@@ -1,20 +1,46 @@
 import os
+import sys
 import shutil
 import uuid
 import subprocess
 import gradio as gr
-import cv2
-import sys
+import cv2  # 用于检查视频帧数
 from glob import glob
-from pathlib import Path
 
-# 获取当前Python解释器路径
-#PYTHON_EXECUTABLE = sys.executable
-PYTHON_EXECUTABLE = "python"
+from huggingface_hub import snapshot_download, hf_hub_download
+
+# Download models
+os.makedirs("pretrained_weights", exist_ok=True)
+
+# List of subdirectories to create inside "checkpoints"
+subfolders = [
+    "stable-video-diffusion-img2vid-xt"
+]
+
+# Create each subdirectory
+for subfolder in subfolders:
+    os.makedirs(os.path.join("pretrained_weights", subfolder), exist_ok=True)
+
+snapshot_download(
+    repo_id="stabilityai/stable-video-diffusion-img2vid",
+    local_dir="./pretrained_weights/stable-video-diffusion-img2vid-xt"
+)
+
+snapshot_download(
+    repo_id="Yhmeng1106/anidoc",
+    local_dir="./pretrained_weights"
+)
+
+hf_hub_download(
+    repo_id="facebook/cotracker",
+    filename="cotracker2.pth",
+    local_dir="./pretrained_weights"
+)
 
 def normalize_path(path: str) -> str:
+    return path
     """标准化路径，将Windows路径转换为正斜杠形式"""
-    return str(Path(path).resolve()).replace('\\', '/')
+    return os.path.abspath(path).replace('\\', '/')
 
 def check_video_frames(video_path: str) -> int:
     """检查视频帧数"""
@@ -33,14 +59,14 @@ def preprocess_video(video_path: str) -> str:
         output_dir = os.path.join(temp_dir, f"processed_{unique_id}")
         output_dir = normalize_path(output_dir)
         os.makedirs(output_dir, exist_ok=True)
-        
+
         print(f"Processing video: {video_path}")
         print(f"Output directory: {output_dir}")
-        
-        # 调用process_video_to_14frames.py处理视频
+
+        # 调用外部脚本处理视频
         result = subprocess.run(
             [
-                PYTHON_EXECUTABLE, "process_video_to_14frames.py",
+                "python", "process_video_to_14frames.py",
                 "--input", video_path,
                 "--output", output_dir
             ],
@@ -48,12 +74,12 @@ def preprocess_video(video_path: str) -> str:
             capture_output=True,
             text=True
         )
-        
+
         if result.stdout:
             print(f"Preprocessing stdout: {result.stdout}")
         if result.stderr:
             print(f"Preprocessing stderr: {result.stderr}")
-        
+
         # 获取处理后的视频路径
         processed_videos = glob(os.path.join(output_dir, "*.mp4"))
         if not processed_videos:
@@ -66,87 +92,52 @@ def preprocess_video(video_path: str) -> str:
         raise gr.Error(f"Error during video preprocessing: {str(e)}")
 
 def generate(control_sequence, ref_image):
+    control_image = control_sequence  # "data_test/sample4.mp4"
+    ref_image = ref_image  # "data_test/sample4.png"
+    unique_id = str(uuid.uuid4())
+    output_dir = f"results_{unique_id}"
+
     try:
-        # 验证输入文件是否存在
-        control_sequence = normalize_path(control_sequence)
-        ref_image = normalize_path(ref_image)
-        
-        if not os.path.exists(control_sequence):
-            raise gr.Error(f"Control sequence file not found: {control_sequence}")
-        if not os.path.exists(ref_image):
-            raise gr.Error(f"Reference image file not found: {ref_image}")
-            
-        # 创建输出目录
-        output_dir = "outputs"
-        os.makedirs(output_dir, exist_ok=True)
-        unique_id = str(uuid.uuid4())
-        result_dir = os.path.join(output_dir, f"results_{unique_id}")
-        result_dir = normalize_path(result_dir)
-        os.makedirs(result_dir, exist_ok=True)
-        
-        print(f"Input control sequence: {control_sequence}")
-        print(f"Input reference image: {ref_image}")
-        print(f"Output directory: {result_dir}")
-        
         # 检查视频帧数
-        frame_count = check_video_frames(control_sequence)
+        frame_count = check_video_frames(control_image)
         if frame_count != 14:
             print(f"Video has {frame_count} frames, preprocessing to 14 frames...")
-            control_sequence = preprocess_video(control_sequence)
-            print(f"Preprocessed video saved to: {control_sequence}")
-        
+            control_image = preprocess_video(control_image)
+            print(f"Preprocessed video saved to: {control_image}")
+
         # 运行推理命令
-        print(f"Running inference...")
-        result = subprocess.run(
+        subprocess.run(
             [
-                PYTHON_EXECUTABLE, "scripts_infer/anidoc_inference.py",
+                "python", "scripts_infer/anidoc_inference.py",
                 "--all_sketch",
                 "--matching",
                 "--tracking",
-                "--control_image", control_sequence,
-                "--ref_image", ref_image,
-                "--output_dir", result_dir,
+                "--control_image", f"{control_image}",
+                "--ref_image", f"{ref_image}",
+                "--output_dir", f"{output_dir}",
                 "--max_point", "10",
             ],
-            check=True,
-            capture_output=True,
-            text=True
+            check=True
         )
-        
-        if result.stdout:
-            print(f"Inference stdout: {result.stdout}")
-        if result.stderr:
-            print(f"Inference stderr: {result.stderr}")
 
         # 搜索输出视频
-        output_video = glob(os.path.join(result_dir, "*.mp4"))
-        print(f"Found output videos: {output_video}")
-        
+        output_video = glob(os.path.join(output_dir, "*.mp4"))
+        print(output_video)
+
         if output_video:
-            output_video_path = normalize_path(output_video[0])
-            print(f"Returning output video: {output_video_path}")
+            output_video_path = output_video[0]  # 获取第一个匹配
         else:
-            raise gr.Error("No output video generated")
-            
-        # 清理临时文件
-        temp_dirs = glob("outputs/processed_*")
-        for temp_dir in temp_dirs:
-            if os.path.isdir(temp_dir):
-                try:
-                    shutil.rmtree(temp_dir)
-                    print(f"Cleaned up temp directory: {temp_dir}")
-                except Exception as e:
-                    print(f"Warning: Failed to clean up temp directory {temp_dir}: {str(e)}")
-        
+            output_video_path = None
+
+        print(output_video_path)
         return output_video_path
-        
+
     except subprocess.CalledProcessError as e:
-        print(f"Inference stderr: {e.stderr}")
-        raise gr.Error(f"Error during inference: {e.stderr}")
+        raise gr.Error(f"Error during inference: {str(e)}")
     except Exception as e:
         raise gr.Error(f"Error: {str(e)}")
 
-css="""
+css = """
 div#col-container{
     margin: 0 auto;
     max-width: 982px;
@@ -168,6 +159,12 @@ with gr.Blocks(css=css) as demo:
             <a href="https://arxiv.org/pdf/2412.14173">
                 <img src='https://img.shields.io/badge/ArXiv-Paper-red'>
             </a>
+            <a href="https://huggingface.co/spaces/fffiloni/AniDoc?duplicate=true">
+                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
+            </a>
+            <a href="https://huggingface.co/fffiloni">
+                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-sm-dark.svg" alt="Follow me on HF">
+            </a>
         </div>
         """)
         with gr.Row():
@@ -179,20 +176,20 @@ with gr.Blocks(css=css) as demo:
                 video_result = gr.Video(label="Result")
 
                 gr.Examples(
-                    examples = [
+                    examples=[
                         ["data_test/sample5.mp4", "data_test/sample5.png"],
                         ["data_test/sample1.mp4", "data_test/sample1.png"],
                         ["data_test/sample2.mp4", "data_test/sample2.png"],
                         ["data_test/sample3.mp4", "data_test/sample3.png"],
                         ["data_test/sample4.mp4", "data_test/sample4.png"]
                     ],
-                    inputs = [control_sequence, ref_image]
-        )
-        
-        submit_btn.click(
-        fn = generate,
-        inputs = [control_sequence, ref_image],
-        outputs = [video_result]
-        )
+                    inputs=[control_sequence, ref_image]
+                )
+
+    submit_btn.click(
+        fn=generate,
+        inputs=[control_sequence, ref_image],
+        outputs=[video_result]
+    )
 
-demo.queue().launch(inbrowser=True,show_api=False, show_error=True, share = True)
+demo.queue().launch(show_api=False, show_error=True, share=True)