Update app.py

app.py

@@ -1,45 +1,50 @@
 import os
 import gradio as gr
 import torch
-import subprocess
 from PIL import Image
 from pathlib import Path
 import io
 import sys
 import traceback
+from huggingface_hub import hf_hub_download
 
 # =========================================
-# 1. Define Hugging Face weights and paths
+# 1. Define Hugging Face dataset + weights
 # =========================================
-from huggingface_hub import hf_hub_download
 
-HF_DATASET_REPO = "roll-ai/FloVD-weights"  # dataset repo ID
+HF_DATASET_REPO = "roll-ai/FloVD-weights"  # your dataset repo on HF
+
 WEIGHT_FILES = {
-    "FVSM/FloVD_FVSM_Controlnet.pt": "FVSM/FloVD_FVSM_Controlnet.pt",
-    "OMSM/selected_blocks.safetensors": "OMSM/selected_blocks.safetensors",
-    "OMSM/pytorch_lora_weights.safetensors": "OMSM/pytorch_lora_weights.safetensors",
-    "others/depth_anything_v2_metric_hypersim_vitb.pth": "others/depth_anything_v2_metric_hypersim_vitb.pth"
+    "ckpt/FVSM/FloVD_FVSM_Controlnet.pt": "FVSM/FloVD_FVSM_Controlnet.pt",
+    "ckpt/OMSM/selected_blocks.safetensors": "OMSM/selected_blocks.safetensors",
+    "ckpt/OMSM/pytorch_lora_weights.safetensors": "OMSM/pytorch_lora_weights.safetensors",
+    "ckpt/others/depth_anything_v2_metric_hypersim_vitb.pth": "others/depth_anything_v2_metric_hypersim_vitb.pth"
 }
+
 def download_weights():
     print("🔄 Downloading model weights via huggingface_hub...")
-    for rel_path in WEIGHT_FILES.values():
-        local_path = Path("ckpt") / rel_path
+    for hf_path, local_rel_path in WEIGHT_FILES.items():
+        local_path = Path("ckpt") / local_rel_path
         if not local_path.exists():
-            print(f"📥 Downloading {rel_path}")
+            print(f"📥 Downloading {hf_path}")
             hf_hub_download(
                 repo_id=HF_DATASET_REPO,
                 repo_type="dataset",
-                filename=rel_path,
-                local_dir="ckpt",
-                local_dir_use_symlinks=False,
+                filename=hf_path,
+                local_dir="ckpt"
             )
         else:
             print(f"✅ Already exists: {local_path}")
+
 download_weights()
+
+# =========================================
+# 2. Import the FloVD generation pipeline
+# =========================================
+
 from inference.flovd_demo import generate_video
 
 def run_inference(prompt, image, pose_type, speed, use_flow_integration, cam_pose_name):
-    # Redirect stdout to capture logs
     log_buffer = io.StringIO()
     sys_stdout = sys.stdout
     sys.stdout = log_buffer
@@ -78,20 +83,19 @@ def run_inference(prompt, image, pose_type, speed, use_flow_integration, cam_pos
         video_path = f"./outputs/generated_videos/{video_name}"
         print(f"✅ Inference complete. Video saved to {video_path}")
 
-    except Exception as e:
+    except Exception:
         print("🔥 Inference failed with exception:")
         traceback.print_exc()
 
-    # Restore stdout and return logs
     sys.stdout = sys_stdout
     logs = log_buffer.getvalue()
     log_buffer.close()
 
     return (video_path if video_path and os.path.exists(video_path) else None), logs
 
-# ========================
-# Gradio Interface
-# ========================
+# =========================================
+# 3. Gradio App Interface
+# =========================================
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🎥 FloVD: Optical Flow + CogVideoX Video Generation")
@@ -113,4 +117,5 @@ with gr.Blocks() as demo:
         inputs=[prompt, image, pose_type, speed, use_flow_integration, cam_pose_name],
         outputs=[output_video, output_logs]
     )
+
 demo.launch(show_error=True)