[diffusion] feat: support LoRA (#13859)

This commit is contained in:
Mick
2025-11-26 00:21:33 +08:00
committed by GitHub
parent 46673b4224
commit dfd7ab9682
29 changed files with 463 additions and 278 deletions

View File

@@ -55,7 +55,7 @@ class skip_init_modules:
cls.reset_parameters = lambda self: None # skip init
def __exit__(self, exc_type, exc_value, traceback):
# Restore originals
# restore originals
for cls, orig in self._orig_reset.items():
cls.reset_parameters = orig
@@ -199,7 +199,7 @@ class TextEncoderLoader(ComponentLoader):
def _get_weights_iterator(
self, source: "Source", to_cpu: bool
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""Get an iterator for the model weights based on the load format."""
"""get an iterator for the model weights based on the load format."""
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
source.model_or_path,
source.fall_back_to_pt,
@@ -214,7 +214,7 @@ class TextEncoderLoader(ComponentLoader):
if self.counter_before_loading_weights == 0.0:
self.counter_before_loading_weights = time.perf_counter()
# Apply the prefix.
# apply the prefix.
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
def _get_all_weights(