Update app.py

app.py

@@ -4,6 +4,9 @@ import torch
 import subprocess
 from PIL import Image
 from pathlib import Path
+import io
+import sys
+import traceback
 
 # =========================================
 # 1. Define Hugging Face weights and paths
@@ -30,16 +33,22 @@ def download_weights():
             print(f"✅ Already exists: {save_path}")
 
 download_weights()
-import gradio as gr
-import torch
-import os
+
 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
+
+    video_path = None
     try:
+        print("🚀 Starting inference...")
         os.makedirs("input_images", exist_ok=True)
         image_path = "input_images/input_image.png"
         image.save(image_path)
+        print(f"📸 Saved input image to {image_path}")
 
         generate_video(
             prompt=prompt,
@@ -62,13 +71,25 @@ def run_inference(prompt, image, pose_type, speed, use_flow_integration, cam_pos
             cam_pose_name=cam_pose_name,
             depth_ckpt_path="./ckpt/others/depth_anything_v2_metric_hypersim_vitb.pth"
         )
-        return f"./outputs/generated_videos/{prompt[:30].strip().replace(' ', '_')}_{cam_pose_name or 'default'}.mp4"
+
+        video_name = f"{prompt[:30].strip().replace(' ', '_')}_{cam_pose_name or 'default'}.mp4"
+        video_path = f"./outputs/generated_videos/{video_name}"
+        print(f"✅ Inference complete. Video saved to {video_path}")
 
     except Exception as e:
-        print("🔥 Inference failed:")
-        import traceback
+        print("🔥 Inference failed with exception:")
         traceback.print_exc()
-        return f"⚠️ Error during inference: {e}"
+
+    # 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
+# ========================
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🎥 FloVD: Optical Flow + CogVideoX Video Generation")
@@ -83,7 +104,12 @@ with gr.Blocks() as demo:
             submit = gr.Button("Generate Video")
         with gr.Column():
             output_video = gr.Video(label="Generated Video")
+            output_logs = gr.Textbox(label="Logs", lines=20, interactive=False)
 
-    submit.click(fn=run_inference, inputs=[prompt, image, pose_type, speed, use_flow_integration, cam_pose_name], outputs=output_video)
+    submit.click(
+        fn=run_inference,
+        inputs=[prompt, image, pose_type, speed, use_flow_integration, cam_pose_name],
+        outputs=[output_video, output_logs]
+    )
 
 demo.launch()