Update app.py

app.py

@@ -103,7 +103,7 @@ if PIPELINE_CONFIG_YAML.get("spatial_upscaler_model_path"):
 target_inference_device = "cuda"
 print(f"Target inference device: {target_inference_device}")
 pipeline_instance.to(target_inference_device)
-if latent_upsampler_instance: 
+if latent_upsampler_instance:
     latent_upsampler_instance.to(target_inference_device)
 
 
@@ -125,22 +125,22 @@ def calculate_new_dimensions(orig_w, orig_h):
         new_h = TARGET_FIXED_SIDE
         aspect_ratio = orig_w / orig_h
         new_w_ideal = new_h * aspect_ratio
-        
+
         # Round to nearest multiple of 32
         new_w = round(new_w_ideal / 32) * 32
-        
+
         # Clamp to [MIN_DIM_SLIDER, MAX_IMAGE_SIZE]
         new_w = max(MIN_DIM_SLIDER, min(new_w, MAX_IMAGE_SIZE))
         # Ensure new_h is also clamped (TARGET_FIXED_SIDE should be within these bounds if configured correctly)
-        new_h = max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE)) 
+        new_h = max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE))
     else:  # Portrait
         new_w = TARGET_FIXED_SIDE
         aspect_ratio = orig_h / orig_w # Use H/W ratio for portrait scaling
         new_h_ideal = new_w * aspect_ratio
-        
+
         # Round to nearest multiple of 32
         new_h = round(new_h_ideal / 32) * 32
-        
+
         # Clamp to [MIN_DIM_SLIDER, MAX_IMAGE_SIZE]
         new_h = max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE))
         # Ensure new_w is also clamped
