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) pipeline = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.7-distilled", torch_dtype=torch.bfloat16) 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" } } # load canny lora pipeline.load_lora_weights( CONTROL_LORAS["canny"]["repo"], weight_name=CONTROL_LORAS["canny"]["weight_name"], adapter_name=CONTROL_LORAS["canny"]["adapter_name"] ) pipeline.set_adapters([CONTROL_LORAS["canny"]["adapter_name"]], adapter_weights=[1.0]) # Initialize MediaPipe pose estimation # mp_drawing = mp.solutions.drawing_utils # mp_drawing_styles = mp.solutions.drawing_styles # mp_pose = mp.solutions.pose canny_processor = CannyDetector() @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_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() ) # Convert back to PIL Image pose_pil = Image.fromarray(pose_frame) pose_video.append(pose_pil) return pose_video def process_video_for_control(reference_video, control_type): video = load_video(reference_video) """Process video based on the selected control type""" if control_type == "canny": processed_video = process_video_for_canny(video) elif control_type == "depth": processed_video = process_video_for_depth(video) elif control_type == "pose": processed_video = process_video_for_pose(video) else: processed_video = reference_video # fps = 24 # with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp2_file: # output2_path = tmp2_file.name # export_to_video(processed_video, output2_path, fps=fps) # return output2_path return processed_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=7, guidance_scale=1.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) # 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 # 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(reference_video, control_type) processed_video = read_video(processed_video) # turns to tensor progress(0.2, desc="Preparing generation parameters...") # 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=0.05, decode_noise_scale=0.025, # 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 = 0.05, guidance_scale=guidance_scale, decode_noise_scale = 0.025, image_cond_noise_scale=0.025, #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, seed except Exception as e: print(e) return None, seed # Create Gradio interface with gr.Blocks() as demo: gr.Markdown( """ # LTX Video Control Canny """ ) # 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="The Joker in his iconic purple suit and green hair, dancing alone in a dimly lit, run-down room. His movements are erratic and unpredictable, shifting between graceful and chaotic as he loses himself in the moment. The camera captures his theatrical gestures, his dance reflecting his unhinged personality. Moody lighting with shadows dancing across the walls, creating an atmosphere of beautiful madness." ) # Control Type Selection control_type = gr.Radio( label="Control Type", choices=["canny", "depth", "pose"], value="canny", visible=False, 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=2.5 ) 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=7 ) with gr.Row(): guidance_scale = gr.Slider( label="Guidance Scale", minimum=1.0, maximum=15.0, step=0.1, value=1.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 ) control_video = gr.Video( label="Control Video", height=400, visible=False ) gr.Examples( examples=[ ["video_assets/vid_1.mp4", "", "canny", 2.5, "", 768, 1152, 7, 1, 0, True], ["video_assets/vid_2.mp4", "", "canny", 2.5, "", 768, 1152, 7, 1, 0, True], ["video_assets/vid_3.mp4", "", "canny", 2.5, "", 768, 1152, 7, 1, 0, True], ["video_assets/vid_4.mp4", "", "canny", 2.5, "", 768, 1152, 7, 1, 0, True], ], inputs=[reference_video, prompt, control_type, # current_lora_state, duration, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed], outputs=[output_video, seed], fn=generate_video, cache_examples="lazy" ) # 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, seed, randomize_seed ], outputs=[output_video, seed], show_progress=True ) if __name__ == "__main__": demo.launch()