app.py

import gradio as gr
import torch
import numpy as np
import tempfile
import os
import spaces
from diffusers import LTXLatentUpsamplePipeline
from pipeline_ltx_condition_control import LTXConditionPipeline
from diffusers.utils import export_to_video, load_video
from torchvision import transforms
import random
from controlnet_aux import CannyDetector
from image_gen_aux import DepthPreprocessor
import mediapipe as mp
from PIL import Image
import cv2

dtype = torch.bfloat16 
device = "cuda" if torch.cuda.is_available() else "cpu"

pipeline = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.7-dev", torch_dtype=dtype)
pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained("Lightricks/ltxv-spatial-upscaler-0.9.7", vae=pipeline.vae, torch_dtype=dtype)
pipeline.to(device)
pipe_upsample.to(device)
pipeline.vae.enable_tiling()

canny_processor = CannyDetector()
depth_processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")

# Initialize MediaPipe pose estimation
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles
mp_pose = mp.solutions.pose

CONTROL_LORAS = {
    "canny": {
        "repo": "Lightricks/LTX-Video-ICLoRA-canny-13b-0.9.7",
        "weight_name": "ltxv-097-ic-lora-canny-control-diffusers.safetensors",
        "adapter_name": "canny_lora"
    },
    "depth": {
        "repo": "Lightricks/LTX-Video-ICLoRA-depth-13b-0.9.7", 
        "weight_name": "ltxv-097-ic-lora-depth-control-diffusers.safetensors",
        "adapter_name": "depth_lora"
    },
    "pose": {
        "repo": "Lightricks/LTX-Video-ICLoRA-pose-13b-0.9.7",
        "weight_name": "ltxv-097-ic-lora-pose-control-diffusers.safetensors", 
        "adapter_name": "pose_lora"
    }
}

@spaces.GPU()
def read_video(video) -> torch.Tensor:
    """
    Reads a video file and converts it into a torch.Tensor with the shape [F, C, H, W].
    """
    
    to_tensor_transform = transforms.ToTensor()
    video_tensor = torch.stack([to_tensor_transform(img) for img in video])
    return video_tensor

def round_to_nearest_resolution_acceptable_by_vae(height, width, vae_temporal_compression_ratio):
    height = height - (height % vae_temporal_compression_ratio)
    width = width - (width % vae_temporal_compression_ratio)
    return height, width

@spaces.GPU()
def load_control_lora(control_type, current_lora_state):
    """Load the specified control LoRA, unloading any previous one"""
    
    if control_type not in CONTROL_LORAS:
        raise ValueError(f"Unknown control type: {control_type}")
    
    # If same LoRA is already loaded, do nothing
    if current_lora_state == control_type:
        print(f"{control_type} LoRA already loaded")
        return current_lora_state
    
    # Unload current LoRA if any
    if current_lora_state is not None:
        try:
            pipeline.unload_lora_weights()
            print(f"Unloaded previous LoRA: {current_lora_state}")
        except Exception as e:
            print(f"Warning: Could not unload previous LoRA: {e}")
    
    # Load new LoRA
    lora_config = CONTROL_LORAS[control_type]
    try:
        pipeline.load_lora_weights(
            lora_config["repo"],
            weight_name=lora_config["weight_name"],
            adapter_name=lora_config["adapter_name"]
        )
        pipeline.set_adapters([lora_config["adapter_name"]], adapter_weights=[1.0])
        new_lora_state = control_type
        print(f"Loaded {control_type} LoRA successfully")
        return new_lora_state
    except Exception as e:
        print(f"Error loading {control_type} LoRA: {e}")
        raise

def process_video_for_canny(video):
    """
    Process video for canny control.
    """
    print("Processing video for canny control...")
    canny_video = []
    for frame in video: 
        # TODO: change resolution logic
        canny_video.append(canny_processor(frame, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024))
     
    return canny_video

@spaces.GPU()
def process_video_for_depth(video):
    """
    Process video for depth control.
    """
    print("Processing video for depth control...")
    depth_video = []
    for frame in video: 
        depth_video.append(depth_processor(frame)[0].convert("RGB"))
    return depth_video

