From 981767154be8918735464943399e5106397dc8a7 Mon Sep 17 00:00:00 2001 From: CrossProduct Date: Wed, 22 Jan 2025 06:56:22 +0000 Subject: [PATCH] modify the forward pass --- .../LTXVideoConditionedTransformer3DModel.py | 137 +++++++++--------- finetrainers/ltx_video/full_finetune.py | 5 +- finetrainers/ltx_video/lora.py | 27 ++++ finetrainers/trainer.py | 14 +- 4 files changed, 111 insertions(+), 72 deletions(-) diff --git a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py b/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py index 6d55cad..cd17de6 100644 --- a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py +++ b/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py @@ -60,6 +60,7 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel): @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): # First create an empty model with the desired architecture + print("Pretrain Loading") model = cls() model_dict = model.state_dict() # # Then load the pretrained weights into it @@ -78,90 +79,90 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel): model.load_state_dict(model_dict) return model -def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_attention_mask: torch.Tensor, - residual_x: torch.Tensor, - num_frames: int, - height: int, - width: int, - rope_interpolation_scale: Optional[Tuple[float, float, float]] = None, - attention_kwargs: Optional[Dict[str, Any]] = None, - return_dict: bool = True - ) -> torch.Tensor: + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_attention_mask: torch.Tensor, + residual_x: torch.Tensor, + num_frames: int, + height: int, + width: int, + rope_interpolation_scale: Optional[Tuple[float, float, float]] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True + ) -> torch.Tensor: - image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale) + 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) + # 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) + batch_size = hidden_states.size(0) - # inject the condition and the residual then project it into the pretrained proj_in - hidden_states = self.adapter(residual_x=residual_x, - conditioned_x=hidden_states) + # inject the condition and the residual then project it into the pretrained proj_in + hidden_states = self.adapter(residual_x=residual_x, + conditioned_x=hidden_states) - hidden_states = self.proj_in(hidden_states) + hidden_states = self.proj_in(hidden_states) - temb, embedded_timestep = self.time_embed( - timestep.flatten(), - batch_size=batch_size, - hidden_dtype=hidden_states.dtype, - ) + temb, embedded_timestep = self.time_embed( + timestep.flatten(), + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) - temb = temb.view(batch_size, -1, temb.size(-1)) - embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) + temb = temb.view(batch_size, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) + 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: + 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) + 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 + 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, - ) + 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] + 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) + 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) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) def apply_rotary_emb(x, freqs): diff --git a/finetrainers/ltx_video/full_finetune.py b/finetrainers/ltx_video/full_finetune.py index ec83802..b5553c9 100644 --- a/finetrainers/ltx_video/full_finetune.py +++ b/finetrainers/ltx_video/full_finetune.py @@ -12,7 +12,8 @@ from .lora import ( prepare_conditions, prepare_latents, validation, - collate_fn_t2v_cond + collate_fn_t2v_cond, + conditioned_forward_pass ) @@ -27,6 +28,6 @@ LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = { "prepare_latents": prepare_latents, "post_latent_preparation": post_latent_preparation, "collate_fn": collate_fn_t2v_cond, - "forward_pass": forward_pass, + "forward_pass": conditioned_forward_pass, "validation": validation, } diff --git a/finetrainers/ltx_video/lora.py b/finetrainers/ltx_video/lora.py index 5fe29df..424dc47 100644 --- a/finetrainers/ltx_video/lora.py +++ b/finetrainers/ltx_video/lora.py @@ -210,6 +210,33 @@ def collate_fn_t2v(batch: List[List[Dict[str, torch.Tensor]]]) -> Dict[str, torc } +def conditioned_forward_pass( + transformer: LTXVideoConditionedTransformer3DModel, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + latents: torch.Tensor, + noisy_latents: torch.Tensor, + timesteps: torch.LongTensor, + num_frames: int, + height: int, + width: int, + **kwargs, +) -> torch.Tensor: + rope_interpolation_scale = [1 / 25, 32, 32] + + denoised_latents = transformer( + hidden_states=noisy_latents, + encoder_hidden_states=prompt_embeds, + timestep=timesteps, + encoder_attention_mask=prompt_attention_mask, + num_frames=num_frames, + height=height, + width=width, + rope_interpolation_scale=rope_interpolation_scale, + return_dict=False, + )[0] + + return {"latents": denoised_latents} def forward_pass( transformer: LTXVideoTransformer3DModel, diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 1c3f8d3..9c17640 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -145,12 +145,14 @@ class Trainer: if not self.args.precompute_conditions: condition_components = self.model_config["load_condition_models"](**load_components_kwargs) latent_components = self.model_config["load_latent_models"](**load_components_kwargs) - diffusion_components = self.model_config["load_diffusion_models"](**load_components_kwargs) + # bug + #diffusion_components = self.model_config["load_diffusion_models"](**load_components_kwargs) components = {} components.update(condition_components) components.update(latent_components) - components.update(diffusion_components) + #bug + #components.update(diffusion_components) self._set_components(components) if self.vae is not None: @@ -817,6 +819,13 @@ class Trainer: patch_size = 1, patch_size_t = 1) + noisy_latents_residual = post_conditioned_latent_patchify(latents=noisy_latents, + num_frames=latent_conditions["num_frames"], + height=latent_conditions["height"], + width=latent_conditions["width"], + patch_size = 1, + patch_size_t = 1) + # pose information latent latent_conditions.update({"noisy_latents": latent_final["latents"]}) else: # Default to flow-matching noise addition @@ -839,6 +848,7 @@ class Trainer: transformer=self.transformer, scheduler=self.scheduler, timesteps=timesteps, + residual_x=noisy_latents_residual, **latent_conditions, **text_conditions, )