Update app.py

app.py

@@ -1,10 +1,9 @@
-import gradio as gr
 import torch
-import spaces
 import numpy as np
 import random
 import os
 import yaml
+import argparse
 from pathlib import Path
 import imageio
 import tempfile
@@ -87,10 +86,9 @@ pipeline_instance.to(target_inference_device)
 if latent_upsampler_instance: 
     latent_upsampler_instance.to(target_inference_device)
 
-
 # --- Helper function for dimension calculation ---
-MIN_DIM_SLIDER = 256  # As defined in the sliders minimum attribute
-TARGET_FIXED_SIDE = 768 # Desired fixed side length as per requirement
+MIN_DIM_SLIDER = 256
+TARGET_FIXED_SIDE = 768
 
 def calculate_new_dimensions(orig_w, orig_h):
     """
@@ -99,59 +97,38 @@ def calculate_new_dimensions(orig_w, orig_h):
     both are multiples of 32, and within [MIN_DIM_SLIDER, MAX_IMAGE_SIZE].
     """
     if orig_w == 0 or orig_h == 0:
-        # Default to TARGET_FIXED_SIDE square if original dimensions are invalid
         return int(TARGET_FIXED_SIDE), int(TARGET_FIXED_SIDE)
 
-    if orig_w >= orig_h:  # Landscape or square
+    if 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)) 
-    else:  # Portrait
+    else:
         new_w = TARGET_FIXED_SIDE
-        aspect_ratio = orig_h / orig_w # Use H/W ratio for portrait scaling
+        aspect_ratio = orig_h / orig_w
         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
         new_w = max(MIN_DIM_SLIDER, min(new_w, MAX_IMAGE_SIZE))
 
     return int(new_h), int(new_w)
 
-def get_duration(prompt, negative_prompt, input_image_filepath, input_video_filepath,
-             height_ui, width_ui, mode,
-             duration_ui, # Removed ui_steps
-             ui_frames_to_use,
-             seed_ui, randomize_seed, ui_guidance_scale, improve_texture_flag,
-             progress):
-    if duration_ui > 7:
-        return 75
-    else:
-        return 60
-
-@spaces.GPU(duration=get_duration)
-def generate(prompt, negative_prompt, 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,
-             progress=gr.Progress(track_tqdm=True)):
+def generate(prompt, negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted", 
+             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=None, improve_texture_flag=True):
 
     if randomize_seed:
         seed_ui = random.randint(0, 2**32 - 1)
     seed_everething(int(seed_ui))
     
+    if ui_guidance_scale is None:
+        ui_guidance_scale = PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0)
+    
     target_frames_ideal = duration_ui * FPS
     target_frames_rounded = round(target_frames_ideal)
     if target_frames_rounded < 1: 
@@ -168,9 +145,7 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
 
     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 
-    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.")
+    num_frames_padded = ((actual_num_frames - 2) // 8 + 1) * 8 + 1
     
     padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
 
@@ -217,7 +192,7 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
             call_kwargs["conditioning_items"] = [ConditioningItem(media_tensor.to(target_inference_device), 0, 1.0)]
         except Exception as e:
             print(f"Error loading image {input_image_filepath}: {e}")
-            raise gr.Error(f"Could not load image: {e}")
+            raise RuntimeError(f"Could not load image: {e}")
     elif mode == "video-to-video" and input_video_filepath:
         try:
             call_kwargs["media_items"] = load_media_file(
@@ -229,7 +204,7 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
             ).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}")
+            raise RuntimeError(f"Could not load video: {e}")
 
     print(f"Moving models to {target_inference_device} for inference (if not already there)...")
     
@@ -240,19 +215,16 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
     result_images_tensor = None
     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.")
+            raise RuntimeError("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()
@@ -269,12 +241,11 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
         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("second_pass", None)
@@ -284,7 +255,7 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
         result_images_tensor = pipeline_instance(**single_pass_call_kwargs).images
 
     if result_images_tensor is None:
-        raise gr.Error("Generation failed.")
+        raise RuntimeError("Generation failed.")
 
     pad_left, pad_right, pad_top, pad_bottom = padding_values
     slice_h_end = -pad_bottom if pad_bottom > 0 else None
@@ -295,192 +266,98 @@ def generate(prompt, negative_prompt, input_image_filepath, input_video_filepath
     ]
 
     video_np = result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy()
-    
     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")
+    timestamp = random.randint(10000, 99999)
+    output_video_path = 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]):
-                progress(frame_idx / video_np.shape[0], desc="Saving video")
                 video_writer.append_data(video_np[frame_idx])
