v5 release with triple architecture support and prompt enhancer
This commit is contained in:
DeepBeepMeep
0
ltx_video/__init__.py
Normal file
0
ltx_video/__init__.py
Normal file
41
ltx_video/configs/ltxv-13b-0.9.7-dev.original.yaml
Normal file
41
ltx_video/configs/ltxv-13b-0.9.7-dev.original.yaml
Normal file
@@ -0,0 +1,41 @@
|
||||
|
||||
pipeline_type: multi-scale
|
||||
checkpoint_path: "ltxv-13b-0.9.7-dev.safetensors"
|
||||
downscale_factor: 0.6666666
|
||||
spatial_upscaler_model_path: "ltxv-spatial-upscaler-0.9.7.safetensors"
|
||||
stg_mode: "attention_values" # options: "attention_values", "attention_skip", "residual", "transformer_block"
|
||||
decode_timestep: 0.05
|
||||
decode_noise_scale: 0.025
|
||||
text_encoder_model_name_or_path: "PixArt-alpha/PixArt-XL-2-1024-MS"
|
||||
sampler: "from_checkpoint" # options: "uniform", "linear-quadratic", "from_checkpoint"
|
||||
prompt_enhancement_words_threshold: 120
|
||||
prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
|
||||
prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
|
||||
stochastic_sampling: false
|
||||
|
||||
|
||||
first_pass:
|
||||
#13b Dynamic
|
||||
guidance_scale: [1, 6, 8, 6, 1, 1]
|
||||
stg_scale: [0, 4, 4, 4, 2, 1]
|
||||
rescaling_scale: [1, 0.5, 0.5, 1, 1, 1]
|
||||
guidance_timesteps: [1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180]
|
||||
skip_block_list: [[11, 25, 35, 39], [22, 35, 39], [28], [28], [28], [28]]
|
||||
num_inference_steps: 30 #default
|
||||
|
||||
|
||||
second_pass:
|
||||
#13b Dynamic
|
||||
guidance_scale: [1, 6, 8, 6, 1, 1]
|
||||
stg_scale: [0, 4, 4, 4, 2, 1]
|
||||
rescaling_scale: [1, 0.5, 0.5, 1, 1, 1]
|
||||
guidance_timesteps: [1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180]
|
||||
skip_block_list: [[11, 25, 35, 39], [22, 35, 39], [28], [28], [28], [28]]
|
||||
#13b Upscale
|
||||
# guidance_scale: [1, 1, 1, 1, 1, 1]
|
||||
# stg_scale: [1, 1, 1, 1, 1, 1]
|
||||
# rescaling_scale: [1, 1, 1, 1, 1, 1]
|
||||
# guidance_timesteps: [1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180]
|
||||
# skip_block_list: [[42], [42], [42], [42], [42], [42]]
|
||||
num_inference_steps: 30 #default
|
||||
strength: 0.85
|
||||
34
ltx_video/configs/ltxv-13b-0.9.7-dev.yaml
Normal file
34
ltx_video/configs/ltxv-13b-0.9.7-dev.yaml
Normal file
@@ -0,0 +1,34 @@
|
||||
pipeline_type: multi-scale
|
||||
checkpoint_path: "ltxv-13b-0.9.7-dev.safetensors"
|
||||
downscale_factor: 0.6666666
|
||||
spatial_upscaler_model_path: "ltxv-spatial-upscaler-0.9.7.safetensors"
|
||||
stg_mode: "attention_values" # options: "attention_values", "attention_skip", "residual", "transformer_block"
|
||||
decode_timestep: 0.05
|
||||
decode_noise_scale: 0.025
|
||||
text_encoder_model_name_or_path: "PixArt-alpha/PixArt-XL-2-1024-MS"
|
||||
precision: "bfloat16"
|
||||
sampler: "from_checkpoint" # options: "uniform", "linear-quadratic", "from_checkpoint"
|
||||
prompt_enhancement_words_threshold: 120
|
||||
prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
|
||||
prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
|
||||
stochastic_sampling: false
|
||||
|
||||
first_pass:
|
||||
guidance_scale: [1, 1, 6, 8, 6, 1, 1]
|
||||
stg_scale: [0, 0, 4, 4, 4, 2, 1]
|
||||
rescaling_scale: [1, 1, 0.5, 0.5, 1, 1, 1]
|
||||
guidance_timesteps: [1.0, 0.996, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180]
|
||||
skip_block_list: [[], [11, 25, 35, 39], [22, 35, 39], [28], [28], [28], [28]]
|
||||
num_inference_steps: 30
|
||||
skip_final_inference_steps: 3
|
||||
cfg_star_rescale: true
|
||||
|
||||
second_pass:
|
||||
guidance_scale: [1]
|
||||
stg_scale: [1]
|
||||
rescaling_scale: [1]
|
||||
guidance_timesteps: [1.0]
|
||||
skip_block_list: [27]
|
||||
num_inference_steps: 30
|
||||
skip_initial_inference_steps: 17
|
||||
cfg_star_rescale: true
|
||||
28
ltx_video/configs/ltxv-13b-0.9.7-distilled.yaml
Normal file
28
ltx_video/configs/ltxv-13b-0.9.7-distilled.yaml
Normal file
@@ -0,0 +1,28 @@
|
||||
pipeline_type: multi-scale
|
||||
checkpoint_path: "ltxv-13b-0.9.7-distilled.safetensors"
|
||||
downscale_factor: 0.6666666
|
||||
spatial_upscaler_model_path: "ltxv-spatial-upscaler-0.9.7.safetensors"
|
||||
stg_mode: "attention_values" # options: "attention_values", "attention_skip", "residual", "transformer_block"
|
||||
decode_timestep: 0.05
|
||||
decode_noise_scale: 0.025
|
||||
text_encoder_model_name_or_path: "PixArt-alpha/PixArt-XL-2-1024-MS"
|
||||
precision: "bfloat16"
|
||||
sampler: "from_checkpoint" # options: "uniform", "linear-quadratic", "from_checkpoint"
|
||||
prompt_enhancement_words_threshold: 120
|
||||
prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
|
||||
prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
|
||||
stochastic_sampling: false
|
||||
|
||||
first_pass:
|
||||
timesteps: [1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250]
|
||||
guidance_scale: 1
|
||||
stg_scale: 0
|
||||
rescaling_scale: 1
|
||||
skip_block_list: [42]
|
||||
|
||||
second_pass:
|
||||
timesteps: [0.9094, 0.7250, 0.4219]
|
||||
guidance_scale: 1
|
||||
stg_scale: 0
|
||||
rescaling_scale: 1
|
||||
skip_block_list: [42]
|
||||
17
ltx_video/configs/ltxv-2b-0.9.6-dev.yaml
Normal file
17
ltx_video/configs/ltxv-2b-0.9.6-dev.yaml
Normal file
@@ -0,0 +1,17 @@
|
||||
pipeline_type: base
|
||||
checkpoint_path: "ltxv-2b-0.9.6-dev-04-25.safetensors"
|
||||
guidance_scale: 3
|
||||
stg_scale: 1
|
||||
rescaling_scale: 0.7
|
||||
skip_block_list: [19]
|
||||
num_inference_steps: 40
|
||||
stg_mode: "attention_values" # options: "attention_values", "attention_skip", "residual", "transformer_block"
|
||||
decode_timestep: 0.05
|
||||
decode_noise_scale: 0.025
|
||||
text_encoder_model_name_or_path: "PixArt-alpha/PixArt-XL-2-1024-MS"
|
||||
precision: "bfloat16"
|
||||
sampler: "from_checkpoint" # options: "uniform", "linear-quadratic", "from_checkpoint"
|
||||
prompt_enhancement_words_threshold: 120
|
||||
prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
|
||||
prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
|
||||
stochastic_sampling: false
|
||||
17
ltx_video/configs/ltxv-2b-0.9.6-distilled.yaml
Normal file
17
ltx_video/configs/ltxv-2b-0.9.6-distilled.yaml
Normal file
@@ -0,0 +1,17 @@
|
||||
pipeline_type: base
|
||||
checkpoint_path: "ltxv-2b-0.9.6-distilled-04-25.safetensors"
|
||||
guidance_scale: 3
|
||||
stg_scale: 1
|
||||
rescaling_scale: 0.7
|
||||
skip_block_list: [19]
|
||||
num_inference_steps: 8
|
||||
stg_mode: "attention_values" # options: "attention_values", "attention_skip", "residual", "transformer_block"
|
||||
decode_timestep: 0.05
|
||||
decode_noise_scale: 0.025
|
||||
text_encoder_model_name_or_path: "PixArt-alpha/PixArt-XL-2-1024-MS"
|
||||
precision: "bfloat16"
|
||||
sampler: "from_checkpoint" # options: "uniform", "linear-quadratic", "from_checkpoint"
|
||||
prompt_enhancement_words_threshold: 120
|
||||
prompt_enhancer_image_caption_model_name_or_path: "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
|
||||
prompt_enhancer_llm_model_name_or_path: "unsloth/Llama-3.2-3B-Instruct"
|
||||
stochastic_sampling: true
|
||||
562
ltx_video/ltxv.py
Normal file
562
ltx_video/ltxv.py
Normal file
@@ -0,0 +1,562 @@
|
||||
from mmgp import offload
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from diffusers.utils import logging
|
||||
from typing import Optional, List, Union
|
||||
import yaml
|
||||
from wan.utils.utils import calculate_new_dimensions
|
||||
import imageio
|
||||
import json
|
||||
import numpy as np
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from PIL import Image
|
||||
from transformers import (
|
||||
T5EncoderModel,
|
||||
T5Tokenizer,
|
||||
AutoModelForCausalLM,
|
||||
AutoProcessor,
|
||||
AutoTokenizer,
|
||||
)
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
from .models.autoencoders.causal_video_autoencoder import (
|
||||
CausalVideoAutoencoder,
|
||||
)
|
||||
from .models.transformers.symmetric_patchifier import SymmetricPatchifier
|
||||
from .models.transformers.transformer3d import Transformer3DModel
|
||||
from .pipelines.pipeline_ltx_video import (
|
||||
ConditioningItem,
|
||||
LTXVideoPipeline,
|
||||
LTXMultiScalePipeline,
|
||||
)
|
||||
from .schedulers.rf import RectifiedFlowScheduler
|
||||
from .utils.skip_layer_strategy import SkipLayerStrategy
|
||||
from .models.autoencoders.latent_upsampler import LatentUpsampler
|
||||
from .pipelines import crf_compressor
|
||||
import cv2
|
||||
|
||||
MAX_HEIGHT = 720
|
||||
MAX_WIDTH = 1280
|
||||
MAX_NUM_FRAMES = 257
|
||||
|
||||
logger = logging.get_logger("LTX-Video")
|
||||
|
||||
|
||||
def get_total_gpu_memory():
|
||||
if torch.cuda.is_available():
|
||||
total_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
|
||||
return total_memory
|
||||
return 0
|
||||
|
||||
|
||||
def get_device():
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
elif torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
return "cpu"
|
||||
|
||||
|
||||
def load_image_to_tensor_with_resize_and_crop(
|
||||
image_input: Union[str, Image.Image],
|
||||
target_height: int = 512,
|
||||
target_width: int = 768,
|
||||
just_crop: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Load and process an image into a tensor.
|
||||
|
||||
Args:
|
||||
image_input: Either a file path (str) or a PIL Image object
|
||||
target_height: Desired height of output tensor
|
||||
target_width: Desired width of output tensor
|
||||
just_crop: If True, only crop the image to the target size without resizing
|
||||
"""
|
||||
if isinstance(image_input, str):
|
||||
image = Image.open(image_input).convert("RGB")
|
||||
elif isinstance(image_input, Image.Image):
|
||||
image = image_input
|
||||
else:
|
||||
raise ValueError("image_input must be either a file path or a PIL Image object")
|
||||
|
||||
input_width, input_height = image.size
|
||||
aspect_ratio_target = target_width / target_height
|
||||
aspect_ratio_frame = input_width / input_height
|
||||
if aspect_ratio_frame > aspect_ratio_target:
|
||||
new_width = int(input_height * aspect_ratio_target)
|
||||
new_height = input_height
|
||||
x_start = (input_width - new_width) // 2
|
||||
y_start = 0
|
||||
else:
|
||||
new_width = input_width
|
||||
new_height = int(input_width / aspect_ratio_target)
|
||||
x_start = 0
|
||||
y_start = (input_height - new_height) // 2
|
||||
|
||||
image = image.crop((x_start, y_start, x_start + new_width, y_start + new_height))
|
||||
if not just_crop:
|
||||
image = image.resize((target_width, target_height))
|
||||
|
||||
image = np.array(image)
|
||||
image = cv2.GaussianBlur(image, (3, 3), 0)
|
||||
frame_tensor = torch.from_numpy(image).float()
|
||||
frame_tensor = crf_compressor.compress(frame_tensor / 255.0) * 255.0
|
||||
frame_tensor = frame_tensor.permute(2, 0, 1)
|
||||
frame_tensor = (frame_tensor / 127.5) - 1.0
|
||||
# Create 5D tensor: (batch_size=1, channels=3, num_frames=1, height, width)
|
||||
return frame_tensor.unsqueeze(0).unsqueeze(2)
|
||||
|
||||
|
||||
|
||||
def calculate_padding(
|
||||
source_height: int, source_width: int, target_height: int, target_width: int
|
||||
) -> tuple[int, int, int, int]:
|
||||
|
||||
# Calculate total padding needed
|
||||
pad_height = target_height - source_height
|
||||
pad_width = target_width - source_width
|
||||
|
||||
# Calculate padding for each side
|
||||
pad_top = pad_height // 2
|
||||
pad_bottom = pad_height - pad_top # Handles odd padding
|
||||
pad_left = pad_width // 2
|
||||
pad_right = pad_width - pad_left # Handles odd padding
|
||||
|
||||
# Return padded tensor
|
||||
# Padding format is (left, right, top, bottom)
|
||||
padding = (pad_left, pad_right, pad_top, pad_bottom)
|
||||
return padding
|
||||
|
||||
|
||||
|
||||
|
||||
def seed_everething(seed: int):
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
if torch.backends.mps.is_available():
|
||||
torch.mps.manual_seed(seed)
|
||||
|
||||
|
||||
class LTXV:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_filepath: str,
|
||||
text_encoder_filepath: str,
|
||||
dtype = torch.bfloat16,
|
||||
VAE_dtype = torch.bfloat16,
|
||||
mixed_precision_transformer = False
|
||||
):
|
||||
|
||||
self.mixed_precision_transformer = mixed_precision_transformer
|
||||
# ckpt_path = Path(ckpt_path)
|
||||
# with safe_open(ckpt_path, framework="pt") as f:
|
||||
# metadata = f.metadata()
|
||||
# config_str = metadata.get("config")
|
||||
# configs = json.loads(config_str)
|
||||
# allowed_inference_steps = configs.get("allowed_inference_steps", None)
|
||||
# transformer = Transformer3DModel.from_pretrained(ckpt_path)
|
||||
# offload.save_model(transformer, "ckpts/ltxv_0.9.7_13B_dev_bf16.safetensors", config_file_path="config_transformer.json")
|
||||
|
||||
# vae = CausalVideoAutoencoder.from_pretrained(ckpt_path)
|
||||
vae = offload.fast_load_transformers_model("ckpts/ltxv_0.9.7_VAE.safetensors", modelClass=CausalVideoAutoencoder)
|
||||
if VAE_dtype == torch.float16:
|
||||
VAE_dtype = torch.bfloat16
|
||||
|
||||
vae = vae.to(VAE_dtype)
|
||||
vae._model_dtype = VAE_dtype
|
||||
# vae = offload.fast_load_transformers_model("vae.safetensors", modelClass=CausalVideoAutoencoder, modelPrefix= "vae", forcedConfigPath="config_vae.json")
|
||||
# offload.save_model(vae, "vae.safetensors", config_file_path="config_vae.json")
|
||||
|
||||
|
||||
transformer = offload.fast_load_transformers_model(model_filepath, modelClass=Transformer3DModel)
|
||||
transformer._model_dtype = dtype
|
||||
if mixed_precision_transformer:
|
||||
transformer._lock_dtype = torch.float
|
||||
|
||||
|
||||
scheduler = RectifiedFlowScheduler.from_pretrained("ckpts/ltxv_scheduler.json")
|
||||
# transformer = offload.fast_load_transformers_model("ltx_13B_quanto_bf16_int8.safetensors", modelClass=Transformer3DModel, modelPrefix= "model.diffusion_model", forcedConfigPath="config_transformer.json")
|
||||
# offload.save_model(transformer, "ltx_13B_quanto_bf16_int8.safetensors", do_quantize= True, config_file_path="config_transformer.json")
|
||||
|
||||
latent_upsampler = LatentUpsampler.from_pretrained("ckpts/ltxv_0.9.7_spatial_upscaler.safetensors").to("cpu").eval()
|
||||
latent_upsampler.to(VAE_dtype)
|
||||
latent_upsampler._model_dtype = VAE_dtype
|
||||
|
||||
allowed_inference_steps = None
|
||||
|
||||
# text_encoder = T5EncoderModel.from_pretrained(
|
||||
# "PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder"
|
||||
# )
|
||||
# text_encoder.to(torch.bfloat16)
|
||||
# offload.save_model(text_encoder, "T5_xxl_1.1_enc_bf16.safetensors", config_file_path="T5_config.json")
|
||||
# offload.save_model(text_encoder, "T5_xxl_1.1_enc_quanto_bf16_int8.safetensors", do_quantize= True, config_file_path="T5_config.json")
|
||||
|
||||
text_encoder = offload.fast_load_transformers_model(text_encoder_filepath)
|
||||
patchifier = SymmetricPatchifier(patch_size=1)
|
||||
tokenizer = T5Tokenizer.from_pretrained( "ckpts/T5_xxl_1.1")
|
||||
|
||||
enhance_prompt = False
|
||||
if enhance_prompt:
|
||||
prompt_enhancer_image_caption_model = AutoModelForCausalLM.from_pretrained( "ckpts/Florence2", trust_remote_code=True)
|
||||
prompt_enhancer_image_caption_processor = AutoProcessor.from_pretrained( "ckpts/Florence2", trust_remote_code=True)
|
||||
prompt_enhancer_llm_model = offload.fast_load_transformers_model("ckpts/Llama3_2_quanto_bf16_int8.safetensors")
|
||||
prompt_enhancer_llm_tokenizer = AutoTokenizer.from_pretrained("ckpts/Llama3_2")
|
||||
else:
|
||||
prompt_enhancer_image_caption_model = None
|
||||
prompt_enhancer_image_caption_processor = None
|
||||
prompt_enhancer_llm_model = None
|
||||
prompt_enhancer_llm_tokenizer = None
|
||||
|
||||
if prompt_enhancer_image_caption_model != None:
|
||||
pipe["prompt_enhancer_image_caption_model"] = prompt_enhancer_image_caption_model
|
||||
prompt_enhancer_image_caption_model._model_dtype = torch.float
|
||||
|
||||
pipe["prompt_enhancer_llm_model"] = prompt_enhancer_llm_model
|
||||
|
||||
# offload.profile(pipe, profile_no=5, extraModelsToQuantize = None, quantizeTransformer = False, budgets = { "prompt_enhancer_llm_model" : 10000, "prompt_enhancer_image_caption_model" : 10000, "vae" : 3000, "*" : 100 }, verboseLevel=2)
|
||||
|
||||
|
||||
# Use submodels for the pipeline
|
||||
submodel_dict = {
|
||||
"transformer": transformer,
|
||||
"patchifier": patchifier,
|
||||
"text_encoder": text_encoder,
|
||||
"tokenizer": tokenizer,
|
||||
"scheduler": scheduler,
|
||||
"vae": vae,
|
||||
"prompt_enhancer_image_caption_model": prompt_enhancer_image_caption_model,
|
||||
"prompt_enhancer_image_caption_processor": prompt_enhancer_image_caption_processor,
|
||||
"prompt_enhancer_llm_model": prompt_enhancer_llm_model,
|
||||
"prompt_enhancer_llm_tokenizer": prompt_enhancer_llm_tokenizer,
|
||||
"allowed_inference_steps": allowed_inference_steps,
|
||||
}
|
||||
pipeline = LTXVideoPipeline(**submodel_dict)
|
||||
pipeline = LTXMultiScalePipeline(pipeline, latent_upsampler=latent_upsampler)
|
||||
|
||||
self.pipeline = pipeline
|
||||
self.model = transformer
|
||||
self.vae = vae
|
||||
# return pipeline, pipe
|
||||
|
||||
def generate(
|
||||
self,
|
||||
input_prompt: str,
|
||||
n_prompt: str,
|
||||
image_start = None,
|
||||
image_end = None,
|
||||
input_video = None,
|
||||
sampling_steps = 50,
|
||||
image_cond_noise_scale: float = 0.15,
|
||||
input_media_path: Optional[str] = None,
|
||||
strength: Optional[float] = 1.0,
|
||||
seed: int = 42,
|
||||
height: Optional[int] = 704,
|
||||
width: Optional[int] = 1216,
|
||||
frame_num: int = 81,
|
||||
frame_rate: int = 30,
|
||||
fit_into_canvas = True,
|
||||
callback=None,
|
||||
device: Optional[str] = None,
|
||||
VAE_tile_size = None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
num_inference_steps1 = sampling_steps
|
||||
num_inference_steps2 = sampling_steps #10
|
||||
conditioning_strengths = None
|
||||
conditioning_media_paths = []
|
||||
conditioning_start_frames = []
|
||||
|
||||
|
||||
if input_video != None:
|
||||
conditioning_media_paths.append(input_video)
|
||||
conditioning_start_frames.append(0)
|
||||
height, width = input_video.shape[-2:]
|
||||
else:
|
||||
if image_start != None:
|
||||
image_start = image_start[0]
|
||||
frame_width, frame_height = image_start.size
|
||||
height, width = calculate_new_dimensions(height, width, frame_height, frame_width, fit_into_canvas, 32)
|
||||
conditioning_media_paths.append(image_start)
|
||||
conditioning_start_frames.append(0)
|
||||
if image_end != None:
|
||||
image_end = image_end[0]
|
||||
conditioning_media_paths.append(image_end)
|
||||
conditioning_start_frames.append(frame_num-1)
|
||||
|
||||
if len(conditioning_media_paths) == 0:
|
||||
conditioning_media_paths = None
|
||||
conditioning_start_frames = None
|
||||
|
||||
pipeline_config = "ltx_video/configs/ltxv-13b-0.9.7-dev.yaml"
|
||||
# check if pipeline_config is a file
|
||||
if not os.path.isfile(pipeline_config):
|
||||
raise ValueError(f"Pipeline config file {pipeline_config} does not exist")
|
||||
with open(pipeline_config, "r") as f:
|
||||
pipeline_config = yaml.safe_load(f)
|
||||
|
||||
|
||||
# Validate conditioning arguments
|
||||
if conditioning_media_paths:
|
||||
# Use default strengths of 1.0
|
||||
if not conditioning_strengths:
|
||||
conditioning_strengths = [1.0] * len(conditioning_media_paths)
|
||||
if not conditioning_start_frames:
|
||||
raise ValueError(
|
||||
"If `conditioning_media_paths` is provided, "
|
||||
"`conditioning_start_frames` must also be provided"
|
||||
)
|
||||
if len(conditioning_media_paths) != len(conditioning_strengths) or len(
|
||||
conditioning_media_paths
|
||||
) != len(conditioning_start_frames):
|
||||
raise ValueError(
|
||||
"`conditioning_media_paths`, `conditioning_strengths`, "
|
||||
"and `conditioning_start_frames` must have the same length"
|
||||
)
|
||||
if any(s < 0 or s > 1 for s in conditioning_strengths):
|
||||
raise ValueError("All conditioning strengths must be between 0 and 1")
|
||||
if any(f < 0 or f >= frame_num for f in conditioning_start_frames):
|
||||
raise ValueError(
|
||||
f"All conditioning start frames must be between 0 and {frame_num-1}"
|
||||
)
|
||||
|
||||
# Adjust dimensions to be divisible by 32 and num_frames to be (N * 8 + 1)
|
||||
height_padded = ((height - 1) // 32 + 1) * 32
|
||||
width_padded = ((width - 1) // 32 + 1) * 32
|
||||
num_frames_padded = ((frame_num - 2) // 8 + 1) * 8 + 1
|
||||
|
||||
padding = calculate_padding(height, width, height_padded, width_padded)
|
||||
|
||||
logger.warning(
|
||||
f"Padded dimensions: {height_padded}x{width_padded}x{num_frames_padded}"
|
||||
)
|
||||
|
||||
|
||||
# prompt_enhancement_words_threshold = pipeline_config[
|
||||
# "prompt_enhancement_words_threshold"
|
||||
# ]
|
||||
|
||||
# prompt_word_count = len(prompt.split())
|
||||
# enhance_prompt = (
|
||||
# prompt_enhancement_words_threshold > 0
|
||||
# and prompt_word_count < prompt_enhancement_words_threshold
|
||||
# )
|
||||
|
||||
# # enhance_prompt = False
|
||||
|
||||
# if prompt_enhancement_words_threshold > 0 and not enhance_prompt:
|
||||
# logger.info(
|
||||
# f"Prompt has {prompt_word_count} words, which exceeds the threshold of {prompt_enhancement_words_threshold}. Prompt enhancement disabled."
|
||||
# )
|
||||
|
||||
|
||||
seed_everething(seed)
|
||||
device = device or get_device()
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
media_item = None
|
||||
if input_media_path:
|
||||
media_item = load_media_file(
|
||||
media_path=input_media_path,
|
||||
height=height,
|
||||
width=width,
|
||||
max_frames=num_frames_padded,
|
||||
padding=padding,
|
||||
)
|
||||
|
||||
conditioning_items = (
|
||||
prepare_conditioning(
|
||||
conditioning_media_paths=conditioning_media_paths,
|
||||
conditioning_strengths=conditioning_strengths,
|
||||
conditioning_start_frames=conditioning_start_frames,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=frame_num,
|
||||
padding=padding,
|
||||
pipeline=self.pipeline,
|
||||
)
|
||||
if conditioning_media_paths
|
||||
else None
|
||||
)
|
||||
|
||||
stg_mode = pipeline_config.get("stg_mode", "attention_values")
|
||||
del pipeline_config["stg_mode"]
|
||||
if stg_mode.lower() == "stg_av" or stg_mode.lower() == "attention_values":
|
||||
skip_layer_strategy = SkipLayerStrategy.AttentionValues
|
||||
elif stg_mode.lower() == "stg_as" or stg_mode.lower() == "attention_skip":
|
||||
skip_layer_strategy = SkipLayerStrategy.AttentionSkip
|
||||
elif stg_mode.lower() == "stg_r" or stg_mode.lower() == "residual":
|
||||
skip_layer_strategy = SkipLayerStrategy.Residual
|
||||
elif stg_mode.lower() == "stg_t" or stg_mode.lower() == "transformer_block":
|
||||
skip_layer_strategy = SkipLayerStrategy.TransformerBlock
|
||||
else:
|
||||
raise ValueError(f"Invalid spatiotemporal guidance mode: {stg_mode}")
|
||||
|
||||
# Prepare input for the pipeline
|
||||
sample = {
|
||||
"prompt": input_prompt,
|
||||
"prompt_attention_mask": None,
|
||||
"negative_prompt": n_prompt,
|
||||
"negative_prompt_attention_mask": None,
|
||||
}
|
||||
|
||||
|
||||
images = self.pipeline(
|
||||
**pipeline_config,
|
||||
ltxv_model = self,
|
||||
num_inference_steps1 = num_inference_steps1,
|
||||
num_inference_steps2 = num_inference_steps2,
|
||||
skip_layer_strategy=skip_layer_strategy,
|
||||
generator=generator,
|
||||
output_type="pt",
|
||||
callback_on_step_end=None,
|
||||
height=height_padded,
|
||||
width=width_padded,
|
||||
num_frames=num_frames_padded,
|
||||
frame_rate=frame_rate,
|
||||
**sample,
|
||||
media_items=media_item,
|
||||
strength=strength,
|
||||
conditioning_items=conditioning_items,
|
||||
is_video=True,
|
||||
vae_per_channel_normalize=True,
|
||||
image_cond_noise_scale=image_cond_noise_scale,
|
||||
mixed_precision=pipeline_config.get("mixed", self.mixed_precision_transformer),
|
||||
callback=callback,
|
||||
VAE_tile_size = VAE_tile_size,
|
||||
device=device,
|
||||
# enhance_prompt=enhance_prompt,
|
||||
)
|
||||
if images == None:
|
||||
return None
|
||||
|
||||
# Crop the padded images to the desired resolution and number of frames
|
||||
(pad_left, pad_right, pad_top, pad_bottom) = padding
|
||||
pad_bottom = -pad_bottom
|
||||
pad_right = -pad_right
|
||||
if pad_bottom == 0:
|
||||
pad_bottom = images.shape[3]
|
||||
if pad_right == 0:
|
||||
pad_right = images.shape[4]
|
||||
images = images[:, :, :frame_num, pad_top:pad_bottom, pad_left:pad_right]
|
||||
images = images.sub_(0.5).mul_(2).squeeze(0)
|
||||
return images
|
||||
|
||||
|
||||
def prepare_conditioning(
|
||||
conditioning_media_paths: List[str],
|
||||
conditioning_strengths: List[float],
|
||||
conditioning_start_frames: List[int],
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
padding: tuple[int, int, int, int],
|
||||
pipeline: LTXVideoPipeline,
|
||||
) -> Optional[List[ConditioningItem]]:
|
||||
"""Prepare conditioning items based on input media paths and their parameters.
|
||||
|
||||
Args:
|
||||
conditioning_media_paths: List of paths to conditioning media (images or videos)
|
||||
conditioning_strengths: List of conditioning strengths for each media item
|
||||
conditioning_start_frames: List of frame indices where each item should be applied
|
||||
height: Height of the output frames
|
||||
width: Width of the output frames
|
||||
num_frames: Number of frames in the output video
|
||||
padding: Padding to apply to the frames
|
||||
pipeline: LTXVideoPipeline object used for condition video trimming
|
||||
|
||||
Returns:
|
||||
A list of ConditioningItem objects.
|
||||
"""
|
||||
conditioning_items = []
|
||||
for path, strength, start_frame in zip(
|
||||
conditioning_media_paths, conditioning_strengths, conditioning_start_frames
|
||||
):
|
||||
if isinstance(path, Image.Image):
|
||||
num_input_frames = orig_num_input_frames = 1
|
||||
else:
|
||||
num_input_frames = orig_num_input_frames = get_media_num_frames(path)
|
||||
if hasattr(pipeline, "trim_conditioning_sequence") and callable(
|
||||
getattr(pipeline, "trim_conditioning_sequence")
|
||||
):
|
||||
num_input_frames = pipeline.trim_conditioning_sequence(
|
||||
start_frame, orig_num_input_frames, num_frames
|
||||
)
|
||||
if num_input_frames < orig_num_input_frames:
|
||||
logger.warning(
|
||||
f"Trimming conditioning video {path} from {orig_num_input_frames} to {num_input_frames} frames."
|
||||
)
|
||||
|
||||
media_tensor = load_media_file(
|
||||
media_path=path,
|
||||
height=height,
|
||||
width=width,
|
||||
max_frames=num_input_frames,
|
||||
padding=padding,
|
||||
just_crop=True,
|
||||
)
|
||||
conditioning_items.append(ConditioningItem(media_tensor, start_frame, strength))
|
||||
return conditioning_items
|
||||
|
||||
|
||||
def get_media_num_frames(media_path: str) -> int:
|
||||
if isinstance(media_path, Image.Image):
|
||||
return 1
|
||||
elif torch.is_tensor(media_path):
|
||||
return media_path.shape[1]
|
||||
elif isinstance(media_path, str) and any( media_path.lower().endswith(ext) for ext in [".mp4", ".avi", ".mov", ".mkv"]):
|
||||
reader = imageio.get_reader(media_path)
|
||||
return min(reader.count_frames(), max_frames)
|
||||
else:
|
||||
raise Exception("video format not supported")
|
||||
|
||||
|
||||
def load_media_file(
|
||||
media_path: str,
|
||||
height: int,
|
||||
width: int,
|
||||
max_frames: int,
|
||||
padding: tuple[int, int, int, int],
|
||||
just_crop: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if isinstance(media_path, Image.Image):
|
||||
# Input image
|
||||
media_tensor = load_image_to_tensor_with_resize_and_crop(
|
||||
media_path, height, width, just_crop=just_crop
|
||||
)
|
||||
media_tensor = torch.nn.functional.pad(media_tensor, padding)
|
||||
|
||||
elif torch.is_tensor(media_path):
|
||||
media_tensor = media_path.unsqueeze(0)
|
||||
num_input_frames = media_tensor.shape[2]
|
||||
elif isinstance(media_path, str) and any( media_path.lower().endswith(ext) for ext in [".mp4", ".avi", ".mov", ".mkv"]):
|
||||
reader = imageio.get_reader(media_path)
|
||||
num_input_frames = min(reader.count_frames(), max_frames)
|
||||
|
||||
# Read and preprocess the relevant frames from the video file.
|
||||
frames = []
|
||||
for i in range(num_input_frames):
|
||||
frame = Image.fromarray(reader.get_data(i))
|
||||
frame_tensor = load_image_to_tensor_with_resize_and_crop(
|
||||
frame, height, width, just_crop=just_crop
|
||||
)
|
||||
frame_tensor = torch.nn.functional.pad(frame_tensor, padding)
|
||||
frames.append(frame_tensor)
|
||||
reader.close()
|
||||
|
||||
# Stack frames along the temporal dimension
|
||||
media_tensor = torch.cat(frames, dim=2)
|
||||
else:
|
||||
raise Exception("video format not supported")
|
||||
return media_tensor
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
0
ltx_video/models/__init__.py
Normal file
0
ltx_video/models/__init__.py
Normal file
0
ltx_video/models/autoencoders/__init__.py
Normal file
0
ltx_video/models/autoencoders/__init__.py
Normal file
63
ltx_video/models/autoencoders/causal_conv3d.py
Normal file
63
ltx_video/models/autoencoders/causal_conv3d.py
Normal file
@@ -0,0 +1,63 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class CausalConv3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size: int = 3,
|
||||
stride: Union[int, Tuple[int]] = 1,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
spatial_padding_mode: str = "zeros",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
kernel_size = (kernel_size, kernel_size, kernel_size)
|
||||
self.time_kernel_size = kernel_size[0]
|
||||
|
||||
dilation = (dilation, 1, 1)
|
||||
|
||||
height_pad = kernel_size[1] // 2
|
||||
width_pad = kernel_size[2] // 2
|
||||
padding = (0, height_pad, width_pad)
|
||||
|
||||
self.conv = nn.Conv3d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
padding_mode=spatial_padding_mode,
|
||||
groups=groups,
|
||||
)
|
||||
|
||||
def forward(self, x, causal: bool = True):
|
||||
if causal:
|
||||
first_frame_pad = x[:, :, :1, :, :].repeat(
|
||||
(1, 1, self.time_kernel_size - 1, 1, 1)
|
||||
)
|
||||
x = torch.concatenate((first_frame_pad, x), dim=2)
|
||||
else:
|
||||
first_frame_pad = x[:, :, :1, :, :].repeat(
|
||||
(1, 1, (self.time_kernel_size - 1) // 2, 1, 1)
|
||||
)
|
||||
last_frame_pad = x[:, :, -1:, :, :].repeat(
|
||||
(1, 1, (self.time_kernel_size - 1) // 2, 1, 1)
|
||||
)
|
||||
x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
return self.conv.weight
|
||||
1405
ltx_video/models/autoencoders/causal_video_autoencoder.py
Normal file
1405
ltx_video/models/autoencoders/causal_video_autoencoder.py
Normal file
File diff suppressed because it is too large
Load Diff
90
ltx_video/models/autoencoders/conv_nd_factory.py
Normal file
90
ltx_video/models/autoencoders/conv_nd_factory.py
Normal file
@@ -0,0 +1,90 @@
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_video.models.autoencoders.dual_conv3d import DualConv3d
|
||||
from ltx_video.models.autoencoders.causal_conv3d import CausalConv3d
|
||||
|
||||
|
||||
def make_conv_nd(
|
||||
dims: Union[int, Tuple[int, int]],
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
causal=False,
|
||||
spatial_padding_mode="zeros",
|
||||
temporal_padding_mode="zeros",
|
||||
):
|
||||
if not (spatial_padding_mode == temporal_padding_mode or causal):
|
||||
raise NotImplementedError("spatial and temporal padding modes must be equal")
|
||||
if dims == 2:
|
||||
return torch.nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif dims == 3:
|
||||
if causal:
|
||||
return CausalConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
spatial_padding_mode=spatial_padding_mode,
|
||||
)
|
||||
return torch.nn.Conv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode,
|
||||
)
|
||||
elif dims == (2, 1):
|
||||
return DualConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
bias=bias,
|
||||
padding_mode=spatial_padding_mode,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
|
||||
|
||||
def make_linear_nd(
|
||||
dims: int,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
bias=True,
|
||||
):
|
||||
if dims == 2:
|
||||
return torch.nn.Conv2d(
|
||||
in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias
|
||||
)
|
||||
elif dims == 3 or dims == (2, 1):
|
||||
return torch.nn.Conv3d(
|
||||
in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unsupported dimensions: {dims}")
|
||||
217
ltx_video/models/autoencoders/dual_conv3d.py
Normal file
217
ltx_video/models/autoencoders/dual_conv3d.py
Normal file
@@ -0,0 +1,217 @@
|
||||
import math
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class DualConv3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride: Union[int, Tuple[int, int, int]] = 1,
|
||||
padding: Union[int, Tuple[int, int, int]] = 0,
|
||||
dilation: Union[int, Tuple[int, int, int]] = 1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
):
|
||||
super(DualConv3d, self).__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.padding_mode = padding_mode
|
||||
# Ensure kernel_size, stride, padding, and dilation are tuples of length 3
|
||||
if isinstance(kernel_size, int):
|
||||
kernel_size = (kernel_size, kernel_size, kernel_size)
|
||||
if kernel_size == (1, 1, 1):
|
||||
raise ValueError(
|
||||
"kernel_size must be greater than 1. Use make_linear_nd instead."
|
||||
)
|
||||
if isinstance(stride, int):
|
||||
stride = (stride, stride, stride)
|
||||
if isinstance(padding, int):
|
||||
padding = (padding, padding, padding)
|
||||
if isinstance(dilation, int):
|
||||
dilation = (dilation, dilation, dilation)
|
||||
|
||||
# Set parameters for convolutions
|
||||
self.groups = groups
|
||||
self.bias = bias
|
||||
|
||||
# Define the size of the channels after the first convolution
|
||||
intermediate_channels = (
|
||||
out_channels if in_channels < out_channels else in_channels
|
||||
)
|
||||
|
||||
# Define parameters for the first convolution
|
||||
self.weight1 = nn.Parameter(
|
||||
torch.Tensor(
|
||||
intermediate_channels,
|
||||
in_channels // groups,
|
||||
1,
|
||||
kernel_size[1],
|
||||
kernel_size[2],
|
||||
)
|
||||
)
|
||||
self.stride1 = (1, stride[1], stride[2])
|
||||
self.padding1 = (0, padding[1], padding[2])
|
||||
self.dilation1 = (1, dilation[1], dilation[2])
|
||||
if bias:
|
||||
self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))
|
||||
else:
|
||||
self.register_parameter("bias1", None)
|
||||
|
||||
# Define parameters for the second convolution
|
||||
self.weight2 = nn.Parameter(
|
||||
torch.Tensor(
|
||||
out_channels, intermediate_channels // groups, kernel_size[0], 1, 1
|
||||
)
|
||||
)
|
||||
self.stride2 = (stride[0], 1, 1)
|
||||
self.padding2 = (padding[0], 0, 0)
|
||||
self.dilation2 = (dilation[0], 1, 1)
|
||||
if bias:
|
||||
self.bias2 = nn.Parameter(torch.Tensor(out_channels))
|
||||
else:
|
||||
self.register_parameter("bias2", None)
|
||||
|
||||
# Initialize weights and biases
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.kaiming_uniform_(self.weight1, a=math.sqrt(5))
|
||||
nn.init.kaiming_uniform_(self.weight2, a=math.sqrt(5))
|
||||
if self.bias:
|
||||
fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)
|
||||
bound1 = 1 / math.sqrt(fan_in1)
|
||||
nn.init.uniform_(self.bias1, -bound1, bound1)
|
||||
fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)
|
||||
bound2 = 1 / math.sqrt(fan_in2)
|
||||
nn.init.uniform_(self.bias2, -bound2, bound2)
|
||||
|
||||
def forward(self, x, use_conv3d=False, skip_time_conv=False):
|
||||
if use_conv3d:
|
||||
return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)
|
||||
else:
|
||||
return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)
|
||||
|
||||
def forward_with_3d(self, x, skip_time_conv):
|
||||
# First convolution
|
||||
x = F.conv3d(
|
||||
x,
|
||||
self.weight1,
|
||||
self.bias1,
|
||||
self.stride1,
|
||||
self.padding1,
|
||||
self.dilation1,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
if skip_time_conv:
|
||||
return x
|
||||
|
||||
# Second convolution
|
||||
x = F.conv3d(
|
||||
x,
|
||||
self.weight2,
|
||||
self.bias2,
|
||||
self.stride2,
|
||||
self.padding2,
|
||||
self.dilation2,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
def forward_with_2d(self, x, skip_time_conv):
|
||||
b, c, d, h, w = x.shape
|
||||
|
||||
# First 2D convolution
|
||||
x = rearrange(x, "b c d h w -> (b d) c h w")
|
||||
# Squeeze the depth dimension out of weight1 since it's 1
|
||||
weight1 = self.weight1.squeeze(2)
|
||||
# Select stride, padding, and dilation for the 2D convolution
|
||||
stride1 = (self.stride1[1], self.stride1[2])
|
||||
padding1 = (self.padding1[1], self.padding1[2])
|
||||
dilation1 = (self.dilation1[1], self.dilation1[2])
|
||||
x = F.conv2d(
|
||||
x,
|
||||
weight1,
|
||||
self.bias1,
|
||||
stride1,
|
||||
padding1,
|
||||
dilation1,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
|
||||
_, _, h, w = x.shape
|
||||
|
||||
if skip_time_conv:
|
||||
x = rearrange(x, "(b d) c h w -> b c d h w", b=b)
|
||||
return x
|
||||
|
||||
# Second convolution which is essentially treated as a 1D convolution across the 'd' dimension
|
||||
x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b)
|
||||
|
||||
# Reshape weight2 to match the expected dimensions for conv1d
|
||||
weight2 = self.weight2.squeeze(-1).squeeze(-1)
|
||||
# Use only the relevant dimension for stride, padding, and dilation for the 1D convolution
|
||||
stride2 = self.stride2[0]
|
||||
padding2 = self.padding2[0]
|
||||
dilation2 = self.dilation2[0]
|
||||
x = F.conv1d(
|
||||
x,
|
||||
weight2,
|
||||
self.bias2,
|
||||
stride2,
|
||||
padding2,
|
||||
dilation2,
|
||||
self.groups,
|
||||
padding_mode=self.padding_mode,
|
||||
)
|
||||
x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w)
|
||||
|
||||
return x
|
||||
|
||||
@property
|
||||
def weight(self):
|
||||
return self.weight2
|
||||
|
||||
|
||||
def test_dual_conv3d_consistency():
|
||||
# Initialize parameters
|
||||
in_channels = 3
|
||||
out_channels = 5
|
||||
kernel_size = (3, 3, 3)
|
||||
stride = (2, 2, 2)
|
||||
padding = (1, 1, 1)
|
||||
|
||||
# Create an instance of the DualConv3d class
|
||||
dual_conv3d = DualConv3d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
# Example input tensor
|
||||
test_input = torch.randn(1, 3, 10, 10, 10)
|
||||
|
||||
# Perform forward passes with both 3D and 2D settings
|
||||
output_conv3d = dual_conv3d(test_input, use_conv3d=True)
|
||||
output_2d = dual_conv3d(test_input, use_conv3d=False)
|
||||
|
||||
# Assert that the outputs from both methods are sufficiently close
|
||||
assert torch.allclose(
|
||||
output_conv3d, output_2d, atol=1e-6
|
||||
), "Outputs are not consistent between 3D and 2D convolutions."
|
||||
203
ltx_video/models/autoencoders/latent_upsampler.py
Normal file
203
ltx_video/models/autoencoders/latent_upsampler.py
Normal file
@@ -0,0 +1,203 @@
|
||||
from typing import Optional, Union
|
||||
from pathlib import Path
|
||||
import os
|
||||
import json
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from diffusers import ConfigMixin, ModelMixin
|
||||
from safetensors.torch import safe_open
|
||||
|
||||
from ltx_video.models.autoencoders.pixel_shuffle import PixelShuffleND
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(
|
||||
self, channels: int, mid_channels: Optional[int] = None, dims: int = 3
|
||||
):
|
||||
super().__init__()
|
||||
if mid_channels is None:
|
||||
mid_channels = channels
|
||||
|
||||
Conv = nn.Conv2d if dims == 2 else nn.Conv3d
|
||||
|
||||
self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1)
|
||||
self.norm1 = nn.GroupNorm(32, mid_channels)
|
||||
self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1)
|
||||
self.norm2 = nn.GroupNorm(32, channels)
|
||||
self.activation = nn.SiLU()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.activation(x)
|
||||
x = self.conv2(x)
|
||||
x = self.norm2(x)
|
||||
x = self.activation(x + residual)
|
||||
return x
|
||||
|
||||
|
||||
class LatentUpsampler(ModelMixin, ConfigMixin):
|
||||
"""
|
||||
Model to spatially upsample VAE latents.
|
||||
|
||||
Args:
|
||||
in_channels (`int`): Number of channels in the input latent
|
||||
mid_channels (`int`): Number of channels in the middle layers
|
||||
num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)
|
||||
dims (`int`): Number of dimensions for convolutions (2 or 3)
|
||||
spatial_upsample (`bool`): Whether to spatially upsample the latent
|
||||
temporal_upsample (`bool`): Whether to temporally upsample the latent
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 128,
|
||||
mid_channels: int = 512,
|
||||
num_blocks_per_stage: int = 4,
|
||||
dims: int = 3,
|
||||
spatial_upsample: bool = True,
|
||||
temporal_upsample: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.mid_channels = mid_channels
|
||||
self.num_blocks_per_stage = num_blocks_per_stage
|
||||
self.dims = dims
|
||||
self.spatial_upsample = spatial_upsample
|
||||
self.temporal_upsample = temporal_upsample
|
||||
|
||||
Conv = nn.Conv2d if dims == 2 else nn.Conv3d
|
||||
|
||||
self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1)
|
||||
self.initial_norm = nn.GroupNorm(32, mid_channels)
|
||||
self.initial_activation = nn.SiLU()
|
||||
|
||||
self.res_blocks = nn.ModuleList(
|
||||
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
|
||||
)
|
||||
|
||||
if spatial_upsample and temporal_upsample:
|
||||
self.upsampler = nn.Sequential(
|
||||
nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
|
||||
PixelShuffleND(3),
|
||||
)
|
||||
elif spatial_upsample:
|
||||
self.upsampler = nn.Sequential(
|
||||
nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
|
||||
PixelShuffleND(2),
|
||||
)
|
||||
elif temporal_upsample:
|
||||
self.upsampler = nn.Sequential(
|
||||
nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
|
||||
PixelShuffleND(1),
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Either spatial_upsample or temporal_upsample must be True"
|
||||
)
|
||||
|
||||
self.post_upsample_res_blocks = nn.ModuleList(
|
||||
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
|
||||
)
|
||||
|
||||
self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, latent: torch.Tensor) -> torch.Tensor:
|
||||
b, c, f, h, w = latent.shape
|
||||
|
||||
if self.dims == 2:
|
||||
x = rearrange(latent, "b c f h w -> (b f) c h w")
|
||||
x = self.initial_conv(x)
|
||||
x = self.initial_norm(x)
|
||||
x = self.initial_activation(x)
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x)
|
||||
|
||||
x = self.upsampler(x)
|
||||
|
||||
for block in self.post_upsample_res_blocks:
|
||||
x = block(x)
|
||||
|
||||
x = self.final_conv(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
|
||||
else:
|
||||
x = self.initial_conv(latent)
|
||||
x = self.initial_norm(x)
|
||||
x = self.initial_activation(x)
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x)
|
||||
|
||||
if self.temporal_upsample:
|
||||
x = self.upsampler(x)
|
||||
x = x[:, :, 1:, :, :]
|
||||
else:
|
||||
x = rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = self.upsampler(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
|
||||
|
||||
for block in self.post_upsample_res_blocks:
|
||||
x = block(x)
|
||||
|
||||
x = self.final_conv(x)
|
||||
|
||||
return x
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config):
|
||||
return cls(
|
||||
in_channels=config.get("in_channels", 4),
|
||||
mid_channels=config.get("mid_channels", 128),
|
||||
num_blocks_per_stage=config.get("num_blocks_per_stage", 4),
|
||||
dims=config.get("dims", 2),
|
||||
spatial_upsample=config.get("spatial_upsample", True),
|
||||
temporal_upsample=config.get("temporal_upsample", False),
|
||||
)
|
||||
|
||||
def config(self):
|
||||
return {
|
||||
"_class_name": "LatentUpsampler",
|
||||
"in_channels": self.in_channels,
|
||||
"mid_channels": self.mid_channels,
|
||||
"num_blocks_per_stage": self.num_blocks_per_stage,
|
||||
"dims": self.dims,
|
||||
"spatial_upsample": self.spatial_upsample,
|
||||
"temporal_upsample": self.temporal_upsample,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
pretrained_model_path: Optional[Union[str, os.PathLike]],
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
pretrained_model_path = Path(pretrained_model_path)
|
||||
if pretrained_model_path.is_file() and str(pretrained_model_path).endswith(
|
||||
".safetensors"
|
||||
):
|
||||
state_dict = {}
|
||||
with safe_open(pretrained_model_path, framework="pt", device="cpu") as f:
|
||||
metadata = f.metadata()
|
||||
for k in f.keys():
|
||||
state_dict[k] = f.get_tensor(k)
|
||||
config = json.loads(metadata["config"])
|
||||
with torch.device("meta"):
|
||||
latent_upsampler = LatentUpsampler.from_config(config)
|
||||
latent_upsampler.load_state_dict(state_dict, assign=True)
|
||||
return latent_upsampler
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
latent_upsampler = LatentUpsampler(num_blocks_per_stage=4, dims=3)
|
||||
print(latent_upsampler)
|
||||
total_params = sum(p.numel() for p in latent_upsampler.parameters())
|
||||
print(f"Total number of parameters: {total_params:,}")
|
||||
latent = torch.randn(1, 128, 9, 16, 16)
|
||||
upsampled_latent = latent_upsampler(latent)
|
||||
print(f"Upsampled latent shape: {upsampled_latent.shape}")
|
||||
12
ltx_video/models/autoencoders/pixel_norm.py
Normal file
12
ltx_video/models/autoencoders/pixel_norm.py
Normal file
@@ -0,0 +1,12 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class PixelNorm(nn.Module):
|
||||
def __init__(self, dim=1, eps=1e-8):
|
||||
super(PixelNorm, self).__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x):
|
||||
return x / torch.sqrt(torch.mean(x**2, dim=self.dim, keepdim=True) + self.eps)
|
||||
33
ltx_video/models/autoencoders/pixel_shuffle.py
Normal file
33
ltx_video/models/autoencoders/pixel_shuffle.py
Normal file
@@ -0,0 +1,33 @@
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class PixelShuffleND(nn.Module):
|
||||
def __init__(self, dims, upscale_factors=(2, 2, 2)):
|
||||
super().__init__()
|
||||
assert dims in [1, 2, 3], "dims must be 1, 2, or 3"
|
||||
self.dims = dims
|
||||
self.upscale_factors = upscale_factors
|
||||
|
||||
def forward(self, x):
|
||||
if self.dims == 3:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
|
||||
p1=self.upscale_factors[0],
|
||||
p2=self.upscale_factors[1],
|
||||
p3=self.upscale_factors[2],
|
||||
)
|
||||
elif self.dims == 2:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1 p2) h w -> b c (h p1) (w p2)",
|
||||
p1=self.upscale_factors[0],
|
||||
p2=self.upscale_factors[1],
|
||||
)
|
||||
elif self.dims == 1:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (c p1) f h w -> b c (f p1) h w",
|
||||
p1=self.upscale_factors[0],
|
||||
)
|
||||
443
ltx_video/models/autoencoders/vae.py
Normal file
443
ltx_video/models/autoencoders/vae.py
Normal file
@@ -0,0 +1,443 @@
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import inspect
|
||||
import math
|
||||
import torch.nn as nn
|
||||
from diffusers import ConfigMixin, ModelMixin
|
||||
from diffusers.models.autoencoders.vae import (
|
||||
DecoderOutput,
|
||||
DiagonalGaussianDistribution,
|
||||
)
|
||||
from diffusers.models.modeling_outputs import AutoencoderKLOutput
|
||||
from ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd
|
||||
|
||||
|
||||
class AutoencoderKLWrapper(ModelMixin, ConfigMixin):
|
||||
"""Variational Autoencoder (VAE) model with KL loss.
|
||||
|
||||
VAE from the paper Auto-Encoding Variational Bayes by Diederik P. Kingma and Max Welling.
|
||||
This model is a wrapper around an encoder and a decoder, and it adds a KL loss term to the reconstruction loss.
|
||||
|
||||
Args:
|
||||
encoder (`nn.Module`):
|
||||
Encoder module.
|
||||
decoder (`nn.Module`):
|
||||
Decoder module.
|
||||
latent_channels (`int`, *optional*, defaults to 4):
|
||||
Number of latent channels.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
encoder: nn.Module,
|
||||
decoder: nn.Module,
|
||||
latent_channels: int = 4,
|
||||
dims: int = 2,
|
||||
sample_size=512,
|
||||
use_quant_conv: bool = True,
|
||||
normalize_latent_channels: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
|
||||
self.per_channel_statistics = nn.Module()
|
||||
std_of_means = torch.zeros( (128,), dtype= torch.bfloat16)
|
||||
|
||||
self.per_channel_statistics.register_buffer("std-of-means", std_of_means)
|
||||
self.per_channel_statistics.register_buffer(
|
||||
"mean-of-means",
|
||||
torch.zeros_like(std_of_means)
|
||||
)
|
||||
|
||||
|
||||
|
||||
# pass init params to Encoder
|
||||
self.encoder = encoder
|
||||
self.use_quant_conv = use_quant_conv
|
||||
self.normalize_latent_channels = normalize_latent_channels
|
||||
|
||||
# pass init params to Decoder
|
||||
quant_dims = 2 if dims == 2 else 3
|
||||
self.decoder = decoder
|
||||
if use_quant_conv:
|
||||
self.quant_conv = make_conv_nd(
|
||||
quant_dims, 2 * latent_channels, 2 * latent_channels, 1
|
||||
)
|
||||
self.post_quant_conv = make_conv_nd(
|
||||
quant_dims, latent_channels, latent_channels, 1
|
||||
)
|
||||
else:
|
||||
self.quant_conv = nn.Identity()
|
||||
self.post_quant_conv = nn.Identity()
|
||||
|
||||
if normalize_latent_channels:
|
||||
if dims == 2:
|
||||
self.latent_norm_out = nn.BatchNorm2d(latent_channels, affine=False)
|
||||
else:
|
||||
self.latent_norm_out = nn.BatchNorm3d(latent_channels, affine=False)
|
||||
else:
|
||||
self.latent_norm_out = nn.Identity()
|
||||
self.use_z_tiling = False
|
||||
self.use_hw_tiling = False
|
||||
self.dims = dims
|
||||
self.z_sample_size = 1
|
||||
|
||||
self.decoder_params = inspect.signature(self.decoder.forward).parameters
|
||||
|
||||
# only relevant if vae tiling is enabled
|
||||
self.set_tiling_params(sample_size=sample_size, overlap_factor=0.25)
|
||||
|
||||
@staticmethod
|
||||
def get_VAE_tile_size(vae_config, device_mem_capacity, mixed_precision):
|
||||
|
||||
z_tile = 4
|
||||
# VAE Tiling
|
||||
if vae_config == 0:
|
||||
if mixed_precision:
|
||||
device_mem_capacity = device_mem_capacity / 1.5
|
||||
if device_mem_capacity >= 24000:
|
||||
use_vae_config = 1
|
||||
elif device_mem_capacity >= 8000:
|
||||
use_vae_config = 2
|
||||
else:
|
||||
use_vae_config = 3
|
||||
else:
|
||||
use_vae_config = vae_config
|
||||
|
||||
if use_vae_config == 1:
|
||||
hw_tile = 0
|
||||
elif use_vae_config == 2:
|
||||
hw_tile = 512
|
||||
else:
|
||||
hw_tile = 256
|
||||
|
||||
return (z_tile, hw_tile)
|
||||
|
||||
def set_tiling_params(self, sample_size: int = 512, overlap_factor: float = 0.25):
|
||||
self.tile_sample_min_size = sample_size
|
||||
num_blocks = len(self.encoder.down_blocks)
|
||||
# self.tile_latent_min_size = int(sample_size / (2 ** (num_blocks - 1)))
|
||||
self.tile_latent_min_size = int(sample_size / 32)
|
||||
self.tile_overlap_factor = overlap_factor
|
||||
|
||||
def enable_z_tiling(self, z_sample_size: int = 4):
|
||||
r"""
|
||||
Enable tiling during VAE decoding.
|
||||
|
||||
When this option is enabled, the VAE will split the input tensor in tiles to compute decoding in several
|
||||
steps. This is useful to save some memory and allow larger batch sizes.
|
||||
"""
|
||||
self.use_z_tiling = z_sample_size > 1
|
||||
self.z_sample_size = z_sample_size
|
||||
assert (
|
||||
z_sample_size % 4 == 0 or z_sample_size == 1
|
||||
), f"z_sample_size must be a multiple of 4 or 1. Got {z_sample_size}."
|
||||
|
||||
def disable_z_tiling(self):
|
||||
r"""
|
||||
Disable tiling during VAE decoding. If `use_tiling` was previously invoked, this method will go back to computing
|
||||
decoding in one step.
|
||||
"""
|
||||
self.use_z_tiling = False
|
||||
|
||||
def enable_hw_tiling(self):
|
||||
r"""
|
||||
Enable tiling during VAE decoding along the height and width dimension.
|
||||
"""
|
||||
self.use_hw_tiling = True
|
||||
|
||||
def disable_hw_tiling(self):
|
||||
r"""
|
||||
Disable tiling during VAE decoding along the height and width dimension.
|
||||
"""
|
||||
self.use_hw_tiling = False
|
||||
|
||||
def _hw_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True):
|
||||
overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
|
||||
blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
|
||||
row_limit = self.tile_latent_min_size - blend_extent
|
||||
|
||||
# Split the image into 512x512 tiles and encode them separately.
|
||||
rows = []
|
||||
for i in range(0, x.shape[3], overlap_size):
|
||||
row = []
|
||||
for j in range(0, x.shape[4], overlap_size):
|
||||
tile = x[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
i : i + self.tile_sample_min_size,
|
||||
j : j + self.tile_sample_min_size,
|
||||
]
|
||||
tile = self.encoder(tile)
|
||||
tile = self.quant_conv(tile)
|
||||
row.append(tile)
|
||||
rows.append(row)
|
||||
result_rows = []
|
||||
for i, row in enumerate(rows):
|
||||
result_row = []
|
||||
for j, tile in enumerate(row):
|
||||
# blend the above tile and the left tile
|
||||
# to the current tile and add the current tile to the result row
|
||||
if i > 0:
|
||||
tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
|
||||
if j > 0:
|
||||
tile = self.blend_h(row[j - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :, :row_limit, :row_limit])
|
||||
result_rows.append(torch.cat(result_row, dim=4))
|
||||
|
||||
moments = torch.cat(result_rows, dim=3)
|
||||
return moments
|
||||
|
||||
def blend_z(
|
||||
self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
|
||||
) -> torch.Tensor:
|
||||
blend_extent = min(a.shape[2], b.shape[2], blend_extent)
|
||||
for z in range(blend_extent):
|
||||
b[:, :, z, :, :] = a[:, :, -blend_extent + z, :, :] * (
|
||||
1 - z / blend_extent
|
||||
) + b[:, :, z, :, :] * (z / blend_extent)
|
||||
return b
|
||||
|
||||
def blend_v(
|
||||
self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
|
||||
) -> torch.Tensor:
|
||||
blend_extent = min(a.shape[3], b.shape[3], blend_extent)
|
||||
for y in range(blend_extent):
|
||||
b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (
|
||||
1 - y / blend_extent
|
||||
) + b[:, :, :, y, :] * (y / blend_extent)
|
||||
return b
|
||||
|
||||
def blend_h(
|
||||
self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
|
||||
) -> torch.Tensor:
|
||||
blend_extent = min(a.shape[4], b.shape[4], blend_extent)
|
||||
for x in range(blend_extent):
|
||||
b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (
|
||||
1 - x / blend_extent
|
||||
) + b[:, :, :, :, x] * (x / blend_extent)
|
||||
return b
|
||||
|
||||
def _hw_tiled_decode(self, z: torch.FloatTensor, target_shape, timestep = None):
|
||||
overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
|
||||
blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
|
||||
row_limit = self.tile_sample_min_size - blend_extent
|
||||
tile_target_shape = (
|
||||
*target_shape[:3],
|
||||
self.tile_sample_min_size,
|
||||
self.tile_sample_min_size,
|
||||
)
|
||||
# Split z into overlapping 64x64 tiles and decode them separately.
|
||||
# The tiles have an overlap to avoid seams between tiles.
|
||||
rows = []
|
||||
for i in range(0, z.shape[3], overlap_size):
|
||||
row = []
|
||||
for j in range(0, z.shape[4], overlap_size):
|
||||
tile = z[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
i : i + self.tile_latent_min_size,
|
||||
j : j + self.tile_latent_min_size,
|
||||
]
|
||||
tile = self.post_quant_conv(tile)
|
||||
decoded = self.decoder(tile, target_shape=tile_target_shape, timestep = timestep)
|
||||
row.append(decoded)
|
||||
rows.append(row)
|
||||
result_rows = []
|
||||
for i, row in enumerate(rows):
|
||||
result_row = []
|
||||
for j, tile in enumerate(row):
|
||||
# blend the above tile and the left tile
|
||||
# to the current tile and add the current tile to the result row
|
||||
if i > 0:
|
||||
tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
|
||||
if j > 0:
|
||||
tile = self.blend_h(row[j - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :, :row_limit, :row_limit])
|
||||
result_rows.append(torch.cat(result_row, dim=4))
|
||||
|
||||
dec = torch.cat(result_rows, dim=3)
|
||||
return dec
|
||||
|
||||
def encode(
|
||||
self, z: torch.FloatTensor, return_dict: bool = True
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
if self.use_z_tiling and z.shape[2] > (self.z_sample_size + 1) > 1:
|
||||
tile_latent_min_tsize = self.z_sample_size
|
||||
tile_sample_min_tsize = tile_latent_min_tsize * 8
|
||||
tile_overlap_factor = 0.25
|
||||
|
||||
B, C, T, H, W = z.shape
|
||||
overlap_size = int(tile_sample_min_tsize * (1 - tile_overlap_factor))
|
||||
blend_extent = int(tile_latent_min_tsize * tile_overlap_factor)
|
||||
t_limit = tile_latent_min_tsize - blend_extent
|
||||
|
||||
row = []
|
||||
for i in range(0, T, overlap_size):
|
||||
tile = z[:, :, i: i + tile_sample_min_tsize + 1, :, :]
|
||||
if self.use_hw_tiling:
|
||||
tile = self._hw_tiled_encode(tile, return_dict)
|
||||
else:
|
||||
tile = self._encode(tile)
|
||||
if i > 0:
|
||||
tile = tile[:, :, 1:, :, :]
|
||||
row.append(tile)
|
||||
result_row = []
|
||||
for i, tile in enumerate(row):
|
||||
if i > 0:
|
||||
tile = self.blend_z(row[i - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :t_limit, :, :])
|
||||
else:
|
||||
result_row.append(tile[:, :, :t_limit + 1, :, :])
|
||||
|
||||
moments = torch.cat(result_row, dim=2)
|
||||
|
||||
|
||||
else:
|
||||
moments = (
|
||||
self._hw_tiled_encode(z, return_dict)
|
||||
if self.use_hw_tiling and z.shape[2] > 1
|
||||
else self._encode(z)
|
||||
)
|
||||
|
||||
posterior = DiagonalGaussianDistribution(moments)
|
||||
if not return_dict:
|
||||
return (posterior,)
|
||||
|
||||
return AutoencoderKLOutput(latent_dist=posterior)
|
||||
|
||||
def _normalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:
|
||||
if isinstance(self.latent_norm_out, nn.BatchNorm3d):
|
||||
_, c, _, _, _ = z.shape
|
||||
z = torch.cat(
|
||||
[
|
||||
self.latent_norm_out(z[:, : c // 2, :, :, :]),
|
||||
z[:, c // 2 :, :, :, :],
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
elif isinstance(self.latent_norm_out, nn.BatchNorm2d):
|
||||
raise NotImplementedError("BatchNorm2d not supported")
|
||||
return z
|
||||
|
||||
def _unnormalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:
|
||||
if isinstance(self.latent_norm_out, nn.BatchNorm3d):
|
||||
running_mean = self.latent_norm_out.running_mean.view(1, -1, 1, 1, 1)
|
||||
running_var = self.latent_norm_out.running_var.view(1, -1, 1, 1, 1)
|
||||
eps = self.latent_norm_out.eps
|
||||
|
||||
z = z * torch.sqrt(running_var + eps) + running_mean
|
||||
elif isinstance(self.latent_norm_out, nn.BatchNorm3d):
|
||||
raise NotImplementedError("BatchNorm2d not supported")
|
||||
return z
|
||||
|
||||
def _encode(self, x: torch.FloatTensor) -> AutoencoderKLOutput:
|
||||
h = self.encoder(x)
|
||||
moments = self.quant_conv(h)
|
||||
moments = self._normalize_latent_channels(moments)
|
||||
return moments
|
||||
|
||||
def _decode(
|
||||
self,
|
||||
z: torch.FloatTensor,
|
||||
target_shape=None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
z = self._unnormalize_latent_channels(z)
|
||||
z = self.post_quant_conv(z)
|
||||
if "timestep" in self.decoder_params:
|
||||
dec = self.decoder(z, target_shape=target_shape, timestep=timestep)
|
||||
else:
|
||||
dec = self.decoder(z, target_shape=target_shape)
|
||||
return dec
|
||||
|
||||
def decode(
|
||||
self,
|
||||
z: torch.FloatTensor,
|
||||
return_dict: bool = True,
|
||||
target_shape=None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
assert target_shape is not None, "target_shape must be provided for decoding"
|
||||
if self.use_z_tiling and z.shape[2] > (self.z_sample_size + 1) > 1:
|
||||
# Split z into overlapping tiles and decode them separately.
|
||||
tile_latent_min_tsize = self.z_sample_size
|
||||
tile_sample_min_tsize = tile_latent_min_tsize * 8
|
||||
tile_overlap_factor = 0.25
|
||||
|
||||
B, C, T, H, W = z.shape
|
||||
overlap_size = int(tile_latent_min_tsize * (1 - tile_overlap_factor))
|
||||
blend_extent = int(tile_sample_min_tsize * tile_overlap_factor)
|
||||
t_limit = tile_sample_min_tsize - blend_extent
|
||||
|
||||
row = []
|
||||
for i in range(0, T, overlap_size):
|
||||
tile = z[:, :, i: i + tile_latent_min_tsize + 1, :, :]
|
||||
target_shape_split = list(target_shape)
|
||||
target_shape_split[2] = tile.shape[2] * 8
|
||||
if self.use_hw_tiling:
|
||||
decoded = self._hw_tiled_decode(tile, target_shape, timestep)
|
||||
else:
|
||||
decoded = self._decode(tile, target_shape=target_shape, timestep=timestep)
|
||||
|
||||
if i > 0:
|
||||
decoded = decoded[:, :, 1:, :, :]
|
||||
row.append(decoded.to(torch.float16).cpu())
|
||||
decoded = None
|
||||
result_row = []
|
||||
for i, tile in enumerate(row):
|
||||
if i > 0:
|
||||
tile = self.blend_z(row[i - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :t_limit, :, :])
|
||||
else:
|
||||
result_row.append(tile[:, :, :t_limit + 1, :, :])
|
||||
|
||||
dec = torch.cat(result_row, dim=2)
|
||||
if not return_dict:
|
||||
return (dec,)
|
||||
|
||||
return DecoderOutput(sample=dec)
|
||||
else:
|
||||
decoded = (
|
||||
self._hw_tiled_decode(z, target_shape, timestep)
|
||||
if self.use_hw_tiling
|
||||
else self._decode(z, target_shape=target_shape, timestep=timestep)
|
||||
)
|
||||
|
||||
if not return_dict:
|
||||
return (decoded,)
|
||||
|
||||
return DecoderOutput(sample=decoded)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
sample_posterior: bool = False,
|
||||
return_dict: bool = True,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> Union[DecoderOutput, torch.FloatTensor]:
|
||||
r"""
|
||||
Args:
|
||||
sample (`torch.FloatTensor`): Input sample.
|
||||
sample_posterior (`bool`, *optional*, defaults to `False`):
|
||||
Whether to sample from the posterior.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether to return a [`DecoderOutput`] instead of a plain tuple.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
Generator used to sample from the posterior.
|
||||
"""
|
||||
x = sample
|
||||
posterior = self.encode(x).latent_dist
|
||||
if sample_posterior:
|
||||
z = posterior.sample(generator=generator)
|
||||
else:
|
||||
z = posterior.mode()
|
||||
dec = self.decode(z, target_shape=sample.shape).sample
|
||||
|
||||
if not return_dict:
|
||||
return (dec,)
|
||||
|
||||
return DecoderOutput(sample=dec)
|
||||
247
ltx_video/models/autoencoders/vae_encode.py
Normal file
247
ltx_video/models/autoencoders/vae_encode.py
Normal file
@@ -0,0 +1,247 @@
|
||||
from typing import Tuple
|
||||
import torch
|
||||
from diffusers import AutoencoderKL
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
from ltx_video.models.autoencoders.causal_video_autoencoder import (
|
||||
CausalVideoAutoencoder,
|
||||
)
|
||||
from ltx_video.models.autoencoders.video_autoencoder import (
|
||||
Downsample3D,
|
||||
VideoAutoencoder,
|
||||
)
|
||||
|
||||
try:
|
||||
import torch_xla.core.xla_model as xm
|
||||
except ImportError:
|
||||
xm = None
|
||||
|
||||
|
||||
def vae_encode(
|
||||
media_items: Tensor,
|
||||
vae: AutoencoderKL,
|
||||
split_size: int = 1,
|
||||
vae_per_channel_normalize=False,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Encodes media items (images or videos) into latent representations using a specified VAE model.
|
||||
The function supports processing batches of images or video frames and can handle the processing
|
||||
in smaller sub-batches if needed.
|
||||
|
||||
Args:
|
||||
media_items (Tensor): A torch Tensor containing the media items to encode. The expected
|
||||
shape is (batch_size, channels, height, width) for images or (batch_size, channels,
|
||||
frames, height, width) for videos.
|
||||
vae (AutoencoderKL): An instance of the `AutoencoderKL` class from the `diffusers` library,
|
||||
pre-configured and loaded with the appropriate model weights.
|
||||
split_size (int, optional): The number of sub-batches to split the input batch into for encoding.
|
||||
If set to more than 1, the input media items are processed in smaller batches according to
|
||||
this value. Defaults to 1, which processes all items in a single batch.
|
||||
|
||||
Returns:
|
||||
Tensor: A torch Tensor of the encoded latent representations. The shape of the tensor is adjusted
|
||||
to match the input shape, scaled by the model's configuration.
|
||||
|
||||
Examples:
|
||||
>>> import torch
|
||||
>>> from diffusers import AutoencoderKL
|
||||
>>> vae = AutoencoderKL.from_pretrained('your-model-name')
|
||||
>>> images = torch.rand(10, 3, 8 256, 256) # Example tensor with 10 videos of 8 frames.
|
||||
>>> latents = vae_encode(images, vae)
|
||||
>>> print(latents.shape) # Output shape will depend on the model's latent configuration.
|
||||
|
||||
Note:
|
||||
In case of a video, the function encodes the media item frame-by frame.
|
||||
"""
|
||||
is_video_shaped = media_items.dim() == 5
|
||||
batch_size, channels = media_items.shape[0:2]
|
||||
|
||||
if channels != 3:
|
||||
raise ValueError(f"Expects tensors with 3 channels, got {channels}.")
|
||||
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
media_items = rearrange(media_items, "b c n h w -> (b n) c h w")
|
||||
if split_size > 1:
|
||||
if len(media_items) % split_size != 0:
|
||||
raise ValueError(
|
||||
"Error: The batch size must be divisible by 'train.vae_bs_split"
|
||||
)
|
||||
encode_bs = len(media_items) // split_size
|
||||
# latents = [vae.encode(image_batch).latent_dist.sample() for image_batch in media_items.split(encode_bs)]
|
||||
latents = []
|
||||
if media_items.device.type == "xla":
|
||||
xm.mark_step()
|
||||
for image_batch in media_items.split(encode_bs):
|
||||
latents.append(vae.encode(image_batch).latent_dist.sample())
|
||||
if media_items.device.type == "xla":
|
||||
xm.mark_step()
|
||||
latents = torch.cat(latents, dim=0)
|
||||
else:
|
||||
latents = vae.encode(media_items).latent_dist.sample()
|
||||
|
||||
latents = normalize_latents(latents, vae, vae_per_channel_normalize)
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
latents = rearrange(latents, "(b n) c h w -> b c n h w", b=batch_size)
|
||||
return latents
|
||||
|
||||
|
||||
def vae_decode(
|
||||
latents: Tensor,
|
||||
vae: AutoencoderKL,
|
||||
is_video: bool = True,
|
||||
split_size: int = 1,
|
||||
vae_per_channel_normalize=False,
|
||||
timestep=None,
|
||||
) -> Tensor:
|
||||
is_video_shaped = latents.dim() == 5
|
||||
batch_size = latents.shape[0]
|
||||
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
latents = rearrange(latents, "b c n h w -> (b n) c h w")
|
||||
if split_size > 1:
|
||||
if len(latents) % split_size != 0:
|
||||
raise ValueError(
|
||||
"Error: The batch size must be divisible by 'train.vae_bs_split"
|
||||
)
|
||||
encode_bs = len(latents) // split_size
|
||||
image_batch = [
|
||||
_run_decoder(
|
||||
latent_batch, vae, is_video, vae_per_channel_normalize, timestep
|
||||
)
|
||||
for latent_batch in latents.split(encode_bs)
|
||||
]
|
||||
images = torch.cat(image_batch, dim=0)
|
||||
else:
|
||||
images = _run_decoder(
|
||||
latents, vae, is_video, vae_per_channel_normalize, timestep
|
||||
)
|
||||
|
||||
if is_video_shaped and not isinstance(
|
||||
vae, (VideoAutoencoder, CausalVideoAutoencoder)
|
||||
):
|
||||
images = rearrange(images, "(b n) c h w -> b c n h w", b=batch_size)
|
||||
return images
|
||||
|
||||
|
||||
def _run_decoder(
|
||||
latents: Tensor,
|
||||
vae: AutoencoderKL,
|
||||
is_video: bool,
|
||||
vae_per_channel_normalize=False,
|
||||
timestep=None,
|
||||
) -> Tensor:
|
||||
if isinstance(vae, (VideoAutoencoder, CausalVideoAutoencoder)):
|
||||
*_, fl, hl, wl = latents.shape
|
||||
temporal_scale, spatial_scale, _ = get_vae_size_scale_factor(vae)
|
||||
latents = latents.to(vae.dtype)
|
||||
vae_decode_kwargs = {}
|
||||
if timestep is not None:
|
||||
vae_decode_kwargs["timestep"] = timestep
|
||||
image = vae.decode(
|
||||
un_normalize_latents(latents, vae, vae_per_channel_normalize),
|
||||
return_dict=False,
|
||||
target_shape=(
|
||||
1,
|
||||
3,
|
||||
fl * temporal_scale if is_video else 1,
|
||||
hl * spatial_scale,
|
||||
wl * spatial_scale,
|
||||
),
|
||||
**vae_decode_kwargs,
|
||||
)[0]
|
||||
else:
|
||||
image = vae.decode(
|
||||
un_normalize_latents(latents, vae, vae_per_channel_normalize),
|
||||
return_dict=False,
|
||||
)[0]
|
||||
return image
|
||||
|
||||
|
||||
def get_vae_size_scale_factor(vae: AutoencoderKL) -> float:
|
||||
if isinstance(vae, CausalVideoAutoencoder):
|
||||
spatial = vae.spatial_downscale_factor
|
||||
temporal = vae.temporal_downscale_factor
|
||||
else:
|
||||
down_blocks = len(
|
||||
[
|
||||
block
|
||||
for block in vae.encoder.down_blocks
|
||||
if isinstance(block.downsample, Downsample3D)
|
||||
]
|
||||
)
|
||||
spatial = vae.config.patch_size * 2**down_blocks
|
||||
temporal = (
|
||||
vae.config.patch_size_t * 2**down_blocks
|
||||
if isinstance(vae, VideoAutoencoder)
|
||||
else 1
|
||||
)
|
||||
|
||||
return (temporal, spatial, spatial)
|
||||
|
||||
|
||||
def latent_to_pixel_coords(
|
||||
latent_coords: Tensor, vae: AutoencoderKL, causal_fix: bool = False
|
||||
) -> Tensor:
|
||||
"""
|
||||
Converts latent coordinates to pixel coordinates by scaling them according to the VAE's
|
||||
configuration.
|
||||
|
||||
Args:
|
||||
latent_coords (Tensor): A tensor of shape [batch_size, 3, num_latents]
|
||||
containing the latent corner coordinates of each token.
|
||||
vae (AutoencoderKL): The VAE model
|
||||
causal_fix (bool): Whether to take into account the different temporal scale
|
||||
of the first frame. Default = False for backwards compatibility.
|
||||
Returns:
|
||||
Tensor: A tensor of pixel coordinates corresponding to the input latent coordinates.
|
||||
"""
|
||||
|
||||
scale_factors = get_vae_size_scale_factor(vae)
|
||||
causal_fix = isinstance(vae, CausalVideoAutoencoder) and causal_fix
|
||||
pixel_coords = latent_to_pixel_coords_from_factors(
|
||||
latent_coords, scale_factors, causal_fix
|
||||
)
|
||||
return pixel_coords
|
||||
|
||||
|
||||
def latent_to_pixel_coords_from_factors(
|
||||
latent_coords: Tensor, scale_factors: Tuple, causal_fix: bool = False
|
||||
) -> Tensor:
|
||||
pixel_coords = (
|
||||
latent_coords
|
||||
* torch.tensor(scale_factors, device=latent_coords.device)[None, :, None]
|
||||
)
|
||||
if causal_fix:
|
||||
# Fix temporal scale for first frame to 1 due to causality
|
||||
pixel_coords[:, 0] = (pixel_coords[:, 0] + 1 - scale_factors[0]).clamp(min=0)
|
||||
return pixel_coords
|
||||
|
||||
|
||||
def normalize_latents(
|
||||
latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False
|
||||
) -> Tensor:
|
||||
return (
|
||||
(latents - vae.mean_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1))
|
||||
/ vae.std_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
|
||||
if vae_per_channel_normalize
|
||||
else latents * vae.config.scaling_factor
|
||||
)
|
||||
|
||||
|
||||
def un_normalize_latents(
|
||||
latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False
|
||||
) -> Tensor:
|
||||
return (
|
||||
latents * vae.std_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
|
||||
+ vae.mean_of_means.to(latents.dtype).to(latents.device).view(1, -1, 1, 1, 1)
|
||||
if vae_per_channel_normalize
|
||||
else latents / vae.config.scaling_factor
|
||||
)
|
||||
1045
ltx_video/models/autoencoders/video_autoencoder.py
Normal file
1045
ltx_video/models/autoencoders/video_autoencoder.py
Normal file
File diff suppressed because it is too large
Load Diff
0
ltx_video/models/transformers/__init__.py
Normal file
0
ltx_video/models/transformers/__init__.py
Normal file
1323
ltx_video/models/transformers/attention.py
Normal file
1323
ltx_video/models/transformers/attention.py
Normal file
File diff suppressed because it is too large
Load Diff
129
ltx_video/models/transformers/embeddings.py
Normal file
129
ltx_video/models/transformers/embeddings.py
Normal file
@@ -0,0 +1,129 @@
|
||||
# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/embeddings.py
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
|
||||
def get_timestep_embedding(
|
||||
timesteps: torch.Tensor,
|
||||
embedding_dim: int,
|
||||
flip_sin_to_cos: bool = False,
|
||||
downscale_freq_shift: float = 1,
|
||||
scale: float = 1,
|
||||
max_period: int = 10000,
|
||||
):
|
||||
"""
|
||||
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
|
||||
|
||||
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
||||
These may be fractional.
|
||||
:param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the
|
||||
embeddings. :return: an [N x dim] Tensor of positional embeddings.
|
||||
"""
|
||||
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
|
||||
|
||||
half_dim = embedding_dim // 2
|
||||
exponent = -math.log(max_period) * torch.arange(
|
||||
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
|
||||
)
|
||||
exponent = exponent / (half_dim - downscale_freq_shift)
|
||||
|
||||
emb = torch.exp(exponent)
|
||||
emb = timesteps[:, None].float() * emb[None, :]
|
||||
|
||||
# scale embeddings
|
||||
emb = scale * emb
|
||||
|
||||
# concat sine and cosine embeddings
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
|
||||
# flip sine and cosine embeddings
|
||||
if flip_sin_to_cos:
|
||||
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
|
||||
|
||||
# zero pad
|
||||
if embedding_dim % 2 == 1:
|
||||
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
|
||||
return emb
|
||||
|
||||
|
||||
def get_3d_sincos_pos_embed(embed_dim, grid, w, h, f):
|
||||
"""
|
||||
grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
|
||||
[1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
|
||||
"""
|
||||
grid = rearrange(grid, "c (f h w) -> c f h w", h=h, w=w)
|
||||
grid = rearrange(grid, "c f h w -> c h w f", h=h, w=w)
|
||||
grid = grid.reshape([3, 1, w, h, f])
|
||||
pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid)
|
||||
pos_embed = pos_embed.transpose(1, 0, 2, 3)
|
||||
return rearrange(pos_embed, "h w f c -> (f h w) c")
|
||||
|
||||
|
||||
def get_3d_sincos_pos_embed_from_grid(embed_dim, grid):
|
||||
if embed_dim % 3 != 0:
|
||||
raise ValueError("embed_dim must be divisible by 3")
|
||||
|
||||
# use half of dimensions to encode grid_h
|
||||
emb_f = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0]) # (H*W*T, D/3)
|
||||
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1]) # (H*W*T, D/3)
|
||||
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2]) # (H*W*T, D/3)
|
||||
|
||||
emb = np.concatenate([emb_h, emb_w, emb_f], axis=-1) # (H*W*T, D)
|
||||
return emb
|
||||
|
||||
|
||||
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
|
||||
"""
|
||||
embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
|
||||
"""
|
||||
if embed_dim % 2 != 0:
|
||||
raise ValueError("embed_dim must be divisible by 2")
|
||||
|
||||
omega = np.arange(embed_dim // 2, dtype=np.float64)
|
||||
omega /= embed_dim / 2.0
|
||||
omega = 1.0 / 10000**omega # (D/2,)
|
||||
|
||||
pos_shape = pos.shape
|
||||
|
||||
pos = pos.reshape(-1)
|
||||
out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
|
||||
out = out.reshape([*pos_shape, -1])[0]
|
||||
|
||||
emb_sin = np.sin(out) # (M, D/2)
|
||||
emb_cos = np.cos(out) # (M, D/2)
|
||||
|
||||
emb = np.concatenate([emb_sin, emb_cos], axis=-1) # (M, D)
|
||||
return emb
|
||||
|
||||
|
||||
class SinusoidalPositionalEmbedding(nn.Module):
|
||||
"""Apply positional information to a sequence of embeddings.
|
||||
|
||||
Takes in a sequence of embeddings with shape (batch_size, seq_length, embed_dim) and adds positional embeddings to
|
||||
them
|
||||
|
||||
Args:
|
||||
embed_dim: (int): Dimension of the positional embedding.
|
||||
max_seq_length: Maximum sequence length to apply positional embeddings
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, embed_dim: int, max_seq_length: int = 32):
|
||||
super().__init__()
|
||||
position = torch.arange(max_seq_length).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim)
|
||||
)
|
||||
pe = torch.zeros(1, max_seq_length, embed_dim)
|
||||
pe[0, :, 0::2] = torch.sin(position * div_term)
|
||||
pe[0, :, 1::2] = torch.cos(position * div_term)
|
||||
self.register_buffer("pe", pe)
|
||||
|
||||
def forward(self, x):
|
||||
_, seq_length, _ = x.shape
|
||||
x = x + self.pe[:, :seq_length]
|
||||
return x
|
||||
84
ltx_video/models/transformers/symmetric_patchifier.py
Normal file
84
ltx_video/models/transformers/symmetric_patchifier.py
Normal file
@@ -0,0 +1,84 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class Patchifier(ConfigMixin, ABC):
|
||||
def __init__(self, patch_size: int):
|
||||
super().__init__()
|
||||
self._patch_size = (1, patch_size, patch_size)
|
||||
|
||||
@abstractmethod
|
||||
def patchify(self, latents: Tensor) -> Tuple[Tensor, Tensor]:
|
||||
raise NotImplementedError("Patchify method not implemented")
|
||||
|
||||
@abstractmethod
|
||||
def unpatchify(
|
||||
self,
|
||||
latents: Tensor,
|
||||
output_height: int,
|
||||
output_width: int,
|
||||
out_channels: int,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
pass
|
||||
|
||||
@property
|
||||
def patch_size(self):
|
||||
return self._patch_size
|
||||
|
||||
def get_latent_coords(
|
||||
self, latent_num_frames, latent_height, latent_width, batch_size, device
|
||||
):
|
||||
"""
|
||||
Return a tensor of shape [batch_size, 3, num_patches] containing the
|
||||
top-left corner latent coordinates of each latent patch.
|
||||
The tensor is repeated for each batch element.
|
||||
"""
|
||||
latent_sample_coords = torch.meshgrid(
|
||||
torch.arange(0, latent_num_frames, self._patch_size[0], device=device),
|
||||
torch.arange(0, latent_height, self._patch_size[1], device=device),
|
||||
torch.arange(0, latent_width, self._patch_size[2], device=device),
|
||||
)
|
||||
latent_sample_coords = torch.stack(latent_sample_coords, dim=0)
|
||||
latent_coords = latent_sample_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1)
|
||||
latent_coords = rearrange(
|
||||
latent_coords, "b c f h w -> b c (f h w)", b=batch_size
|
||||
)
|
||||
return latent_coords
|
||||
|
||||
|
||||
class SymmetricPatchifier(Patchifier):
|
||||
def patchify(self, latents: Tensor) -> Tuple[Tensor, Tensor]:
|
||||
b, _, f, h, w = latents.shape
|
||||
latent_coords = self.get_latent_coords(f, h, w, b, latents.device)
|
||||
latents = rearrange(
|
||||
latents,
|
||||
"b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
|
||||
p1=self._patch_size[0],
|
||||
p2=self._patch_size[1],
|
||||
p3=self._patch_size[2],
|
||||
)
|
||||
return latents, latent_coords
|
||||
|
||||
def unpatchify(
|
||||
self,
|
||||
latents: Tensor,
|
||||
output_height: int,
|
||||
output_width: int,
|
||||
out_channels: int,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
output_height = output_height // self._patch_size[1]
|
||||
output_width = output_width // self._patch_size[2]
|
||||
latents = rearrange(
|
||||
latents,
|
||||
"b (f h w) (c p q) -> b c f (h p) (w q)",
|
||||
h=output_height,
|
||||
w=output_width,
|
||||
p=self._patch_size[1],
|
||||
q=self._patch_size[2],
|
||||
)
|
||||
return latents
|
||||
507
ltx_video/models/transformers/transformer3d.py
Normal file
507
ltx_video/models/transformers/transformer3d.py
Normal file
@@ -0,0 +1,507 @@
|
||||
# Adapted from: https://github.com/huggingface/diffusers/blob/v0.26.3/src/diffusers/models/transformers/transformer_2d.py
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import os
|
||||
import json
|
||||
import glob
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.embeddings import PixArtAlphaTextProjection
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.models.normalization import AdaLayerNormSingle
|
||||
from diffusers.utils import BaseOutput, is_torch_version
|
||||
from diffusers.utils import logging
|
||||
from torch import nn
|
||||
from safetensors import safe_open
|
||||
from ltx_video.models.transformers.attention import BasicTransformerBlock, reshape_hidden_states, restore_hidden_states_shape
|
||||
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
|
||||
|
||||
from ltx_video.utils.diffusers_config_mapping import (
|
||||
diffusers_and_ours_config_mapping,
|
||||
make_hashable_key,
|
||||
TRANSFORMER_KEYS_RENAME_DICT,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Transformer3DModelOutput(BaseOutput):
|
||||
"""
|
||||
The output of [`Transformer2DModel`].
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
|
||||
The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
|
||||
distributions for the unnoised latent pixels.
|
||||
"""
|
||||
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
class Transformer3DModel(ModelMixin, ConfigMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_attention_heads: int = 16,
|
||||
attention_head_dim: int = 88,
|
||||
in_channels: Optional[int] = None,
|
||||
out_channels: Optional[int] = None,
|
||||
num_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
norm_num_groups: int = 32,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
num_vector_embeds: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
use_linear_projection: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
double_self_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
adaptive_norm: str = "single_scale_shift", # 'single_scale_shift' or 'single_scale'
|
||||
standardization_norm: str = "layer_norm", # 'layer_norm' or 'rms_norm'
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_eps: float = 1e-5,
|
||||
attention_type: str = "default",
|
||||
caption_channels: int = None,
|
||||
use_tpu_flash_attention: bool = False, # if True uses the TPU attention offload ('flash attention')
|
||||
qk_norm: Optional[str] = None,
|
||||
positional_embedding_type: str = "rope",
|
||||
positional_embedding_theta: Optional[float] = None,
|
||||
positional_embedding_max_pos: Optional[List[int]] = None,
|
||||
timestep_scale_multiplier: Optional[float] = None,
|
||||
causal_temporal_positioning: bool = False, # For backward compatibility, will be deprecated
|
||||
):
|
||||
super().__init__()
|
||||
self.use_tpu_flash_attention = (
|
||||
use_tpu_flash_attention # FIXME: push config down to the attention modules
|
||||
)
|
||||
self.use_linear_projection = use_linear_projection
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
self.inner_dim = inner_dim
|
||||
self.patchify_proj = nn.Linear(in_channels, inner_dim, bias=True)
|
||||
self.positional_embedding_type = positional_embedding_type
|
||||
self.positional_embedding_theta = positional_embedding_theta
|
||||
self.positional_embedding_max_pos = positional_embedding_max_pos
|
||||
self.use_rope = self.positional_embedding_type == "rope"
|
||||
self.timestep_scale_multiplier = timestep_scale_multiplier
|
||||
|
||||
if self.positional_embedding_type == "absolute":
|
||||
raise ValueError("Absolute positional embedding is no longer supported")
|
||||
elif self.positional_embedding_type == "rope":
|
||||
if positional_embedding_theta is None:
|
||||
raise ValueError(
|
||||
"If `positional_embedding_type` type is rope, `positional_embedding_theta` must also be defined"
|
||||
)
|
||||
if positional_embedding_max_pos is None:
|
||||
raise ValueError(
|
||||
"If `positional_embedding_type` type is rope, `positional_embedding_max_pos` must also be defined"
|
||||
)
|
||||
|
||||
# 3. Define transformers blocks
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
BasicTransformerBlock(
|
||||
inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
dropout=dropout,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
num_embeds_ada_norm=num_embeds_ada_norm,
|
||||
attention_bias=attention_bias,
|
||||
only_cross_attention=only_cross_attention,
|
||||
double_self_attention=double_self_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
adaptive_norm=adaptive_norm,
|
||||
standardization_norm=standardization_norm,
|
||||
norm_elementwise_affine=norm_elementwise_affine,
|
||||
norm_eps=norm_eps,
|
||||
attention_type=attention_type,
|
||||
use_tpu_flash_attention=use_tpu_flash_attention,
|
||||
qk_norm=qk_norm,
|
||||
use_rope=self.use_rope,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 4. Define output layers
|
||||
self.out_channels = in_channels if out_channels is None else out_channels
|
||||
self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
|
||||
self.scale_shift_table = nn.Parameter(
|
||||
torch.randn(2, inner_dim) / inner_dim**0.5
|
||||
)
|
||||
self.proj_out = nn.Linear(inner_dim, self.out_channels)
|
||||
|
||||
self.adaln_single = AdaLayerNormSingle(
|
||||
inner_dim, use_additional_conditions=False
|
||||
)
|
||||
if adaptive_norm == "single_scale":
|
||||
self.adaln_single.linear = nn.Linear(inner_dim, 4 * inner_dim, bias=True)
|
||||
|
||||
self.caption_projection = None
|
||||
if caption_channels is not None:
|
||||
self.caption_projection = PixArtAlphaTextProjection(
|
||||
in_features=caption_channels, hidden_size=inner_dim
|
||||
)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def set_use_tpu_flash_attention(self):
|
||||
r"""
|
||||
Function sets the flag in this object and propagates down the children. The flag will enforce the usage of TPU
|
||||
attention kernel.
|
||||
"""
|
||||
logger.info("ENABLE TPU FLASH ATTENTION -> TRUE")
|
||||
self.use_tpu_flash_attention = True
|
||||
# push config down to the attention modules
|
||||
for block in self.transformer_blocks:
|
||||
block.set_use_tpu_flash_attention()
|
||||
|
||||
def create_skip_layer_mask(
|
||||
self,
|
||||
batch_size: int,
|
||||
num_conds: int,
|
||||
ptb_index: int,
|
||||
skip_block_list: Optional[List[int]] = None,
|
||||
):
|
||||
if skip_block_list is None or len(skip_block_list) == 0:
|
||||
return None
|
||||
num_layers = len(self.transformer_blocks)
|
||||
mask = torch.ones(
|
||||
(num_layers, batch_size * num_conds), device=self.device, dtype=self.dtype
|
||||
)
|
||||
for block_idx in skip_block_list:
|
||||
mask[block_idx, ptb_index::num_conds] = 0
|
||||
return mask
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def get_fractional_positions(self, indices_grid):
|
||||
fractional_positions = torch.stack(
|
||||
[
|
||||
indices_grid[:, i] / self.positional_embedding_max_pos[i]
|
||||
for i in range(3)
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
return fractional_positions
|
||||
|
||||
def precompute_freqs_cis(self, indices_grid, spacing="exp"):
|
||||
dtype = torch.float32 # We need full precision in the freqs_cis computation.
|
||||
dim = self.inner_dim
|
||||
theta = self.positional_embedding_theta
|
||||
|
||||
fractional_positions = self.get_fractional_positions(indices_grid)
|
||||
|
||||
start = 1
|
||||
end = theta
|
||||
device = fractional_positions.device
|
||||
if spacing == "exp":
|
||||
indices = theta ** (
|
||||
torch.linspace(
|
||||
math.log(start, theta),
|
||||
math.log(end, theta),
|
||||
dim // 6,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
)
|
||||
indices = indices.to(dtype=dtype)
|
||||
elif spacing == "exp_2":
|
||||
indices = 1.0 / theta ** (torch.arange(0, dim, 6, device=device) / dim)
|
||||
indices = indices.to(dtype=dtype)
|
||||
elif spacing == "linear":
|
||||
indices = torch.linspace(start, end, dim // 6, device=device, dtype=dtype)
|
||||
elif spacing == "sqrt":
|
||||
indices = torch.linspace(
|
||||
start**2, end**2, dim // 6, device=device, dtype=dtype
|
||||
).sqrt()
|
||||
|
||||
indices = indices * math.pi / 2
|
||||
|
||||
if spacing == "exp_2":
|
||||
freqs = (
|
||||
(indices * fractional_positions.unsqueeze(-1))
|
||||
.transpose(-1, -2)
|
||||
.flatten(2)
|
||||
)
|
||||
else:
|
||||
freqs = (
|
||||
(indices * (fractional_positions.unsqueeze(-1) * 2 - 1))
|
||||
.transpose(-1, -2)
|
||||
.flatten(2)
|
||||
)
|
||||
|
||||
cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
|
||||
sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
|
||||
if dim % 6 != 0:
|
||||
cos_padding = torch.ones_like(cos_freq[:, :, : dim % 6])
|
||||
sin_padding = torch.zeros_like(cos_freq[:, :, : dim % 6])
|
||||
cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
|
||||
sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
|
||||
return cos_freq.to(self.dtype), sin_freq.to(self.dtype)
|
||||
|
||||
def load_state_dict(
|
||||
self,
|
||||
state_dict: Dict,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
if any([key.startswith("model.diffusion_model.") for key in state_dict.keys()]):
|
||||
state_dict = {
|
||||
key.replace("model.diffusion_model.", ""): value
|
||||
for key, value in state_dict.items()
|
||||
if key.startswith("model.diffusion_model.")
|
||||
}
|
||||
return super().load_state_dict(state_dict, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
pretrained_model_path: Optional[Union[str, os.PathLike]],
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
pretrained_model_path = Path(pretrained_model_path)
|
||||
if pretrained_model_path.is_dir():
|
||||
config_path = pretrained_model_path / "transformer" / "config.json"
|
||||
with open(config_path, "r") as f:
|
||||
config = make_hashable_key(json.load(f))
|
||||
|
||||
assert config in diffusers_and_ours_config_mapping, (
|
||||
"Provided diffusers checkpoint config for transformer is not suppported. "
|
||||
"We only support diffusers configs found in Lightricks/LTX-Video."
|
||||
)
|
||||
|
||||
config = diffusers_and_ours_config_mapping[config]
|
||||
state_dict = {}
|
||||
ckpt_paths = (
|
||||
pretrained_model_path
|
||||
/ "transformer"
|
||||
/ "diffusion_pytorch_model*.safetensors"
|
||||
)
|
||||
dict_list = glob.glob(str(ckpt_paths))
|
||||
for dict_path in dict_list:
|
||||
part_dict = {}
|
||||
with safe_open(dict_path, framework="pt", device="cpu") as f:
|
||||
for k in f.keys():
|
||||
part_dict[k] = f.get_tensor(k)
|
||||
state_dict.update(part_dict)
|
||||
|
||||
for key in list(state_dict.keys()):
|
||||
new_key = key
|
||||
for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
|
||||
new_key = new_key.replace(replace_key, rename_key)
|
||||
state_dict[new_key] = state_dict.pop(key)
|
||||
|
||||
with torch.device("meta"):
|
||||
transformer = cls.from_config(config)
|
||||
transformer.load_state_dict(state_dict, assign=True, strict=True)
|
||||
elif pretrained_model_path.is_file() and str(pretrained_model_path).endswith(
|
||||
".safetensors"
|
||||
):
|
||||
comfy_single_file_state_dict = {}
|
||||
with safe_open(pretrained_model_path, framework="pt", device="cpu") as f:
|
||||
metadata = f.metadata()
|
||||
for k in f.keys():
|
||||
comfy_single_file_state_dict[k] = f.get_tensor(k)
|
||||
configs = json.loads(metadata["config"])
|
||||
transformer_config = configs["transformer"]
|
||||
with torch.device("meta"):
|
||||
transformer = Transformer3DModel.from_config(transformer_config)
|
||||
transformer.load_state_dict(comfy_single_file_state_dict, assign=True)
|
||||
return transformer
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
freqs_cis: list,
|
||||
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||||
skip_layer_mask: Optional[torch.Tensor] = None,
|
||||
skip_layer_strategy: Optional[SkipLayerStrategy] = None,
|
||||
latent_shape = None,
|
||||
joint_pass = True,
|
||||
ltxv_model = None,
|
||||
mixed = False,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
"""
|
||||
The [`Transformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
|
||||
Input `hidden_states`.
|
||||
indices_grid (`torch.LongTensor` of shape `(batch size, 3, num latent pixels)`):
|
||||
encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
|
||||
Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
|
||||
self-attention.
|
||||
timestep ( `torch.LongTensor`, *optional*):
|
||||
Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
|
||||
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
|
||||
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
|
||||
`AdaLayerZeroNorm`.
|
||||
cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
attention_mask ( `torch.Tensor`, *optional*):
|
||||
An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
|
||||
is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
|
||||
negative values to the attention scores corresponding to "discard" tokens.
|
||||
encoder_attention_mask ( `torch.Tensor`, *optional*):
|
||||
Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
|
||||
|
||||
* Mask `(batch, sequence_length)` True = keep, False = discard.
|
||||
* Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
|
||||
|
||||
If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
|
||||
above. This bias will be added to the cross-attention scores.
|
||||
skip_layer_mask ( `torch.Tensor`, *optional*):
|
||||
A mask of shape `(num_layers, batch)` that indicates which layers to skip. `0` at position
|
||||
`layer, batch_idx` indicates that the layer should be skipped for the corresponding batch index.
|
||||
skip_layer_strategy ( `SkipLayerStrategy`, *optional*, defaults to `None`):
|
||||
Controls which layers are skipped when calculating a perturbed latent for spatiotemporal guidance.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
# for tpu attention offload 2d token masks are used. No need to transform.
|
||||
if not self.use_tpu_flash_attention:
|
||||
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
|
||||
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
|
||||
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
|
||||
# expects mask of shape:
|
||||
# [batch, key_tokens]
|
||||
# adds singleton query_tokens dimension:
|
||||
# [batch, 1, key_tokens]
|
||||
# this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
|
||||
# [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
|
||||
# [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
|
||||
if attention_mask is not None and attention_mask.ndim == 2:
|
||||
# assume that mask is expressed as:
|
||||
# (1 = keep, 0 = discard)
|
||||
# convert mask into a bias that can be added to attention scores:
|
||||
# (keep = +0, discard = -10000.0)
|
||||
attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# convert encoder_attention_mask to a bias the same way we do for attention_mask
|
||||
if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
|
||||
encoder_attention_mask = (
|
||||
1 - encoder_attention_mask.to(hidden_states.dtype)
|
||||
) * -10000.0
|
||||
encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
|
||||
|
||||
# 1. Input
|
||||
hidden_states = self.patchify_proj(hidden_states)
|
||||
|
||||
if self.timestep_scale_multiplier:
|
||||
timestep = self.timestep_scale_multiplier * timestep
|
||||
|
||||
if timestep.shape[-1] > 1:
|
||||
timestep = timestep.reshape(timestep.shape[0], -1, latent_shape[-2] * latent_shape[-1] )
|
||||
timestep = timestep[:, :, 0]
|
||||
|
||||
batch_size = hidden_states.shape[0]
|
||||
timestep, embedded_timestep = self.adaln_single(
|
||||
timestep.flatten(),
|
||||
{"resolution": None, "aspect_ratio": None},
|
||||
batch_size=batch_size,
|
||||
hidden_dtype=hidden_states.dtype,
|
||||
)
|
||||
# Second dimension is 1 or number of tokens (if timestep_per_token)
|
||||
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
|
||||
embedded_timestep = embedded_timestep.view(
|
||||
batch_size, -1, embedded_timestep.shape[-1]
|
||||
)
|
||||
if mixed:
|
||||
timestep = timestep.float()
|
||||
embedded_timestep = embedded_timestep.float()
|
||||
hidden_states = hidden_states.float()
|
||||
|
||||
|
||||
# 2. Blocks
|
||||
if self.caption_projection is not None:
|
||||
batch_size = hidden_states.shape[0]
|
||||
encoder_hidden_states = self.caption_projection(encoder_hidden_states)
|
||||
encoder_hidden_states = encoder_hidden_states.view(
|
||||
batch_size, -1, hidden_states.shape[-1]
|
||||
)
|
||||
|
||||
|
||||
if joint_pass:
|
||||
for block_idx, block in enumerate(self.transformer_blocks):
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
freqs_cis=freqs_cis,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
timestep=timestep,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
class_labels=class_labels,
|
||||
skip_layer_mask= None if skip_layer_mask is None else skip_layer_mask[block_idx],
|
||||
skip_layer_strategy=skip_layer_strategy,
|
||||
)
|
||||
if ltxv_model._interrupt:
|
||||
return [None]
|
||||
|
||||
else:
|
||||
for block_idx, block in enumerate(self.transformer_blocks):
|
||||
for i, (one_hidden_states, one_encoder_hidden_states, one_encoder_attention_mask,one_timestep) in enumerate(zip(hidden_states, encoder_hidden_states,encoder_attention_mask,timestep)):
|
||||
hidden_states[i][...] = block(
|
||||
one_hidden_states.unsqueeze(0),
|
||||
freqs_cis=freqs_cis,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=one_encoder_hidden_states.unsqueeze(0),
|
||||
encoder_attention_mask=one_encoder_attention_mask.unsqueeze(0),
|
||||
timestep=one_timestep.unsqueeze(0),
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
class_labels=class_labels,
|
||||
skip_layer_mask= None if skip_layer_mask is None else skip_layer_mask[block_idx, i],
|
||||
skip_layer_strategy=skip_layer_strategy,
|
||||
)
|
||||
if ltxv_model._interrupt:
|
||||
return [None]
|
||||
|
||||
# 3. Output
|
||||
scale_shift_values = (
|
||||
self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
|
||||
)
|
||||
shift, scale = scale_shift_values[:, :, 0].unsqueeze(-2), scale_shift_values[:, :, 1].unsqueeze(-2)
|
||||
hidden_states = self.norm_out(hidden_states)
|
||||
# Modulation
|
||||
|
||||
|
||||
hidden_states = reshape_hidden_states(hidden_states, scale.shape[1])
|
||||
# hidden_states = hidden_states * (1 + scale)
|
||||
hidden_states *= 1 + scale
|
||||
hidden_states += shift
|
||||
hidden_states = restore_hidden_states_shape(hidden_states)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
if not return_dict:
|
||||
return (hidden_states,)
|
||||
|
||||
return Transformer3DModelOutput(sample=hidden_states)
|
||||
0
ltx_video/pipelines/__init__.py
Normal file
0
ltx_video/pipelines/__init__.py
Normal file
50
ltx_video/pipelines/crf_compressor.py
Normal file
50
ltx_video/pipelines/crf_compressor.py
Normal file
@@ -0,0 +1,50 @@
|
||||
import av
|
||||
import torch
|
||||
import io
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _encode_single_frame(output_file, image_array: np.ndarray, crf):
|
||||
container = av.open(output_file, "w", format="mp4")
|
||||
try:
|
||||
stream = container.add_stream(
|
||||
"libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"}
|
||||
)
|
||||
stream.height = image_array.shape[0]
|
||||
stream.width = image_array.shape[1]
|
||||
av_frame = av.VideoFrame.from_ndarray(image_array, format="rgb24").reformat(
|
||||
format="yuv420p"
|
||||
)
|
||||
container.mux(stream.encode(av_frame))
|
||||
container.mux(stream.encode())
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
def _decode_single_frame(video_file):
|
||||
container = av.open(video_file)
|
||||
try:
|
||||
stream = next(s for s in container.streams if s.type == "video")
|
||||
frame = next(container.decode(stream))
|
||||
finally:
|
||||
container.close()
|
||||
return frame.to_ndarray(format="rgb24")
|
||||
|
||||
|
||||
def compress(image: torch.Tensor, crf=29):
|
||||
if crf == 0:
|
||||
return image
|
||||
|
||||
image_array = (
|
||||
(image[: (image.shape[0] // 2) * 2, : (image.shape[1] // 2) * 2] * 255.0)
|
||||
.byte()
|
||||
.cpu()
|
||||
.numpy()
|
||||
)
|
||||
with io.BytesIO() as output_file:
|
||||
_encode_single_frame(output_file, image_array, crf)
|
||||
video_bytes = output_file.getvalue()
|
||||
with io.BytesIO(video_bytes) as video_file:
|
||||
image_array = _decode_single_frame(video_file)
|
||||
tensor = torch.tensor(image_array, dtype=image.dtype, device=image.device) / 255.0
|
||||
return tensor
|
||||
1903
ltx_video/pipelines/pipeline_ltx_video.py
Normal file
1903
ltx_video/pipelines/pipeline_ltx_video.py
Normal file
File diff suppressed because it is too large
Load Diff
0
ltx_video/schedulers/__init__.py
Normal file
0
ltx_video/schedulers/__init__.py
Normal file
392
ltx_video/schedulers/rf.py
Normal file
392
ltx_video/schedulers/rf.py
Normal file
@@ -0,0 +1,392 @@
|
||||
import math
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Optional, Tuple, Union
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
||||
from diffusers.utils import BaseOutput
|
||||
from torch import Tensor
|
||||
from safetensors import safe_open
|
||||
|
||||
|
||||
from ltx_video.utils.torch_utils import append_dims
|
||||
|
||||
from ltx_video.utils.diffusers_config_mapping import (
|
||||
diffusers_and_ours_config_mapping,
|
||||
make_hashable_key,
|
||||
)
|
||||
|
||||
|
||||
def linear_quadratic_schedule(num_steps, threshold_noise=0.025, linear_steps=None):
|
||||
if num_steps == 1:
|
||||
return torch.tensor([1.0])
|
||||
if linear_steps is None:
|
||||
linear_steps = num_steps // 2
|
||||
linear_sigma_schedule = [
|
||||
i * threshold_noise / linear_steps for i in range(linear_steps)
|
||||
]
|
||||
threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
|
||||
quadratic_steps = num_steps - linear_steps
|
||||
quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
|
||||
linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (
|
||||
quadratic_steps**2
|
||||
)
|
||||
const = quadratic_coef * (linear_steps**2)
|
||||
quadratic_sigma_schedule = [
|
||||
quadratic_coef * (i**2) + linear_coef * i + const
|
||||
for i in range(linear_steps, num_steps)
|
||||
]
|
||||
sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
|
||||
sigma_schedule = [1.0 - x for x in sigma_schedule]
|
||||
return torch.tensor(sigma_schedule[:-1])
|
||||
|
||||
|
||||
def simple_diffusion_resolution_dependent_timestep_shift(
|
||||
samples_shape: torch.Size,
|
||||
timesteps: Tensor,
|
||||
n: int = 32 * 32,
|
||||
) -> Tensor:
|
||||
if len(samples_shape) == 3:
|
||||
_, m, _ = samples_shape
|
||||
elif len(samples_shape) in [4, 5]:
|
||||
m = math.prod(samples_shape[2:])
|
||||
else:
|
||||
raise ValueError(
|
||||
"Samples must have shape (b, t, c), (b, c, h, w) or (b, c, f, h, w)"
|
||||
)
|
||||
snr = (timesteps / (1 - timesteps)) ** 2
|
||||
shift_snr = torch.log(snr) + 2 * math.log(m / n)
|
||||
shifted_timesteps = torch.sigmoid(0.5 * shift_snr)
|
||||
|
||||
return shifted_timesteps
|
||||
|
||||
|
||||
def time_shift(mu: float, sigma: float, t: Tensor):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
|
||||
def get_normal_shift(
|
||||
n_tokens: int,
|
||||
min_tokens: int = 1024,
|
||||
max_tokens: int = 4096,
|
||||
min_shift: float = 0.95,
|
||||
max_shift: float = 2.05,
|
||||
) -> Callable[[float], float]:
|
||||
m = (max_shift - min_shift) / (max_tokens - min_tokens)
|
||||
b = min_shift - m * min_tokens
|
||||
return m * n_tokens + b
|
||||
|
||||
|
||||
def strech_shifts_to_terminal(shifts: Tensor, terminal=0.1):
|
||||
"""
|
||||
Stretch a function (given as sampled shifts) so that its final value matches the given terminal value
|
||||
using the provided formula.
|
||||
|
||||
Parameters:
|
||||
- shifts (Tensor): The samples of the function to be stretched (PyTorch Tensor).
|
||||
- terminal (float): The desired terminal value (value at the last sample).
|
||||
|
||||
Returns:
|
||||
- Tensor: The stretched shifts such that the final value equals `terminal`.
|
||||
"""
|
||||
if shifts.numel() == 0:
|
||||
raise ValueError("The 'shifts' tensor must not be empty.")
|
||||
|
||||
# Ensure terminal value is valid
|
||||
if terminal <= 0 or terminal >= 1:
|
||||
raise ValueError("The terminal value must be between 0 and 1 (exclusive).")
|
||||
|
||||
# Transform the shifts using the given formula
|
||||
one_minus_z = 1 - shifts
|
||||
scale_factor = one_minus_z[-1] / (1 - terminal)
|
||||
stretched_shifts = 1 - (one_minus_z / scale_factor)
|
||||
|
||||
return stretched_shifts
|
||||
|
||||
|
||||
def sd3_resolution_dependent_timestep_shift(
|
||||
samples_shape: torch.Size,
|
||||
timesteps: Tensor,
|
||||
target_shift_terminal: Optional[float] = None,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Shifts the timestep schedule as a function of the generated resolution.
|
||||
|
||||
In the SD3 paper, the authors empirically how to shift the timesteps based on the resolution of the target images.
|
||||
For more details: https://arxiv.org/pdf/2403.03206
|
||||
|
||||
In Flux they later propose a more dynamic resolution dependent timestep shift, see:
|
||||
https://github.com/black-forest-labs/flux/blob/87f6fff727a377ea1c378af692afb41ae84cbe04/src/flux/sampling.py#L66
|
||||
|
||||
|
||||
Args:
|
||||
samples_shape (torch.Size): The samples batch shape (batch_size, channels, height, width) or
|
||||
(batch_size, channels, frame, height, width).
|
||||
timesteps (Tensor): A batch of timesteps with shape (batch_size,).
|
||||
target_shift_terminal (float): The target terminal value for the shifted timesteps.
|
||||
|
||||
Returns:
|
||||
Tensor: The shifted timesteps.
|
||||
"""
|
||||
if len(samples_shape) == 3:
|
||||
_, m, _ = samples_shape
|
||||
elif len(samples_shape) in [4, 5]:
|
||||
m = math.prod(samples_shape[2:])
|
||||
else:
|
||||
raise ValueError(
|
||||
"Samples must have shape (b, t, c), (b, c, h, w) or (b, c, f, h, w)"
|
||||
)
|
||||
|
||||
shift = get_normal_shift(m)
|
||||
time_shifts = time_shift(shift, 1, timesteps)
|
||||
if target_shift_terminal is not None: # Stretch the shifts to the target terminal
|
||||
time_shifts = strech_shifts_to_terminal(time_shifts, target_shift_terminal)
|
||||
return time_shifts
|
||||
|
||||
|
||||
class TimestepShifter(ABC):
|
||||
@abstractmethod
|
||||
def shift_timesteps(self, samples_shape: torch.Size, timesteps: Tensor) -> Tensor:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class RectifiedFlowSchedulerOutput(BaseOutput):
|
||||
"""
|
||||
Output class for the scheduler's step function output.
|
||||
|
||||
Args:
|
||||
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
|
||||
denoising loop.
|
||||
pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
The predicted denoised sample (x_{0}) based on the model output from the current timestep.
|
||||
`pred_original_sample` can be used to preview progress or for guidance.
|
||||
"""
|
||||
|
||||
prev_sample: torch.FloatTensor
|
||||
pred_original_sample: Optional[torch.FloatTensor] = None
|
||||
|
||||
|
||||
class RectifiedFlowScheduler(SchedulerMixin, ConfigMixin, TimestepShifter):
|
||||
order = 1
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_train_timesteps=1000,
|
||||
shifting: Optional[str] = None,
|
||||
base_resolution: int = 32**2,
|
||||
target_shift_terminal: Optional[float] = None,
|
||||
sampler: Optional[str] = "Uniform",
|
||||
shift: Optional[float] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.init_noise_sigma = 1.0
|
||||
self.num_inference_steps = None
|
||||
self.sampler = sampler
|
||||
self.shifting = shifting
|
||||
self.base_resolution = base_resolution
|
||||
self.target_shift_terminal = target_shift_terminal
|
||||
self.timesteps = self.sigmas = self.get_initial_timesteps(
|
||||
num_train_timesteps, shift=shift
|
||||
)
|
||||
self.shift = shift
|
||||
|
||||
def get_initial_timesteps(
|
||||
self, num_timesteps: int, shift: Optional[float] = None
|
||||
) -> Tensor:
|
||||
if self.sampler == "Uniform":
|
||||
return torch.linspace(1, 1 / num_timesteps, num_timesteps)
|
||||
elif self.sampler == "LinearQuadratic":
|
||||
return linear_quadratic_schedule(num_timesteps)
|
||||
elif self.sampler == "Constant":
|
||||
assert (
|
||||
shift is not None
|
||||
), "Shift must be provided for constant time shift sampler."
|
||||
return time_shift(
|
||||
shift, 1, torch.linspace(1, 1 / num_timesteps, num_timesteps)
|
||||
)
|
||||
|
||||
def shift_timesteps(self, samples_shape: torch.Size, timesteps: Tensor) -> Tensor:
|
||||
if self.shifting == "SD3":
|
||||
return sd3_resolution_dependent_timestep_shift(
|
||||
samples_shape, timesteps, self.target_shift_terminal
|
||||
)
|
||||
elif self.shifting == "SimpleDiffusion":
|
||||
return simple_diffusion_resolution_dependent_timestep_shift(
|
||||
samples_shape, timesteps, self.base_resolution
|
||||
)
|
||||
return timesteps
|
||||
|
||||
def set_timesteps(
|
||||
self,
|
||||
num_inference_steps: Optional[int] = None,
|
||||
samples_shape: Optional[torch.Size] = None,
|
||||
timesteps: Optional[Tensor] = None,
|
||||
device: Union[str, torch.device] = None,
|
||||
):
|
||||
"""
|
||||
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
|
||||
If `timesteps` are provided, they will be used instead of the scheduled timesteps.
|
||||
|
||||
Args:
|
||||
num_inference_steps (`int` *optional*): The number of diffusion steps used when generating samples.
|
||||
samples_shape (`torch.Size` *optional*): The samples batch shape, used for shifting.
|
||||
timesteps ('torch.Tensor' *optional*): Specific timesteps to use instead of scheduled timesteps.
|
||||
device (`Union[str, torch.device]`, *optional*): The device to which the timesteps tensor will be moved.
|
||||
"""
|
||||
if timesteps is not None and num_inference_steps is not None:
|
||||
raise ValueError(
|
||||
"You cannot provide both `timesteps` and `num_inference_steps`."
|
||||
)
|
||||
if timesteps is None:
|
||||
num_inference_steps = min(
|
||||
self.config.num_train_timesteps, num_inference_steps
|
||||
)
|
||||
timesteps = self.get_initial_timesteps(
|
||||
num_inference_steps, shift=self.shift
|
||||
).to(device)
|
||||
timesteps = self.shift_timesteps(samples_shape, timesteps)
|
||||
else:
|
||||
timesteps = torch.Tensor(timesteps).to(device)
|
||||
num_inference_steps = len(timesteps)
|
||||
self.timesteps = timesteps
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.sigmas = self.timesteps
|
||||
|
||||
@staticmethod
|
||||
def from_pretrained(pretrained_model_path: Union[str, os.PathLike]):
|
||||
with open(pretrained_model_path, "r", encoding="utf-8") as reader:
|
||||
text = reader.read()
|
||||
|
||||
config = json.loads(text)
|
||||
return RectifiedFlowScheduler.from_config(config)
|
||||
|
||||
pretrained_model_path = Path(pretrained_model_path)
|
||||
if pretrained_model_path.is_file():
|
||||
comfy_single_file_state_dict = {}
|
||||
with safe_open(pretrained_model_path, framework="pt", device="cpu") as f:
|
||||
metadata = f.metadata()
|
||||
for k in f.keys():
|
||||
comfy_single_file_state_dict[k] = f.get_tensor(k)
|
||||
configs = json.loads(metadata["config"])
|
||||
config = configs["scheduler"]
|
||||
del comfy_single_file_state_dict
|
||||
|
||||
elif pretrained_model_path.is_dir():
|
||||
diffusers_noise_scheduler_config_path = (
|
||||
pretrained_model_path / "scheduler" / "scheduler_config.json"
|
||||
)
|
||||
|
||||
with open(diffusers_noise_scheduler_config_path, "r") as f:
|
||||
scheduler_config = json.load(f)
|
||||
hashable_config = make_hashable_key(scheduler_config)
|
||||
if hashable_config in diffusers_and_ours_config_mapping:
|
||||
config = diffusers_and_ours_config_mapping[hashable_config]
|
||||
return RectifiedFlowScheduler.from_config(config)
|
||||
|
||||
def scale_model_input(
|
||||
self, sample: torch.FloatTensor, timestep: Optional[int] = None
|
||||
) -> torch.FloatTensor:
|
||||
# pylint: disable=unused-argument
|
||||
"""
|
||||
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
|
||||
current timestep.
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor`): input sample
|
||||
timestep (`int`, optional): current timestep
|
||||
|
||||
Returns:
|
||||
`torch.FloatTensor`: scaled input sample
|
||||
"""
|
||||
return sample
|
||||
|
||||
def step(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: torch.FloatTensor,
|
||||
sample: torch.FloatTensor,
|
||||
return_dict: bool = True,
|
||||
stochastic_sampling: Optional[bool] = False,
|
||||
**kwargs,
|
||||
) -> Union[RectifiedFlowSchedulerOutput, Tuple]:
|
||||
"""
|
||||
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
||||
process from the learned model outputs (most often the predicted noise).
|
||||
z_{t_1} = z_t - \Delta_t * v
|
||||
The method finds the next timestep that is lower than the input timestep(s) and denoises the latents
|
||||
to that level. The input timestep(s) are not required to be one of the predefined timesteps.
|
||||
|
||||
Args:
|
||||
model_output (`torch.FloatTensor`):
|
||||
The direct output from learned diffusion model - the velocity,
|
||||
timestep (`float`):
|
||||
The current discrete timestep in the diffusion chain (global or per-token).
|
||||
sample (`torch.FloatTensor`):
|
||||
A current latent tokens to be de-noised.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`.
|
||||
stochastic_sampling (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use stochastic sampling for the sampling process.
|
||||
|
||||
Returns:
|
||||
[`~schedulers.scheduling_utils.RectifiedFlowSchedulerOutput`] or `tuple`:
|
||||
If return_dict is `True`, [`~schedulers.rf_scheduler.RectifiedFlowSchedulerOutput`] is returned,
|
||||
otherwise a tuple is returned where the first element is the sample tensor.
|
||||
"""
|
||||
if self.num_inference_steps is None:
|
||||
raise ValueError(
|
||||
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
|
||||
)
|
||||
t_eps = 1e-6 # Small epsilon to avoid numerical issues in timestep values
|
||||
|
||||
timesteps_padded = torch.cat(
|
||||
[self.timesteps, torch.zeros(1, device=self.timesteps.device)]
|
||||
)
|
||||
|
||||
# Find the next lower timestep(s) and compute the dt from the current timestep(s)
|
||||
if timestep.ndim == 0:
|
||||
# Global timestep case
|
||||
lower_mask = timesteps_padded < timestep - t_eps
|
||||
lower_timestep = timesteps_padded[lower_mask][0] # Closest lower timestep
|
||||
dt = timestep - lower_timestep
|
||||
|
||||
else:
|
||||
# Per-token case
|
||||
assert timestep.ndim == 2
|
||||
lower_mask = timesteps_padded[:, None, None] < timestep[None] - t_eps
|
||||
lower_timestep = lower_mask * timesteps_padded[:, None, None]
|
||||
lower_timestep, _ = lower_timestep.max(dim=0)
|
||||
dt = (timestep - lower_timestep)[..., None]
|
||||
|
||||
# Compute previous sample
|
||||
if stochastic_sampling:
|
||||
x0 = sample - timestep[..., None] * model_output
|
||||
next_timestep = timestep[..., None] - dt
|
||||
prev_sample = self.add_noise(x0, torch.randn_like(sample), next_timestep)
|
||||
else:
|
||||
prev_sample = sample - dt * model_output
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample,)
|
||||
|
||||
return RectifiedFlowSchedulerOutput(prev_sample=prev_sample)
|
||||
|
||||
def add_noise(
|
||||
self,
|
||||
original_samples: torch.FloatTensor,
|
||||
noise: torch.FloatTensor,
|
||||
timesteps: torch.FloatTensor,
|
||||
) -> torch.FloatTensor:
|
||||
sigmas = timesteps
|
||||
sigmas = append_dims(sigmas, original_samples.ndim)
|
||||
alphas = 1 - sigmas
|
||||
noisy_samples = alphas * original_samples + sigmas * noise
|
||||
return noisy_samples
|
||||
0
ltx_video/utils/__init__.py
Normal file
0
ltx_video/utils/__init__.py
Normal file
174
ltx_video/utils/diffusers_config_mapping.py
Normal file
174
ltx_video/utils/diffusers_config_mapping.py
Normal file
@@ -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",
|
||||
}
|
||||
214
ltx_video/utils/prompt_enhance_utils.py
Normal file
214
ltx_video/utils/prompt_enhance_utils.py
Normal file
@@ -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
|
||||
8
ltx_video/utils/skip_layer_strategy.py
Normal file
8
ltx_video/utils/skip_layer_strategy.py
Normal file
@@ -0,0 +1,8 @@
|
||||
from enum import Enum, auto
|
||||
|
||||
|
||||
class SkipLayerStrategy(Enum):
|
||||
AttentionSkip = auto()
|
||||
AttentionValues = auto()
|
||||
Residual = auto()
|
||||
TransformerBlock = auto()
|
||||
25
ltx_video/utils/torch_utils.py
Normal file
25
ltx_video/utils/torch_utils.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user