github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add a method to have an inference only lora

Browse Source
This commit is contained in:
Jaret Burkett
2024-10-04 10:06:53 -06:00
parent 28e6f00790
commit a800c9d19e
2 changed files with 49 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
toolkit/config_modules.py
View File

@@ -410,6 +410,7 @@ class ModelConfig:
self.lora_path = kwargs.get('lora_path', None) self.lora_path = kwargs.get('lora_path', None)
# mainly for decompression loras for distilled models # mainly for decompression loras for distilled models
self.assistant_lora_path = kwargs.get('assistant_lora_path', None) self.assistant_lora_path = kwargs.get('assistant_lora_path', None)
self.inference_lora_path = kwargs.get('inference_lora_path', None)
self.latent_space_version = kwargs.get('latent_space_version', None) self.latent_space_version = kwargs.get('latent_space_version', None)
# only for SDXL models for now # only for SDXL models for now

64
toolkit/stable_diffusion_model.py
View File

@@ -497,33 +497,46 @@ class StableDiffusion:
transformer.to(torch.device(self.quantize_device), dtype=dtype) transformer.to(torch.device(self.quantize_device), dtype=dtype)
flush() flush()
if self.model_config.assistant_lora_path is not None: if self.model_config.assistant_lora_path is not None or self.model_config.inference_lora_path is not None:
if self.model_config.inference_lora_path is not None and self.model_config.assistant_lora_path is not None:
raise ValueError("Cannot load both assistant lora and inference lora at the same time")
if self.model_config.lora_path: if self.model_config.lora_path:
raise ValueError("Cannot load both assistant lora and lora at the same time") raise ValueError("Cannot load both assistant lora and lora at the same time")
if not self.is_flux: if not self.is_flux:
raise ValueError("Assistant lora is only supported for flux models currently") raise ValueError("Assistant/ inference lora is only supported for flux models currently")
if os.path.isdir(self.model_config.assistant_lora_path): load_lora_path = self.model_config.inference_lora_path
self.model_config.assistant_lora_path = os.path.join( if load_lora_path is None:
self.model_config.assistant_lora_path, "pytorch_lora_weights.safetensors" load_lora_path = self.model_config.assistant_lora_path
if os.path.isdir(load_lora_path):
load_lora_path = os.path.join(
load_lora_path, "pytorch_lora_weights.safetensors"
) )
elif not os.path.exists(self.model_config.assistant_lora_path): elif not os.path.exists(load_lora_path):
print(f"Grabbing assistant lora from the hub: {self.model_config.assistant_lora_path}") print(f"Grabbing lora from the hub: {load_lora_path}")
new_lora_path = hf_hub_download( new_lora_path = hf_hub_download(
self.model_config.assistant_lora_path, load_lora_path,
filename="pytorch_lora_weights.safetensors" filename="pytorch_lora_weights.safetensors"
) )
# replace the path # replace the path
self.model_config.assistant_lora_path = new_lora_path load_lora_path = new_lora_path
# for flux, we assume it is flux schnell. We cannot merge in the assistant lora and unmerge it on if self.model_config.inference_lora_path is not None:
# quantized weights so it had to process unmerged (slow). Since schnell samples in just 4 steps self.model_config.inference_lora_path = new_lora_path
# it is better to merge it in now, and sample slowly later, otherwise training is slowed in half if self.model_config.assistant_lora_path is not None:
# so we will merge in now and sample with -1 weight later self.model_config.assistant_lora_path = new_lora_path
self.invert_assistant_lora = True
# trigger it to get merged in if self.model_config.assistant_lora_path is not None:
self.model_config.lora_path = self.model_config.assistant_lora_path # for flux, we assume it is flux schnell. We cannot merge in the assistant lora and unmerge it on
# quantized weights so it had to process unmerged (slow). Since schnell samples in just 4 steps
# it is better to merge it in now, and sample slowly later, otherwise training is slowed in half
# so we will merge in now and sample with -1 weight later
self.invert_assistant_lora = True
# trigger it to get merged in
self.model_config.lora_path = self.model_config.assistant_lora_path
if self.model_config.lora_path is not None: if self.model_config.lora_path is not None:
print("Fusing in LoRA") print("Fusing in LoRA")
@@ -764,6 +777,13 @@ class StableDiffusion:
self.assistant_lora.multiplier = -1.0 self.assistant_lora.multiplier = -1.0
self.assistant_lora.is_active = False self.assistant_lora.is_active = False
if self.model_config.inference_lora_path is not None:
print("Loading inference lora")
self.assistant_lora: 'LoRASpecialNetwork' = load_assistant_lora_from_path(
self.model_config.inference_lora_path, self)
# disable during training
self.assistant_lora.is_active = False
if self.is_pixart and self.vae_scale_factor == 16: if self.is_pixart and self.vae_scale_factor == 16:
# TODO make our own pipeline? # TODO make our own pipeline?
# we generate an image 2x larger, so we need to copy the sizes from larger ones down # we generate an image 2x larger, so we need to copy the sizes from larger ones down
@@ -841,6 +861,12 @@ class StableDiffusion:
else: else:
self.assistant_lora.is_active = False self.assistant_lora.is_active = False
if self.model_config.inference_lora_path is not None:
print("Loading inference lora")
self.assistant_lora.is_active = True
# move weights on to the device
self.assistant_lora.force_to(self.device_torch, self.torch_dtype)
if self.network is not None: if self.network is not None:
self.network.eval() self.network.eval()
network = self.network network = self.network
@@ -1357,6 +1383,12 @@ class StableDiffusion:
else: else:
self.assistant_lora.is_active = True self.assistant_lora.is_active = True
if self.model_config.inference_lora_path is not None:
print("Unloading inference lora")
self.assistant_lora.is_active = False
# move weights off the device
self.assistant_lora.force_to('cpu', self.torch_dtype)
flush() flush()
def get_latent_noise( def get_latent_noise(