+                if frame_idx % 10 == 0:
+                    print(f"Saving frame {frame_idx + 1}/{video_np.shape[0]}")
     except Exception as e:
         print(f"Error saving video with macro_block_size=1: {e}")
         try:
             with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], format='FFMPEG', codec='libx264', quality=8) as video_writer:
-                 for frame_idx in range(video_np.shape[0]):
-                    progress(frame_idx / video_np.shape[0], desc="Saving video (fallback ffmpeg)")
+                for frame_idx in range(video_np.shape[0]):
                     video_writer.append_data(video_np[frame_idx])
+                    if frame_idx % 10 == 0:
+                        print(f"Saving frame {frame_idx + 1}/{video_np.shape[0]} (fallback)")
         except Exception as e2:
             print(f"Fallback video saving error: {e2}")
-            raise gr.Error(f"Failed to save video: {e2}")
+            raise RuntimeError(f"Failed to save video: {e2}")
             
     return output_video_path, seed_ui
 
-def update_task_image():
-    return "image-to-video"
-
-def update_task_text():
-    return "text-to-video"
-
-def update_task_video():
-    return "video-to-video"
-
-# --- Gradio UI Definition ---
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 900px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("# LTX Video 0.9.7 Distilled")
-    gr.Markdown("Fast high quality video generation. [Model](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-distilled.safetensors) [GitHub](https://github.com/Lightricks/LTX-Video) [Diffusers](#)")
+def main():
+    parser = argparse.ArgumentParser(description="LTX Video Generation from Command Line")
+    parser.add_argument("--prompt", required=True, help="Text prompt for video generation")
+    parser.add_argument("--negative-prompt", default="worst quality, inconsistent motion, blurry, jittery, distorted", 
+                        help="Negative prompt")
+    parser.add_argument("--mode", choices=["text-to-video", "image-to-video", "video-to-video"], 
+                        default="text-to-video", help="Generation mode")
+    parser.add_argument("--input-image", help="Input image path for image-to-video mode")
+    parser.add_argument("--input-video", help="Input video path for video-to-video mode")
+    parser.add_argument("--duration", type=float, default=2.0, help="Video duration in seconds (0.3-8.5)")
+    parser.add_argument("--height", type=int, default=512, help="Video height (must be divisible by 32)")
+    parser.add_argument("--width", type=int, default=704, help="Video width (must be divisible by 32)")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--randomize-seed", action="store_true", help="Use random seed")
+    parser.add_argument("--guidance-scale", type=float, help="Guidance scale for generation")
+    parser.add_argument("--no-improve-texture", action="store_true", help="Disable texture improvement (faster)")
+    parser.add_argument("--frames-to-use", type=int, default=9, help="Frames to use from input video (for video-to-video)")
     
