Update app.py

app.py

@@ -11,8 +11,8 @@ import tempfile
 from PIL import Image
 from huggingface_hub import hf_hub_download
 import shutil
+import math # For math.round, though built-in round works for floats
 
-# --- Import necessary classes from the provided files ---
 from inference import (
     create_ltx_video_pipeline,
     create_latent_upsampler,
@@ -25,38 +25,9 @@ from inference import (
 from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem, LTXMultiScalePipeline, LTXVideoPipeline
 from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
 
-# --- Global constants from user's request and YAML ---
-YAML_CONFIG_STRING = """
-pipeline_type: multi-scale
-checkpoint_path: "ltxv-13b-0.9.7-distilled.safetensors" # This will be replaced by the rc3 version
-downscale_factor: 0.6666666
-spatial_upscaler_model_path: "ltxv-spatial-upscaler-0.9.7.safetensors"
-stg_mode: "attention_values" # options: "attention_values", "attention_skip", "residual", "transformer_block"
-decode_timestep: 0.05
-decode_noise_scale: 0.025
-text_encoder_model_name_or_path: "PixArt-alpha/PixArt-XL-2-1024-MS"
-precision: "bfloat16"
-sampler: "from_checkpoint" # options: "uniform", "linear-quadratic", "from_checkpoint"
-prompt_enhancement_words_threshold: 120
-prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
-prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
-stochastic_sampling: false
-
-first_pass:
-  timesteps: [1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250]
-  guidance_scale: 1
-  stg_scale: 0
-  rescaling_scale: 1
-  skip_block_list: [42]
-
-second_pass:
-  timesteps: [0.9094, 0.7250, 0.4219]
-  guidance_scale: 1
-  stg_scale: 0
-  rescaling_scale: 1
-  skip_block_list: [42]
-"""
-PIPELINE_CONFIG_YAML = yaml.safe_load(YAML_CONFIG_STRING)
+config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
+with open(config_file_path, "r") as file:
+    PIPELINE_CONFIG_YAML = yaml.safe_load(file)
 
 # Model specific paths (to be downloaded)
 DISTILLED_MODEL_REPO = "LTX-Colab/LTX-Video-Preview"
@@ -66,6 +37,7 @@ UPSCALER_REPO = "Lightricks/LTX-Video"
 
 MAX_IMAGE_SIZE = PIPELINE_CONFIG_YAML.get("max_resolution", 1280)
 MAX_NUM_FRAMES = 257
+FPS = 30.0  # Frames per second for duration calculation
 
 # --- Global variables for loaded models ---
 pipeline_instance = None
@@ -117,12 +89,13 @@ 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)
-latent_upsampler_instance.to(target_inference_device)
+if latent_upsampler_instance: # Check if it was created before moving
+    latent_upsampler_instance.to(target_inference_device)
 
 @spaces.GPU
 def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath,
              height_ui, width_ui, mode,
-             ui_steps, num_frames_ui,
+             ui_steps, duration_ui, # << CHANGED from num_frames_ui
              ui_frames_to_use,
              seed_ui, randomize_seed, ui_guidance_scale, improve_texture_flag,
              progress=gr.Progress(track_tqdm=True)):
@@ -131,13 +104,33 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
         seed_ui = random.randint(0, 2**32 - 1)
     seed_everething(int(seed_ui))
     
+    # Convert duration_ui (seconds) to actual_num_frames (N*8+1 format)
+    target_frames_ideal = duration_ui * FPS
+    target_frames_rounded = round(target_frames_ideal)
+    if target_frames_rounded < 1: # ensure positive for calculation
+        target_frames_rounded = 1
+    
+    # Calculate N for N*8+1, ensuring it's rounded to the nearest integer
+    # (target_frames_rounded - 1) could be float if target_frames_rounded is float
+    n_val = round((float(target_frames_rounded) - 1.0) / 8.0)
+    actual_num_frames = int(n_val * 8 + 1)
+
+    # Clamp to the allowed min (9) and max (MAX_NUM_FRAMES) N*8+1 values
+    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)
-    actual_num_frames = int(num_frames_ui)
+    # actual_num_frames is now calculated above
 
     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