@@ -165,37 +165,55 @@ def get_duration(prompt, negative_prompt, input_image_filepath, input_video_file
     else:
         return 45
 
+# --- NEW ---
+# This function handles the logic for extracting the last frame of the output video.
+def use_last_frame_as_input(video_filepath):
+    """
+    Extracts the last frame from a video, saves it as a temp image,
+    and returns its path to update the image-to-video tab.
+    Also returns an update to switch to the image-to-video tab.
+    """
+    if not video_filepath or not os.path.exists(video_filepath):
+        gr.Warning("No video available to get the last frame from.")
+        return None, gr.update()
+
+    try:
+        print(f"Extracting last frame from {video_filepath}")
+        with imageio.get_reader(video_filepath) as reader:
+            # Iterating is a robust way to get the last frame if len() is not available
+            last_frame_np = None
+            for frame in reader:
+                last_frame_np = frame
+
+            if last_frame_np is None:
+                raise ValueError("Could not read any frames from the video.")
+
+        pil_image = Image.fromarray(last_frame_np)
+
+        # Save to a temporary file
+        temp_dir = tempfile.mkdtemp()
+        timestamp = random.randint(10000, 99999)
+        output_image_path = os.path.join(temp_dir, f"last_frame_{timestamp}.png")
+        pil_image.save(output_image_path)
+        print(f"Saved last frame to {output_image_path}")
+
+        # Return the path to the new image and an update to select the i2v tab
+        return output_image_path, gr.Tab(selected=True)
+
+    except Exception as e:
+        print(f"Error extracting last frame: {e}")
+        gr.Error(f"Failed to extract the last frame: {e}")
+        return None, gr.update()
+
 @spaces.GPU(duration=get_duration)
 def generate(prompt, negative_prompt, input_image_filepath=None, input_video_filepath=None,
              height_ui=512, width_ui=704, mode="text-to-video",
-             duration_ui=2.0, 
+             duration_ui=2.0,
              ui_frames_to_use=9,
              seed_ui=42, randomize_seed=True, ui_guidance_scale=3.0, improve_texture_flag=True, num_steps=20, fps=30.0,
              progress=gr.Progress(track_tqdm=True)):
-    """
-    Generate high-quality videos using LTX Video model with support for text-to-video, image-to-video, and video-to-video modes.
-
-    Args:
-        prompt (str): Text description of the desired video content. Required for all modes.
-        negative_prompt (str): Text describing what to avoid in the generated video. Optional, can be empty string.
-        input_image_filepath (str or None): Path to input image file. Required for image-to-video mode, None for other modes.
-        input_video_filepath (str or None): Path to input video file. Required for video-to-video mode, None for other modes.
-        height_ui (int): Height of the output video in pixels, must be divisible by 32. Default: 512.
-        width_ui (int): Width of the output video in pixels, must be divisible by 32. Default: 704.
-        mode (str): Generation mode. Required. One of "text-to-video", "image-to-video", or "video-to-video". Default: "text-to-video".
-        duration_ui (float): Duration of the output video in seconds. Range: 0.3 to 8.5. Default: 2.0.
-        ui_frames_to_use (int): Number of frames to use from input video. Only used in video-to-video mode. Must be N*8+1. Default: 9.
-        seed_ui (int): Random seed for reproducible generation. Range: 0 to 2^32-1. Default: 42.
-        randomize_seed (bool): Whether to use a random seed instead of seed_ui. Default: True.
-        ui_guidance_scale (float): CFG scale controlling prompt influence. Range: 1.0 to 10.0. Higher values = stronger prompt influence. Default: 3.0.
-        improve_texture_flag (bool): Whether to use multi-scale generation for better texture quality. Slower but higher quality. Default: True.
-        progress (gr.Progress): Progress tracker for the generation process. Optional, used for UI updates.
-
-    Returns:
-        tuple: A tuple containing (output_video_path, used_seed) where output_video_path is the path to the generated video file and used_seed is the actual seed used for generation.
-    """
 
-    # Validate mode-specific required parameters
+    # ... (the beginning of the generate function is unchanged)
     if mode == "image-to-video":
         if not input_image_filepath:
             raise gr.Error("input_image_filepath is required for image-to-video mode")
@@ -203,7 +221,6 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
         if not input_video_filepath:
             raise gr.Error("input_video_filepath is required for video-to-video mode")
     elif mode == "text-to-video":
-        # No additional file inputs required for text-to-video
         pass
     else:
         raise gr.Error(f"Invalid mode: {mode}. Must be one of: text-to-video, image-to-video, video-to-video")
@@ -211,27 +228,27 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
     if randomize_seed:
         seed_ui = random.randint(0, 2**32 - 1)
     seed_everething(int(seed_ui))
-    
+
     target_frames_ideal = duration_ui * fps
     target_frames_rounded = round(target_frames_ideal)
-    if target_frames_rounded < 1: 
+    if target_frames_rounded < 1:
         target_frames_rounded = 1
-    
+
     n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
     actual_num_frames = int(n_val * 8 + 1)
 
     actual_num_frames = max(9, actual_num_frames)
     actual_num_frames = min(MAX_NUM_FRAMES, actual_num_frames)
-    
+
     actual_height = int(height_ui)
     actual_width = int(width_ui)
 
     height_padded = ((actual_height - 1) // 32 + 1) * 32
     width_padded = ((actual_width - 1) // 32 + 1) * 32
-    num_frames_padded = ((actual_num_frames - 2) // 8 + 1) * 8 + 1 
+    num_frames_padded = ((actual_num_frames - 2) // 8 + 1) * 8 + 1
     if num_frames_padded != actual_num_frames:
         print(f"Warning: actual_num_frames ({actual_num_frames}) and num_frames_padded ({num_frames_padded}) differ. Using num_frames_padded for pipeline.")
-    
+
     padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
 
     call_kwargs = {
@@ -239,11 +256,11 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
         "negative_prompt": negative_prompt,
         "height": height_padded,
         "width": width_padded,
-        "num_frames": num_frames_padded, 
+        "num_frames": num_frames_padded,
         "num_inference_steps": num_steps,
-        "frame_rate": int(fps), 
+        "frame_rate": int(fps),
         "generator": torch.Generator(device=target_inference_device).manual_seed(int(seed_ui)),
-        "output_type": "pt", 
+        "output_type": "pt",
         "conditioning_items": None,
         "media_items": None,
         "decode_timestep": PIPELINE_CONFIG_YAML["decode_timestep"],
@@ -283,9 +300,9 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
         try:
             call_kwargs["media_items"] = load_media_file(
                 media_path=input_video_filepath,
-                height=actual_height, 
+                height=actual_height,
                 width=actual_width,
-                max_frames=int(ui_frames_to_use), 
+                max_frames=int(ui_frames_to_use),
                 padding=padding_values
             ).to(target_inference_device)
         except Exception as e:
@@ -293,7 +310,7 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
             raise gr.Error(f"Could not load video: {e}")
 
     print(f"Moving models to {target_inference_device} for inference (if not already there)...")
-    
+
     active_latent_upsampler = None
     if improve_texture_flag and latent_upsampler_instance:
         active_latent_upsampler = latent_upsampler_instance
@@ -302,27 +319,24 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
     if improve_texture_flag:
         if not active_latent_upsampler:
             raise gr.Error("Spatial upscaler model not loaded or improve_texture not selected, cannot use multi-scale.")
-        
+
         multi_scale_pipeline_obj = LTXMultiScalePipeline(pipeline_instance, active_latent_upsampler)
-        
+
         first_pass_args = PIPELINE_CONFIG_YAML.get("first_pass", {}).copy()
-        first_pass_args["guidance_scale"] = float(ui_guidance_scale) # UI overrides YAML
-        # num_inference_steps will be derived from len(timesteps) in the pipeline
+        first_pass_args["guidance_scale"] = float(ui_guidance_scale)
         first_pass_args.pop("num_inference_steps", None)
 
-
         second_pass_args = PIPELINE_CONFIG_YAML.get("second_pass", {}).copy()
-        second_pass_args["guidance_scale"] = float(ui_guidance_scale) # UI overrides YAML
-        # num_inference_steps will be derived from len(timesteps) in the pipeline
+        second_pass_args["guidance_scale"] = float(ui_guidance_scale)
         second_pass_args.pop("num_inference_steps", None)
-        
+
         multi_scale_call_kwargs = call_kwargs.copy()
         multi_scale_call_kwargs.update({
             "downscale_factor": PIPELINE_CONFIG_YAML["downscale_factor"],
             "first_pass": first_pass_args,
             "second_pass": second_pass_args,
         })
-        
+
         print(f"Calling multi-scale pipeline (eff. HxW: {actual_height}x{actual_width}, Frames: {actual_num_frames} -> Padded: {num_frames_padded}) on {target_inference_device}")
         result_images_tensor = multi_scale_pipeline_obj(**multi_scale_call_kwargs).images
     else:
@@ -330,17 +344,16 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
         first_pass_config_from_yaml = PIPELINE_CONFIG_YAML.get("first_pass", {})
 
         single_pass_call_kwargs["timesteps"] = first_pass_config_from_yaml.get("timesteps")
-        single_pass_call_kwargs["guidance_scale"] = float(ui_guidance_scale) # UI overrides YAML
+        single_pass_call_kwargs["guidance_scale"] = float(ui_guidance_scale)
         single_pass_call_kwargs["stg_scale"] = first_pass_config_from_yaml.get("stg_scale")
         single_pass_call_kwargs["rescaling_scale"] = first_pass_config_from_yaml.get("rescaling_scale")
         single_pass_call_kwargs["skip_block_list"] = first_pass_config_from_yaml.get("skip_block_list")
-        
-        # Remove keys that might conflict or are not used in single pass / handled by above
-        single_pass_call_kwargs.pop("num_inference_steps", None) 
-        single_pass_call_kwargs.pop("first_pass", None) 
+
+        single_pass_call_kwargs.pop("num_inference_steps", None)
+        single_pass_call_kwargs.pop("first_pass", None)
         single_pass_call_kwargs.pop("second_pass", None)
         single_pass_call_kwargs.pop("downscale_factor", None)
-        
+
         print(f"Calling base pipeline (padded HxW: {height_padded}x{width_padded}, Frames: {actual_num_frames} -> Padded: {num_frames_padded}) on {target_inference_device}")
         result_images_tensor = pipeline_instance(**single_pass_call_kwargs).images
 
@@ -350,20 +363,20 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
     pad_left, pad_right, pad_top, pad_bottom = padding_values
     slice_h_end = -pad_bottom if pad_bottom > 0 else None
     slice_w_end = -pad_right if pad_right > 0 else None
-    
+
     result_images_tensor = result_images_tensor[
         :, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end
     ]
 
     video_np = result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy()
-    
-    video_np = np.clip(video_np, 0, 1) 
+
+    video_np = np.clip(video_np, 0, 1)
     video_np = (video_np * 255).astype(np.uint8)
 
     temp_dir = tempfile.mkdtemp()
     timestamp = random.randint(10000,99999)
     output_video_path = os.path.join(temp_dir, f"output_{timestamp}.mp4")
-    
+
     try:
         with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], macro_block_size=1) as video_writer:
             for frame_idx in range(video_np.shape[0]):
@@ -379,8 +392,10 @@ def generate(prompt, negative_prompt, input_image_filepath=None, input_video_fil
         except Exception as e2:
             print(f"Fallback video saving error: {e2}")
             raise gr.Error(f"Failed to save video: {e2}")
-            
-    return output_video_path, seed_ui
+
+    # --- MODIFIED ---
+    # The function now returns a third value: an update to make the new button visible.
+    return output_video_path, seed_ui, gr.update(visible=True)
 
 def update_task_image():
     return "image-to-video"
@@ -401,11 +416,10 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown("# LTX Video 0.9.8 13B Distilled")
     gr.Markdown("Fast high quality video generation.**Update (17/07):** now with the new v0.9.8 for improved prompt understanding and detail generation" )
     gr.Markdown("[Model](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.8-distilled.safetensors) [GitHub](https://github.com/Lightricks/LTX-Video) [Diffusers](https://huggingface.co/Lightricks/LTX-Video-0.9.8-13B-distilled#diffusers-🧨)")
-    
+
     with gr.Row():
         with gr.Column():
             with gr.Tab("image-to-video") as image_tab:
-                # The hidden textboxes are a good way to manage state for each tab
                 video_i_hidden = gr.Textbox(label="video_i", visible=False, value=None)
                 image_i2v = gr.Image(label="Input Image", type="filepath", sources=["upload", "webcam", "clipboard"])
                 i2v_prompt = gr.Textbox(label="Prompt", value="The creature from the image starts to move", lines=3)
@@ -416,8 +430,7 @@ with gr.Blocks(css=css) as demo:
                 video_n_hidden = gr.Textbox(label="video_n", visible=False, value=None)
                 t2v_prompt = gr.Textbox(label="Prompt", value="A majestic dragon flying over a medieval castle", lines=3)
                 t2v_button = gr.Button("Generate Text-to-Video", variant="primary")
-            
-            # This tab was set to visible=False, making it inaccessible. We need to change that.
+
             with gr.Tab("video-to-video") as video_tab:
                 image_v_hidden = gr.Textbox(label="image_v", visible=False, value=None)
                 video_v2v = gr.Video(label="Input Video", sources=["upload", "webcam"])
@@ -426,20 +439,22 @@ with gr.Blocks(css=css) as demo:
                 v2v_button = gr.Button("Generate Video-to-Video", variant="primary")
 
             duration_input = gr.Slider(
-                label="Video Duration (seconds)", 
-                minimum=3.0, 
-                maximum=60.0, 
-                value=5.0,  
-                step=0.1, 
+                label="Video Duration (seconds)",
+                minimum=3.0,
+                maximum=60.0,
+                value=5.0,
+                step=0.1,
                 info=f"Target video duration (0.3s to 8.5s)"
             )
             improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=True, visible=True, info="Uses a two-pass generation for better quality, but is slower. Recommended for final output.")
 
         with gr.Column():
             output_video = gr.Video(label="Generated Video", interactive=False)
-            
+            # --- NEW ---
+            # This is the new button, initially hidden.
+            use_last_frame_button = gr.Button("Use Last Frame as Input Image", visible=False)
+
     with gr.Accordion("Advanced settings", open=False):
-        # We'll use this dropdown to track the currently selected mode from the tabs
         mode = gr.Dropdown(["text-to-video", "image-to-video", "video-to-video"], label="task", value="image-to-video", visible=False)
         negative_prompt_input = gr.Textbox(label="Negative Prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted", lines=2)
         with gr.Row():
@@ -471,11 +486,11 @@ with gr.Blocks(css=css) as demo:
         if not video_filepath:
             return gr.update(value=current_h), gr.update(value=current_w)
         try:
-            video_filepath_str = str(video_filepath) 
+            video_filepath_str = str(video_filepath)
             if not os.path.exists(video_filepath_str):
                 print(f"Video file path does not exist for dimension update: {video_filepath_str}")
                 return gr.update(value=current_h), gr.update(value=current_w)
-            
+
             orig_w, orig_h = -1, -1
             with imageio.get_reader(video_filepath_str) as reader:
                 meta = reader.get_meta_data()
@@ -488,7 +503,7 @@ with gr.Blocks(css=css) as demo:
                     except Exception as e_frame:
                         print(f"Could not get video size from metadata or first frame: {e_frame}")
                         return gr.update(value=current_h), gr.update(value=current_w)
-            
+
             if orig_w == -1 or orig_h == -1:
                  print(f"Could not determine dimensions for video: {video_filepath_str}")
                  return gr.update(value=current_h), gr.update(value=current_w)
@@ -498,7 +513,7 @@ with gr.Blocks(css=css) as demo:
         except Exception as e:
             print(f"Error processing video for dimension update: {e} (Path: {video_filepath}, Type: {type(video_filepath)})")
             return gr.update(value=current_h), gr.update(value=current_w)
-    
+
     image_i2v.upload(
         fn=handle_image_upload_for_dims,
         inputs=[image_i2v, height_input, width_input],
@@ -518,34 +533,50 @@ with gr.Blocks(css=css) as demo:
         fn=update_task_text,
         outputs=[mode]
     )
-    
-    # This is the new, crucial event handler for the video tab
     video_tab.select(
         fn=update_task_video,
         outputs=[mode]
     )
 
+    # --- MODIFIED ---
+    # The outputs for the generate buttons now include `use_last_frame_button` to control its visibility.
+    # We also use the `.then()` method to chain events. First, we instantly hide the button,
+    # then we run the main (and slower) video generation process.
     t2v_inputs = [t2v_prompt, negative_prompt_input, image_n_hidden, video_n_hidden,
                   height_input, width_input, mode,
-                  duration_input, frames_to_use, 
+                  duration_input, frames_to_use,
                   seed_input, randomize_seed_input, guidance_scale_input, improve_texture, num_steps, fps]
-    
     i2v_inputs = [i2v_prompt, negative_prompt_input, image_i2v, video_i_hidden,
                   height_input, width_input, mode,
-                  duration_input, frames_to_use, 
+                  duration_input, frames_to_use,
                   seed_input, randomize_seed_input, guidance_scale_input, improve_texture, num_steps, fps]
-
     v2v_inputs = [v2v_prompt, negative_prompt_input, image_v_hidden, video_v2v,
                   height_input, width_input, mode,
-                  duration_input, frames_to_use, 
+                  duration_input, frames_to_use,
                   seed_input, randomize_seed_input, guidance_scale_input, improve_texture, num_steps, fps]
 
-    t2v_button.click(fn=generate, inputs=t2v_inputs, outputs=[output_video, seed_input], api_name="text_to_video")
-    i2v_button.click(fn=generate, inputs=i2v_inputs, outputs=[output_video, seed_input], api_name="image_to_video")
-    v2v_button.click(fn=generate, inputs=v2v_inputs, outputs=[output_video, seed_input], api_name="video_to_video")
+    t2v_button.click(fn=lambda: gr.update(visible=False), outputs=[use_last_frame_button], queue=False).then(
+        fn=generate, inputs=t2v_inputs, outputs=[output_video, seed_input, use_last_frame_button], api_name="text_to_video"
+    )
+    i2v_button.click(fn=lambda: gr.update(visible=False), outputs=[use_last_frame_button], queue=False).then(
+        fn=generate, inputs=i2v_inputs, outputs=[output_video, seed_input, use_last_frame_button], api_name="image_to_video"
+    )
+    v2v_button.click(fn=lambda: gr.update(visible=False), outputs=[use_last_frame_button], queue=False).then(
+        fn=generate, inputs=v2v_inputs, outputs=[output_video, seed_input, use_last_frame_button], api_name="video_to_video"
+    )
+
+    # --- NEW ---
+    # This is the event handler for our new button.
+    # It takes the generated video as input.
+    # It updates the image component in the first tab and also selects that tab.
+    use_last_frame_button.click(
+        fn=use_last_frame_as_input,
+        inputs=[output_video],
+        outputs=[image_i2v, image_tab]
+    )
 
 if __name__ == "__main__":
     if os.path.exists(models_dir) and os.path.isdir(models_dir):
         print(f"Model directory: {Path(models_dir).resolve()}")
-    
+
     demo.queue().launch(debug=True, share=False, mcp_server=True)