Update app.py

app.py

@@ -1,7 +1,4 @@
-# --- NEW ---
-# Add moviepy for video stitching. Make sure to install it: pip install moviepy
 from moviepy.editor import VideoFileClip, concatenate_videoclips
-
 import spaces
 import os
 os.putenv('PYTORCH_NVML_BASED_CUDA_CHECK','1')
@@ -56,7 +53,6 @@ LTX_REPO = "Lightricks/LTX-Video"
 MAX_IMAGE_SIZE = PIPELINE_CONFIG_YAML.get("max_resolution", 1280)
 MAX_NUM_FRAMES = 900
 
-# --- Global variables for loaded models ---
 pipeline_instance = None
 latent_upsampler_instance = None
 models_dir = "downloaded_models_gradio_cpu_init"
@@ -110,8 +106,6 @@ pipeline_instance.to(target_inference_device)
 if latent_upsampler_instance:
     latent_upsampler_instance.to(target_inference_device)
 
-
-# --- Helper functions ---
 def calculate_new_dimensions(orig_w, orig_h):
     if orig_w == 0 or orig_h == 0:
         return int(768), int(768)
@@ -123,7 +117,7 @@ def calculate_new_dimensions(orig_w, orig_h):
         new_h = round((new_w * (orig_h / orig_w)) / 32) * 32
     return int(max(256, min(new_h, MAX_IMAGE_SIZE))), int(max(256, min(new_w, MAX_IMAGE_SIZE)))
 
-def get_duration(*args, **kwargs): # Simplified for brevity
+def get_duration(*args, **kwargs):
     duration_ui = kwargs.get('duration_ui', 5.0)
     if duration_ui > 20.0: return 120
     if duration_ui > 13.0: return 90
@@ -137,7 +131,12 @@ def use_last_frame_as_input(video_filepath):
         return None, gr.update()
     try:
         with imageio.get_reader(video_filepath) as reader:
-            last_frame_np = next(reversed(list(reader)))
+            # A robust way to get the last frame
+            frame_list = list(reader)
+            if not frame_list:
+                 raise ValueError("Video file is empty or unreadable.")
+            last_frame_np = frame_list[-1]
+
         pil_image = Image.fromarray(last_frame_np)
         output_image_path = os.path.join(tempfile.mkdtemp(), f"last_frame_{random.randint(10000,99999)}.png")
         pil_image.save(output_image_path)
@@ -146,123 +145,105 @@ def use_last_frame_as_input(video_filepath):
         gr.Error(f"Failed to extract the last frame: {e}")
         return None, gr.update()
 
-# --- NEW ---
-# Function to stitch video clips together using moviepy
 def stitch_videos(clips_list):
     if not clips_list or len(clips_list) < 2:
         raise gr.Error("You need at least two clips to stitch them together!")
-
     print(f"Stitching {len(clips_list)} clips...")
     try:
         video_clips = [VideoFileClip(clip_path) for clip_path in clips_list]
         final_clip = concatenate_videoclips(video_clips, method="compose")
-        
         final_output_path = os.path.join(tempfile.mkdtemp(), f"stitched_video_{random.randint(10000,99999)}.mp4")
         final_clip.write_videofile(final_output_path, codec="libx264", audio=False, threads=4, preset='ultrafast')
-        
-        # Close all clip objects to release file handles
         for clip in video_clips:
             clip.close()
-            
         print(f"Final video saved to {final_output_path}")
         return final_output_path
     except Exception as e:
         print(f"Error during video stitching: {e}")
         raise gr.Error(f"Failed to stitch videos: {e}")
 
-# --- NEW ---
-# Function to clear the clip list and reset the UI
 def clear_clips():
     return [], "Clips created: 0", None, None
 
+# --- MODIFIED ---
+# The function signature has been reordered to match the way the inputs are assembled later.
 @spaces.GPU(duration=get_duration)
-def generate(prompt, negative_prompt, clips_list, # --- MODIFIED --- added clips_list
-             input_image_filepath=None, input_video_filepath=None,
-             height_ui=512, width_ui=704, mode="text-to-video",
-             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,
+def generate(prompt, negative_prompt, clips_list, input_image_filepath, input_video_filepath,
+             height_ui, width_ui, mode, duration_ui, ui_frames_to_use,
+             seed_ui, randomize_seed, ui_guidance_scale, improve_texture_flag, num_steps, fps,
              progress=gr.Progress(track_tqdm=True)):
-    
-    # ... (most of the generate function logic is unchanged) ...
+
+    if mode not in ["text-to-video", "image-to-video", "video-to-video"]:
+        raise gr.Error(f"Invalid mode: {mode}.")
     if mode == "image-to-video" and not input_image_filepath:
         raise gr.Error("input_image_filepath is required for image-to-video mode")
     elif mode == "video-to-video" and not input_video_filepath:
         raise gr.Error("input_video_filepath is required for video-to-video mode")
-    elif mode not in ["text-to-video", "image-to-video", "video-to-video"]:
-        raise gr.Error(f"Invalid mode: {mode}.")
+
     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 = max(1, round(target_frames_ideal))
     n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
     actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(n_val * 8 + 1)))