+    # This padding ensures the model gets a frame count that is N*8+1
+    # Since actual_num_frames is already N*8+1, this should preserve it.
+    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.")
+        # This case should ideally not happen if actual_num_frames is correctly N*8+1 and >= 9.
     
     padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
 
@@ -146,10 +139,10 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
         "negative_prompt": negative_prompt,
         "height": height_padded,
         "width": width_padded,
-        "num_frames": num_frames_padded,
-        "frame_rate": 30, 
+        "num_frames": num_frames_padded, # Use the padded value for the model
+        "frame_rate": int(FPS), 
         "generator": torch.Generator(device=target_inference_device).manual_seed(int(seed_ui)),
-        "output_type": "pt", # Crucial: pipeline will output [0,1] range tensors
+        "output_type": "pt", 
         "conditioning_items": None,
         "media_items": None,
         "decode_timestep": PIPELINE_CONFIG_YAML["decode_timestep"],
@@ -191,14 +184,18 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
                 media_path=input_video_filepath,
                 height=actual_height, 
                 width=actual_width,
-                max_frames=int(ui_frames_to_use),
+                max_frames=int(ui_frames_to_use), # This is from a separate slider for V2V
                 padding=padding_values
             ).to(target_inference_device)
         except Exception as e:
             print(f"Error loading video {input_video_filepath}: {e}")
             raise gr.Error(f"Could not load video: {e}")
 
-    print(f"Moving models to {target_inference_device} for inference...")
+    print(f"Moving models to {target_inference_device} for inference (if not already there)...")
+    # Models are moved globally once, no need to move per call unless strategy changes.
+    # pipeline_instance.to(target_inference_device)
+    # if latent_upsampler_instance:
+    #    latent_upsampler_instance.to(target_inference_device)
     
     active_latent_upsampler = None
     if improve_texture_flag and latent_upsampler_instance:
@@ -227,17 +224,18 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
             "second_pass": second_pass_args,
         })
         
-        print(f"Calling multi-scale pipeline (eff. HxW: {actual_height}x{actual_width}) on {target_inference_device}")
+        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:
         single_pass_call_kwargs = call_kwargs.copy()
         single_pass_call_kwargs["guidance_scale"] = float(ui_guidance_scale)
         single_pass_call_kwargs["num_inference_steps"] = int(ui_steps)
-        single_pass_call_kwargs.pop("first_pass", None)
+        # These keys might not exist if improve_texture_flag is false from the start of call_kwargs
+        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}) on {target_inference_device}")
+        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
 
     if result_images_tensor is None:
@@ -246,14 +244,14 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
     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
+    
+    # Crop to actual_num_frames, which is the desired output length
     result_images_tensor = result_images_tensor[
         :, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end
     ]
 
-    # The pipeline with output_type="pt" should return tensors in the [0, 1] range.
     video_np = result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy()
     
-    # Clip to ensure values are indeed in [0, 1] before scaling to uint8
     video_np = np.clip(video_np, 0, 1) 
     video_np = (video_np * 255).astype(np.uint8)
 
@@ -278,13 +276,13 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
             raise gr.Error(f"Failed to save video: {e2}")
 
     if isinstance(input_image_filepath, tempfile._TemporaryFileWrapper):
-        if os.path.exists(input_image_filepath.name): # Check if it's already closed by Gradio
+        if os.path.exists(input_image_filepath.name):
              try:
                 input_image_filepath.close()
                 os.remove(input_image_filepath.name)
-             except: pass # May already be closed/removed
+             except: pass 
     elif input_image_filepath and os.path.exists(input_image_filepath) and input_image_filepath.startswith(tempfile.gettempdir()):
-        try: os.remove(input_image_filepath) # If Gradio passed a path to a temp file
+        try: os.remove(input_image_filepath)
         except: pass
 
     if isinstance(input_video_filepath, tempfile._TemporaryFileWrapper):
@@ -312,29 +310,35 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown("Generates a short video based on text prompt, image, or existing video. Models are moved to GPU during generation and back to CPU afterwards to save VRAM.")
     with gr.Row():
         with gr.Column():
