Update app.py

app.py

@@ -35,6 +35,60 @@ import aoti
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
 warnings.filterwarnings("ignore")
 
+# --- FRAME EXTRACTION JS & LOGIC ---
+
+# JS to grab timestamp from the output video
+get_timestamp_js = """
+function() {
+    // Select the video element specifically inside the component with id 'generated-video'
+    const video = document.querySelector('#generated-video video');
+    
+    if (video) {
+        console.log("Video found! Time: " + video.currentTime);
+        return video.currentTime;
+    } else {
+        console.log("No video element found.");
+        return 0;
+    }
+}
+"""
+
+def extract_frame(video_path, timestamp):
+    # Safety check: if no video is present
+    if not video_path:
+        return None
+    
+    print(f"Extracting frame at timestamp: {timestamp}") 
+    
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        return None
+
+    # Calculate frame number
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    target_frame_num = int(float(timestamp) * fps)
+    
+    # Cap total frames to prevent errors at the very end of video
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if target_frame_num >= total_frames:
+        target_frame_num = total_frames - 1
+    
+    # Set position
+    cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame_num)
+    ret, frame = cap.read()
+    cap.release()
+    
+    if ret:
+        # Convert from BGR (OpenCV) to RGB (Gradio)
+        # Gradio Image component handles Numpy array -> PIL conversion automatically
+        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    
+    return None
+
+# --- END FRAME EXTRACTION LOGIC ---
+
+
 # RIFE
 if not os.path.exists("RIFEv4.26_0921.zip"):
     print("Downloading RIFE Model...")