-    actual_height = int(height_ui)
-    actual_width = int(width_ui)
+    actual_height, actual_width = int(height_ui), 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
     padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
+    
     call_kwargs = {
         "prompt": prompt, "negative_prompt": negative_prompt, "height": height_padded, "width": width_padded,
-        "num_frames": num_frames_padded, "num_inference_steps": num_steps, "frame_rate": int(fps),
+        "num_frames": max(9, ((actual_num_frames - 2) // 8 + 1) * 8 + 1), "num_inference_steps": num_steps, "frame_rate": int(fps),
         "generator": torch.Generator(device=target_inference_device).manual_seed(int(seed_ui)), "output_type": "pt",
-        "decode_timestep": PIPELINE_CONFIG_YAML["decode_timestep"], "decode_noise_scale": PIPELINE_CONFIG_YAML["decode_noise_scale"],
-        "stochastic_sampling": PIPELINE_CONFIG_YAML["stochastic_sampling"], "image_cond_noise_scale": 0.15, "is_video": True,
-        "vae_per_channel_normalize": True, "mixed_precision": (PIPELINE_CONFIG_YAML["precision"] == "mixed_precision"),
-        "offload_to_cpu": False, "enhance_prompt": False, "conditioning_items": None, "media_items": None
+        **PIPELINE_CONFIG_YAML,
+        "is_video": True, "mixed_precision": (PIPELINE_CONFIG_YAML["precision"] == "mixed_precision"),
+        "offload_to_cpu": False, "enhance_prompt": False,
     }
-    stg_mode_str = PIPELINE_CONFIG_YAML.get("stg_mode", "attention_values")
+    
+    # This logic can be simplified and organized
+    stg_mode_str = PIPELINE_CONFIG_YAML.get("stg_mode", "attention_values").lower()
     stg_map = {
         "stg_av": SkipLayerStrategy.AttentionValues, "attention_values": SkipLayerStrategy.AttentionValues,
         "stg_as": SkipLayerStrategy.AttentionSkip, "attention_skip": SkipLayerStrategy.AttentionSkip,
         "stg_r": SkipLayerStrategy.Residual, "residual": SkipLayerStrategy.Residual,
         "stg_t": SkipLayerStrategy.TransformerBlock, "transformer_block": SkipLayerStrategy.TransformerBlock,
     }
-    call_kwargs["skip_layer_strategy"] = stg_map.get(stg_mode_str.lower())
-    if call_kwargs["skip_layer_strategy"] is None:
-        raise ValueError(f"Invalid stg_mode: {stg_mode_str}")
+    call_kwargs["skip_layer_strategy"] = stg_map.get(stg_mode_str, SkipLayerStrategy.AttentionValues)
+
     if mode == "image-to-video":
         media_tensor = load_image_to_tensor_with_resize_and_crop(input_image_filepath, actual_height, actual_width)
-        media_tensor = torch.nn.functional.pad(media_tensor, padding_values)
-        call_kwargs["conditioning_items"] = [ConditioningItem(media_tensor.to(target_inference_device), 0, 1.0)]
+        call_kwargs["conditioning_items"] = [ConditioningItem(torch.nn.functional.pad(media_tensor, padding_values).to(target_inference_device), 0, 1.0)]
     elif mode == "video-to-video":
         call_kwargs["media_items"] = load_media_file(media_path=input_video_filepath, height=actual_height, width=actual_width, max_frames=int(ui_frames_to_use), padding=padding_values).to(target_inference_device)
-    active_latent_upsampler = latent_upsampler_instance if improve_texture_flag and latent_upsampler_instance else None
-    if improve_texture_flag:
-        if not active_latent_upsampler:
-            raise gr.Error("Spatial upscaler model not loaded or improve_texture not selected.")
-        multi_scale_pipeline_obj = LTXMultiScalePipeline(pipeline_instance, active_latent_upsampler)
-        first_pass_args = {**PIPELINE_CONFIG_YAML.get("first_pass", {}), "guidance_scale": float(ui_guidance_scale)}
-        second_pass_args = {**PIPELINE_CONFIG_YAML.get("second_pass", {}), "guidance_scale": float(ui_guidance_scale)}
-        multi_scale_call_kwargs = {**call_kwargs, "downscale_factor": PIPELINE_CONFIG_YAML["downscale_factor"], "first_pass": first_pass_args, "second_pass": second_pass_args}
-        result_images_tensor = multi_scale_pipeline_obj(**multi_scale_call_kwargs).images
+
+    if improve_texture_flag and latent_upsampler_instance:
+        multi_scale_pipeline = LTXMultiScalePipeline(pipeline_instance, latent_upsampler_instance)
+        pass_args = {"guidance_scale": float(ui_guidance_scale)}
+        multi_scale_kwargs = {**call_kwargs, "first_pass": {**PIPELINE_CONFIG_YAML.get("first_pass", {}), **pass_args}, "second_pass": {**PIPELINE_CONFIG_YAML.get("second_pass", {}), **pass_args}}
+        result_images_tensor = multi_scale_pipeline(**multi_scale_kwargs).images
     else:
-        first_pass_config = PIPELINE_CONFIG_YAML.get("first_pass", {})
-        single_pass_call_kwargs = {**call_kwargs, "guidance_scale": float(ui_guidance_scale), **first_pass_config}
-        result_images_tensor = pipeline_instance(**single_pass_call_kwargs).images
+        single_pass_kwargs = {**call_kwargs, "guidance_scale": float(ui_guidance_scale), **PIPELINE_CONFIG_YAML.get("first_pass", {})}
+        result_images_tensor = pipeline_instance(**single_pass_kwargs).images
+
     if result_images_tensor is None:
         raise gr.Error("Generation failed.")
-    pad_left, pad_right, pad_top, pad_bottom = padding_values
-    result_images_tensor = result_images_tensor[:, :, :actual_num_frames, pad_top:(-pad_bottom or None), pad_left:(-pad_right or None)]
+    
+    pad_l, pad_r, pad_t, pad_b = padding_values
+    result_images_tensor = result_images_tensor[:, :, :actual_num_frames, pad_t:(-pad_b or None), pad_l:(-pad_r or None)]
     video_np = (np.clip(result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy(), 0, 1) * 255).astype(np.uint8)
+    
     output_video_path = os.path.join(tempfile.mkdtemp(), f"output_{random.randint(10000,99999)}.mp4")
-    try:
-        with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], macro_block_size=1) as video_writer:
-            for frame_idx, frame in enumerate(video_np):
-                progress(frame_idx / len(video_np), desc="Saving video")
-                video_writer.append_data(frame)
-    except Exception as e:
-        gr.Error(f"Failed to save video: {e}")
+    with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], macro_block_size=1) as video_writer:
+        for idx, frame in enumerate(video_np):
+            progress(idx / len(video_np), desc="Saving video clip...")
+            video_writer.append_data(frame)
 
