碎碎念
终于找到了一个模型符合规范可以分割的模型和项目了，总之可以勉强在两张2080ti上面跑起来，不过要手动修改模型的分割逻辑，把一个大模型切割放在显卡里，虽然慢的要死是六分钟左右生成一个分辨率512*512 4秒的视频，不过好歹是跑起来了，补足了工作流的最后一环。
开搞
LTX-video 地址：https://github.com/Lightricks/LTX-Video
配置环境
老规矩还是先创建环境：
conda create -n ltxvideo python=3.10
conda activate ltxvideo然后建议先手动下载好配置文件LTX-Video/configs/ltxv-13b-0.9.8-distilled.yaml里面的模型（如果要使用其他模型请修改对应的配置文件），代码内下载也可以，主要看你的网络环境吧
‘huggingface-cli login
export HF_ENDPOINT=https://hf-mirror.com
# 下载 PixArt-alpha Text Encoder 模型
huggingface-cli download PixArt-alpha/PixArt-XL-2-1024-MS --local-dir ./PixArt-alpha-PixArt-XL-2-1024-MS
# 下载 Florence-2 PromptGen 模型
huggingface-cli download MiaoshouAI/Florence-2-large-PromptGen-v2.0 --local-dir ./MiaoshouAI-Florence-2-large-PromptGen-v2.0
# 下载 Llama-3.2 LLM 模型
huggingface-cli download unsloth/Llama-3.2-3B-Instruct --local-dir ./unsloth-Llama-3.2-3B-Instruct然后对以下几个文件进行修改
LTX-Video/configs/ltxv-13b-0.9.8-distilled.yaml修改
LTX-Video/configs/ltxv-13b-0.9.8-distilled.yaml
pipeline_type: multi-scale
checkpoint_path: "ltxv-13b-0.9.8-distilled.safetensors"
downscale_factor: 0.6666666
spatial_upscaler_model_path: "ltxv-spatial-upscaler-0.9.8.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: 0
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]
tone_map_compression_ratio: 0.6
AIvideo/LTX-Video/ltx_video/models/autoencoders/vae_encode.py修改
LTX-Video/ltx_video/models/autoencoders/vae_encode.py
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:
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 = []
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:
# ==================================================================================
# START OF PROACTIVE FIX
# This block ensures the data is moved to the VAE's device before decoding.
# ==================================================================================
# Get the device of the VAE model in a robust way (works with accelerate)
vae_device = next(vae.parameters()).device
# First, un-normalize the latents. The fix in that function handles its own device mismatches.
un_normalized_latents = un_normalize_latents(latents, vae, vae_per_channel_normalize)
# Move the result to the VAE's device before decoding. This is the crucial step.
latents_for_decode = un_normalized_latents.to(device=vae_device, dtype=vae.dtype)
if isinstance(vae, (VideoAutoencoder, CausalVideoAutoencoder)):
*_, fl, hl, wl = latents.shape
temporal_scale, spatial_scale, _ = get_vae_size_scale_factor(vae)
vae_decode_kwargs = {}
if timestep is not None:
# Also ensure the timestep tensor is on the correct device
vae_decode_kwargs["timestep"] = timestep.to(vae_device)
image = vae.decode(
latents_for_decode,
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(
latents_for_decode,
return_dict=False,
)[0]
# ==================================================================================
# END OF PROACTIVE FIX
# ==================================================================================
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:
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:
if vae_per_channel_normalize:
# This part is correctly modified
mean_on_device = vae.mean_of_means.to(device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
std_on_device = vae.std_of_means.to(device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
return (latents - mean_on_device) / std_on_device
else:
# This branch is safe
return latents * vae.config.scaling_factor
def un_normalize_latents(
latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False
) -> Tensor:
if vae_per_channel_normalize:
# This part is correctly modified
std_on_device = vae.std_of_means.to(device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
mean_on_device = vae.mean_of_means.to(device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1)
return (latents * std_on_device) + mean_on_device
else:
# This branch is safe
return latents / vae.config.scaling_factorLTX-Video/ltx_video/models/transformers/transformer3d.py修改
LTX-Video/ltx_video/models/transformers/transformer3d.py
# 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
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.")
}
super().load_state_dict(state_dict, *args, **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,
indices_grid: torch.Tensor,
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,
return_dict: bool = True,
):
"""
The [`Transformer2DModel`] forward method.
(Docstring remains the same)
"""
if not self.use_tpu_flash_attention:
if attention_mask is not None and attention_mask.ndim == 2:
attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
attention_mask = attention_mask.unsqueeze(1)
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 Projection
hidden_states = hidden_states.to(self.patchify_proj.weight.device)
hidden_states = self.patchify_proj(hidden_states)
if self.timestep_scale_multiplier:
timestep = self.timestep_scale_multiplier * timestep
freqs_cis = self.precompute_freqs_cis(indices_grid)
batch_size = hidden_states.shape[0]
# 2. Timestep Embedding
timestep = timestep.to(next(self.adaln_single.parameters()).device)
timestep, embedded_timestep = self.adaln_single(
timestep.flatten(),
{"resolution": None, "aspect_ratio": None},
batch_size=batch_size,
hidden_dtype=hidden_states.dtype,
)
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
embedded_timestep = embedded_timestep.view(
batch_size, -1, embedded_timestep.shape[-1]
)
# 3. Caption Projection
if self.caption_projection is not None:
encoder_hidden_states = encoder_hidden_states.to(next(self.caption_projection.parameters()).device)
encoder_hidden_states = self.caption_projection(encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states.view(
batch_size, -1, hidden_states.shape[-1]
)
# 4. Transformer Blocks Loop
for block_idx, block in enumerate(self.transformer_blocks):
block_device = next(block.parameters()).device
hidden_states = hidden_states.to(block_device)
freqs_cis = tuple(t.to(block_device) for t in freqs_cis)
if encoder_hidden_states is not None:
encoder_hidden_states = encoder_hidden_states.to(block_device)
if timestep is not None:
timestep = timestep.to(block_device)
if attention_mask is not None:
attention_mask = attention_mask.to(block_device)
if encoder_attention_mask is not None:
encoder_attention_mask = encoder_attention_mask.to(block_device)
if self.training and self.gradient_checkpointing:
def create_custom_forward(module, return_dict=None):
def custom_forward(*inputs):
if return_dict is not None:
return module(*inputs, return_dict=return_dict)
else:
return module(*inputs)
return custom_forward
ckpt_kwargs: Dict[str, Any] = (
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
)
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
freqs_cis,
attention_mask,
encoder_hidden_states,
encoder_attention_mask,
timestep,
cross_attention_kwargs,
class_labels,
(
skip_layer_mask[block_idx]
if skip_layer_mask is not None
else None
),
skip_layer_strategy,
**ckpt_kwargs,
)
else:
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=(
skip_layer_mask[block_idx]
if skip_layer_mask is not None
else None
),
skip_layer_strategy=skip_layer_strategy,
)
# 5. Output Layers
# Get device from proj_out, as norm_out has no parameters (elementwise_affine=False).
final_device = self.proj_out.weight.device
hidden_states = hidden_states.to(final_device)
embedded_timestep = embedded_timestep.to(final_device)
scale_shift_values = (
self.scale_shift_table.to(final_device)[None, None] + embedded_timestep[:, :, None]
)
shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
# norm_out might be on a different device, but since it has no learnable params,
# it's just a calculation. The data (hidden_states) is already on the target device.
hidden_states = self.norm_out(hidden_states)
hidden_states = hidden_states * (1 + scale) + shift
# hidden_states is already on the correct device for proj_out
hidden_states = self.proj_out(hidden_states)
if not return_dict:
return (hidden_states,)
return Transformer3DModelOutput(sample=hidden_states)LTX-Video/ltx_video/pipelines/pipeline_ltx_video.py修改
LTX-Video/ltx_video/pipelines/pipeline_ltx_video.py
# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
import copy
import inspect
import math
import re
from contextlib import nullcontext
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
from diffusers.image_processor import VaeImageProcessor
from diffusers.models import AutoencoderKL
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.schedulers import DPMSolverMultistepScheduler
from diffusers.utils import deprecate, logging
from diffusers.utils.torch_utils import randn_tensor
from einops import rearrange
from transformers import (
T5EncoderModel,
T5Tokenizer,
AutoModelForCausalLM,
AutoProcessor,
AutoTokenizer,
)
from ltx_video.models.autoencoders.causal_video_autoencoder import (
CausalVideoAutoencoder,
)
from ltx_video.models.autoencoders.vae_encode import (
get_vae_size_scale_factor,
latent_to_pixel_coords,
vae_decode,
vae_encode,
)
from ltx_video.models.transformers.symmetric_patchifier import Patchifier
from ltx_video.models.transformers.transformer3d import Transformer3DModel
from ltx_video.schedulers.rf import TimestepShifter
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
from ltx_video.utils.prompt_enhance_utils import generate_cinematic_prompt
from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
from ltx_video.models.autoencoders.vae_encode import (
un_normalize_latents,
normalize_latents,
)
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
ASPECT_RATIO_1024_BIN = {
"0.25": [512.0, 2048.0],
"0.28": [512.0, 1856.0],
"0.32": [576.0, 1792.0],
"0.33": [576.0, 1728.0],
"0.35": [576.0, 1664.0],
"0.4": [640.0, 1600.0],
"0.42": [640.0, 1536.0],
"0.48": [704.0, 1472.0],
"0.5": [704.0, 1408.0],
"0.52": [704.0, 1344.0],
"0.57": [768.0, 1344.0],
"0.6": [768.0, 1280.0],
"0.68": [832.0, 1216.0],
"0.72": [832.0, 1152.0],
"0.78": [896.0, 1152.0],
"0.82": [896.0, 1088.0],
"0.88": [960.0, 1088.0],
"0.94": [960.0, 1024.0],
"1.0": [1024.0, 1024.0],
"1.07": [1024.0, 960.0],
"1.13": [1088.0, 960.0],
"1.21": [1088.0, 896.0],
"1.29": [1152.0, 896.0],
"1.38": [1152.0, 832.0],
"1.46": [1216.0, 832.0],
"1.67": [1280.0, 768.0],
"1.75": [1344.0, 768.0],
"2.0": [1408.0, 704.0],
"2.09": [1472.0, 704.0],
"2.4": [1536.0, 640.0],
"2.5": [1600.0, 640.0],
"3.0": [1728.0, 576.0],
"4.0": [2048.0, 512.0],
}
ASPECT_RATIO_512_BIN = {
"0.25": [256.0, 1024.0],
"0.28": [256.0, 928.0],
"0.32": [288.0, 896.0],
"0.33": [288.0, 864.0],
"0.35": [288.0, 832.0],
"0.4": [320.0, 800.0],
"0.42": [320.0, 768.0],
"0.48": [352.0, 736.0],
"0.5": [352.0, 704.0],
"0.52": [352.0, 672.0],
"0.57": [384.0, 672.0],
"0.6": [384.0, 640.0],
"0.68": [416.0, 608.0],
"0.72": [416.0, 576.0],
"0.78": [448.0, 576.0],
"0.82": [448.0, 544.0],
"0.88": [480.0, 544.0],
"0.94": [480.0, 512.0],
"1.0": [512.0, 512.0],
"1.07": [512.0, 480.0],
"1.13": [544.0, 480.0],
"1.21": [544.0, 448.0],
"1.29": [576.0, 448.0],
"1.38": [576.0, 416.0],
"1.46": [608.0, 416.0],
"1.67": [640.0, 384.0],
"1.75": [672.0, 384.0],
"2.0": [704.0, 352.0],
"2.09": [736.0, 352.0],
"2.4": [768.0, 320.0],
"2.5": [800.0, 320.0],
"3.0": [864.0, 288.0],
"4.0": [1024.0, 256.0],
}
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
scheduler,
num_inference_steps: Optional[int] = None,
device: Optional[Union[str, torch.device]] = None,
timesteps: Optional[List[int]] = None,
skip_initial_inference_steps: int = 0,
skip_final_inference_steps: int = 0,
**kwargs,
):
"""
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
Args:
scheduler (`SchedulerMixin`):
The scheduler to get timesteps from.
num_inference_steps (`int`):
The number of diffusion steps used when generating samples with a pre-trained model. If used,
`timesteps` must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
timesteps (`List[int]`, *optional*):
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
must be `None`.
max_timestep ('float', *optional*, defaults to 1.0):
The initial noising level for image-to-image/video-to-video. The list if timestamps will be
truncated to start with a timestamp greater or equal to this.
Returns:
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None:
accepts_timesteps = "timesteps" in set(
inspect.signature(scheduler.set_timesteps).parameters.keys()
)
if not accepts_timesteps:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
)
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
timesteps = scheduler.timesteps
num_inference_steps = len(timesteps)
else:
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
timesteps = scheduler.timesteps
if (
skip_initial_inference_steps < 0
or skip_final_inference_steps < 0
or skip_initial_inference_steps + skip_final_inference_steps
>= num_inference_steps
):
raise ValueError(
"invalid skip inference step values: must be non-negative and the sum of skip_initial_inference_steps and skip_final_inference_steps must be less than the number of inference steps"
)
timesteps = timesteps[
skip_initial_inference_steps : len(timesteps) - skip_final_inference_steps
]
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
num_inference_steps = len(timesteps)
return timesteps, num_inference_steps
@dataclass
class ConditioningItem:
"""
Defines a single frame-conditioning item - a single frame or a sequence of frames.
Attributes:
media_item (torch.Tensor): shape=(b, 3, f, h, w). The media item to condition on.
media_frame_number (int): The start-frame number of the media item in the generated video.
conditioning_strength (float): The strength of the conditioning (1.0 = full conditioning).
media_x (Optional[int]): Optional left x coordinate of the media item in the generated frame.
media_y (Optional[int]): Optional top y coordinate of the media item in the generated frame.
"""
media_item: torch.Tensor
media_frame_number: int
conditioning_strength: float
media_x: Optional[int] = None
media_y: Optional[int] = None
class LTXVideoPipeline(DiffusionPipeline):
r"""
Pipeline for text-to-image generation using LTX-Video.
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
text_encoder ([`T5EncoderModel`]):
Frozen text-encoder. This uses
[T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
[t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant.
tokenizer (`T5Tokenizer`):
Tokenizer of class
[T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
transformer ([`Transformer2DModel`]):
A text conditioned `Transformer2DModel` to denoise the encoded image latents.
scheduler ([`SchedulerMixin`]):
A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
"""
bad_punct_regex = re.compile(
r"["
+ "#®•©™&@·º½¾¿¡§~"
+ r"\)"
+ r"\("
+ r"\]"
+ r"\["
+ r"\}"
+ r"\{"
+ r"\|"
+ "\\"
+ r"\/"
+ r"\*"
+ r"]{1,}"
) # noqa
_optional_components = [
"tokenizer",
"text_encoder",
"prompt_enhancer_image_caption_model",
"prompt_enhancer_image_caption_processor",
"prompt_enhancer_llm_model",
"prompt_enhancer_llm_tokenizer",
]
model_cpu_offload_seq = "prompt_enhancer_image_caption_model->prompt_enhancer_llm_model->text_encoder->transformer->vae"
def __init__(
self,
tokenizer: T5Tokenizer,
text_encoder: T5EncoderModel,
vae: AutoencoderKL,
transformer: Transformer3DModel,
scheduler: DPMSolverMultistepScheduler,
patchifier: Patchifier,
prompt_enhancer_image_caption_model: AutoModelForCausalLM,
prompt_enhancer_image_caption_processor: AutoProcessor,
prompt_enhancer_llm_model: AutoModelForCausalLM,
prompt_enhancer_llm_tokenizer: AutoTokenizer,
allowed_inference_steps: Optional[List[float]] = None,
):
super().__init__()
self.register_modules(
tokenizer=tokenizer,
text_encoder=text_encoder,
vae=vae,
transformer=transformer,
scheduler=scheduler,
patchifier=patchifier,
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,
)
self.video_scale_factor, self.vae_scale_factor, _ = get_vae_size_scale_factor(
self.vae
)
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
self.allowed_inference_steps = allowed_inference_steps
def mask_text_embeddings(self, emb, mask):
if emb.shape[0] == 1:
keep_index = mask.sum().item()
return emb[:, :, :keep_index, :], keep_index
else:
masked_feature = emb * mask[:, None, :, None]
return masked_feature, emb.shape[2]
# Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
def encode_prompt(
self,
prompt: Union[str, List[str]],
do_classifier_free_guidance: bool = True,
negative_prompt: str = "",
num_images_per_prompt: int = 1,
device: Optional[torch.device] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
prompt_attention_mask: Optional[torch.FloatTensor] = None,
negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
text_encoder_max_tokens: int = 256,
**kwargs,
):
# ================== 验证代码 ==================
print(">>> NOW RUNNING THE CORRECTLY MODIFIED 'encode_prompt' FUNCTION. <<<")
# ===============================================
r"""
Encodes the prompt into text encoder hidden states.
(Docstring remains the same)
"""
if "mask_feature" in kwargs:
deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
if device is None:
device = self._execution_device
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
max_length = text_encoder_max_tokens
if prompt_embeds is None:
assert (
self.text_encoder is not None
), "You should provide either prompt_embeds or self.text_encoder should not be None,"
# 关键修复点 1：获取 text_encoder 所在的设备
text_encoder_device = self.text_encoder.get_input_embeddings().weight.device
prompt = self._text_preprocessing(prompt)
text_inputs = self.tokenizer(
prompt,
padding="max_length",
max_length=max_length,
truncation=True,
add_special_tokens=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
untruncated_ids = self.tokenizer(
prompt, padding="longest", return_tensors="pt"
).input_ids
if untruncated_ids.shape[-1] >= text_input_ids.shape[
-1
] and not torch.equal(text_input_ids, untruncated_ids):
removed_text = self.tokenizer.batch_decode(
untruncated_ids[:, max_length - 1 : -1]
)
logger.warning(
"The following part of your input was truncated because CLIP can only handle sequences up to"
f" {max_length} tokens: {removed_text}"
)
prompt_attention_mask = text_inputs.attention_mask.to(text_encoder_device)
# 关键修复点 2：在调用 text_encoder 之前，将 input_ids 移动到正确的设备
prompt_embeds = self.text_encoder(
text_input_ids.to(text_encoder_device), attention_mask=prompt_attention_mask
)
prompt_embeds = prompt_embeds[0]
if self.text_encoder is not None:
dtype = self.text_encoder.dtype
elif self.transformer is not None:
dtype = self.transformer.dtype
else:
dtype = None
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
bs_embed, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds = prompt_embeds.view(
bs_embed * num_images_per_prompt, seq_len, -1
)
prompt_attention_mask = prompt_attention_mask.to(device) # Move mask to final device
prompt_attention_mask = prompt_attention_mask.repeat(1, num_images_per_prompt)
prompt_attention_mask = prompt_attention_mask.view(
bs_embed * num_images_per_prompt, -1
)
if do_classifier_free_guidance and negative_prompt_embeds is None:
uncond_tokens = self._text_preprocessing(negative_prompt)
uncond_tokens = uncond_tokens * batch_size
max_length = prompt_embeds.shape[1]
uncond_input = self.tokenizer(
uncond_tokens,
padding="max_length",
max_length=max_length,
truncation=True,
return_attention_mask=True,
add_special_tokens=True,
return_tensors="pt",
)
negative_prompt_attention_mask = uncond_input.attention_mask
# 关键修复点 3：同样为 negative_prompt 获取设备并移动
text_encoder_device = self.text_encoder.get_input_embeddings().weight.device
negative_prompt_attention_mask = negative_prompt_attention_mask.to(text_encoder_device)
negative_prompt_embeds = self.text_encoder(
uncond_input.input_ids.to(text_encoder_device),
attention_mask=negative_prompt_attention_mask,
)
negative_prompt_embeds = negative_prompt_embeds[0]
if do_classifier_free_guidance:
seq_len = negative_prompt_embeds.shape[1]
negative_prompt_embeds = negative_prompt_embeds.to(
dtype=dtype, device=device
)
negative_prompt_embeds = negative_prompt_embeds.repeat(
1, num_images_per_prompt, 1
)
negative_prompt_embeds = negative_prompt_embeds.view(
batch_size * num_images_per_prompt, seq_len, -1
)
negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(
1, num_images_per_prompt
)
negative_prompt_attention_mask = negative_prompt_attention_mask.view(
bs_embed * num_images_per_prompt, -1
)
else:
negative_prompt_embeds = None
negative_prompt_attention_mask = None
return (
prompt_embeds,
prompt_attention_mask,
negative_prompt_embeds,
negative_prompt_attention_mask,
)
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
def prepare_extra_step_kwargs(self, generator, eta):
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
# and should be between [0, 1]
accepts_eta = "eta" in set(
inspect.signature(self.scheduler.step).parameters.keys()
)
extra_step_kwargs = {}
if accepts_eta:
extra_step_kwargs["eta"] = eta
# check if the scheduler accepts generator
accepts_generator = "generator" in set(
inspect.signature(self.scheduler.step).parameters.keys()
)
if accepts_generator:
extra_step_kwargs["generator"] = generator
return extra_step_kwargs
def check_inputs(
self,
prompt,
height,
width,
negative_prompt,
prompt_embeds=None,
negative_prompt_embeds=None,
prompt_attention_mask=None,
negative_prompt_attention_mask=None,
enhance_prompt=False,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(
f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
)
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elif prompt is not None and (
not isinstance(prompt, str) and not isinstance(prompt, list)
):
raise ValueError(
f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
)
if prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
if prompt_embeds is not None and prompt_attention_mask is None:
raise ValueError(
"Must provide `prompt_attention_mask` when specifying `prompt_embeds`."
)
if (
negative_prompt_embeds is not None
and negative_prompt_attention_mask is None
):
raise ValueError(
"Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`."
)
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
raise ValueError(
"`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
f" {negative_prompt_attention_mask.shape}."
)
if enhance_prompt:
assert (
self.prompt_enhancer_image_caption_model is not None
), "Image caption model must be initialized if enhance_prompt is True"
assert (
self.prompt_enhancer_image_caption_processor is not None
), "Image caption processor must be initialized if enhance_prompt is True"
assert (
self.prompt_enhancer_llm_model is not None
), "Text prompt enhancer model must be initialized if enhance_prompt is True"
assert (
self.prompt_enhancer_llm_tokenizer is not None
), "Text prompt enhancer tokenizer must be initialized if enhance_prompt is True"
def _text_preprocessing(self, text):
if not isinstance(text, (tuple, list)):
text = [text]
def process(text: str):
text = text.strip()
return text
return [process(t) for t in text]
@staticmethod
def add_noise_to_image_conditioning_latents(
t: float,
init_latents: torch.Tensor,
latents: torch.Tensor,
noise_scale: float,
conditioning_mask: torch.Tensor,
generator,
eps=1e-6,
):
"""
Add timestep-dependent noise to the hard-conditioning latents.
This helps with motion continuity, especially when conditioned on a single frame.
"""
noise = randn_tensor(
latents.shape,
generator=generator,
device=latents.device,
dtype=latents.dtype,
)
# Add noise only to hard-conditioning latents (conditioning_mask = 1.0)
need_to_noise = (conditioning_mask > 1.0 - eps).unsqueeze(-1)
noised_latents = init_latents + noise_scale * noise * (t**2)
latents = torch.where(need_to_noise, noised_latents, latents)
return latents
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
def prepare_latents(
self,
latents: torch.Tensor | None,
media_items: torch.Tensor | None,
timestep: float,
latent_shape: torch.Size | Tuple[Any, ...],
dtype: torch.dtype,
device: torch.device,
generator: torch.Generator | List[torch.Generator],
vae_per_channel_normalize: bool = True,
):
"""
Prepare the initial latent tensor to be denoised.
The latents are either pure noise or a noised version of the encoded media items.
Args:
latents (`torch.FloatTensor` or `None`):
The latents to use (provided by the user) or `None` to create new latents.
media_items (`torch.FloatTensor` or `None`):
An image or video to be updated using img2img or vid2vid. The media item is encoded and noised.
timestep (`float`):
The timestep to noise the encoded media_items to.
latent_shape (`torch.Size`):
The target latent shape.
dtype (`torch.dtype`):
The target dtype.
device (`torch.device`):
The target device.
generator (`torch.Generator` or `List[torch.Generator]`):
Generator(s) to be used for the noising process.
vae_per_channel_normalize ('bool'):
When encoding the media_items, whether to normalize the latents per-channel.
Returns:
`torch.FloatTensor`: The latents to be used for the denoising process. This is a tensor of shape
(batch_size, num_channels, height, width).
"""
if isinstance(generator, list) and len(generator) != latent_shape[0]:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {latent_shape[0]}. Make sure the batch size matches the length of the generators."
)
# Initialize the latents with the given latents or encoded media item, if provided
assert (
latents is None or media_items is None
), "Cannot provide both latents and media_items. Please provide only one of the two."
assert (
latents is None and media_items is None or timestep < 1.0
), "Input media_item or latents are provided, but they will be replaced with noise."
if media_items is not None:
latents = vae_encode(
media_items.to(dtype=self.vae.dtype, device=self.vae.device),
self.vae,
vae_per_channel_normalize=vae_per_channel_normalize,
)
if latents is not None:
assert (
latents.shape == latent_shape
), f"Latents have to be of shape {latent_shape} but are {latents.shape}."
latents = latents.to(device=device, dtype=dtype)
# For backward compatibility, generate in the "patchified" shape and rearrange
b, c, f, h, w = latent_shape
noise = randn_tensor(
(b, f * h * w, c), generator=generator, device=device, dtype=dtype
)
noise = rearrange(noise, "b (f h w) c -> b c f h w", f=f, h=h, w=w)
# scale the initial noise by the standard deviation required by the scheduler
noise = noise * self.scheduler.init_noise_sigma
if latents is None:
latents = noise
else:
# Noise the latents to the required (first) timestep
latents = timestep * noise + (1 - timestep) * latents
return latents
@staticmethod
def classify_height_width_bin(
height: int, width: int, ratios: dict
) -> Tuple[int, int]:
"""Returns binned height and width."""
ar = float(height / width)
closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
default_hw = ratios[closest_ratio]
return int(default_hw[0]), int(default_hw[1])
@staticmethod
def resize_and_crop_tensor(
samples: torch.Tensor, new_width: int, new_height: int
) -> torch.Tensor:
n_frames, orig_height, orig_width = samples.shape[-3:]
# Check if resizing is needed
if orig_height != new_height or orig_width != new_width:
ratio = max(new_height / orig_height, new_width / orig_width)
resized_width = int(orig_width * ratio)
resized_height = int(orig_height * ratio)
# Resize
samples = LTXVideoPipeline.resize_tensor(
samples, resized_height, resized_width
)
# Center Crop
start_x = (resized_width - new_width) // 2
end_x = start_x + new_width
start_y = (resized_height - new_height) // 2
end_y = start_y + new_height
samples = samples[..., start_y:end_y, start_x:end_x]
return samples
@staticmethod
def resize_tensor(media_items, height, width):
n_frames = media_items.shape[2]
if media_items.shape[-2:] != (height, width):
media_items = rearrange(media_items, "b c n h w -> (b n) c h w")
media_items = F.interpolate(
media_items,
size=(height, width),
mode="bilinear",
align_corners=False,
)
media_items = rearrange(media_items, "(b n) c h w -> b c n h w", n=n_frames)
return media_items
@torch.no_grad()
def __call__(
self,
height: int,
width: int,
num_frames: int,
frame_rate: float,
prompt: Union[str, List[str]] = None,
negative_prompt: str = "",
num_inference_steps: int = 20,
skip_initial_inference_steps: int = 0,
skip_final_inference_steps: int = 0,
timesteps: List[int] = None,
guidance_scale: Union[float, List[float]] = 4.5,
cfg_star_rescale: bool = False,
skip_layer_strategy: Optional[SkipLayerStrategy] = None,
skip_block_list: Optional[Union[List[List[int]], List[int]]] = None,
stg_scale: Union[float, List[float]] = 1.0,
rescaling_scale: Union[float, List[float]] = 0.7,
guidance_timesteps: Optional[List[int]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
prompt_attention_mask: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
conditioning_items: Optional[List[ConditioningItem]] = None,
decode_timestep: Union[List[float], float] = 0.0,
decode_noise_scale: Optional[List[float]] = None,
mixed_precision: bool = False,
offload_to_cpu: bool = False,
enhance_prompt: bool = False,
text_encoder_max_tokens: int = 256,
stochastic_sampling: bool = False,
media_items: Optional[torch.Tensor] = None,
tone_map_compression_ratio: float = 0.0,
**kwargs,
) -> Union[ImagePipelineOutput, Tuple]:
"""
Function invoked when calling the pipeline for generation.
Args:
prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
num_inference_steps (`int`, *optional*, defaults to 100):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. If `timesteps` is provided, this parameter is ignored.
skip_initial_inference_steps (`int`, *optional*, defaults to 0):
The number of initial timesteps to skip. After calculating the timesteps, this number of timesteps will
be removed from the beginning of the timesteps list. Meaning the highest-timesteps values will not run.
skip_final_inference_steps (`int`, *optional*, defaults to 0):
The number of final timesteps to skip. After calculating the timesteps, this number of timesteps will
be removed from the end of the timesteps list. Meaning the lowest-timesteps values will not run.
timesteps (`List[int]`, *optional*):
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
timesteps are used. Must be in descending order.
guidance_scale (`float`, *optional*, defaults to 4.5):
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
`guidance_scale` is defined as `w` of equation 2. of [Imagen
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
usually at the expense of lower image quality.
cfg_star_rescale (`bool`, *optional*, defaults to `False`):
If set to `True`, applies the CFG star rescale. Scales the negative prediction according to dot
product between positive and negative.
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
height (`int`, *optional*, defaults to self.unet.config.sample_size):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to self.unet.config.sample_size):
The width in pixels of the generated image.
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
[`schedulers.DDIMScheduler`], will be ignored for others.
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will ge generated by sampling using the supplied random `generator`.
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative text embeddings. This negative prompt should be "". If not
provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
Pre-generated attention mask for negative text embeddings.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generate image. Choose between
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
callback_on_step_end (`Callable`, *optional*):
A function that calls at the end of each denoising steps during the inference. The function is called
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
`callback_on_step_end_tensor_inputs`.
use_resolution_binning (`bool` defaults to `True`):
If set to `True`, the requested height and width are first mapped to the closest resolutions using
`ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
the requested resolution. Useful for generating non-square images.
enhance_prompt (`bool`, *optional*, defaults to `False`):
If set to `True`, the prompt is enhanced using a LLM model.
text_encoder_max_tokens (`int`, *optional*, defaults to `256`):
The maximum number of tokens to use for the text encoder.
stochastic_sampling (`bool`, *optional*, defaults to `False`):
If set to `True`, the sampling is stochastic. If set to `False`, the sampling is deterministic.
media_items ('torch.Tensor', *optional*):
The input media item used for image-to-image / video-to-video.
tone_map_compression_ratio: compression ratio for tone mapping, defaults to 0.0.
If set to 0.0, no tone mapping is applied. If set to 1.0 - full compression is applied.
Examples:
Returns:
[`~pipelines.ImagePipelineOutput`] or `tuple`:
If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
returned where the first element is a list with the generated images
"""
if "mask_feature" in kwargs:
deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
is_video = kwargs.get("is_video", False)
self.check_inputs(
prompt,
height,
width,
negative_prompt,
prompt_embeds,
negative_prompt_embeds,
prompt_attention_mask,
negative_prompt_attention_mask,
)
# 2. Default height and width to transformer
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
self.video_scale_factor = self.video_scale_factor if is_video else 1
vae_per_channel_normalize = kwargs.get("vae_per_channel_normalize", True)
image_cond_noise_scale = kwargs.get("image_cond_noise_scale", 0.0)
latent_height = height // self.vae_scale_factor
latent_width = width // self.vae_scale_factor
latent_num_frames = num_frames // self.video_scale_factor
if isinstance(self.vae, CausalVideoAutoencoder) and is_video:
latent_num_frames += 1
latent_shape = (
batch_size * num_images_per_prompt,
self.transformer.config.in_channels,
latent_num_frames,
latent_height,
latent_width,
)
# Prepare the list of denoising time-steps
retrieve_timesteps_kwargs = {}
if isinstance(self.scheduler, TimestepShifter):
retrieve_timesteps_kwargs["samples_shape"] = latent_shape
assert (
skip_initial_inference_steps == 0
or latents is not None
or media_items is not None
), (
f"skip_initial_inference_steps ({skip_initial_inference_steps}) is used for image-to-image/video-to-video - "
"media_item or latents should be provided."
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
num_inference_steps,
device,
timesteps,
skip_initial_inference_steps=skip_initial_inference_steps,
skip_final_inference_steps=skip_final_inference_steps,
**retrieve_timesteps_kwargs,
)
if self.allowed_inference_steps is not None:
for timestep in [round(x, 4) for x in timesteps.tolist()]:
assert (
timestep in self.allowed_inference_steps
), f"Invalid inference timestep {timestep}. Allowed timesteps are {self.allowed_inference_steps}."
if guidance_timesteps:
guidance_mapping = []
for timestep in timesteps:
indices = [
i for i, val in enumerate(guidance_timesteps) if val <= timestep
]
# assert len(indices) > 0, f"No guidance timestep found for {timestep}"
guidance_mapping.append(
indices[0] if len(indices) > 0 else (len(guidance_timesteps) - 1)
)
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
if not isinstance(guidance_scale, List):
guidance_scale = [guidance_scale] * len(timesteps)
else:
guidance_scale = [
guidance_scale[guidance_mapping[i]] for i in range(len(timesteps))
]
if not isinstance(stg_scale, List):
stg_scale = [stg_scale] * len(timesteps)
else:
stg_scale = [stg_scale[guidance_mapping[i]] for i in range(len(timesteps))]
if not isinstance(rescaling_scale, List):
rescaling_scale = [rescaling_scale] * len(timesteps)
else:
rescaling_scale = [
rescaling_scale[guidance_mapping[i]] for i in range(len(timesteps))
]
# Normalize skip_block_list to always be None or a list of lists matching timesteps
if skip_block_list is not None:
# Convert single list to list of lists if needed
if len(skip_block_list) == 0 or not isinstance(skip_block_list[0], list):
skip_block_list = [skip_block_list] * len(timesteps)
else:
new_skip_block_list = []
for i, timestep in enumerate(timesteps):
new_skip_block_list.append(skip_block_list[guidance_mapping[i]])
skip_block_list = new_skip_block_list
if enhance_prompt:
self.prompt_enhancer_image_caption_model = (
self.prompt_enhancer_image_caption_model.to(self._execution_device)
)
self.prompt_enhancer_llm_model = self.prompt_enhancer_llm_model.to(
self._execution_device
)
prompt = generate_cinematic_prompt(
self.prompt_enhancer_image_caption_model,
self.prompt_enhancer_image_caption_processor,
self.prompt_enhancer_llm_model,
self.prompt_enhancer_llm_tokenizer,
prompt,
conditioning_items,
max_new_tokens=text_encoder_max_tokens,
)
# 3. Encode input prompt
if self.text_encoder is not None:
(
prompt_embeds,
prompt_attention_mask,
negative_prompt_embeds,
negative_prompt_attention_mask,
) = self.encode_prompt(
prompt,
True,
negative_prompt=negative_prompt,
num_images_per_prompt=num_images_per_prompt,
device=device,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
text_encoder_max_tokens=text_encoder_max_tokens,
)
# if offload_to_cpu and self.text_encoder is not None:
# self.text_encoder = self.text_encoder.cpu()
#self.transformer = self.transformer.to(self._execution_device)
prompt_embeds_batch = prompt_embeds
prompt_attention_mask_batch = prompt_attention_mask
negative_prompt_embeds = (
torch.zeros_like(prompt_embeds)
if negative_prompt_embeds is None
else negative_prompt_embeds
)
negative_prompt_attention_mask = (
torch.zeros_like(prompt_attention_mask)
if negative_prompt_attention_mask is None
else negative_prompt_attention_mask
)
prompt_embeds_batch = torch.cat(
[negative_prompt_embeds, prompt_embeds, prompt_embeds], dim=0
)
prompt_attention_mask_batch = torch.cat(
[
negative_prompt_attention_mask,
prompt_attention_mask,
prompt_attention_mask,
],
dim=0,
)
# 4. Prepare the initial latents using the provided media and conditioning items
# Prepare the initial latents tensor, shape = (b, c, f, h, w)
latents = self.prepare_latents(
latents=latents,
media_items=media_items,
timestep=timesteps[0],
latent_shape=latent_shape,
dtype=prompt_embeds.dtype,
device=device,
generator=generator,
vae_per_channel_normalize=vae_per_channel_normalize,
)
# Update the latents with the conditioning items and patchify them into (b, n, c)
latents, pixel_coords, conditioning_mask, num_cond_latents = (
self.prepare_conditioning(
conditioning_items=conditioning_items,
init_latents=latents,
num_frames=num_frames,
height=height,
width=width,
vae_per_channel_normalize=vae_per_channel_normalize,
generator=generator,
)
)
init_latents = latents.clone() # Used for image_cond_noise_update
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 7. Denoising loop
num_warmup_steps = max(
len(timesteps) - num_inference_steps * self.scheduler.order, 0
)
orig_conditioning_mask = conditioning_mask
# Befor compiling this code please be aware:
# This code might generate different input shapes if some timesteps have no STG or CFG.
# This means that the codes might need to be compiled mutliple times.
# To avoid that, use the same STG and CFG values for all timesteps.
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
do_classifier_free_guidance = guidance_scale[i] > 1.0
do_spatio_temporal_guidance = stg_scale[i] > 0
do_rescaling = rescaling_scale[i] != 1.0
num_conds = 1
if do_classifier_free_guidance:
num_conds += 1
if do_spatio_temporal_guidance:
num_conds += 1
if do_classifier_free_guidance and do_spatio_temporal_guidance:
indices = slice(batch_size * 0, batch_size * 3)
elif do_classifier_free_guidance:
indices = slice(batch_size * 0, batch_size * 2)
elif do_spatio_temporal_guidance:
indices = slice(batch_size * 1, batch_size * 3)
else:
indices = slice(batch_size * 1, batch_size * 2)
# Prepare skip layer masks
skip_layer_mask: Optional[torch.Tensor] = None
if do_spatio_temporal_guidance:
if skip_block_list is not None:
skip_layer_mask = self.transformer.create_skip_layer_mask(
batch_size, num_conds, num_conds - 1, skip_block_list[i]
)
batch_pixel_coords = torch.cat([pixel_coords] * num_conds)
conditioning_mask = orig_conditioning_mask
if conditioning_mask is not None and is_video:
assert num_images_per_prompt == 1
conditioning_mask = torch.cat([conditioning_mask] * num_conds)
fractional_coords = batch_pixel_coords.to(torch.float32)
fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)
if conditioning_mask is not None and image_cond_noise_scale > 0.0:
latents = self.add_noise_to_image_conditioning_latents(
t,
init_latents,
latents,
image_cond_noise_scale,
orig_conditioning_mask,
generator,
)
latent_model_input = (
torch.cat([latents] * num_conds) if num_conds > 1 else latents
)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t
)
current_timestep = t
if not torch.is_tensor(current_timestep):
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
# This would be a good case for the `match` statement (Python 3.10+)
is_mps = latent_model_input.device.type == "mps"
if isinstance(current_timestep, float):
dtype = torch.float32 if is_mps else torch.float64
else:
dtype = torch.int32 if is_mps else torch.int64
current_timestep = torch.tensor(
[current_timestep],
dtype=dtype,
device=latent_model_input.device,
)
elif len(current_timestep.shape) == 0:
current_timestep = current_timestep[None].to(
latent_model_input.device
)
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
current_timestep = current_timestep.expand(
latent_model_input.shape[0]
).unsqueeze(-1)
if conditioning_mask is not None:
# Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)
# and will start to be denoised when the current timestep is lower than their conditioning timestep.
current_timestep = torch.min(
current_timestep, 1.0 - conditioning_mask
)
# Choose the appropriate context manager based on `mixed_precision`
if mixed_precision:
context_manager = torch.autocast(device.type, dtype=torch.bfloat16)
else:
context_manager = nullcontext() # Dummy context manager
# predict noise model_output
with context_manager:
noise_pred = self.transformer(
latent_model_input.to(self.transformer.dtype),
indices_grid=fractional_coords,
encoder_hidden_states=prompt_embeds_batch[indices].to(
self.transformer.dtype
),
encoder_attention_mask=prompt_attention_mask_batch[indices],
timestep=current_timestep,
skip_layer_mask=skip_layer_mask,
skip_layer_strategy=skip_layer_strategy,
return_dict=False,
)[0]
noise_pred = noise_pred.to(latents.device)
# perform guidance
if do_spatio_temporal_guidance:
noise_pred_text, noise_pred_text_perturb = noise_pred.chunk(
num_conds
)[-2:]
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(num_conds)[:2]
if cfg_star_rescale:
# Rescales the unconditional noise prediction using the projection of the conditional prediction onto it:
# α = (⟨ε_text, ε_uncond⟩ / ||ε_uncond||²), then ε_uncond ← α * ε_uncond
# where ε_text is the conditional noise prediction and ε_uncond is the unconditional one.
positive_flat = noise_pred_text.view(batch_size, -1)
negative_flat = noise_pred_uncond.view(batch_size, -1)
dot_product = torch.sum(
positive_flat * negative_flat, dim=1, keepdim=True
)
squared_norm = (
torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
)
alpha = dot_product / squared_norm
noise_pred_uncond = alpha * noise_pred_uncond
noise_pred = noise_pred_uncond + guidance_scale[i] * (
noise_pred_text - noise_pred_uncond
)
elif do_spatio_temporal_guidance:
noise_pred = noise_pred_text
if do_spatio_temporal_guidance:
noise_pred = noise_pred + stg_scale[i] * (
noise_pred_text - noise_pred_text_perturb
)
if do_rescaling and stg_scale[i] > 0.0:
noise_pred_text_std = noise_pred_text.view(batch_size, -1).std(
dim=1, keepdim=True
)
noise_pred_std = noise_pred.view(batch_size, -1).std(
dim=1, keepdim=True
)
factor = noise_pred_text_std / noise_pred_std
factor = rescaling_scale[i] * factor + (1 - rescaling_scale[i])
noise_pred = noise_pred * factor.view(batch_size, 1, 1)
current_timestep = current_timestep[:1]
# learned sigma
if (
self.transformer.config.out_channels // 2
== self.transformer.config.in_channels
):
noise_pred = noise_pred.chunk(2, dim=1)[0]
# compute previous image: x_t -> x_t-1
latents = self.denoising_step(
latents,
noise_pred,
current_timestep,
orig_conditioning_mask,
t,
extra_step_kwargs,
stochastic_sampling=stochastic_sampling,
)
# call the callback, if provided
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
):
progress_bar.update()
if callback_on_step_end is not None:
callback_on_step_end(self, i, t, {})
if offload_to_cpu:
self.transformer = self.transformer.cpu()
if self._execution_device == "cuda":
torch.cuda.empty_cache()
# Remove the added conditioning latents
latents = latents[:, num_cond_latents:]
latents = self.patchifier.unpatchify(
latents=latents,
output_height=latent_height,
output_width=latent_width,
out_channels=self.transformer.in_channels
// math.prod(self.patchifier.patch_size),
)
if output_type != "latent":
if self.vae.decoder.timestep_conditioning:
noise = torch.randn_like(latents)
if not isinstance(decode_timestep, list):
decode_timestep = [decode_timestep] * latents.shape[0]
if decode_noise_scale is None:
decode_noise_scale = decode_timestep
elif not isinstance(decode_noise_scale, list):
decode_noise_scale = [decode_noise_scale] * latents.shape[0]
decode_timestep = torch.tensor(decode_timestep).to(latents.device)
decode_noise_scale = torch.tensor(decode_noise_scale).to(
latents.device
)[:, None, None, None, None]
latents = (
latents * (1 - decode_noise_scale) + noise * decode_noise_scale
)
else:
decode_timestep = None
latents = self.tone_map_latents(latents, tone_map_compression_ratio)
image = vae_decode(
latents,
self.vae,
is_video,
vae_per_channel_normalize=kwargs["vae_per_channel_normalize"],
timestep=decode_timestep,
)
image = self.image_processor.postprocess(image, output_type=output_type)
else:
image = latents
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)
return ImagePipelineOutput(images=image)
def denoising_step(
self,
latents: torch.Tensor,
noise_pred: torch.Tensor,
current_timestep: torch.Tensor,
conditioning_mask: torch.Tensor,
t: float,
extra_step_kwargs,
t_eps=1e-6,
stochastic_sampling=False,
):
"""
Perform the denoising step for the required tokens, based on the current timestep and
conditioning mask:
Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)
and will start to be denoised when the current timestep is equal or lower than their
conditioning timestep.
(hard-conditioning latents with conditioning_mask = 1.0 are never denoised)
"""
# Denoise the latents using the scheduler
denoised_latents = self.scheduler.step(
noise_pred,
t if current_timestep is None else current_timestep,
latents,
**extra_step_kwargs,
return_dict=False,
stochastic_sampling=stochastic_sampling,
)[0]
if conditioning_mask is None:
return denoised_latents
tokens_to_denoise_mask = (t - t_eps < (1.0 - conditioning_mask)).unsqueeze(-1)
return torch.where(tokens_to_denoise_mask, denoised_latents, latents)
def prepare_conditioning(
self,
conditioning_items: Optional[List[ConditioningItem]],
init_latents: torch.Tensor,
num_frames: int,
height: int,
width: int,
vae_per_channel_normalize: bool = False,
generator=None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
"""
Prepare conditioning tokens based on the provided conditioning items.
This method encodes provided conditioning items (video frames or single frames) into latents
and integrates them with the initial latent tensor. It also calculates corresponding pixel
coordinates, a mask indicating the influence of conditioning latents, and the total number of
conditioning latents.
Args:
conditioning_items (Optional[List[ConditioningItem]]): A list of ConditioningItem objects.
init_latents (torch.Tensor): The initial latent tensor of shape (b, c, f_l, h_l, w_l), where
`f_l` is the number of latent frames, and `h_l` and `w_l` are latent spatial dimensions.
num_frames, height, width: The dimensions of the generated video.
vae_per_channel_normalize (bool, optional): Whether to normalize channels during VAE encoding.
Defaults to `False`.
generator: The random generator
Returns:
Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
- `init_latents` (torch.Tensor): The updated latent tensor including conditioning latents,
patchified into (b, n, c) shape.
- `init_pixel_coords` (torch.Tensor): The pixel coordinates corresponding to the updated
latent tensor.
- `conditioning_mask` (torch.Tensor): A mask indicating the conditioning-strength of each
latent token.
- `num_cond_latents` (int): The total number of latent tokens added from conditioning items.
Raises:
AssertionError: If input shapes, dimensions, or conditions for applying conditioning are invalid.
"""
assert isinstance(self.vae, CausalVideoAutoencoder)
if conditioning_items:
batch_size, _, num_latent_frames = init_latents.shape[:3]
init_conditioning_mask = torch.zeros(
init_latents[:, 0, :, :, :].shape,
dtype=torch.float32,
device=init_latents.device,
)
extra_conditioning_latents = []
extra_conditioning_pixel_coords = []
extra_conditioning_mask = []
extra_conditioning_num_latents = 0 # Number of extra conditioning latents added (should be removed before decoding)
# Process each conditioning item
for conditioning_item in conditioning_items:
conditioning_item = self._resize_conditioning_item(
conditioning_item, height, width
)
media_item = conditioning_item.media_item
media_frame_number = conditioning_item.media_frame_number
strength = conditioning_item.conditioning_strength
assert media_item.ndim == 5 # (b, c, f, h, w)
b, c, n_frames, h, w = media_item.shape
assert (
height == h and width == w
) or media_frame_number == 0, f"Dimensions do not match: {height}x{width} != {h}x{w} - allowed only when media_frame_number == 0"
assert n_frames % 8 == 1
assert (
media_frame_number >= 0
and media_frame_number + n_frames <= num_frames
)
# Encode the provided conditioning media item
media_item_latents = vae_encode(
media_item.to(dtype=self.vae.dtype, device=self.vae.device),
self.vae,
vae_per_channel_normalize=vae_per_channel_normalize,
).to(dtype=init_latents.dtype)
# Handle the different conditioning cases
if media_frame_number == 0:
# Get the target spatial position of the latent conditioning item
media_item_latents, l_x, l_y = self._get_latent_spatial_position(
media_item_latents,
conditioning_item,
height,
width,
strip_latent_border=True,
)
b, c_l, f_l, h_l, w_l = media_item_latents.shape
media_item_latents_on_device = media_item_latents.to(init_latents.device)
# First frame or sequence - just update the initial noise latents and the mask
init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l] = (
torch.lerp(
init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l],
media_item_latents_on_device,
strength,
)
)
init_conditioning_mask[
:, :f_l, l_y : l_y + h_l, l_x : l_x + w_l
] = strength
else:
# Non-first frame or sequence
if n_frames > 1:
# Handle non-first sequence.
# Encoded latents are either fully consumed, or the prefix is handled separately below.
(
init_latents,
init_conditioning_mask,
media_item_latents,
) = self._handle_non_first_conditioning_sequence(
init_latents,
init_conditioning_mask,
media_item_latents,
media_frame_number,
strength,
)
# Single frame or sequence-prefix latents
if media_item_latents is not None:
noise = randn_tensor(
media_item_latents.shape,
generator=generator,
device=media_item_latents.device,
dtype=media_item_latents.dtype,
)
media_item_latents = torch.lerp(
noise, media_item_latents, strength
)
# Patchify the extra conditioning latents and calculate their pixel coordinates
media_item_latents, latent_coords = self.patchifier.patchify(
latents=media_item_latents
)
pixel_coords = latent_to_pixel_coords(
latent_coords,
self.vae,
causal_fix=self.transformer.config.causal_temporal_positioning,
)
# Update the frame numbers to match the target frame number
pixel_coords[:, 0] += media_frame_number
extra_conditioning_num_latents += media_item_latents.shape[1]
conditioning_mask = torch.full(
media_item_latents.shape[:2],
strength,
dtype=torch.float32,
device=init_latents.device,
)
extra_conditioning_latents.append(media_item_latents)
extra_conditioning_pixel_coords.append(pixel_coords)
extra_conditioning_mask.append(conditioning_mask)
# Patchify the updated latents and calculate their pixel coordinates
init_latents, init_latent_coords = self.patchifier.patchify(
latents=init_latents
)
init_pixel_coords = latent_to_pixel_coords(
init_latent_coords,
self.vae,
causal_fix=self.transformer.config.causal_temporal_positioning,
)
if not conditioning_items:
return init_latents, init_pixel_coords, None, 0
init_conditioning_mask, _ = self.patchifier.patchify(
latents=init_conditioning_mask.unsqueeze(1)
)
init_conditioning_mask = init_conditioning_mask.squeeze(-1)
if extra_conditioning_latents:
# Stack the extra conditioning latents, pixel coordinates and mask
init_latents = torch.cat([*extra_conditioning_latents, init_latents], dim=1)
init_pixel_coords = torch.cat(
[*extra_conditioning_pixel_coords, init_pixel_coords], dim=2
)
init_conditioning_mask = torch.cat(
[*extra_conditioning_mask, init_conditioning_mask], dim=1
)
if self.transformer.use_tpu_flash_attention:
# When flash attention is used, keep the original number of tokens by removing
# tokens from the end.
init_latents = init_latents[:, :-extra_conditioning_num_latents]
init_pixel_coords = init_pixel_coords[
:, :, :-extra_conditioning_num_latents
]
init_conditioning_mask = init_conditioning_mask[
:, :-extra_conditioning_num_latents
]
return (
init_latents,
init_pixel_coords,
init_conditioning_mask,
extra_conditioning_num_latents,
)
@staticmethod
def _resize_conditioning_item(
conditioning_item: ConditioningItem,
height: int,
width: int,
):
if conditioning_item.media_x or conditioning_item.media_y:
raise ValueError(
"Provide media_item in the target size for spatial conditioning."
)
new_conditioning_item = copy.copy(conditioning_item)
new_conditioning_item.media_item = LTXVideoPipeline.resize_tensor(
conditioning_item.media_item, height, width
)
return new_conditioning_item
def _get_latent_spatial_position(
self,
latents: torch.Tensor,
conditioning_item: ConditioningItem,
height: int,
width: int,
strip_latent_border,
):
"""
Get the spatial position of the conditioning item in the latent space.
If requested, strip the conditioning latent borders that do not align with target borders.
(border latents look different then other latents and might confuse the model)
"""
scale = self.vae_scale_factor
h, w = conditioning_item.media_item.shape[-2:]
assert (
h <= height and w <= width
), f"Conditioning item size {h}x{w} is larger than target size {height}x{width}"
assert h % scale == 0 and w % scale == 0
# Compute the start and end spatial positions of the media item
x_start, y_start = conditioning_item.media_x, conditioning_item.media_y
x_start = (width - w) // 2 if x_start is None else x_start
y_start = (height - h) // 2 if y_start is None else y_start
x_end, y_end = x_start + w, y_start + h
assert (
x_end <= width and y_end <= height
), f"Conditioning item {x_start}:{x_end}x{y_start}:{y_end} is out of bounds for target size {width}x{height}"
if strip_latent_border:
# Strip one latent from left/right and/or top/bottom, update x, y accordingly
if x_start > 0:
x_start += scale
latents = latents[:, :, :, :, 1:]
if y_start > 0:
y_start += scale
latents = latents[:, :, :, 1:, :]
if x_end < width:
latents = latents[:, :, :, :, :-1]
if y_end < height:
latents = latents[:, :, :, :-1, :]
return latents, x_start // scale, y_start // scale
@staticmethod
def _handle_non_first_conditioning_sequence(
init_latents: torch.Tensor,
init_conditioning_mask: torch.Tensor,
latents: torch.Tensor,
media_frame_number: int,
strength: float,
num_prefix_latent_frames: int = 2,
prefix_latents_mode: str = "concat",
prefix_soft_conditioning_strength: float = 0.15,
):
"""
Special handling for a conditioning sequence that does not start on the first frame.
The special handling is required to allow a short encoded video to be used as middle
(or last) sequence in a longer video.
Args:
init_latents (torch.Tensor): The initial noise latents to be updated.
init_conditioning_mask (torch.Tensor): The initial conditioning mask to be updated.
latents (torch.Tensor): The encoded conditioning item.
media_frame_number (int): The target frame number of the first frame in the conditioning sequence.
strength (float): The conditioning strength for the conditioning latents.
num_prefix_latent_frames (int, optional): The length of the sequence prefix, to be handled
separately. Defaults to 2.
prefix_latents_mode (str, optional): Special treatment for prefix (boundary) latents.
- "drop": Drop the prefix latents.
- "soft": Use the prefix latents, but with soft-conditioning
- "concat": Add the prefix latents as extra tokens (like single frames)
prefix_soft_conditioning_strength (float, optional): The strength of the soft-conditioning for
the prefix latents, relevant if `prefix_latents_mode` is "soft". Defaults to 0.1.
"""
f_l = latents.shape[2]
f_l_p = num_prefix_latent_frames
assert f_l >= f_l_p
assert media_frame_number % 8 == 0
if f_l > f_l_p:
# Insert the conditioning latents **excluding the prefix** into the sequence
f_l_start = media_frame_number // 8 + f_l_p
f_l_end = f_l_start + f_l - f_l_p
init_latents[:, :, f_l_start:f_l_end] = torch.lerp(
init_latents[:, :, f_l_start:f_l_end],
latents[:, :, f_l_p:],
strength,
)
# Mark these latent frames as conditioning latents
init_conditioning_mask[:, f_l_start:f_l_end] = strength
# Handle the prefix-latents
if prefix_latents_mode == "soft":
if f_l_p > 1:
# Drop the first (single-frame) latent and soft-condition the remaining prefix
f_l_start = media_frame_number // 8 + 1
f_l_end = f_l_start + f_l_p - 1
strength = min(prefix_soft_conditioning_strength, strength)
init_latents[:, :, f_l_start:f_l_end] = torch.lerp(
init_latents[:, :, f_l_start:f_l_end],
latents[:, :, 1:f_l_p],
strength,
)
# Mark these latent frames as conditioning latents
init_conditioning_mask[:, f_l_start:f_l_end] = strength
latents = None # No more latents to handle
elif prefix_latents_mode == "drop":
# Drop the prefix latents
latents = None
elif prefix_latents_mode == "concat":
# Pass-on the prefix latents to be handled as extra conditioning frames
latents = latents[:, :, :f_l_p]
else:
raise ValueError(f"Invalid prefix_latents_mode: {prefix_latents_mode}")
return (
init_latents,
init_conditioning_mask,
latents,
)
def trim_conditioning_sequence(
self, start_frame: int, sequence_num_frames: int, target_num_frames: int
):
"""
Trim a conditioning sequence to the allowed number of frames.
Args:
start_frame (int): The target frame number of the first frame in the sequence.
sequence_num_frames (int): The number of frames in the sequence.
target_num_frames (int): The target number of frames in the generated video.
Returns:
int: updated sequence length
"""
scale_factor = self.video_scale_factor
num_frames = min(sequence_num_frames, target_num_frames - start_frame)
# Trim down to a multiple of temporal_scale_factor frames plus 1
num_frames = (num_frames - 1) // scale_factor * scale_factor + 1
return num_frames
@staticmethod
def tone_map_latents(
latents: torch.Tensor,
compression: float,
) -> torch.Tensor:
"""
Applies a non-linear tone-mapping function to latent values to reduce their dynamic range
in a perceptually smooth way using a sigmoid-based compression.
This is useful for regularizing high-variance latents or for conditioning outputs
during generation, especially when controlling dynamic behavior with a `compression` factor.
Parameters:
----------
latents : torch.Tensor
Input latent tensor with arbitrary shape. Expected to be roughly in [-1, 1] or [0, 1] range.
compression : float
Compression strength in the range [0, 1].
- 0.0: No tone-mapping (identity transform)
- 1.0: Full compression effect
Returns:
-------
torch.Tensor
The tone-mapped latent tensor of the same shape as input.
"""
if not (0 <= compression <= 1):
raise ValueError("Compression must be in the range [0, 1]")
# Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot
scale_factor = compression * 0.75
abs_latents = torch.abs(latents)
# Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0
# When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect
sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0))
scales = 1.0 - 0.8 * scale_factor * sigmoid_term
filtered = latents * scales
return filtered
def adain_filter_latent(
latents: torch.Tensor, reference_latents: torch.Tensor, factor=1.0
):
"""
Applies Adaptive Instance Normalization (AdaIN) to a latent tensor based on
statistics from a reference latent tensor.
Args:
latent (torch.Tensor): Input latents to normalize
reference_latent (torch.Tensor): The reference latents providing style statistics.
factor (float): Blending factor between original and transformed latent.
Range: -10.0 to 10.0, Default: 1.0
Returns:
torch.Tensor: The transformed latent tensor
"""
result = latents.clone()
for i in range(latents.size(0)):
for c in range(latents.size(1)):
r_sd, r_mean = torch.std_mean(
reference_latents[i, c], dim=None
) # index by original dim order
i_sd, i_mean = torch.std_mean(result[i, c], dim=None)
result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean
result = torch.lerp(latents, result, factor)
return result
class LTXMultiScalePipeline:
def _upsample_latents(
self, latest_upsampler: LatentUpsampler, latents: torch.Tensor
):
assert latents.device == latest_upsampler.device
latents = un_normalize_latents(
latents, self.vae, vae_per_channel_normalize=True
)
upsampled_latents = latest_upsampler(latents)
upsampled_latents = normalize_latents(
upsampled_latents, self.vae, vae_per_channel_normalize=True
)
return upsampled_latents
def __init__(
self, video_pipeline: LTXVideoPipeline, latent_upsampler: LatentUpsampler
):
self.video_pipeline = video_pipeline
self.vae = video_pipeline.vae
self.latent_upsampler = latent_upsampler
def __call__(
self,
downscale_factor: float,
first_pass: dict,
second_pass: dict,
*args: Any,
**kwargs: Any,
) -> Any:
original_kwargs = kwargs.copy()
original_output_type = kwargs["output_type"]
original_width = kwargs["width"]
original_height = kwargs["height"]
x_width = int(kwargs["width"] * downscale_factor)
downscaled_width = x_width - (x_width % self.video_pipeline.vae_scale_factor)
x_height = int(kwargs["height"] * downscale_factor)
downscaled_height = x_height - (x_height % self.video_pipeline.vae_scale_factor)
kwargs["output_type"] = "latent"
kwargs["width"] = downscaled_width
kwargs["height"] = downscaled_height
kwargs.update(**first_pass)
result = self.video_pipeline(*args, **kwargs)
latents = result.images
upsampled_latents = self._upsample_latents(self.latent_upsampler, latents)
upsampled_latents = adain_filter_latent(
latents=upsampled_latents, reference_latents=latents
)
kwargs = original_kwargs
kwargs["latents"] = upsampled_latents
kwargs["output_type"] = original_output_type
kwargs["width"] = downscaled_width * 2
kwargs["height"] = downscaled_height * 2
kwargs.update(**second_pass)
result = self.video_pipeline(*args, **kwargs)
if original_output_type != "latent":
num_frames = result.images.shape[2]
videos = rearrange(result.images, "b c f h w -> (b f) c h w")
videos = F.interpolate(
videos,
size=(original_height, original_width),
mode="bilinear",
align_corners=False,
)
videos = rearrange(videos, "(b f) c h w -> b c f h w", f=num_frames)
result.images = videos
return result
LTX-Video/ltx_video/inference.py
LTX-Video/ltx_video/inference.py
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
import imageio
import json
import numpy as np
import torch
from safetensors import safe_open
from PIL import Image
import torchvision.transforms.functional as TVF
from transformers import (
T5EncoderModel,
T5Tokenizer,
AutoModelForCausalLM,
AutoProcessor,
AutoTokenizer,
)
from huggingface_hub import hf_hub_download
from dataclasses import dataclass, field
# Make sure you have accelerate installed: pip install accelerate
from ltx_video.models.autoencoders.causal_video_autoencoder import (
CausalVideoAutoencoder,
)
from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier
from ltx_video.models.transformers.transformer3d import Transformer3DModel
from ltx_video.pipelines.pipeline_ltx_video import (
ConditioningItem,
LTXVideoPipeline,
LTXMultiScalePipeline,
)
from ltx_video.schedulers.rf import RectifiedFlowScheduler
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
import ltx_video.pipelines.crf_compressor as crf_compressor
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))
frame_tensor = TVF.to_tensor(image) # PIL -> tensor (C, H, W), [0,1]
frame_tensor = TVF.gaussian_blur(frame_tensor, kernel_size=3, sigma=1.0)
frame_tensor_hwc = frame_tensor.permute(1, 2, 0) # (C, H, W) -> (H, W, C)
frame_tensor_hwc = crf_compressor.compress(frame_tensor_hwc)
frame_tensor = frame_tensor_hwc.permute(2, 0, 1) * 255.0 # (H, W, C) -> (C, H, W)
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 convert_prompt_to_filename(text: str, max_len: int = 20) -> str:
# Remove non-letters and convert to lowercase
clean_text = "".join(
char.lower() for char in text if char.isalpha() or char.isspace()
)
# Split into words
words = clean_text.split()
# Build result string keeping track of length
result = []
current_length = 0
for word in words:
# Add word length plus 1 for underscore (except for first word)
new_length = current_length + len(word)
if new_length <= max_len:
result.append(word)
current_length += len(word)
else:
break
return "-".join(result)
# Generate output video name
def get_unique_filename(
base: str,
ext: str,
prompt: str,
seed: int,
resolution: tuple[int, int, int],
dir: Path,
endswith=None,
index_range=1000,
) -> Path:
base_filename = f"{base}_{convert_prompt_to_filename(prompt, max_len=30)}_{seed}_{resolution[0]}x{resolution[1]}x{resolution[2]}"
for i in range(index_range):
filename = dir / f"{base_filename}_{i}{endswith if endswith else ''}{ext}"
if not os.path.exists(filename):
return filename
raise FileExistsError(
f"Could not find a unique filename after {index_range} attempts."
)
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)
# ==============================================================================
# START OF MODIFIED FUNCTIONS
# ==============================================================================
def create_transformer(ckpt_path: str, precision: str, device_map) -> Transformer3DModel:
"""
Loads the Transformer3DModel using a device_map provided by the caller.
This allows for fine-grained control over device placement and offloading.
"""
torch_dtype = torch.float32
if precision == "bfloat16":
torch_dtype = torch.bfloat16
elif precision == "float16":
torch_dtype = torch.float16
return Transformer3DModel.from_pretrained(
ckpt_path,
torch_dtype=torch_dtype,
device_map=device_map # Use the device_map passed from the pipeline function
)
def create_ltx_video_pipeline(
ckpt_path: str,
precision: str,
text_encoder_model_name_or_path: str,
sampler: Optional[str] = None,
device: Optional[str] = None, # The 'device' argument is no longer necessary for loading
enhance_prompt: bool = False,
prompt_enhancer_image_caption_model_name_or_path: Optional[str] = None,
prompt_enhancer_llm_model_name_or_path: Optional[str] = None,
) -> LTXVideoPipeline:
"""
Creates the LTX Video pipeline using accelerate's automatic device mapping
to efficiently distribute the models across available GPUs and CPU.
"""
ckpt_path = Path(ckpt_path)
assert os.path.exists(
ckpt_path
), f"Ckpt path provided (--ckpt_path) {ckpt_path} does not exist"
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)
# ‼️ Key Change: Use "auto" for device_map.
# This lets accelerate distribute the models across GPUs and offload to CPU/disk if needed.
# We create an offload folder to store model parts that are moved to the CPU.
offload_folder = Path("./offload")
offload_folder.mkdir(exist_ok=True)
transformer = Transformer3DModel.from_pretrained(
ckpt_path,
torch_dtype=torch.bfloat16 if precision == "bfloat16" else torch.float16,
device_map="auto",
offload_folder=offload_folder,
)
vae = CausalVideoAutoencoder.from_pretrained(
ckpt_path,
device_map="auto",
offload_folder=offload_folder,
)
text_encoder = T5EncoderModel.from_pretrained(
text_encoder_model_name_or_path,
subfolder="text_encoder",
device_map="auto",
offload_folder=offload_folder,
)
if sampler == "from_checkpoint" or not sampler:
scheduler = RectifiedFlowScheduler.from_pretrained(ckpt_path)
else:
scheduler = RectifiedFlowScheduler(
sampler=("Uniform" if sampler.lower() == "uniform" else "LinearQuadratic")
)
patchifier = SymmetricPatchifier(patch_size=1)
tokenizer = T5Tokenizer.from_pretrained(
text_encoder_model_name_or_path, subfolder="tokenizer"
)
# The rest of the function can remain the same, but ensure device_map="auto" is used for other models if they are large
if enhance_prompt:
prompt_enhancer_image_caption_model = AutoModelForCausalLM.from_pretrained(
prompt_enhancer_image_caption_model_name_or_path,
trust_remote_code=True,
device_map="auto",
offload_folder=offload_folder,
)
prompt_enhancer_image_caption_processor = AutoProcessor.from_pretrained(
prompt_enhancer_image_caption_model_name_or_path, trust_remote_code=True
)
prompt_enhancer_llm_model = AutoModelForCausalLM.from_pretrained(
prompt_enhancer_llm_model_name_or_path,
torch_dtype=torch.bfloat16,
device_map="auto",
offload_folder=offload_folder,
)
prompt_enhancer_llm_tokenizer = AutoTokenizer.from_pretrained(
prompt_enhancer_llm_model_name_or_path,
)
else:
prompt_enhancer_image_caption_model = None
prompt_enhancer_image_caption_processor = None
prompt_enhancer_llm_model = None
prompt_enhancer_llm_tokenizer = None
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)
return pipeline
# ============================================================================
# END OF MODIFIED FUNCTIONS
# (The rest of your file remains the same)
# ============================================================================
def create_latent_upsampler(latent_upsampler_model_path: str, device: str):
latent_upsampler = LatentUpsampler.from_pretrained(latent_upsampler_model_path)
latent_upsampler.to(device)
latent_upsampler.eval()
return latent_upsampler
def load_pipeline_config(pipeline_config: str):
current_file = Path(__file__)
path = None
if os.path.isfile(current_file.parent / pipeline_config):
path = current_file.parent / pipeline_config
elif os.path.isfile(pipeline_config):
path = pipeline_config
else:
raise ValueError(f"Pipeline config file {pipeline_config} does not exist")
with open(path, "r") as f:
return yaml.safe_load(f)
@dataclass
class InferenceConfig:
prompt: str = field(metadata={"help": "Prompt for the generation"})
output_path: str = field(
default_factory=lambda: Path(
f"outputs/{datetime.today().strftime('%Y-%m-%d')}"
),
metadata={"help": "Path to the folder to save the output video"},
)
# Pipeline settings
pipeline_config: str = field(
default="configs/ltxv-13b-0.9.7-dev.yaml",
metadata={"help": "Path to the pipeline config file"},
)
seed: int = field(
default=171198, metadata={"help": "Random seed for the inference"}
)
height: int = field(
default=704, metadata={"help": "Height of the output video frames"}
)
width: int = field(
default=1216, metadata={"help": "Width of the output video frames"}
)
num_frames: int = field(
default=121,
metadata={"help": "Number of frames to generate in the output video"},
)
frame_rate: int = field(
default=30, metadata={"help": "Frame rate for the output video"}
)
offload_to_cpu: bool = field(
default=False, metadata={"help": "Offloading unnecessary computations to CPU."}
)
negative_prompt: str = field(
default="worst quality, inconsistent motion, blurry, jittery, distorted",
metadata={"help": "Negative prompt for undesired features"},
)
# Video-to-video arguments
input_media_path: Optional[str] = field(
default=None,
metadata={
"help": "Path to the input video (or image) to be modified using the video-to-video pipeline"
},
)
# Conditioning
image_cond_noise_scale: float = field(
default=0.15,
metadata={"help": "Amount of noise to add to the conditioned image"},
)
conditioning_media_paths: Optional[List[str]] = field(
default=None,
metadata={
"help": "List of paths to conditioning media (images or videos). Each path will be used as a conditioning item."
},
)
conditioning_strengths: Optional[List[float]] = field(
default=None,
metadata={
"help": "List of conditioning strengths (between 0 and 1) for each conditioning item. Must match the number of conditioning items."
},
)
conditioning_start_frames: Optional[List[int]] = field(
default=None,
metadata={
"help": "List of frame indices where each conditioning item should be applied. Must match the number of conditioning items."
},
)
def infer(config: InferenceConfig):
pipeline_config = load_pipeline_config(config.pipeline_config)
ltxv_model_name_or_path = pipeline_config["checkpoint_path"]
if not os.path.isfile(ltxv_model_name_or_path):
ltxv_model_path = hf_hub_download(
repo_id="Lightricks/LTX-Video",
filename=ltxv_model_name_or_path,
repo_type="model",
)
else:
ltxv_model_path = ltxv_model_name_or_path
spatial_upscaler_model_name_or_path = pipeline_config.get(
"spatial_upscaler_model_path"
)
if spatial_upscaler_model_name_or_path and not os.path.isfile(
spatial_upscaler_model_name_or_path
):
spatial_upscaler_model_path = hf_hub_download(
repo_id="Lightricks/LTX-Video",
filename=spatial_upscaler_model_name_or_path,
repo_type="model",
)
else:
spatial_upscaler_model_path = spatial_upscaler_model_name_or_path
conditioning_media_paths = config.conditioning_media_paths
conditioning_strengths = config.conditioning_strengths
conditioning_start_frames = config.conditioning_start_frames
# 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 >= config.num_frames for f in conditioning_start_frames):
raise ValueError(
f"All conditioning start frames must be between 0 and {config.num_frames-1}"
)
seed_everething(config.seed)
if config.offload_to_cpu and not torch.cuda.is_available():
logger.warning(
"offload_to_cpu is set to True, but offloading will not occur since the model is already running on CPU."
)
offload_to_cpu = False
else:
offload_to_cpu = config.offload_to_cpu and get_total_gpu_memory() < 30
provided_path = Path(config.output_path)
# 检查用户提供的路径是否包含文件后缀 (e.g., .mp4, .png)
if provided_path.suffix:
# 如果是文件路径, 则将其作为最终输出文件名
final_output_path = provided_path
# 输出目录是该文件的父目录
output_dir = final_output_path.parent
else:
# 如果是目录路径 (原始行为), 则将其设为输出目录
output_dir = provided_path
# 最终文件名将在稍后生成，先设为 None
final_output_path = None
# 确保输出目录存在
output_dir.mkdir(parents=True, exist_ok=True)
# Adjust dimensions to be divisible by 32 and num_frames to be (N * 8 + 1)
height_padded = ((config.height - 1) // 32 + 1) * 32
width_padded = ((config.width - 1) // 32 + 1) * 32
num_frames_padded = ((config.num_frames - 2) // 8 + 1) * 8 + 1
padding = calculate_padding(
config.height, config.width, height_padded, width_padded
)
logger.warning(
f"Padded dimensions: {height_padded}x{width_padded}x{num_frames_padded}"
)
device = get_device()
prompt_enhancement_words_threshold = pipeline_config[
"prompt_enhancement_words_threshold"
]
prompt_word_count = len(config.prompt.split())
enhance_prompt = (
prompt_enhancement_words_threshold > 0
and prompt_word_count < prompt_enhancement_words_threshold
)
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."
)
precision = pipeline_config["precision"]
text_encoder_model_name_or_path = pipeline_config["text_encoder_model_name_or_path"]
sampler = pipeline_config.get("sampler", None)
prompt_enhancer_image_caption_model_name_or_path = pipeline_config[
"prompt_enhancer_image_caption_model_name_or_path"
]
prompt_enhancer_llm_model_name_or_path = pipeline_config[
"prompt_enhancer_llm_model_name_or_path"
]
pipeline = create_ltx_video_pipeline(
ckpt_path=ltxv_model_path,
precision=precision,
text_encoder_model_name_or_path=text_encoder_model_name_or_path,
sampler=sampler,
device=device,
enhance_prompt=enhance_prompt,
prompt_enhancer_image_caption_model_name_or_path=prompt_enhancer_image_caption_model_name_or_path,
prompt_enhancer_llm_model_name_or_path=prompt_enhancer_llm_model_name_or_path,
)
if pipeline_config.get("pipeline_type", None) == "multi-scale":
if not spatial_upscaler_model_path:
raise ValueError(
"spatial upscaler model path is missing from pipeline config file and is required for multi-scale rendering"
)
latent_upsampler = create_latent_upsampler(
spatial_upscaler_model_path, pipeline.device
)
pipeline = LTXMultiScalePipeline(pipeline, latent_upsampler=latent_upsampler)
media_item = None
if config.input_media_path:
media_item = load_media_file(
media_path=config.input_media_path,
height=config.height,
width=config.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=config.height,
width=config.width,
num_frames=config.num_frames,
padding=padding,
pipeline=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": config.prompt,
"prompt_attention_mask": None,
"negative_prompt": config.negative_prompt,
"negative_prompt_attention_mask": None,
}
# The `device` for the generator should be the main GPU where inference happens
if hasattr(pipeline, "transformer"):
generator_device = pipeline.transformer.device
elif hasattr(pipeline, "inner_pipeline") and hasattr(pipeline.inner_pipeline, "transformer"):
generator_device = pipeline.inner_pipeline.transformer.device
else:
generator_device = get_device()
generator = torch.Generator(device=generator_device).manual_seed(config.seed)
images = pipeline(
**pipeline_config,
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=config.frame_rate,
**sample,
media_items=media_item,
conditioning_items=conditioning_items,
is_video=True,
vae_per_channel_normalize=True,
image_cond_noise_scale=config.image_cond_noise_scale,
mixed_precision=(precision == "mixed_precision"),
offload_to_cpu=offload_to_cpu,
device=generator_device, # Pass the main device to the pipeline
enhance_prompt=enhance_prompt,
).images
# 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[:, :, : config.num_frames, pad_top:pad_bottom, pad_left:pad_right]
for i in range(images.shape[0]):
# ...
video_np = images[i].permute(1, 2, 3, 0).cpu().float().numpy()
video_np = (video_np * 255).astype(np.uint8)
fps = config.frame_rate
height, width = video_np.shape[1:3]
# --- 开始修改 ---
# 确定要使用的输出文件名
if final_output_path:
# 如果在函数开头已经确定了完整路径, 直接使用它
# 注意：在多批次处理时，为避免覆盖，可以添加索引 i
if images.shape[0] > 1:
# 给后续文件添加索引以防万一
output_filename = final_output_path.with_name(f"{final_output_path.stem}_{i}{final_output_path.suffix}")
else:
output_filename = final_output_path
else:
# 否则 (如果输入的是目录), 调用原始函数生成唯一文件名
if video_np.shape[0] == 1: # 图片情况
output_filename = get_unique_filename(
f"image_output_{i}", ".png", prompt=config.prompt,
seed=config.seed, resolution=(height, width, config.num_frames), dir=output_dir
)
else: # 视频情况
output_filename = get_unique_filename(
f"video_output_{i}", ".mp4", prompt=config.prompt,
seed=config.seed, resolution=(height, width, config.num_frames), dir=output_dir
)
# 写入文件（这部分逻辑不变）
if video_np.shape[0] == 1:
imageio.imwrite(output_filename, video_np[0])
else:
with imageio.get_writer(output_filename, fps=fps) as video:
for frame in video_np:
video.append_data(frame)
# --- 结束修改 ---
logger.warning(f"Output saved to {output_filename}")
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
):
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:
is_video = any(
media_path.lower().endswith(ext) for ext in [".mp4", ".avi", ".mov", ".mkv"]
)
num_frames = 1
if is_video:
reader = imageio.get_reader(media_path)
num_frames = reader.count_frames()
reader.close()
return num_frames
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:
is_video = any(
media_path.lower().endswith(ext) for ext in [".mp4", ".avi", ".mov", ".mkv"]
)
if is_video:
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: # 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)
return media_tensor一键启动脚本（加换脸）
# ==============================================================================
# --- STAGE 2: BATCH VIDEO GENERATION ---
# ==============================================================================
echo ""
echo "=============================================================================="
echo "🎬 STAGE 2: KICKING OFF BATCH VIDEO GENERATION"
echo "=============================================================================="
export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_OFFLINE=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export CUDA_VISIBLE_DEVICES=0,1
cd /home/yanchang/DATA/AIvideo/LTX-Video
PROMPT="A girl"
NEGATIVE_PROMPT="(((deformed))), blurry, over saturation, bad anatomy, disfigured, poorly drawn face, mutation, mutated, (extra_limb), (ugly), (poorly drawn hands), fused fingers, messy drawing, broken legsm,(((Wearing clothes))),((No genitals shown))"
HEIGHT=512
WIDTH=512
NUM_FRAMES=65
FRAME_RATE=16
# SEED=42 # <--- 修改点 1: 删除固定的种子
PIPELINE_CONFIG="configs/ltxv-13b-0.9.8-distilled.yaml"
INPUT_DIR='/home/yanchang/DATA/AIvideo/LTX-Video/target'
OUTPUT_DIR='/home/yanchang/DATA/AIvideo/LTX-Video/outputs'
if [ ! -d "$INPUT_DIR" ]; then
echo "错误: 输入目录 '$INPUT_DIR' 不存在。"
exit 1
fi
mkdir -p "$OUTPUT_DIR"
echo "视频输出将被保存到: $OUTPUT_DIR"
echo "开始批量处理图片生成视频..."
STAGE2_START_TIME=$(date +%s)
video_count=0
total_video_processing_time=0
while IFS= read -r image_path; do
if [ -f "$image_path" ]; then
video_count=$((video_count + 1))
# <--- 修改点 2: 在每次循环开始时生成一个新的随机种子
current_seed=$RANDOM
echo ""
echo "================================================================"
echo "🎬 正在处理图片 #$video_count: $(basename "$image_path")"
echo "================================================================"
image_filename_with_ext=$(basename "$image_path")
video_basename="${image_filename_with_ext%.*}"
output_video_path="${OUTPUT_DIR}/${video_basename}.mp4"
echo " ↳ 准备生成视频: ${output_video_path}"
# <--- 修改点 3 (推荐): 打印出当前使用的种子，方便复现
echo " 🌱 使用随机种子: ${current_seed}"
start_time_video=$(date +%s)
PYTHONUNBUFFERED=1 conda run -n ltxvideo_test python inference.py \
--prompt "$PROMPT" \
--negative_prompt "$NEGATIVE_PROMPT" \
--conditioning_media_paths "$image_path" \
--conditioning_start_frames 0 \
--height $HEIGHT \
--width $WIDTH \
--num_frames $NUM_FRAMES \
--frame_rate $FRAME_RATE \
--seed "$current_seed" \
--pipeline_config "$PIPELINE_CONFIG" \
--output_path "$output_video_path" \
--offload_to_cpu True
end_time_video=$(date +%s)
duration_video=$((end_time_video - start_time_video))
total_video_processing_time=$((total_video_processing_time + duration_video))
if [ $? -eq 0 ]; then
echo "✅ 成功生成视频. 耗时: ${duration_video} 秒."
else
echo "❌ 生成视频时发生错误. 耗时: ${duration_video} 秒."
fi
fi
done < <(find "$INPUT_DIR" -maxdepth 1 -type f \( -iname "*.png" -o -iname "*.jpg" -o -iname "*.jpeg" -o -iname "*.webp" \))
echo ""
echo "🎉 批量视频处理完成！"
STAGE2_END_TIME=$(date +%s)
STAGE2_TOTAL_DURATION=$((STAGE2_END_TIME - STAGE2_START_TIME))
STAGE2_AVG_TIME=0
if [ "$video_count" -gt 0 ]; then
STAGE2_AVG_TIME=$((total_video_processing_time / video_count))
fi
echo "--------------------------------------------------------"
echo "📊 STAGE 2: VIDEO GENERATION SUMMARY (CORRECTED)"
echo "--------------------------------------------------------"
echo "Total videos processed: ${video_count}"
echo "Total wall-clock time for video generation: ${STAGE2_TOTAL_DURATION} seconds."
echo "Sum of individual video processing times: ${total_video_processing_time} seconds."
echo "True average time per video: ${STAGE2_AVG_TIME} seconds."
echo "--------------------------------------------------------"
# ==============================================================================
# --- STAGE 3: FACEFUSION PROCESSING ---
# ==============================================================================
echo ""
echo "=============================================================================="
echo "🎭 STAGE 3: KICKING OFF FACEFUSION PROCESSING"
echo "=============================================================================="
cd /home/yanchang/DATA/facefusion
echo "切换到 $(pwd) 目录。"
echo "正在使用 'face' Conda 环境运行 start.py..."
STAGE3_START_TIME=$(date +%s)
# MODIFICATION: Added 'PYTHONUNBUFFERED=1' to show conda logs in real-time.
PYTHONUNBUFFERED=1 conda run -n face python start.py
STAGE3_END_TIME=$(date +%s)
STAGE3_TOTAL_DURATION=$((STAGE3_END_TIME - STAGE3_START_TIME))
echo "start.py 脚本已完成。"
echo "--------------------------------------------------------"
echo "📊 STAGE 3: FACEFUSION SUMMARY"
echo "--------------------------------------------------------"
echo "Total time for facefusion processing: ${STAGE3_TOTAL_DURATION} seconds."
echo "结果保存在: /home/yanchang/DATA/facefusion/data/output"
echo "--------------------------------------------------------"
# ==============================================================================
# --- FINAL SCRIPT SUMMARY ---
# ==============================================================================
GRAND_END_TIME=$(date +%s)
GRAND_TOTAL_DURATION=$((GRAND_END_TIME - GRAND_START_TIME))
echo ""
echo "=============================================================================="
echo "🏁 ENTIRE SCRIPT FINISHED 🏁"
echo "=============================================================================="
echo "Total execution time for all stages: ${GRAND_TOTAL_DURATION} seconds."
echo "=============================================================================="整体项目修改后的压缩包：
一定先去克隆原来的仓库然后再去下载压缩包替换文件