From 532625e22190cef43213e754b366a469d7d2b79a Mon Sep 17 00:00:00 2001 From: dxqb Date: Mon, 25 May 2026 18:11:02 +0200 Subject: [PATCH 01/10] Make gradient checkpointing and offloading per-component; centralize model composition in ModelType - Gradient checkpointing and layer offloading are now configured per component (text encoder, transformer, VAE) rather than globally - ModelType centralizes model composition and training method associations Co-Authored-By: Claude Sonnet 4.6 --- modules/model/ChromaModel.py | 6 +- modules/model/ErnieModel.py | 6 +- modules/model/Flux2Model.py | 6 +- modules/model/FluxModel.py | 6 +- modules/model/HiDreamModel.py | 9 +- modules/model/HunyuanVideoModel.py | 6 +- modules/model/PixArtAlphaModel.py | 6 +- modules/model/QwenModel.py | 6 +- modules/model/SanaModel.py | 6 +- modules/model/StableDiffusion3Model.py | 6 +- modules/model/ZImageModel.py | 6 +- modules/modelSetup/BaseChromaSetup.py | 10 +- modules/modelSetup/BaseErnieSetup.py | 5 +- modules/modelSetup/BaseFlux2Setup.py | 17 +- modules/modelSetup/BaseFluxSetup.py | 14 +- modules/modelSetup/BaseHiDreamSetup.py | 23 +- modules/modelSetup/BaseHunyuanVideoSetup.py | 14 +- modules/modelSetup/BasePixArtAlphaSetup.py | 10 +- modules/modelSetup/BaseQwenSetup.py | 10 +- modules/modelSetup/BaseSanaSetup.py | 10 +- .../modelSetup/BaseStableDiffusion3Setup.py | 18 +- .../modelSetup/BaseStableDiffusionSetup.py | 10 +- .../modelSetup/BaseStableDiffusionXLSetup.py | 10 +- modules/modelSetup/BaseWuerstchenSetup.py | 8 +- modules/modelSetup/BaseZImageSetup.py | 10 +- modules/ui/ModelTab.py | 335 +++--------------- modules/ui/OffloadingWindow.py | 75 ---- modules/ui/TopBar.py | 38 +- modules/ui/TrainUI.py | 4 + modules/ui/TrainingTab.py | 161 ++++++--- modules/util/LayerOffloadConductor.py | 11 +- modules/util/checkpointing_util.py | 136 ++++--- modules/util/config/TrainConfig.py | 76 +++- modules/util/create.py | 9 +- .../util/enum/GradientCheckpointingMethod.py | 17 - modules/util/enum/ModelType.py | 55 +++ training_presets/#chroma Finetune 16GB.json | 7 +- training_presets/#chroma Finetune 8GB.json | 7 +- training_presets/#chroma LoRA 8GB.json | 7 +- training_presets/#ernie LoRA 8GB.json | 7 +- training_presets/#flux2 Finetune 16GB.json | 7 +- training_presets/#flux2 LoRA 8GB.json | 10 +- training_presets/#hidream LoRA.json | 11 +- training_presets/#hunyuan video LoRA.json | 8 +- training_presets/#qwen Finetune 16GB.json | 5 +- training_presets/#qwen Finetune 24GB.json | 5 +- training_presets/#qwen LoRA 16GB.json | 5 +- training_presets/#qwen LoRA 24GB.json | 5 +- .../#z-image DeTurbo LoRA 8GB.json | 5 +- training_presets/#z-image Finetune 16GB.json | 5 +- training_presets/#z-image LoRA 8GB.json | 5 +- 51 files changed, 525 insertions(+), 729 deletions(-) delete mode 100644 modules/ui/OffloadingWindow.py delete mode 100644 modules/util/enum/GradientCheckpointingMethod.py diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index d62a8d21b..228ccc424 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -111,8 +111,7 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -121,8 +120,7 @@ def text_encoder_to(self, device: torch.device): self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/ErnieModel.py b/modules/model/ErnieModel.py index fb44425e8..3d15e5385 100644 --- a/modules/model/ErnieModel.py +++ b/modules/model/ErnieModel.py @@ -72,15 +72,13 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/Flux2Model.py b/modules/model/Flux2Model.py index 004c70bee..a501bb116 100644 --- a/modules/model/Flux2Model.py +++ b/modules/model/Flux2Model.py @@ -122,15 +122,13 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index 4b85f1272..22baeec67 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -145,8 +145,7 @@ def text_encoder_1_to(self, device: torch.device): def text_encoder_2_to(self, device: torch.device): if self.text_encoder_2 is not None: - if self.text_encoder_2_offload_conductor is not None and \ - self.text_encoder_2_offload_conductor.layer_offload_activated(): + if self.text_encoder_2_offload_conductor is not None: self.text_encoder_2_offload_conductor.to(device) else: self.text_encoder_2.to(device=device) @@ -155,8 +154,7 @@ def text_encoder_2_to(self, device: torch.device): self.text_encoder_2_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index ca888f4f0..1b45d1488 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -220,8 +220,7 @@ def text_encoder_2_to(self, device: torch.device): def text_encoder_3_to(self, device: torch.device): if self.text_encoder_3 is not None: - if self.text_encoder_3_offload_conductor is not None and \ - self.text_encoder_3_offload_conductor.layer_offload_activated(): + if self.text_encoder_3_offload_conductor is not None: self.text_encoder_3_offload_conductor.to(device) else: self.text_encoder_3.to(device=device) @@ -231,8 +230,7 @@ def text_encoder_3_to(self, device: torch.device): def text_encoder_4_to(self, device: torch.device): if self.text_encoder_4 is not None: - if self.text_encoder_4_offload_conductor is not None and \ - self.text_encoder_4_offload_conductor.layer_offload_activated(): + if self.text_encoder_4_offload_conductor is not None: self.text_encoder_4_offload_conductor.to(device) else: self.text_encoder_4.to(device=device) @@ -241,8 +239,7 @@ def text_encoder_4_to(self, device: torch.device): self.text_encoder_4_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 146569abc..075c406c7 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -157,8 +157,7 @@ def text_encoder_to(self, device: torch.device): def text_encoder_1_to(self, device: torch.device): if self.text_encoder_1 is not None: - if self.text_encoder_1_offload_conductor is not None and \ - self.text_encoder_1_offload_conductor.layer_offload_activated(): + if self.text_encoder_1_offload_conductor is not None: self.text_encoder_1_offload_conductor.to(device) else: self.text_encoder_1.to(device=device) @@ -174,8 +173,7 @@ def text_encoder_2_to(self, device: torch.device): self.text_encoder_2_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index 248b308bb..ab0155fd3 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -112,8 +112,7 @@ def vae_to(self, device: torch.device): self.vae.to(device=device) def text_encoder_to(self, device: torch.device): - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -122,8 +121,7 @@ def text_encoder_to(self, device: torch.device): self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/QwenModel.py b/modules/model/QwenModel.py index afa6c24fe..1aadcd467 100644 --- a/modules/model/QwenModel.py +++ b/modules/model/QwenModel.py @@ -81,8 +81,7 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): #TODO share more code between models if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -91,8 +90,7 @@ def text_encoder_to(self, device: torch.device): #TODO share more code between m self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index d00d75fa2..39ac83310 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -114,8 +114,7 @@ def vae_to(self, device: torch.device): self.vae.to(device=device) def text_encoder_to(self, device: torch.device): - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) @@ -124,8 +123,7 @@ def text_encoder_to(self, device: torch.device): self.text_encoder_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index 65f29ac49..378bb78f8 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -174,8 +174,7 @@ def text_encoder_2_to(self, device: torch.device): def text_encoder_3_to(self, device: torch.device): if self.text_encoder_3 is not None: - if self.text_encoder_3_offload_conductor is not None and \ - self.text_encoder_3_offload_conductor.layer_offload_activated(): + if self.text_encoder_3_offload_conductor is not None: self.text_encoder_3_offload_conductor.to(device) else: self.text_encoder_3.to(device=device) @@ -184,8 +183,7 @@ def text_encoder_3_to(self, device: torch.device): self.text_encoder_3_lora.to(device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/model/ZImageModel.py b/modules/model/ZImageModel.py index 7fd9e52cb..167b58f47 100644 --- a/modules/model/ZImageModel.py +++ b/modules/model/ZImageModel.py @@ -83,15 +83,13 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): #TODO share more code between models if self.text_encoder is not None: - if self.text_encoder_offload_conductor is not None and \ - self.text_encoder_offload_conductor.layer_offload_activated(): + if self.text_encoder_offload_conductor is not None: self.text_encoder_offload_conductor.to(device) else: self.text_encoder.to(device=device) def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None and \ - self.transformer_offload_conductor.layer_offload_activated(): + if self.transformer_offload_conductor is not None: self.transformer_offload_conductor.to(device) else: self.transformer.to(device=device) diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 0eb623399..f8cc33187 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -49,12 +49,10 @@ def setup_optimizations( model: ChromaModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_chroma_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_chroma_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c164dc0bc..72f8ee2d8 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -42,9 +42,8 @@ def setup_optimizations( model: ErnieModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_ernie_transformer(model.transformer, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_ernie_transformer(model.transformer, config, config.transformer) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index b91cd8af8..6e6e7320f 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -45,16 +45,13 @@ def setup_optimizations( model: Flux2Model, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_flux2_transformer(model.transformer, config) - if model.text_encoder is not None: - if model.is_dev(): - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config) - else: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_flux2_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + if model.is_dev(): + model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config, config.text_encoder) + else: + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 9bae83cde..e6638ddb2 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -49,14 +49,12 @@ def setup_optimizations( model: FluxModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_flux_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - model.text_encoder_2_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_flux_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + model.text_encoder_2_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 17fbcc0d6..b5a856b73 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -49,19 +49,16 @@ def setup_optimizations( model: HiDreamModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_hi_dream_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) - if model.text_encoder_3 is not None: - model.text_encoder_3_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config) - if model.text_encoder_4 is not None: - model.text_encoder_4_offload_conductor = \ - enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_hi_dream_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) + if model.text_encoder_3 is not None and config.text_encoder_3.checkpointing_or_offloading_enabled(): + model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) + if model.text_encoder_4 is not None and config.text_encoder_4.checkpointing_or_offloading_enabled(): + model.text_encoder_4_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config, config.text_encoder_4) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index b072bf4ba..2e32165e7 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -49,14 +49,12 @@ def setup_optimizations( model: HunyuanVideoModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - model.text_encoder_1_offload_conductor = \ - enable_checkpointing_for_llama_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_1_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 3069b4884..4935ebbfd 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -51,12 +51,10 @@ def setup_optimizations( model: PixArtAlphaModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.vae.enable_gradient_checkpointing() - model.transformer_offload_conductor = \ - enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, offload_enabled=True) - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, config.transformer, offload_enabled=True) + if config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index a8a7be8f6..82004d8ac 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -46,12 +46,10 @@ def setup_optimizations( model: QwenModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_qwen_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 0c9ea6da0..c26acbc82 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -52,12 +52,10 @@ def setup_optimizations( config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - # model.vae.enable_gradient_checkpointing() - model.transformer_offload_conductor = \ - enable_checkpointing_for_sana_transformer(model.transformer, config) - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_gemma_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_sana_transformer(model.transformer, config, config.transformer) + if config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_gemma_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index de5dc04e8..efe761427 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -48,16 +48,14 @@ def setup_optimizations( model: StableDiffusion3Model, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) - if model.text_encoder_3 is not None: - model.text_encoder_3_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) + if model.text_encoder_3 is not None and config.text_encoder_3.checkpointing_or_offloading_enabled(): + model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 8cf63ac07..a16c72cb8 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -51,11 +51,13 @@ def setup_optimizations( model: StableDiffusionModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.vae.enable_gradient_checkpointing() + if config.unet.checkpointing_or_offloading_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, offload_enabled=False) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + if config.vae.checkpointing_enabled(): + model.vae.enable_gradient_checkpointing() + if config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 61cb1e457..7fe35e1e1 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -48,11 +48,13 @@ def setup_optimizations( model: StableDiffusionXLModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): + if config.unet.checkpointing_or_offloading_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, offload_enabled=False) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + if config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if config.text_encoder_2.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index 48cd05942..abc77bcdd 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -57,13 +57,13 @@ def setup_optimizations( model: WuerstchenModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): + if config.prior.checkpointing_or_offloading_enabled(): if model.model_type.is_wuerstchen_v2(): model.prior_prior.enable_gradient_checkpointing() - enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config) elif model.model_type.is_stable_cascade(): - enable_checkpointing_for_stable_cascade_blocks(model.prior_prior, config) - enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config) + enable_checkpointing_for_stable_cascade_blocks(model.prior_prior, config, config.prior) + if config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.decoder_vqgan) diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index b822d7304..d64e7e88a 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -47,12 +47,10 @@ def setup_optimizations( model: ZImageModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_z_image_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_z_image_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index ff17ea3ba..5bc3092d3 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -45,307 +45,66 @@ def refresh_ui(self): base_frame.grid_columnconfigure(3, weight=0) base_frame.grid_columnconfigure(4, weight=1) - if self.train_config.model_type.is_stable_diffusion(): #TODO simplify - self.__setup_stable_diffusion_ui(base_frame) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(base_frame) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(base_frame) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(base_frame) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(base_frame) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(base_frame) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(base_frame) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(base_frame) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(base_frame) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(base_frame) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(base_frame) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(base_frame) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(base_frame) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(base_frame) - - def __setup_stable_diffusion_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method in [ - TrainingMethod.FINE_TUNE, - TrainingMethod.FINE_TUNE_VAE, - ], - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_3_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_2_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + self.__setup_ui(base_frame) - def __setup_z_image_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + def __setup_ui(self, frame): + model_type = self.train_config.model_type + training_method = self.train_config.training_method + parts = model_type.model_parts() - def __setup_ernie_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + # The transformer override path exists only for these architectures; SD3, PixArt, Sana + # and HiDream have a transformer but expose no override field. + allow_override_transformer = ( + model_type.is_flux() + or model_type.is_z_image() + or model_type.is_ernie() + or model_type.is_chroma() + or model_type.is_qwen() + or model_type.is_hunyuan_video() ) - def __setup_chroma_ui(self, frame): row = 0 row = self.__create_base_dtype_components(frame, row) row = self.__create_base_components( frame, row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + has_unet="unet" in parts, + has_prior="prior" in parts, + allow_override_prior=model_type.is_stable_cascade(), + has_transformer="transformer" in parts, + allow_override_transformer=allow_override_transformer, + has_text_encoder=not model_type.has_multiple_text_encoders(), + has_text_encoder_1=model_type.has_multiple_text_encoders(), + has_text_encoder_2="text_encoder_2" in parts, + has_text_encoder_3="text_encoder_3" in parts, + has_text_encoder_4="text_encoder_4" in parts, + allow_override_text_encoder_4="text_encoder_4" in parts, + has_vae="vae" in parts, + ) + if "effnet_encoder" in parts: + row = self.__create_effnet_encoder_components(frame, row) + if "decoder" in parts: + row = self.__create_decoder_components(frame, row, "decoder_text_encoder" in parts) + + if model_type.is_sana(): + allow_safetensors = training_method != TrainingMethod.FINE_TUNE + elif model_type.is_wuerstchen(): + allow_safetensors = training_method != TrainingMethod.FINE_TUNE \ + or model_type.is_stable_cascade() + else: + allow_safetensors = True + + if model_type.is_stable_diffusion(): + allow_diffusers = training_method in [TrainingMethod.FINE_TUNE, TrainingMethod.FINE_TUNE_VAE] + else: + allow_diffusers = training_method == TrainingMethod.FINE_TUNE - def __setup_qwen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_xl_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_wuerstchen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_prior=True, - allow_override_prior=self.train_config.model_type.is_stable_cascade(), - has_text_encoder=True, - ) - row = self.__create_effnet_encoder_components(frame, row) - row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE - or self.train_config.model_type.is_stable_cascade(), - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_pixart_alpha_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_sana_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hunyuan_video_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hi_dream_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_text_encoder_4=True, - allow_override_text_encoder_4=True, - has_vae=True, - ) row = self.__create_output_components( frame, row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_safetensors=allow_safetensors, + allow_diffusers=allow_diffusers, + allow_legacy_safetensors=training_method == TrainingMethod.LORA, ) def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: diff --git a/modules/ui/OffloadingWindow.py b/modules/ui/OffloadingWindow.py deleted file mode 100644 index 54035e121..000000000 --- a/modules/ui/OffloadingWindow.py +++ /dev/null @@ -1,75 +0,0 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.GradientCheckpointingMethod import ( - GradientCheckpointingMethod, -) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import customtkinter as ctk - - -class OffloadingWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None - - self.title("Offloading") - self.geometry("800x400") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing") - - # gradient checkpointing layer offloading - components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") - - frame.pack(fill="both", expand=1) - return frame - - def __ok(self): - self.destroy() diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py index 820fdb71a..d2ca78f7b 100644 --- a/modules/ui/TopBar.py +++ b/modules/ui/TopBar.py @@ -112,37 +112,13 @@ def __create_training_method(self): if self.training_method: self.training_method.destroy() - values = [] - #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), - ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ] + labels = { + TrainingMethod.FINE_TUNE: "Fine Tune", + TrainingMethod.LORA: "LoRA", + TrainingMethod.EMBEDDING: "Embedding", + TrainingMethod.FINE_TUNE_VAE: "Fine Tune VAE", + } + values = [(labels[m], m) for m in self.train_config.model_type.supported_training_methods()] # training method self.training_method = components.options_kv( diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index ba90d2e64..f822895d6 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -308,6 +308,10 @@ def create_general_tab(self, master): tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) + components.label(frame, 11, 2, "Async Offloading", + tooltip="Overlaps CPU<->GPU transfers with computation using CUDA streams. Applies to every offloaded component") + components.switch(frame, 11, 3, self.ui_state, "async_offloading") + components.label(frame, 12, 0, "Multi-GPU", tooltip="Enable multi-GPU training") components.switch(frame, 12, 1, self.ui_state, "multi_gpu") diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index bcca11ae9..2d24d968c 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -1,4 +1,3 @@ -from modules.ui.OffloadingWindow import OffloadingWindow from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow @@ -6,13 +5,13 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod from modules.util.enum.LearningRateScaler import LearningRateScaler from modules.util.enum.LearningRateScheduler import LearningRateScheduler from modules.util.enum.LossScaler import LossScaler from modules.util.enum.LossWeight import LossWeight from modules.util.enum.Optimizer import Optimizer from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.optimizer_util import change_optimizer from modules.util.ui import components from modules.util.ui.UIState import UIState @@ -99,6 +98,8 @@ def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): self.__create_unet_frame(column_1, 1) self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) + if self.train_config.training_method == TrainingMethod.FINE_TUNE_VAE: + self.__create_vae_frame(column_2, 0) self.__create_masked_frame(column_2, 1) self.__create_loss_frame(column_2, 2) self.__create_layer_frame(column_2, 3) @@ -375,19 +376,6 @@ def __create_base2_frame(self, master, row, video_training_enabled: bool=False, components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") row += 1 - # gradient checkpointing - components.label(frame, row, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing", adv_command=self.__open_offloading_window) - row += 1 - - # gradient checkpointing layer offloading - components.label(frame, row, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") - row += 1 - # train dtype components.label(frame, row, 0, "Train Data Type", tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") @@ -434,38 +422,70 @@ def __create_base2_frame(self, master, row, video_training_enabled: bool=False, tooltip="Enables circular padding for all conv layers to better train seamless images") components.switch(frame, row, 1, self.ui_state, "force_circular_padding") + def __create_offloading_widgets(self, frame, row, part, supports_checkpointing=True, supports_activation_offloading=False): + if supports_checkpointing: + components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for this component. Reduces VRAM usage at the cost of training speed") + components.switch(frame, row, 1, self.ui_state, f"{part}.gradient_checkpointing") + row += 1 + + components.label(frame, row, 0, "Layer Offload Fraction", + tooltip="Fraction of this component's layers to offload to CPU to reduce VRAM usage. Increases training time and RAM usage. 0=disabled, 1=all layers") + components.entry(frame, row, 1, self.ui_state, f"{part}.offload_fraction") + row += 1 + + if supports_activation_offloading: + components.label(frame, row, 0, "Offload Activations", + tooltip="Offloads this component's activations to CPU during training to reduce VRAM usage") + components.switch(frame, row, 1, self.ui_state, f"{part}.activation_offloading") + row += 1 + + return row + def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 if supports_training: - components.label(frame, 0, 0, "Train Text Encoder", + components.label(frame, row, 0, "Train Text Encoder", tooltip="Enables training the text encoder model") - components.switch(frame, 0, 1, self.ui_state, "text_encoder.train") + components.switch(frame, row, 1, self.ui_state, "text_encoder.train") + row += 1 + else: + # no Train switch to act as the frame's header, so add an explicit one + components.label(frame, row, 0, "Text Encoder") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "text_encoder", supports_checkpointing=supports_training) # dropout - components.label(frame, 1, 0, "Caption Dropout Probability", + components.label(frame, row, 0, "Caption Dropout Probability", tooltip="The Probability for dropping the text encoder conditioning") - components.entry(frame, 1, 1, self.ui_state, "text_encoder.dropout_probability") + components.entry(frame, row, 1, self.ui_state, "text_encoder.dropout_probability") + row += 1 if supports_training: # train text encoder epochs - components.label(frame, 2, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the text encoder") - components.time_entry(frame, 2, 1, self.ui_state, "text_encoder.stop_training_after", + components.time_entry(frame, row, 1, self.ui_state, "text_encoder.stop_training_after", "text_encoder.stop_training_after_unit", supports_time_units=False) + row += 1 # text encoder learning rate - components.label(frame, 3, 0, "Text Encoder Learning Rate", + components.label(frame, row, 0, "Text Encoder Learning Rate", tooltip="The learning rate of the text encoder. Overrides the base learning rate") - components.entry(frame, 3, 1, self.ui_state, "text_encoder.learning_rate") + components.entry(frame, row, 1, self.ui_state, "text_encoder.learning_rate") + row += 1 if supports_clip_skip: # text encoder layer skip (clip skip) - components.label(frame, 4, 0, "Clip Skip", + components.label(frame, row, 0, "Clip Skip", tooltip="The number of additional clip layers to skip. 0 = the model default") - components.entry(frame, 4, 1, self.ui_state, "text_encoder_layer_skip") + components.entry(frame, row, 1, self.ui_state, "text_encoder_layer_skip") + row += 1 if supports_sequence_length: # text encoder sequence length @@ -503,6 +523,8 @@ def __create_text_encoder_n_frame( components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") row += 1 + row = self.__create_offloading_widgets(frame, row, f"text_encoder{suffix}") + # train text encoder embedding components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", tooltip=f"Enables training embeddings for the text encoder {i} model") @@ -560,82 +582,119 @@ def __create_unet_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train unet - components.label(frame, 0, 0, "Train UNet", + components.label(frame, row, 0, "Train UNet", tooltip="Enables training the UNet model") - components.switch(frame, 0, 1, self.ui_state, "unet.train") + components.switch(frame, row, 1, self.ui_state, "unet.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "unet", supports_activation_offloading=True) # train unet epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the UNet") - components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", supports_time_units=False) + row += 1 # unet learning rate - components.label(frame, 2, 0, "UNet Learning Rate", + components.label(frame, row, 0, "UNet Learning Rate", tooltip="The learning rate of the UNet. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") + components.entry(frame, row, 1, self.ui_state, "unet.learning_rate") + row += 1 # rescale noise scheduler to zero terminal SNR - rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", + rescale_label = components.label(frame, row, 0, "Rescale Noise Scheduler + V-pred", tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") rescale_label.configure(wraplength=130, justify="left") - components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + components.switch(frame, row, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + row += 1 + + def __create_vae_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + row = 0 + + components.label(frame, row, 0, "Train VAE", + tooltip="Enables training the VAE model") + components.switch(frame, row, 1, self.ui_state, "vae.train") + row += 1 + + components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for the VAE. Reduces VRAM usage at the cost of training speed") + components.switch(frame, row, 1, self.ui_state, "vae.gradient_checkpointing") + row += 1 def __create_prior_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train prior - components.label(frame, 0, 0, "Train Prior", + components.label(frame, row, 0, "Train Prior", tooltip="Enables training the Prior model") - components.switch(frame, 0, 1, self.ui_state, "prior.train") + components.switch(frame, row, 1, self.ui_state, "prior.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "prior", supports_activation_offloading=True) # train prior epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the Prior") - components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", supports_time_units=False) + row += 1 # prior learning rate - components.label(frame, 2, 0, "Prior Learning Rate", + components.label(frame, row, 0, "Prior Learning Rate", tooltip="The learning rate of the Prior. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") + components.entry(frame, row, 1, self.ui_state, "prior.learning_rate") + row += 1 def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train transformer - components.label(frame, 0, 0, "Train Transformer", + components.label(frame, row, 0, "Train Transformer", tooltip="Enables training the Transformer model") - components.switch(frame, 0, 1, self.ui_state, "transformer.train") + components.switch(frame, row, 1, self.ui_state, "transformer.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "transformer", supports_activation_offloading=True) # train transformer epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the Transformer") - components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", supports_time_units=False) + row += 1 # transformer learning rate - components.label(frame, 2, 0, "Transformer Learning Rate", + components.label(frame, row, 0, "Transformer Learning Rate", tooltip="The learning rate of the Transformer. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") + components.entry(frame, row, 1, self.ui_state, "transformer.learning_rate") + row += 1 if supports_force_attention_mask: # transformer learning rate - components.label(frame, 3, 0, "Force Attention Mask", + components.label(frame, row, 0, "Force Attention Mask", tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") - components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") + components.switch(frame, row, 1, self.ui_state, "transformer.attention_mask") + row += 1 if supports_guidance_scale: # guidance scale - components.label(frame, 4, 0, "Guidance Scale", + components.label(frame, row, 0, "Guidance Scale", tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") - components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") + components.entry(frame, row, 1, self.ui_state, "transformer.guidance_scale") + row += 1 def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): frame = ctk.CTkFrame(master=master, corner_radius=5) @@ -838,10 +897,6 @@ def __open_timestep_distribution_window(self): window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) self.master.wait_window(window) - def __open_offloading_window(self): - window = OffloadingWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - def __restore_optimizer_config(self, *args): optimizer_config = change_optimizer(self.train_config) self.ui_state.get_var("optimizer").update(optimizer_config) diff --git a/modules/util/LayerOffloadConductor.py b/modules/util/LayerOffloadConductor.py index 0b7e6c7ca..062d1391e 100644 --- a/modules/util/LayerOffloadConductor.py +++ b/modules/util/LayerOffloadConductor.py @@ -2,7 +2,7 @@ import random from typing import Any -from modules.util.config.TrainConfig import TrainConfig +from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig from modules.util.quantization_util import get_offload_tensor_bytes, offload_quantized from modules.util.torch_util import ( create_stream_context, @@ -566,6 +566,7 @@ def __init__( self, module: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, ): super().__init__() @@ -573,16 +574,16 @@ def __init__( self.__layers = [] self.__layer_device_map = [] - self.__layer_offload_fraction = config.layer_offload_fraction + self.__layer_offload_fraction = part.offload_fraction self.__layer_activations_included_offload_param_indices_map = [] self.__train_device = torch.device(config.train_device) self.__temp_device = torch.device(config.temp_device) - self.__offload_activations = config.gradient_checkpointing.offload() and config.enable_activation_offloading - self.__offload_layers = config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 - self.__async_transfer = self.__train_device.type == "cuda" and config.enable_async_offloading + self.__offload_activations = part.activation_offloading + self.__offload_layers = part.offload_fraction > 0 + self.__async_transfer = self.__train_device.type == "cuda" and config.async_offloading if self.__async_transfer: self.__train_stream = torch.cuda.default_stream(self.__train_device) diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index ba3f7fec2..d5b9f765c 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -3,7 +3,7 @@ from typing import Any from modules.util.compile_util import init_compile -from modules.util.config.TrainConfig import TrainConfig +from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig from modules.util.LayerOffloadConductor import LayerOffloadConductor from modules.util.torch_util import add_dummy_grad_fn_, has_grad_fn @@ -74,21 +74,25 @@ def __init__(self, *args, **kwargs): class CheckpointLayer(BaseCheckpointLayer): - def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device): + def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, checkpointing: bool = True): super().__init__() assert (orig_module is None or orig_forward is None) and not (orig_module is None and orig_forward is None) self.checkpoint = orig_module self.orig_forward = orig_forward + self.checkpointing = checkpointing # dummy tensor that requires grad is needed for checkpointing to work when training a LoRA self.dummy = torch.zeros((1,), device=train_device, requires_grad=True) - def __checkpointing_forward(self, dummy: torch.Tensor, *args, **kwargs): + def __orig(self, *args, **kwargs): return self.orig_forward(*args, **kwargs) if self.checkpoint is None else self.checkpoint(*args, **kwargs) + def __checkpointing_forward(self, dummy: torch.Tensor, *args, **kwargs): + return self.__orig(*args, **kwargs) + def forward(self, *args, **kwargs): - if torch.is_grad_enabled(): + if self.checkpointing and torch.is_grad_enabled(): return torch.utils.checkpoint.checkpoint( self.__checkpointing_forward, self.dummy, @@ -97,7 +101,7 @@ def forward(self, *args, **kwargs): use_reentrant=False ) else: - return self.orig_forward(*args, **kwargs) if self.checkpoint is None else self.checkpoint(*args, **kwargs) + return self.__orig(*args, **kwargs) class OffloadCheckpointLayer(BaseCheckpointLayer): def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, conductor: LayerOffloadConductor, layer_index: int): @@ -153,6 +157,7 @@ def create_checkpoint( train_device: torch.device, include_from_offload_param_names: list[str] = None, conductor: LayerOffloadConductor | None = None, + checkpointing: bool = True, layer_index: int = 0, compile: bool = False, ) -> Callable: @@ -164,6 +169,9 @@ def create_checkpoint( conductor.add_layer(orig_module, included_offload_param_indices) if conductor is not None and conductor.offload_activated(): + # offloading is structurally coupled to use_reentrant=True checkpointing during the back pass + # (the recompute is what fires before_layer/after_layer in the backward direction), so the offload + # layer always checkpoints when grad is enabled, regardless of the part's gradient_checkpointing flag. if compile: layer = OffloadCheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, conductor=conductor, layer_index=layer_index) #don't compile the checkpointing layer - offloading cannot be compiled: @@ -176,12 +184,12 @@ def create_checkpoint( return orig_module else: if compile: - layer = CheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device) + layer = CheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, checkpointing=checkpointing) #do compile the checkpointing layer - slightly faster layer.compile(fullgraph=True) return layer else: - layer = CheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device) + layer = CheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device, checkpointing=checkpointing) orig_module.forward = layer.forward return orig_module @@ -189,6 +197,7 @@ def _create_checkpoints_for_module_list( module_list: nn.ModuleList, include_from_offload_param_names: list[str], conductor: LayerOffloadConductor, + checkpointing: bool, train_device: torch.device, layer_index: int, compile: bool, @@ -200,7 +209,7 @@ def _create_checkpoints_for_module_list( module_list[i] = create_checkpoint( layer, train_device, include_from_offload_param_names, - conductor, layer_index, compile=compile, + conductor, checkpointing, layer_index, compile=compile, ) layer_index += 1 return layer_index @@ -213,11 +222,15 @@ def _remove_checkpoint_keys(module, state_dict, prefix, local_metadata): def enable_checkpointing( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, compile: bool, lists, # if there are multiple entries in this list, they must be in the exact order they are executed - otherwise offloading fails offload_enabled: bool = True, -) -> LayerOffloadConductor: - conductor = LayerOffloadConductor(model, config) +) -> LayerOffloadConductor | None: + # a conductor exists iff this part actually offloads (and the component supports conductor offloading) + offload = offload_enabled and part.offloading_enabled() + conductor = LayerOffloadConductor(model, config, part) if offload else None + checkpointing = part.checkpointing_enabled() layer_index = 0 for type_or_list, param_names in lists: @@ -228,7 +241,8 @@ def enable_checkpointing( layer_index = _create_checkpoints_for_module_list( module_list, param_names, - conductor if offload_enabled else None, + conductor, + checkpointing, torch.device(config.train_device), layer_index, compile = compile, @@ -242,7 +256,8 @@ def enable_checkpointing( layer_index = _create_checkpoints_for_module_list( module_list, param_names, - conductor if offload_enabled else None, + conductor, + checkpointing, torch.device(config.train_device), layer_index, compile = compile, @@ -253,9 +268,10 @@ def enable_checkpointing( def enable_checkpointing_for_basic_transformer_blocks( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, offload_enabled: bool, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (BasicTransformerBlock , []), ], offload_enabled = offload_enabled, @@ -264,16 +280,18 @@ def enable_checkpointing_for_basic_transformer_blocks( def enable_checkpointing_for_clip_encoder_layers( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, ): - return enable_checkpointing(model, config, False, [ + return enable_checkpointing(model, config, part, False, [ (CLIPEncoderLayer, []), # No activation offloading for text encoders, because the output might be taken from the middle of the network ]) def enable_checkpointing_for_stable_cascade_blocks( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (SDCascadeResBlock, []), (SDCascadeAttnBlock, []), (SDCascadeTimestepBlock, []), @@ -282,8 +300,9 @@ def enable_checkpointing_for_stable_cascade_blocks( def enable_checkpointing_for_t5_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (T5Block, []), ]) @@ -291,8 +310,9 @@ def enable_checkpointing_for_t5_encoder_layers( def enable_checkpointing_for_gemma_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (Gemma2DecoderLayer, []), ]) @@ -300,17 +320,19 @@ def enable_checkpointing_for_gemma_layers( def enable_checkpointing_for_llama_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (LlamaDecoderLayer, []), ]) def enable_checkpointing_for_mistral_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ - (MistralDecoderLayer, []), + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (MistralDecoderLayer, []), # no activation offloading: this encoder is never trained ]) @@ -318,32 +340,36 @@ def enable_checkpointing_for_mistral_encoder_layers( def enable_checkpointing_for_qwen25vl_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (Qwen2_5_VLDecoderLayer, []), # TODO No activation offloading for other encoders, see above. But clip skip is not implemented for QwenVL. Then do activation offloading? ]) def enable_checkpointing_for_qwen3_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ - (Qwen3DecoderLayer, []), # No activation offloading, because hidden states are taken from the middle of the network by Flux2 + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (Qwen3DecoderLayer, []), # no activation offloading: this encoder is never trained ]) def enable_checkpointing_for_stable_diffusion_3_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (JointTransformerBlock, ["hidden_states", "encoder_hidden_states"]), ]) def enable_checkpointing_for_flux_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -351,8 +377,9 @@ def enable_checkpointing_for_flux_transformer( def enable_checkpointing_for_flux2_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -361,8 +388,9 @@ def enable_checkpointing_for_flux2_transformer( def enable_checkpointing_for_chroma_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -371,16 +399,18 @@ def enable_checkpointing_for_chroma_transformer( def enable_checkpointing_for_qwen_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), ]) def enable_checkpointing_for_z_image_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.noise_refiner, ["x"]), (model.context_refiner, ["x"]), (model.layers, ["x"]), @@ -390,16 +420,18 @@ def enable_checkpointing_for_z_image_transformer( def enable_checkpointing_for_sana_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (SanaTransformerBlock, ["hidden_states"]), ]) def enable_checkpointing_for_hunyuan_video_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (HunyuanVideoIndividualTokenRefinerBlock, ["hidden_states" ]), (HunyuanVideoTransformerBlock, ["hidden_states", "encoder_hidden_states"]), (HunyuanVideoSingleTransformerBlock, ["hidden_states" ]), @@ -408,8 +440,9 @@ def enable_checkpointing_for_hunyuan_video_transformer( def enable_checkpointing_for_hi_dream_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (HiDreamImageTransformerBlock, ["hidden_states", "encoder_hidden_states"]), (HiDreamImageSingleTransformerBlock, ["hidden_states" ]), ]) @@ -417,7 +450,8 @@ def enable_checkpointing_for_hi_dream_transformer( def enable_checkpointing_for_ernie_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.layers, ["x"]), ]) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index bbc70a030..125b53669 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -13,7 +13,6 @@ from modules.util.enum.ConfigPart import ConfigPart from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod from modules.util.enum.GradientReducePrecision import GradientReducePrecision from modules.util.enum.ImageFormat import ImageFormat from modules.util.enum.LearningRateScaler import LearningRateScaler @@ -267,10 +266,27 @@ class TrainModelPartConfig(BaseConfig): train_embedding: bool attention_mask: bool guidance_scale: float + gradient_checkpointing: bool + offload_fraction: float + activation_offloading: bool def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) + def offloading_enabled(self) -> bool: + # a conductor should exist iff this is True. Layer offloading applies even to frozen parts (to fit + # them in VRAM), but activation offloading only does work during a backward pass, so it only applies + # when the part is trained -- even if activation_offloading is True in the config. + return self.offload_fraction > 0 or (self.activation_offloading and self.train) + + def checkpointing_enabled(self) -> bool: + # the inner torch checkpoint() should run iff this is True + return self.gradient_checkpointing and self.train + + def checkpointing_or_offloading_enabled(self) -> bool: + # whether the checkpoint layer wrapper needs to be installed for this part at all + return self.checkpointing_enabled() or self.offloading_enabled() + @staticmethod def default_values(): data = [] @@ -287,6 +303,9 @@ def default_values(): data.append(("train_embedding", True, bool, False)) data.append(("attention_mask", False, bool, False)) data.append(("guidance_scale", 1.0, float, False)) + data.append(("gradient_checkpointing", True, bool, False)) + data.append(("offload_fraction", 0.0, float, False)) + data.append(("activation_offloading", True, bool, False)) return TrainModelPartConfig(data) @@ -374,10 +393,7 @@ class TrainConfig(BaseConfig): output_dtype: DataType output_model_format: ModelFormat output_model_destination: str - gradient_checkpointing: GradientCheckpointingMethod - enable_async_offloading: bool - enable_activation_offloading: bool - layer_offload_fraction: float + async_offloading: bool force_circular_padding: bool compile: bool @@ -558,7 +574,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=10, + config_version=11, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -570,6 +586,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, + 10: self.__migration_10, } ) @@ -716,12 +733,14 @@ def __migration_3(self, data: dict) -> dict: def __migration_4(self, data: dict) -> dict: migrated_data = data.copy() + # Translate the old bool form of gradient_checkpointing into the v5..v10 + # string/enum form. __migration_10 later fans this out per-component. gradient_checkpointing = migrated_data.pop("gradient_checkpointing", True) if gradient_checkpointing: - migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.ON + migrated_data["gradient_checkpointing"] = "ON" else: - migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.OFF + migrated_data["gradient_checkpointing"] = "OFF" return migrated_data @@ -789,6 +808,42 @@ def replace_dtype(part: str): return migrated_data + def __migration_10(self, data: dict) -> dict: + migrated_data = data.copy() + + # Fan the four old global offload/checkpointing settings out per-component. + # After __migration_4 gradient_checkpointing is a string "OFF"/"ON"/"CPU_OFFLOADED". + gc = migrated_data.pop("gradient_checkpointing", "ON") + act = migrated_data.pop("enable_activation_offloading", True) + frac = migrated_data.pop("layer_offload_fraction", 0.0) + migrated_data["async_offloading"] = migrated_data.pop("enable_async_offloading", True) + + def fan_out(part: str): + if part in migrated_data: + migrated_data[part]["gradient_checkpointing"] = gc != "OFF" + migrated_data[part]["activation_offloading"] = (gc == "CPU_OFFLOADED") and act + migrated_data[part]["offload_fraction"] = frac if gc == "CPU_OFFLOADED" else 0.0 + + fan_out("unet") + fan_out("prior") + fan_out("transformer") + fan_out("text_encoder") + fan_out("text_encoder_2") + fan_out("text_encoder_3") + fan_out("text_encoder_4") + fan_out("vae") + fan_out("effnet_encoder") + fan_out("decoder") + fan_out("decoder_text_encoder") + fan_out("decoder_vqgan") + + return migrated_data + + def model_part_configs(self) -> list[TrainModelPartConfig]: + # the per-part configs for the components this model_type actually has. Avoids "phantom" parts whose + # fields keep their defaults (train=True) or migrated offload values but don't exist in the model. + return [getattr(self, name) for name in self.model_type.model_parts()] + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, @@ -958,10 +1013,7 @@ def default_values() -> 'TrainConfig': data.append(("output_dtype", DataType.FLOAT_32, DataType, False)) data.append(("output_model_format", ModelFormat.SAFETENSORS, ModelFormat, False)) data.append(("output_model_destination", "models/model.safetensors", str, False)) - data.append(("gradient_checkpointing", GradientCheckpointingMethod.ON, GradientCheckpointingMethod, False)) - data.append(("enable_async_offloading", True, bool, False)) - data.append(("enable_activation_offloading", True, bool, False)) - data.append(("layer_offload_fraction", 0.0, float, False)) + data.append(("async_offloading", True, bool, False)) data.append(("force_circular_padding", False, bool, False)) data.append(("compile", False, bool, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 7c0194da8..0c0adba16 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -110,7 +110,11 @@ def create_data_loader( train_progress: TrainProgress | None = None, is_validation: bool = False ) -> BaseDataLoader | None: - if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 and config.dataloader_threads > 1: + # Layer offloading uses a non-thread-safe conductor. This check is too broad: it trips whenever any model + # part does layer offloading, even though only a component that is actually cached really runs in the + # dataloader worker threads. + # TODO: narrow this to the cached components only. + if config.dataloader_threads > 1 and any(part.offload_fraction > 0 for part in config.model_part_configs()): raise RuntimeError('layer offloading can not be activated if "dataloader_threads" > 1') if train_progress is None: @@ -133,7 +137,8 @@ def create_optimizer( if optimizer_config.optimizer is None: return None - if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0: + # a trained, layer-offloaded part has its params evicted during the back pass, so it needs fused_back_pass + if any(part.offload_fraction > 0 and part.train for part in config.model_part_configs()): if (not optimizer_config.optimizer.supports_fused_back_pass() or not optimizer_config.fused_back_pass) \ and config.training_method == TrainingMethod.FINE_TUNE: raise RuntimeError('layer offloading can only be used for fine tuning when using an optimizer that supports "fused_back_pass"') diff --git a/modules/util/enum/GradientCheckpointingMethod.py b/modules/util/enum/GradientCheckpointingMethod.py deleted file mode 100644 index d3f05666a..000000000 --- a/modules/util/enum/GradientCheckpointingMethod.py +++ /dev/null @@ -1,17 +0,0 @@ -from enum import Enum - - -class GradientCheckpointingMethod(Enum): - OFF = 'OFF' - ON = 'ON' - CPU_OFFLOADED = 'CPU_OFFLOADED' - - def __str__(self): - return self.value - - def enabled(self): - return self == GradientCheckpointingMethod.ON \ - or self == GradientCheckpointingMethod.CPU_OFFLOADED - - def offload(self): - return self == GradientCheckpointingMethod.CPU_OFFLOADED diff --git a/modules/util/enum/ModelType.py b/modules/util/enum/ModelType.py index 5b6658501..4df90ba52 100644 --- a/modules/util/enum/ModelType.py +++ b/modules/util/enum/ModelType.py @@ -1,5 +1,7 @@ from enum import Enum +from modules.util.enum.TrainingMethod import TrainingMethod + class ModelType(Enum): STABLE_DIFFUSION_15 = 'STABLE_DIFFUSION_15' @@ -166,6 +168,59 @@ def is_flow_matching(self) -> bool: def is_video_model(self) -> bool: return self.is_hunyuan_video() #incase we add more video models in the future + def model_parts(self) -> tuple[str, ...]: + return _MODEL_PARTS[self] + + def supported_training_methods(self) -> tuple[TrainingMethod, ...]: + if self.is_stable_diffusion(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING, TrainingMethod.FINE_TUNE_VAE) + if self.is_stable_diffusion_3() \ + or self.is_stable_diffusion_xl() \ + or self.is_wuerstchen() \ + or self.is_pixart() \ + or self.is_flux_1() \ + or self.is_sana() \ + or self.is_hunyuan_video() \ + or self.is_hi_dream() \ + or self.is_chroma(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING) + if self.is_qwen() or self.is_z_image() or self.is_flux_2() or self.is_ernie(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA) + raise ValueError(f"No supported training methods defined for model type {self}") + + +# The first text encoder is always "text_encoder" here (matching the config field), even for +# multi-encoder models that refer to it as "text_encoder_1" elsewhere in the code. +_MODEL_PARTS: dict[ModelType, tuple[str, ...]] = { + ModelType.STABLE_DIFFUSION_15: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_15_INPAINTING: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_BASE: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_INPAINTING: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_DEPTH: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_21: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_21_BASE: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_3: ("text_encoder", "text_encoder_2", "text_encoder_3", "transformer", "vae"), + ModelType.STABLE_DIFFUSION_35: ("text_encoder", "text_encoder_2", "text_encoder_3", "transformer", "vae"), + ModelType.STABLE_DIFFUSION_XL_10_BASE: ("text_encoder", "text_encoder_2", "unet", "vae"), + ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING: ("text_encoder", "text_encoder_2", "unet", "vae"), + # Only Würstchen v2's decoder has its own text encoder; Stable Cascade's decoder does not. + ModelType.WUERSTCHEN_2: ("text_encoder", "prior", "effnet_encoder", "decoder", "decoder_text_encoder", "decoder_vqgan"), + ModelType.STABLE_CASCADE_1: ("text_encoder", "prior", "effnet_encoder", "decoder", "decoder_vqgan"), + ModelType.PIXART_ALPHA: ("text_encoder", "transformer", "vae"), + ModelType.PIXART_SIGMA: ("text_encoder", "transformer", "vae"), + ModelType.FLUX_DEV_1: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.FLUX_FILL_DEV_1: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.FLUX_2: ("text_encoder", "transformer", "vae"), + ModelType.SANA: ("text_encoder", "transformer", "vae"), + ModelType.HUNYUAN_VIDEO: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.HI_DREAM_FULL: ("text_encoder", "text_encoder_2", "text_encoder_3", "text_encoder_4", "transformer", "vae"), + ModelType.CHROMA_1: ("text_encoder", "transformer", "vae"), + ModelType.QWEN: ("text_encoder", "transformer", "vae"), + ModelType.Z_IMAGE: ("text_encoder", "transformer", "vae"), + ModelType.ERNIE: ("text_encoder", "transformer", "vae"), +} + class PeftType(Enum): LORA = 'LORA' diff --git a/training_presets/#chroma Finetune 16GB.json b/training_presets/#chroma Finetune 16GB.json index 2dacbee20..145ab5b04 100644 --- a/training_presets/#chroma Finetune 16GB.json +++ b/training_presets/#chroma Finetune 16GB.json @@ -4,16 +4,15 @@ "learning_rate": 1e-5, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.4, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.4 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { diff --git a/training_presets/#chroma Finetune 8GB.json b/training_presets/#chroma Finetune 8GB.json index 508410995..29b36c84b 100644 --- a/training_presets/#chroma Finetune 8GB.json +++ b/training_presets/#chroma Finetune 8GB.json @@ -4,16 +4,15 @@ "learning_rate": 1e-5, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.85, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.85 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { diff --git a/training_presets/#chroma LoRA 8GB.json b/training_presets/#chroma LoRA 8GB.json index 78027aac4..437ac71d7 100644 --- a/training_presets/#chroma LoRA 8GB.json +++ b/training_presets/#chroma LoRA 8GB.json @@ -4,16 +4,15 @@ "learning_rate": 0.0003, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "LORA", "vae": { diff --git a/training_presets/#ernie LoRA 8GB.json b/training_presets/#ernie LoRA 8GB.json index bc1e4a82b..a63d34ece 100644 --- a/training_presets/#ernie LoRA 8GB.json +++ b/training_presets/#ernie LoRA 8GB.json @@ -7,7 +7,8 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "INT_W8A8" + "weight_dtype": "INT_W8A8", + "offload_fraction": 0.7 }, "text_encoder": { "train": false, @@ -26,7 +27,5 @@ "layer_filter_preset": "blocks" }, "timestep_distribution": "LOGIT_NORMAL", - "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.7 + "dataloader_threads": 1 } diff --git a/training_presets/#flux2 Finetune 16GB.json b/training_presets/#flux2 Finetune 16GB.json index 6e3addbe8..55feaf992 100644 --- a/training_presets/#flux2 Finetune 16GB.json +++ b/training_presets/#flux2 Finetune 16GB.json @@ -7,11 +7,12 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { @@ -25,8 +26,6 @@ }, "timestep_distribution": "LOGIT_NORMAL", "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "optimizer": { "optimizer": "ADAFACTOR" }, diff --git a/training_presets/#flux2 LoRA 8GB.json b/training_presets/#flux2 LoRA 8GB.json index dfdc77839..dd3992123 100644 --- a/training_presets/#flux2 LoRA 8GB.json +++ b/training_presets/#flux2 LoRA 8GB.json @@ -7,11 +7,13 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "INT_W8A8" + "weight_dtype": "INT_W8A8", + "offload_fraction": 0.7 }, "text_encoder": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.7 }, "training_method": "LORA", "vae": { @@ -26,7 +28,5 @@ "layer_filter_preset": "blocks" }, "timestep_distribution": "LOGIT_NORMAL", - "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.7 + "dataloader_threads": 1 } diff --git a/training_presets/#hidream LoRA.json b/training_presets/#hidream LoRA.json index 2eb588763..8b025a628 100644 --- a/training_presets/#hidream LoRA.json +++ b/training_presets/#hidream LoRA.json @@ -2,8 +2,6 @@ "backup_after": 10, "base_model_name": "HiDream-ai/HiDream-I1-Full", "batch_size": 4, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "learning_rate": 0.0003, "model_type": "HI_DREAM_FULL", @@ -16,7 +14,8 @@ "training_method": "LORA", "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, @@ -28,11 +27,13 @@ }, "text_encoder_3": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder_4": { "model_name": "meta-llama/Llama-3.1-8B-Instruct", "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 } } diff --git a/training_presets/#hunyuan video LoRA.json b/training_presets/#hunyuan video LoRA.json index 754550155..96bbb80e7 100644 --- a/training_presets/#hunyuan video LoRA.json +++ b/training_presets/#hunyuan video LoRA.json @@ -2,8 +2,6 @@ "backup_after": 10, "base_model_name": "hunyuanvideo-community/HunyuanVideo", "batch_size": 4, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "learning_rate": 0.0003, "model_type": "HUNYUAN_VIDEO", @@ -16,11 +14,13 @@ "training_method": "LORA", "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder_2": { "train": false, diff --git a/training_presets/#qwen Finetune 16GB.json b/training_presets/#qwen Finetune 16GB.json index 811d7e0b1..2224993f4 100644 --- a/training_presets/#qwen Finetune 16GB.json +++ b/training_presets/#qwen Finetune 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.75, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.75 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen Finetune 24GB.json b/training_presets/#qwen Finetune 24GB.json index 8bee3cd3f..1cf9dc09f 100644 --- a/training_presets/#qwen Finetune 24GB.json +++ b/training_presets/#qwen Finetune 24GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.55, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.55 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen LoRA 16GB.json b/training_presets/#qwen LoRA 16GB.json index b4e0d7e88..0eda34d5d 100644 --- a/training_presets/#qwen LoRA 16GB.json +++ b/training_presets/#qwen LoRA 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen LoRA 24GB.json b/training_presets/#qwen LoRA 24GB.json index 696648a42..cd7b7216e 100644 --- a/training_presets/#qwen LoRA 24GB.json +++ b/training_presets/#qwen LoRA 24GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.1, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.1 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image DeTurbo LoRA 8GB.json b/training_presets/#z-image DeTurbo LoRA 8GB.json index cc38e60eb..957f08b80 100644 --- a/training_presets/#z-image DeTurbo LoRA 8GB.json +++ b/training_presets/#z-image DeTurbo LoRA 8GB.json @@ -4,13 +4,12 @@ "learning_rate": 0.0003, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "compile": true, "transformer": { "train": true, "weight_dtype": "INT_W8A8", - "model_name": "https://huggingface.co/ostris/Z-Image-De-Turbo/blob/main/z_image_de_turbo_v1_bf16.safetensors" + "model_name": "https://huggingface.co/ostris/Z-Image-De-Turbo/blob/main/z_image_de_turbo_v1_bf16.safetensors", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image Finetune 16GB.json b/training_presets/#z-image Finetune 16GB.json index 0d23d3992..3911c866c 100644 --- a/training_presets/#z-image Finetune 16GB.json +++ b/training_presets/#z-image Finetune 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.1, "compile": true, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.1 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image LoRA 8GB.json b/training_presets/#z-image LoRA 8GB.json index 78b4b05cc..8dbce7b21 100644 --- a/training_presets/#z-image LoRA 8GB.json +++ b/training_presets/#z-image LoRA 8GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "compile": true, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, From 2f0620be0bb367abbd60225224bebfdda112719d Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 5 Jun 2026 21:00:59 +0200 Subject: [PATCH 02/10] =?UTF-8?q?Revert=20"Revert=20"Upgrade=20transformer?= =?UTF-8?q?s=20to=205.9=20and=20huggingface-hub=20to=201.16=20(#147?= =?UTF-8?q?=E2=80=A6"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 574ec55ab9a174267768f1222489103666276820. --- modules/dataLoader/ErnieBaseDataLoader.py | 3 -- modules/dataLoader/Flux2BaseDataLoader.py | 5 --- modules/dataLoader/ZImageBaseDataLoader.py | 3 -- modules/model/ChromaModel.py | 6 ++- modules/model/FluxModel.py | 10 +++-- modules/model/HiDreamModel.py | 10 ++--- modules/model/HunyuanVideoModel.py | 6 +-- modules/model/PixArtAlphaModel.py | 9 ++-- modules/model/SanaModel.py | 6 ++- modules/model/StableDiffusion3Model.py | 14 ++++-- modules/model/StableDiffusionModel.py | 11 +++-- modules/model/StableDiffusionXLModel.py | 10 +++-- modules/model/WuerstchenModel.py | 11 +++-- modules/modelLoader/ErnieModelLoader.py | 28 +----------- modules/modelLoader/Flux2ModelLoader.py | 13 ------ modules/modelLoader/ZImageModelLoader.py | 13 ------ .../modelLoader/chroma/ChromaModelLoader.py | 14 +----- modules/modelLoader/flux/FluxModelLoader.py | 22 +++------- .../modelLoader/hiDream/HiDreamModelLoader.py | 28 ------------ .../hunyuanVideo/HunyuanVideoModelLoader.py | 18 -------- .../modelLoader/mixin/HFModelLoaderMixin.py | 24 ++--------- .../pixartAlpha/PixArtAlphaModelLoader.py | 2 + modules/modelLoader/qwen/QwenModelLoader.py | 12 ------ modules/modelLoader/sana/SanaModelLoader.py | 2 + .../StableDiffusionModelLoader.py | 4 ++ .../StableDiffusion3ModelLoader.py | 9 +++- .../StableDiffusionXLModelLoader.py | 7 +++ .../wuerstchen/WuerstchenModelLoader.py | 2 + modules/modelSampler/HiDreamSampler.py | 2 +- modules/modelSampler/HunyuanVideoSampler.py | 2 +- modules/modelSaver/chroma/ChromaModelSaver.py | 2 +- modules/modelSaver/flux/FluxModelSaver.py | 2 +- .../modelSaver/hidream/HiDreamModelSaver.py | 2 +- .../hunyuanVideo/HunyuanVideoModelSaver.py | 2 +- .../pixartAlpha/PixArtAlphaModelSaver.py | 2 +- modules/modelSaver/sana/SanaModelSaver.py | 2 +- .../StableDiffusionModelSaver.py | 2 +- .../StableDiffusion3ModelSaver.py | 2 +- .../StableDiffusionXLModelSaver.py | 2 +- .../wuerstchen/WuerstchenModelSaver.py | 2 +- modules/modelSetup/ChromaEmbeddingSetup.py | 1 - modules/modelSetup/ChromaFineTuneSetup.py | 1 - modules/modelSetup/ChromaLoRASetup.py | 1 - modules/modelSetup/FluxEmbeddingSetup.py | 2 - modules/modelSetup/FluxFineTuneSetup.py | 2 - modules/modelSetup/FluxLoRASetup.py | 2 - .../modelSetup/PixArtAlphaEmbeddingSetup.py | 1 - .../modelSetup/PixArtAlphaFineTuneSetup.py | 1 - modules/modelSetup/PixArtAlphaLoRASetup.py | 1 - modules/modelSetup/SanaEmbeddingSetup.py | 1 - modules/modelSetup/SanaFineTuneSetup.py | 1 - modules/modelSetup/SanaLoRASetup.py | 1 - .../StableDiffusion3EmbeddingSetup.py | 3 -- .../StableDiffusion3FineTuneSetup.py | 3 -- .../modelSetup/StableDiffusion3LoRASetup.py | 3 -- .../StableDiffusionEmbeddingSetup.py | 1 - .../StableDiffusionFineTuneSetup.py | 1 - .../modelSetup/StableDiffusionLoRASetup.py | 1 - .../StableDiffusionXLEmbeddingSetup.py | 2 - .../StableDiffusionXLFineTuneSetup.py | 2 - .../modelSetup/StableDiffusionXLLoRASetup.py | 2 - .../modelSetup/WuerstchenEmbeddingSetup.py | 1 - modules/modelSetup/WuerstchenFineTuneSetup.py | 1 - modules/modelSetup/WuerstchenLoRASetup.py | 1 - .../mixin/ModelSetupEmbeddingMixin.py | 14 +----- modules/trainer/GenericTrainer.py | 5 +-- modules/ui/ConceptWindow.py | 3 +- modules/util/thread_safety.py | 43 ------------------- requirements-global.txt | 4 +- run-cmd.sh | 5 --- start-ui.bat | 8 ---- start-ui.sh | 5 --- 72 files changed, 115 insertions(+), 334 deletions(-) delete mode 100644 modules/util/thread_safety.py diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index d26568cb5..edd634e79 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -7,7 +7,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -32,8 +31,6 @@ def _preparation_modules(self, config: TrainConfig, model: ErnieModel): image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125) tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673, unclear if Mistral is affected encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=HIDDEN_STATES_LAYER, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index 587f480c5..0aa4522d7 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -14,7 +14,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -43,8 +42,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: mistral_format_input([caption], MISTRAL_SYSTEM_MESSAGE), apply_chat_template_kwargs = {'add_generation_prompt': False}, ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), hidden_state_output_index=MISTRAL_HIDDEN_STATES_LAYERS, @@ -53,8 +50,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length, apply_chat_template = lambda caption: qwen3_format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': False} ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=QWEN3_HIDDEN_STATES_LAYERS, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 03934dd4d..66cfea16f 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -9,7 +9,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.thread_safety import apply_thread_safe_forward from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -38,8 +37,6 @@ def _preparation_modules(self, config: TrainConfig, model: ZImageModel): tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH, apply_chat_template = lambda caption: format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': True} ) - if config.dataloader_threads > 1: - apply_thread_safe_forward(model.text_encoder) # workaround for transformers#42673 encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', text_encoder=model.text_encoder, hidden_state_output_index=-2, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state') diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index 19146bd2b..59967c8dc 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -40,6 +40,7 @@ def __init__( class ChromaModel(BaseModel): # base model data tokenizer: T5Tokenizer | None + orig_tokenizer: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder: T5EncoderModel | None vae: AutoencoderKL | None @@ -72,6 +73,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -141,13 +143,13 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return ChromaPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer, ) def add_text_encoder_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index e4bfaa066..b981865c4 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -50,7 +50,9 @@ def __init__( class FluxModel(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: T5Tokenizer | None + orig_tokenizer_2: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder_1: CLIPTextModel | None text_encoder_2: T5EncoderModel | None @@ -86,7 +88,9 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -177,15 +181,15 @@ def eval(self): self.text_encoder_2.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return FluxPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index 121f5fd7c..049b0e6d3 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -267,19 +267,19 @@ def eval(self): self.text_encoder_4.eval() self.transformer.eval() - def create_pipeline(self, use_original_modules: bool) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HiDreamImagePipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.orig_tokenizer_1 if use_original_modules else self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.orig_tokenizer_2 if use_original_modules else self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, text_encoder_3=self.text_encoder_3, - tokenizer_3=self.orig_tokenizer_3 if use_original_modules else self.tokenizer_3, + tokenizer_3=self.orig_tokenizer_3 if use_original_tokenizers else self.tokenizer_3, text_encoder_4=self.text_encoder_4, - tokenizer_4=self.orig_tokenizer_4 if use_original_modules else self.tokenizer_4, + tokenizer_4=self.orig_tokenizer_4 if use_original_tokenizers else self.tokenizer_4, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 49c40cc15..2107e55d7 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -196,15 +196,15 @@ def eval(self): self.text_encoder_2.eval() self.transformer.eval() - def create_pipeline(self, use_original_modules: bool) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HunyuanVideoPipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.orig_tokenizer_1 if use_original_modules else self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.orig_tokenizer_2 if use_original_modules else self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index 97f87e5ca..466cc61f9 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -42,6 +42,7 @@ def __init__( class PixArtAlphaModel(BaseModel): # base model data tokenizer: T5Tokenizer | None + orig_tokenizer: T5Tokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: T5EncoderModel | None vae: AutoencoderKL | None @@ -74,6 +75,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -141,11 +143,12 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + tokenizer = self.orig_tokenizer if use_original_tokenizers else self.tokenizer match self.model_type: case ModelType.PIXART_ALPHA: return PixArtAlphaPipeline( - tokenizer=self.tokenizer, + tokenizer=tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, @@ -153,7 +156,7 @@ def create_pipeline(self) -> DiffusionPipeline: ) case ModelType.PIXART_SIGMA: return PixArtSigmaPipeline( - tokenizer=self.tokenizer, + tokenizer=tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index 506afe741..9e8008219 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -41,6 +41,7 @@ def __init__( class SanaModel(BaseModel): # base model data tokenizer: GemmaTokenizer | None + orig_tokenizer: GemmaTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: Gemma2Model | None vae: AutoencoderDC | None @@ -74,6 +75,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -143,9 +145,9 @@ def eval(self): self.text_encoder.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return SanaPipeline( - tokenizer=self.tokenizer, + tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer, text_encoder=self.text_encoder, vae=self.vae, transformer=self.transformer, diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index bcaf46fcb..8f6cf5818 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -58,8 +58,11 @@ def __init__( class StableDiffusion3Model(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: CLIPTokenizer | None + orig_tokenizer_2: CLIPTokenizer | None tokenizer_3: T5Tokenizer | None + orig_tokenizer_3: T5Tokenizer | None noise_scheduler: FlowMatchEulerDiscreteScheduler | None text_encoder_1: CLIPTextModelWithProjection | None text_encoder_2: CLIPTextModelWithProjection | None @@ -98,8 +101,11 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.tokenizer_3 = None + self.orig_tokenizer_3 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -208,17 +214,17 @@ def eval(self): self.text_encoder_3.eval() self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusion3Pipeline( transformer=self.transformer, scheduler=self.noise_scheduler, vae=self.vae, text_encoder=self.text_encoder_1, - tokenizer=self.tokenizer_1, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, text_encoder_2=self.text_encoder_2, - tokenizer_2=self.tokenizer_2, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, text_encoder_3=self.text_encoder_3, - tokenizer_3=self.tokenizer_3, + tokenizer_3=self.orig_tokenizer_3 if use_original_tokenizers else self.tokenizer_3, ) def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: diff --git a/modules/model/StableDiffusionModel.py b/modules/model/StableDiffusionModel.py index 2bee0fd70..4b1b11a2c 100644 --- a/modules/model/StableDiffusionModel.py +++ b/modules/model/StableDiffusionModel.py @@ -43,6 +43,7 @@ def __init__( class StableDiffusionModel(BaseModel): # base model data tokenizer: CLIPTokenizer | None + orig_tokenizer: CLIPTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder: CLIPTextModel | None vae: AutoencoderKL | None @@ -72,6 +73,7 @@ def __init__( ) self.tokenizer = None + self.orig_tokenizer = None self.noise_scheduler = None self.text_encoder = None self.vae = None @@ -136,12 +138,13 @@ def eval(self): self.text_encoder.eval() self.unet.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + tokenizer = self.orig_tokenizer if use_original_tokenizers else self.tokenizer if self.model_type.has_depth_input(): return StableDiffusionDepth2ImgPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, depth_estimator=self.depth_estimator, @@ -151,7 +154,7 @@ def create_pipeline(self) -> DiffusionPipeline: return StableDiffusionInpaintPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, safety_checker=None, @@ -162,7 +165,7 @@ def create_pipeline(self) -> DiffusionPipeline: return StableDiffusionPipeline( vae=self.vae, text_encoder=self.text_encoder, - tokenizer=self.tokenizer, + tokenizer=tokenizer, unet=self.unet, scheduler=self.noise_scheduler, safety_checker=None, diff --git a/modules/model/StableDiffusionXLModel.py b/modules/model/StableDiffusionXLModel.py index b68df2d23..43edead8e 100644 --- a/modules/model/StableDiffusionXLModel.py +++ b/modules/model/StableDiffusionXLModel.py @@ -45,7 +45,9 @@ def __init__( class StableDiffusionXLModel(BaseModel): # base model data tokenizer_1: CLIPTokenizer | None + orig_tokenizer_1: CLIPTokenizer | None tokenizer_2: CLIPTokenizer | None + orig_tokenizer_2: CLIPTokenizer | None noise_scheduler: DDIMScheduler | None text_encoder_1: CLIPTextModel | None text_encoder_2: CLIPTextModelWithProjection | None @@ -81,7 +83,9 @@ def __init__( ) self.tokenizer_1 = None + self.orig_tokenizer_1 = None self.tokenizer_2 = None + self.orig_tokenizer_2 = None self.noise_scheduler = None self.text_encoder_1 = None self.text_encoder_2 = None @@ -166,13 +170,13 @@ def eval(self): self.text_encoder_2.eval() self.unet.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusionXLPipeline( vae=self.vae, text_encoder=self.text_encoder_1, text_encoder_2=self.text_encoder_2, - tokenizer=self.tokenizer_1, - tokenizer_2=self.tokenizer_2, + tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1, + tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2, unet=self.unet, scheduler=self.noise_scheduler, ) diff --git a/modules/model/WuerstchenModel.py b/modules/model/WuerstchenModel.py index b844d6eb1..7b50d81ed 100644 --- a/modules/model/WuerstchenModel.py +++ b/modules/model/WuerstchenModel.py @@ -69,6 +69,7 @@ class WuerstchenModel(BaseModel): decoder_vqgan: PaellaVQModel | None effnet_encoder: WuerstchenEfficientNetEncoder | None prior_tokenizer: CLIPTokenizer | None + orig_prior_tokenizer: CLIPTokenizer | None prior_text_encoder: CLIPTextModel | None prior_noise_scheduler: DDPMWuerstchenScheduler | None prior_prior: WuerstchenPrior | StableCascadeUNet | None @@ -105,6 +106,7 @@ def __init__( self.decoder_vqgan = None self.effnet_encoder = None self.prior_tokenizer = None + self.orig_prior_tokenizer = None self.prior_text_encoder = None self.prior_noise_scheduler = None self.prior_prior = None @@ -179,7 +181,8 @@ def eval(self): self.prior_text_encoder.eval() self.prior_prior.eval() - def create_pipeline(self) -> DiffusionPipeline: + def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: + prior_tokenizer = self.orig_prior_tokenizer if use_original_tokenizers else self.prior_tokenizer if self.model_type.is_wuerstchen_v2(): return WuerstchenCombinedPipeline( tokenizer=self.decoder_tokenizer, @@ -187,19 +190,19 @@ def create_pipeline(self) -> DiffusionPipeline: decoder=self.decoder_decoder, scheduler=self.decoder_noise_scheduler, vqgan=self.decoder_vqgan, - prior_tokenizer=self.prior_tokenizer, + prior_tokenizer=prior_tokenizer, prior_text_encoder=self.prior_text_encoder, prior_prior=self.prior_prior, prior_scheduler=self.prior_noise_scheduler, ) elif self.model_type.is_stable_cascade(): return StableCascadeCombinedPipeline( - tokenizer=self.prior_tokenizer, + tokenizer=prior_tokenizer, text_encoder=self.prior_text_encoder, decoder=self.decoder_decoder, scheduler=self.decoder_noise_scheduler, vqgan=self.decoder_vqgan, - prior_tokenizer=self.prior_tokenizer, + prior_tokenizer=prior_tokenizer, prior_text_encoder=self.prior_text_encoder, prior_prior=self.prior_prior, prior_scheduler=self.prior_noise_scheduler, diff --git a/modules/modelLoader/ErnieModelLoader.py b/modules/modelLoader/ErnieModelLoader.py index 8e8365981..c9d623463 100644 --- a/modules/modelLoader/ErnieModelLoader.py +++ b/modules/modelLoader/ErnieModelLoader.py @@ -1,4 +1,3 @@ -import logging import os import traceback @@ -22,8 +21,7 @@ FlowMatchEulerDiscreteScheduler, GGUFQuantizationConfig, ) -from transformers import Mistral3Model, MistralConfig, PreTrainedTokenizerFast -from transformers.models.auto.configuration_auto import CONFIG_MAPPING +from transformers import AutoTokenizer, Mistral3Model class ErnieModelLoader( @@ -60,23 +58,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - # transformers < 5.x doesn't register "ministral3"; patch it so Mistral3Config can parse its text_config - if "ministral3" not in CONFIG_MAPPING: - CONFIG_MAPPING.register("ministral3", MistralConfig) - - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if transformer_model_name: transformer = ErnieImageTransformer2DModel.from_single_file( transformer_model_name, @@ -98,15 +79,10 @@ def __load_diffusers( quantization, ) - # TokenizersBackend is the Rust tokenizers library backend, not a transformers class — warning is a false alarm - tokenization_logger = logging.getLogger("transformers.tokenization_utils_base") - prev_level = tokenization_logger.level - tokenization_logger.setLevel(logging.ERROR) - tokenizer = PreTrainedTokenizerFast.from_pretrained( + tokenizer = AutoTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - tokenization_logger.setLevel(prev_level) text_encoder = self._load_transformers_sub_module( Mistral3Model, diff --git a/modules/modelLoader/Flux2ModelLoader.py b/modules/modelLoader/Flux2ModelLoader.py index 414e43232..670bacab3 100644 --- a/modules/modelLoader/Flux2ModelLoader.py +++ b/modules/modelLoader/Flux2ModelLoader.py @@ -62,19 +62,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if transformer_model_name: transformer = Flux2Transformer2DModel.from_single_file( transformer_model_name, diff --git a/modules/modelLoader/ZImageModelLoader.py b/modules/modelLoader/ZImageModelLoader.py index 0fe6e3c00..104d04e7b 100644 --- a/modules/modelLoader/ZImageModelLoader.py +++ b/modules/modelLoader/ZImageModelLoader.py @@ -60,19 +60,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = ["text_encoder"] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/chroma/ChromaModelLoader.py b/modules/modelLoader/chroma/ChromaModelLoader.py index 9c5a5d543..7dcbef794 100644 --- a/modules/modelLoader/chroma/ChromaModelLoader.py +++ b/modules/modelLoader/chroma/ChromaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -52,18 +53,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=["text_encoder"], - ) - tokenizer = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", @@ -120,6 +109,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/flux/FluxModelLoader.py b/modules/modelLoader/flux/FluxModelLoader.py index 7747a9f3e..d4f21ea2a 100644 --- a/modules/modelLoader/flux/FluxModelLoader.py +++ b/modules/modelLoader/flux/FluxModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -58,23 +59,6 @@ def __load_diffusers( include_text_encoder_2: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if include_text_encoder_1: tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, @@ -156,7 +140,9 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -231,7 +217,9 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = pipeline.scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/hiDream/HiDreamModelLoader.py b/modules/modelLoader/hiDream/HiDreamModelLoader.py index 639156a38..b3e20f23c 100644 --- a/modules/modelLoader/hiDream/HiDreamModelLoader.py +++ b/modules/modelLoader/hiDream/HiDreamModelLoader.py @@ -67,34 +67,6 @@ def __load_diffusers( include_text_encoder_4: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if include_text_encoder_3: - transformers_sub.append("text_encoder_3") - if include_text_encoder_4: - if text_encoder_4_model_name: - self._prepare_sub_modules( - text_encoder_4_model_name, - transformers_modules=[""], - diffusers_modules=[], - ) - else: - transformers_sub.append("text_encoder_4") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py index 0e9d577cc..85c91699b 100644 --- a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py +++ b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py @@ -59,24 +59,6 @@ def __load_diffusers( include_text_encoder_2: bool, quantization: QuantizationConfig, ): - diffusers_sub = [] - transformers_sub = [] - - if not transformer_model_name: - diffusers_sub.append("transformer") - if include_text_encoder_1: - transformers_sub.append("text_encoder") - if include_text_encoder_2: - transformers_sub.append("text_encoder_2") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=transformers_sub, - ) - if include_text_encoder_1: tokenizer_1 = LlamaTokenizerFast.from_pretrained( base_model_name, diff --git a/modules/modelLoader/mixin/HFModelLoaderMixin.py b/modules/modelLoader/mixin/HFModelLoaderMixin.py index 9180014bb..eae11cd3b 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -1,7 +1,7 @@ import json +import logging import os import re -import traceback from abc import ABCMeta from itertools import repeat @@ -20,6 +20,9 @@ from huggingface_hub.utils import EntryNotFoundError from safetensors.torch import load_file +# huggingface_hub 1.16+ uses httpx, which logs every HTTP request/response at INFO level. +logging.getLogger("httpx").setLevel(logging.WARNING) + class HFModelLoaderMixin(metaclass=ABCMeta): def __init__(self): @@ -298,22 +301,3 @@ def _convert_diffusers_sub_module_to_dtype( None, quantization, ) - - def _prepare_sub_modules(self, pretrained_model_name_or_path: str, diffusers_modules: list[str], transformers_modules: list[str]): - is_local = os.path.isdir(pretrained_model_name_or_path) - if is_local: - return - - diffusers_paths = [((folder + "/") if folder else "") + "diffusion_pytorch_model*" for folder in diffusers_modules] - transformers_paths = [((folder + "/") if folder else "") + "model*" for folder in transformers_modules] - transformers_paths.extend([((folder + "/") if folder else "") + "pytorch_model*" for folder in transformers_modules]) - try: - huggingface_hub.snapshot_download( - pretrained_model_name_or_path, - allow_patterns=diffusers_paths + transformers_paths, - ) - except huggingface_hub.errors.HFValidationError: - pass - except Exception: - traceback.print_exc() - print("Error during bulk preloading of Huggingface model repository, proceeding without preloading") diff --git a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py index f8e75907d..467c29c8a 100644 --- a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py +++ b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -86,6 +87,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/qwen/QwenModelLoader.py b/modules/modelLoader/qwen/QwenModelLoader.py index 498b925af..953f15bfb 100644 --- a/modules/modelLoader/qwen/QwenModelLoader.py +++ b/modules/modelLoader/qwen/QwenModelLoader.py @@ -52,18 +52,6 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - diffusers_sub = [] - if not transformer_model_name: - diffusers_sub.append("transformer") - if not vae_model_name: - diffusers_sub.append("vae") - - self._prepare_sub_modules( - base_model_name, - diffusers_modules=diffusers_sub, - transformers_modules=["text_encoder"], - ) - tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", diff --git a/modules/modelLoader/sana/SanaModelLoader.py b/modules/modelLoader/sana/SanaModelLoader.py index 508579ba2..a904e3996 100644 --- a/modules/modelLoader/sana/SanaModelLoader.py +++ b/modules/modelLoader/sana/SanaModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -86,6 +87,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py index 4333897ae..aa610f485 100644 --- a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py +++ b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -132,6 +133,7 @@ def __load_diffusers( model.model_type = model_type model.tokenizer = tokenizer + model.orig_tokenizer = copy.deepcopy(tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae @@ -206,6 +208,7 @@ def __load_ckpt( model.model_type = model_type model.tokenizer = pipeline.tokenizer + model.orig_tokenizer = copy.deepcopy(pipeline.tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae @@ -262,6 +265,7 @@ def __load_safetensors( model.model_type = model_type model.tokenizer = pipeline.tokenizer + model.orig_tokenizer = copy.deepcopy(pipeline.tokenizer) model.noise_scheduler = noise_scheduler model.text_encoder = text_encoder model.vae = vae diff --git a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py index 17e55696d..47d87da74 100644 --- a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py +++ b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -50,8 +51,6 @@ def __load_diffusers( include_text_encoder_3: bool, quantization: QuantizationConfig, ): - #no call to self._prepare_sub_modules, because SAI polluted their sd3 / sd3.5 medium repo text encoders with fp16 files - if include_text_encoder_1: tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, @@ -141,8 +140,11 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.tokenizer_3 = tokenizer_3 + model.orig_tokenizer_3 = copy.deepcopy(tokenizer_3) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -230,8 +232,11 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.tokenizer_3 = tokenizer_3 + model.orig_tokenizer_3 = copy.deepcopy(tokenizer_3) model.noise_scheduler = pipeline.scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py index 6d31516fe..afbab6581 100644 --- a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py +++ b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py @@ -1,3 +1,4 @@ +import copy import os import traceback @@ -125,7 +126,9 @@ def __load_diffusers( model.model_type = model_type model.tokenizer_1 = tokenizer_1 + model.orig_tokenizer_1 = copy.deepcopy(tokenizer_1) model.tokenizer_2 = tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -167,7 +170,9 @@ def __load_ckpt( model.model_type = model_type model.tokenizer_1 = pipeline.tokenizer + model.orig_tokenizer_1 = copy.deepcopy(pipeline.tokenizer) model.tokenizer_2 = pipeline.tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(pipeline.tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 @@ -227,7 +232,9 @@ def __load_safetensors( model.model_type = model_type model.tokenizer_1 = pipeline.tokenizer + model.orig_tokenizer_1 = copy.deepcopy(pipeline.tokenizer) model.tokenizer_2 = pipeline.tokenizer_2 + model.orig_tokenizer_2 = copy.deepcopy(pipeline.tokenizer_2) model.noise_scheduler = noise_scheduler model.text_encoder_1 = text_encoder_1 model.text_encoder_2 = text_encoder_2 diff --git a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py index 8fec87b29..188107a2c 100644 --- a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py +++ b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py @@ -1,3 +1,4 @@ +import copy import json import os.path import traceback @@ -192,6 +193,7 @@ def __load_diffusers( model.decoder_vqgan = decoder_vqgan model.effnet_encoder = effnet_encoder model.prior_tokenizer = prior_tokenizer + model.orig_prior_tokenizer = copy.deepcopy(prior_tokenizer) model.prior_text_encoder = prior_text_encoder model.prior_noise_scheduler = prior_noise_scheduler model.prior_prior = prior_prior diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py index d02197a35..dbafdae81 100644 --- a/modules/modelSampler/HiDreamSampler.py +++ b/modules/modelSampler/HiDreamSampler.py @@ -31,7 +31,7 @@ def __init__( self.model = model self.model_type = model_type - self.pipeline = model.create_pipeline(use_original_modules=False) + self.pipeline = model.create_pipeline() @torch.no_grad() def __sample_base( diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py index 275626341..020df6a39 100644 --- a/modules/modelSampler/HunyuanVideoSampler.py +++ b/modules/modelSampler/HunyuanVideoSampler.py @@ -32,7 +32,7 @@ def __init__( self.model = model self.model_type = model_type - self.pipeline = model.create_pipeline(use_original_modules=False) + self.pipeline = model.create_pipeline() @torch.no_grad() def __sample_base( diff --git a/modules/modelSaver/chroma/ChromaModelSaver.py b/modules/modelSaver/chroma/ChromaModelSaver.py index f8129558d..2dee69bfc 100644 --- a/modules/modelSaver/chroma/ChromaModelSaver.py +++ b/modules/modelSaver/chroma/ChromaModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/flux/FluxModelSaver.py b/modules/modelSaver/flux/FluxModelSaver.py index dc632f2ae..b5ac966d5 100644 --- a/modules/modelSaver/flux/FluxModelSaver.py +++ b/modules/modelSaver/flux/FluxModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/hidream/HiDreamModelSaver.py b/modules/modelSaver/hidream/HiDreamModelSaver.py index 4bfe56ebc..b243166ec 100644 --- a/modules/modelSaver/hidream/HiDreamModelSaver.py +++ b/modules/modelSaver/hidream/HiDreamModelSaver.py @@ -24,7 +24,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline(use_original_modules=True) + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py index 866d68617..3fad1c7a0 100644 --- a/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py +++ b/modules/modelSaver/hunyuanVideo/HunyuanVideoModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline(use_original_modules=True) + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py index 83e8af4a9..d6ef4d580 100644 --- a/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py +++ b/modules/modelSaver/pixartAlpha/PixArtAlphaModelSaver.py @@ -25,7 +25,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/sana/SanaModelSaver.py b/modules/modelSaver/sana/SanaModelSaver.py index 1dede8c8f..29301adbb 100644 --- a/modules/modelSaver/sana/SanaModelSaver.py +++ b/modules/modelSaver/sana/SanaModelSaver.py @@ -22,7 +22,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py b/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py index e0131f52f..f24e6b5f9 100644 --- a/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py +++ b/modules/modelSaver/stableDiffusion/StableDiffusionModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py index bcfabab67..095b9e75f 100644 --- a/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py +++ b/modules/modelSaver/stableDiffusion3/StableDiffusion3ModelSaver.py @@ -27,7 +27,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. diff --git a/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py b/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py index 09170e67b..24625fc52 100644 --- a/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py +++ b/modules/modelSaver/stableDiffusionXL/StableDiffusionXLModelSaver.py @@ -26,7 +26,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline() + pipeline = model.create_pipeline(use_original_tokenizers=True) pipeline.to("cpu") if dtype is not None: diff --git a/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py b/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py index 3bbe1b55d..cee5e28b4 100644 --- a/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py +++ b/modules/modelSaver/wuerstchen/WuerstchenModelSaver.py @@ -25,7 +25,7 @@ def __save_diffusers( dtype: torch.dtype | None, ): # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. - pipeline = model.create_pipeline().prior_pipe + pipeline = model.create_pipeline(use_original_tokenizers=True).prior_pipe original_device = pipeline.device pipeline.to("cpu") pipeline_copy = copy.deepcopy(pipeline) diff --git a/modules/modelSetup/ChromaEmbeddingSetup.py b/modules/modelSetup/ChromaEmbeddingSetup.py index 9fa65b69f..63a6d75af 100644 --- a/modules/modelSetup/ChromaEmbeddingSetup.py +++ b/modules/modelSetup/ChromaEmbeddingSetup.py @@ -61,7 +61,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/ChromaFineTuneSetup.py b/modules/modelSetup/ChromaFineTuneSetup.py index 371c414ec..c78e4e700 100644 --- a/modules/modelSetup/ChromaFineTuneSetup.py +++ b/modules/modelSetup/ChromaFineTuneSetup.py @@ -70,7 +70,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/ChromaLoRASetup.py b/modules/modelSetup/ChromaLoRASetup.py index 1396b4a52..514eca8e8 100644 --- a/modules/modelSetup/ChromaLoRASetup.py +++ b/modules/modelSetup/ChromaLoRASetup.py @@ -98,7 +98,6 @@ def setup_model( if model.text_encoder is not None: model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxEmbeddingSetup.py b/modules/modelSetup/FluxEmbeddingSetup.py index 0418390da..98dd12579 100644 --- a/modules/modelSetup/FluxEmbeddingSetup.py +++ b/modules/modelSetup/FluxEmbeddingSetup.py @@ -71,8 +71,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxFineTuneSetup.py b/modules/modelSetup/FluxFineTuneSetup.py index a2dc78057..3afbdf31c 100644 --- a/modules/modelSetup/FluxFineTuneSetup.py +++ b/modules/modelSetup/FluxFineTuneSetup.py @@ -80,8 +80,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/FluxLoRASetup.py b/modules/modelSetup/FluxLoRASetup.py index e3de941fd..e6204565f 100644 --- a/modules/modelSetup/FluxLoRASetup.py +++ b/modules/modelSetup/FluxLoRASetup.py @@ -122,8 +122,6 @@ def setup_model( if model.text_encoder_2 is not None: model.text_encoder_2.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py index 5d94c076d..5b842e267 100644 --- a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py +++ b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py @@ -58,7 +58,6 @@ def setup_model( ): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaFineTuneSetup.py b/modules/modelSetup/PixArtAlphaFineTuneSetup.py index 277f39004..31b92378e 100644 --- a/modules/modelSetup/PixArtAlphaFineTuneSetup.py +++ b/modules/modelSetup/PixArtAlphaFineTuneSetup.py @@ -74,7 +74,6 @@ def setup_model( if config.train_any_embedding(): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/PixArtAlphaLoRASetup.py b/modules/modelSetup/PixArtAlphaLoRASetup.py index 5c865b224..1e18dc673 100644 --- a/modules/modelSetup/PixArtAlphaLoRASetup.py +++ b/modules/modelSetup/PixArtAlphaLoRASetup.py @@ -94,7 +94,6 @@ def setup_model( model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.transformer_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaEmbeddingSetup.py b/modules/modelSetup/SanaEmbeddingSetup.py index 7704c3afd..c31c2aa11 100644 --- a/modules/modelSetup/SanaEmbeddingSetup.py +++ b/modules/modelSetup/SanaEmbeddingSetup.py @@ -58,7 +58,6 @@ def setup_model( ): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaFineTuneSetup.py b/modules/modelSetup/SanaFineTuneSetup.py index 11faa747a..1c334c44f 100644 --- a/modules/modelSetup/SanaFineTuneSetup.py +++ b/modules/modelSetup/SanaFineTuneSetup.py @@ -68,7 +68,6 @@ def setup_model( if config.train_any_embedding(): model.text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/SanaLoRASetup.py b/modules/modelSetup/SanaLoRASetup.py index f481cd445..45cb2f7bb 100644 --- a/modules/modelSetup/SanaLoRASetup.py +++ b/modules/modelSetup/SanaLoRASetup.py @@ -94,7 +94,6 @@ def setup_model( model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.transformer_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py index 1ff229928..e9209bdc4 100644 --- a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py @@ -81,9 +81,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3FineTuneSetup.py b/modules/modelSetup/StableDiffusion3FineTuneSetup.py index 3bbe41821..3ce400a59 100644 --- a/modules/modelSetup/StableDiffusion3FineTuneSetup.py +++ b/modules/modelSetup/StableDiffusion3FineTuneSetup.py @@ -90,9 +90,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusion3LoRASetup.py b/modules/modelSetup/StableDiffusion3LoRASetup.py index ae326cbbe..0d802c18a 100644 --- a/modules/modelSetup/StableDiffusion3LoRASetup.py +++ b/modules/modelSetup/StableDiffusion3LoRASetup.py @@ -148,9 +148,6 @@ def setup_model( if model.text_encoder_3 is not None: model.text_encoder_3.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_3) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionEmbeddingSetup.py b/modules/modelSetup/StableDiffusionEmbeddingSetup.py index 6e3a41c2b..f8cc956f2 100644 --- a/modules/modelSetup/StableDiffusionEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionEmbeddingSetup.py @@ -62,7 +62,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionFineTuneSetup.py b/modules/modelSetup/StableDiffusionFineTuneSetup.py index 8c42e59df..9f6fff8f3 100644 --- a/modules/modelSetup/StableDiffusionFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneSetup.py @@ -76,7 +76,6 @@ def setup_model( elif config.force_epsilon_prediction: model.force_epsilon_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionLoRASetup.py b/modules/modelSetup/StableDiffusionLoRASetup.py index 9b9e9a613..d82147611 100644 --- a/modules/modelSetup/StableDiffusionLoRASetup.py +++ b/modules/modelSetup/StableDiffusionLoRASetup.py @@ -99,7 +99,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py index 7d9ef442f..62f78cba2 100644 --- a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py @@ -71,8 +71,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py index ec6dbe16c..56a5bfe40 100644 --- a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py @@ -86,8 +86,6 @@ def setup_model( elif config.force_epsilon_prediction: model.force_epsilon_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/StableDiffusionXLLoRASetup.py b/modules/modelSetup/StableDiffusionXLLoRASetup.py index 6b193edc5..a0db29b07 100644 --- a/modules/modelSetup/StableDiffusionXLLoRASetup.py +++ b/modules/modelSetup/StableDiffusionXLLoRASetup.py @@ -120,8 +120,6 @@ def setup_model( model.rescale_noise_scheduler_to_zero_terminal_snr() model.force_v_prediction() - self._remove_added_embeddings_from_tokenizer(model.tokenizer_1) - self._remove_added_embeddings_from_tokenizer(model.tokenizer_2) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenEmbeddingSetup.py b/modules/modelSetup/WuerstchenEmbeddingSetup.py index 37dc7f85f..d5dd8a83b 100644 --- a/modules/modelSetup/WuerstchenEmbeddingSetup.py +++ b/modules/modelSetup/WuerstchenEmbeddingSetup.py @@ -62,7 +62,6 @@ def setup_model( ): model.prior_text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenFineTuneSetup.py b/modules/modelSetup/WuerstchenFineTuneSetup.py index 1291125dc..111ce016f 100644 --- a/modules/modelSetup/WuerstchenFineTuneSetup.py +++ b/modules/modelSetup/WuerstchenFineTuneSetup.py @@ -71,7 +71,6 @@ def setup_model( if config.train_any_embedding(): model.prior_text_encoder.get_input_embeddings().to(dtype=config.embedding_weight_dtype.torch_dtype()) - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/WuerstchenLoRASetup.py b/modules/modelSetup/WuerstchenLoRASetup.py index b26f57436..812bb43e6 100644 --- a/modules/modelSetup/WuerstchenLoRASetup.py +++ b/modules/modelSetup/WuerstchenLoRASetup.py @@ -98,7 +98,6 @@ def setup_model( model.prior_prior_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) model.prior_prior_lora.hook_to_module() - self._remove_added_embeddings_from_tokenizer(model.prior_tokenizer) self._setup_embeddings(model, config) self._setup_embedding_wrapper(model, config) diff --git a/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py b/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py index 33b6d34cd..24f726230 100644 --- a/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py +++ b/modules/modelSetup/mixin/ModelSetupEmbeddingMixin.py @@ -13,27 +13,15 @@ CLIPTextModelWithProjection, Gemma2Model, LlamaModel, + PreTrainedTokenizer, T5EncoderModel, ) -from transformers.tokenization_utils import PreTrainedTokenizer, Trie class ModelSetupEmbeddingMixin(metaclass=ABCMeta): def __init__(self): super().__init__() - def _remove_added_embeddings_from_tokenizer( - self, - tokenizer: PreTrainedTokenizer, - ): - if tokenizer: - added_tokens = list(filter(lambda item: not item[1].special, tokenizer._added_tokens_decoder.items())) - for key, added_token in added_tokens: - tokenizer._added_tokens_decoder.pop(key) - tokenizer._added_tokens_encoder.pop(added_token.content) - tokenizer.tokens_trie = Trie() - tokenizer._update_trie() - def _create_new_embedding( self, model: BaseModel, diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index f7c97c132..605994e57 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -118,10 +118,7 @@ def start(self): if self.config.secrets.huggingface_token != "": self.callbacks.on_update_status("logging into Hugging Face") with contextlib.suppress(ConnectionError): - huggingface_hub.login( - token = self.config.secrets.huggingface_token, - new_session = False, - ) + huggingface_hub.login(token=self.config.secrets.huggingface_token) self.callbacks.on_update_status("loading the model") diff --git a/modules/ui/ConceptWindow.py b/modules/ui/ConceptWindow.py index f58879d5f..824ff43e4 100644 --- a/modules/ui/ConceptWindow.py +++ b/modules/ui/ConceptWindow.py @@ -575,7 +575,8 @@ def get_concept_path(path: str) -> str | None: def __download_dataset(self): try: - huggingface_hub.login(token=self.train_config.secrets.huggingface_token, new_session=False) + if self.train_config.secrets.huggingface_token != "": + huggingface_hub.login(token=self.train_config.secrets.huggingface_token) huggingface_hub.snapshot_download(repo_id=self.concept.path, repo_type="dataset") except Exception: traceback.print_exc() diff --git a/modules/util/thread_safety.py b/modules/util/thread_safety.py deleted file mode 100644 index 365076c2a..000000000 --- a/modules/util/thread_safety.py +++ /dev/null @@ -1,43 +0,0 @@ -import functools -import threading - -import torch - -_THREAD_SAFE_FORWARD_ATTR = "_thread_safe_forward_lock" - - -def apply_thread_safe_forward(model: torch.nn.Module) -> None: - """ - Wrap ``model.forward()`` with a per-instance ``threading.Lock`` to - serialize concurrent calls. - - This is a workaround for a thread-safety bug in the transformers library's - ``check_model_inputs`` decorator, which monkey-patches child module - ``.forward()`` methods during execution and is not safe for concurrent use - from multiple dataloader threads. - - See: https://github.com/huggingface/transformers/issues/42673 - Fix: https://github.com/huggingface/transformers/pull/43765 (v5 only) - - This patch can be removed when upgrading to transformers v5+. - - The lock is per-model-instance so different model instances do not block - each other. The function is idempotent: calling it twice on the same model - is a no-op. - - Args: - model: The ``nn.Module`` whose ``forward()`` should be made thread-safe. - """ - if hasattr(model, _THREAD_SAFE_FORWARD_ATTR): - return - - lock = threading.Lock() - original_forward = model.forward - - @functools.wraps(original_forward) - def locked_forward(*args, **kwargs): - with lock: - return original_forward(*args, **kwargs) - - model.forward = locked_forward - setattr(model, _THREAD_SAFE_FORWARD_ATTR, lock) diff --git a/requirements-global.txt b/requirements-global.txt index 91014c6ff..299f765b0 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -5,7 +5,7 @@ pillow==12.2.0 imagesize==1.4.1 #for concept statistics tqdm==4.67.1 PyYAML==6.0.2 -huggingface-hub==0.34.4 +huggingface-hub==1.16.1 scipy==1.15.3 matplotlib==3.10.3 av==16.1.0 @@ -22,7 +22,7 @@ tensorboard==2.20.0 # diffusion models -e git+https://github.com/huggingface/diffusers.git@0f1abc4#egg=diffusers gguf==0.17.1 -transformers==4.57.6 +transformers==5.9.0 sentencepiece==0.2.1 # transitive dependency of transformers for tokenizer loading omegaconf==2.3.0 # needed to load stable diffusion from single ckpt files invisible-watermark==0.2.0 # needed for the SDXL pipeline diff --git a/run-cmd.sh b/run-cmd.sh index 2b1c9d14f..8fc77cb02 100755 --- a/run-cmd.sh +++ b/run-cmd.sh @@ -2,11 +2,6 @@ set -e -# Xet is buggy. Disabled by default unless already defined - https://github.com/Nerogar/OneTrainer/issues/949 -if [[ -z "${HF_HUB_DISABLE_XET+x}" ]]; then - export HF_HUB_DISABLE_XET=1 -fi - source "${BASH_SOURCE[0]%/*}/lib.include.sh" # Fetch and validate the name of the target script. diff --git a/start-ui.bat b/start-ui.bat index 5edc8e6c4..05881410b 100644 --- a/start-ui.bat +++ b/start-ui.bat @@ -35,14 +35,6 @@ set PYTHON="%VENV_DIR%\Scripts\python.exe" -X utf8 if defined PROFILE (set PYTHON=%PYTHON% -m scalene --off --cpu --gpu --profile-all --no-browser) echo Using Python %PYTHON% -REM Disable HF_HUB_DISABLE_XET, buggy; default disables Xet (set to 0 to enable) - https://github.com/Nerogar/OneTrainer/issues/949 -if not defined HF_HUB_DISABLE_XET ( - set "HF_HUB_DISABLE_XET=1" -) -echo HF_HUB_DISABLE_XET=%HF_HUB_DISABLE_XET% -echo. -echo NOTE: Xet disabled, to enable it set as 0 before launch - :check_python_version echo Checking Python version... %PYTHON% --version diff --git a/start-ui.sh b/start-ui.sh index 2f0eecc0d..b2960c262 100755 --- a/start-ui.sh +++ b/start-ui.sh @@ -4,11 +4,6 @@ set -e source "${BASH_SOURCE[0]%/*}/lib.include.sh" -# Xet is buggy. Disabled by default unless already defined - https://github.com/Nerogar/OneTrainer/issues/949 -if [[ -z "${HF_HUB_DISABLE_XET+x}" ]]; then - export HF_HUB_DISABLE_XET=1 -fi - prepare_runtime_environment run_python_in_active_env "scripts/train_ui.py" "$@" From 18a35524383570b3a1ea4ce498950f15446fa5d0 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 6 Jun 2026 13:32:04 +0200 Subject: [PATCH 03/10] =?UTF-8?q?Add=20on-demand=20text-encoder=20loading?= =?UTF-8?q?=20/=20rename=20BaseModel.to()=20=E2=86=92=20release()?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces OnDemandModule, a persistent delegating proxy for text encoders that must be loaded on demand and freed after use rather than parked on the CPU temp device. Adds load_on_demand per-component config and four text_encoder_N_on_demand() resolvers in TrainConfig. BaseModel.to(device) is removed as an abstract method; release() is now the sole abstract method for parking a model. Each concrete model reads self.train_config.temp_device directly. Call sites in modelSetup, dataLoader, trainer, and SampleWindow are updated to model.release(). Co-Authored-By: dxqb <183307934+dxqb@users.noreply.github.com> Co-Authored-By: Claude Sonnet 4.6 --- .../StableDiffusionFineTuneVaeDataLoader.py | 2 +- .../dataLoader/WuerstchenBaseDataLoader.py | 4 +- .../mixin/DataLoaderText2ImageMixin.py | 4 +- modules/model/BaseModel.py | 5 +- modules/model/ChromaModel.py | 8 +-- modules/model/ErnieModel.py | 8 +-- modules/model/Flux2Model.py | 8 +-- modules/model/FluxModel.py | 8 +-- modules/model/HiDreamModel.py | 8 +-- modules/model/HunyuanVideoModel.py | 8 +-- modules/model/PixArtAlphaModel.py | 8 +-- modules/model/QwenModel.py | 8 +-- modules/model/SanaModel.py | 8 +-- modules/model/StableDiffusion3Model.py | 8 +-- modules/model/StableDiffusionModel.py | 10 +-- modules/model/StableDiffusionXLModel.py | 8 +-- modules/model/WuerstchenModel.py | 14 ++-- modules/model/ZImageModel.py | 8 +-- modules/modelSetup/BaseChromaSetup.py | 2 +- modules/modelSetup/BaseErnieSetup.py | 2 +- modules/modelSetup/BaseFlux2Setup.py | 2 +- modules/modelSetup/BaseFluxSetup.py | 2 +- modules/modelSetup/BaseHiDreamSetup.py | 2 +- modules/modelSetup/BaseHunyuanVideoSetup.py | 2 +- modules/modelSetup/BasePixArtAlphaSetup.py | 2 +- modules/modelSetup/BaseQwenSetup.py | 2 +- modules/modelSetup/BaseSanaSetup.py | 2 +- .../modelSetup/BaseStableDiffusion3Setup.py | 2 +- .../modelSetup/BaseStableDiffusionSetup.py | 2 +- .../modelSetup/BaseStableDiffusionXLSetup.py | 2 +- modules/modelSetup/BaseWuerstchenSetup.py | 2 +- modules/modelSetup/BaseZImageSetup.py | 2 +- modules/trainer/GenericTrainer.py | 12 ++-- modules/ui/SampleWindow.py | 2 +- modules/util/ModelNames.py | 8 +++ modules/util/OnDemandModule.py | 67 +++++++++++++++++++ modules/util/config/TrainConfig.py | 28 ++++++++ 37 files changed, 192 insertions(+), 88 deletions(-) create mode 100644 modules/util/OnDemandModule.py diff --git a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py index ed5dd32b8..27f91a2e5 100644 --- a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py +++ b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py @@ -47,7 +47,7 @@ def _setup_cache_device( temp_device: torch.device, config: TrainConfig, ): - model.to(self.temp_device) + model.release() model.vae_to(train_device) diff --git a/modules/dataLoader/WuerstchenBaseDataLoader.py b/modules/dataLoader/WuerstchenBaseDataLoader.py index f5cf2f41a..9a42e57f6 100644 --- a/modules/dataLoader/WuerstchenBaseDataLoader.py +++ b/modules/dataLoader/WuerstchenBaseDataLoader.py @@ -72,7 +72,7 @@ def _cache_modules(self, config: TrainConfig, model: WuerstchenModel, model_setu ] def before_cache_image_fun(): - model.to(self.temp_device) + model.release() model.effnet_encoder_to(self.train_device) model.eval() torch_gc() @@ -107,7 +107,7 @@ def _output_modules(self, config: TrainConfig, model: WuerstchenModel, model_set output_names.append('pooled_text_encoder_output') def before_cache_image_fun(): - model.to(self.temp_device) + model.release() model.effnet_encoder_to(self.train_device) model.eval() torch_gc() diff --git a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py b/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py index 2654bdd19..af79459d6 100644 --- a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py +++ b/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py @@ -274,7 +274,7 @@ def _output_modules_from_out_names( ): if before_cache_image_fun is None: def prepare_vae(): - model.to(self.temp_device) + model.release() model.vae_to(self.train_device) model.eval() torch_gc() @@ -340,7 +340,7 @@ def _cache_modules_from_names( if before_cache_image_fun is None: def prepare_vae(): - model.to(self.temp_device) + model.release() model.vae_to(self.train_device) model.eval() torch_gc() diff --git a/modules/model/BaseModel.py b/modules/model/BaseModel.py index fef1fa9e1..e1caca247 100644 --- a/modules/model/BaseModel.py +++ b/modules/model/BaseModel.py @@ -94,8 +94,11 @@ def __init__( self.autocast_context = nullcontext() self.train_dtype = DataType.FLOAT_32 + #park the whole model on the temp device to free VRAM. Models with on-demand components + #(which cannot be parked, only discarded and rebuilt) override this to free those components + #instead of moving them. @abstractmethod - def to(self, device: torch.device): + def release(self): pass @abstractmethod diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index bff51d3e2..72f134818 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -128,10 +128,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/ErnieModel.py b/modules/model/ErnieModel.py index 3d15e5385..8a78d5b78 100644 --- a/modules/model/ErnieModel.py +++ b/modules/model/ErnieModel.py @@ -86,10 +86,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/Flux2Model.py b/modules/model/Flux2Model.py index a501bb116..cd02747fe 100644 --- a/modules/model/Flux2Model.py +++ b/modules/model/Flux2Model.py @@ -136,10 +136,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index 76a18e8f9..5cd261de0 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -162,10 +162,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index 4bad3b62c..e1a4bbee1 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -247,10 +247,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 0b1e4b990..936d9645a 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -181,10 +181,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index b0cc8c541..89b74c417 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -129,10 +129,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/QwenModel.py b/modules/model/QwenModel.py index 1aadcd467..361d7ee2f 100644 --- a/modules/model/QwenModel.py +++ b/modules/model/QwenModel.py @@ -98,10 +98,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index b226d4463..3a210eefc 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -131,10 +131,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index 529209e8b..fedf448bc 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -191,10 +191,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/StableDiffusionModel.py b/modules/model/StableDiffusionModel.py index 2bee0fd70..e114201e4 100644 --- a/modules/model/StableDiffusionModel.py +++ b/modules/model/StableDiffusionModel.py @@ -123,11 +123,11 @@ def unet_to(self, device: torch.device): if self.unet_lora is not None: self.unet_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.depth_estimator_to(device) - self.text_encoder_to(device) - self.unet_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.depth_estimator_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.unet_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/StableDiffusionXLModel.py b/modules/model/StableDiffusionXLModel.py index b68df2d23..f4da4905f 100644 --- a/modules/model/StableDiffusionXLModel.py +++ b/modules/model/StableDiffusionXLModel.py @@ -155,10 +155,10 @@ def unet_to(self, device: torch.device): if self.unet_lora is not None: self.unet_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.unet_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.unet_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/WuerstchenModel.py b/modules/model/WuerstchenModel.py index b844d6eb1..726621ba9 100644 --- a/modules/model/WuerstchenModel.py +++ b/modules/model/WuerstchenModel.py @@ -161,14 +161,14 @@ def prior_prior_to(self, device: torch.device): if self.prior_prior_lora is not None: self.prior_prior_lora.to(device) - def to(self, device: torch.device): + def release(self): if self.model_type.is_wuerstchen_v2(): - self.decoder_text_encoder_to(device) - self.decoder_decoder_to(device) - self.decoder_vqgan_to(device) - self.effnet_encoder_to(device) - self.prior_text_encoder_to(device) - self.prior_prior_to(device) + self.decoder_text_encoder_to(self.train_config.temp_device) + self.decoder_decoder_to(self.train_config.temp_device) + self.decoder_vqgan_to(self.train_config.temp_device) + self.effnet_encoder_to(self.train_config.temp_device) + self.prior_text_encoder_to(self.train_config.temp_device) + self.prior_prior_to(self.train_config.temp_device) def eval(self): if self.model_type.is_wuerstchen_v2(): diff --git a/modules/model/ZImageModel.py b/modules/model/ZImageModel.py index 167b58f47..33e09ffd7 100644 --- a/modules/model/ZImageModel.py +++ b/modules/model/ZImageModel.py @@ -97,10 +97,10 @@ def transformer_to(self, device: torch.device): if self.transformer_lora is not None: self.transformer_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.transformer_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index f8cc33187..806693fe6 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -277,7 +277,7 @@ def calculate_loss( def prepare_text_caching(self, model: ChromaModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index 72f8ee2d8..7034ae285 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -164,7 +164,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: ErnieModel, config: TrainConfig): - model.to(self.temp_device) + model.release() model.text_encoder_to(self.train_device) model.eval() torch_gc() diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index 6e6e7320f..577e26a51 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -191,7 +191,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: FluxModel, config: TrainConfig): - model.to(self.temp_device) + model.release() model.text_encoder_to(self.train_device) model.eval() torch_gc() diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index e6638ddb2..bebf3b6c4 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -337,7 +337,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: FluxModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index b5a856b73..806833ff6 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -429,7 +429,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: HiDreamModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index 2e32165e7..d8b097b89 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -306,7 +306,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: HunyuanVideoModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 4935ebbfd..6ba03034e 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -338,7 +338,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: PixArtAlphaModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index 82004d8ac..9d7e20d15 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -187,7 +187,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: QwenModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index c26acbc82..7efb015f0 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -261,7 +261,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: SanaModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index efe761427..637dd0d38 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -359,7 +359,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusion3Model, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index a16c72cb8..e9c6407df 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -342,7 +342,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusionModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 7fe35e1e1..e23c4dfd3 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -390,7 +390,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusionXLModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index abc77bcdd..662f9a7bb 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -362,7 +362,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: WuerstchenModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index d64e7e88a..daa172f10 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -170,7 +170,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: ZImageModel, config: TrainConfig): - model.to(self.temp_device) + model.release() model.text_encoder_to(self.train_device) model.eval() diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index f7c97c132..10939f96a 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -142,8 +142,7 @@ def start(self): self.model_setup.setup_optimizations(self.model, self.config) self.model_setup.setup_train_device(self.model, self.config) self.model_setup.setup_model(self.model, self.config) - self.model.to(self.temp_device) - self.model.eval() + self.model.release() torch_gc() self.callbacks.on_update_status("creating the data loader/caching") @@ -256,8 +255,7 @@ def on_sample_custom(sampler_output: ModelSamplerOutput): on_sample = on_sample_custom if is_custom_sample else on_sample_default on_update_progress = self.callbacks.on_update_sample_custom_progress if is_custom_sample else self.callbacks.on_update_sample_default_progress - self.model.to(self.temp_device) - self.model.eval() + self.model.release() sample_config = copy.copy(sample_config) sample_config.from_train_config(self.config) @@ -719,7 +717,7 @@ def sample_commands_fun(): backup = self.commands.get_and_reset_backup_command() save = self.commands.get_and_reset_save_command() if multi.is_master() and (backup or save): - self.model.to(self.temp_device) + self.model.release() if backup: self.__backup(train_progress, True, step_tqdm.write) if save: @@ -844,7 +842,7 @@ def sample_commands_fun(): def end(self): if self.one_step_trained: - self.model.to(self.temp_device) + self.model.release() if self.config.backup_before_save and multi.is_master(): self.__backup(self.model.train_progress) @@ -877,7 +875,7 @@ def end(self): ) if self.model is not None: - self.model.to(self.temp_device) + self.model.release() if multi.is_master(): self.tensorboard.close() diff --git a/modules/ui/SampleWindow.py b/modules/ui/SampleWindow.py index 0f91ad2fa..bdd664d84 100644 --- a/modules/ui/SampleWindow.py +++ b/modules/ui/SampleWindow.py @@ -144,7 +144,7 @@ def __load_model(self) -> BaseModel: model_setup.setup_optimizations(model, self.initial_train_config) model_setup.setup_train_device(model, self.initial_train_config) model_setup.setup_model(model, self.initial_train_config) - model.to(torch.device(self.initial_train_config.temp_device)) + model.release() return model diff --git a/modules/util/ModelNames.py b/modules/util/ModelNames.py index 8dec9a9bf..c58a2d284 100644 --- a/modules/util/ModelNames.py +++ b/modules/util/ModelNames.py @@ -25,6 +25,10 @@ def __init__( include_text_encoder_2: bool = True, include_text_encoder_3: bool = True, include_text_encoder_4: bool = True, + text_encoder_on_demand: bool = False, + text_encoder_2_on_demand: bool = False, + text_encoder_3_on_demand: bool = False, + text_encoder_4_on_demand: bool = False, ): self.base_model = base_model self.prior_model = prior_model @@ -40,6 +44,10 @@ def __init__( self.include_text_encoder_2 = include_text_encoder_2 self.include_text_encoder_3 = include_text_encoder_3 self.include_text_encoder_4 = include_text_encoder_4 + self.text_encoder_on_demand = text_encoder_on_demand + self.text_encoder_2_on_demand = text_encoder_2_on_demand + self.text_encoder_3_on_demand = text_encoder_3_on_demand + self.text_encoder_4_on_demand = text_encoder_4_on_demand def all_embedding(self): if self.embedding is not None: diff --git a/modules/util/OnDemandModule.py b/modules/util/OnDemandModule.py new file mode 100644 index 000000000..e8eba8cbf --- /dev/null +++ b/modules/util/OnDemandModule.py @@ -0,0 +1,67 @@ +from collections.abc import Callable + +from modules.util.torch_util import torch_gc + +from torch import nn + + +# A persistent delegating proxy for a module loaded on demand and discarded +# (weights freed) after use. Always truthy, so presence gates and captured +# references (e.g. MGDS Encode*Text nodes) stay valid while no weights are loaded. +# On-demand modules are frozen and inference-only by definition, so the proxy +# enforces that invariant: eval/requires_grad_(False) are accepted, while +# train(True)/requires_grad_(True) raise. Other attribute access delegates to the +# inner. The loader loads to cpu; the caller (model) is responsible for any +# quantization and for moving the materialized module to the accelerator. +# +# Intentionally a plain object, not an nn.Module, to avoid submodule/parameter +# registration fighting the materialize/discard swap. +class OnDemandModule: + def __init__(self, loader: Callable[[], nn.Module]): + self._inner: nn.Module | None = None + self._loader = loader + + def __bool__(self): + return True + + @property + def inner(self) -> nn.Module | None: + return self._inner + + def __call__(self, *args, **kwargs): + return self._inner(*args, **kwargs) + + def __getattr__(self, name): + if self._inner is None: + raise AttributeError( + f"OnDemandModule has no materialized inner module; cannot access '{name}'" + ) + return getattr(self._inner, name) + + def eval(self): + return self.train(False) + + def train(self, mode: bool = True): + if mode: + raise RuntimeError("OnDemandModule is inference-only; train mode is not allowed") + return self + + def requires_grad_(self, requires_grad: bool = True): + if requires_grad: + raise RuntimeError("OnDemandModule is frozen; requires_grad_(True) is not allowed") + return self + + # No-op: the materialized module is placed on the accelerator by the caller and + # freed by discard(). A blanket model.to(temp_device) must never move or reload it. + def to(self, *args, **kwargs): + return self + + def materialize(self): + if self._inner is None: + self._inner = self._loader() + self._inner.eval() + self._inner.requires_grad_(False) + + def discard(self): + self._inner = None + torch_gc() diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 92528a1c2..68281b8dd 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -269,6 +269,7 @@ class TrainModelPartConfig(BaseConfig): gradient_checkpointing: bool offload_fraction: float activation_offloading: bool + load_on_demand: bool def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -306,6 +307,7 @@ def default_values(): data.append(("gradient_checkpointing", True, bool, False)) data.append(("offload_fraction", 0.0, float, False)) data.append(("activation_offloading", True, bool, False)) + data.append(("load_on_demand", False, bool, False)) return TrainModelPartConfig(data) @@ -893,6 +895,10 @@ def model_names(self) -> ModelNames: include_text_encoder_2=self.text_encoder_2.include, include_text_encoder_3=self.text_encoder_3.include, include_text_encoder_4=self.text_encoder_4.include, + text_encoder_on_demand=self.text_encoder_on_demand(), + text_encoder_2_on_demand=self.text_encoder_2_on_demand(), + text_encoder_3_on_demand=self.text_encoder_3_on_demand(), + text_encoder_4_on_demand=self.text_encoder_4_on_demand(), ) def train_any_embedding(self) -> bool: @@ -927,6 +933,28 @@ def train_text_encoder_4_or_embedding(self) -> bool: or ((self.text_encoder_4.train_embedding or not self.model_type.has_multiple_text_encoders()) and self.train_any_embedding()) + #an encoder is loaded on demand only when it is requested, frozen (not trained, no embedding + #training) and its conditioning is cached -- otherwise it is needed resident every step. + def text_encoder_on_demand(self) -> bool: + return self.text_encoder.load_on_demand \ + and not self.train_text_encoder_or_embedding() \ + and self.latent_caching + + def text_encoder_2_on_demand(self) -> bool: + return self.text_encoder_2.load_on_demand \ + and not self.train_text_encoder_2_or_embedding() \ + and self.latent_caching + + def text_encoder_3_on_demand(self) -> bool: + return self.text_encoder_3.load_on_demand \ + and not self.train_text_encoder_3_or_embedding() \ + and self.latent_caching + + def text_encoder_4_on_demand(self) -> bool: + return self.text_encoder_4.load_on_demand \ + and not self.train_text_encoder_4_or_embedding() \ + and self.latent_caching + def all_embedding_configs(self): if self.training_method == TrainingMethod.EMBEDDING: return self.additional_embeddings + [self.embedding] From 0fffc53fb5457aeb6f10c993d72e8b0b04fbcc09 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 6 Jun 2026 14:28:16 +0200 Subject: [PATCH 04/10] Add Lens model (LoRA + Fine-Tune training + sampling) Co-Authored-By: Claude Sonnet 4.6 --- modules/dataLoader/LensBaseDataLoader.py | 161 ++++++++++ modules/model/Flux2Model.py | 2 +- modules/model/LensModel.py | 297 +++++++++++++++++++ modules/modelLoader/Flux2ModelLoader.py | 3 +- modules/modelLoader/LensModelLoader.py | 212 +++++++++++++ modules/modelSampler/LensSampler.py | 189 ++++++++++++ modules/modelSaver/LensFineTuneModelSaver.py | 11 + modules/modelSaver/LensLoRAModelSaver.py | 11 + modules/modelSaver/lens/LensLoRASaver.py | 38 +++ modules/modelSaver/lens/LensModelSaver.py | 87 ++++++ modules/modelSetup/BaseLensSetup.py | 180 +++++++++++ modules/modelSetup/LensFineTuneSetup.py | 96 ++++++ modules/modelSetup/LensLoRASetup.py | 108 +++++++ modules/ui/ConvertModelUI.py | 1 + modules/ui/ModelTab.py | 4 +- modules/ui/TopBar.py | 1 + modules/ui/TrainingTab.py | 31 +- modules/util/checkpointing_util.py | 21 ++ modules/util/config/SampleConfig.py | 7 + modules/util/config/TrainConfig.py | 2 + modules/util/enum/ModelType.py | 11 +- modules/util/optimizer/muon_util.py | 2 +- requirements-global.txt | 5 +- resources/sd_model_spec/lens-lora.json | 6 + resources/sd_model_spec/lens.json | 6 + training_presets/#lens Finetune 16GB.json | 47 +++ training_presets/#lens LoRA 16GB.json | 30 ++ 27 files changed, 1554 insertions(+), 15 deletions(-) create mode 100644 modules/dataLoader/LensBaseDataLoader.py create mode 100644 modules/model/LensModel.py create mode 100644 modules/modelLoader/LensModelLoader.py create mode 100644 modules/modelSampler/LensSampler.py create mode 100644 modules/modelSaver/LensFineTuneModelSaver.py create mode 100644 modules/modelSaver/LensLoRAModelSaver.py create mode 100644 modules/modelSaver/lens/LensLoRASaver.py create mode 100644 modules/modelSaver/lens/LensModelSaver.py create mode 100644 modules/modelSetup/BaseLensSetup.py create mode 100644 modules/modelSetup/LensFineTuneSetup.py create mode 100644 modules/modelSetup/LensLoRASetup.py create mode 100644 resources/sd_model_spec/lens-lora.json create mode 100644 resources/sd_model_spec/lens.json create mode 100644 training_presets/#lens Finetune 16GB.json create mode 100644 training_presets/#lens LoRA 16GB.json diff --git a/modules/dataLoader/LensBaseDataLoader.py b/modules/dataLoader/LensBaseDataLoader.py new file mode 100644 index 000000000..9c3d2816e --- /dev/null +++ b/modules/dataLoader/LensBaseDataLoader.py @@ -0,0 +1,161 @@ +import os + +from modules.dataLoader.BaseDataLoader import BaseDataLoader +from modules.dataLoader.mixin.DataLoaderText2ImageMixin import DataLoaderText2ImageMixin +from modules.model.LensModel import ( + PROMPT_MAX_LENGTH, + PROMPT_TEMPLATE_CROP_START, + LensModel, + make_lens_conversation, +) +from modules.modelSetup.BaseLensSetup import BaseLensSetup +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.TrainProgress import TrainProgress + +from mgds.pipelineModules.DecodeTokens import DecodeTokens +from mgds.pipelineModules.DecodeVAE import DecodeVAE +from mgds.pipelineModules.EncodeLensText import EncodeLensText +from mgds.pipelineModules.EncodeVAE import EncodeVAE +from mgds.pipelineModules.RescaleImageChannels import RescaleImageChannels +from mgds.pipelineModules.SampleVAEDistribution import SampleVAEDistribution +from mgds.pipelineModules.SaveImage import SaveImage +from mgds.pipelineModules.SaveText import SaveText +from mgds.pipelineModules.ScaleImage import ScaleImage +from mgds.pipelineModules.Tokenize import Tokenize + + +class LensBaseDataLoader( + BaseDataLoader, + DataLoaderText2ImageMixin, +): + def _preparation_modules(self, config: TrainConfig, model: LensModel): + rescale_image = RescaleImageChannels(image_in_name='image', image_out_name='image', in_range_min=0, in_range_max=1, out_range_min=-1, out_range_max=1) + encode_image = EncodeVAE(in_name='image', out_name='latent_image_distribution', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) + image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') + downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125) + + tokenize_prompt = Tokenize( + in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', + tokenizer=model.tokenizer, + max_token_length=PROMPT_MAX_LENGTH + PROMPT_TEMPLATE_CROP_START, + apply_chat_template=make_lens_conversation, + apply_chat_template_kwargs={'add_generation_prompt': False}, + apply_chat_template_post_process=lambda t: t.split("<|return|>")[0], + ) + + encode_prompt = EncodeLensText( + tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', + hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', + text_encoder=model.text_encoder, + crop_start=PROMPT_TEMPLATE_CROP_START, + autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), + ) + + + modules = [rescale_image, encode_image, image_sample] + if config.masked_training or config.model_type.has_mask_input(): + modules.append(downscale_mask) + + modules += [tokenize_prompt, encode_prompt] + return modules + + def _cache_modules(self, config: TrainConfig, model: LensModel, model_setup: BaseLensSetup): + image_split_names = ['latent_image', 'original_resolution', 'crop_offset'] + + if config.masked_training or config.model_type.has_mask_input(): + image_split_names.append('latent_mask') + + image_aggregate_names = ['crop_resolution', 'image_path'] + + text_split_names = [] + + sort_names = image_aggregate_names + image_split_names + [ + 'prompt', 'tokens', 'tokens_mask', 'text_encoder_hidden_state', + 'concept' + ] + + text_split_names += ['tokens', 'tokens_mask', 'text_encoder_hidden_state'] + + return self._cache_modules_from_names( + model, model_setup, + image_split_names=image_split_names, + image_aggregate_names=image_aggregate_names, + text_split_names=text_split_names, + sort_names=sort_names, + config=config, + text_caching=True, + ) + + def _output_modules(self, config: TrainConfig, model: LensModel, model_setup: BaseLensSetup): + output_names = [ + 'image_path', 'latent_image', + 'prompt', + 'tokens', + 'tokens_mask', + 'original_resolution', 'crop_resolution', 'crop_offset', + ] + + if config.masked_training or config.model_type.has_mask_input(): + output_names.append('latent_mask') + + output_names.append('text_encoder_hidden_state') + + return self._output_modules_from_out_names( + model, model_setup, + output_names=output_names, + config=config, + use_conditioning_image=False, + vae=model.vae, + autocast_context=[model.autocast_context], + train_dtype=model.train_dtype, + ) + + def _debug_modules(self, config: TrainConfig, model: LensModel): + debug_dir = os.path.join(config.debug_dir, "dataloader") + + def before_save_fun(): + model.vae_to(self.train_device) + + decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) + upscale_mask = ScaleImage(in_name='latent_mask', out_name='decoded_mask', factor=8) + decode_prompt = DecodeTokens(in_name='tokens', out_name='decoded_prompt', tokenizer=model.tokenizer) + save_image = SaveImage(image_in_name='decoded_image', original_path_in_name='image_path', path=debug_dir, in_range_min=-1, in_range_max=1, before_save_fun=before_save_fun) + # SaveImage(image_in_name='latent_mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1, before_save_fun=before_save_fun) + save_mask = SaveImage(image_in_name='decoded_mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1, before_save_fun=before_save_fun) + save_prompt = SaveText(text_in_name='decoded_prompt', original_path_in_name='image_path', path=debug_dir, before_save_fun=before_save_fun) + + # These modules don't really work, since they are inserted after a sorting operation that does not include this data + # SaveImage(image_in_name='mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1), + # SaveImage(image_in_name='image', original_path_in_name='image_path', path=debug_dir, in_range_min=-1, in_range_max=1), + + modules = [] + + modules.append(decode_image) + modules.append(save_image) + + if config.masked_training or config.model_type.has_mask_input(): + modules.append(upscale_mask) + modules.append(save_mask) + + modules.append(decode_prompt) + modules.append(save_prompt) + + return modules + + def _create_dataset( + self, + config: TrainConfig, + model: LensModel, + model_setup: BaseLensSetup, + train_progress: TrainProgress, + is_validation: bool = False, + ): + return DataLoaderText2ImageMixin._create_dataset(self, + config, model, model_setup, train_progress, is_validation, + aspect_bucketing_quantization=64, + ) + + +factory.register(BaseDataLoader, LensBaseDataLoader, ModelType.LENS) diff --git a/modules/model/Flux2Model.py b/modules/model/Flux2Model.py index cd02747fe..58b67d3f4 100644 --- a/modules/model/Flux2Model.py +++ b/modules/model/Flux2Model.py @@ -162,7 +162,7 @@ def encode_text( train_device: torch.device, batch_size: int = 1, #TODO unused rand: Random | None = None, - text: str = None, + text: str | list[str] = None, tokens: Tensor = None, tokens_mask: Tensor = None, text_encoder_sequence_length: int | None = None, diff --git a/modules/model/LensModel.py b/modules/model/LensModel.py new file mode 100644 index 000000000..525aebf4f --- /dev/null +++ b/modules/model/LensModel.py @@ -0,0 +1,297 @@ +import math +from contextlib import nullcontext +from random import Random + +from modules.model.BaseModel import BaseModel +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util.enum.ModelType import ModelType +from modules.util.LayerOffloadConductor import LayerOffloadConductor +from modules.util.OnDemandModule import OnDemandModule +from modules.util.torch_util import torch_gc + +import torch +from torch import Tensor + +from diffusers import ( + AutoencoderKLFlux2, + DiffusionPipeline, + FlowMatchEulerDiscreteScheduler, +) +from transformers import PreTrainedTokenizerFast + +from lens.pipeline import LensPipeline +from lens.text_encoder import LensGptOssEncoder +from lens.transformer import LensTransformer2DModel + +# Chat template constants, matching lens/pipeline.py +CHAT_SYSTEM = ( + "Describe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background." +) +CHAT_ASSISTANT_THINKING = "Need to generate one image according to the description." +PROMPT_TEMPLATE_CROP_START = 97 # tokens consumed by the chat template prefix +PROMPT_MAX_LENGTH = 512 # caption token budget (chat template tokens are added on top) + + +def make_lens_conversation(caption: str) -> list[dict]: + return [ + {"role": "system", "content": CHAT_SYSTEM, "thinking": None}, + {"role": "user", "content": caption, "thinking": None}, + {"role": "assistant", "thinking": CHAT_ASSISTANT_THINKING, "content": ""}, + ] + + +class LensModel(BaseModel): + # base model data + tokenizer: PreTrainedTokenizerFast | None + noise_scheduler: FlowMatchEulerDiscreteScheduler | None + text_encoder: LensGptOssEncoder | OnDemandModule | None + text_encoder_hidden_size: int | None # cached so encode_text() works after encoder is deleted + vae: AutoencoderKLFlux2 | None + transformer: LensTransformer2DModel | None + + # autocast context + text_encoder_autocast_context: torch.autocast | nullcontext + + transformer_offload_conductor: LayerOffloadConductor | None + + transformer_lora: LoRAModuleWrapper | None + lora_state_dict: dict | None + + def __init__( + self, + model_type: ModelType, + ): + super().__init__( + model_type=model_type, + ) + + self.tokenizer = None + self.noise_scheduler = None + self.text_encoder = None + self.text_encoder_hidden_size = None + self.vae = None + self.transformer = None + + self.text_encoder_autocast_context = nullcontext() + + self.transformer_offload_conductor = None + + self.transformer_lora = None + self.lora_state_dict = None + + def adapters(self) -> list[LoRAModuleWrapper]: + return [a for a in [ + self.transformer_lora, + ] if a is not None] + + def vae_to(self, device: torch.device): + self.vae.to(device=device) + + def text_encoder_to(self, device: torch.device): + # only moves a resident/materialized encoder. The on-demand proxy's to() is a no-op; + # its placement is owned by materialize_text_encoder / release_text_encoder. + if self.text_encoder is not None: + self.text_encoder.to(device=device) + + # Lens always loads on demand (see TrainConfig.text_encoder_on_demand): the GPT-OSS encoder is + # MXFP4-quantized and cannot be parked on the CPU temp device, so it is rebuilt straight onto the + # accelerator when needed and discarded afterwards. MXFP4 quantization happens inside + # from_pretrained (the loader lambda), so no quantize_layers call is needed here. + def materialize_text_encoder(self, device: torch.device): + if isinstance(self.text_encoder, OnDemandModule): + self.text_encoder.materialize() + self.text_encoder.inner.to(device) + else: + self.text_encoder_to(device) # resident: just move + + # load an on-demand encoder for saving. resident encoders already hold their weights. + def materialize_text_encoder_for_save(self): + if isinstance(self.text_encoder, OnDemandModule): + self.text_encoder.materialize() + + def release_text_encoder(self): + if isinstance(self.text_encoder, OnDemandModule): + self.text_encoder.discard() # free the weights + torch_gc() + else: + self.text_encoder_to(torch.device(self.train_config.temp_device)) # resident: park on temp + + def transformer_to(self, device: torch.device): + if self.transformer_offload_conductor is not None and \ + self.transformer_offload_conductor.layer_offload_activated(): + self.transformer_offload_conductor.to(device) + else: + self.transformer.to(device=device) + + if self.transformer_lora is not None: + self.transformer_lora.to(device) + + def to(self, device: torch.device): + self.vae_to(device) + self.text_encoder_to(device) + self.transformer_to(device) + + # park resident components on the temp device and discard the on-demand text encoder + def release(self): + temp_device = torch.device(self.train_config.temp_device) + self.vae_to(temp_device) + self.release_text_encoder() + self.transformer_to(temp_device) + + def eval(self): + self.vae.eval() + if self.text_encoder is not None: + self.text_encoder.eval() + self.transformer.eval() + + # the real module behind the (possibly on-demand) text_encoder: the inner module when + # materialized, None when discarded, or the resident module. The pipeline and save path + # need the unwrapped module, not the proxy. + def _resolved_text_encoder(self) -> torch.nn.Module | None: + if isinstance(self.text_encoder, OnDemandModule): + return self.text_encoder.inner + return self.text_encoder + + def create_pipeline(self) -> DiffusionPipeline: + return LensPipeline( + scheduler=self.noise_scheduler, + vae=self.vae, + text_encoder=self._resolved_text_encoder(), + tokenizer=self.tokenizer, + transformer=self.transformer, + ) + + def encode_text( + self, + train_device: torch.device, + batch_size: int = 1, + rand: Random | None = None, + text: str | list[str] = None, + tokens: Tensor = None, + tokens_mask: Tensor = None, + text_encoder_dropout_probability: float | None = None, + text_encoder_output: list[Tensor] | None = None, + ) -> tuple[list[Tensor] | None, Tensor | None]: + + # the on-demand proxy is always truthy, so resolve the real (materialized) encoder here and + # gate the fresh-encode path on it being present rather than on `self.text_encoder`. + text_encoder = self._resolved_text_encoder() + + if tokens is None and text is not None: + if isinstance(text, str): + text = [text] + + rendered = [] + for prompt in text: + t = self.tokenizer.apply_chat_template( + make_lens_conversation(prompt), tokenize=False, add_generation_prompt=False + ) + rendered.append(t.split("<|return|>")[0]) + + tokenizer_output = self.tokenizer( + rendered, + max_length=PROMPT_MAX_LENGTH + PROMPT_TEMPLATE_CROP_START, + padding='max_length', + truncation=True, + return_tensors="pt", + add_special_tokens=True, + ) + tokens = tokenizer_output.input_ids.to(text_encoder.device) + tokens_mask = tokenizer_output.attention_mask.to(text_encoder.device) + + if text_encoder_output is None and text_encoder is not None: + with self.text_encoder_autocast_context: + # encode_layers() is used instead of the standard output_hidden_states=True API + # because transformers' @capture_outputs decorator applies tie_last_hidden_states, + # which silently replaces hidden_states[-1] with the norm-applied final output. + # GPT-OSS has 24 layers and selected_layer_index ends at layer 23 (the last), + # so output_hidden_states=True returns the normed version for that layer while + # encode_layers() correctly returns the pre-norm hidden state. + layer_outputs = text_encoder.encode_layers(tokens, tokens_mask) + if tokens.shape[1] > PROMPT_TEMPLATE_CROP_START: + text_encoder_output = [feat[:, PROMPT_TEMPLATE_CROP_START:, :].contiguous() for feat in layer_outputs] + tokens_mask = tokens_mask[:, PROMPT_TEMPLATE_CROP_START:] + else: + #TODO can this ever happen? max_length=PROMPT_MAX_LENGTH+PROMPT_TEMPLATE_CROP_START, + #so the tokenizer output should always be longer than PROMPT_TEMPLATE_CROP_START. + #upstream pipeline has the same guard, presumably as a safety net. + zero_shape = (tokens.shape[0], 0, layer_outputs[0].shape[-1]) + text_encoder_output = [layer_outputs[0].new_zeros(zero_shape) for _ in layer_outputs] + tokens_mask = torch.zeros((tokens.shape[0], 0), dtype=torch.bool, device=tokens.device) + + elif text_encoder_output is not None: + # Cached: EncodeLensText concatenates layers along dim=-1; split back into per-layer list. + hidden_dim = self.text_encoder_hidden_size + text_encoder_output = list(text_encoder_output.split(hidden_dim, dim=-1)) + + if text_encoder_dropout_probability is not None and text_encoder_dropout_probability > 0.0: + raise NotImplementedError # https://github.com/Nerogar/OneTrainer/issues/957 + + #prune tokens that are masked in all batch samples: + seq_lengths = tokens_mask.sum(dim=1) + max_seq_length = int(seq_lengths.max().item()) + + #pad to 16 because attention processors and/or torch.compile can have issues with uneven sequence lengths, but only pad if an attention mask has to be used anyway: + if max_seq_length % 16 > 0 and (seq_lengths != max_seq_length).any(): + max_seq_length += (16 - max_seq_length % 16) + + text_encoder_output = [feat[:, :max_seq_length, :] for feat in text_encoder_output] + tokens_mask = tokens_mask[:, :max_seq_length].bool() + + return text_encoder_output, tokens_mask + + def calculate_timestep_shift(self, latent_height: int, latent_width: int) -> float: + base_seq_len = self.noise_scheduler.config.base_image_seq_len + max_seq_len = self.noise_scheduler.config.max_image_seq_len + base_shift = self.noise_scheduler.config.base_shift + max_shift = self.noise_scheduler.config.max_shift + patch_size = 2 + + image_seq_len = (latent_width // patch_size) * (latent_height // patch_size) + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return math.exp(mu) + + #packing and unpacking on patchified latents + @staticmethod + def pack_latents(latents) -> Tensor: + batch_size, num_channels, height, width = latents.shape + return latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1) + + @staticmethod + def unpack_latents(latents, height: int, width: int) -> Tensor: + batch_size, seq_len, num_channels = latents.shape + return latents.reshape(batch_size, height, width, num_channels).permute(0, 3, 1, 2) + + @staticmethod + def patchify_latents(latents: torch.Tensor) -> torch.Tensor: + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2) + return latents + + @staticmethod + def unpatchify_latents(latents: torch.Tensor) -> torch.Tensor: + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2) + return latents + + #scaling on patchified latents + def scale_latents(self, latents: Tensor) -> Tensor: + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + return (latents - latents_bn_mean) / latents_bn_std + + def unscale_latents(self, latents: Tensor) -> Tensor: + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + return latents * latents_bn_std + latents_bn_mean diff --git a/modules/modelLoader/Flux2ModelLoader.py b/modules/modelLoader/Flux2ModelLoader.py index 670bacab3..1b621422d 100644 --- a/modules/modelLoader/Flux2ModelLoader.py +++ b/modules/modelLoader/Flux2ModelLoader.py @@ -202,8 +202,7 @@ def __init__(self): super().__init__() def _get_convert_key_sets(self, model: BaseModel) -> list[LoraConversionKeySet] | None: - return None #TODO - #return convert_flux_lora_key_sets() + return None def load( self, diff --git a/modules/modelLoader/LensModelLoader.py b/modules/modelLoader/LensModelLoader.py new file mode 100644 index 000000000..d1dba3561 --- /dev/null +++ b/modules/modelLoader/LensModelLoader.py @@ -0,0 +1,212 @@ +import os +import traceback + +from modules.model.BaseModel import BaseModel +from modules.model.LensModel import LensModel +from modules.modelLoader.GenericFineTuneModelLoader import make_fine_tune_model_loader +from modules.modelLoader.GenericLoRAModelLoader import make_lora_model_loader +from modules.modelLoader.mixin.HFModelLoaderMixin import HFModelLoaderMixin +from modules.modelLoader.mixin.LoRALoaderMixin import LoRALoaderMixin +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.convert.lora.convert_lora_util import LoraConversionKeySet +from modules.util.enum.ModelType import ModelType +from modules.util.ModelNames import ModelNames +from modules.util.ModelWeightDtypes import ModelWeightDtypes +from modules.util.OnDemandModule import OnDemandModule + +import torch + +from diffusers import ( + AutoencoderKLFlux2, + FlowMatchEulerDiscreteScheduler, + GGUFQuantizationConfig, +) +from transformers import PreTrainedTokenizerFast + +from lens.text_encoder import LensGptOssEncoder +from lens.transformer import LensTransformer2DModel + + +class LensModelLoader( + HFModelLoaderMixin, +): + def __init__(self): + super().__init__() + + def __load_internal( + self, + model: LensModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + transformer_model_name: str, + vae_model_name: str, + quantization: QuantizationConfig, + text_encoder_on_demand: bool, + ): + if os.path.isfile(os.path.join(base_model_name, "meta.json")): + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name, quantization, text_encoder_on_demand, + ) + else: + raise Exception("not an internal model") + + def __load_diffusers( + self, + model: LensModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + transformer_model_name: str, + vae_model_name: str, + quantization: QuantizationConfig, + text_encoder_on_demand: bool, + ): + if transformer_model_name: + transformer = LensTransformer2DModel.from_single_file( + transformer_model_name, + config=base_model_name, + subfolder="transformer", + #avoid loading the transformer in float32: + torch_dtype=torch.bfloat16 if weight_dtypes.transformer.torch_dtype() is None else weight_dtypes.transformer.torch_dtype(), + quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16) if weight_dtypes.transformer.is_gguf() else None, + ) + transformer = self._convert_diffusers_sub_module_to_dtype( + transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization, + ) + else: + transformer = self._load_diffusers_sub_module( + LensTransformer2DModel, + weight_dtypes.transformer, + weight_dtypes.train_dtype, + base_model_name, + "transformer", + quantization, + ) + + #TODO verify whether the TokenizersBackend warning actually appears for Lens; if so, uncomment the log suppression below (see ErnieModelLoader for the pattern) + #tokenization_logger = logging.getLogger("transformers.tokenization_utils_base") + #prev_level = tokenization_logger.level + #tokenization_logger.setLevel(logging.ERROR) + tokenizer = PreTrainedTokenizerFast.from_pretrained( + base_model_name, + subfolder="tokenizer", + ) + #tokenization_logger.setLevel(prev_level) + + selected_layer_index = transformer.config.selected_layer_index + + def load_text_encoder(): + text_encoder = LensGptOssEncoder.from_pretrained( + base_model_name, + subfolder="text_encoder", + ) + # set_selected_layers must be called before encode_layers(); the upstream does this in + # LensPipeline.__init__ — we do it here since OneTrainer loads components separately. + text_encoder.set_selected_layers(selected_layer_index) + return text_encoder + + # Lens always loads on demand: its MXFP4 encoder cannot be parked on the CPU temp device, so it + # is built straight onto the accelerator when needed and discarded afterwards (see + # TrainConfig.text_encoder_on_demand / LensModel.materialize_text_encoder). + text_encoder = OnDemandModule(load_text_encoder) if text_encoder_on_demand else load_text_encoder() + + noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + base_model_name, + subfolder="scheduler", + ) + + if vae_model_name: + vae = self._load_diffusers_sub_module( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + vae_model_name, + ) + else: + vae = self._load_diffusers_sub_module( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + "vae", + ) + + model.model_type = model_type + model.tokenizer = tokenizer + model.noise_scheduler = noise_scheduler + model.text_encoder = text_encoder + # read hidden_size from the config only (no weights), so an on-demand encoder stays unmaterialized + model.text_encoder_hidden_size = LensGptOssEncoder.config_class.from_pretrained( + base_model_name, subfolder="text_encoder", + ).hidden_size + model.vae = vae + model.transformer = transformer + + def load( + self, + model: LensModel, + model_type: ModelType, + model_names: ModelNames, + weight_dtypes: ModelWeightDtypes, + quantization: QuantizationConfig, + ): + stacktraces = [] + + try: + self.__load_internal( + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, model_names.text_encoder_on_demand, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + try: + self.__load_diffusers( + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, model_names.text_encoder_on_demand, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + for stacktrace in stacktraces: + print(stacktrace) + raise Exception("could not load model: " + model_names.base_model) + + + +class LensLoRALoader( + LoRALoaderMixin +): + def __init__(self): + super().__init__() + + def _get_convert_key_sets(self, model: BaseModel) -> list[LoraConversionKeySet] | None: + return None #TODO + + def load( + self, + model: LensModel, + model_names: ModelNames, + ): + return self._load(model, model_names) + + +LensLoRAModelLoader = make_lora_model_loader( + model_spec_map={ + ModelType.LENS: "resources/sd_model_spec/lens-lora.json", + }, + model_class=LensModel, + model_loader_class=LensModelLoader, + lora_loader_class=LensLoRALoader, + embedding_loader_class=None, +) + +LensFineTuneModelLoader = make_fine_tune_model_loader( + model_spec_map={ + ModelType.LENS: "resources/sd_model_spec/lens.json", + }, + model_class=LensModel, + model_loader_class=LensModelLoader, + embedding_loader_class=None, +) diff --git a/modules/modelSampler/LensSampler.py b/modules/modelSampler/LensSampler.py new file mode 100644 index 000000000..22751d199 --- /dev/null +++ b/modules/modelSampler/LensSampler.py @@ -0,0 +1,189 @@ +import copy +import inspect +from collections.abc import Callable + +from modules.model.LensModel import LensModel +from modules.modelSampler.BaseModelSampler import BaseModelSampler, ModelSamplerOutput +from modules.util import factory +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.AudioFormat import AudioFormat +from modules.util.enum.FileType import FileType +from modules.util.enum.ImageFormat import ImageFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.NoiseScheduler import NoiseScheduler +from modules.util.enum.VideoFormat import VideoFormat +from modules.util.torch_util import torch_gc + +import torch + +import numpy as np +from lens.pipeline import LensPipeline, compute_empirical_mu +from tqdm import tqdm + + +class LensSampler(BaseModelSampler): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + model: LensModel, + model_type: ModelType, + ): + super().__init__(train_device, temp_device) + + self.model = model + self.model_type = model_type + + @torch.no_grad() + def __sample_base( + self, + prompt: str, + negative_prompt: str, + height: int, + width: int, + seed: int, + random_seed: bool, + diffusion_steps: int, + cfg_scale: float, + noise_scheduler: NoiseScheduler, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ) -> ModelSamplerOutput: + with self.model.autocast_context: + generator = torch.Generator(device=self.train_device) + if random_seed: + generator.seed() + else: + generator.manual_seed(seed) + + noise_scheduler = copy.deepcopy(self.model.noise_scheduler) + transformer = self.model.transformer + vae = self.model.vae + + vae_scale_factor = 8 + num_latent_channels = 32 + patch_size = 2 + + # prepare prompt + self.model.materialize_text_encoder(self.train_device) + + batch_size = 2 if cfg_scale > 1.0 else 1 + (prompt_features, prompt_mask) = self.model.encode_text( + text=[prompt, negative_prompt] if batch_size == 2 else [prompt], + train_device=self.train_device, + ) + + self.model.release_text_encoder() + torch_gc() + + # prepare latent image + latent_image = torch.randn( + size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), + generator=generator, + device=self.train_device, + dtype=torch.float32, + ) + + latent_image = self.model.patchify_latents(latent_image) + + latent_image = self.model.pack_latents(latent_image) + image_seq_len = latent_image.shape[1] + mu = compute_empirical_mu(image_seq_len, diffusion_steps) + + # prepare timesteps + #TODO for other models, too? This is different than with sigmas=None + sigmas = np.linspace(1.0, 1 / diffusion_steps, diffusion_steps) + noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device, mu=mu, sigmas=sigmas) + timesteps = noise_scheduler.timesteps + + # denoising loop + extra_step_kwargs = {} #TODO remove + if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): + extra_step_kwargs["generator"] = generator + + # img_shapes: (frame, h_lat, w_lat) for RoPE + img_shapes = [(1, height // vae_scale_factor // patch_size, width // vae_scale_factor // patch_size)] + + self.model.transformer_to(self.train_device) + for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): + latent_model_input = torch.cat([latent_image] * batch_size) + expanded_timestep = timestep.expand(latent_model_input.shape[0]) + + + noise_pred = transformer( + hidden_states=latent_model_input.to(dtype=self.model.train_dtype.torch_dtype()), + encoder_hidden_states=[f.to(dtype=self.model.train_dtype.torch_dtype()) for f in prompt_features], # list of 4 per-GPT-OSS-layer tensors, each [B, S_txt, dim] + encoder_hidden_states_mask=prompt_mask, + timestep=expanded_timestep / 1000, + img_shapes=img_shapes, + ) + + if batch_size == 2: + cond, uncond = noise_pred.chunk(2) + # Norm-rescaled CFG: rescale the combined prediction to match ||cond||, + # preventing magnitude blowup at high guidance scales. Used by the upstream + # pipeline even though not documented in the Lens paper (arxiv:2605.21573). + comb = uncond + cfg_scale * (cond - uncond) + cond_norm = torch.norm(cond, dim=-1, keepdim=True) + comb_norm = torch.norm(comb, dim=-1, keepdim=True) + scale = torch.where(comb_norm > 0, cond_norm / comb_norm.clamp_min(1e-12), torch.ones_like(comb_norm)) + noise_pred = comb * scale + + latent_image = noise_scheduler.step(noise_pred, timestep, latent_image, return_dict=False, **extra_step_kwargs)[0] + + on_update_progress(i + 1, len(timesteps)) + + self.model.transformer_to(self.temp_device) + torch_gc() + self.model.vae_to(self.train_device) + + latent_image = self.model.unpack_latents( + latent_image, + height // vae_scale_factor // patch_size, + width // vae_scale_factor // patch_size, + ) + latents = self.model.unscale_latents(latent_image) + latents = self.model.unpatchify_latents(latents) + + decoded = vae.decode(latents, return_dict=False)[0] + + image = LensPipeline._to_pil(decoded) + + self.model.vae_to(self.temp_device) + torch_gc() + + return ModelSamplerOutput( + file_type=FileType.IMAGE, + data=image[0], + ) + + def sample( + self, + sample_config: SampleConfig, + destination: str, + image_format: ImageFormat | None = None, + video_format: VideoFormat | None = None, + audio_format: AudioFormat | None = None, + on_sample: Callable[[ModelSamplerOutput], None] = lambda _: None, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ): + sampler_output = self.__sample_base( + prompt=sample_config.prompt, + negative_prompt=sample_config.negative_prompt, + height=self.quantize_resolution(sample_config.height, 64), + width=self.quantize_resolution(sample_config.width, 64), + seed=sample_config.seed, + random_seed=sample_config.random_seed, + diffusion_steps=sample_config.diffusion_steps, + cfg_scale=sample_config.cfg_scale, + noise_scheduler=sample_config.noise_scheduler, + on_update_progress=on_update_progress, + ) + + self.save_sampler_output( + sampler_output, destination, + image_format, video_format, audio_format, + ) + + on_sample(sampler_output) + +factory.register(BaseModelSampler, LensSampler, ModelType.LENS) diff --git a/modules/modelSaver/LensFineTuneModelSaver.py b/modules/modelSaver/LensFineTuneModelSaver.py new file mode 100644 index 000000000..300c18f3a --- /dev/null +++ b/modules/modelSaver/LensFineTuneModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.LensModel import LensModel +from modules.modelSaver.GenericFineTuneModelSaver import make_fine_tune_model_saver +from modules.modelSaver.lens.LensModelSaver import LensModelSaver +from modules.util.enum.ModelType import ModelType + +LensFineTuneModelSaver = make_fine_tune_model_saver( + ModelType.LENS, + model_class=LensModel, + model_saver_class=LensModelSaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/LensLoRAModelSaver.py b/modules/modelSaver/LensLoRAModelSaver.py new file mode 100644 index 000000000..54656d48c --- /dev/null +++ b/modules/modelSaver/LensLoRAModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.LensModel import LensModel +from modules.modelSaver.GenericLoRAModelSaver import make_lora_model_saver +from modules.modelSaver.lens.LensLoRASaver import LensLoRASaver +from modules.util.enum.ModelType import ModelType + +LensLoRAModelSaver = make_lora_model_saver( + ModelType.LENS, + model_class=LensModel, + lora_saver_class=LensLoRASaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/lens/LensLoRASaver.py b/modules/modelSaver/lens/LensLoRASaver.py new file mode 100644 index 000000000..97846bed3 --- /dev/null +++ b/modules/modelSaver/lens/LensLoRASaver.py @@ -0,0 +1,38 @@ +from modules.model.LensModel import LensModel +from modules.modelSaver.mixin.LoRASaverMixin import LoRASaverMixin +from modules.util.convert.lora.convert_lora_util import LoraConversionKeySet +from modules.util.enum.ModelFormat import ModelFormat + +import torch +from torch import Tensor + + +class LensLoRASaver( + LoRASaverMixin, +): + def __init__(self): + super().__init__() + + def _get_convert_key_sets(self, model: LensModel) -> list[LoraConversionKeySet] | None: + return None + + def _get_state_dict( + self, + model: LensModel, + ) -> dict[str, Tensor]: + state_dict = {} + if model.transformer_lora is not None: + state_dict |= model.transformer_lora.state_dict() + if model.lora_state_dict is not None: + state_dict |= model.lora_state_dict + + return state_dict + + def save( + self, + model: LensModel, + output_model_format: ModelFormat, + output_model_destination: str, + dtype: torch.dtype | None, + ): + self._save(model, output_model_format, output_model_destination, dtype) diff --git a/modules/modelSaver/lens/LensModelSaver.py b/modules/modelSaver/lens/LensModelSaver.py new file mode 100644 index 000000000..26bfdabd7 --- /dev/null +++ b/modules/modelSaver/lens/LensModelSaver.py @@ -0,0 +1,87 @@ +import copy +import os.path +from pathlib import Path + +from modules.model.LensModel import LensModel +from modules.modelSaver.mixin.DtypeModelSaverMixin import DtypeModelSaverMixin +from modules.util.enum.ModelFormat import ModelFormat + +import torch + +from safetensors.torch import save_file + + +class LensModelSaver( + DtypeModelSaverMixin, +): + def __init__(self): + super().__init__() + + def __save_diffusers( + self, + model: LensModel, + destination: str, + dtype: torch.dtype | None, + ): + # this is the only save path that embeds the base text encoder, so the on-demand encoder is + # loaded here (and freed again afterwards). resident encoders are untouched. + model.materialize_text_encoder_for_save() + try: + # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. + pipeline = model.create_pipeline() + pipeline.to("cpu") + if dtype is not None: #TODO necessary? + # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. + # the tokenizer tries to reload from the file system otherwise + tokenizer = pipeline.tokenizer + tokenizer.__deepcopy__ = lambda memo: tokenizer + + save_pipeline = copy.deepcopy(pipeline) + save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) + + delattr(tokenizer, '__deepcopy__') + else: + save_pipeline = pipeline + + os.makedirs(Path(destination).absolute(), exist_ok=True) + save_pipeline.save_pretrained(destination) + + if dtype is not None: + del save_pipeline + finally: + model.release_text_encoder() + + def __save_safetensors( + self, + model: LensModel, + destination: str, + dtype: torch.dtype | None, + ): + # Lens transformer uses diffusers-format keys; no key conversion needed. + state_dict = model.transformer.state_dict() + save_state_dict = self._convert_state_dict_dtype(state_dict, dtype) + self._convert_state_dict_to_contiguous(save_state_dict) + os.makedirs(Path(destination).parent.absolute(), exist_ok=True) + save_file(save_state_dict, destination, self._create_safetensors_header(model, save_state_dict)) + + def __save_internal( + self, + model: LensModel, + destination: str, + ): + self.__save_diffusers(model, destination, None) + + def save( + self, + model: LensModel, + output_model_format: ModelFormat, + output_model_destination: str, + dtype: torch.dtype | None, + ): + match output_model_format: + case ModelFormat.DIFFUSERS: + self.__save_diffusers(model, output_model_destination, dtype) + case ModelFormat.SAFETENSORS: + self.__save_safetensors(model, output_model_destination, dtype) + case ModelFormat.INTERNAL: + self.__save_internal(model, output_model_destination) diff --git a/modules/modelSetup/BaseLensSetup.py b/modules/modelSetup/BaseLensSetup.py new file mode 100644 index 000000000..de868cbc3 --- /dev/null +++ b/modules/modelSetup/BaseLensSetup.py @@ -0,0 +1,180 @@ +from abc import ABCMeta +from random import Random + +import modules.util.multi_gpu_util as multi +from modules.model.LensModel import LensModel +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.modelSetup.mixin.ModelSetupDebugMixin import ModelSetupDebugMixin +from modules.modelSetup.mixin.ModelSetupDiffusionLossMixin import ModelSetupDiffusionLossMixin +from modules.modelSetup.mixin.ModelSetupEmbeddingMixin import ModelSetupEmbeddingMixin +from modules.modelSetup.mixin.ModelSetupFlowMatchingMixin import ModelSetupFlowMatchingMixin +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin +from modules.util.checkpointing_util import ( + enable_checkpointing_for_lens_transformer, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.quantization_util import quantize_layers +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress + +import torch +from torch import Tensor + + +class BaseLensSetup( + BaseModelSetup, + ModelSetupDiffusionLossMixin, + ModelSetupDebugMixin, + ModelSetupNoiseMixin, + ModelSetupFlowMatchingMixin, + ModelSetupEmbeddingMixin, + metaclass=ABCMeta +): + LAYER_PRESETS = { + "attn-mlp": ["attn", "mlp"], + "attn-only": ["attn"], + "blocks": ["transformer_block"], + "full": [], + } + + def setup_optimizations( + self, + model: LensModel, + config: TrainConfig, + ): + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = \ + enable_checkpointing_for_lens_transformer(model.transformer, config, config.transformer) + + model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ + config.weight_dtypes().transformer, + config.weight_dtypes().text_encoder, + config.weight_dtypes().vae, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], config.enable_autocast_cache) + + model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ + disable_fp16_autocast_context( + self.train_device, + config.train_dtype, + config.fallback_train_dtype, + [ + config.weight_dtypes().text_encoder, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], + config.enable_autocast_cache, + ) + + quantize_layers(model.vae, self.train_device, model.train_dtype, config) + quantize_layers(model.transformer, self.train_device, model.train_dtype, config) + #model.text_encoder is not quantized - GPT OSS is loaded in MXFP4 + + def predict( + self, + model: LensModel, + batch: dict, + config: TrainConfig, + train_progress: TrainProgress, + *, + deterministic: bool = False, + ) -> dict: + with model.autocast_context: + batch_seed = 0 if deterministic else train_progress.global_step * multi.world_size() + multi.rank() + generator = torch.Generator(device=config.train_device) + generator.manual_seed(batch_seed) + rand = Random(batch_seed) + + text_encoder_output, tokens_mask = model.encode_text( + train_device=self.train_device, + batch_size=batch['latent_image'].shape[0], + rand=rand, + tokens=batch.get("tokens"), + tokens_mask=batch.get("tokens_mask"), + text_encoder_output=batch.get('text_encoder_hidden_state'), + text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, + ) + latent_image = model.patchify_latents(batch['latent_image'].float()) + latent_height = latent_image.shape[-2] + latent_width = latent_image.shape[-1] + scaled_latent_image = model.scale_latents(latent_image) + + latent_noise = self._create_noise(scaled_latent_image, config, generator) + + shift = model.calculate_timestep_shift(latent_height, latent_width) + timestep = self._get_timestep_discrete( + model.noise_scheduler.config['num_train_timesteps'], + deterministic, + generator, + scaled_latent_image.shape[0], + config, + shift = shift if config.dynamic_timestep_shifting else config.timestep_shift, + ) + + scaled_noisy_latent_image, sigma = self._add_noise_discrete( + scaled_latent_image, + latent_noise, + timestep, + model.noise_scheduler.timesteps, + ) + + img_shapes = [(1, latent_height, latent_width)] + packed_latent_input = model.pack_latents(scaled_noisy_latent_image) + + packed_predicted_flow = model.transformer( + hidden_states=packed_latent_input.to(dtype=model.train_dtype.torch_dtype()), + encoder_hidden_states=[f.to(dtype=model.train_dtype.torch_dtype()) for f in text_encoder_output], + encoder_hidden_states_mask=tokens_mask, + timestep=timestep / 1000, + img_shapes=img_shapes, + ) + + predicted_flow = model.unpack_latents( + packed_predicted_flow, + latent_height, + latent_width, + ) + + flow = latent_noise - scaled_latent_image + model_output_data = { + 'loss_type': 'target', + 'timestep': timestep, + #unpatchify, to make the shape match the mask shape of masked training: + 'predicted': model.unpatchify_latents(predicted_flow), + 'target': model.unpatchify_latents(flow), + } + + if config.debug_mode: + with torch.no_grad(): + predicted_scaled_latent_image = scaled_noisy_latent_image - predicted_flow * sigma + self._save_tokens("7-prompt", batch['tokens'], model.tokenizer, config, train_progress) + self._save_latent("1-noise", latent_noise, config, train_progress) + self._save_latent("2-noisy_image", scaled_noisy_latent_image, config, train_progress) + self._save_latent("3-predicted_flow", predicted_flow, config, train_progress) + self._save_latent("4-flow", flow, config, train_progress) + self._save_latent("5-predicted_image", predicted_scaled_latent_image, config, train_progress) + self._save_latent("6-image", scaled_latent_image, config, train_progress) + + return model_output_data + + def calculate_loss( + self, + model: LensModel, + batch: dict, + data: dict, + config: TrainConfig, + ) -> Tensor: + return self._flow_matching_losses( + batch=batch, + data=data, + config=config, + train_device=self.train_device, + sigmas=model.noise_scheduler.sigmas, + ).mean() + + def prepare_text_caching(self, model: LensModel, config: TrainConfig): + model.release() + model.materialize_text_encoder(self.train_device) + model.eval() + torch_gc() diff --git a/modules/modelSetup/LensFineTuneSetup.py b/modules/modelSetup/LensFineTuneSetup.py new file mode 100644 index 000000000..01ac4efdb --- /dev/null +++ b/modules/modelSetup/LensFineTuneSetup.py @@ -0,0 +1,96 @@ +from modules.model.LensModel import LensModel +from modules.modelSetup.BaseLensSetup import BaseLensSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModuleFilter import ModuleFilter +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +class LensFineTuneSetup( + BaseLensSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: LensModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + + self._create_model_part_parameters(parameter_group_collection, "transformer", model.transformer, config.transformer, + freeze=ModuleFilter.create(config), debug=config.debug_mode) + return parameter_group_collection + + def __setup_requires_grad( + self, + model: LensModel, + config: TrainConfig, + ): + self._setup_model_part_requires_grad("transformer", model.transformer, config.transformer, model.train_progress) + model.vae.requires_grad_(False) + if model.text_encoder is not None: + model.text_encoder.requires_grad_(False) + + + def setup_model( + self, + model: LensModel, + config: TrainConfig, + ): + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: LensModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + text_encoder_on_train_device = not config.latent_caching + + # the encoder is always on-demand for Lens, so materialize it when training needs it and + # discard it (freeing VRAM) otherwise -- text_encoder_to would be a no-op on the proxy. + if text_encoder_on_train_device: + model.materialize_text_encoder(self.train_device) + else: + model.release_text_encoder() + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + + if model.text_encoder is not None: + model.text_encoder.eval() + model.vae.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: LensModel, + config: TrainConfig, + train_progress: TrainProgress + ): + self.__setup_requires_grad(model, config) + +factory.register(BaseModelSetup, LensFineTuneSetup, ModelType.LENS, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/LensLoRASetup.py b/modules/modelSetup/LensLoRASetup.py new file mode 100644 index 000000000..c481e72a6 --- /dev/null +++ b/modules/modelSetup/LensLoRASetup.py @@ -0,0 +1,108 @@ +from modules.model.LensModel import LensModel +from modules.modelSetup.BaseLensSetup import BaseLensSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +class LensLoRASetup( + BaseLensSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: LensModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + + self._create_model_part_parameters(parameter_group_collection, "transformer", model.transformer_lora, config.transformer) + return parameter_group_collection + + def __setup_requires_grad( + self, + model: LensModel, + config: TrainConfig, + ): + if model.text_encoder is not None: + model.text_encoder.requires_grad_(False) + model.transformer.requires_grad_(False) + model.vae.requires_grad_(False) + + self._setup_model_part_requires_grad("transformer", model.transformer_lora, config.transformer, model.train_progress) + + def setup_model( + self, + model: LensModel, + config: TrainConfig, + ): + model.transformer_lora = LoRAModuleWrapper( + model.transformer, "transformer", config, config.layer_filter.split(",") + ) + + if model.lora_state_dict: + model.transformer_lora.load_state_dict(model.lora_state_dict) + model.lora_state_dict = None + + model.transformer_lora.set_dropout(config.dropout_probability) + model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) + model.transformer_lora.hook_to_module() + + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: LensModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + text_encoder_on_train_device = not config.latent_caching + + # the encoder is always on-demand for Lens, so materialize it when training needs it and + # discard it (freeing VRAM) otherwise -- text_encoder_to would be a no-op on the proxy. + if text_encoder_on_train_device: + model.materialize_text_encoder(self.train_device) + else: + model.release_text_encoder() + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + + if model.text_encoder is not None: + model.text_encoder.eval() + model.vae.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: LensModel, + config: TrainConfig, + train_progress: TrainProgress + ): + self.__setup_requires_grad(model, config) + +factory.register(BaseModelSetup, LensLoRASetup, ModelType.LENS, TrainingMethod.LORA) diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/ConvertModelUI.py index 6cb1b507a..a32987de4 100644 --- a/modules/ui/ConvertModelUI.py +++ b/modules/ui/ConvertModelUI.py @@ -67,6 +67,7 @@ def main_frame(self, master): ("Flux Dev", ModelType.FLUX_DEV_1), ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), ("Flux 2", ModelType.FLUX_2), + ("Lens", ModelType.LENS), ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index 5bc3092d3..73ce614a3 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -61,6 +61,7 @@ def __setup_ui(self, frame): or model_type.is_chroma() or model_type.is_qwen() or model_type.is_hunyuan_video() + or model_type.is_lens() ) row = 0 @@ -73,7 +74,8 @@ def __setup_ui(self, frame): allow_override_prior=model_type.is_stable_cascade(), has_transformer="transformer" in parts, allow_override_transformer=allow_override_transformer, - has_text_encoder=not model_type.has_multiple_text_encoders(), + # Lens: GPT-OSS dtype is fixed by MXFP4 + checkpoint; not user-configurable. + has_text_encoder=not model_type.has_multiple_text_encoders() and not model_type.is_lens(), has_text_encoder_1=model_type.has_multiple_text_encoders(), has_text_encoder_2="text_encoder_2" in parts, has_text_encoder_3="text_encoder_3" in parts, diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py index d2ca78f7b..e416ca9dd 100644 --- a/modules/ui/TopBar.py +++ b/modules/ui/TopBar.py @@ -95,6 +95,7 @@ def __init__( ("Flux Dev.1", ModelType.FLUX_DEV_1), ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Lens", ModelType.LENS), ("Sana", ModelType.SANA), ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), ("HiDream Full", ModelType.HI_DREAM_FULL), diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index 2d24d968c..2d0724e5d 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -73,6 +73,8 @@ def refresh_ui(self): self.__setup_flux_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_flux_2(): self.__setup_flux_2_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_lens(): + self.__setup_lens_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_chroma(): self.__setup_chroma_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_qwen(): @@ -185,6 +187,19 @@ def __setup_flux_2_ui(self, column_0, column_1, column_2): self.__create_loss_frame(column_2, 2) self.__create_layer_frame(column_2, 3) + def __setup_lens_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + # the GPT-OSS encoder is always on-demand (MXFP4); layer offloading is incompatible with that, so hide it + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=False, supports_offloading=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + def __setup_chroma_ui(self, column_0, column_1, column_2): self.__create_base_frame(column_0, 0) self.__create_text_encoder_frame(column_0, 1) @@ -422,17 +437,19 @@ def __create_base2_frame(self, master, row, video_training_enabled: bool=False, tooltip="Enables circular padding for all conv layers to better train seamless images") components.switch(frame, row, 1, self.ui_state, "force_circular_padding") - def __create_offloading_widgets(self, frame, row, part, supports_checkpointing=True, supports_activation_offloading=False): + def __create_offloading_widgets(self, frame, row, part, supports_checkpointing=True, supports_activation_offloading=False, supports_layer_offloading=True): + # per-component offloading / checkpointing controls (bound to the model part config) if supports_checkpointing: components.label(frame, row, 0, "Gradient Checkpointing", tooltip="Enables gradient checkpointing for this component. Reduces VRAM usage at the cost of training speed") components.switch(frame, row, 1, self.ui_state, f"{part}.gradient_checkpointing") row += 1 - components.label(frame, row, 0, "Layer Offload Fraction", - tooltip="Fraction of this component's layers to offload to CPU to reduce VRAM usage. Increases training time and RAM usage. 0=disabled, 1=all layers") - components.entry(frame, row, 1, self.ui_state, f"{part}.offload_fraction") - row += 1 + if supports_layer_offloading: + components.label(frame, row, 0, "Layer Offload Fraction", + tooltip="Fraction of this component's layers to offload to CPU to reduce VRAM usage. Increases training time and RAM usage. 0=disabled, 1=all layers") + components.entry(frame, row, 1, self.ui_state, f"{part}.offload_fraction") + row += 1 if supports_activation_offloading: components.label(frame, row, 0, "Offload Activations", @@ -442,7 +459,7 @@ def __create_offloading_widgets(self, frame, row, part, supports_checkpointing=T return row - def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): + def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False, supports_offloading=True): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) @@ -458,7 +475,7 @@ def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supp components.label(frame, row, 0, "Text Encoder") row += 1 - row = self.__create_offloading_widgets(frame, row, "text_encoder", supports_checkpointing=supports_training) + row = self.__create_offloading_widgets(frame, row, "text_encoder", supports_checkpointing=supports_training, supports_layer_offloading=supports_offloading) # dropout components.label(frame, row, 0, "Caption Dropout Probability", diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index d5b9f765c..ae81e0d58 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -24,6 +24,7 @@ from diffusers.models.unets.unet_stable_cascade import SDCascadeAttnBlock, SDCascadeResBlock, SDCascadeTimestepBlock from transformers.models.clip.modeling_clip import CLIPEncoderLayer from transformers.models.gemma2.modeling_gemma2 import Gemma2DecoderLayer +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssDecoderLayer from transformers.models.llama.modeling_llama import LlamaDecoderLayer from transformers.models.mistral.modeling_mistral import MistralDecoderLayer from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLDecoderLayer @@ -455,3 +456,23 @@ def enable_checkpointing_for_ernie_transformer( return enable_checkpointing(model, config, part, config.compile, [ (model.layers, ["x"]), ]) + + +def enable_checkpointing_for_lens_transformer( + model: nn.Module, + config: TrainConfig, + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ + (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), + ]) + + +def enable_checkpointing_for_gpt_oss_encoder_layers( + model: nn.Module, + config: TrainConfig, + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (GptOssDecoderLayer, []), # No activation offloading: hidden states are taken from intermediate layers by encode_layers() + ]) diff --git a/modules/util/config/SampleConfig.py b/modules/util/config/SampleConfig.py index 38e99e182..9e079f4e7 100644 --- a/modules/util/config/SampleConfig.py +++ b/modules/util/config/SampleConfig.py @@ -79,6 +79,13 @@ def _get_model_defaults(model_type) -> dict: "diffusion_steps": 30, "cfg_scale": 4.0, }) + elif model_type.is_lens(): + defaults.update({ + "width": 1024, + "height": 1024, + "diffusion_steps": 30, + "cfg_scale": 5.0, + }) elif model_type.is_qwen(): defaults.update({ "width": 1024, diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 68281b8dd..d9e801a6f 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -936,6 +936,8 @@ def train_text_encoder_4_or_embedding(self) -> bool: #an encoder is loaded on demand only when it is requested, frozen (not trained, no embedding #training) and its conditioning is cached -- otherwise it is needed resident every step. def text_encoder_on_demand(self) -> bool: + if self.model_type.is_lens(): + return True return self.text_encoder.load_on_demand \ and not self.train_text_encoder_or_embedding() \ and self.latent_caching diff --git a/modules/util/enum/ModelType.py b/modules/util/enum/ModelType.py index a6e6a2a96..8f1e61bde 100644 --- a/modules/util/enum/ModelType.py +++ b/modules/util/enum/ModelType.py @@ -43,6 +43,8 @@ class ModelType(Enum): ERNIE = 'ERNIE' + LENS = 'LENS' + def __str__(self): return self.value @@ -114,6 +116,9 @@ def is_z_image(self): def is_ernie(self): return self == ModelType.ERNIE + def is_lens(self): + return self == ModelType.LENS + def has_mask_input(self) -> bool: return self == ModelType.STABLE_DIFFUSION_15_INPAINTING \ or self == ModelType.STABLE_DIFFUSION_20_INPAINTING \ @@ -163,7 +168,8 @@ def is_flow_matching(self) -> bool: or self.is_hunyuan_video() \ or self.is_hi_dream() \ or self.is_z_image() \ - or self.is_ernie() + or self.is_ernie() \ + or self.is_lens() def is_video_model(self) -> bool: return self.is_hunyuan_video() #incase we add more video models in the future @@ -184,7 +190,7 @@ def supported_training_methods(self) -> tuple[TrainingMethod, ...]: or self.is_hi_dream() \ or self.is_chroma(): return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING) - if self.is_qwen() or self.is_z_image() or self.is_flux_2() or self.is_ernie(): + if self.is_qwen() or self.is_z_image() or self.is_flux_2() or self.is_ernie() or self.is_lens(): return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA) raise ValueError(f"No supported training methods defined for model type {self}") @@ -219,6 +225,7 @@ def supported_training_methods(self) -> tuple[TrainingMethod, ...]: ModelType.QWEN: ("text_encoder", "transformer", "vae"), ModelType.Z_IMAGE: ("text_encoder", "transformer", "vae"), ModelType.ERNIE: ("text_encoder", "transformer", "vae"), + ModelType.LENS: ("text_encoder", "transformer", "vae"), } diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index b2630c16b..cd39c06bb 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -35,7 +35,7 @@ def build_muon_adam_key_fn( 'block', # UNet 'text_model.encoder.layers', # TEs (CLIPs) ] - case ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA: + case ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA | ModelType.LENS: default_patterns = [ 'transformer_blocks', 'encoder.block', # TE (T5) diff --git a/requirements-global.txt b/requirements-global.txt index 299f765b0..5e2e5820f 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -28,11 +28,14 @@ omegaconf==2.3.0 # needed to load stable diffusion from single ckpt files invisible-watermark==0.2.0 # needed for the SDXL pipeline # other models +git+https://github.com/dxqb/Lens.git@daba3a8#egg=lens +kernels==0.14.1 # Triton kernel loader for Lens GPT-OSS MXFP4 text encoder +einops==0.8.2 pooch==1.8.2 open-clip-torch==2.32.0 # data loader --e git+https://github.com/Nerogar/mgds.git@9320a69#egg=mgds +-e git+https://github.com/dxqb/mgds.git@f3b1bd1#egg=mgds # optimizers dadaptation==3.2 # dadaptation optimizers diff --git a/resources/sd_model_spec/lens-lora.json b/resources/sd_model_spec/lens-lora.json new file mode 100644 index 000000000..31fede12a --- /dev/null +++ b/resources/sd_model_spec/lens-lora.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Lens/lora", + "modelspec.implementation": "https://github.com/microsoft/Lens", + "modelspec.title": "Lens LoRA" +} diff --git a/resources/sd_model_spec/lens.json b/resources/sd_model_spec/lens.json new file mode 100644 index 000000000..57c2681aa --- /dev/null +++ b/resources/sd_model_spec/lens.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Lens", + "modelspec.implementation": "https://github.com/microsoft/Lens", + "modelspec.title": "Lens" +} diff --git a/training_presets/#lens Finetune 16GB.json b/training_presets/#lens Finetune 16GB.json new file mode 100644 index 000000000..d1dd235ab --- /dev/null +++ b/training_presets/#lens Finetune 16GB.json @@ -0,0 +1,47 @@ +{ + "base_model_name": "microsoft/Lens", + "batch_size": 2, + "learning_rate": 1e-6, + "model_type": "LENS", + "resolution": "512", + "compile": true, + "transformer": { + "train": true, + "weight_dtype": "BFLOAT_16" + }, + "text_encoder": { + "train": false, + "weight_dtype": "BFLOAT_16" + }, + "training_method": "FINE_TUNE", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "quantization": { + "layer_filter": "transformer_block", + "layer_filter_preset": "blocks" + }, + "timestep_distribution": "LOGIT_NORMAL", + "dataloader_threads": 1, + "optimizer": { + "optimizer": "ADAFACTOR" + }, + "optimizer_defaults": { + "ADAFACTOR": { + "optimizer": "ADAFACTOR", + "fused_back_pass": true, + "beta1": null, + "clip_threshold": 1.0, + "decay_rate": -0.8, + "eps": 1e-30, + "eps2": 0.001, + "relative_step": false, + "scale_parameter": false, + "stochastic_rounding": true, + "warmup_init": false, + "weight_decay": 0.0 + } + } +} diff --git a/training_presets/#lens LoRA 16GB.json b/training_presets/#lens LoRA 16GB.json new file mode 100644 index 000000000..5400f22df --- /dev/null +++ b/training_presets/#lens LoRA 16GB.json @@ -0,0 +1,30 @@ +{ + "base_model_name": "microsoft/Lens", + "batch_size": 2, + "learning_rate": 3e-4, + "model_type": "LENS", + "resolution": "512", + "compile": true, + "transformer": { + "train": true, + "weight_dtype": "INT_W8A8" + }, + "text_encoder": { + "train": false, + "weight_dtype": "BFLOAT_16" + }, + "training_method": "LORA", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "layer_filter": "attn,mlp", + "layer_filter_preset": "attn-mlp", + "quantization": { + "layer_filter": "attn,mlp", + "layer_filter_preset": "attn-mlp" + }, + "timestep_distribution": "LOGIT_NORMAL", + "dataloader_threads": 1 +} From 6fc0261a3c87acf8966c37c8fbc10387ae0aa379 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sun, 7 Jun 2026 14:12:52 +0200 Subject: [PATCH 05/10] fix: initialize torch.compile config in autograd worker threads Since torch 2.12, dynamo config overrides are stored in ContextVars (pytorch/pytorch#173568), which threads don't inherit. The recompile limit raised by init_compile() therefore no longer applied to the autograd worker threads, which recompute checkpointed modules during the backward pass - training crashed with FailOnRecompileLimitHit once enough input shapes accumulated, despite the raised limit. Call init_compile() in OffloadCheckpointLayer.__checkpointing_forward so the recomputation thread sets its own config. The non-offload CheckpointLayer is unaffected: it is compiled as a whole, so its recomputation runs inside the compiled backward without re-entering dynamo. The Mod.eval patch is applied once at import time instead. Co-Authored-By: Claude Opus 4.8 --- modules/util/checkpointing_util.py | 3 +++ modules/util/compile_util.py | 14 ++++++++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index ba3f7fec2..a94302f32 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -112,6 +112,9 @@ def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.dev self.layer_index = layer_index def __checkpointing_forward(self, dummy: torch.Tensor, call_id: int, *args): + # during the backward pass, this runs on an autograd worker thread, which needs its own + # dynamo config initialization (see init_compile) + init_compile() if self.layer_index == 0 and not torch.is_grad_enabled(): self.conductor.start_forward(True) diff --git a/modules/util/compile_util.py b/modules/util/compile_util.py index 3fd310cd2..43cacfe9d 100644 --- a/modules/util/compile_util.py +++ b/modules/util/compile_util.py @@ -1,11 +1,15 @@ import torch +import torch.utils._sympy.functions from sympy import S #code from https://github.com/pytorch/pytorch/blob/ed82d5fcfd80110565f69130f286c7bfec6db2dc/torch/utils/_sympy/functions.py#L481 #but accepts negative numbers, to avoid https://github.com/Nerogar/OneTrainer/issues/1126 -#can be removed once https://github.com/pytorch/pytorch/pull/169726 is merged into a torch version we use +#torch 2.12 (https://github.com/pytorch/pytorch/pull/169726) fixed the call path of that issue, but Mod.eval +#itself still raises on negative numbers, and inductor's tiling_utils still calls sympy is_constant() (which +#substitutes negative test values) on expressions that can contain Mod. can be removed once Mod.eval itself +#accepts negative numbers in a torch version we use @classmethod def Mod_patched_eval(cls, p, q): # This was adapted from: sympy/core/mod.py @@ -58,5 +62,11 @@ def Mod_patched_eval(cls, p, q): def init_compile(): # cache_size_limit and recompile_limit are aliases for the same dynamo config value. + # since torch 2.12, dynamo config overrides are stored in ContextVars (pytorch/pytorch#173568), + # which threads don't inherit. init_compile() must therefore be called in every thread that can + # trigger a compilation - including the autograd worker threads, which recompute checkpointed + # modules during the backward pass. torch._dynamo.config.cache_size_limit = 8192 - torch.utils._sympy.functions.Mod.eval = Mod_patched_eval + + +torch.utils._sympy.functions.Mod.eval = Mod_patched_eval From 31285e82d9f40567c5d362b46e6340d0f19707e1 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sun, 7 Jun 2026 14:19:02 +0200 Subject: [PATCH 06/10] Refine comments on Mod.eval negative number handling Updated comments regarding the handling of negative numbers in Mod.eval and related issues. --- modules/util/compile_util.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/modules/util/compile_util.py b/modules/util/compile_util.py index 43cacfe9d..c1a876962 100644 --- a/modules/util/compile_util.py +++ b/modules/util/compile_util.py @@ -6,10 +6,11 @@ #code from https://github.com/pytorch/pytorch/blob/ed82d5fcfd80110565f69130f286c7bfec6db2dc/torch/utils/_sympy/functions.py#L481 #but accepts negative numbers, to avoid https://github.com/Nerogar/OneTrainer/issues/1126 -#torch 2.12 (https://github.com/pytorch/pytorch/pull/169726) fixed the call path of that issue, but Mod.eval -#itself still raises on negative numbers, and inductor's tiling_utils still calls sympy is_constant() (which +#torch 2.12 merged the fix PR https://github.com/pytorch/pytorch/pull/169726, but according to Claude's analysis this bug +#still exists. TODO confirm manually by testing https://github.com/Nerogar/OneTrainer/issues/1126: +#Claude: "Mod.eval itself still raises on negative numbers, and inductor's tiling_utils still calls sympy is_constant() (which #substitutes negative test values) on expressions that can contain Mod. can be removed once Mod.eval itself -#accepts negative numbers in a torch version we use +#accepts negative numbers in a torch version we use" @classmethod def Mod_patched_eval(cls, p, q): # This was adapted from: sympy/core/mod.py From cb79456eb2001c5f1b409d86b0b2bc7a765a653e Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 13 Jun 2026 23:38:35 +0200 Subject: [PATCH 07/10] Centralize checkpointing/offloading gate inside enable_checkpointing Move the per-callsite checkpointing_or_offloading_enabled() guard into enable_checkpointing() itself, so every Base*Setup can call enable_checkpointing_for_* unconditionally. Also extend the central gate to allow a compile-only path (no checkpointing/offloading, but still per-layer torch.compile wrapping) when config.compile is set. Three direct diffusers enable_gradient_checkpointing() calls (SD/SDXL unet, Wuerstchen v2 prior) keep their explicit guard since they bypass this central mechanism. --- modules/modelSetup/BaseChromaSetup.py | 5 ++--- modules/modelSetup/BaseErnieSetup.py | 3 +-- modules/modelSetup/BaseFlux2Setup.py | 5 ++--- modules/modelSetup/BaseFluxSetup.py | 7 +++---- modules/modelSetup/BaseHiDreamSetup.py | 11 +++++------ modules/modelSetup/BaseHunyuanVideoSetup.py | 7 +++---- modules/modelSetup/BasePixArtAlphaSetup.py | 6 ++---- modules/modelSetup/BaseQwenSetup.py | 5 ++--- modules/modelSetup/BaseSanaSetup.py | 6 ++---- modules/modelSetup/BaseStableDiffusion3Setup.py | 9 ++++----- modules/modelSetup/BaseStableDiffusionSetup.py | 3 +-- modules/modelSetup/BaseStableDiffusionXLSetup.py | 6 ++---- modules/modelSetup/BaseWuerstchenSetup.py | 11 +++++------ modules/modelSetup/BaseZImageSetup.py | 5 ++--- modules/util/checkpointing_util.py | 3 +++ 15 files changed, 39 insertions(+), 53 deletions(-) diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index f8cc33187..cf2918bc1 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -49,9 +49,8 @@ def setup_optimizations( model: ChromaModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_chroma_transformer(model.transformer, config, config.transformer) - if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_chroma_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index 72f8ee2d8..083f74c34 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -42,8 +42,7 @@ def setup_optimizations( model: ErnieModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_ernie_transformer(model.transformer, config, config.transformer) + model.transformer_offload_conductor = enable_checkpointing_for_ernie_transformer(model.transformer, config, config.transformer) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index 6e6e7320f..019f0f179 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -45,9 +45,8 @@ def setup_optimizations( model: Flux2Model, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_flux2_transformer(model.transformer, config, config.transformer) - if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_flux2_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: if model.is_dev(): model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config, config.text_encoder) else: diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index e6638ddb2..8eae6211d 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -49,11 +49,10 @@ def setup_optimizations( model: FluxModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_flux_transformer(model.transformer, config, config.transformer) - if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_flux_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) - if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + if model.text_encoder_2 is not None: model.text_encoder_2_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index b5a856b73..b5d742fb8 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -49,15 +49,14 @@ def setup_optimizations( model: HiDreamModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_hi_dream_transformer(model.transformer, config, config.transformer) - if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_hi_dream_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) - if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + if model.text_encoder_2 is not None: enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) - if model.text_encoder_3 is not None and config.text_encoder_3.checkpointing_or_offloading_enabled(): + if model.text_encoder_3 is not None: model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) - if model.text_encoder_4 is not None and config.text_encoder_4.checkpointing_or_offloading_enabled(): + if model.text_encoder_4 is not None: model.text_encoder_4_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config, config.text_encoder_4) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index 2e32165e7..b1dbebca8 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -49,11 +49,10 @@ def setup_optimizations( model: HunyuanVideoModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config, config.transformer) - if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: model.text_encoder_1_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_1, config, config.text_encoder) - if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + if model.text_encoder_2 is not None: enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 4935ebbfd..0d6dc2000 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -51,10 +51,8 @@ def setup_optimizations( model: PixArtAlphaModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, config.transformer, offload_enabled=True) - if config.text_encoder.checkpointing_or_offloading_enabled(): - model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) + model.transformer_offload_conductor = enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, config.transformer, offload_enabled=True) + model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index 82004d8ac..0805358be 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -46,9 +46,8 @@ def setup_optimizations( model: QwenModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer) - if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: model.text_encoder_offload_conductor = enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index c26acbc82..d265fee2a 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -52,10 +52,8 @@ def setup_optimizations( config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_sana_transformer(model.transformer, config, config.transformer) - if config.text_encoder.checkpointing_or_offloading_enabled(): - model.text_encoder_offload_conductor = enable_checkpointing_for_gemma_layers(model.text_encoder, config, config.text_encoder) + model.transformer_offload_conductor = enable_checkpointing_for_sana_transformer(model.transformer, config, config.transformer) + model.text_encoder_offload_conductor = enable_checkpointing_for_gemma_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index efe761427..ee70847fa 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -48,13 +48,12 @@ def setup_optimizations( model: StableDiffusion3Model, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config, config.transformer) - if model.text_encoder_1 is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) - if model.text_encoder_2 is not None and config.text_encoder_2.checkpointing_or_offloading_enabled(): + if model.text_encoder_2 is not None: enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) - if model.text_encoder_3 is not None and config.text_encoder_3.checkpointing_or_offloading_enabled(): + if model.text_encoder_3 is not None: model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index a16c72cb8..3060816f7 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -56,8 +56,7 @@ def setup_optimizations( enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) if config.vae.checkpointing_enabled(): model.vae.enable_gradient_checkpointing() - if config.text_encoder.checkpointing_or_offloading_enabled(): - enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config, config.text_encoder) + enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 7fe35e1e1..2d0057b8b 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -51,10 +51,8 @@ def setup_optimizations( if config.unet.checkpointing_or_offloading_enabled(): model.unet.enable_gradient_checkpointing() enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) - if config.text_encoder.checkpointing_or_offloading_enabled(): - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) - if config.text_encoder_2.checkpointing_or_offloading_enabled(): - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index abc77bcdd..c2025537c 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -57,13 +57,12 @@ def setup_optimizations( model: WuerstchenModel, config: TrainConfig, ): - if config.prior.checkpointing_or_offloading_enabled(): - if model.model_type.is_wuerstchen_v2(): + if model.model_type.is_wuerstchen_v2(): + if config.prior.checkpointing_or_offloading_enabled(): model.prior_prior.enable_gradient_checkpointing() - elif model.model_type.is_stable_cascade(): - enable_checkpointing_for_stable_cascade_blocks(model.prior_prior, config, config.prior) - if config.text_encoder.checkpointing_or_offloading_enabled(): - enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config, config.text_encoder) + elif model.model_type.is_stable_cascade(): + enable_checkpointing_for_stable_cascade_blocks(model.prior_prior, config, config.prior) + enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.decoder_vqgan) diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index d64e7e88a..2a0287c19 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -47,9 +47,8 @@ def setup_optimizations( model: ZImageModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = enable_checkpointing_for_z_image_transformer(model.transformer, config, config.transformer) - if model.text_encoder is not None and config.text_encoder.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_z_image_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index 48bd9d816..4cdc4ce00 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -230,6 +230,9 @@ def enable_checkpointing( lists, # if there are multiple entries in this list, they must be in the exact order they are executed - otherwise offloading fails offload_enabled: bool = True, ) -> LayerOffloadConductor | None: + if not part.checkpointing_or_offloading_enabled() and not compile: + return None + # a conductor exists iff this part actually offloads (and the component supports conductor offloading) offload = offload_enabled and part.offloading_enabled() conductor = LayerOffloadConductor(model, config, part) if offload else None From cf5b106f769ac8bb0643076541275c9d3fe4a5e5 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 13 Jun 2026 23:46:40 +0200 Subject: [PATCH 08/10] Make BaseLensSetup checkpointing call unconditional, matching other models The gating now lives inside enable_checkpointing() after the split-offload merge, so the per-callsite guard here is redundant. --- modules/modelSetup/BaseLensSetup.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/modules/modelSetup/BaseLensSetup.py b/modules/modelSetup/BaseLensSetup.py index de868cbc3..1f884e3f0 100644 --- a/modules/modelSetup/BaseLensSetup.py +++ b/modules/modelSetup/BaseLensSetup.py @@ -44,9 +44,8 @@ def setup_optimizations( model: LensModel, config: TrainConfig, ): - if config.transformer.checkpointing_or_offloading_enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_lens_transformer(model.transformer, config, config.transformer) + model.transformer_offload_conductor = \ + enable_checkpointing_for_lens_transformer(model.transformer, config, config.transformer) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, From 58879e8774e6031b86e54ba4a13f55463a09fc6e Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 18 Jun 2026 01:36:57 +0200 Subject: [PATCH 09/10] Use decorator form for factory.register() in Lens model files Mirrors upstream commit 75a44d2bc, which converted the rest of the codebase from the trailing factory.register() call to the @factory.register decorator form. --- modules/dataLoader/LensBaseDataLoader.py | 4 +--- modules/modelSampler/LensSampler.py | 3 +-- modules/modelSetup/LensFineTuneSetup.py | 3 +-- modules/modelSetup/LensLoRASetup.py | 3 +-- 4 files changed, 4 insertions(+), 9 deletions(-) diff --git a/modules/dataLoader/LensBaseDataLoader.py b/modules/dataLoader/LensBaseDataLoader.py index 9c3d2816e..c4ad8088d 100644 --- a/modules/dataLoader/LensBaseDataLoader.py +++ b/modules/dataLoader/LensBaseDataLoader.py @@ -26,6 +26,7 @@ from mgds.pipelineModules.Tokenize import Tokenize +@factory.register(BaseDataLoader, ModelType.LENS) class LensBaseDataLoader( BaseDataLoader, DataLoaderText2ImageMixin, @@ -156,6 +157,3 @@ def _create_dataset( config, model, model_setup, train_progress, is_validation, aspect_bucketing_quantization=64, ) - - -factory.register(BaseDataLoader, LensBaseDataLoader, ModelType.LENS) diff --git a/modules/modelSampler/LensSampler.py b/modules/modelSampler/LensSampler.py index 22751d199..2d4ccb96b 100644 --- a/modules/modelSampler/LensSampler.py +++ b/modules/modelSampler/LensSampler.py @@ -21,6 +21,7 @@ from tqdm import tqdm +@factory.register(BaseModelSampler, ModelType.LENS) class LensSampler(BaseModelSampler): def __init__( self, @@ -185,5 +186,3 @@ def sample( ) on_sample(sampler_output) - -factory.register(BaseModelSampler, LensSampler, ModelType.LENS) diff --git a/modules/modelSetup/LensFineTuneSetup.py b/modules/modelSetup/LensFineTuneSetup.py index 01ac4efdb..d2b08cc0b 100644 --- a/modules/modelSetup/LensFineTuneSetup.py +++ b/modules/modelSetup/LensFineTuneSetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.LENS, TrainingMethod.FINE_TUNE) class LensFineTuneSetup( BaseLensSetup, ): @@ -92,5 +93,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, LensFineTuneSetup, ModelType.LENS, TrainingMethod.FINE_TUNE) diff --git a/modules/modelSetup/LensLoRASetup.py b/modules/modelSetup/LensLoRASetup.py index c481e72a6..760307fad 100644 --- a/modules/modelSetup/LensLoRASetup.py +++ b/modules/modelSetup/LensLoRASetup.py @@ -13,6 +13,7 @@ import torch +@factory.register(BaseModelSetup, ModelType.LENS, TrainingMethod.LORA) class LensLoRASetup( BaseLensSetup, ): @@ -104,5 +105,3 @@ def after_optimizer_step( train_progress: TrainProgress ): self.__setup_requires_grad(model, config) - -factory.register(BaseModelSetup, LensLoRASetup, ModelType.LENS, TrainingMethod.LORA) From 18e38f835496fddb182082b5892d1fb590a92d2a Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 19 Jun 2026 08:20:54 +0200 Subject: [PATCH 10/10] Strip padding from text latent cache for Lens Mirrors #1520 (Ernie): tokens_mask/text_encoder_hidden_state are fixed at PROMPT_MAX_LENGTH after EncodeLensText's crop, so caching wastes disk space on padding for shorter prompts. Prune before caching, pad back to PROMPT_MAX_LENGTH on load. Co-Authored-By: Claude Sonnet 4.6 --- modules/dataLoader/LensBaseDataLoader.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/modules/dataLoader/LensBaseDataLoader.py b/modules/dataLoader/LensBaseDataLoader.py index c4ad8088d..ade6ccfb0 100644 --- a/modules/dataLoader/LensBaseDataLoader.py +++ b/modules/dataLoader/LensBaseDataLoader.py @@ -18,6 +18,8 @@ from mgds.pipelineModules.DecodeVAE import DecodeVAE from mgds.pipelineModules.EncodeLensText import EncodeLensText from mgds.pipelineModules.EncodeVAE import EncodeVAE +from mgds.pipelineModules.PadMaskedTokens import PadMaskedTokens +from mgds.pipelineModules.PruneMaskedTokens import PruneMaskedTokens from mgds.pipelineModules.RescaleImageChannels import RescaleImageChannels from mgds.pipelineModules.SampleVAEDistribution import SampleVAEDistribution from mgds.pipelineModules.SaveImage import SaveImage @@ -53,13 +55,17 @@ def _preparation_modules(self, config: TrainConfig, model: LensModel): crop_start=PROMPT_TEMPLATE_CROP_START, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), ) - + prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state') modules = [rescale_image, encode_image, image_sample] if config.masked_training or config.model_type.has_mask_input(): modules.append(downscale_mask) modules += [tokenize_prompt, encode_prompt] + + if config.latent_caching: + modules.append(prune_masked_tokens) + return modules def _cache_modules(self, config: TrainConfig, model: LensModel, model_setup: BaseLensSetup): @@ -90,6 +96,8 @@ def _cache_modules(self, config: TrainConfig, model: LensModel, model_setup: Bas ) def _output_modules(self, config: TrainConfig, model: LensModel, model_setup: BaseLensSetup): + pad_masked_tokens = PadMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state', max_length=PROMPT_MAX_LENGTH) + output_names = [ 'image_path', 'latent_image', 'prompt', @@ -103,7 +111,7 @@ def _output_modules(self, config: TrainConfig, model: LensModel, model_setup: Ba output_names.append('text_encoder_hidden_state') - return self._output_modules_from_out_names( + output_module_list = self._output_modules_from_out_names( model, model_setup, output_names=output_names, config=config, @@ -113,6 +121,11 @@ def _output_modules(self, config: TrainConfig, model: LensModel, model_setup: Ba train_dtype=model.train_dtype, ) + if config.latent_caching: + output_module_list = [pad_masked_tokens] + output_module_list + + return output_module_list + def _debug_modules(self, config: TrainConfig, model: LensModel): debug_dir = os.path.join(config.debug_dir, "dataloader")