Update app.py

app.py

@@ -18,25 +18,15 @@ snapshot_download(
     token=hf_token,
 )
 
-print("✅ Download complete.")
-print("📂 Contents of ./ckpt directory:")
-for root, dirs, files in os.walk("./ckpt"):
-    level = root.replace("./ckpt", "").count(os.sep)
-    indent = " " * 4 * level
-    print(f"{indent}{os.path.basename(root)}/")
-    subindent = " " * 4 * (level + 1)
-    for f in files:
-        print(f"{subindent}{f}")
-
 # -----------------------------------
-# Step 2: Define inference function
+# Step 2: Setup paths and config
 # -----------------------------------
-
 FVSM_PATH = "./ckpt/FVSM/FloVD_FVSM_Controlnet.pt"
 OMSM_PATH = "./ckpt/OMSM/"
 DEPTH_CKPT_PATH = "./ckpt/others/depth_anything_v2_metric_hypersim_vitb.pth"
 OUTPUT_PATH = "./results/"
 GEN_VID_DIR = os.path.join(OUTPUT_PATH, "generated_videos")
+
 POSE_TYPE = "re10k"
 CONTROLNET_GUIDANCE_END = 0.4
 SPEED = 1.0
@@ -46,14 +36,19 @@ INFER_STEPS = 50
 
 os.makedirs(GEN_VID_DIR, exist_ok=True)
 
+# -----------------------------------
+# Helper Functions
+# -----------------------------------
+
 def list_generated_videos():
     try:
-        files = os.listdir(GEN_VID_DIR)
-        print(f"\n📂 Contents of {GEN_VID_DIR}:")
-        for f in files:
-            print(f"   - {f}")
+        return sorted([
+            f for f in os.listdir(GEN_VID_DIR)
+            if f.endswith(".mp4")
+        ])
     except Exception as e:
         print(f"[⚠️] Could not list contents: {str(e)}")
+        return []
 
 def run_flovd(prompt, image, cam_pose_name):
     try:
@@ -61,8 +56,6 @@ def run_flovd(prompt, image, cam_pose_name):
         print("🚀 Starting video generation")
         print("----------------------------")
 
-        list_generated_videos()
-
         image_path = "./temp_input.png"
         image.save(image_path)
         print(f"📸 Image saved at {image_path}")
@@ -85,9 +78,6 @@ def run_flovd(prompt, image, cam_pose_name):
             num_inference_steps=INFER_STEPS,
         )
 
-        print("\n✅ Generation finished.")
-        list_generated_videos()
-
         prompt_short = prompt[:30].strip().replace(" ", "_").replace(".", "").replace(",", "")
         video_filename = f"{prompt_short}_{cam_pose_name}.mp4"
         video_path = os.path.join(GEN_VID_DIR, video_filename)
@@ -95,30 +85,51 @@ def run_flovd(prompt, image, cam_pose_name):
         print(f"\n📍 Looking for generated video at: {video_path}")
         if os.path.exists(video_path):
             print("✅ Video file found.")
-            return video_path
+            return video_path, list_generated_videos(), None
         else:
             print("❌ File missing:", video_path)
-            list_generated_videos()
-            return f"[ERROR] File not found at: {video_path}"
+            return f"[ERROR] File not found at: {video_path}", list_generated_videos(), None
 
     except Exception as e:
         print(f"🔥 Exception occurred: {str(e)}")
-        return f"An error occurred during generation: {str(e)}"
+        return f"[EXCEPTION] {str(e)}", list_generated_videos(), None
+
+def get_video_file(video_name):
+    video_path = os.path.join(GEN_VID_DIR, video_name)
+    if os.path.exists(video_path):
+        return video_path
+    else:
+        return None
 
 # -----------------------------------
-# Step 3: Launch Gradio UI
+# Step 3: Launch Gradio Interface
 # -----------------------------------
 
-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 (e.g., 1593596b99e2dde9.txt)"),
-    ],
-    outputs=gr.Video(label="Generated Video or Error Message"),
-    title="🎥 FloVD + CogVideoX - Camera Motion Guided Video Generation",
-    description="Upload an image, enter a prompt, and provide a camera pose filename to generate a video."
-)
-
-iface.launch(server_name="0.0.0.0", server_port=7860)
+with gr.Blocks(title="FloVD + CogVideoX") as demo:
+    gr.Markdown("## 🎥 FloVD - Camera Motion Guided Video Generation + Downloader")
+
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt")
+            image = gr.Image(type="pil", label="Input Image")
+            pose_file = gr.Textbox(label="Camera Pose Filename (e.g. abc.txt)")
+            run_btn = gr.Button("Generate Video")
+        with gr.Column():
+            output_video = gr.Video(label="Generated Video")
+            video_selector = gr.Dropdown(choices=list_generated_videos(), label="Select Video to Download")
+            download_btn = gr.Button("Download Selected")
+            file_output = gr.File(label="Download Link")
+
+    run_btn.click(
+        fn=run_flovd,
+        inputs=[prompt, image, pose_file],
+        outputs=[output_video, video_selector, file_output],
+    )
+
+    download_btn.click(
+        fn=get_video_file,
+        inputs=[video_selector],
+        outputs=[file_output],
+    )
+
+demo.launch(server_name="0.0.0.0", server_port=7860)