v5 release with triple architecture support and prompt enhancer

ltx_video/utils/diffusers_config_mapping.py

@@ -0,0 +1,174 @@
+def make_hashable_key(dict_key):
+    def convert_value(value):
+        if isinstance(value, list):
+            return tuple(value)
+        elif isinstance(value, dict):
+            return tuple(sorted((k, convert_value(v)) for k, v in value.items()))
+        else:
+            return value
+
+    return tuple(sorted((k, convert_value(v)) for k, v in dict_key.items()))
+
+
+DIFFUSERS_SCHEDULER_CONFIG = {
+    "_class_name": "FlowMatchEulerDiscreteScheduler",
+    "_diffusers_version": "0.32.0.dev0",
+    "base_image_seq_len": 1024,
+    "base_shift": 0.95,
+    "invert_sigmas": False,
+    "max_image_seq_len": 4096,
+    "max_shift": 2.05,
+    "num_train_timesteps": 1000,
+    "shift": 1.0,
+    "shift_terminal": 0.1,
+    "use_beta_sigmas": False,
+    "use_dynamic_shifting": True,
+    "use_exponential_sigmas": False,
+    "use_karras_sigmas": False,
+}
+DIFFUSERS_TRANSFORMER_CONFIG = {
+    "_class_name": "LTXVideoTransformer3DModel",
+    "_diffusers_version": "0.32.0.dev0",
+    "activation_fn": "gelu-approximate",
+    "attention_bias": True,
+    "attention_head_dim": 64,
+    "attention_out_bias": True,
+    "caption_channels": 4096,
+    "cross_attention_dim": 2048,
+    "in_channels": 128,
+    "norm_elementwise_affine": False,
+    "norm_eps": 1e-06,
+    "num_attention_heads": 32,
+    "num_layers": 28,
+    "out_channels": 128,
+    "patch_size": 1,
+    "patch_size_t": 1,
+    "qk_norm": "rms_norm_across_heads",
+}
+DIFFUSERS_VAE_CONFIG = {
+    "_class_name": "AutoencoderKLLTXVideo",
+    "_diffusers_version": "0.32.0.dev0",
+    "block_out_channels": [128, 256, 512, 512],
+    "decoder_causal": False,
+    "encoder_causal": True,
+    "in_channels": 3,
+    "latent_channels": 128,
+    "layers_per_block": [4, 3, 3, 3, 4],
+    "out_channels": 3,
+    "patch_size": 4,
+    "patch_size_t": 1,
+    "resnet_norm_eps": 1e-06,
+    "scaling_factor": 1.0,
+    "spatio_temporal_scaling": [True, True, True, False],
+}
+
+OURS_SCHEDULER_CONFIG = {
+    "_class_name": "RectifiedFlowScheduler",
+    "_diffusers_version": "0.25.1",
+    "num_train_timesteps": 1000,
+    "shifting": "SD3",
+    "base_resolution": None,
+    "target_shift_terminal": 0.1,
+}
+
+OURS_TRANSFORMER_CONFIG = {
+    "_class_name": "Transformer3DModel",
+    "_diffusers_version": "0.25.1",
+    "_name_or_path": "PixArt-alpha/PixArt-XL-2-256x256",
+    "activation_fn": "gelu-approximate",
+    "attention_bias": True,
+    "attention_head_dim": 64,
+    "attention_type": "default",
+    "caption_channels": 4096,
+    "cross_attention_dim": 2048,
+    "double_self_attention": False,
+    "dropout": 0.0,
+    "in_channels": 128,
+    "norm_elementwise_affine": False,
+    "norm_eps": 1e-06,
+    "norm_num_groups": 32,
+    "num_attention_heads": 32,
+    "num_embeds_ada_norm": 1000,
+    "num_layers": 28,
+    "num_vector_embeds": None,
+    "only_cross_attention": False,
+    "out_channels": 128,
+    "project_to_2d_pos": True,
+    "upcast_attention": False,
+    "use_linear_projection": False,
+    "qk_norm": "rms_norm",
+    "standardization_norm": "rms_norm",
+    "positional_embedding_type": "rope",
+    "positional_embedding_theta": 10000.0,
+    "positional_embedding_max_pos": [20, 2048, 2048],
+    "timestep_scale_multiplier": 1000,
+}
+OURS_VAE_CONFIG = {
+    "_class_name": "CausalVideoAutoencoder",
+    "dims": 3,
+    "in_channels": 3,
+    "out_channels": 3,
+    "latent_channels": 128,
+    "blocks": [
+        ["res_x", 4],
+        ["compress_all", 1],
+        ["res_x_y", 1],
+        ["res_x", 3],
+        ["compress_all", 1],
+        ["res_x_y", 1],
+        ["res_x", 3],
+        ["compress_all", 1],
+        ["res_x", 3],
+        ["res_x", 4],
+    ],
+    "scaling_factor": 1.0,
+    "norm_layer": "pixel_norm",
+    "patch_size": 4,
+    "latent_log_var": "uniform",
+    "use_quant_conv": False,
+    "causal_decoder": False,
+}
+
+
+diffusers_and_ours_config_mapping = {
+    make_hashable_key(DIFFUSERS_SCHEDULER_CONFIG): OURS_SCHEDULER_CONFIG,
+    make_hashable_key(DIFFUSERS_TRANSFORMER_CONFIG): OURS_TRANSFORMER_CONFIG,
+    make_hashable_key(DIFFUSERS_VAE_CONFIG): OURS_VAE_CONFIG,
+}
+
+
+TRANSFORMER_KEYS_RENAME_DICT = {
+    "proj_in": "patchify_proj",
+    "time_embed": "adaln_single",
+    "norm_q": "q_norm",
+    "norm_k": "k_norm",
+}
+
+
+VAE_KEYS_RENAME_DICT = {
+    "decoder.up_blocks.3.conv_in": "decoder.up_blocks.7",
+    "decoder.up_blocks.3.upsamplers.0": "decoder.up_blocks.8",
+    "decoder.up_blocks.3": "decoder.up_blocks.9",
+    "decoder.up_blocks.2.upsamplers.0": "decoder.up_blocks.5",
+    "decoder.up_blocks.2.conv_in": "decoder.up_blocks.4",
+    "decoder.up_blocks.2": "decoder.up_blocks.6",
+    "decoder.up_blocks.1.upsamplers.0": "decoder.up_blocks.2",
+    "decoder.up_blocks.1": "decoder.up_blocks.3",
+    "decoder.up_blocks.0": "decoder.up_blocks.1",
+    "decoder.mid_block": "decoder.up_blocks.0",
+    "encoder.down_blocks.3": "encoder.down_blocks.8",
+    "encoder.down_blocks.2.downsamplers.0": "encoder.down_blocks.7",
+    "encoder.down_blocks.2": "encoder.down_blocks.6",
+    "encoder.down_blocks.1.downsamplers.0": "encoder.down_blocks.4",
+    "encoder.down_blocks.1.conv_out": "encoder.down_blocks.5",
+    "encoder.down_blocks.1": "encoder.down_blocks.3",
+    "encoder.down_blocks.0.conv_out": "encoder.down_blocks.2",
+    "encoder.down_blocks.0.downsamplers.0": "encoder.down_blocks.1",
+    "encoder.down_blocks.0": "encoder.down_blocks.0",
+    "encoder.mid_block": "encoder.down_blocks.9",
+    "conv_shortcut.conv": "conv_shortcut",
+    "resnets": "res_blocks",
+    "norm3": "norm3.norm",
+    "latents_mean": "per_channel_statistics.mean-of-means",
+    "latents_std": "per_channel_statistics.std-of-means",
+}

