ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 03:01:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix issue that prevented ramtorch layer offloading with z_image

Browse Source
This commit is contained in:
Jaret Burkett
2025-12-02 16:14:34 -07:00
parent d42f5af2fc
commit e6c5aead3b
2 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
extensions_built_in/diffusion_models/z_image/z_image.py
View File

@@ -124,6 +124,7 @@ class ZImageModel(BaseModel):
is_transformer=True, is_transformer=True,
target_lin_modules=self.target_lora_modules, target_lin_modules=self.target_lora_modules,
is_assistant_adapter=True, is_assistant_adapter=True,
is_ara=True,
) )
network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True) network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True)
self.print_and_status_update("Merging in assistant LoRA") self.print_and_status_update("Merging in assistant LoRA")
@@ -189,6 +190,10 @@ class ZImageModel(BaseModel):
transformer, transformer,
self.device_torch, self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent, offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[
transformer.x_pad_token,
transformer.cap_pad_token,
]
) )
if self.model_config.low_vram: if self.model_config.low_vram:

5
toolkit/config_modules.py
View File

@@ -1308,6 +1308,11 @@ def validate_configs(
if train_config.bypass_guidance_embedding and train_config.do_guidance_loss: if train_config.bypass_guidance_embedding and train_config.do_guidance_loss:
raise ValueError("Cannot bypass guidance embedding and do guidance loss at the same time. " raise ValueError("Cannot bypass guidance embedding and do guidance loss at the same time. "
"Please set bypass_guidance_embedding to False or do_guidance_loss to False.") "Please set bypass_guidance_embedding to False or do_guidance_loss to False.")
if model_config.accuracy_recovery_adapter is not None:
if model_config.assistant_lora_path is not None:
raise ValueError("Cannot use accuracy recovery adapter and assistant lora at the same time. "
"Please set one of them to None.")
# see if any datasets are caching text embeddings # see if any datasets are caching text embeddings
is_caching_text_embeddings = any(dataset.cache_text_embeddings for dataset in dataset_configs) is_caching_text_embeddings = any(dataset.cache_text_embeddings for dataset in dataset_configs)