@@ -347,23 +401,22 @@ def run_inference(
     raw_frames_np = result.frames[0] # Returns (T, H, W, C) float32
     pipe.scheduler = original_scheduler
 
-    start = time.time()
     if frame_multiplier > 1:
+        start = time.time()
         print(f"Processing frames (RIFE Multiplier: {frame_multiplier}x)...")
         final_frames = interpolate_bits(raw_frames_np, multiplier=int(frame_multiplier))
+        print("Interpolation time passed:", time.time() - start)
     else:
         final_frames = list(raw_frames_np)
-    print("Interpolation time passed:", time.time() - start)
 
     final_fps = FIXED_FPS * int(frame_multiplier)
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
 
-    print(f"Exporting video at {final_fps} FPS...")
     start = time.time()
     export_to_video(final_frames, video_path, fps=final_fps, quality=quality)
-    print("Export time passed:", time.time() - start)
+    print(f"Export time passed, {final_fps} FPS:", time.time() - start)
     
     return video_path
 
@@ -382,6 +435,7 @@ def generate_video(
     scheduler="UniPCMultistep",
     flow_shift=6.0,
     frame_multiplier=1,
+    video_component=True,
     progress=gr.Progress(track_tqdm=True),
 ):
     """
@@ -412,6 +466,8 @@ def generate_video(
         scheduler (str, optional): The name of the scheduler to use for inference. Defaults to "UniPCMultistep".
         flow_shift (float, optional): The flow shift value for compatible schedulers. Defaults to 6.0.
         frame_multiplier (int, optional): The int value for fps enhancer
+        video_component(bool, optional): Show video player in output.
+            Defaults to True.
         progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
     Returns:
         tuple: A tuple containing:
@@ -457,12 +513,26 @@ def generate_video(
     )
     print("GPU complete")
 
-    return video_path, video_path, current_seed
+    return (video_path if video_component else None), video_path, current_seed
+
+
+CSS = """
+#hidden-timestamp {
+    opacity: 0;
+    height: 0px;
+    width: 0px;
+    margin: 0px;
+    padding: 0px;
+    overflow: hidden;
+    position: absolute;
+    pointer-events: none;
+}
+"""
 
 
 with gr.Blocks(delete_cache=(3600, 10800)) as demo:
-    gr.Markdown("# WAMU - Wan 2.2 I2V (14B)")
-    gr.Markdown("## ℹ️ **A Note on Performance:** This version prioritizes a straightforward setup over maximum speed, so performance may vary.")
+    gr.Markdown("## WAMU - Wan 2.2 I2V (14B)")
+    gr.Markdown("#### ℹ️ **A Note on Performance:** This version prioritizes a straightforward setup over maximum speed, so performance may vary.")
     gr.Markdown("run Wan 2.2 in just 4-8 steps, fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU⚡️")
 
     with gr.Row():
@@ -475,16 +545,16 @@ with gr.Blocks(delete_cache=(3600, 10800)) as demo:
                 choices=[1, 2, 4, 8], 
                 value=1, 
                 label="Frame Rate Enhancer (Interpolation)",
-                info="2 = Double FPS (e.g. 16 -> 32). Higher multipliers create more intermediate frames."
+                info="Increases video fluidity. Example: 2x converts 16 FPS -> 32 FPS."
             )
             with gr.Accordion("Advanced Settings", open=False):
                 last_image_component = gr.Image(type="pil", label="Last Image (Optional)")
                 negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, info="Used if any Guidance Scale > 1.", lines=3)
-                quality_slider = gr.Slider(minimum=1, maximum=10, step=1, value=6, label="Video Quality")
+                quality_slider = gr.Slider(minimum=1, maximum=10, step=1, value=6, label="Video Quality", info="If set to 10, the generated video may be too large and won't play in the Gradio preview.")
                 seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                 randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
-                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage")
+                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage", info="Values above 1 increase GPU usage and may take longer to process.")
+                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage", info="Values above 1 increase GPU usage and may take longer to process.")
                 scheduler_dropdown = gr.Dropdown(
                     label="Scheduler",
                     choices=list(SCHEDULER_MAP.keys()),
@@ -492,18 +562,29 @@ with gr.Blocks(delete_cache=(3600, 10800)) as demo:
                     info="Select a custom scheduler."
                 )
                 flow_shift_slider = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=3.0, label="Flow Shift")
+                play_result_video = gr.Checkbox(label="Display result", value=True, interactive=True)
+                gr.Markdown("[ZeroGPU Help, Tips, and Troubleshooting](https://huggingface.co/datasets/TestOrganizationPleaseIgnore/help/blob/main/gpu_help.md)")
 
             generate_button = gr.Button("Generate Video", variant="primary")
 
         with gr.Column():
-            video_output = gr.Video(label="Generated Video", autoplay=True)
+            # ASSIGNED elem_id="generated-video" so JS can find it
+            video_output = gr.Video(label="Generated Video", autoplay=True, elem_id="generated-video")
+            
+            # --- Frame Grabbing UI ---
+            with gr.Row():
+                grab_frame_btn = gr.Button("📸 Use Current Frame as Input", variant="secondary")
+                timestamp_box = gr.Number(value=0, label="Timestamp", visible=True, elem_id="hidden-timestamp")
+            # -------------------------
+            
             file_output = gr.File(label="Download Video")
 
     ui_inputs = [
         input_image_component, last_image_component, prompt_input, steps_slider,
         negative_prompt_input, duration_seconds_input,
         guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox,
-        quality_slider, scheduler_dropdown, flow_shift_slider, frame_multi
+        quality_slider, scheduler_dropdown, flow_shift_slider, frame_multi,
+        play_result_video
     ]
     
     generate_button.click(
@@ -511,6 +592,25 @@ with gr.Blocks(delete_cache=(3600, 10800)) as demo:
         inputs=ui_inputs, 
         outputs=[video_output, file_output, seed_input]
     )
+    
+    # --- Frame Grabbing Events ---
+    # 1. Click button -> JS runs -> puts time in hidden number box
+    grab_frame_btn.click(
+        fn=None,
+        inputs=None,
+        outputs=[timestamp_box],
+        js=get_timestamp_js
+    )
+    
+    # 2. Hidden number box changes -> Python runs -> puts frame in Input Image
+    timestamp_box.change(
+        fn=extract_frame,
+        inputs=[video_output, timestamp_box],
+        outputs=[input_image_component]
+    )
 
 if __name__ == "__main__":
-    demo.queue().launch(mcp_server=True)
+    demo.queue().launch(
+        mcp_server=True,
+        css=CSS,
+    )