-    with gr.Row():
-        with gr.Column():
-            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", "clipboard"])
-                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", visible=False) 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"]) # type defaults to filepath
-                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="Video 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.DeepLinkButton()
-
-    with gr.Accordion("Advanced settings", open=False):
-        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():
-            seed_input = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=2**32-1)
-            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
-        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.")
-        with gr.Row():
-            height_input = gr.Slider(label="Height", value=512, step=32, minimum=MIN_DIM_SLIDER, maximum=MAX_IMAGE_SIZE, info="Must be divisible by 32.")
-            width_input = gr.Slider(label="Width", value=704, step=32, minimum=MIN_DIM_SLIDER, maximum=MAX_IMAGE_SIZE, info="Must be divisible by 32.")
-
-
-    # --- Event handlers for updating dimensions on upload ---
-    def handle_image_upload_for_dims(image_filepath, current_h, current_w):
-        if not image_filepath:  # Image cleared or no image initially
-            # Keep current slider values if image is cleared or no input
-            return gr.update(value=current_h), gr.update(value=current_w)
-        try:
-            img = Image.open(image_filepath)
-            orig_w, orig_h = img.size
-            new_h, new_w = calculate_new_dimensions(orig_w, orig_h)
-            return gr.update(value=new_h), gr.update(value=new_w)
-        except Exception as e:
-            print(f"Error processing image for dimension update: {e}")
-            # Keep current slider values on error
-            return gr.update(value=current_h), gr.update(value=current_w)
-
-    def handle_video_upload_for_dims(video_filepath, current_h, current_w):
-        if not video_filepath:  # Video cleared or no video initially
-            return gr.update(value=current_h), gr.update(value=current_w)
-        try:
-            # Ensure video_filepath is a string for os.path.exists and imageio
-            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()
-                if 'size' in meta:
-                    orig_w, orig_h = meta['size']
-                else:
-                    # Fallback: read first frame if 'size' not in metadata
-                    try:
-                        first_frame = reader.get_data(0)
-                        # Shape is (h, w, c) for frames
-                        orig_h, orig_w = first_frame.shape[0], first_frame.shape[1]
-                    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: # If dimensions couldn't be determined
-                 print(f"Could not determine dimensions for video: {video_filepath_str}")
-                 return gr.update(value=current_h), gr.update(value=current_w)
-
-            new_h, new_w = calculate_new_dimensions(orig_w, orig_h)
-            return gr.update(value=new_h), gr.update(value=new_w)
-        except Exception as e:
-            # Log type of video_filepath for debugging if it's not a path-like string
-            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)
-
+    args = parser.parse_args()
     
-    image_i2v.upload(
-        fn=handle_image_upload_for_dims,
-        inputs=[image_i2v, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
-    video_v2v.upload(
-        fn=handle_video_upload_for_dims,
-        inputs=[video_v2v, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
-
-    image_tab.select(
-        fn=update_task_image,
-        outputs=[mode]
-    )
-    text_tab.select(
-        fn=update_task_text,
-        outputs=[mode]
-    )
+    # Validate parameters
+    if args.mode == "image-to-video" and not args.input_image:
+        print("Error: --input-image is required for image-to-video mode")
+        return
     
-    t2v_inputs = [t2v_prompt, negative_prompt_input, image_n_hidden, video_n_hidden,
-                  height_input, width_input, mode,
-                  duration_input, frames_to_use, 
-                  seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
+    if args.mode == "video-to-video" and not args.input_video:
+        print("Error: --input-video is required for video-to-video mode")
+        return
     
-    i2v_inputs = [i2v_prompt, negative_prompt_input, image_i2v, video_i_hidden,
-                  height_input, width_input, mode,
-                  duration_input, frames_to_use, 
-                  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, mode,
-                  duration_input, frames_to_use, 
-                  seed_input, randomize_seed_input, guidance_scale_input, improve_texture]
-
-    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")
+    # Ensure dimensions are divisible by 32
+    args.height = ((args.height - 1) // 32 + 1) * 32
+    args.width = ((args.width - 1) // 32 + 1) * 32
+    
+    print(f"Starting video generation...")
+    print(f"Prompt: {args.prompt}")
+    print(f"Mode: {args.mode}")
+    print(f"Duration: {args.duration}s")
+    print(f"Resolution: {args.width}x{args.height}")
+    
+    try:
+        output_path, used_seed = generate(
+            prompt=args.prompt,
+            negative_prompt=args.negative_prompt,
+            input_image_filepath=args.input_image,
+            input_video_filepath=args.input_video,
+            height_ui=args.height,
+            width_ui=args.width,
+            mode=args.mode,
+            duration_ui=args.duration,
+            ui_frames_to_use=args.frames_to_use,
+            seed_ui=args.seed,
+            randomize_seed=args.randomize_seed,
+            ui_guidance_scale=args.guidance_scale,
+            improve_texture_flag=not args.no_improve_texture
+        )
+        
+        print(f"\nVideo generation completed!")
+        print(f"Output saved to: {output_path}")
+        print(f"Used seed: {used_seed}")
+        
+    except Exception as e:
+        print(f"Error during generation: {e}")
+        raise
 
 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)
+    main()