ltx_video/utils/prompt_enhance_utils.py

@@ -0,0 +1,214 @@
+import logging
+from typing import Union, List, Optional
+
+import torch
+from PIL import Image
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+T2V_CINEMATIC_PROMPT = """You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.
+Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.
+Start directly with the action, and keep descriptions literal and precise.
+Think like a cinematographer describing a shot list.
+Do not change the user input intent, just enhance it.
+Keep within 150 words.
+For best results, build your prompts using this structure:
+Start with main action in a single sentence
+Add specific details about movements and gestures
+Describe character/object appearances precisely
+Include background and environment details
+Specify camera angles and movements
+Describe lighting and colors
+Note any changes or sudden events
+Do not exceed the 150 word limit!
+Output the enhanced prompt only.
+"""
+
+I2V_CINEMATIC_PROMPT = """You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.
+Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.
+Start directly with the action, and keep descriptions literal and precise.
+Think like a cinematographer describing a shot list.
+Keep within 150 words.
+For best results, build your prompts using this structure:
+Describe the image first and then add the user input. Image description should be in first priority! Align to the image caption if it contradicts the user text input.
+Start with main action in a single sentence
+Add specific details about movements and gestures
+Describe character/object appearances precisely
+Include background and environment details
+Specify camera angles and movements
+Describe lighting and colors
+Note any changes or sudden events
+Align to the image caption if it contradicts the user text input.
+Do not exceed the 150 word limit!
+Output the enhanced prompt only.
+"""
+
+
+def tensor_to_pil(tensor):
+    # Ensure tensor is in range [-1, 1]
+    assert tensor.min() >= -1 and tensor.max() <= 1
+
+    # Convert from [-1, 1] to [0, 1]
+    tensor = (tensor + 1) / 2
+
+    # Rearrange from [C, H, W] to [H, W, C]
+    tensor = tensor.permute(1, 2, 0)
+
+    # Convert to numpy array and then to uint8 range [0, 255]
+    numpy_image = (tensor.cpu().numpy() * 255).astype("uint8")
+
+    # Convert to PIL Image
+    return Image.fromarray(numpy_image)
+
+
+def generate_cinematic_prompt(
+    image_caption_model,
+    image_caption_processor,
+    prompt_enhancer_model,
+    prompt_enhancer_tokenizer,
+    prompt: Union[str, List[str]],
+    images: Optional[List] = None,
+    max_new_tokens: int = 256,
+) -> List[str]:
+    prompts = [prompt] if isinstance(prompt, str) else prompt
+
+    if images is None:
+        prompts = _generate_t2v_prompt(
+            prompt_enhancer_model,
+            prompt_enhancer_tokenizer,
+            prompts,
+            max_new_tokens,
+            T2V_CINEMATIC_PROMPT,
+        )
+    else:
+
+        prompts = _generate_i2v_prompt(
+            image_caption_model,
+            image_caption_processor,
+            prompt_enhancer_model,
+            prompt_enhancer_tokenizer,
+            prompts,
+            images,
+            max_new_tokens,
+            I2V_CINEMATIC_PROMPT,
+        )
+
+    return prompts
+
+
+def _get_first_frames_from_conditioning_item(conditioning_item) -> List[Image.Image]:
+    frames_tensor = conditioning_item.media_item
+    return [
+        tensor_to_pil(frames_tensor[i, :, 0, :, :])
+        for i in range(frames_tensor.shape[0])
+    ]
+
+
+def _generate_t2v_prompt(
+    prompt_enhancer_model,
+    prompt_enhancer_tokenizer,
+    prompts: List[str],
+    max_new_tokens: int,
+    system_prompt: str,
+) -> List[str]:
+    messages = [
+        [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"user_prompt: {p}"},
+        ]
+        for p in prompts
+    ]
+
+    texts = [
+        prompt_enhancer_tokenizer.apply_chat_template(
+            m, tokenize=False, add_generation_prompt=True
+        )
+        for m in messages
+    ]
+    model_inputs = prompt_enhancer_tokenizer(texts, return_tensors="pt").to(
+        prompt_enhancer_model.device
+    )
+
+    return _generate_and_decode_prompts(
+        prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens
+    )
+
+
+def _generate_i2v_prompt(
+    image_caption_model,
+    image_caption_processor,
+    prompt_enhancer_model,
+    prompt_enhancer_tokenizer,
+    prompts: List[str],
+    first_frames: List[Image.Image],
+    max_new_tokens: int,
+    system_prompt: str,
+) -> List[str]:
+    image_captions = _generate_image_captions(
+        image_caption_model, image_caption_processor, first_frames
+    )
+    if len(image_captions) == 1 and len(image_captions) < len(prompts):
+        image_captions *= len(prompts)
+    messages = [
+        [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": f"user_prompt: {p}\nimage_caption: {c}"},
+        ]
+        for p, c in zip(prompts, image_captions)
+    ]
+
+    texts = [
+        prompt_enhancer_tokenizer.apply_chat_template(
+            m, tokenize=False, add_generation_prompt=True
+        )
+        for m in messages
+    ]
+    out_prompts = []
+    for text in texts:
+        model_inputs = prompt_enhancer_tokenizer(text, return_tensors="pt").to(
+            prompt_enhancer_model.device
+        )
+        out_prompts.append(_generate_and_decode_prompts(prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens)[0])
+
+    return out_prompts
+
+
+def _generate_image_captions(
+    image_caption_model,
+    image_caption_processor,
+    images: List[Image.Image],
+    system_prompt: str = "<DETAILED_CAPTION>",
+) -> List[str]:
+    image_caption_prompts = [system_prompt] * len(images)
+    inputs = image_caption_processor(
+        image_caption_prompts, images, return_tensors="pt"
+    ).to("cuda") #.to(image_caption_model.device)
+
+    with torch.inference_mode():
+        generated_ids = image_caption_model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=1024,
+            do_sample=False,
+            num_beams=3,
+        )
+
+    return image_caption_processor.batch_decode(generated_ids, skip_special_tokens=True)
+
+
+def _generate_and_decode_prompts(
+    prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens: int
+) -> List[str]:
+    with torch.inference_mode():
+        outputs = prompt_enhancer_model.generate(
+            **model_inputs, max_new_tokens=max_new_tokens
+        )
+        generated_ids = [
+            output_ids[len(input_ids) :]
+            for input_ids, output_ids in zip(model_inputs.input_ids, outputs)
+        ]
+        decoded_prompts = prompt_enhancer_tokenizer.batch_decode(
+            generated_ids, skip_special_tokens=True
+        )
+
+    return decoded_prompts

ltx_video/utils/skip_layer_strategy.py

@@ -0,0 +1,8 @@
+from enum import Enum, auto
+
+
+class SkipLayerStrategy(Enum):
+    AttentionSkip = auto()
+    AttentionValues = auto()
+    Residual = auto()
+    TransformerBlock = auto()

ltx_video/utils/torch_utils.py

@@ -0,0 +1,25 @@
+import torch
+from torch import nn
+
+
+def append_dims(x: torch.Tensor, target_dims: int) -> torch.Tensor:
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(
+            f"input has {x.ndim} dims but target_dims is {target_dims}, which is less"
+        )
+    elif dims_to_append == 0:
+        return x
+    return x[(...,) + (None,) * dims_to_append]
+
+
+class Identity(nn.Module):
+    """A placeholder identity operator that is argument-insensitive."""
+
+    def __init__(self, *args, **kwargs) -> None:  # pylint: disable=unused-argument
+        super().__init__()
+
+    # pylint: disable=unused-argument
+    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        return x