Create app.py

app.py

@@ -0,0 +1,431 @@
+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
+
+
+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()
+
+
+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_path: str) -> torch.Tensor:
+    """
+    Reads a video file and converts it into a torch.Tensor with the shape [F, C, H, W].
+    """
+    pil_images = load_video(video_path)
+    to_tensor_transform = transforms.ToTensor()
+    video_tensor = torch.stack([to_tensor_transform(img) for img in pil_images])
+    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_tensor):
+    """
+    Process video for canny control.
+    Placeholder function - will return video as-is for now.
+    TODO: Implement canny edge detection processing
+    """
+    print("Processing video for canny control...")
+   
+    return video_tensor
+
+def process_video_for_depth(video_tensor):
+    """
+    Process video for depth control.
+    Placeholder function - will return video as-is for now.
+    TODO: Implement depth estimation processing
+    """
+    print("Processing video for depth control...")
+   
+    return video_tensor
+
+def process_video_for_pose(video_tensor):
+    """
+    Process video for pose control.
+    Placeholder function - will return video as-is for now.
+    TODO: Implement pose estimation processing
+    """
+    print("Processing video for pose control...")
+   
+    return video_tensor
+
+def process_video_for_control(video_tensor, control_type):
+    """Process video based on the selected control type"""
+    if control_type == "canny":
+        return process_video_for_canny(video_tensor)
+    elif control_type == "depth":
+        return process_video_for_depth(video_tensor)
+    elif control_type == "pose":
+        return process_video_for_pose(video_tensor)
+    else:
+        return video_tensor
+        
+@spaces.GPU()
+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)
+        
+        progress(0.1, desc="Loading reference video...")
+        
+        # Read the reference video
+        video = read_video(reference_video)
+        
+        progress(0.15, desc="Processing video for control...")
+        
+        # Process video based on control type
+        processed_video = process_video_for_control(video, control_type)
+        
+        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().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:
+        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
+                    )
+                
+                with gr.Row():
+                    decode_timestep = gr.Slider(
+                        label="Decode Timestep",
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.01,
+                        value=0.05
+                    )
+                    decode_noise_scale = gr.Slider(
+                        label="Decode Noise Scale",
+                        minimum=0.0,
+                        maximum=0.1,
+                        step=0.005,
+                        value=0.025
+                    )
+                
+                image_cond_noise_scale = gr.Slider(
+                    label="Image Condition Noise Scale",
+                    minimum=0.0,
+                    maximum=0.5,
+                    step=0.01,
+                    value=0.0
+                )
+                
+                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()