storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-28 08:57:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

modify the forward pass

Browse Source
This commit is contained in:
CrossProduct
2025-01-22 06:56:22 +00:00
parent 07e755c2b3
commit 981767154b
4 changed files with 111 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py
+69 -68
View File
@@ -60,6 +60,7 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
@classmethod @classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
# First create an empty model with the desired architecture # First create an empty model with the desired architecture
print("Pretrain Loading")
model = cls() model = cls()
model_dict = model.state_dict() model_dict = model.state_dict()
# # Then load the pretrained weights into it # # Then load the pretrained weights into it
@@ -78,90 +79,90 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
model.load_state_dict(model_dict) model.load_state_dict(model_dict)
return model return model
def forward( def forward(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor, timestep: torch.LongTensor,
encoder_attention_mask: torch.Tensor, encoder_attention_mask: torch.Tensor,
residual_x: torch.Tensor, residual_x: torch.Tensor,
num_frames: int, num_frames: int,
height: int, height: int,
width: int, width: int,
rope_interpolation_scale: Optional[Tuple[float, float, float]] = None, rope_interpolation_scale: Optional[Tuple[float, float, float]] = None,
attention_kwargs: Optional[Dict[str, Any]] = None, attention_kwargs: Optional[Dict[str, Any]] = None,
return_dict: bool = True return_dict: bool = True
) -> torch.Tensor: ) -> 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 # 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: 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 = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
encoder_attention_mask = encoder_attention_mask.unsqueeze(1) 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 # inject the condition and the residual then project it into the pretrained proj_in
hidden_states = self.adapter(residual_x=residual_x, hidden_states = self.adapter(residual_x=residual_x,
conditioned_x=hidden_states) conditioned_x=hidden_states)
hidden_states = self.proj_in(hidden_states) hidden_states = self.proj_in(hidden_states)
temb, embedded_timestep = self.time_embed( temb, embedded_timestep = self.time_embed(
timestep.flatten(), timestep.flatten(),
batch_size=batch_size, batch_size=batch_size,
hidden_dtype=hidden_states.dtype, hidden_dtype=hidden_states.dtype,
) )
temb = temb.view(batch_size, -1, temb.size(-1)) temb = temb.view(batch_size, -1, temb.size(-1))
embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.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 = self.caption_projection(encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
for block in self.transformer_blocks: for block in self.transformer_blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing: if torch.is_grad_enabled() and self.gradient_checkpointing:
def create_custom_forward(module, return_dict=None): def create_custom_forward(module, return_dict=None):
def custom_forward(*inputs): def custom_forward(*inputs):
if return_dict is not None: if return_dict is not None:
return module(*inputs, return_dict=return_dict) return module(*inputs, return_dict=return_dict)
else: else:
return module(*inputs) 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 {} ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
hidden_states = torch.utils.checkpoint.checkpoint( hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block), create_custom_forward(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, **ckpt_kwargs,
) )
else: else:
hidden_states = block( hidden_states = block(
hidden_states=hidden_states, hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states, encoder_hidden_states=encoder_hidden_states,
temb=temb, temb=temb,
image_rotary_emb=image_rotary_emb, image_rotary_emb=image_rotary_emb,
encoder_attention_mask=encoder_attention_mask, encoder_attention_mask=encoder_attention_mask,
) )
scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None] scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
hidden_states = self.norm_out(hidden_states) hidden_states = self.norm_out(hidden_states)
hidden_states = hidden_states * (1 + scale) + shift hidden_states = hidden_states * (1 + scale) + shift
output = self.proj_out(hidden_states) output = self.proj_out(hidden_states)
if not return_dict: if not return_dict:
return (output,) return (output,)
return Transformer2DModelOutput(sample=output) return Transformer2DModelOutput(sample=output)
def apply_rotary_emb(x, freqs): def apply_rotary_emb(x, freqs):
finetrainers/ltx_video/full_finetune.py
+3 -2
View File
@@ -12,7 +12,8 @@ from .lora import (
prepare_conditions, prepare_conditions,
prepare_latents, prepare_latents,
validation, 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, "prepare_latents": prepare_latents,
"post_latent_preparation": post_latent_preparation, "post_latent_preparation": post_latent_preparation,
"collate_fn": collate_fn_t2v_cond, "collate_fn": collate_fn_t2v_cond,
"forward_pass": forward_pass, "forward_pass": conditioned_forward_pass,
"validation": validation, "validation": validation,
} }
finetrainers/ltx_video/lora.py
+27
View File
@@ -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( def forward_pass(
transformer: LTXVideoTransformer3DModel, transformer: LTXVideoTransformer3DModel,
finetrainers/trainer.py
+12 -2
View File
@@ -145,12 +145,14 @@ class Trainer:
if not self.args.precompute_conditions: if not self.args.precompute_conditions:
condition_components = self.model_config["load_condition_models"](**load_components_kwargs) condition_components = self.model_config["load_condition_models"](**load_components_kwargs)
latent_components = self.model_config["load_latent_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 = {}
components.update(condition_components) components.update(condition_components)
components.update(latent_components) components.update(latent_components)
components.update(diffusion_components) #bug
#components.update(diffusion_components)
self._set_components(components) self._set_components(components)
if self.vae is not None: if self.vae is not None:
@@ -817,6 +819,13 @@ class Trainer:
patch_size = 1, patch_size = 1,
patch_size_t = 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"]}) latent_conditions.update({"noisy_latents": latent_final["latents"]})
else: else:
# Default to flow-matching noise addition # Default to flow-matching noise addition
@@ -839,6 +848,7 @@ class Trainer:
transformer=self.transformer, transformer=self.transformer,
scheduler=self.scheduler, scheduler=self.scheduler,
timesteps=timesteps, timesteps=timesteps,
residual_x=noisy_latents_residual,
**latent_conditions, **latent_conditions,
**text_conditions, **text_conditions,
) )