diff --git a/finetrainers/dataset.py b/finetrainers/dataset.py index 53b9378..b1be6ef 100644 --- a/finetrainers/dataset.py +++ b/finetrainers/dataset.py @@ -20,7 +20,7 @@ import decord # isort:skip decord.bridge.set_bridge("torch") -from .constants import PRECOMPUTED_DIR_NAME, PRECOMPUTED_CONDITIONS_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME +from .constants import PRECOMPUTED_CONDITIONS_DIR_NAME, PRECOMPUTED_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME # noqa logger = get_logger(__name__) diff --git a/finetrainers/hunyuan_video/hunyuan_video_lora.py b/finetrainers/hunyuan_video/hunyuan_video_lora.py index ec623d5..719425f 100644 --- a/finetrainers/hunyuan_video/hunyuan_video_lora.py +++ b/finetrainers/hunyuan_video/hunyuan_video_lora.py @@ -9,8 +9,8 @@ from diffusers import ( HunyuanVideoPipeline, HunyuanVideoTransformer3DModel, ) -from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizer from PIL import Image +from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizer logger = get_logger("finetrainers") # pylint: disable=invalid-name diff --git a/finetrainers/ltx_video/ltx_video_lora.py b/finetrainers/ltx_video/ltx_video_lora.py index 59a121f..e104a8c 100644 --- a/finetrainers/ltx_video/ltx_video_lora.py +++ b/finetrainers/ltx_video/ltx_video_lora.py @@ -4,9 +4,8 @@ import torch import torch.nn as nn from accelerate.logging import get_logger from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel -from diffusers.utils import logging -from transformers import T5EncoderModel, T5Tokenizer from PIL import Image +from transformers import T5EncoderModel, T5Tokenizer logger = get_logger("finetrainers") # pylint: disable=invalid-name diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 5523999..22ae323 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -4,8 +4,8 @@ import math import os import random from datetime import timedelta -from typing import Any, Dict from pathlib import Path +from typing import Any, Dict import diffusers import torch @@ -18,37 +18,37 @@ from accelerate.utils import ( DistributedDataParallelKwargs, InitProcessGroupKwargs, ProjectConfiguration, - set_seed, gather_object, + set_seed, ) +from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.optimization import get_scheduler from diffusers.training_utils import ( cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3, ) -from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.utils import export_to_video, load_image, load_video from huggingface_hub import create_repo, upload_folder from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict from tqdm import tqdm -from .args import Args, validate_args, _INVERSE_DTYPE_MAP +from .args import _INVERSE_DTYPE_MAP, Args, validate_args from .constants import ( FINETRAINERS_LOG_LEVEL, - PRECOMPUTED_DIR_NAME, PRECOMPUTED_CONDITIONS_DIR_NAME, + PRECOMPUTED_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME, ) from .dataset import BucketSampler, PrecomputedDataset, VideoDatasetWithResizing from .models import get_config_from_model_name from .state import State +from .utils.checkpointing import get_intermediate_ckpt_path, get_latest_ckpt_path_to_resume_from from .utils.data_utils import should_perform_precomputation from .utils.file_utils import string_to_filename +from .utils.memory_utils import free_memory, get_memory_statistics, make_contiguous from .utils.optimizer_utils import get_optimizer -from .utils.memory_utils import get_memory_statistics, free_memory, make_contiguous -from .utils.torch_utils import unwrap_model, align_device_and_dtype, expand_tensor_to_dims -from .utils.checkpointing import get_latest_ckpt_path_to_resume_from, get_intermediate_ckpt_path +from .utils.torch_utils import align_device_and_dtype, expand_tensor_to_dims, unwrap_model logger = get_logger("finetrainers") @@ -934,7 +934,6 @@ class Trainer: step=step, ) - # Remove all hooks that might have been added during pipeline initialization to the models pipeline.remove_all_hooks() del pipeline diff --git a/finetrainers/utils/__init__.py b/finetrainers/utils/__init__.py index bdf27f5..85ffd2f 100644 --- a/finetrainers/utils/__init__.py +++ b/finetrainers/utils/__init__.py @@ -1,5 +1,5 @@ -from .file_utils import find_files, delete_files -from .diffusion_utils import resolution_dependant_timestep_flow_shift, default_flow_shift -from .memory_utils import get_memory_statistics, bytes_to_gigabytes, free_memory, make_contiguous -from .torch_utils import unwrap_model +from .diffusion_utils import default_flow_shift, resolution_dependant_timestep_flow_shift +from .file_utils import delete_files, find_files +from .memory_utils import bytes_to_gigabytes, free_memory, get_memory_statistics, make_contiguous from .optimizer_utils import get_optimizer, gradient_norm, max_gradient +from .torch_utils import unwrap_model diff --git a/finetrainers/utils/checkpointing.py b/finetrainers/utils/checkpointing.py index ba7f3d6..be0124b 100644 --- a/finetrainers/utils/checkpointing.py +++ b/finetrainers/utils/checkpointing.py @@ -1,8 +1,11 @@ import os from typing import Tuple + from accelerate.logging import get_logger + from ..constants import FINETRAINERS_LOG_LEVEL -from ..utils.file_utils import find_files, delete_files +from ..utils.file_utils import delete_files, find_files + logger = get_logger("finetrainers") logger.setLevel(FINETRAINERS_LOG_LEVEL) diff --git a/finetrainers/utils/data_utils.py b/finetrainers/utils/data_utils.py index b05c49e..c40f4ef 100644 --- a/finetrainers/utils/data_utils.py +++ b/finetrainers/utils/data_utils.py @@ -3,7 +3,7 @@ from typing import Union from accelerate.logging import get_logger -from ..constants import PRECOMPUTED_DIR_NAME, PRECOMPUTED_CONDITIONS_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME +from ..constants import PRECOMPUTED_CONDITIONS_DIR_NAME, PRECOMPUTED_DIR_NAME, PRECOMPUTED_LATENTS_DIR_NAME logger = get_logger("finetrainers") diff --git a/finetrainers/utils/file_utils.py b/finetrainers/utils/file_utils.py index 682563d..eb73177 100644 --- a/finetrainers/utils/file_utils.py +++ b/finetrainers/utils/file_utils.py @@ -2,7 +2,7 @@ import logging import os import shutil from pathlib import Path -from typing import Any, Dict, List, Union +from typing import List, Union logger = logging.getLogger("finetrainers") diff --git a/finetrainers/utils/memory_utils.py b/finetrainers/utils/memory_utils.py index 3742d46..add8e0b 100644 --- a/finetrainers/utils/memory_utils.py +++ b/finetrainers/utils/memory_utils.py @@ -1,5 +1,4 @@ import gc -import logging from typing import Any, Dict, Union import torch diff --git a/finetrainers/utils/optimizer_utils.py b/finetrainers/utils/optimizer_utils.py index e988687..84c215b 100644 --- a/finetrainers/utils/optimizer_utils.py +++ b/finetrainers/utils/optimizer_utils.py @@ -1,7 +1,8 @@ import inspect -from accelerate.logging import get_logger import torch +from accelerate.logging import get_logger + logger = get_logger("finetrainers") @@ -45,7 +46,7 @@ def get_optimizer( if (use_torchao and (use_8bit or use_4bit)) or use_cpu_offload_optimizer: try: - import torchao + import torchao # noqa except ImportError: raise ImportError( diff --git a/finetrainers/utils/torch_utils.py b/finetrainers/utils/torch_utils.py index 17c5f9c..13989ad 100644 --- a/finetrainers/utils/torch_utils.py +++ b/finetrainers/utils/torch_utils.py @@ -32,4 +32,4 @@ def align_device_and_dtype( def expand_tensor_to_dims(tensor, ndim): while len(tensor.shape) < ndim: tensor = tensor.unsqueeze(-1) - return tensor \ No newline at end of file + return tensor diff --git a/tests/test_lora_inference.py b/tests/test_lora_inference.py index 31428a4..d1c0439 100644 --- a/tests/test_lora_inference.py +++ b/tests/test_lora_inference.py @@ -8,6 +8,7 @@ python3 test_lora_inference.py --prompt "A girl is ridding a bike." --model_path """ import argparse + import torch from diffusers import CogVideoXPipeline from diffusers.utils import export_to_video