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

trying training again but broke something so going back.

Browse Source
This commit is contained in:
CrossProduct
2025-01-28 22:36:09 +00:00
parent dfa03f1ed1
commit 8ec35d58fb
3 changed files with 110 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/conditioning/conditioned_pipeline.py
+56 -46
View File
@@ -19,7 +19,7 @@ from transformers import T5EncoderModel, T5TokenizerFast
from typing import Any, Callable, Dict, List, Optional, Union from typing import Any, Callable, Dict, List, Optional, Union
from finetrainers.conditioning import condition_latents_prepare,post_conditioned_latent_patchify from finetrainers.conditioning import condition_latents_prepare,post_conditioned_latent_patchify
from finetrainers.utils.torch_utils import randn_tensor from diffusers.utils.torch_utils import randn_tensor
if is_torch_xla_available(): if is_torch_xla_available():
import torch_xla.core.xla_model as xm import torch_xla.core.xla_model as xm
@@ -95,12 +95,12 @@ class LTXConditionedPipeline(LTXPipeline):
callback_on_step_end_tensor_inputs: List[str] = ["latents"], callback_on_step_end_tensor_inputs: List[str] = ["latents"],
max_sequence_length: int = 128, max_sequence_length: int = 128,
pose_video=None, pose_video=None,
image_ref_video=None, img_ref_video=None,
): ):
# These are the raw videos # These are the raw videos
if pose_video is None: if pose_video is None:
raise ValueError("pose_video cannot be None.") raise ValueError("pose_video cannot be None.")
if image_ref_video is None: if img_ref_video is None:
raise ValueError("image_ref_video cannot be None.") raise ValueError("image_ref_video cannot be None.")
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
@@ -159,20 +159,20 @@ class LTXConditionedPipeline(LTXPipeline):
num_channels_latents = self.transformer.config.in_channels num_channels_latents = self.transformer.config.in_channels
# TODO this is the noise latent patchified. # # TODO this is the noise latent patchified.
# we have to use the image size by default. # # we have to use the image size by default.
latents = self.prepare_latents( # latents = self.prepare_latents(
batch_size * num_videos_per_prompt, # batch_size * num_videos_per_prompt,
num_channels_latents, # num_channels_latents,
height, # height,
width, # width,
num_frames, # num_frames,
torch.float32, # torch.float32,
device, # device,
generator, # generator,
latents, # latents,
) # creates noise tensor the size suggested by the user # dtype=self.text_encoder.dtype
# ) # creates noise tensor the size suggested by the user
noise_latent = self.noise_condition_latent_prepare( noise_latent = self.noise_condition_latent_prepare(
batch_size * num_videos_per_prompt, batch_size * num_videos_per_prompt,
@@ -185,54 +185,59 @@ class LTXConditionedPipeline(LTXPipeline):
generator, generator,
) )
# device=device,
# dtype=weight_dtype,
# 4. create conditioning latents from the video. # 4. create conditioning latents from the video.
pose_latent = condition_latents_prepare( pose_latent = condition_latents_prepare.prepare_latents_for_conditioning(
vae=self.vae, vae=self.vae,
image_or_video=pose_video, image_or_video=pose_video,
patch_size=self.transformer_config.patch_size, patch_size=self.transformer.config.patch_size,
patch_size_t=self.transformer_config.patch_size_t, patch_size_t=self.transformer.config.patch_size_t,
device=device, device=device,
dtype=None, dtype=torch.float32,
generator=generator, generator=generator,
)["latents"] )["latents"]
img_ref_latent = condition_latents_prepare( img_ref_latent = condition_latents_prepare.prepare_latents_for_conditioning(
vae=self.vae, vae=self.vae,
image_or_video=pose_video, image_or_video=pose_video,
patch_size=self.transformer_config.patch_size, patch_size=self.transformer.config.patch_size,
patch_size_t=self.transformer_config.patch_size_t, patch_size_t=self.transformer.config.patch_size_t,
device=device, device=device,
dtype=None, dtype=torch.float32,
generator=generator, generator=generator,
)["latents"] )["latents"]
# pose template noisey input [cat] img_ref + pose video
# pose template ref video latent + patchify video latent.
# img ref patchify video latent +
# add them together as input
# residual x latent
# add pose information to both channels # add pose information to both channels
pose_noisy_latents = noise_latent + pose_latent noisy_latents = noise_latent # + pose_latent
pose_img_ref_latents = img_ref_latent + pose_latent pose_img_ref_latents = img_ref_latent + pose_latent
# expand channel information # B x 2C latent will be projected to adapter to scale it back to 128d using adapter. # expand channel information # B x 2C latent will be projected to adapter to scale it back to 128d using adapter.
pose_img_ref_cond_latents = torch.cat([pose_img_ref_latents,pose_noisy_latents], dim=1) condition_latent = torch.cat([pose_img_ref_latents,noisy_latents], dim=1)
noisy_latent_tokens = post_conditioned_latent_patchify(latents=noisy_latents, num_frames=num_frames,
height=height,
width=width,
patch_size = 1,
patch_size_t = 1)["latents"]
pose_img_ref_latents = post_conditioned_latent_patchify(latents=pose_img_ref_latents, condition_tokens = post_conditioned_latent_patchify(latents=condition_latent,
num_frames=num_frames, num_frames=num_frames,
height=height, height=height,
width=width, width=width,
patch_size = 1, patch_size = 1,
patch_size_t = 1) patch_size_t = 1)["latents"]
# target video as a input residual # target video as a input residual
noisy_latents_residual = post_conditioned_latent_patchify(latents=noise_latent, noisy_residual_tokens = post_conditioned_latent_patchify(latents=noisy_latents,
num_frames=num_frames, num_frames=num_frames,
height=height, height=height,
width=width, width=width,
patch_size = 1, patch_size = 1,
patch_size_t = 1) patch_size_t = 1)["latents"]
# Need to change the latents to # Need to change the latents to
# 5. Prepare timesteps # 5. Prepare timesteps
@@ -272,15 +277,18 @@ class LTXConditionedPipeline(LTXPipeline):
for i, t in enumerate(timesteps): for i, t in enumerate(timesteps):
if self.interrupt: if self.interrupt:
continue continue
# change this
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents # latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
# latent_model_input = latent_model_input.to(prompt_embeds.dtype)
latent_model_input = noisy_latent_tokens
latent_model_input = latent_model_input.to(prompt_embeds.dtype) latent_model_input = latent_model_input.to(prompt_embeds.dtype)
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timestep = t.expand(latent_model_input.shape[0]) timestep = t.expand(latent_model_input.shape[0])
noise_pred = self.transformer( noise_pred = self.transformer(
hidden_states=latent_model_input, hidden_states=condition_tokens,
encoder_hidden_states=prompt_embeds, encoder_hidden_states=prompt_embeds,
timestep=timestep, timestep=timestep,
encoder_attention_mask=prompt_attention_mask, encoder_attention_mask=prompt_attention_mask,
@@ -290,7 +298,9 @@ class LTXConditionedPipeline(LTXPipeline):
rope_interpolation_scale=rope_interpolation_scale, rope_interpolation_scale=rope_interpolation_scale,
attention_kwargs=attention_kwargs, attention_kwargs=attention_kwargs,
return_dict=False, return_dict=False,
residual_x=noisy_residual_tokens,
)[0] )[0]
noise_pred = noise_pred.float() noise_pred = noise_pred.float()
if self.do_classifier_free_guidance: if self.do_classifier_free_guidance:
@@ -298,7 +308,7 @@ class LTXConditionedPipeline(LTXPipeline):
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1 # compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] noisy_latents = self.scheduler.step(noise_pred, t, noisy_latents, return_dict=False)[0]
if callback_on_step_end is not None: if callback_on_step_end is not None:
callback_kwargs = {} callback_kwargs = {}
@@ -306,7 +316,7 @@ class LTXConditionedPipeline(LTXPipeline):
callback_kwargs[k] = locals()[k] callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents) noisy_latents = callback_outputs.pop("latents", noisy_latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
# call the callback, if provided # call the callback, if provided
@@ -319,18 +329,18 @@ class LTXConditionedPipeline(LTXPipeline):
if output_type == "latent": if output_type == "latent":
video = latents video = latents
else: else:
latents = self._unpack_latents( noisy_latents = self._unpack_latents(
latents, noisy_latents,
latent_num_frames, latent_num_frames,
latent_height, latent_height,
latent_width, latent_width,
self.transformer_spatial_patch_size, self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size, self.transformer_temporal_patch_size,
) )
latents = self._denormalize_latents( noisy_latents = self._denormalize_latents(
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor noisy_latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
) )
latents = latents.to(prompt_embeds.dtype) noisy_latents = noisy_latents.to(prompt_embeds.dtype)
if not self.vae.config.timestep_conditioning: if not self.vae.config.timestep_conditioning:
timestep = None timestep = None
finetrainers/ltx_video/full_finetune_condition.py
+3 -8
View File
@@ -137,7 +137,7 @@ def conditional_validation(
image: Optional[Image.Image] = None, image: Optional[Image.Image] = None,
video: Optional[List[Image.Image]] = None, video: Optional[List[Image.Image]] = None,
pose_video: Optional[List[Image.Image]] = None, pose_video: Optional[List[Image.Image]] = None,
image_ref_video: Optional[List[Image.Image]] = None, img_ref_video: Optional[List[Image.Image]] = None,
height: Optional[int] = None, height: Optional[int] = None,
width: Optional[int] = None, width: Optional[int] = None,
num_frames: Optional[int] = None, num_frames: Optional[int] = None,
@@ -157,15 +157,10 @@ def conditional_validation(
"return_dict": True, "return_dict": True,
"output_type": "pil", "output_type": "pil",
"pose_video": pose_video, "pose_video": pose_video,
"image_ref_video": image_ref_video, "img_ref_video": img_ref_video,
} }
# pose template noisey input [cat] img_ref + pose video
# pose template ref video latent + patchify video latent.
# img ref patchify video latent +
# add them together as input
# residual x latent
generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None} generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}
video = pipeline(**generation_kwargs).frames[0] video = pipeline(**generation_kwargs).frames[0]
finetrainers/trainer.py
+51 -4
View File
@@ -6,7 +6,7 @@ import random
from datetime import datetime, timedelta from datetime import datetime, timedelta
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List from typing import Any, Dict, List
from finetrainers.dataset import ImageOrVideoDataset
import diffusers import diffusers
import torch import torch
import torch.backends import torch.backends
@@ -21,6 +21,10 @@ from accelerate.utils import (
gather_object, gather_object,
set_seed, set_seed,
) )
import decord
from typing import Tuple,Optional
from diffusers import DiffusionPipeline from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict from diffusers.configuration_utils import FrozenDict
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
@@ -816,6 +820,7 @@ class Trainer:
pose_img_ref_latents = torch.cat([pose_img_ref_latents,pose_noisy_latents], dim=1) pose_img_ref_latents = torch.cat([pose_img_ref_latents,pose_noisy_latents], dim=1)
# project this turn into patches # project this turn into patches
# concat tensor
pose_img_ref_latents = post_conditioned_latent_patchify(latents=pose_img_ref_latents, pose_img_ref_latents = post_conditioned_latent_patchify(latents=pose_img_ref_latents,
num_frames=latent_conditions["num_frames"], num_frames=latent_conditions["num_frames"],
height=latent_conditions["height"], height=latent_conditions["height"],
@@ -829,7 +834,7 @@ class Trainer:
width=latent_conditions["width"], width=latent_conditions["width"],
patch_size = 1, patch_size = 1,
patch_size_t = 1) patch_size_t = 1)
# Target Latent Patchified # Target Latent Patchified concat tensor
# pose template noisey input [cat] img_ref + pose video # pose template noisey input [cat] img_ref + pose video
latent_conditions.update({"noisy_latents": pose_img_ref_latents["latents"]}) latent_conditions.update({"noisy_latents": pose_img_ref_latents["latents"]})
# input video noise at level residual information to adapter # input video noise at level residual information to adapter
@@ -1027,9 +1032,9 @@ class Trainer:
if self.pose_condition: if self.pose_condition:
if pose_video is not None: if pose_video is not None:
pose_video = load_video(pose_video) pose_video = self.preprocess_condition_video(pose_video,max_num_frames=self.dataset.max_num_frames)
if img_ref_video is not None: if img_ref_video is not None:
img_ref_video = load_video(img_ref_video) img_ref_video = self.preprocess_condition_video_image_reference_video(path=img_ref_video,resolution_buckets=self.dataset.resolution_buckets,max_num_frames=self.dataset.max_num_frames)
logger.debug( logger.debug(
f"Validating sample {i + 1}/{num_validation_samples} on process {accelerator.process_index}. Prompt: {prompt}", f"Validating sample {i + 1}/{num_validation_samples} on process {accelerator.process_index}. Prompt: {prompt}",
@@ -1053,6 +1058,7 @@ class Trainer:
self.args.seed if self.args.seed is not None else 0 self.args.seed if self.args.seed is not None else 0
), ),
) )
else: else:
validation_artifacts = self.model_config["validation"]( validation_artifacts = self.model_config["validation"](
pipeline=pipeline, pipeline=pipeline,
@@ -1335,3 +1341,44 @@ class Trainer:
for component in components: for component in components:
if component is not None: if component is not None:
component.requires_grad_(True) component.requires_grad_(True)
# Two conditioning functions
@staticmethod
def preprocess_condition_video(path: Path,max_num_frames:int) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
r"""
Loads a single video, or latent and prompt embedding, based on initialization parameters.
Returns a [F, C, H, W] video tensor.
"""
video_reader = decord.VideoReader(uri=Path(path).as_posix())
video_num_frames = len(video_reader)
indices = list(range(0, video_num_frames, video_num_frames // max_num_frames))
frames = video_reader.get_batch(indices)
frames = frames[: max_num_frames].float()
frames = frames.permute(0, 3, 1, 2).contiguous()
frames = torch.stack([frame for frame in frames], dim=0)
return frames
@staticmethod
def preprocess_condition_video_image_reference_video(path: Path,resolution_buckets,max_num_frames) -> torch.Tensor:
video_reader = decord.VideoReader(uri=Path(path).as_posix())
video_num_frames = len(video_reader)
nearest_frame_bucket = min(
[bucket for bucket in resolution_buckets if bucket[0] <= video_num_frames],
key=lambda x: abs(x[0] - min(video_num_frames, max_num_frames)),
default=1,
)[0]
frame_indices = [0 for _ in range(video_num_frames)]
frames = video_reader.get_batch(frame_indices)
frames = frames[:nearest_frame_bucket].float()
frames = frames.permute(0, 3, 1, 2).contiguous()
# nearest_res = self._find_nearest_resolution(frames.shape[2], frames.shape[3])
# frames_resized = torch.stack([frame for frame in frames], dim=0)
frames = torch.stack([frame for frame in frames], dim=0)
return frames