|
from typing import TYPE_CHECKING |
|
|
|
import torch |
|
|
|
from .dependencies.diffusers.peft import load_lora_weights |
|
|
|
|
|
if TYPE_CHECKING: |
|
from finetrainers.args import BaseArgsType |
|
from finetrainers.parallel import ParallelBackendType |
|
|
|
|
|
def perform_patches_for_training(args: "BaseArgsType", parallel_backend: "ParallelBackendType") -> None: |
|
|
|
from finetrainers.config import ModelType, TrainingType |
|
|
|
from .dependencies.diffusers import patch |
|
|
|
|
|
patch_scaled_dot_product_attention() |
|
|
|
patch.patch_diffusers_rms_norm_forward() |
|
|
|
|
|
if args.model_name == ModelType.LTX_VIDEO: |
|
from .models.ltx_video import patch |
|
|
|
patch.patch_transformer_forward() |
|
if parallel_backend.tensor_parallel_enabled: |
|
patch.patch_apply_rotary_emb_for_tp_compatibility() |
|
|
|
|
|
if args.model_name == ModelType.WAN and "transformer" in args.layerwise_upcasting_modules: |
|
from .models.wan import patch |
|
|
|
patch.patch_time_text_image_embedding_forward() |
|
|
|
|
|
if args.training_type == TrainingType.LORA and len(args.layerwise_upcasting_modules) > 0: |
|
from .dependencies.peft import patch |
|
|
|
patch.patch_peft_move_adapter_to_device_of_base_layer() |
|
|
|
|
|
def perform_patches_for_inference(args: "BaseArgsType", parallel_backend: "ParallelBackendType") -> None: |
|
|
|
from .dependencies.diffusers import patch |
|
|
|
|
|
patch_scaled_dot_product_attention() |
|
|
|
patch.patch_diffusers_rms_norm_forward() |
|
|
|
|
|
def patch_scaled_dot_product_attention(): |
|
from finetrainers.models.attention_dispatch import attention_dispatch |
|
|
|
torch.nn.functional.scaled_dot_product_attention = attention_dispatch |
|
|