@spaces.GPU()
def process_video_for_pose(video):
    """
    Process video for pose control using MediaPipe pose estimation.
    Returns video frames with pose landmarks drawn on black background.
    """
    print("Processing video for pose control...")
    pose_video = []
    
    with mp_pose.Pose(
        static_image_mode=True,
        model_complexity=1,
        enable_segmentation=False,
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5
    ) as pose:
        
        for frame in video:
            # Convert PIL image to numpy array
            frame_np = np.array(frame)
            
            # Convert RGB to BGR for MediaPipe
            frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)
            
            # Process the frame
            results = pose.process(frame_bgr)
            
            # Create black background with same dimensions
            pose_frame = np.zeros_like(frame_np)
            
            # Draw pose landmarks if detected
            if results.pose_landmarks:
                mp_drawing.draw_landmarks(
                    pose_frame,
                    results.pose_landmarks,
                    mp_pose.POSE_CONNECTIONS,
                    landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style(),
                    connection_drawing_spec=mp_drawing_styles.get_default_pose_connections_style()
                )
            
            # Convert back to PIL Image
            pose_pil = Image.fromarray(pose_frame)
            pose_video.append(pose_pil)
    
    return pose_video

def process_video_for_control(video, control_type):
    """Process video based on the selected control type"""
    if control_type == "canny":
        return process_video_for_canny(video)
    elif control_type == "depth":
        return process_video_for_depth(video)
    elif control_type == "pose":
        return process_video_for_pose(video)
    else:
        return video
        
