Update app.py

app.py

@@ -5,66 +5,78 @@ from PIL import Image
 from inference.flovd_demo import generate_video
 import requests
 import shutil
+import os
+from huggingface_hub import snapshot_download
 
-# --------- SETUP CODE: Download models if not found ---------
-FVSM_PATH = "./ckpt/FVSM/FloVD_FVSM_Controlnet.pt"
-OMSM_PATH = "./ckpt/OMSM"
-POSE_DIR = "./assets/manual_poses"
-EXAMPLE_POSE = os.path.join(POSE_DIR, "example.txt")
+hf_token = os.getenv("HF_TOKEN")
 
-def download_if_missing():
-    os.makedirs("ckpt/FVSM", exist_ok=True)
-    os.makedirs("ckpt/OMSM", exist_ok=True)
-    os.makedirs("assets/manual_poses", exist_ok=True)
-    os.makedirs("output/generated_videos", exist_ok=True)
+snapshot_download(
+    repo_id="roll-ai/FloVD-weights",
+    repo_type="dataset",
+    local_dir="./ckpt",
+    allow_patterns="ckpt/**",
+    token=hf_token
+)
 
-    # Download FVSM model
-    if not os.path.exists(FVSM_PATH):
-        print("Downloading FVSM model...")
-        url = "https://huggingface.co/datasets/mutqa/FloVD-HF-Assets/resolve/main/FloVD_FVSM_Controlnet.pt"
-        r = requests.get(url, stream=True)
-        with open(FVSM_PATH, 'wb') as f:
-            shutil.copyfileobj(r.raw, f)
+# # --------- SETUP CODE: Download models if not found ---------
+# FVSM_PATH = "./ckpt/FVSM/FloVD_FVSM_Controlnet.pt"
+# OMSM_PATH = "./ckpt/OMSM"
+# POSE_DIR = "./assets/manual_poses"
+# EXAMPLE_POSE = os.path.join(POSE_DIR, "example.txt")
 
-    # Download OMSM weights (as folder with .safetensors inside)
-    if not os.listdir(OMSM_PATH):
-        print("Cloning OMSM weights...")
-        os.system("git clone https://huggingface.co/datasets/mutqa/FloVD-HF-OMSM ckpt/OMSM")
+# def download_if_missing():
+#     os.makedirs("ckpt/FVSM", exist_ok=True)
+#     os.makedirs("ckpt/OMSM", exist_ok=True)
+#     os.makedirs("assets/manual_poses", exist_ok=True)
+#     os.makedirs("output/generated_videos", exist_ok=True)
 
-    # Download example camera pose
-    if not os.path.exists(EXAMPLE_POSE):
-        print("Downloading example pose...")
-        url = "https://huggingface.co/datasets/mutqa/FloVD-HF-Assets/resolve/main/example.txt"
-        r = requests.get(url)
-        with open(EXAMPLE_POSE, "w") as f:
-            f.write(r.text)
+#     # Download FVSM model
+#     if not os.path.exists(FVSM_PATH):
+#         print("Downloading FVSM model...")
+#         url = "https://huggingface.co/datasets/mutqa/FloVD-HF-Assets/resolve/main/FloVD_FVSM_Controlnet.pt"
+#         r = requests.get(url, stream=True)
+#         with open(FVSM_PATH, 'wb') as f:
+#             shutil.copyfileobj(r.raw, f)
 
-# --------- UI Function ---------
-def run_flovd(prompt, image, cam_pose_name):
-    download_if_missing()
-    image_path = "./temp_input.png"
-    image.save(image_path)
+#     # Download OMSM weights (as folder with .safetensors inside)
+#     if not os.listdir(OMSM_PATH):
+#         print("Cloning OMSM weights...")
+#         os.system("git clone https://huggingface.co/datasets/mutqa/FloVD-HF-OMSM ckpt/OMSM")
 
-    generate_video(
-        prompt=prompt,
-        fvsm_path=FVSM_PATH,
-        omsm_path=OMSM_PATH,
-        image_path=image_path,
-        cam_pose_name=cam_pose_name,
-        output_path="./output/",
-        dtype=torch.float16,
-    )
-    return "./output/generated_videos/your_video.mp4"
+#     # Download example camera pose
+#     if not os.path.exists(EXAMPLE_POSE):
+#         print("Downloading example pose...")
+#         url = "https://huggingface.co/datasets/mutqa/FloVD-HF-Assets/resolve/main/example.txt"
+#         r = requests.get(url)
+#         with open(EXAMPLE_POSE, "w") as f:
+#             f.write(r.text)
 
-# --------- Launch Gradio ---------
-iface = gr.Interface(
-    fn=run_flovd,
-    inputs=[
-        gr.Textbox(label="Prompt"),
-        gr.Image(type="pil", label="Input Image"),
-        gr.Textbox(label="Camera Pose File Name", value="example.txt"),
-    ],
-    outputs=gr.Video(label="Generated Video"),
-    title="FloVD - Optical Flow Video Diffusion with Camera Motion",
-)
-iface.launch(server_name="0.0.0.0", server_port=7860)
+# # --------- UI Function ---------
+# def run_flovd(prompt, image, cam_pose_name):
+#     download_if_missing()
+#     image_path = "./temp_input.png"
+#     image.save(image_path)
+
+#     generate_video(
+#         prompt=prompt,
+#         fvsm_path=FVSM_PATH,
+#         omsm_path=OMSM_PATH,
+#         image_path=image_path,
+#         cam_pose_name=cam_pose_name,
+#         output_path="./output/",
+#         dtype=torch.float16,
+#     )
+#     return "./output/generated_videos/your_video.mp4"
+
+# # --------- Launch Gradio ---------
+# iface = gr.Interface(
+#     fn=run_flovd,
+#     inputs=[
+#         gr.Textbox(label="Prompt"),
+#         gr.Image(type="pil", label="Input Image"),
+#         gr.Textbox(label="Camera Pose File Name", value="example.txt"),
+#     ],
+#     outputs=gr.Video(label="Generated Video"),
+#     title="FloVD - Optical Flow Video Diffusion with Camera Motion",
+# )
+# iface.launch(server_name="0.0.0.0", server_port=7860)