Update app.py

app.py

@@ -151,17 +151,20 @@ def detect_aspect_ratio(image) -> str:
     return min(candidates, key=lambda k: abs(ratio - candidates[k]))
 
 
-def on_image_upload(image, high_res):
-    """Auto-set resolution when image is uploaded."""
-    aspect = detect_aspect_ratio(image)
+def on_image_upload(first_image, last_image, high_res):
+    """Auto-set resolution when an image is uploaded."""
+    # Use first image for aspect ratio detection, fall back to last image
+    ref_image = first_image if first_image is not None else last_image
+    aspect = detect_aspect_ratio(ref_image)
     tier = "high" if high_res else "low"
     w, h = RESOLUTIONS[tier][aspect]
     return gr.update(value=w), gr.update(value=h)
 
 
-def on_highres_toggle(image, high_res):
+def on_highres_toggle(first_image, last_image, high_res):
     """Update resolution when high-res toggle changes."""
-    aspect = detect_aspect_ratio(image)
+    ref_image = first_image if first_image is not None else last_image
+    aspect = detect_aspect_ratio(ref_image)
     tier = "high" if high_res else "low"
     w, h = RESOLUTIONS[tier][aspect]
     return gr.update(value=w), gr.update(value=h)
@@ -170,7 +173,8 @@ def on_highres_toggle(image, high_res):
 @spaces.GPU(duration=75)
 @torch.inference_mode()
 def generate_video(
-    input_image,
+    first_image,
+    last_image,
     prompt: str,
     duration: float,
     enhance_prompt: bool = True,
@@ -193,15 +197,24 @@ def generate_video(
         print(f"Generating: {height}x{width}, {num_frames} frames ({duration}s), seed={current_seed}")
 
         images = []
-        if input_image is not None:
-            output_dir = Path("outputs")
-            output_dir.mkdir(exist_ok=True)
-            temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"
-            if hasattr(input_image, "save"):
-                input_image.save(temp_image_path)
+        output_dir = Path("outputs")
+        output_dir.mkdir(exist_ok=True)
+
+        if first_image is not None:
+            temp_first_path = output_dir / f"temp_first_{current_seed}.jpg"
+            if hasattr(first_image, "save"):
+                first_image.save(temp_first_path)
+            else:
+                temp_first_path = Path(first_image)
+            images.append(ImageConditioningInput(path=str(temp_first_path), frame_idx=0, strength=1.0))
+
+        if last_image is not None:
+            temp_last_path = output_dir / f"temp_last_{current_seed}.jpg"
+            if hasattr(last_image, "save"):
+                last_image.save(temp_last_path)
             else:
-                temp_image_path = Path(input_image)
-            images = [ImageConditioningInput(path=str(temp_image_path), frame_idx=0, strength=1.0)]
+                temp_last_path = Path(last_image)
+            images.append(ImageConditioningInput(path=str(temp_last_path), frame_idx=num_frames - 1, strength=1.0))
 
         tiling_config = TilingConfig.default()
         video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
@@ -251,7 +264,9 @@ with gr.Blocks(title="LTX-2.3 Distilled") as demo:
 
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(label="Input Image (Optional)", type="pil")
+            with gr.Row():
+                first_image = gr.Image(label="First Frame (Optional)", type="pil")
+                last_image = gr.Image(label="Last Frame (Optional)", type="pil")
             prompt = gr.Textbox(
                 label="Prompt",
                 info="for best results - make it as elaborate as possible",
@@ -279,23 +294,29 @@ with gr.Blocks(title="LTX-2.3 Distilled") as demo:
             output_video = gr.Video(label="Generated Video", autoplay=True)
 
     # Auto-detect aspect ratio from uploaded image and set resolution
-    input_image.change(
+    first_image.change(
         fn=on_image_upload,
-        inputs=[input_image, high_res],
+        inputs=[first_image, last_image, high_res],
+        outputs=[width, height],
+    )
+
+    last_image.change(
+        fn=on_image_upload,
+        inputs=[first_image, last_image, high_res],
         outputs=[width, height],
     )
 
     # Update resolution when high-res toggle changes
     high_res.change(
         fn=on_highres_toggle,
-        inputs=[input_image, high_res],
+        inputs=[first_image, last_image, high_res],
         outputs=[width, height],
     )
 
     generate_btn.click(
         fn=generate_video,
         inputs=[
-            input_image, prompt, duration, enhance_prompt,
+            first_image, last_image, prompt, duration, enhance_prompt,
             seed, randomize_seed, height, width,
         ],
         outputs=[output_video, seed],
@@ -307,5 +328,4 @@ css = """
 """
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Citrus(), css=css)
-
+    demo.launch(theme=gr.themes.Citrus(), css=css)