-    # --- MODIFIED ---
-    # Append the new clip to the list and prepare the updated state and counter text
     updated_clips_list = clips_list + [output_video_path]
     counter_text = f"Clips created: {len(updated_clips_list)}"
-    
     return output_video_path, seed_ui, gr.update(visible=True), updated_clips_list, counter_text
 
-# ... (update_task functions are unchanged) ...
 def update_task_image(): return "image-to-video"
 def update_task_text(): return "text-to-video"
 def update_task_video(): return "video-to-video"
@@ -270,18 +251,13 @@ def update_task_video(): return "video-to-video"
 css="""#col-container{margin:0 auto;max-width:900px;}"""
 
 with gr.Blocks(css=css) as demo:
-    # --- NEW ---
-    # Add a state component to store the list of clip paths
     clips_state = gr.State([])
-
-    gr.Markdown("# LTX Video 0.9.8 13B Distilled")
+    gr.Markdown("# LTX Video Clip Stitcher")
     gr.Markdown("Generate short video clips and stitch them together to create a longer animation.")
-    
     with gr.Row():
         with gr.Column():
             with gr.Tabs() as tabs:
                 with gr.Tab("image-to-video", id="i2v_tab") as image_tab:
-                    # ... (tab content is unchanged) ...
                     video_i_hidden = gr.Textbox(visible=False)
                     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)
@@ -297,25 +273,18 @@ with gr.Blocks(css=css) as demo:
                     frames_to_use = gr.Slider(label="Frames to use from input video", minimum=9, maximum=120, value=9, step=8, info="Must be N*8+1.")
                     v2v_prompt = gr.Textbox(label="Prompt", value="Change the style to cinematic anime", lines=3)
                     v2v_button = gr.Button("Generate Video-to-Video Clip", variant="primary")
