mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-19 07:43:28 +00:00
Testing the original code path.
Need to check rectified linear flow vs flow matching.
This commit is contained in:
CrossProduct
@@ -0,0 +1 @@
|
||||
from .condition_latents_prepare import post_conditioned_latent_patchify, prepare_latents_for_conditioning
|
||||
@@ -0,0 +1,80 @@
|
||||
from typing import Optional
|
||||
import torch
|
||||
from diffusers import AutoencoderKLLTXVideo
|
||||
|
||||
def _normalize_latents(
|
||||
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
|
||||
) -> torch.Tensor:
|
||||
# Normalize latents across the channel dimension [B, C, F, H, W]
|
||||
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
latents = (latents - latents_mean) * scaling_factor / latents_std
|
||||
return latents
|
||||
|
||||
def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
|
||||
# Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
|
||||
# The patch dimensions are then permuted and collapsed into the channel dimension of shape:
|
||||
# [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor).
|
||||
# dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features
|
||||
batch_size, num_channels, num_frames, height, width = latents.shape
|
||||
post_patch_num_frames = num_frames // patch_size_t
|
||||
post_patch_height = height // patch_size
|
||||
post_patch_width = width // patch_size
|
||||
|
||||
dim1 = num_frames // patch_size_t * height // patch_size * width // patch_size
|
||||
dim2 = num_channels * patch_size_t * patch_size * patch_size
|
||||
|
||||
latents = latents.reshape(
|
||||
batch_size,
|
||||
-1,
|
||||
post_patch_num_frames,
|
||||
patch_size_t,
|
||||
post_patch_height,
|
||||
patch_size,
|
||||
post_patch_width,
|
||||
patch_size,
|
||||
)
|
||||
latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
|
||||
return latents
|
||||
|
||||
def post_conditioned_latent_patchify(
|
||||
latents: torch.Tensor,
|
||||
latents_mean: torch.Tensor,
|
||||
latents_std: torch.Tensor,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
patch_size: int = 1,
|
||||
patch_size_t: int = 1,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
latents = _normalize_latents(latents, latents_mean, latents_std)
|
||||
latents = _pack_latents(latents, patch_size, patch_size_t)
|
||||
return {"latents": latents, "num_frames": num_frames, "height": height, "width": width}
|
||||
|
||||
def prepare_latents_for_conditioning(
|
||||
vae: AutoencoderKLLTXVideo,
|
||||
image_or_video: torch.Tensor,
|
||||
patch_size: int = 1,
|
||||
patch_size_t: int = 1,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
|
||||
) -> torch.Tensor:
|
||||
device = device or vae.device
|
||||
|
||||
if image_or_video.ndim == 4:
|
||||
image_or_video = image_or_video.unsqueeze(2)
|
||||
assert image_or_video.ndim == 5, f"Expected 5D tensor, got {image_or_video.ndim}D tensor"
|
||||
|
||||
image_or_video = image_or_video.to(device=device, dtype=vae.dtype)
|
||||
image_or_video = image_or_video.permute(0, 2, 1, 3, 4).contiguous() # [B, C, F, H, W] -> [B, F, C, H, W]
|
||||
|
||||
latents = vae.encode(image_or_video).latent_dist.sample(generator=generator)
|
||||
|
||||
latents = latents.to(dtype=dtype)
|
||||
_, _, num_frames, height, width = latents.shape
|
||||
latents = _normalize_latents(latents, vae.latents_mean, vae.latents_std)
|
||||
latents = _pack_latents(latents, patch_size, patch_size_t)
|
||||
return {"latents": latents, "num_frames": num_frames, "height": height, "width": width}
|
||||
@@ -0,0 +1,85 @@
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
|
||||
# This class takes the conditioned multichannel video tensor in patchified
|
||||
# Downsamples it using a bottle neck layer
|
||||
# To adapt the input into the diffusion transformer.
|
||||
# You want to use the patchified tensor non conditioned tensor to pass in as a residual before passing this off to the transformer.
|
||||
class ConditionedResidualAdapterBottleneck(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
output_dim: int, # New parameter for output dimension
|
||||
bottleneck_dim: int = None,
|
||||
adapter_dropout: float = 0.1,
|
||||
adapter_init_scale: float = 1e-3,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input_dim: Size of input dimension
|
||||
output_dim: Size of output dimension
|
||||
adapter_dropout: Dropout probability
|
||||
adapter_init_scale: Initial scale for adapter layer parameters
|
||||
use_residual: Whether to use residual connection (only if input_dim == output_dim)
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Down projection
|
||||
self.down_proj = nn.Linear(input_dim, bottleneck_dim)
|
||||
|
||||
# Activation and dropout
|
||||
self.activation = nn.GELU()
|
||||
self.dropout = nn.Dropout(adapter_dropout)
|
||||
|
||||
# Up projection (now projects to output_dim instead of input_dim)
|
||||
self.up_proj = nn.Linear(bottleneck_dim, output_dim)
|
||||
|
||||
# Initialize weights
|
||||
self.down_proj.weight.data.normal_(mean=0.0, std=adapter_init_scale)
|
||||
self.down_proj.bias.data.zero_()
|
||||
self.up_proj.weight.data.normal_(mean=0.0, std=adapter_init_scale)
|
||||
self.up_proj.bias.data.zero_()
|
||||
|
||||
def forward(self,residual_x:torch.tensor, conditioned_x: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
# Down projection
|
||||
hidden_states = self.down_proj(conditioned_x)
|
||||
|
||||
# Activation and dropout
|
||||
hidden_states = self.activation(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
# Up projection to new dimension
|
||||
hidden_states = self.up_proj(hidden_states)
|
||||
|
||||
|
||||
hidden_states = hidden_states + residual_x
|
||||
|
||||
return hidden_states
|
||||
|
||||
# Example usage:
|
||||
def example_usage():
|
||||
# Create a sample input tensor
|
||||
batch_size, seq_length, input_dim = 1, 4086, 256
|
||||
|
||||
output_dim = 128 # Reduced output dimension
|
||||
x = torch.randn(batch_size, seq_length, input_dim)
|
||||
|
||||
res_x = torch.randn(1,4086,128)
|
||||
# Initialize adapter
|
||||
adapter = ConditionedResidualAdapterBottleneck(
|
||||
input_dim=input_dim,
|
||||
output_dim=output_dim,
|
||||
bottleneck_dim=64,
|
||||
adapter_dropout=0.1,
|
||||
adapter_init_scale=1e-3
|
||||
)
|
||||
adapter.requires_grad_(True)
|
||||
|
||||
# Forward pass
|
||||
output = adapter(residual_x=res_x,conditioned_x=x)
|
||||
print(f"Input shape: {x.shape}")
|
||||
print(f"Output shape: {output.shape}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
example_usage()
|
||||
@@ -0,0 +1,35 @@
|
||||
import torch
|
||||
from diffusers import LTXPipeline
|
||||
from diffusers import LTXVideoTransformer3DModel
|
||||
from diffusers.utils import export_to_video
|
||||
from diffusers.utils import testing_utils
|
||||
testing_utils.enable_full_determinism()
|
||||
|
||||
transformer = LTXVideoTransformer3DModel.from_pretrained("/home/ubuntu/upstream/finetrainers/ltx-video/ltxv_strip/checkpoint-38000/transformer/",torch_dtype=torch.bfloat16).to("cuda")
|
||||
pipe = LTXPipeline.from_pretrained(
|
||||
"Lightricks/LTX-Video", torch_dtype=torch.bfloat16,
|
||||
# transformer=transformer
|
||||
).to("cuda")
|
||||
|
||||
prompt = "A woman semi nude in red lingerie with black straps and garter belts walks through a modern, minimalist interior featuring a white sofa, black chairs, and a glass partition. Her confident and alluring demeanor is highlighted by the soft lighting and the room's neutral tones. As she strides forward, her bare feet touch the carpet, and her blonde hair cascades over her shoulders. The setting includes a sleek black dining table and a bookshelf filled with books and decorative items, creating an atmosphere of sophistication and intimacy."
|
||||
|
||||
g = torch.Generator(device="cuda")
|
||||
g.manual_seed(42)
|
||||
|
||||
for seconds in range(3, 8):
|
||||
num_frames = seconds * 24
|
||||
|
||||
result = pipe(
|
||||
prompt=prompt,
|
||||
width=768,
|
||||
height=512,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=50,
|
||||
generator=g
|
||||
)
|
||||
video = result.frames[0]
|
||||
export_to_video(video, f"out_{seconds}.mp4", fps=24)
|
||||
|
||||
|
||||
# pipe.load_lora_weights("/home/ubuntu/finetrainers/ltx-video/ltxv_strip/checkpoint-10000/pytorch_lora_weights.safetensors", adapter_name="ltxv-lora")
|
||||
# pipe.set_adapters(["ltxv-lora"], [1.0])
|
||||
@@ -0,0 +1,72 @@
|
||||
import torch
|
||||
|
||||
# single channel latent 4608x128 patchified
|
||||
# doubled channel ...
|
||||
# torch.Size([1, 4608, 256])
|
||||
# modify the linear layer ... (4608x256 and 128x2048)
|
||||
|
||||
# Pipeline
|
||||
# Video Input
|
||||
# Video tensor torch.Size([1, 96, 3, 512, 768]) B F C H W
|
||||
# to
|
||||
# Video tensor permuted torch.Size([1, 3, 96, 512, 768]) B C F H W
|
||||
|
||||
# Latent Encode temporal frames 8 compression 32 compression spatio
|
||||
# channels go to 128
|
||||
# torch.Size([1, 128, 12, 16, 24])
|
||||
|
||||
# Creating a Multi Channel Single Video
|
||||
a = torch.randn(1,2,2,2) # video
|
||||
print(a)
|
||||
|
||||
a = a.flatten(0,3)
|
||||
print(a)
|
||||
print(a.shape)
|
||||
|
||||
# b = torch.randn(1,2,4,4) # conditioning
|
||||
# c = torch.cat([a,b],dim=1)
|
||||
|
||||
# Final shape.
|
||||
# print(c.shape)
|
||||
# print(c)
|
||||
|
||||
|
||||
#_prepare_latents
|
||||
# prepare_latents
|
||||
|
||||
|
||||
# ImageOrVideoDataset Shape result
|
||||
# _preprocess_video [F, C, H, W] video tensor 4
|
||||
|
||||
# _preprocess_image [1, C, H, W] (1-frame video)
|
||||
|
||||
# VAE Latent Shape Results Dims
|
||||
# 4 second video 24 fps
|
||||
# torch.Size([1, 128, 12, 16, 24])
|
||||
|
||||
|
||||
# VideoProcessor latent output shape preprocess_video
|
||||
|
||||
# Video Tensor Shape
|
||||
|
||||
# Latent Tensor Shape
|
||||
|
||||
# VAE Patchification Result Shape
|
||||
|
||||
# Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
|
||||
# The patch dimensions are then permuted and collapsed into the channel dimension of shape:
|
||||
# [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor).
|
||||
# dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features
|
||||
|
||||
# encode and then decode and decouple the videos
|
||||
|
||||
# check transformer hidden state, should still be same sequence or slightly larger sequence.
|
||||
|
||||
def patchify_multichannel_latent(video_latent:torch.tensor,patch_size:int):# B, C, F, W, H
|
||||
pass
|
||||
|
||||
def unpatchify_multichannel_laten(video_latent:torch.tensor,patch_size:int): # B, C ,F, W, H
|
||||
pass
|
||||
|
||||
def mask_multichannel_latent(mask_shape:torch.tensor):
|
||||
pass
|
||||
+67
-46
@@ -57,7 +57,7 @@ from .utils.model_utils import resolve_vae_cls_from_ckpt_path
|
||||
from .utils.optimizer_utils import get_optimizer
|
||||
from .utils.torch_utils import align_device_and_dtype, expand_tensor_dims, unwrap_model
|
||||
|
||||
|
||||
from .conditioning import condition_latents_prepare,post_conditioned_latent_patchify
|
||||
logger = get_logger("finetrainers")
|
||||
logger.setLevel(FINETRAINERS_LOG_LEVEL)
|
||||
|
||||
@@ -673,16 +673,26 @@ class Trainer:
|
||||
if self.args.caption_dropout_technique == "empty":
|
||||
if random.random() < self.args.caption_dropout_p:
|
||||
prompts = [""] * batch_size
|
||||
|
||||
latent_conditions = self.model_config["prepare_latents"](
|
||||
vae=self.vae,
|
||||
image_or_video=videos,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=self.state.generator,
|
||||
)
|
||||
if self.pose_condition == True:
|
||||
latent_conditions = condition_latents_prepare.prepare_latents_for_conditioning(
|
||||
vae=self.vae,
|
||||
image_or_video=videos,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=self.state.generator,
|
||||
)
|
||||
else:
|
||||
latent_conditions = self.model_config["prepare_latents"](
|
||||
vae=self.vae,
|
||||
image_or_video=videos,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=self.state.generator,
|
||||
)
|
||||
text_conditions = self.model_config["prepare_conditions"](
|
||||
tokenizer=self.tokenizer,
|
||||
text_encoder=self.text_encoder,
|
||||
@@ -714,27 +724,28 @@ class Trainer:
|
||||
latent_conditions = make_contiguous(latent_conditions)
|
||||
|
||||
if self.pose_conditioning:
|
||||
pose_video_latents = self.model_config["prepare_latents"](
|
||||
vae=self.vae,
|
||||
image_or_video=poses,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=generator,
|
||||
)
|
||||
# VAE output not patchified yet.
|
||||
pose_video_latents = condition_latents_prepare.prepare_latents_for_conditioning(
|
||||
vae=self.vae,
|
||||
image_or_video=poses,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=self.state.generator,
|
||||
)
|
||||
|
||||
pose_video_latents = make_contiguous(pose_video_latents)
|
||||
|
||||
img_refs_latents = self.model_config["prepare_latents"](
|
||||
vae=self.vae,
|
||||
image_or_video=img_refs,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=generator,
|
||||
)
|
||||
img_refs_latents = condition_latents_prepare.prepare_latents_for_conditioning(
|
||||
vae=self.vae,
|
||||
image_or_video=img_refs_latents,
|
||||
patch_size=self.transformer_config.patch_size,
|
||||
patch_size_t=self.transformer_config.patch_size_t,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
generator=self.state.generator,
|
||||
)
|
||||
|
||||
img_refs_latents = make_contiguous(img_refs_latents)
|
||||
|
||||
@@ -762,13 +773,24 @@ class Trainer:
|
||||
generator=self.state.generator,
|
||||
)
|
||||
timesteps = (sigmas * 1000.0).long()
|
||||
if self.pose_conditioning:
|
||||
noise = torch.randn(
|
||||
latent_conditions["latents"].shape,
|
||||
generator=self.state.generator,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
)
|
||||
# B x 2C latent
|
||||
imf_ref_noise = torch.cat([img_refs_latents,noise])
|
||||
# create noise for the img ref and concat it to latent to be shape B, 2c, etc.
|
||||
else:
|
||||
noise = torch.randn(
|
||||
latent_conditions["latents"].shape,
|
||||
generator=self.state.generator,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
)
|
||||
|
||||
noise = torch.randn(
|
||||
latent_conditions["latents"].shape,
|
||||
generator=self.state.generator,
|
||||
device=accelerator.device,
|
||||
dtype=weight_dtype,
|
||||
)
|
||||
sigmas = expand_tensor_dims(sigmas, ndim=noise.ndim)
|
||||
|
||||
# TODO(aryan): We probably don't need calculate_noisy_latents because we can determine the type of
|
||||
@@ -781,19 +803,20 @@ class Trainer:
|
||||
timesteps=timesteps,
|
||||
)
|
||||
else:
|
||||
if self.pose_conditioning:
|
||||
|
||||
noisy_latents = (1.0 - sigmas) * img_refs_latents["latents"] + sigmas * noise
|
||||
|
||||
noisy_latents = noisy_latents + pose_video_latents["latents"]
|
||||
img_refs_latents.update({"noisy_latents": noisy_latents})
|
||||
if self.pose_condition:
|
||||
# Default to flow-matching noise addition
|
||||
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise
|
||||
noisy_latents = noisy_latents.to(latent_conditions["latents"].dtype)
|
||||
|
||||
# this is used for whatever reason to pass into the
|
||||
latent_conditions.update({"noisy_latents": noisy_latents})
|
||||
else:
|
||||
# Default to flow-matching noise addition
|
||||
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise
|
||||
noisy_latents = noisy_latents.to(latent_conditions["latents"].dtype)
|
||||
|
||||
# this is used for whatever reason to pass into the
|
||||
latent_conditions.update({"noisy_latents": noisy_latents})
|
||||
|
||||
# this is used for whatever reason to pass into the
|
||||
latent_conditions.update({"noisy_latents": noisy_latents})
|
||||
|
||||
weights = prepare_loss_weights(
|
||||
scheduler=self.scheduler,
|
||||
@@ -821,8 +844,6 @@ class Trainer:
|
||||
**text_conditions,
|
||||
)
|
||||
|
||||
|
||||
|
||||
target = prepare_target(
|
||||
scheduler=self.scheduler, noise=noise, latents=latent_conditions["latents"]
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user