-            with gr.Group():
-                with gr.Tab("image-to-video") as image_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"])
-                    i2v_prompt = gr.Textbox(label="Prompt", value="The creature from the image starts to move", lines=3)
-                    i2v_button = gr.Button("Generate Image-to-Video", variant="primary")
-                with gr.Tab("text-to-video") as text_tab:
-                    image_n_hidden = gr.Textbox(label="image_n", visible=False, value=None)
-                    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")
-                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"])
-                    frames_to_use = gr.Slider(label="Frames to use from input video", minimum=9, maximum=MAX_NUM_FRAMES, value=9, step=8, info="Number of initial frames to use for conditioning/transformation. 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", variant="primary")
-
+            with gr.Tab("image-to-video") as image_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"])
+                i2v_prompt = gr.Textbox(label="Prompt", value="The creature from the image starts to move", lines=3)
+                i2v_button = gr.Button("Generate Image-to-Video", variant="primary")
+            with gr.Tab("text-to-video") as text_tab:
+                image_n_hidden = gr.Textbox(label="image_n", visible=False, value=None)
+                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")
+            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"])
+                frames_to_use = gr.Slider(label="Frames to use from input video", minimum=9, maximum=MAX_NUM_FRAMES, value=9, step=8, info="Number of initial frames to use for conditioning/transformation. 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", variant="primary")
+
+            duration_input = gr.Slider(
+                label="Duration (seconds)", 
+                minimum=0.3, 
+                maximum=8.5, 
+                value=2,  
+                step=0.1, 
+                info=f"Target video duration (0.3s to 8.5s)"
+            )
             improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=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)
-            gr.Markdown("Note: Generation can take a few minutes depending on settings and hardware.")
 
     with gr.Accordion("Advanced settings", open=False):
         negative_prompt_input = gr.Textbox(label="Negative Prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted", lines=2)
@@ -343,27 +347,26 @@ with gr.Blocks(css=css) as demo:
             randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=False)
         with gr.Row():
             guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=1.0, maximum=10.0, value=PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0), step=0.1, info="Controls how much the prompt influences the output. Higher values = stronger influence.")
-            default_steps = len(PIPELINE_CONFIG_YAML.get("first_pass", {}).get("timesteps", [1]*7))
+            default_steps = len(PIPELINE_CONFIG_YAML.get("first_pass", {}).get("timesteps", [1]*7)) # Default to 7 if not found
             steps_input = gr.Slider(label="Inference Steps (for first pass if multi-scale)", minimum=1, maximum=30, value=default_steps, step=1, info="Number of denoising steps. More steps can improve quality but increase time. If YAML defines 'timesteps' for a pass, this UI value is ignored for that pass.")
-        with gr.Row():
-            num_frames_input = gr.Slider(label="Number of Frames to Generate", minimum=9, maximum=MAX_NUM_FRAMES, value=25, step=8, info="Total frames in the output video. Should be N*8+1 (e.g., 9, 17, 25...).")
         with gr.Row():
             height_input = gr.Slider(label="Height", value=512, step=32, minimum=256, maximum=MAX_IMAGE_SIZE, info="Must be divisible by 32.")
             width_input = gr.Slider(label="Width", value=704, step=32, minimum=256, maximum=MAX_IMAGE_SIZE, info="Must be divisible by 32.")
     
+    # --- UPDATED INPUT LISTS ---
     t2v_inputs = [t2v_prompt, negative_prompt_input, image_n_hidden, video_n_hidden,
                   height_input, width_input, gr.State("text-to-video"),
-                  steps_input, num_frames_input, gr.State(0),
+                  steps_input, duration_input, gr.State(0), # Replaced num_frames_input with duration_input
                   seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
     
     i2v_inputs = [i2v_prompt, negative_prompt_input, image_i2v, video_i_hidden,
                   height_input, width_input, gr.State("image-to-video"),
-                  steps_input, num_frames_input, gr.State(0),
+                  steps_input, duration_input, gr.State(0), # Replaced num_frames_input with duration_input
                   seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
 
     v2v_inputs = [v2v_prompt, negative_prompt_input, image_v_hidden, video_v2v,
                   height_input, width_input, gr.State("video-to-video"),
-                  steps_input, num_frames_input, frames_to_use,
+                  steps_input, duration_input, frames_to_use, # Replaced num_frames_input with duration_input
                   seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
 
     t2v_button.click(fn=generate, inputs=t2v_inputs, outputs=[output_video], api_name="text_to_video")