@spaces.GPU(duration=160)
def generate_video(
    reference_video,
    prompt,
    control_type,
    current_lora_state,
    duration=3.0,
    negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted",
    height=768,
    width=1152,
    num_inference_steps=30,
    guidance_scale=5.0,
    guidance_rescale=0.7,
    decode_timestep=0.05,
    decode_noise_scale=0.025,
    image_cond_noise_scale=0.0,
    seed=0,
    randomize_seed=False,
    progress=gr.Progress()
):
    try:
        # Initialize models if needed
        # Models are already loaded at startup
        
        if reference_video is None:
            return None, "Please upload a reference video."
        
        if not prompt.strip():
            return None, "Please enter a prompt."
        
        # Handle seed
        if randomize_seed:
            seed = random.randint(0, 2**32 - 1)
        
        progress(0.05, desc="Loading control LoRA...")
        
        # Load the appropriate control LoRA and update state
        updated_lora_state = load_control_lora(control_type, current_lora_state)
        
        # Loads video into a list of pil images 
        video = load_video(reference_video)
        progress(0.1, desc="Processing video for control...")
        
        # Process video based on control type
        processed_video = process_video_for_control(video, control_type)
        processed_video = read_video(processed_video) # turns to tensor
        
        progress(0.2, desc="Preparing generation parameters...")
        
        # Calculate number of frames from duration (24 fps)
        fps = 24
        num_frames = int(duration * fps) + 1  # +1 for proper frame count
        # Ensure num_frames is valid for the model (multiple of temporal compression + 1)
        temporal_compression = pipeline.vae_temporal_compression_ratio
        num_frames = ((num_frames - 1) // temporal_compression) * temporal_compression + 1
        
        # Calculate downscaled dimensions
        downscale_factor = 2 / 3
        downscaled_height = int(height * downscale_factor)
        downscaled_width = int(width * downscale_factor)
        downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(
            downscaled_height, downscaled_width, pipeline.vae_temporal_compression_ratio
        )
        
        progress(0.3, desc="Generating video at lower resolution...")
        
        # 1. Generate video at smaller resolution
        latents = pipeline(
            reference_video=processed_video,  # Use processed video
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=downscaled_width,
            height=downscaled_height,
            num_frames=num_frames,
            num_inference_steps=num_inference_steps,
            decode_timestep=decode_timestep,
            decode_noise_scale=decode_noise_scale,
            image_cond_noise_scale=image_cond_noise_scale,
            guidance_scale=guidance_scale,
            guidance_rescale=guidance_rescale,
            generator=torch.Generator().manual_seed(seed),
            output_type="latent",
        ).frames

        progress(0.6, desc="Upscaling video...")
        
        # 2. Upscale generated video
        upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
        upscaled_latents = pipe_upsample(
            latents=latents,
            output_type="latent"
        ).frames

        progress(0.8, desc="Final denoising and processing...")
        
        # 3. Denoise the upscaled video
        video_output = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=upscaled_width,
            height=upscaled_height,
            num_frames=num_frames,
            denoise_strength=0.4,
            num_inference_steps=10,
            latents=upscaled_latents,
            decode_timestep=decode_timestep,
            decode_noise_scale=decode_noise_scale,
            image_cond_noise_scale=image_cond_noise_scale,
            guidance_scale=guidance_scale,
            guidance_rescale=guidance_rescale,
            generator=torch.Generator(device="cuda").manual_seed(seed),
            output_type="pil",
        ).frames[0]

        progress(0.9, desc="Finalizing output...")
        
        # 4. Downscale to expected resolution
        video_output = [frame.resize((width, height)) for frame in video_output]

        # Export to temporary file
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_file:
            output_path = tmp_file.name
            export_to_video(video_output, output_path, fps=fps)
        
        progress(1.0, desc="Complete!")
        
        return output_path, updated_lora_state
        
    except Exception as e:
        print(e)
        return None, current_lora_state

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # LTX Video Control
        """
    )
    
    # State variable for tracking current LoRA
    current_lora_state = gr.State(value=None)
    
    with gr.Row():
        with gr.Column(scale=1):
    
            reference_video = gr.Video(
                label="Reference Video",
                height=300
            )
            
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the video you want to generate...",
                lines=3,
                value="A graceful pink swan gliding smoothly across a serene lake, its elegant neck curved as it moves through the calm water. The swan's soft pink feathers shimmer in the gentle sunlight, creating ripples that spread outward in concentric circles. The lake is surrounded by lush green trees reflected in the still water. Shot from a side angle, the camera slowly follows the swan's peaceful movement across the frame. Cinematic lighting, 4K quality, smooth motion."
            )
            
            # Control Type Selection
            control_type = gr.Radio(
                label="Control Type",
                choices=["canny", "depth", "pose"],
                value="canny",
                info="Choose the type of control guidance for video generation"
            )
            
            duration = gr.Slider(
                label="Duration (seconds)",
                minimum=1.0,
                maximum=10.0,
                step=0.5,
                value=3.0
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="What you don't want in the video...",
                lines=2,
                value="worst quality, inconsistent motion, blurry, jittery, distorted"
            )
            
            # Advanced Settings
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=1024,
                        step=32,
                        value=768
                    )
                    width = gr.Slider(
                        label="Width", 
                        minimum=256,
                        maximum=1536,
                        step=32,
                        value=1152
                    )
                
                num_inference_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=50,
                    step=1,
                    value=30
                )
                
                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        minimum=1.0,
                        maximum=15.0,
                        step=0.1,
                        value=5.0
                    )
                    guidance_rescale = gr.Slider(
                        label="Guidance Rescale",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.7,
                        visible=False
                    )
                
                with gr.Row():
                    decode_timestep = gr.Slider(
                        label="Decode Timestep",
                        minimum=0.0,
                        maximum=1.0,
                        step=0.01,
                        value=0.05,
                        visible=False
                    )
                    decode_noise_scale = gr.Slider(
                        label="Decode Noise Scale",
                        minimum=0.0,
                        maximum=0.1,
                        step=0.005,
                        value=0.025,
                        visible=False
                    )
                
                image_cond_noise_scale = gr.Slider(
                    label="Image Condition Noise Scale",
                    minimum=0.0,
                    maximum=0.5,
                    step=0.01,
                    value=0.0,
                    visible=False
                )
                
                with gr.Row():
                    randomize_seed = gr.Checkbox(
                        label="Randomize Seed",
                        value=False
                    )
                    seed = gr.Number(
                        label="Seed",
                        value=0,
                        precision=0
                    )
            
            generate_btn = gr.Button(
                "Generate",
            )
        
        with gr.Column(scale=1):          
            output_video = gr.Video(
                label="Generated Video",
                height=400
            )
        
            
    
    # Event handlers
    generate_btn.click(
        fn=generate_video,
        inputs=[
            reference_video,
            prompt,
            control_type,
            current_lora_state,
            duration,
            negative_prompt,
            height,
            width,
            num_inference_steps,
            guidance_scale,
            guidance_rescale,
            decode_timestep,
            decode_noise_scale,
            image_cond_noise_scale,
            seed,
            randomize_seed
        ],
        outputs=[output_video, current_lora_state],
        show_progress=True
    )

if __name__ == "__main__":
    demo.launch()