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

Full Finetuning for LTX possibily extended to other models. (#192)

Browse Source 
* Full Finetuning for LTX possibily extended to other models.

* Change name of the flag

* Used disable grad for component on lora fine tuning enabled

* Suggestions Addressed
Renamed to SFT
Added 2 other models.
Testing required.

* Switching to Full FineTuning

* Run linter.

* parse subfolder when needed.

* tackle saving and loading hooks.

* tackle validation.

* fix subfolder bug.

* remove __class__.

* refactor

* remove unnecessary changes

* handle saving of final model weights correctly

* remove unnecessary changes

* LTX uses a default frame rate of 24 FPS
We need to modify the output validation framerate to match that value.
Add Framerate args.
Add Update video output and inference frame rate

* There was a results_args mapping that needed to be modified.

* update

* update README

* Update README.md

* update docs

* add training configuration in cogvideox

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Aryan <aryan@huggingface.co>
Co-authored-by: Aryan <contact.aryanvs@gmail.com>
This commit is contained in:
ArEnSc
2025-01-13 13:25:39 -05:00
committed by GitHub
parent a3898075d1
commit e5df80cc36
17 changed files with 365 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/args.py
+28 -2
View File
@@ -207,6 +207,8 @@ class Args:
Perform validation every `n` training steps.
enable_model_cpu_offload (`bool`, defaults to `False`):
Whether or not to offload different modeling components to CPU during validation.
validation_frame_rate (`int`, defaults to `25`):
Frame rate to use for the validation videos. This value is defaulted to 25, as used in LTX Video pipeline.
MISCELLANEOUS ARGUMENTS
-----------------------
@@ -319,6 +321,7 @@ class Args:
validation_every_n_epochs: Optional[int] = None
validation_every_n_steps: Optional[int] = None
enable_model_cpu_offload: bool = False
validation_frame_rate: int = 25
# Miscellaneous arguments
tracker_name: str = "finetrainers"
@@ -417,6 +420,7 @@ class Args:
"validation_every_n_epochs": self.validation_every_n_epochs,
"validation_every_n_steps": self.validation_every_n_steps,
"enable_model_cpu_offload": self.enable_model_cpu_offload,
"validation_frame_rate": self.validation_frame_rate,
},
"miscellaneous_arguments": {
"tracker_name": self.tracker_name,
@@ -460,6 +464,7 @@ def parse_arguments() -> Args:
def validate_args(args: Args):
_validate_training_args(args)
_validate_validation_args(args)
@@ -678,8 +683,9 @@ def _add_training_arguments(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--training_type",
type=str,
choices=["lora", "full-finetune"],
required=True,
help="Type of training to perform. Choose between ['lora']",
help="Type of training to perform. Choose between ['lora', 'full-finetune']",
)
parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
parser.add_argument(
@@ -713,7 +719,11 @@ def _add_training_arguments(parser: argparse.ArgumentParser) -> None:
help="The lora_alpha to compute scaling factor (lora_alpha / rank) for LoRA matrices.",
)
parser.add_argument(
"--target_modules", type=str, default="to_k to_q to_v to_out.0", nargs="+", help="The target modules for LoRA."
"--target_modules",
type=str,
default=["to_k", "to_q", "to_v", "to_out.0"],
nargs="+",
help="The target modules for LoRA.",
)
parser.add_argument(
"--gradient_accumulation_steps",
@@ -890,6 +900,12 @@ def _add_validation_arguments(parser: argparse.ArgumentParser) -> None:
default=None,
help="Run validation every X training steps. Validation consists of running the validation prompt `args.num_validation_videos` times.",
)
parser.add_argument(
"--validation_frame_rate",
type=int,
default=25,
help="Frame rate to use for the validation videos.",
)
parser.add_argument(
"--enable_model_cpu_offload",
action="store_true",
@@ -1085,6 +1101,7 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args:
result_args.validation_every_n_epochs = args.validation_epochs
result_args.validation_every_n_steps = args.validation_steps
result_args.enable_model_cpu_offload = args.enable_model_cpu_offload
result_args.validation_frame_rate = args.validation_frame_rate
# Miscellaneous arguments
result_args.tracker_name = args.tracker_name
@@ -1100,6 +1117,15 @@ def _map_to_args_type(args: Dict[str, Any]) -> Args:
return result_args
def _validate_training_args(args: Args):
if args.training_type == "lora":
assert args.rank is not None, "Rank is required for LoRA training"
assert args.lora_alpha is not None, "LoRA alpha is required for LoRA training"
assert (
args.target_modules is not None and len(args.target_modules) > 0
), "Target modules are required for LoRA training"
def _validate_validation_args(args: Args):
assert args.validation_prompts is not None, "Validation prompts are required for validation"
if args.validation_images is not None:
finetrainers/cogvideox/__init__.py
+1
View File
@@ -1 +1,2 @@
from .cogvideox_lora import COGVIDEOX_T2V_LORA_CONFIG
from .full_finetune import COGVIDEOX_T2V_FULL_FINETUNE_CONFIG
finetrainers/cogvideox/cogvideox_lora.py
+1
View File
@@ -311,6 +311,7 @@ def _pad_frames(latents: torch.Tensor, patch_size_t: int):
return latents
# TODO(aryan): refactor into model specs for better re-use
COGVIDEOX_T2V_LORA_CONFIG = {
"pipeline_cls": CogVideoXPipeline,
"load_condition_models": load_condition_models,
finetrainers/cogvideox/full_finetune.py
+32
View File
@@ -0,0 +1,32 @@
from diffusers import CogVideoXPipeline
from .cogvideox_lora import (
calculate_noisy_latents,
collate_fn_t2v,
forward_pass,
initialize_pipeline,
load_condition_models,
load_diffusion_models,
load_latent_models,
post_latent_preparation,
prepare_conditions,
prepare_latents,
validation,
)
# TODO(aryan): refactor into model specs for better re-use
COGVIDEOX_T2V_FULL_FINETUNE_CONFIG = {
"pipeline_cls": CogVideoXPipeline,
"load_condition_models": load_condition_models,
"load_latent_models": load_latent_models,
"load_diffusion_models": load_diffusion_models,
"initialize_pipeline": initialize_pipeline,
"prepare_conditions": prepare_conditions,
"prepare_latents": prepare_latents,
"post_latent_preparation": post_latent_preparation,
"collate_fn": collate_fn_t2v,
"calculate_noisy_latents": calculate_noisy_latents,
"forward_pass": forward_pass,
"validation": validation,
}
finetrainers/hunyuan_video/__init__.py
+1
View File
@@ -1 +1,2 @@
from .full_finetune import HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG
from .hunyuan_video_lora import HUNYUAN_VIDEO_T2V_LORA_CONFIG
finetrainers/hunyuan_video/full_finetune.py
+30
View File
@@ -0,0 +1,30 @@
from diffusers import HunyuanVideoPipeline
from .hunyuan_video_lora import (
collate_fn_t2v,
forward_pass,
initialize_pipeline,
load_condition_models,
load_diffusion_models,
load_latent_models,
post_latent_preparation,
prepare_conditions,
prepare_latents,
validation,
)
# TODO(aryan): refactor into model specs for better re-use
HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG = {
"pipeline_cls": HunyuanVideoPipeline,
"load_condition_models": load_condition_models,
"load_latent_models": load_latent_models,
"load_diffusion_models": load_diffusion_models,
"initialize_pipeline": initialize_pipeline,
"prepare_conditions": prepare_conditions,
"prepare_latents": prepare_latents,
"post_latent_preparation": post_latent_preparation,
"collate_fn": collate_fn_t2v,
"forward_pass": forward_pass,
"validation": validation,
}
finetrainers/hunyuan_video/hunyuan_video_lora.py
+1
View File
@@ -345,6 +345,7 @@ def _get_clip_prompt_embeds(
return {"pooled_prompt_embeds": prompt_embeds}
# TODO(aryan): refactor into model specs for better re-use
HUNYUAN_VIDEO_T2V_LORA_CONFIG = {
"pipeline_cls": HunyuanVideoPipeline,
"load_condition_models": load_condition_models,
finetrainers/ltx_video/__init__.py
+1
View File
@@ -1 +1,2 @@
from .full_finetune import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG
from .ltx_video_lora import LTX_VIDEO_T2V_LORA_CONFIG
finetrainers/ltx_video/full_finetune.py
+30
View File
@@ -0,0 +1,30 @@
from diffusers import LTXPipeline
from .ltx_video_lora import (
collate_fn_t2v,
forward_pass,
initialize_pipeline,
load_condition_models,
load_diffusion_models,
load_latent_models,
post_latent_preparation,
prepare_conditions,
prepare_latents,
validation,
)
# TODO(aryan): refactor into model specs for better re-use
LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG = {
"pipeline_cls": LTXPipeline,
"load_condition_models": load_condition_models,
"load_latent_models": load_latent_models,
"load_diffusion_models": load_diffusion_models,
"initialize_pipeline": initialize_pipeline,
"prepare_conditions": prepare_conditions,
"prepare_latents": prepare_latents,
"post_latent_preparation": post_latent_preparation,
"collate_fn": collate_fn_t2v,
"forward_pass": forward_pass,
"validation": validation,
}
finetrainers/ltx_video/ltx_video_lora.py
+1 -1
View File
@@ -225,7 +225,7 @@ def validation(
height: Optional[int] = None,
width: Optional[int] = None,
num_frames: Optional[int] = None,
frame_rate: int = 25,
frame_rate: int = 24,
num_videos_per_prompt: int = 1,
generator: Optional[torch.Generator] = None,
**kwargs,
finetrainers/models.py
+6 -3
View File
@@ -1,19 +1,22 @@
from typing import Any, Dict
from .cogvideox import COGVIDEOX_T2V_LORA_CONFIG
from .hunyuan_video import HUNYUAN_VIDEO_T2V_LORA_CONFIG
from .ltx_video import LTX_VIDEO_T2V_LORA_CONFIG
from .cogvideox import COGVIDEOX_T2V_FULL_FINETUNE_CONFIG, COGVIDEOX_T2V_LORA_CONFIG
from .hunyuan_video import HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG, HUNYUAN_VIDEO_T2V_LORA_CONFIG
from .ltx_video import LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG, LTX_VIDEO_T2V_LORA_CONFIG
SUPPORTED_MODEL_CONFIGS = {
"hunyuan_video": {
"lora": HUNYUAN_VIDEO_T2V_LORA_CONFIG,
"full-finetune": HUNYUAN_VIDEO_T2V_FULL_FINETUNE_CONFIG,
},
"ltx_video": {
"lora": LTX_VIDEO_T2V_LORA_CONFIG,
"full-finetune": LTX_VIDEO_T2V_FULL_FINETUNE_CONFIG,
},
"cogvideox": {
"lora": COGVIDEOX_T2V_LORA_CONFIG,
"full-finetune": COGVIDEOX_T2V_FULL_FINETUNE_CONFIG,
},
}
finetrainers/trainer.py
+145 -107
View File
@@ -5,7 +5,7 @@ import os
import random
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict
from typing import Any, Dict, List
import diffusers
import torch
@@ -21,6 +21,7 @@ from accelerate.utils import (
gather_object,
set_seed,
)
from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from diffusers.optimization import get_scheduler
@@ -242,16 +243,7 @@ class Trainer:
condition_components = self.model_config["load_condition_models"](**self._get_load_components_kwargs())
self._set_components(condition_components)
self._move_components_to_device()
# TODO(aryan): refactor later. for now only lora is supported
components_to_disable_grads = [
self.text_encoder,
self.text_encoder_2,
self.text_encoder_3,
]
for component in components_to_disable_grads:
if component is not None:
component.requires_grad_(False)
self._disable_grad_for_components([self.text_encoder, self.text_encoder_2, self.text_encoder_3])
if self.args.caption_dropout_p > 0 and self.args.caption_dropout_technique == "empty":
logger.warning(
@@ -305,12 +297,7 @@ class Trainer:
latent_components = self.model_config["load_latent_models"](**self._get_load_components_kwargs())
self._set_components(latent_components)
self._move_components_to_device()
# TODO(aryan): refactor later
components_to_disable_grads = [self.vae]
for component in components_to_disable_grads:
if component is not None:
component.requires_grad_(False)
self._disable_grad_for_components([self.vae])
if self.vae is not None:
if self.args.enable_slicing:
@@ -371,24 +358,22 @@ class Trainer:
diffusion_components = self.model_config["load_diffusion_models"](**self._get_load_components_kwargs())
self._set_components(diffusion_components)
# TODO(aryan): refactor later. for now only lora is supported
components_to_disable_grads = [
self.text_encoder,
self.text_encoder_2,
self.text_encoder_3,
self.transformer,
self.vae,
]
for component in components_to_disable_grads:
if component is not None:
component.requires_grad_(False)
components = [self.text_encoder, self.text_encoder_2, self.text_encoder_3, self.vae]
self._disable_grad_for_components(components)
if self.args.training_type == "full-finetune":
logger.info("Finetuning transformer with no additional parameters")
self._enable_grad_for_components([self.transformer])
else:
logger.info("Finetuning transformer with PEFT parameters")
self._disable_grad_for_components([self.transformer])
# For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
# as these weights are only used for inference, keeping weights in full precision is not required.
weight_dtype = self._get_training_dtype(accelerator=self.state.accelerator)
if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
# due to pytorch#99272, MPS does not yet support bfloat16.
# Due to pytorch#99272, MPS does not yet support bfloat16.
raise ValueError(
"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
)
@@ -406,13 +391,16 @@ class Trainer:
if self.args.gradient_checkpointing:
self.transformer.enable_gradient_checkpointing()
transformer_lora_config = LoraConfig(
r=self.args.rank,
lora_alpha=self.args.lora_alpha,
init_lora_weights=True,
target_modules=self.args.target_modules,
)
self.transformer.add_adapter(transformer_lora_config)
if self.args.training_type == "lora":
transformer_lora_config = LoraConfig(
r=self.args.rank,
lora_alpha=self.args.lora_alpha,
init_lora_weights=True,
target_modules=self.args.target_modules,
)
self.transformer.add_adapter(transformer_lora_config)
else:
transformer_lora_config = None
# Enable TF32 for faster training on Ampere GPUs: https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
if self.args.allow_tf32 and torch.cuda.is_available():
@@ -432,7 +420,8 @@ class Trainer:
type(unwrap_model(self.state.accelerator, self.transformer)),
):
model = unwrap_model(self.state.accelerator, model)
transformer_lora_layers_to_save = get_peft_model_state_dict(model)
if self.args.training_type == "lora":
transformer_lora_layers_to_save = get_peft_model_state_dict(model)
else:
raise ValueError(f"Unexpected save model: {model.__class__}")
@@ -440,10 +429,18 @@ class Trainer:
if weights:
weights.pop()
self.model_config["pipeline_cls"].save_lora_weights(
output_dir,
transformer_lora_layers=transformer_lora_layers_to_save,
)
if self.args.training_type == "lora":
self.model_config["pipeline_cls"].save_lora_weights(
output_dir,
transformer_lora_layers=transformer_lora_layers_to_save,
)
else:
model.save_pretrained(os.path.join(output_dir, "transformer"))
# In some cases, the scheduler needs to be loaded with specific config (e.g. in CogVideoX). Since we need
# to able to load all diffusion components from a specific checkpoint folder during validation, we need to
# ensure the scheduler config is serialized as well.
self.scheduler.save_pretrained(os.path.join(output_dir, "scheduler"))
def load_model_hook(models, input_dir):
if not self.state.accelerator.distributed_type == DistributedType.DEEPSPEED:
@@ -459,33 +456,39 @@ class Trainer:
f"Unexpected save model: {unwrap_model(self.state.accelerator, model).__class__}"
)
else:
transformer_ = unwrap_model(self.state.accelerator, self.transformer).__class__.from_pretrained(
self.args.pretrained_model_name_or_path, subfolder="transformer"
)
transformer_.add_adapter(transformer_lora_config)
transformer_cls_ = unwrap_model(self.state.accelerator, self.transformer).__class__
lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir)
transformer_state_dict = {
f'{k.replace("transformer.", "")}': v
for k, v in lora_state_dict.items()
if k.startswith("transformer.")
}
incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")
if incompatible_keys is not None:
# check only for unexpected keys
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
if unexpected_keys:
logger.warning(
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
f" {unexpected_keys}. "
if self.args.training_type == "lora":
transformer_ = transformer_cls_.from_pretrained(
self.args.pretrained_model_name_or_path, subfolder="transformer"
)
transformer_.add_adapter(transformer_lora_config)
lora_state_dict = self.model_config["pipeline_cls"].lora_state_dict(input_dir)
transformer_state_dict = {
f'{k.replace("transformer.", "")}': v
for k, v in lora_state_dict.items()
if k.startswith("transformer.")
}
incompatible_keys = set_peft_model_state_dict(
transformer_, transformer_state_dict, adapter_name="default"
)
if incompatible_keys is not None:
# check only for unexpected keys
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
if unexpected_keys:
logger.warning(
f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
f" {unexpected_keys}. "
)
# Make sure the trainable params are in float32. This is again needed since the base models
# are in `weight_dtype`. More details:
# https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
if self.args.mixed_precision == "fp16":
# only upcast trainable parameters (LoRA) into fp32
cast_training_params([transformer_])
# Make sure the trainable params are in float32. This is again needed since the base models
# are in `weight_dtype`. More details:
# https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
if self.args.mixed_precision == "fp16" and self.args.training_type == "lora":
# only upcast trainable parameters (LoRA) into fp32
cast_training_params([transformer_], dtype=torch.float32)
else:
transformer_ = transformer_cls_.from_pretrained(os.path.join(input_dir, "transformer"))
self.state.accelerator.register_save_state_pre_hook(save_model_hook)
self.state.accelerator.register_load_state_pre_hook(load_model_hook)
@@ -497,7 +500,7 @@ class Trainer:
self.state.train_steps = self.args.train_steps
# Make sure the trainable params are in float32
if self.args.mixed_precision == "fp16":
if self.args.mixed_precision == "fp16" and self.args.training_type == "lora":
# only upcast trainable parameters (LoRA) into fp32
cast_training_params([self.transformer], dtype=torch.float32)
@@ -510,13 +513,13 @@ class Trainer:
* self.state.accelerator.num_processes
)
transformer_lora_parameters = list(filter(lambda p: p.requires_grad, self.transformer.parameters()))
transformer_trainable_parameters = list(filter(lambda p: p.requires_grad, self.transformer.parameters()))
transformer_parameters_with_lr = {
"params": transformer_lora_parameters,
"params": transformer_trainable_parameters,
"lr": self.state.learning_rate,
}
params_to_optimize = [transformer_parameters_with_lr]
self.state.num_trainable_parameters = sum(p.numel() for p in transformer_lora_parameters)
self.state.num_trainable_parameters = sum(p.numel() for p in transformer_trainable_parameters)
use_deepspeed_opt = (
self.state.accelerator.state.deepspeed_plugin is not None
@@ -608,6 +611,12 @@ class Trainer:
)
self.vae_config = FrozenDict(**vae_config)
# In some cases, the scheduler needs to be loaded with specific config (e.g. in CogVideoX). Since we need
# to able to load all diffusion components from a specific checkpoint folder during validation, we need to
# ensure the scheduler config is serialized as well.
if self.args.training_type == "full-finetune":
self.scheduler.save_pretrained(os.path.join(self.args.output_dir, "scheduler"))
self.state.train_batch_size = (
self.args.batch_size * self.state.accelerator.num_processes * self.args.gradient_accumulation_steps
)
@@ -872,14 +881,17 @@ class Trainer:
accelerator.wait_for_everyone()
if accelerator.is_main_process:
# TODO: consider factoring this out when supporting other types of training algos.
self.transformer = unwrap_model(accelerator, self.transformer)
transformer_lora_layers = get_peft_model_state_dict(self.transformer)
transformer = unwrap_model(accelerator, self.transformer)
self.model_config["pipeline_cls"].save_lora_weights(
save_directory=self.args.output_dir,
transformer_lora_layers=transformer_lora_layers,
)
if self.args.training_type == "lora":
transformer_lora_layers = get_peft_model_state_dict(transformer)
self.model_config["pipeline_cls"].save_lora_weights(
save_directory=self.args.output_dir,
transformer_lora_layers=transformer_lora_layers,
)
else:
transformer.save_pretrained(os.path.join(self.args.output_dir, "transformer"))
self.validate(step=global_step, final_validation=True)
@@ -910,35 +922,7 @@ class Trainer:
memory_statistics = get_memory_statistics()
logger.info(f"Memory before validation start: {json.dumps(memory_statistics, indent=4)}")
if not final_validation:
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
tokenizer=self.tokenizer,
text_encoder=self.text_encoder,
tokenizer_2=self.tokenizer_2,
text_encoder_2=self.text_encoder_2,
transformer=unwrap_model(accelerator, self.transformer),
vae=self.vae,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
else:
# `torch_dtype` is manually set within `initialize_pipeline()`.
self._delete_components()
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
pipeline.load_lora_weights(self.args.output_dir)
pipeline = self._get_and_prepare_pipeline_for_validation(final_validation=final_validation)
all_processes_artifacts = []
prompts_to_filenames = {}
@@ -953,7 +937,7 @@ class Trainer:
height = self.args.validation_heights[i]
width = self.args.validation_widths[i]
num_frames = self.args.validation_num_frames[i]
frame_rate = self.args.validation_frame_rate
if image is not None:
image = load_image(image)
if video is not None:
@@ -971,6 +955,7 @@ class Trainer:
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_videos_per_prompt=self.args.num_validation_videos_per_prompt,
generator=torch.Generator(device=accelerator.device).manual_seed(
self.args.seed if self.args.seed is not None else 0
@@ -1010,7 +995,7 @@ class Trainer:
elif artifact_type == "video":
logger.debug(f"Saving video to {filename}")
# TODO: this should be configurable here as well as in validation runs where we call the pipeline that has `fps`.
export_to_video(artifact_value, filename, fps=15)
export_to_video(artifact_value, filename, fps=frame_rate)
artifact_value = wandb.Video(filename, caption=prompt)
all_processes_artifacts.append(artifact_value)
@@ -1144,3 +1129,56 @@ class Trainer:
elif self.state.accelerator.mixed_precision == "bf16":
weight_dtype = torch.bfloat16
return weight_dtype
def _get_and_prepare_pipeline_for_validation(self, final_validation: bool = False) -> DiffusionPipeline:
accelerator = self.state.accelerator
if not final_validation:
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
tokenizer=self.tokenizer,
text_encoder=self.text_encoder,
tokenizer_2=self.tokenizer_2,
text_encoder_2=self.text_encoder_2,
transformer=unwrap_model(accelerator, self.transformer),
vae=self.vae,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
else:
self._delete_components()
# Load the transformer weights from the final checkpoint if performing full-finetune
transformer = None
if self.args.training_type == "full-finetune":
transformer = self.model_config["load_diffusion_models"](model_id=self.args.output_dir)["transformer"]
pipeline = self.model_config["initialize_pipeline"](
model_id=self.args.pretrained_model_name_or_path,
transformer=transformer,
device=accelerator.device,
revision=self.args.revision,
cache_dir=self.args.cache_dir,
enable_slicing=self.args.enable_slicing,
enable_tiling=self.args.enable_tiling,
enable_model_cpu_offload=self.args.enable_model_cpu_offload,
)
# Load the LoRA weights if performing LoRA finetuning
if self.args.training_type == "lora":
pipeline.load_lora_weights(self.args.output_dir)
return pipeline
def _disable_grad_for_components(self, components: List[torch.nn.Module]):
for component in components:
if component is not None:
component.requires_grad_(False)
def _enable_grad_for_components(self, components: List[torch.nn.Module]):
for component in components:
if component is not None:
component.requires_grad_(True)