-
             duration_input = gr.Slider(label="Clip Duration (seconds)", minimum=1.0, maximum=10.0, value=2.0, step=0.1)
             improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=True)
-
         with gr.Column():
             output_video = gr.Video(label="Last Generated Clip", interactive=False)
             use_last_frame_button = gr.Button("Use Last Frame as Input Image", visible=False)
-
-            # --- NEW ---
-            # Add UI components for stitching
             with gr.Accordion("Stitching Controls", open=True):
                 clip_counter_display = gr.Markdown("Clips created: 0")
                 with gr.Row():
                     stitch_button = gr.Button("🎬 Stitch All Clips")
                     clear_button = gr.Button("🗑️ Clear All Clips")
                 final_video_output = gr.Video(label="Final Stitched Video", interactive=False)
-
     with gr.Accordion("Advanced settings", open=False):
-        # ... (advanced settings are unchanged) ...
         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():
@@ -329,7 +298,6 @@ with gr.Blocks(css=css) as demo:
             num_steps = gr.Slider(label="Steps", value=20, step=1, minimum=1, maximum=420)
             fps = gr.Slider(label="FPS", value=30.0, step=1.0, minimum=4.0, maximum=60.0)
     
-    # ... (event handlers for uploads and tab changes are unchanged) ...
     def handle_image_upload_for_dims(f, h, w):
         if not f: return gr.update(value=h), gr.update(value=w)
         img = Image.open(f)
@@ -342,6 +310,7 @@ with gr.Blocks(css=css) as demo:
             orig_w, orig_h = meta.get('size', (reader.get_data(0).shape[1], reader.get_data(0).shape[0]))
         new_h, new_w = calculate_new_dimensions(orig_w, orig_h)
         return gr.update(value=new_h), gr.update(value=new_w)
+    
     image_i2v.upload(handle_image_upload_for_dims, [image_i2v, height_input, width_input], [height_input, width_input])
     video_v2v.upload(handle_video_upload_for_dims, [video_v2v, height_input, width_input], [height_input, width_input])
     image_tab.select(update_task_image, outputs=[mode])
@@ -349,14 +318,14 @@ with gr.Blocks(css=css) as demo:
     video_tab.select(update_task_video, outputs=[mode])
 
     # --- MODIFIED ---
-    # The inputs and outputs for the generate buttons now include the clips_state and clip_counter_display
-    base_inputs = [negative_prompt_input, clips_state,
-                   height_input, width_input, mode, duration_input, frames_to_use, 
-                   seed_input, randomize_seed_input, guidance_scale_input, improve_texture, num_steps, fps]
-    
-    t2v_inputs = [t2v_prompt] + base_inputs + [image_n_hidden, video_n_hidden]
-    i2v_inputs = [i2v_prompt] + base_inputs + [image_i2v, video_i_hidden]
-    v2v_inputs = [v2v_prompt] + base_inputs + [image_v_hidden, video_v2v]
+    # The input lists are now defined explicitly and in the correct order to match the `generate` function signature.
+    common_params = [
+        height_input, width_input, mode, duration_input, frames_to_use,
+        seed_input, randomize_seed_input, guidance_scale_input, improve_texture, num_steps, fps
+    ]
+    t2v_inputs = [t2v_prompt, negative_prompt_input, clips_state, image_n_hidden, video_n_hidden] + common_params
+    i2v_inputs = [i2v_prompt, negative_prompt_input, clips_state, image_i2v, video_i_hidden] + common_params
+    v2v_inputs = [v2v_prompt, negative_prompt_input, clips_state, image_v_hidden, video_v2v] + common_params
     
     gen_outputs = [output_video, seed_input, use_last_frame_button, clips_state, clip_counter_display]
 
@@ -366,13 +335,9 @@ with gr.Blocks(css=css) as demo:
     v2v_button.click(hide_btn, outputs=[use_last_frame_button], queue=False).then(fn=generate, inputs=v2v_inputs, outputs=gen_outputs, api_name="video_to_video")
 
     use_last_frame_button.click(fn=use_last_frame_as_input, inputs=[output_video], outputs=[image_i2v, tabs])
-
-    # --- NEW ---
-    # Add event handlers for the new stitching and clearing buttons
     stitch_button.click(fn=stitch_videos, inputs=[clips_state], outputs=[final_video_output])
     clear_button.click(fn=clear_clips, outputs=[clips_state, clip_counter_display, output_video, final_video_output])
 
-
 if __name__ == "__main__":
     if os.path.exists(models_dir): print(f"Model directory: {Path(models_dir).resolve()}")
     demo.queue().launch(debug=True, share=False, mcp_server=True)