duration instead

app.py

@@ -26,7 +26,7 @@ try:
     print("download sam")
     sam_dir = snapshot_download(repo_id="alexnasa/sam2")
 
-    @spaces.GPU(duration=450)
+    @spaces.GPU(duration=500)
     def install_sam():
         os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
         sh(f"cd {sam_dir} && python setup.py build_ext --inplace && pip install -e .")
@@ -62,7 +62,7 @@ rc_mapping = {
 }
 
 
-def preprocess_video(input_video_path, session_id=None):
+def preprocess_video(input_video_path, duration, session_id=None):
     
     if session_id is None:
         session_id = uuid.uuid4().hex
@@ -72,7 +72,7 @@ def preprocess_video(input_video_path, session_id=None):
 
     process_video_path = os.path.join(output_dir, 'input_video.mp4')
 
-    convert_video_to_30fps_and_clip(input_video_path, process_video_path, crop_width=720, crop_height=1280)
+    convert_video_to_30fps_and_clip(input_video_path, process_video_path, duration_s=duration, crop_width=720, crop_height=1280)
 
     return process_video_path
 
@@ -207,9 +207,9 @@ def get_frames_count(video_file):
 
     return orig_frame_count
 
-def calculate_time_required(input_video, rc_bool):
+def calculate_time_required(max_duration_s, rc_bool):
 
-    frames_count = get_frames_count(input_video)
+    frames_count = 30 * max_duration_s
 
     chunks = frames_count // 77 + 1
 
@@ -218,32 +218,29 @@ def calculate_time_required(input_video, rc_bool):
         pose2d_tracking_duration_s = 75
         iteration_per_step_s = 13
     else:
-        pose2d_tracking_duration_s = 75
+        pose2d_tracking_duration_s = 50
         iteration_per_step_s = 12
         
     time_required = pose2d_tracking_duration_s + iteration_per_step_s * 10 * chunks
     print(f'for frames_count:{frames_count} doing {chunks} chunks the time_required is {time_required}')
     return time_required
 
-def update_time_required(input_video, rc_str):
+def update_time_required(max_duration_s, rc_str):
 
-    if input_video is None:
-        return gr.update(value="⌚ Zero GPU Required: --")
-    
     rc_bool = rc_mapping[rc_str]
 
-    duration_s = calculate_time_required(input_video, rc_bool)
+    duration_s = calculate_time_required(max_duration_s, rc_bool)
     duration_m = duration_s / 60
     
     return gr.update(value=f"⌚ Zero GPU Required: ~{duration_s}.0s ({duration_m:.1f} mins)")
 
-def get_duration(input_video, edited_frame, rc_bool, session_id, progress):
+def get_duration(input_video, max_duration_s, edited_frame, rc_bool, session_id, progress):
 
-    return calculate_time_required(input_video, rc_bool)
+    return calculate_time_required(max_duration_s, rc_bool)
     
 
 @spaces.GPU(duration=get_duration)
-def _animate(input_video, edited_frame, rc_bool, session_id = None, progress=gr.Progress(track_tqdm=True),):
+def _animate(input_video, max_duration_s, edited_frame, rc_bool, session_id = None, progress=gr.Progress(track_tqdm=True),):
 
     if session_id is None:
         session_id = uuid.uuid4().hex
@@ -301,7 +298,7 @@ def _animate(input_video, edited_frame, rc_bool, session_id = None, progress=gr.
 
     return output_video_path
 
-def animate_scene(input_video, edited_frame, rc_str, session_id = None, progress=gr.Progress(track_tqdm=True),):
+def animate_scene(input_video, max_duration_s, edited_frame, rc_str, session_id = None, progress=gr.Progress(track_tqdm=True),):
 
     if not input_video:
         raise gr.Error("Please provide an video")
@@ -312,8 +309,11 @@ def animate_scene(input_video, edited_frame, rc_str, session_id = None, progress
     if session_id is None:
         session_id = uuid.uuid4().hex
 
+    input_video = preprocess_video(input_video, max_duration_s, session_id)
+
     rc_bool = rc_mapping[rc_str]
     
+
     output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
     os.makedirs(output_dir, exist_ok=True)
 
@@ -324,8 +324,8 @@ def animate_scene(input_video, edited_frame, rc_str, session_id = None, progress
     edited_frame_png = os.path.join(output_dir, 'edited_frame.png')
     edited_frame_img = Image.open(edited_frame)
     edited_frame_img.save(edited_frame_png)
-
-    output_video_path = _animate(input_video, edited_frame, rc_bool, session_id, progress)
+    
+    output_video_path = _animate(input_video, max_duration_s, edited_frame_png, rc_bool, session_id, progress)
 
     final_video_path = os.path.join(output_dir, 'final_result.mp4')
 
@@ -345,6 +345,7 @@ def animate_scene(input_video, edited_frame, rc_str, session_id = None, progress
         combine_video_and_audio_ffmpeg(output_video_path, input_audio_path, final_video_path)
     else:
         final_video_path = output_video_path
+
     return final_video_path, pose_video, bg_video, mask_video, face_video
 
 css = """
@@ -480,7 +481,8 @@ with gr.Blocks(css=css, title="Wan 2.2 Animate --replace", theme=gr.themes.Ocean
                 </div>
                 """)
                 input_video = gr.Video(label="Input Video", height=512)
-                trim_button = gr.Button("Trim to 2s")
+                max_duration_slider = gr.Slider(2, 8, 2, step=2, label="Max Duration")
+
 
             with gr.Column(elem_id="step-column"):
                 gr.HTML("""
@@ -520,32 +522,32 @@ with gr.Blocks(css=css, title="Wan 2.2 Animate --replace", theme=gr.themes.Ocean
 
                 gr.Examples(
                     examples=[
-                        
+
                         [
                             "./examples/paul.mp4", 
+                            2,
                             "./examples/man.png",      
                             "Video → Ref Image"
                         ],
 
                         [
                             "./examples/desi.mp4", 
+                            2,
                             "./examples/desi.png",     
                             "Video ← Ref Image"
                         ],
 
 
-
                     ],
-                    inputs=[input_video, edited_frame, replace_character_string],
+                    inputs=[input_video, max_duration_slider, edited_frame, replace_character_string],
                     outputs=[output_video, pose_video, bg_video, mask_video, face_video],
                     fn=animate_scene,
                     cache_examples=True,
                 )
 
-    action_button.click(fn=animate_scene, inputs=[input_video, edited_frame, replace_character_string, session_state], outputs=[output_video, pose_video, bg_video, mask_video, face_video])
-    input_video.change(update_time_required, inputs=[input_video, replace_character_string], outputs=[time_required])
-    trim_button.click(preprocess_video, inputs=[input_video, session_state], outputs=[input_video])
-    replace_character_string.change(update_time_required, inputs=[input_video, replace_character_string], outputs=[time_required])
+    action_button.click(fn=animate_scene, inputs=[input_video, max_duration_slider, edited_frame, replace_character_string, session_state], outputs=[output_video, pose_video, bg_video, mask_video, face_video])
+    max_duration_slider.change(update_time_required, inputs=[max_duration_slider, replace_character_string], outputs=[time_required])
+    replace_character_string.change(update_time_required, inputs=[max_duration_slider, replace_character_string], outputs=[time_required])
                    
 if __name__ == "__main__":
     demo.queue()