File size: 5,200 Bytes
84669a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | from pathlib import Path, PurePath
from typing import Dict, List, Optional, Union
from diffusers.loaders.lora_pipeline import _fetch_state_dict
from diffusers.loaders.lora_conversion_utils import _convert_hunyuan_video_lora_to_diffusers
from diffusers.utils.peft_utils import set_weights_and_activate_adapters
from diffusers.loaders.peft import _SET_ADAPTER_SCALE_FN_MAPPING
import torch
def load_lora(transformer, lora_path: Path, weight_name: Optional[str] = "pytorch_lora_weights.safetensors"):
"""
Load LoRA weights into the transformer model.
Args:
transformer: The transformer model to which LoRA weights will be applied.
lora_path (Path): Path to the LoRA weights file.
weight_name (Optional[str]): Name of the weight to load.
"""
state_dict = _fetch_state_dict(
lora_path,
weight_name,
True,
True,
None,
None,
None,
None,
None,
None,
None,
None)
state_dict = _convert_hunyuan_video_lora_to_diffusers(state_dict)
# should weight_name even be Optional[str] or just str?
# For now, we assume it is never None
# The module name in the state_dict must not include a . in the name
# See https://github.com/pytorch/pytorch/pull/6639/files#diff-4be56271f7bfe650e3521c81fd363da58f109cd23ee80d243156d2d6ccda6263R133-R134
adapter_name = str(PurePath(weight_name).with_suffix('')).replace('.', '_DOT_')
if '_DOT_' in adapter_name:
print(
f"LoRA file '{weight_name}' contains a '.' in the name. " +
'This may cause issues. Consider renaming the file.' +
f" Using '{adapter_name}' as the adapter name to be safe."
)
# Check if adapter already exists and delete it if it does
if hasattr(transformer, 'peft_config') and adapter_name in transformer.peft_config:
print(f"Adapter '{adapter_name}' already exists. Removing it before loading again.")
# Use delete_adapters (plural) instead of delete_adapter
transformer.delete_adapters([adapter_name])
# Load the adapter with the original name
transformer.load_lora_adapter(state_dict, network_alphas=None, adapter_name=adapter_name)
print(f"LoRA weights '{adapter_name}' loaded successfully.")
return transformer
def unload_all_loras(transformer):
"""
Completely unload all LoRA adapters from the transformer model.
"""
if hasattr(transformer, 'peft_config') and transformer.peft_config:
# Get all adapter names
adapter_names = list(transformer.peft_config.keys())
if adapter_names:
print(f"Removing all LoRA adapters: {', '.join(adapter_names)}")
# Delete all adapters
transformer.delete_adapters(adapter_names)
# Force cleanup of any remaining adapter references
if hasattr(transformer, 'active_adapter'):
transformer.active_adapter = None
# Clear any cached states
for module in transformer.modules():
if hasattr(module, 'lora_A'):
if isinstance(module.lora_A, dict):
module.lora_A.clear()
if hasattr(module, 'lora_B'):
if isinstance(module.lora_B, dict):
module.lora_B.clear()
if hasattr(module, 'scaling'):
if isinstance(module.scaling, dict):
module.scaling.clear()
print("All LoRA adapters have been completely removed.")
else:
print("No LoRA adapters found to remove.")
else:
print("Model doesn't have any LoRA adapters or peft_config.")
# Force garbage collection
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
return transformer
# TODO(neph1): remove when HunyuanVideoTransformer3DModelPacked is in _SET_ADAPTER_SCALE_FN_MAPPING
def set_adapters(
transformer,
adapter_names: Union[List[str], str],
weights: Optional[Union[float, Dict, List[float], List[Dict], List[None]]] = None,
):
adapter_names = [adapter_names] if isinstance(adapter_names, str) else adapter_names
# Expand weights into a list, one entry per adapter
# examples for e.g. 2 adapters: [{...}, 7] -> [7,7] ; None -> [None, None]
if not isinstance(weights, list):
weights = [weights] * len(adapter_names)
if len(adapter_names) != len(weights):
raise ValueError(
f"Length of adapter names {len(adapter_names)} is not equal to the length of their weights {len(weights)}."
)
# Set None values to default of 1.0
# e.g. [{...}, 7] -> [{...}, 7] ; [None, None] -> [1.0, 1.0]
weights = [w if w is not None else 1.0 for w in weights]
# e.g. [{...}, 7] -> [{expanded dict...}, 7]
scale_expansion_fn = _SET_ADAPTER_SCALE_FN_MAPPING["HunyuanVideoTransformer3DModel"]
weights = scale_expansion_fn(transformer, weights)
set_weights_and_activate_adapters(transformer, adapter_names, weights)
|