Spaces:
Running
Running
| from typing import Any, Dict, Optional, Tuple | |
| import diffusers | |
| import torch | |
| from diffusers import LTXVideoTransformer3DModel | |
| from diffusers.models.modeling_outputs import Transformer2DModelOutput | |
| from diffusers.utils.import_utils import is_torch_version | |
| def patch_transformer_forward() -> None: | |
| _perform_ltx_transformer_forward_patch() | |
| def patch_apply_rotary_emb_for_tp_compatibility() -> None: | |
| _perform_ltx_apply_rotary_emb_tensor_parallel_compatibility_patch() | |
| def _perform_ltx_transformer_forward_patch() -> None: | |
| LTXVideoTransformer3DModel.forward = _patched_LTXVideoTransformer3Dforward | |
| def _perform_ltx_apply_rotary_emb_tensor_parallel_compatibility_patch() -> None: | |
| def apply_rotary_emb(x, freqs): | |
| cos, sin = freqs | |
| # ======== THIS IS CHANGED FROM THE ORIGINAL IMPLEMENTATION ======== | |
| # The change is made due to unsupported DTensor operation aten.ops.unbind | |
| # FIXME: Once aten.ops.unbind support lands, this will no longer be required | |
| # x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) # [B, S, H, D // 2] | |
| x_real, x_imag = x.unflatten(2, (-1, 2)).chunk(2, dim=-1) # [B, S, H, D // 2] | |
| # ================================================================== | |
| x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) | |
| out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) | |
| return out | |
| diffusers.models.transformers.transformer_ltx.apply_rotary_emb = apply_rotary_emb | |
| def _patched_LTXVideoTransformer3Dforward( | |
| self, | |
| hidden_states: torch.Tensor, | |
| encoder_hidden_states: torch.Tensor, | |
| timestep: torch.LongTensor, | |
| encoder_attention_mask: torch.Tensor, | |
| num_frames: int, | |
| height: int, | |
| width: int, | |
| rope_interpolation_scale: Optional[Tuple[float, float, float]] = None, | |
| return_dict: bool = True, | |
| *args, | |
| **kwargs, | |
| ) -> torch.Tensor: | |
| image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale) | |
| # convert encoder_attention_mask to a bias the same way we do for attention_mask | |
| if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: | |
| encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 | |
| encoder_attention_mask = encoder_attention_mask.unsqueeze(1) | |
| batch_size = hidden_states.size(0) | |
| # ===== This is modified compared to Diffusers ===== | |
| # This is done because the Diffusers pipeline will pass in a 1D tensor for timestep | |
| if timestep.ndim == 1: | |
| timestep = timestep.view(-1, 1, 1).expand(-1, *hidden_states.shape[1:-1], -1) | |
| # ================================================== | |
| temb, embedded_timestep = self.time_embed( | |
| timestep.flatten(), | |
| batch_size=batch_size, | |
| hidden_dtype=hidden_states.dtype, | |
| ) | |
| # ===== This is modified compared to Diffusers ===== | |
| # temb = temb.view(batch_size, -1, temb.size(-1)) | |
| # embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) | |
| # ================================================== | |
| # This is done to make it possible to use per-token timestep embedding | |
| temb = temb.view(batch_size, *hidden_states.shape[1:-1], temb.size(-1)) | |
| embedded_timestep = embedded_timestep.view(batch_size, *hidden_states.shape[1:-1], embedded_timestep.size(-1)) | |
| # ================================================== | |
| hidden_states = self.proj_in(hidden_states) | |
| encoder_hidden_states = self.caption_projection(encoder_hidden_states) | |
| encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) | |
| for block in self.transformer_blocks: | |
| if torch.is_grad_enabled() 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, | |
| encoder_hidden_states, | |
| temb, | |
| image_rotary_emb, | |
| encoder_attention_mask, | |
| **ckpt_kwargs, | |
| ) | |
| else: | |
| hidden_states = block( | |
| hidden_states=hidden_states, | |
| encoder_hidden_states=encoder_hidden_states, | |
| temb=temb, | |
| image_rotary_emb=image_rotary_emb, | |
| encoder_attention_mask=encoder_attention_mask, | |
| ) | |
| scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None] | |
| shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] | |
| hidden_states = self.norm_out(hidden_states) | |
| hidden_states = hidden_states * (1 + scale) + shift | |
| output = self.proj_out(hidden_states) | |
| if not return_dict: | |
| return (output,) | |
| return Transformer2DModelOutput(sample=output) | |