Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from gradio_client import Client, handle_file
 from google import genai
 import os
-from typing import Optional, List
+from typing import Optional, List, Tuple
 from huggingface_hub import whoami
 from PIL import Image
 from io import BytesIO
@@ -35,28 +35,35 @@ def _extract_image_data_from_response(response) -> Optional[bytes]:
                 return part.inline_data.data
     return None
 
-def _get_framerate(video_path: str) -> float:
-    """Instantly gets the framerate of a video using ffprobe."""
+def _get_video_info(video_path: str) -> Tuple[float, Tuple[int, int]]:
+    """Instantly gets the framerate and (width, height) of a video using ffprobe."""
     probe = ffmpeg.probe(video_path)
-    video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
-    if video_stream is None:
-        raise ValueError("Could not find video stream in the file.")
-    return eval(video_stream['avg_frame_rate'])
+    video_stream = next((s for s in probe['streams'] if s['codec_type'] == 'video'), None)
+    if not video_stream:
+        raise ValueError("No video stream found in the file.")
+    framerate = eval(video_stream['avg_frame_rate'])
+    resolution = (int(video_stream['width']), int(video_stream['height']))
+    return framerate, resolution
+
+def _resize_image(image_path: str, target_size: Tuple[int, int]) -> str:
+    """Resizes an image to a target size and saves it to a new temp file."""
+    with Image.open(image_path) as img:
+        if img.size == target_size:
+            return image_path
+        resized_img = img.resize(target_size, Image.Resampling.LANCZOS)
+        suffix = os.path.splitext(image_path)[1] or ".png"
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
+            resized_img.save(tmp_file.name)
+            return tmp_file.name
 
 def _trim_first_frame_fast(video_path: str) -> str:
-    """
-    Removes exactly the first frame of a video without re-encoding.
-    This is the frame-accurate and fast method.
-    """
+    """Removes exactly the first frame of a video without re-encoding."""
     with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_output_file:
         output_path = tmp_output_file.name
-    
     try:
-        framerate = _get_framerate(video_path)
+        framerate, _ = _get_video_info(video_path)
         if framerate == 0: raise ValueError("Framerate cannot be zero.")
         start_time = 1 / framerate
-
-        # The key is placing -ss AFTER -i for accuracy, combined with -c copy for speed.
         (
             ffmpeg
             .input(video_path, ss=start_time)
@@ -68,17 +75,14 @@ def _trim_first_frame_fast(video_path: str) -> str:
         raise RuntimeError(f"FFmpeg trim error: {e}")
 
 def _combine_videos_simple(video1_path: str, video2_path: str) -> str:
-    """
-    Combines two videos using the fast concat demuxer. Assumes video2 is already trimmed.
-    """
+    """Combines two videos using the fast concat demuxer."""
+
     with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".txt") as tmp_list_file:
         tmp_list_file.write(f"file '{os.path.abspath(video1_path)}'\n")
         tmp_list_file.write(f"file '{os.path.abspath(video2_path)}'\n")
         list_file_path = tmp_list_file.name
-
     with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_output_file:
         output_path = tmp_output_file.name
