Update app.py

app.py

@@ -169,17 +169,24 @@ def process_video_for_pose(video):
     
     return pose_video
 
-def process_video_for_control(video, control_type):
+def process_video_for_control(reference_video, control_type):
+    video = load_video(reference_video)
     """Process video based on the selected control type"""
     if control_type == "canny":
-        return process_video_for_canny(video)
+        processed_video = process_video_for_canny(video)
     elif control_type == "depth":
-        return process_video_for_depth(video)
+        processed_video = process_video_for_depth(video)
     elif control_type == "pose":
-        return process_video_for_pose(video)
+        processed_video = process_video_for_pose(video)
     else:
-        return video
-        
+        processed_video = reference_video
+    fps = 24
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp2_file:
+        output2_path = tmp2_file.name
+        export_to_video(processed_video, output2_path, fps=fps)
+    return output2_path
+
+
 @spaces.GPU(duration=160)
 def generate_video(
     reference_video,
@@ -213,6 +220,13 @@ def generate_video(
         # Handle seed
         if randomize_seed:
             seed = random.randint(0, 2**32 - 1)
+
+        # Calculate number of frames from duration (24 fps)
+        fps = 24
+        num_frames = int(duration * fps) + 1  # +1 for proper frame count
+        # Ensure num_frames is valid for the model (multiple of temporal compression + 1)
+        temporal_compression = pipeline.vae_temporal_compression_ratio
+        num_frames = ((num_frames - 1) // temporal_compression) * temporal_compression + 1
         
         progress(0.05, desc="Loading control LoRA...")
         
@@ -221,20 +235,14 @@ def generate_video(
         
         # Loads video into a list of pil images 
         video = load_video(reference_video)
-        progress(0.1, desc="Processing video for control...")
+        # progress(0.1, desc="Processing video for control...")
         
         # Process video based on control type
-        processed_video = process_video_for_control(video, control_type)
-        processed_video = read_video(processed_video) # turns to tensor
+        #processed_video = process_video_for_control(video, control_type)
         
-        progress(0.2, desc="Preparing generation parameters...")
+        processed_video = read_video(video) # turns to tensor
         
-        # Calculate number of frames from duration (24 fps)
-        fps = 24
-        num_frames = int(duration * fps) + 1  # +1 for proper frame count
-        # Ensure num_frames is valid for the model (multiple of temporal compression + 1)
-        temporal_compression = pipeline.vae_temporal_compression_ratio
-        num_frames = ((num_frames - 1) // temporal_compression) * temporal_compression + 1
+        progress(0.2, desc="Preparing generation parameters...")
         
         # Calculate downscaled dimensions
         downscale_factor = 2 / 3
@@ -451,14 +459,20 @@ with gr.Blocks() as demo:
                 label="Generated Video",
                 height=400
             )
+            control_video = gr.Video(
+                label="Control Video",
+                height=400
+            )
         
             
     
     # Event handlers
     generate_btn.click(
+        fn = process_video_for_control,
+        inputs = [reference_video, control_type], outputs = [control_video]).then(
         fn=generate_video,
         inputs=[
-            reference_video,
+            control_video,
             prompt,
             control_type,
             current_lora_state,