-
     try:
         (
             ffmpeg
@@ -104,9 +108,6 @@ def _generate_video_segment(input_image_path: str, output_image_path: str, promp
     return result[0]["video"]
 
 def unified_image_generator(prompt: str, images: Optional[List[str]], previous_video_path: Optional[str], oauth_token: Optional[gr.OAuthToken]) -> tuple:
-    """
-    Handles image generation and determines the visibility of video creation buttons.
-    """
     if not verify_pro_status(oauth_token): raise gr.Error("Access Denied.")
     try:
         contents = [Image.open(image_path[0]) for image_path in images] if images else []
@@ -114,25 +115,16 @@ def unified_image_generator(prompt: str, images: Optional[List[str]], previous_v
         response = client.models.generate_content(model=GEMINI_MODEL_NAME, contents=contents)
         image_data = _extract_image_data_from_response(response)
         if not image_data: raise ValueError("No image data in response.")
-        
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
             Image.open(BytesIO(image_data)).save(tmp.name)
             output_path = tmp.name
-        
         can_create_video = bool(images and len(images) == 1)
         can_extend_video = can_create_video and bool(previous_video_path)
-        
-        return (
-            output_path,
-            gr.update(visible=can_create_video),
-            gr.update(visible=can_extend_video),
-            gr.update(visible=False)
-        )
+        return (output_path, gr.update(visible=can_create_video), gr.update(visible=can_extend_video), gr.update(visible=False))
     except Exception as e:
         raise gr.Error(f"Image generation failed: {e}")
 
 def create_new_video(input_image_gallery: List[str], prompt_input: str, output_image: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
-    """Starts a NEW video chain, overwriting any previous video state."""
     if not verify_pro_status(oauth_token): raise gr.Error("Access Denied.")
     if not input_image_gallery or not output_image: raise gr.Error("Input/output images required.")
     try:
@@ -142,12 +134,14 @@ def create_new_video(input_image_gallery: List[str], prompt_input: str, output_i
         raise gr.Error(f"Video creation failed: {e}")
 
 def extend_existing_video(input_image_gallery: List[str], prompt_input: str, output_image: str, previous_video_path: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
-    """Extends an existing video with a new segment."""
     if not verify_pro_status(oauth_token): raise gr.Error("Access Denied.")
     if not previous_video_path: raise gr.Error("No previous video to extend.")
     if not input_image_gallery or not output_image: raise gr.Error("Input/output images required.")
     try:
-        new_segment_path = _generate_video_segment(input_image_gallery[0][0], output_image, prompt_input, oauth_token.token)
+        _, target_resolution = _get_video_info(previous_video_path)
+        resized_input_path = _resize_image(input_image_gallery[0][0], target_resolution)
+        resized_output_path = _resize_image(output_image, target_resolution)
+        new_segment_path = _generate_video_segment(resized_input_path, resized_output_path, prompt_input, oauth_token.token)
         trimmed_segment_path = _trim_first_frame_fast(new_segment_path)
         final_video_path = _combine_videos_simple(previous_video_path, trimmed_segment_path)
         return final_video_path, final_video_path
@@ -189,8 +183,8 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                 output_image = gr.Image(label="Output", interactive=False, elem_id="output", type="filepath")
                 use_image_button = gr.Button("♻️ Use this Image for Next Edit", variant="primary")
                 with gr.Row():
-                    create_video_button = gr.Button("Create video between the two images 🎥", variant="secondary", visible=False)
-                    extend_video_button = gr.Button("Extend previous video with new scene 🎞️", variant="secondary", visible=False)
+                    create_video_button = gr.Button("Create a video between the two images 🎥", variant="secondary", visible=False)
+                    extend_video_button = gr.Button("Extend video with new scene 🎞️", variant="secondary", visible=False)
                 with gr.Group(visible=False) as video_group:
                     video_output = gr.Video(label="Generated Video", show_download_button=True, autoplay=True)
                     gr.Markdown("Generate more with [Wan 2.2 first-last-frame](https://huggingface.co/spaces/multimodalart/wan-2-2-first-last-frame)", elem_id="wan_ad")
@@ -204,19 +198,14 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
         inputs=[prompt_input, image_input_gallery, previous_video_state],
         outputs=[output_image, create_video_button, extend_video_button, video_group]
     )
-
     use_image_button.click(
         fn=lambda img: (
-            [img] if img else None,
-            None,
-            gr.update(visible=False),
-            gr.update(visible=False),
-            gr.update(visible=False)
+            [img] if img else None, None, gr.update(visible=False),
+            gr.update(visible=False), gr.update(visible=False)
         ),
         inputs=[output_image],
         outputs=[image_input_gallery, output_image, create_video_button, extend_video_button, video_group]
     )
-
     create_video_button.click(
         fn=lambda: gr.update(visible=True), outputs=[video_group]
     ).then(
@@ -224,7 +213,6 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
         inputs=[image_input_gallery, prompt_input, output_image],
         outputs=[video_output, previous_video_state],
     )
-
     extend_video_button.click(
         fn=lambda: gr.update(visible=True), outputs=[video_group]
     ).then(