From 11506607538115fb24c7ad7eb7acec4fcb9becd6 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:07:18 +0200 Subject: [PATCH 01/10] Replace per-model device-placement methods with a generic materialize()/evict() API Replaces the per-model `{part}_to(device)` methods across all model classes, plus scattered call sites in dataLoader/modelSetup/modelSampler/GenericTrainer, with generic BaseModel methods driven by the existing ModelType.model_parts() registry: materialize(*parts), evict(*parts), and materialize_only(*parts) (evict everything else, then materialize the given parts - the swap-in/swap-out pattern used throughout the Samplers and text-caching setup). eval() and adapters() are likewise made concrete on BaseModel instead of hand-written per model. Models whose component names diverge (Wuerstchen) or that have components outside model_parts() (SD's depth_estimator, Anima's text_conditioner) override the relevant methods directly. Also fixes multi-TE samplers (Flux/SD3/SDXL/HiDream/HunyuanVideo) that previously evicted all but the first text encoder and ran encode_text with the rest still on temp_device. Co-Authored-By: Claude Sonnet 5 Co-Authored-By: Claude Opus 4.8 --- modules/dataLoader/AnimaBaseDataLoader.py | 2 +- modules/dataLoader/ChromaBaseDataLoader.py | 2 +- modules/dataLoader/ErnieBaseDataLoader.py | 2 +- modules/dataLoader/Flux2BaseDataLoader.py | 2 +- modules/dataLoader/FluxBaseDataLoader.py | 2 +- modules/dataLoader/HiDreamBaseDataLoader.py | 2 +- .../dataLoader/HunyuanVideoBaseDataLoader.py | 2 +- modules/dataLoader/IdeogramBaseDataLoader.py | 2 +- modules/dataLoader/Krea2BaseDataLoader.py | 2 +- .../dataLoader/PixArtAlphaBaseDataLoader.py | 2 +- modules/dataLoader/QwenBaseDataLoader.py | 2 +- modules/dataLoader/SanaBaseDataLoader.py | 2 +- .../StableDiffusion3BaseDataLoader.py | 2 +- .../StableDiffusionBaseDataLoader.py | 2 +- .../StableDiffusionFineTuneVaeDataLoader.py | 8 +- .../StableDiffusionXLBaseDataLoader.py | 2 +- .../dataLoader/WuerstchenBaseDataLoader.py | 9 +-- modules/dataLoader/ZImageBaseDataLoader.py | 2 +- .../mixin/DataLoaderText2ImageMixin.py | 9 +-- modules/model/AnimaModel.py | 42 +++------- modules/model/BaseModel.py | 79 +++++++++++++++++-- modules/model/ChromaModel.py | 37 --------- modules/model/ErnieModel.py | 35 -------- modules/model/Flux2Model.py | 35 -------- modules/model/FluxModel.py | 53 ------------- modules/model/HiDreamModel.py | 78 ------------------ modules/model/HunyuanVideoModel.py | 53 ------------- modules/model/IdeogramModel.py | 44 ----------- modules/model/Krea2Model.py | 33 -------- modules/model/PixArtAlphaModel.py | 37 --------- modules/model/QwenModel.py | 37 --------- modules/model/SanaModel.py | 37 --------- modules/model/StableDiffusion3Model.py | 64 --------------- modules/model/StableDiffusionModel.py | 36 +-------- modules/model/StableDiffusionXLModel.py | 49 ------------ modules/model/WuerstchenModel.py | 53 +++++-------- modules/model/ZImageModel.py | 35 -------- modules/modelSampler/AnimaSampler.py | 16 +--- modules/modelSampler/ChromaSampler.py | 16 +--- modules/modelSampler/ErnieSampler.py | 15 +--- modules/modelSampler/Flux2Sampler.py | 17 +--- modules/modelSampler/FluxSampler.py | 32 +++----- modules/modelSampler/HiDreamSampler.py | 16 +--- modules/modelSampler/HunyuanVideoSampler.py | 16 +--- modules/modelSampler/IdeogramSampler.py | 18 ++--- modules/modelSampler/Krea2Sampler.py | 16 +--- modules/modelSampler/PixArtAlphaSampler.py | 16 +--- modules/modelSampler/QwenSampler.py | 16 +--- modules/modelSampler/SanaSampler.py | 16 +--- .../modelSampler/StableDiffusion3Sampler.py | 16 +--- .../modelSampler/StableDiffusionSampler.py | 36 +++------ .../modelSampler/StableDiffusionVaeSampler.py | 4 +- .../modelSampler/StableDiffusionXLSampler.py | 36 +++------ modules/modelSampler/WuerstchenSampler.py | 31 ++------ modules/modelSampler/ZImageSampler.py | 15 +--- modules/modelSetup/AnimaFineTuneSetup.py | 26 ++---- modules/modelSetup/AnimaLoRASetup.py | 26 ++---- modules/modelSetup/BaseAnimaSetup.py | 4 +- modules/modelSetup/BaseChromaSetup.py | 8 +- modules/modelSetup/BaseErnieSetup.py | 5 +- modules/modelSetup/BaseFlux2Setup.py | 5 +- modules/modelSetup/BaseFluxSetup.py | 9 +-- modules/modelSetup/BaseHiDreamSetup.py | 17 ++-- modules/modelSetup/BaseHunyuanVideoSetup.py | 11 +-- modules/modelSetup/BaseIdeogramSetup.py | 5 +- modules/modelSetup/BaseKrea2Setup.py | 8 +- modules/modelSetup/BasePixArtAlphaSetup.py | 8 +- modules/modelSetup/BaseQwenSetup.py | 8 +- modules/modelSetup/BaseSanaSetup.py | 8 +- .../modelSetup/BaseStableDiffusion3Setup.py | 14 ++-- .../modelSetup/BaseStableDiffusionSetup.py | 10 +-- .../modelSetup/BaseStableDiffusionXLSetup.py | 11 +-- modules/modelSetup/BaseWuerstchenSetup.py | 8 +- modules/modelSetup/BaseZImageSetup.py | 5 +- modules/modelSetup/ChromaEmbeddingSetup.py | 23 ++---- modules/modelSetup/ChromaFineTuneSetup.py | 23 ++---- modules/modelSetup/ChromaLoRASetup.py | 23 ++---- modules/modelSetup/ErnieFineTuneSetup.py | 23 ++---- modules/modelSetup/ErnieLoRASetup.py | 23 ++---- modules/modelSetup/Flux2FineTuneSetup.py | 23 ++---- modules/modelSetup/Flux2LoRASetup.py | 23 ++---- modules/modelSetup/FluxEmbeddingSetup.py | 26 ++---- modules/modelSetup/FluxFineTuneSetup.py | 26 ++---- modules/modelSetup/FluxLoRASetup.py | 26 ++---- modules/modelSetup/HiDreamEmbeddingSetup.py | 32 +++----- modules/modelSetup/HiDreamFineTuneSetup.py | 32 +++----- modules/modelSetup/HiDreamLoRASetup.py | 32 +++----- .../modelSetup/HunyuanVideoEmbeddingSetup.py | 26 ++---- .../modelSetup/HunyuanVideoFineTuneSetup.py | 24 ++---- modules/modelSetup/HunyuanVideoLoRASetup.py | 26 ++---- modules/modelSetup/IdeogramFineTuneSetup.py | 27 ++----- modules/modelSetup/IdeogramLoRASetup.py | 27 ++----- modules/modelSetup/Krea2FineTuneSetup.py | 23 ++---- modules/modelSetup/Krea2LoRASetup.py | 23 ++---- .../modelSetup/PixArtAlphaEmbeddingSetup.py | 21 +---- .../modelSetup/PixArtAlphaFineTuneSetup.py | 23 ++---- modules/modelSetup/PixArtAlphaLoRASetup.py | 23 ++---- modules/modelSetup/QwenFineTuneSetup.py | 23 ++---- modules/modelSetup/QwenLoRASetup.py | 23 ++---- modules/modelSetup/SanaEmbeddingSetup.py | 21 +---- modules/modelSetup/SanaFineTuneSetup.py | 23 ++---- modules/modelSetup/SanaLoRASetup.py | 23 ++---- .../StableDiffusion3EmbeddingSetup.py | 29 +++---- .../StableDiffusion3FineTuneSetup.py | 29 +++---- .../modelSetup/StableDiffusion3LoRASetup.py | 29 +++---- .../StableDiffusionEmbeddingSetup.py | 24 ++---- .../StableDiffusionFineTuneSetup.py | 26 ++---- .../StableDiffusionFineTuneVaeSetup.py | 16 +--- .../modelSetup/StableDiffusionLoRASetup.py | 26 ++---- .../StableDiffusionXLEmbeddingSetup.py | 26 ++---- .../StableDiffusionXLFineTuneSetup.py | 26 ++---- .../modelSetup/StableDiffusionXLLoRASetup.py | 26 ++---- .../modelSetup/WuerstchenEmbeddingSetup.py | 28 ++----- modules/modelSetup/WuerstchenFineTuneSetup.py | 31 ++------ modules/modelSetup/WuerstchenLoRASetup.py | 31 ++------ modules/modelSetup/ZImageFineTuneSetup.py | 23 ++---- modules/modelSetup/ZImageLoRASetup.py | 23 ++---- modules/trainer/GenericTrainer.py | 11 ++- modules/ui/SampleWindowController.py | 2 +- 119 files changed, 599 insertions(+), 1937 deletions(-) diff --git a/modules/dataLoader/AnimaBaseDataLoader.py b/modules/dataLoader/AnimaBaseDataLoader.py index 9dba5d433..9acdbcb6a 100644 --- a/modules/dataLoader/AnimaBaseDataLoader.py +++ b/modules/dataLoader/AnimaBaseDataLoader.py @@ -109,7 +109,7 @@ def _debug_modules(self, config: TrainConfig, model: AnimaModel): #TODO clean up debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/ChromaBaseDataLoader.py b/modules/dataLoader/ChromaBaseDataLoader.py index 6b89457e6..660d9ecd0 100644 --- a/modules/dataLoader/ChromaBaseDataLoader.py +++ b/modules/dataLoader/ChromaBaseDataLoader.py @@ -120,7 +120,7 @@ def _debug_modules(self, config: TrainConfig, model: ChromaModel): #TODO clean u debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index 032180bc3..3da261272 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -110,7 +110,7 @@ def _debug_modules(self, config: TrainConfig, model: ErnieModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index a6bc3a05d..11815e8df 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -117,7 +117,7 @@ def _debug_modules(self, config: TrainConfig, model: Flux2Model): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/FluxBaseDataLoader.py b/modules/dataLoader/FluxBaseDataLoader.py index d23dee3a8..1184a3d17 100644 --- a/modules/dataLoader/FluxBaseDataLoader.py +++ b/modules/dataLoader/FluxBaseDataLoader.py @@ -143,7 +143,7 @@ def _debug_modules(self, config: TrainConfig, model: FluxModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/HiDreamBaseDataLoader.py b/modules/dataLoader/HiDreamBaseDataLoader.py index f080943bf..02a35899f 100644 --- a/modules/dataLoader/HiDreamBaseDataLoader.py +++ b/modules/dataLoader/HiDreamBaseDataLoader.py @@ -180,7 +180,7 @@ def _debug_modules(self, config: TrainConfig, model: HiDreamModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/HunyuanVideoBaseDataLoader.py b/modules/dataLoader/HunyuanVideoBaseDataLoader.py index 38f7e6a8a..5c0e5ece9 100644 --- a/modules/dataLoader/HunyuanVideoBaseDataLoader.py +++ b/modules/dataLoader/HunyuanVideoBaseDataLoader.py @@ -136,7 +136,7 @@ def _debug_modules(self, config: TrainConfig, model: HunyuanVideoModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/IdeogramBaseDataLoader.py b/modules/dataLoader/IdeogramBaseDataLoader.py index 28de23e40..f6b3c8f26 100644 --- a/modules/dataLoader/IdeogramBaseDataLoader.py +++ b/modules/dataLoader/IdeogramBaseDataLoader.py @@ -122,7 +122,7 @@ def _debug_modules(self, config: TrainConfig, model: IdeogramModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/Krea2BaseDataLoader.py b/modules/dataLoader/Krea2BaseDataLoader.py index 5d243f637..40ee02e27 100644 --- a/modules/dataLoader/Krea2BaseDataLoader.py +++ b/modules/dataLoader/Krea2BaseDataLoader.py @@ -131,7 +131,7 @@ def _debug_modules(self, config: TrainConfig, model: Krea2Model): #TODO clean up debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/PixArtAlphaBaseDataLoader.py b/modules/dataLoader/PixArtAlphaBaseDataLoader.py index bb110dc2e..1e7ff0e3f 100644 --- a/modules/dataLoader/PixArtAlphaBaseDataLoader.py +++ b/modules/dataLoader/PixArtAlphaBaseDataLoader.py @@ -121,7 +121,7 @@ def _debug_modules(self, config: TrainConfig, model: PixArtAlphaModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/QwenBaseDataLoader.py b/modules/dataLoader/QwenBaseDataLoader.py index 9a4a962a9..944dd70a9 100644 --- a/modules/dataLoader/QwenBaseDataLoader.py +++ b/modules/dataLoader/QwenBaseDataLoader.py @@ -124,7 +124,7 @@ def _debug_modules(self, config: TrainConfig, model: QwenModel): #TODO clean up debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/SanaBaseDataLoader.py b/modules/dataLoader/SanaBaseDataLoader.py index a44ff8130..2ebad9c6f 100644 --- a/modules/dataLoader/SanaBaseDataLoader.py +++ b/modules/dataLoader/SanaBaseDataLoader.py @@ -113,7 +113,7 @@ def _debug_modules(self, config: TrainConfig, model: SanaModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.train_dtype.torch_dtype()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/StableDiffusion3BaseDataLoader.py b/modules/dataLoader/StableDiffusion3BaseDataLoader.py index 55a0d9001..fd929b672 100644 --- a/modules/dataLoader/StableDiffusion3BaseDataLoader.py +++ b/modules/dataLoader/StableDiffusion3BaseDataLoader.py @@ -160,7 +160,7 @@ def _debug_modules(self, config: TrainConfig, model: StableDiffusion3Model): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/StableDiffusionBaseDataLoader.py b/modules/dataLoader/StableDiffusionBaseDataLoader.py index 63ad57cac..ee6a246be 100644 --- a/modules/dataLoader/StableDiffusionBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionBaseDataLoader.py @@ -130,7 +130,7 @@ def _debug_modules(self, config: TrainConfig, model: StableDiffusionModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py index ad0c890b0..7e1da8b10 100644 --- a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py +++ b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py @@ -8,7 +8,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType from modules.util.enum.TrainingMethod import TrainingMethod -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from mgds.OutputPipelineModule import OutputPipelineModule @@ -55,12 +54,9 @@ def _setup_cache_device( temp_device: torch.device, config: TrainConfig, ): - model.to(self.temp_device) - - model.vae_to(train_device) + model.materialize_only("vae") model.eval() - torch_gc() def __enumerate_input_modules(self, config: TrainConfig) -> list: supported_extensions = path_util.supported_image_extensions() @@ -248,7 +244,7 @@ def __debug_modules(self, config: TrainConfig, model: StableDiffusionModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae) upscale_mask = ScaleImage(in_name='latent_mask', out_name='decoded_mask', factor=8) diff --git a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py index 12739adf8..e481afba8 100644 --- a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py @@ -138,7 +138,7 @@ def _debug_modules(self, config: TrainConfig, model: StableDiffusionXLModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.vae_train_dtype.torch_dtype()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.vae_train_dtype.torch_dtype()) diff --git a/modules/dataLoader/WuerstchenBaseDataLoader.py b/modules/dataLoader/WuerstchenBaseDataLoader.py index 4689b09a4..e6f27d621 100644 --- a/modules/dataLoader/WuerstchenBaseDataLoader.py +++ b/modules/dataLoader/WuerstchenBaseDataLoader.py @@ -10,7 +10,6 @@ from modules.util import factory from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.ModelType import ModelType -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from mgds.pipelineModules.DecodeTokens import DecodeTokens @@ -74,10 +73,8 @@ def _cache_modules(self, config: TrainConfig, model: WuerstchenModel, model_setu ] def before_cache_image_fun(): - model.to(self.temp_device) - model.effnet_encoder_to(self.train_device) + model.materialize_only("effnet_encoder") model.eval() - torch_gc() return self._cache_modules_from_names( model, model_setup, @@ -109,10 +106,8 @@ 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.effnet_encoder_to(self.train_device) + model.materialize_only("effnet_encoder") model.eval() - torch_gc() return self._output_modules_from_out_names( model, model_setup, diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 23863d625..4ed9e1e54 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -116,7 +116,7 @@ def _debug_modules(self, config: TrainConfig, model: ZImageModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.vae_to(self.train_device) + model.materialize("vae") 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) diff --git a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py b/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py index 2654bdd19..a9ed0970d 100644 --- a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py +++ b/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py @@ -10,7 +10,6 @@ from modules.util import path_util from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.DataType import DataType -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress from mgds.OutputPipelineModule import OutputPipelineModule @@ -274,10 +273,8 @@ def _output_modules_from_out_names( ): if before_cache_image_fun is None: def prepare_vae(): - model.to(self.temp_device) - model.vae_to(self.train_device) + model.materialize_only("vae") model.eval() - torch_gc() before_cache_image_fun = prepare_vae sort_names = output_names + ['concept'] @@ -340,10 +337,8 @@ def _cache_modules_from_names( if before_cache_image_fun is None: def prepare_vae(): - model.to(self.temp_device) - model.vae_to(self.train_device) + model.materialize_only("vae") model.eval() - torch_gc() before_cache_image_fun = prepare_vae def before_cache_text_fun(): diff --git a/modules/model/AnimaModel.py b/modules/model/AnimaModel.py index a23d2998d..9c365273d 100644 --- a/modules/model/AnimaModel.py +++ b/modules/model/AnimaModel.py @@ -72,11 +72,6 @@ def __init__( 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 _diffusers_to_dit(self) -> list: # the netless diffusers CosmosTransformer3DModel -> Anima DiT rename (the inverse of diffusers' # scripts/convert_anima_to_diffusers.py transformer rename). These are the bare module names kohya-ss @@ -130,35 +125,22 @@ def lora_diffusers_to_kohya(self) -> list | None: # kohya-ss loads the DiT with the net. wrapper stripped -> the netless body. return self._diffusers_to_dit() - 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: - self.text_encoder_offload_conductor.to(device) - else: - self.text_encoder.to(device=device) - self.text_conditioner.to(device=device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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 materialize(self, *parts: str): + super().materialize(*parts) + # text_conditioner isn't in ModelType.model_parts(); it always travels with text_encoder. + if "text_encoder" in parts: + self.text_conditioner.to(device=self.train_device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) + def evict(self, *parts: str): + super().evict(*parts) + # evict() with no parts means "evict all", which includes text_encoder. + if not parts or "text_encoder" in parts: + self.text_conditioner.to(device=self.temp_device) def eval(self): - self.vae.eval() - self.text_encoder.eval() + super().eval() + # text_conditioner isn't in ModelType.model_parts(); it always travels with text_encoder. self.text_conditioner.eval() - self.transformer.eval() def create_pipeline(self): pipe = AnimaAutoBlocks().init_pipeline() diff --git a/modules/model/BaseModel.py b/modules/model/BaseModel.py index c346a60fe..d7227cced 100644 --- a/modules/model/BaseModel.py +++ b/modules/model/BaseModel.py @@ -1,4 +1,4 @@ -from abc import ABCMeta, abstractmethod +from abc import ABCMeta from contextlib import nullcontext from uuid import uuid4 @@ -11,6 +11,7 @@ from modules.util.enum.ModelType import ModelType from modules.util.modelSpec.ModelSpec import ModelSpec from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -96,17 +97,79 @@ def __init__( self.autocast_context = nullcontext() self.train_dtype = DataType.FLOAT_32 - @abstractmethod - def to(self, device: torch.device): - pass + @property + def train_device(self) -> torch.device: + return torch.device(self.train_config.train_device) + + @property + def temp_device(self) -> torch.device: + return torch.device(self.train_config.temp_device) + + def materialize(self, *parts: str): + # Move `parts` onto train_device. + for part in parts: + self._move_part(part, self.train_device) + + def evict(self, *parts: str): + # Move `parts` onto temp_device. No parts given -> every component in ModelType.model_parts(). + for part in parts or self.model_type.model_parts(): + self._move_part(part, self.temp_device) + torch_gc() + + def materialize_only(self, *parts: str): + # Materialize exactly `parts` on train_device; evict every other component in ModelType.model_parts() + # to temp_device. Lets a caller state what it needs now without tracking what to evict first. + # Evicts before materializing, so the two sets are never resident on train_device at once. + # Skipped (rather than passed as evict()) when empty, since evict() with no parts means "evict all". + to_evict = [part for part in self.model_type.model_parts() if part not in parts] + if to_evict: + self.evict(*to_evict) + self.materialize(*parts) + + def materialize_only_text_encoders(self): + # Materialize all of this model's text encoders on train_device, evicting everything else. Samplers + # call this before encode_text, which reads every text encoder the model has. + self.materialize_only(*self.model_type.text_encoder_parts()) + + def _move_part(self, part: str, device: torch.device): + # The generic per-component move: `part` (or `part_1` for the first of several split text encoders), + # its LoRA (`{part}_lora`), and its layer-offload conductor (`{part}_offload_conductor`), if present. + stem = f"{part}_1" if hasattr(self, f"{part}_1") else part + + conductor = getattr(self, f"{stem}_offload_conductor", None) + if conductor is not None: + conductor.to(device) + else: + component = getattr(self, stem) # raises if `part` doesn't name a real attribute + # None when the part is excluded from training (e.g. a text encoder with include_text_encoder off): + # it stays in model_parts() but the loader never populated it, so there is nothing to move. + if component is not None: + component.to(device=device) + + lora = getattr(self, f"{stem}_lora", None) + if lora is not None: + lora.to(device) - @abstractmethod def eval(self): - pass + # Put every present component on eval(); driven by the same part registry as materialize()/evict(). + # A model whose component names diverge (Wuerstchen) or that has a component outside model_parts() + # (SD's depth_estimator, Anima's text_conditioner) overrides this. + for part in self.model_type.model_parts(): + stem = f"{part}_1" if hasattr(self, f"{part}_1") else part + component = getattr(self, stem) + if component is not None: + component.eval() - @abstractmethod def adapters(self) -> list[LoRAModuleWrapper]: - pass + # Every LoRA adapter present on a model part, in model_parts() order. Parts without a LoRA + # (e.g. the vae, or an untrained component) contribute nothing. + result = [] + for part in self.model_type.model_parts(): + stem = f"{part}_1" if hasattr(self, f"{part}_1") else part + lora = getattr(self, f"{stem}_lora", None) + if lora is not None: + result.append(lora) + return result def diffusers_to_original(self) -> list | None: # the canonical(diffusers) -> native key-conversion BODY (rename only) for this model's denoising diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index 64279f982..4fbcf430b 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -95,12 +95,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_lora, - self.transformer_lora, - ] if a is not None] - def fusion_groups(self) -> list | None: # Chroma is Flux-structured: double blocks fuse img q/k/v and txt q/k/v separately; single blocks # fuse q/k/v + mlp into one linear1. @@ -164,37 +158,6 @@ def all_text_encoder_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_embedding] if self.embedding is not None else []) - 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: - self.text_encoder_offload_conductor.to(device) - else: - self.text_encoder.to(device=device) - - if self.text_encoder_lora is not None: - self.text_encoder_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - self.text_encoder.eval() - self.transformer.eval() - def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return ChromaPipeline( transformer=self.transformer, diff --git a/modules/model/ErnieModel.py b/modules/model/ErnieModel.py index 3d15e5385..42bb1be74 100644 --- a/modules/model/ErnieModel.py +++ b/modules/model/ErnieModel.py @@ -62,41 +62,6 @@ def __init__( 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): - if self.text_encoder is not None: - 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: - 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) - - def eval(self): - self.vae.eval() - if self.text_encoder is not None: - self.text_encoder.eval() - self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: return ErnieImagePipeline( transformer=self.transformer, diff --git a/modules/model/Flux2Model.py b/modules/model/Flux2Model.py index 15c64bcc3..79eb02006 100644 --- a/modules/model/Flux2Model.py +++ b/modules/model/Flux2Model.py @@ -71,11 +71,6 @@ def __init__( 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 fusion_groups(self) -> list | None: # Only the two double-block qkv groups -- Flux2's single block is the already-fused attn.to_qkv_mlp_proj. # NOTE: the fused suffix (attn.qkv / attn.added_qkv) is OneTrainer's name for the fused module on the @@ -126,36 +121,6 @@ def diffusers_to_original(self) -> list | None: ]), ] - def vae_to(self, device: torch.device): - self.vae.to(device=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: - 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: - 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) - - def eval(self): - self.vae.eval() - if self.text_encoder is not None: - self.text_encoder.eval() - self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: klass = Flux2Pipeline if self.is_dev() else Flux2KleinPipeline return klass( diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index 02631c3a3..5236b3498 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -116,13 +116,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_1_lora, - self.text_encoder_2_lora, - self.transformer_lora, - ] if a is not None] - def fusion_groups(self) -> list | None: # Flux fuses img qkv, txt qkv, and the single-block qkv+mlp (4 leaves -> linear1). return [ @@ -206,52 +199,6 @@ def all_text_encoder_2_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_2_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_2_embedding] if self.embedding is not None else []) - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def text_encoder_to(self, device: torch.device): - self.text_encoder_1_to(device=device) - self.text_encoder_2_to(device=device) - - def text_encoder_1_to(self, device: torch.device): - if self.text_encoder_1 is not None: - self.text_encoder_1.to(device=device) - - if self.text_encoder_1_lora is not None: - self.text_encoder_1_lora.to(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: - self.text_encoder_2_offload_conductor.to(device) - else: - self.text_encoder_2.to(device=device) - - if self.text_encoder_2_lora is not None: - self.text_encoder_2_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - if self.text_encoder_1 is not None: - self.text_encoder_1.eval() - if self.text_encoder_2 is not None: - self.text_encoder_2.eval() - self.transformer.eval() - def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return FluxPipeline( transformer=self.transformer, diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index a68feaf9d..f99b5116b 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -167,15 +167,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_1_lora, - self.text_encoder_2_lora, - self.text_encoder_3_lora, - self.text_encoder_4_lora, - self.transformer_lora, - ] if a is not None] - def lora_text_encoders(self) -> list[tuple[torch.nn.Module | None, dict[ModelFormat, str]]]: # HiDream's four TEs: clip_l + clip_g + t5xxl + llama (Comfy's HiDreamTEModel). Any can be absent, so # only the TEs actually present are declared. @@ -226,75 +217,6 @@ def all_text_encoder_4_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_4_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_4_embedding] if self.embedding is not None else []) - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def text_encoder_to(self, device: torch.device): - self.text_encoder_1_to(device=device) - self.text_encoder_2_to(device=device) - self.text_encoder_3_to(device=device) - self.text_encoder_4_to(device=device) - - def text_encoder_1_to(self, device: torch.device): - if self.text_encoder_1 is not None: - self.text_encoder_1.to(device=device) - - if self.text_encoder_1_lora is not None: - self.text_encoder_1_lora.to(device) - - def text_encoder_2_to(self, device: torch.device): - if self.text_encoder_2 is not None: - self.text_encoder_2.to(device=device) - - if self.text_encoder_2_lora is not None: - self.text_encoder_2_lora.to(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: - self.text_encoder_3_offload_conductor.to(device) - else: - self.text_encoder_3.to(device=device) - - if self.text_encoder_3_lora is not None: - self.text_encoder_3_lora.to(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: - self.text_encoder_4_offload_conductor.to(device) - else: - self.text_encoder_4.to(device=device) - - if self.text_encoder_4_lora is not None: - self.text_encoder_4_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - if self.text_encoder_1 is not None: - self.text_encoder_1.eval() - if self.text_encoder_2 is not None: - self.text_encoder_2.eval() - if self.text_encoder_3 is not None: - self.text_encoder_3.eval() - if self.text_encoder_4 is not None: - self.text_encoder_4.eval() - self.transformer.eval() - def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HiDreamImagePipeline( transformer=self.transformer, diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index eb3e8271b..9ec6b4516 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -131,13 +131,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_1_lora, - self.text_encoder_2_lora, - self.transformer_lora, - ] if a is not None] - def fusion_groups(self) -> list | None: # HunyuanVideo fuses qkv in three places: the context-embedder token-refiner blocks, the double # blocks (img + txt), and the single blocks (qkv+mlp). @@ -230,52 +223,6 @@ def all_text_encoder_2_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_2_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_2_embedding] if self.embedding is not None else []) - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def text_encoder_to(self, device: torch.device): - self.text_encoder_1_to(device=device) - self.text_encoder_2_to(device=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: - self.text_encoder_1_offload_conductor.to(device) - else: - self.text_encoder_1.to(device=device) - - if self.text_encoder_1_lora is not None: - self.text_encoder_1_lora.to(device) - - def text_encoder_2_to(self, device: torch.device): - if self.text_encoder_2 is not None: - self.text_encoder_2.to(device=device) - - if self.text_encoder_2_lora is not None: - self.text_encoder_2_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - if self.text_encoder_1 is not None: - self.text_encoder_1.eval() - if self.text_encoder_2 is not None: - self.text_encoder_2.eval() - self.transformer.eval() - def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return HunyuanVideoPipeline( transformer=self.transformer, diff --git a/modules/model/IdeogramModel.py b/modules/model/IdeogramModel.py index d388615f8..70fae8444 100644 --- a/modules/model/IdeogramModel.py +++ b/modules/model/IdeogramModel.py @@ -67,12 +67,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - # only the conditional transformer is trainable; the unconditional transformer never sees the concept - return [a for a in [ - self.transformer_lora, - ] if a is not None] - def fusion_groups(self) -> list | None: # Ideogram4 fuses q/k/v into one qkv Linear per block; everything else in the transformer -- including # the output projection (to_out.0 -> o, see diffusers_to_original) -- already matches the original @@ -91,44 +85,6 @@ def diffusers_to_original(self) -> list | None: ("{path}", "{path}"), ] - 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: - 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: - 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 unconditional_transformer_to(self, device: torch.device): - if self.unconditional_transformer is not None: - if self.unconditional_transformer_offload_conductor is not None: - self.unconditional_transformer_offload_conductor.to(device) - else: - self.unconditional_transformer.to(device=device) - - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.transformer_to(device) - self.unconditional_transformer_to(device) - - def eval(self): - self.vae.eval() - self.text_encoder.eval() - self.transformer.eval() - if self.unconditional_transformer is not None: - self.unconditional_transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: return Ideogram4Pipeline( transformer=self.transformer, diff --git a/modules/model/Krea2Model.py b/modules/model/Krea2Model.py index 8bd614848..cf5ca0d69 100644 --- a/modules/model/Krea2Model.py +++ b/modules/model/Krea2Model.py @@ -80,11 +80,6 @@ def __init__( 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 diffusers_to_original(self) -> list | None: # Krea 2's native checkpoint (krea/Krea-2-Raw's raw.safetensors) is a pure rename of the diffusers # Krea2Transformer2DModel state dict -- q/k/v are already split in both namespaces, so no qkv fusion @@ -127,34 +122,6 @@ def table_mod(t): return t.reshape(6, -1) ]), ] - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def text_encoder_to(self, device: torch.device): #TODO share more code between models - 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: - 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) - - def eval(self): - self.vae.eval() - self.text_encoder.eval() - self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: return Krea2Pipeline( transformer=self.transformer, diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index d53989da9..bf36d1a21 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -97,12 +97,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_lora, - self.transformer_lora, - ] if a is not None] - def fusion_groups(self) -> list | None: # PixArt fuses TWO attentions differently: self-attention fuses q/k/v (3 leaves), while cross-attention # fuses ONLY k/v (2 leaves) into kv_linear. @@ -158,37 +152,6 @@ def all_text_encoder_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_embedding] if self.embedding is not None else []) - 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: - self.text_encoder_offload_conductor.to(device) - else: - self.text_encoder.to(device=device) - - if self.text_encoder_lora is not None: - self.text_encoder_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - self.text_encoder.eval() - self.transformer.eval() - 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: diff --git a/modules/model/QwenModel.py b/modules/model/QwenModel.py index 775143063..3b69d8f8b 100644 --- a/modules/model/QwenModel.py +++ b/modules/model/QwenModel.py @@ -71,12 +71,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_lora, - self.transformer_lora, - ] if a is not None] - def lora_text_encoders(self) -> list[tuple[torch.nn.Module | None, dict[ModelFormat, str]]]: # Single Qwen2.5-VL TE (Comfy's QwenImageTEModel is a single qwen25_7b). return [ @@ -87,37 +81,6 @@ def lora_text_encoders(self) -> list[tuple[torch.nn.Module | None, dict[ModelFor }), ] - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def text_encoder_to(self, device: torch.device): #TODO share more code between models - if self.text_encoder_offload_conductor is not None: - self.text_encoder_offload_conductor.to(device) - else: - self.text_encoder.to(device=device) - - if self.text_encoder_lora is not None: - self.text_encoder_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - self.text_encoder.eval() - self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: return QwenImagePipeline( transformer=self.transformer, diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index 75dae101c..f43e98ed5 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -99,12 +99,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_lora, - self.transformer_lora, - ] if a is not None] - def lora_text_encoders(self) -> list[tuple[torch.nn.Module | None, dict[ModelFormat, str]]]: # Single Gemma2 TE. No COMFY_LORA name -- ComfyUI cannot load Sana, so the COMFY format refuses to # write its TE keys. @@ -123,37 +117,6 @@ def all_text_encoder_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_embedding] if self.embedding is not None else []) - 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: - self.text_encoder_offload_conductor.to(device) - else: - self.text_encoder.to(device=device) - - if self.text_encoder_lora is not None: - self.text_encoder_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - self.text_encoder.eval() - self.transformer.eval() - def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return SanaPipeline( tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer, diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index 7de0c4125..e9324bb1b 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -134,14 +134,6 @@ def __init__( self.transformer_lora = None self.lora_state_dict = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_1_lora, - self.text_encoder_2_lora, - self.text_encoder_3_lora, - self.transformer_lora, - ] if a is not None] - def fusion_groups(self) -> list | None: # SD3 fuses TWO joint streams (x_block + context_block) plus a dual-attention attn2 that only exists # in some SD3.5 blocks (the group fires per-block only where all its leaves are present). @@ -231,62 +223,6 @@ def all_text_encoder_3_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_3_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_3_embedding] if self.embedding is not None else []) - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def text_encoder_to(self, device: torch.device): - self.text_encoder_1_to(device=device) - self.text_encoder_2_to(device=device) - self.text_encoder_3_to(device=device) - - def text_encoder_1_to(self, device: torch.device): - if self.text_encoder_1 is not None: - self.text_encoder_1.to(device=device) - - if self.text_encoder_1_lora is not None: - self.text_encoder_1_lora.to(device) - - def text_encoder_2_to(self, device: torch.device): - if self.text_encoder_2 is not None: - self.text_encoder_2.to(device=device) - - if self.text_encoder_2_lora is not None: - self.text_encoder_2_lora.to(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: - self.text_encoder_3_offload_conductor.to(device) - else: - self.text_encoder_3.to(device=device) - - if self.text_encoder_3_lora is not None: - self.text_encoder_3_lora.to(device) - - def transformer_to(self, device: torch.device): - if self.transformer_offload_conductor is not None: - 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) - - def eval(self): - self.vae.eval() - if self.text_encoder_1 is not None: - self.text_encoder_1.eval() - if self.text_encoder_2 is not None: - self.text_encoder_2.eval() - if self.text_encoder_3 is not None: - self.text_encoder_3.eval() - self.transformer.eval() - def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusion3Pipeline( transformer=self.transformer, diff --git a/modules/model/StableDiffusionModel.py b/modules/model/StableDiffusionModel.py index be3a8af8b..697d2420a 100644 --- a/modules/model/StableDiffusionModel.py +++ b/modules/model/StableDiffusionModel.py @@ -104,12 +104,6 @@ def __init__( self.sd_config = None self.sd_config_filename = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_lora, - self.unet_lora, - ] if a is not None] - def diffusers_to_original(self) -> list | None: # SD1.5/2.x UNet diffusers -> original/sgm key map, convert()-native (bare sgm names, no top prefix). # SD has NO qkv fusion and NO add_embedding, so this is a pure key rename. Spatial-transformer @@ -181,37 +175,11 @@ def all_text_encoder_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_embedding] if self.embedding is not None else []) - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def depth_estimator_to(self, device: torch.device): - if self.depth_estimator is not None: - self.depth_estimator.to(device=device) - - def text_encoder_to(self, device: torch.device): - self.text_encoder.to(device=device) - - if self.text_encoder_lora is not None: - self.text_encoder_lora.to(device) - - def unet_to(self, device: torch.device): - self.unet.to(device=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 eval(self): - self.vae.eval() + super().eval() + # depth_estimator isn't in ModelType.model_parts(); only the depth model variants have it. if self.depth_estimator is not None: self.depth_estimator.eval() - self.text_encoder.eval() - self.unet.eval() def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: tokenizer = self.orig_tokenizer if use_original_tokenizers else self.tokenizer diff --git a/modules/model/StableDiffusionXLModel.py b/modules/model/StableDiffusionXLModel.py index 79af4e55f..7f2881d70 100644 --- a/modules/model/StableDiffusionXLModel.py +++ b/modules/model/StableDiffusionXLModel.py @@ -121,13 +121,6 @@ def __init__( self.sd_config = None self.sd_config_filename = None - def adapters(self) -> list[LoRAModuleWrapper]: - return [a for a in [ - self.text_encoder_1_lora, - self.text_encoder_2_lora, - self.unet_lora, - ] if a is not None] - def diffusers_to_original(self) -> list | None: # SDXL UNet diffusers -> original/sgm key map, convert()-native (bare sgm names, no top prefix). # SDXL has NO qkv fusion, so this is a pure key rename. Spatial-transformer (attention) blocks have @@ -199,48 +192,6 @@ def all_text_encoder_2_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_2_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.text_encoder_2_embedding] if self.embedding is not None else []) - def vae_to(self, device: torch.device): - self.vae.to(device=device) - - def text_encoder_to(self, device: torch.device): - self.text_encoder_1.to(device=device) - self.text_encoder_2.to(device=device) - - if self.text_encoder_1_lora is not None: - self.text_encoder_1_lora.to(device) - - if self.text_encoder_2_lora is not None: - self.text_encoder_2_lora.to(device) - - def text_encoder_1_to(self, device: torch.device): - self.text_encoder_1.to(device=device) - - if self.text_encoder_1_lora is not None: - self.text_encoder_1_lora.to(device) - - def text_encoder_2_to(self, device: torch.device): - self.text_encoder_2.to(device=device) - - if self.text_encoder_2_lora is not None: - self.text_encoder_2_lora.to(device) - - def unet_to(self, device: torch.device): - self.unet.to(device=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 eval(self): - self.vae.eval() - self.text_encoder_1.eval() - self.text_encoder_2.eval() - self.unet.eval() - def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline: return StableDiffusionXLPipeline( vae=self.vae, diff --git a/modules/model/WuerstchenModel.py b/modules/model/WuerstchenModel.py index 5268a9c16..f04c5b8c7 100644 --- a/modules/model/WuerstchenModel.py +++ b/modules/model/WuerstchenModel.py @@ -174,38 +174,27 @@ def all_prior_text_encoder_embeddings(self) -> list[BaseModelEmbedding]: return [embedding.text_encoder_embedding for embedding in self.additional_embeddings] \ + ([self.embedding.prior_text_encoder_embedding] if self.embedding is not None else []) - def decoder_text_encoder_to(self, device: torch.device): - self.decoder_text_encoder.to(device=device) - - def decoder_decoder_to(self, device: torch.device): - self.decoder_decoder.to(device=device) - - def decoder_vqgan_to(self, device: torch.device): - self.decoder_vqgan.to(device=device) - - def effnet_encoder_to(self, device: torch.device): - self.effnet_encoder.to(device=device) - - def prior_text_encoder_to(self, device: torch.device): - self.prior_text_encoder.to(device=device) - - if self.prior_text_encoder_lora is not None: - self.prior_text_encoder_lora.to(device) - - def prior_prior_to(self, device: torch.device): - self.prior_prior.to(device=device) - - if self.prior_prior_lora is not None: - self.prior_prior_lora.to(device) - - def to(self, device: torch.device): - 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) + def materialize(self, *parts: str): + super().materialize(*self._translate_parts(parts)) + + def evict(self, *parts: str): + # evict() with no parts means "evict all"; translate against the model's own full part list. + super().evict(*self._translate_parts(parts or self.model_type.model_parts())) + + def _translate_parts(self, parts: tuple[str, ...]) -> tuple[str, ...]: + # The prior stage's own diffusion module and the decoder stage's own diffusion module are each + # named after their stage, and the main text encoder belongs to the prior stage. + translated = [] + for part in parts: + if part == "prior": + translated.append("prior_prior") + elif part == "decoder": + translated.append("decoder_decoder") + elif part == "text_encoder": + translated.append("prior_text_encoder") + else: + translated.append(part) + return tuple(translated) def eval(self): if self.model_type.is_wuerstchen_v2(): diff --git a/modules/model/ZImageModel.py b/modules/model/ZImageModel.py index ba8511dce..8b60f2c28 100644 --- a/modules/model/ZImageModel.py +++ b/modules/model/ZImageModel.py @@ -74,11 +74,6 @@ def __init__( 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 checkpoint_diffusers_to_comfy(self) -> list | None: # Full-model COMFY_TRANSFORMER conversion: Z-Image is the one model whose Comfy checkpoint layout # diverges from diffusers/original (ComfyUI #12303). Only these keys change -- everything else passes @@ -97,36 +92,6 @@ def checkpoint_diffusers_to_comfy(self) -> list | None: ("{p}.attention.to_out.0", "{p}.attention.out"), ] - def vae_to(self, device: torch.device): - self.vae.to(device=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: - 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: - 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) - - def eval(self): - self.vae.eval() - if self.text_encoder is not None: - self.text_encoder.eval() - self.transformer.eval() - def create_pipeline(self) -> DiffusionPipeline: return ZImagePipeline( transformer=self.transformer, diff --git a/modules/modelSampler/AnimaSampler.py b/modules/modelSampler/AnimaSampler.py index 2c7322816..974d3482a 100644 --- a/modules/modelSampler/AnimaSampler.py +++ b/modules/modelSampler/AnimaSampler.py @@ -12,7 +12,6 @@ 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 @@ -66,7 +65,7 @@ def __sample_base( num_latent_channels = 16 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") batch_size = 2 if cfg_scale > 1.0 else 1 combined_prompt_embedding = self.model.encode_text( @@ -75,9 +74,6 @@ def __sample_base( train_device=self.train_device, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, 1, height // vae_scale_factor, width // vae_scale_factor), @@ -99,7 +95,7 @@ def __sample_base( if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = generator - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * batch_size) expanded_timestep = timestep.expand(batch_size) / noise_scheduler.config.num_train_timesteps @@ -121,11 +117,8 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = self.model.unscale_latents(latent_image) image = vae.decode(latents, return_dict=False)[0][:, :, 0] @@ -133,8 +126,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = self.image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/ChromaSampler.py b/modules/modelSampler/ChromaSampler.py index bc89e5cab..15441b164 100644 --- a/modules/modelSampler/ChromaSampler.py +++ b/modules/modelSampler/ChromaSampler.py @@ -12,7 +12,6 @@ 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 @@ -64,7 +63,7 @@ def __sample_base( num_latent_channels = 16 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") combined_prompt_embedding, text_attention_mask = self.model.encode_text( text=[prompt, negative_prompt], @@ -73,9 +72,6 @@ def __sample_base( text_encoder_layer_skip=text_encoder_layer_skip, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), @@ -109,7 +105,7 @@ def __sample_base( image_attention_mask = torch.full((2, image_seq_len), True, dtype=torch.bool, device=text_attention_mask.device) attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1) - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * 2) expanded_timestep = timestep.expand(2) @@ -134,9 +130,6 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - latent_image = self.model.unpack_latents( latent_image, height // vae_scale_factor, @@ -144,7 +137,7 @@ def __sample_base( ) # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = (latent_image / vae.config.scaling_factor) + vae.config.shift_factor image = vae.decode(latents, return_dict=False)[0] @@ -152,8 +145,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/ErnieSampler.py b/modules/modelSampler/ErnieSampler.py index 2122e5343..78aee2405 100644 --- a/modules/modelSampler/ErnieSampler.py +++ b/modules/modelSampler/ErnieSampler.py @@ -11,7 +11,6 @@ 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 @@ -63,7 +62,7 @@ def __sample_base( num_latent_channels = 32 # encode text - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") batch_size = 2 if cfg_scale > 1.0 else 1 text_bth, text_lens = self.model.encode_text( @@ -72,9 +71,6 @@ def __sample_base( ) dtype = self.model.train_dtype.torch_dtype() - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latents latent_image = torch.randn( size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), @@ -88,7 +84,7 @@ def __sample_base( noise_scheduler.set_timesteps(sigmas=sigmas, device=self.train_device) timesteps = noise_scheduler.timesteps - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") transformer = self.pipeline.transformer for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): @@ -112,9 +108,7 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") # unscale and unpatchify latents = self.model.unscale_latents(latent_image) @@ -126,8 +120,7 @@ def __sample_base( image = image.cpu().permute(0, 2, 3, 1).float().numpy() image = [PILImage.fromarray((img * 255).astype(np.uint8)) for img in image] - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/Flux2Sampler.py b/modules/modelSampler/Flux2Sampler.py index 0a4cca9b9..6653840fd 100644 --- a/modules/modelSampler/Flux2Sampler.py +++ b/modules/modelSampler/Flux2Sampler.py @@ -12,7 +12,6 @@ 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 @@ -69,7 +68,7 @@ def __sample_base( patch_size = 2 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") batch_size = 2 if cfg_scale > 1.0 and not transformer.config.guidance_embeds else 1 prompt_embedding = self.model.encode_text( @@ -78,9 +77,6 @@ def __sample_base( text_encoder_sequence_length=text_encoder_sequence_length, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), @@ -109,15 +105,13 @@ def __sample_base( text_ids = self.model.prepare_text_ids(prompt_embedding) - - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") guidance = (torch.tensor([cfg_scale], device=self.train_device, dtype=self.model.train_dtype.torch_dtype()) if transformer.config.guidance_embeds else None) 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()), timestep=expanded_timestep / 1000, @@ -137,9 +131,7 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latent_image = self.model.unpack_latents( latent_image, @@ -153,8 +145,7 @@ def __sample_base( image = image_processor.postprocess(image, output_type='pil') - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/FluxSampler.py b/modules/modelSampler/FluxSampler.py index fbe532053..a7fb6a34d 100644 --- a/modules/modelSampler/FluxSampler.py +++ b/modules/modelSampler/FluxSampler.py @@ -14,7 +14,6 @@ from modules.util.enum.NoiseScheduler import NoiseScheduler from modules.util.enum.VideoFormat import VideoFormat from modules.util.image_util import load_image -from modules.util.torch_util import torch_gc import torch from torch import nn @@ -72,7 +71,7 @@ def __sample_base( num_latent_channels = 16 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only_text_encoders() prompt_embedding, pooled_prompt_embedding = self.model.encode_text( text=prompt, @@ -83,9 +82,6 @@ def __sample_base( apply_attention_mask=transformer_attention_mask, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), @@ -115,7 +111,7 @@ def __sample_base( text_ids = torch.zeros(prompt_embedding.shape[1], 3, device=self.train_device) - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image]) expanded_timestep = timestep.expand(latent_model_input.shape[0]) @@ -147,8 +143,6 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() latent_image = self.model.unpack_latents( latent_image, height // vae_scale_factor, @@ -156,7 +150,7 @@ def __sample_base( ) # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = (latent_image / vae.config.scaling_factor) + vae.config.shift_factor image = vae.decode(latents, return_dict=False)[0] @@ -164,8 +158,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] #TODO remove and test, from Flux and other models. True is the default image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, @@ -222,7 +215,7 @@ def __sample_inpainting( num_latent_channels = 16 # prepare conditioning image - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") if sample_inpainting: t = transforms.Compose([ @@ -296,7 +289,7 @@ def __sample_inpainting( ) # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only_text_encoders() prompt_embedding, pooled_prompt_embedding = self.model.encode_text( text=prompt, @@ -307,9 +300,6 @@ def __sample_inpainting( apply_attention_mask=transformer_attention_mask, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), @@ -337,7 +327,7 @@ def __sample_inpainting( text_ids = torch.zeros(prompt_embedding.shape[1], 3, device=self.train_device) - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image]) latent_model_input = torch.concat( @@ -372,9 +362,6 @@ def __sample_inpainting( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - latent_image = self.model.unpack_latents( latent_image, height // vae_scale_factor, @@ -382,7 +369,7 @@ def __sample_inpainting( ) # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = (latent_image / vae.config.scaling_factor) + vae.config.shift_factor image = vae.decode(latents, return_dict=False)[0] @@ -390,8 +377,7 @@ def __sample_inpainting( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py index ee7569ccf..eefc3a358 100644 --- a/modules/modelSampler/HiDreamSampler.py +++ b/modules/modelSampler/HiDreamSampler.py @@ -12,7 +12,6 @@ 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 @@ -65,7 +64,7 @@ def __sample_base( num_latent_channels = 16 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only_text_encoders() text_encoder_3_prompt_embedding, text_encoder_4_prompt_embedding, pooled_prompt_embedding = \ self.model.combine_text_encoder_output( @@ -92,9 +91,6 @@ def __sample_base( combined_pooled_prompt_embedding = torch.cat( [negative_pooled_prompt_embedding, pooled_prompt_embedding], dim=0) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), @@ -111,7 +107,7 @@ def __sample_base( if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = generator - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * 2) expanded_timestep = timestep.expand(latent_model_input.shape[0]) @@ -142,11 +138,8 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = (latent_image / vae.config.scaling_factor) + vae.config.shift_factor image = vae.decode(latents, return_dict=False)[0] @@ -154,8 +147,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py index c12056c92..54f508f1d 100644 --- a/modules/modelSampler/HunyuanVideoSampler.py +++ b/modules/modelSampler/HunyuanVideoSampler.py @@ -12,7 +12,6 @@ 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 @@ -69,7 +68,7 @@ def __sample_base( num_latent_channels = 16 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only_text_encoders() prompt_embedding, pooled_prompt_embedding, prompt_attention_mask = self.model.encode_text( text=prompt, @@ -78,9 +77,6 @@ def __sample_base( text_encoder_2_layer_skip=text_encoder_2_layer_skip, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image num_latent_frames = (num_frames - 1) // vae_temporal_scale_factor + 1 latent_image = torch.randn( @@ -108,7 +104,7 @@ def __sample_base( if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = generator - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image]) expanded_timestep = timestep.expand(latent_model_input.shape[0]) @@ -139,19 +135,15 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = latent_image / vae.config.scaling_factor image = vae.decode(latents, return_dict=False)[0] image = video_processor.postprocess(image, output_type='pt') - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() is_image = image.shape[2] == 1 if is_image: diff --git a/modules/modelSampler/IdeogramSampler.py b/modules/modelSampler/IdeogramSampler.py index 0a5dd420b..cfc5862b0 100644 --- a/modules/modelSampler/IdeogramSampler.py +++ b/modules/modelSampler/IdeogramSampler.py @@ -11,7 +11,6 @@ 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 @@ -92,7 +91,7 @@ def pack_conditioning(text_features: torch.Tensor, text_lengths: torch.Tensor) - return max_text_tokens, position_ids, segment_ids, indicator, llm_features, text_z_padding # encode text (conditional branch, and the empty-prompt negative branch if needed) - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") text_features, text_lengths = self.model.encode_text( train_device=self.train_device, text=prompt, @@ -111,8 +110,6 @@ def pack_conditioning(text_features: torch.Tensor, text_lengths: torch.Tensor) - neg_text_z_padding, ) = pack_conditioning(neg_text_features, neg_text_lengths) del neg_text_features - self.model.text_encoder_to(self.temp_device) - torch_gc() if use_unconditional_transformer: # unconditional (image-only) branch: zeroed text features over the image-region slices of the layout @@ -141,9 +138,8 @@ def pack_conditioning(text_features: torch.Tensor, text_lengths: torch.Tensor) - timesteps = noise_scheduler.timesteps num_train_timesteps = noise_scheduler.config.num_train_timesteps - self.model.transformer_to(self.train_device) - if use_unconditional_transformer: - self.model.unconditional_transformer_to(self.train_device) + transformer_parts = ("transformer", "unconditional_transformer") if use_unconditional_transformer else ("transformer",) + self.model.materialize_only(*transformer_parts) for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): # scheduler stores num_train_timesteps-scaled timesteps; convert back to model time (0=noise, 1=data) @@ -190,10 +186,7 @@ def pack_conditioning(text_features: torch.Tensor, text_lengths: torch.Tensor) - on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - self.model.unconditional_transformer_to(self.temp_device) - torch_gc() - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") # bn-denormalize the packed latents and unpatchify back to (B, C, H, W) before VAE decode latents = self.model.unscale_latents(latent_image) @@ -205,8 +198,7 @@ def pack_conditioning(text_features: torch.Tensor, text_lengths: torch.Tensor) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() image = [PILImage.fromarray((img * 255).astype(np.uint8)) for img in image] - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/Krea2Sampler.py b/modules/modelSampler/Krea2Sampler.py index 20aecd5b1..e5f2a68c4 100644 --- a/modules/modelSampler/Krea2Sampler.py +++ b/modules/modelSampler/Krea2Sampler.py @@ -13,7 +13,6 @@ 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 @@ -67,7 +66,7 @@ def __sample_base( num_latent_channels = 16 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") batch_size = 2 if cfg_scale > 1.0 else 1 combined_prompt_embedding, text_attention_mask = self.model.encode_text( @@ -76,9 +75,6 @@ def __sample_base( train_device=self.train_device, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, 1, height // vae_scale_factor, width // vae_scale_factor), @@ -104,7 +100,7 @@ def __sample_base( if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = generator - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * batch_size) expanded_timestep = timestep.expand(batch_size) @@ -125,16 +121,13 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - - torch_gc() latent_image = self.model.unpack_latents( latent_image, height // vae_scale_factor, width // vae_scale_factor, ) - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = self.model.unscale_latents(latent_image) image = vae.decode(latents, return_dict=False)[0].squeeze(-3) @@ -142,8 +135,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/PixArtAlphaSampler.py b/modules/modelSampler/PixArtAlphaSampler.py index 9adddf3f2..4f21e5250 100644 --- a/modules/modelSampler/PixArtAlphaSampler.py +++ b/modules/modelSampler/PixArtAlphaSampler.py @@ -11,7 +11,6 @@ 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 @@ -63,7 +62,7 @@ def __sample_base( vae_scale_factor = self.pipeline.vae_scale_factor # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") prompt_embedding, tokens_attention_mask = self.model.encode_text( text=prompt, @@ -80,9 +79,6 @@ def __sample_base( combined_prompt_embedding = torch.cat([negative_prompt_embedding, prompt_embedding]) combined_prompt_attention_mask = torch.cat([negative_tokens_attention_mask, tokens_attention_mask]) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = noise_scheduler.timesteps @@ -113,7 +109,7 @@ def __sample_base( added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio} # denoising loop - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * 2) latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep) @@ -143,11 +139,8 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latent_image = latent_image.to(dtype=vae.dtype) image = vae.decode(latent_image / vae.config.scaling_factor, return_dict=False)[0] @@ -155,8 +148,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/QwenSampler.py b/modules/modelSampler/QwenSampler.py index 4ca604102..4ac4c3c17 100644 --- a/modules/modelSampler/QwenSampler.py +++ b/modules/modelSampler/QwenSampler.py @@ -13,7 +13,6 @@ 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 @@ -65,7 +64,7 @@ def __sample_base( num_latent_channels = 16 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") #unlike other models, Qwen benefits from CFG but is still quite good at CFG 1. Optimize for that: batch_size = 2 if cfg_scale > 1.0 else 1 @@ -75,9 +74,6 @@ def __sample_base( train_device=self.train_device, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, 1, height // vae_scale_factor, width // vae_scale_factor), @@ -110,7 +106,7 @@ def __sample_base( if torch.all(text_attention_mask): text_attention_mask = None - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * batch_size) expanded_timestep = timestep.expand(batch_size) @@ -134,9 +130,6 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - - torch_gc() latent_image = self.model.unpack_latents( latent_image, height // vae_scale_factor, @@ -144,7 +137,7 @@ def __sample_base( ) # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = self.model.unscale_latents(latent_image) image = vae.decode(latents, return_dict=False)[0].squeeze(-3) @@ -152,8 +145,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/SanaSampler.py b/modules/modelSampler/SanaSampler.py index 6251ab87e..17d0645b6 100644 --- a/modules/modelSampler/SanaSampler.py +++ b/modules/modelSampler/SanaSampler.py @@ -12,7 +12,6 @@ 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 @@ -63,7 +62,7 @@ def __sample_base( vae_scale_factor = self.pipeline.vae_scale_factor # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") prompt_embedding, tokens_attention_mask = self.model.encode_text( text=prompt, @@ -80,9 +79,6 @@ def __sample_base( combined_prompt_embedding = torch.cat([negative_prompt_embedding, prompt_embedding]) combined_prompt_attention_mask = torch.cat([negative_tokens_attention_mask, tokens_attention_mask]) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = noise_scheduler.timesteps @@ -102,7 +98,7 @@ def __sample_base( extra_step_kwargs["generator"] = generator # denoising loop - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * 2) @@ -127,11 +123,8 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latent_image = latent_image.to(dtype=vae.dtype) with self.model.vae_autocast_context: @@ -140,8 +133,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/StableDiffusion3Sampler.py b/modules/modelSampler/StableDiffusion3Sampler.py index fe159e914..eaea37f06 100644 --- a/modules/modelSampler/StableDiffusion3Sampler.py +++ b/modules/modelSampler/StableDiffusion3Sampler.py @@ -12,7 +12,6 @@ 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 @@ -67,7 +66,7 @@ def __sample_base( vae_scale_factor = self.pipeline.vae_scale_factor # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only_text_encoders() prompt_embedding, pooled_prompt_embedding = self.model.combine_text_encoder_output( *self.model.encode_text( @@ -93,9 +92,6 @@ def __sample_base( combined_pooled_prompt_embedding = torch.cat( [negative_pooled_prompt_embedding, pooled_prompt_embedding], dim=0) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = noise_scheduler.timesteps @@ -114,7 +110,7 @@ def __sample_base( if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = generator - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * 2) expanded_timestep = timestep.expand(latent_model_input.shape[0]) @@ -140,11 +136,8 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = (latent_image / vae.config.scaling_factor) + vae.config.shift_factor image = vae.decode(latents, return_dict=False)[0] @@ -152,8 +145,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/StableDiffusionSampler.py b/modules/modelSampler/StableDiffusionSampler.py index 92bfcd759..1881a15bf 100644 --- a/modules/modelSampler/StableDiffusionSampler.py +++ b/modules/modelSampler/StableDiffusionSampler.py @@ -12,7 +12,6 @@ from modules.util.enum.NoiseScheduler import NoiseScheduler from modules.util.enum.VideoFormat import VideoFormat from modules.util.image_util import load_image -from modules.util.torch_util import torch_gc import torch from torch import nn @@ -74,7 +73,7 @@ def __sample_base( vae_scale_factor = self.pipeline.vae_scale_factor # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") prompt_embedding = self.model.encode_text( text=prompt, @@ -90,9 +89,6 @@ def __sample_base( combined_prompt_embedding = torch.cat([negative_prompt_embedding, prompt_embedding]) \ .to(dtype=self.model.train_dtype.torch_dtype()) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = noise_scheduler.timesteps @@ -121,7 +117,7 @@ def __sample_base( extra_step_kwargs["generator"] = generator # denoising loop - self.model.unet_to(self.train_device) + self.model.materialize_only("unet") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * 2) latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep) @@ -155,11 +151,8 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.unet_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latent_image = latent_image.to(dtype=vae.dtype) image = vae.decode(latent_image / vae.config.scaling_factor, return_dict=False)[0] @@ -167,8 +160,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, @@ -223,7 +215,7 @@ def __sample_inpainting( vae_scale_factor = self.pipeline.vae_scale_factor # prepare conditioning image - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") if sample_inpainting: t = transforms.Compose([ @@ -277,11 +269,8 @@ def __sample_inpainting( device=self.train_device ) - self.model.vae_to(self.temp_device) - torch_gc() - # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") prompt_embedding = self.model.encode_text( text=prompt, @@ -297,9 +286,6 @@ def __sample_inpainting( combined_prompt_embedding = torch.cat([negative_prompt_embedding, prompt_embedding]) \ .to(dtype=self.model.train_dtype.torch_dtype()) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = noise_scheduler.timesteps @@ -328,7 +314,7 @@ def __sample_inpainting( extra_step_kwargs["generator"] = generator # denoising loop - self.model.unet_to(self.train_device) + self.model.materialize_only("unet") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = noise_scheduler.scale_model_input(latent_image, timestep) latent_model_input = torch.concat( @@ -365,11 +351,8 @@ def __sample_inpainting( on_update_progress(i + 1, len(timesteps)) - self.model.unet_to(self.temp_device) - torch_gc() - #decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latent_image = latent_image.to(dtype=vae.dtype) image = vae.decode(latent_image / vae.config.scaling_factor, return_dict=False)[0] @@ -377,8 +360,7 @@ def __sample_inpainting( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/StableDiffusionVaeSampler.py b/modules/modelSampler/StableDiffusionVaeSampler.py index 254c73fdd..325a1d667 100644 --- a/modules/modelSampler/StableDiffusionVaeSampler.py +++ b/modules/modelSampler/StableDiffusionVaeSampler.py @@ -63,13 +63,13 @@ def sample( image_tensor = t_in(image).to(device=self.train_device, dtype=self.model.vae.dtype) image_tensor = image_tensor * 2 - 1 - self.model.vae_to(self.train_device) + self.model.materialize("vae") with torch.no_grad(): latent_image_tensor = self.model.vae.encode(image_tensor.unsqueeze(0)).latent_dist.mean image_tensor = self.model.vae.decode(latent_image_tensor).sample.squeeze() - self.model.vae_to(self.temp_device) + self.model.evict("vae") image_tensor = (image_tensor + 1) * 0.5 image_tensor = image_tensor.clamp(0, 1) diff --git a/modules/modelSampler/StableDiffusionXLSampler.py b/modules/modelSampler/StableDiffusionXLSampler.py index 1f066268c..e351dac78 100644 --- a/modules/modelSampler/StableDiffusionXLSampler.py +++ b/modules/modelSampler/StableDiffusionXLSampler.py @@ -12,7 +12,6 @@ from modules.util.enum.NoiseScheduler import NoiseScheduler from modules.util.enum.VideoFormat import VideoFormat from modules.util.image_util import load_image -from modules.util.torch_util import torch_gc import torch from torch import nn @@ -69,7 +68,7 @@ def __sample_base( vae_scale_factor = self.pipeline.vae_scale_factor # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only_text_encoders() prompt_embedding, pooled_text_encoder_2_output = self.model.combine_text_encoder_output(*self.model.encode_text( text=prompt, @@ -88,9 +87,6 @@ def __sample_base( combined_prompt_embedding = torch.cat([negative_prompt_embedding, prompt_embedding]) \ .to(dtype=self.model.train_dtype.torch_dtype()) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = noise_scheduler.timesteps @@ -144,7 +140,7 @@ def __sample_base( extra_step_kwargs["generator"] = generator # denoising loop - self.model.unet_to(self.train_device) + self.model.materialize_only("unet") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = torch.cat([latent_image] * 2) latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep) @@ -177,11 +173,8 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.unet_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latent_image = latent_image.to(dtype=self.model.vae_train_dtype.torch_dtype()) with self.model.vae_autocast_context: @@ -190,8 +183,7 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, @@ -247,7 +239,7 @@ def __sample_inpainting( vae_scale_factor = self.pipeline.vae_scale_factor # prepare conditioning image - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") with self.model.vae_autocast_context: if sample_inpainting: @@ -303,11 +295,8 @@ def __sample_inpainting( device=self.train_device ) - self.model.vae_to(self.temp_device) - torch_gc() - # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only_text_encoders() prompt_embedding, pooled_text_encoder_2_output = self.model.combine_text_encoder_output( *self.model.encode_text( @@ -328,9 +317,6 @@ def __sample_inpainting( combined_prompt_embedding = torch.cat([negative_prompt_embedding, prompt_embedding]) \ .to(dtype=self.model.train_dtype.torch_dtype()) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = noise_scheduler.timesteps @@ -392,7 +378,7 @@ def __sample_inpainting( extra_step_kwargs["generator"] = generator # denoising loop - self.model.unet_to(self.train_device) + self.model.materialize_only("unet") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = noise_scheduler.scale_model_input(latent_image, timestep) latent_model_input = torch.concat( @@ -428,11 +414,8 @@ def __sample_inpainting( on_update_progress(i + 1, len(timesteps)) - self.model.unet_to(self.temp_device) - torch_gc() - # decode - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latent_image = latent_image.to(dtype=self.model.vae_train_dtype.torch_dtype()) with self.model.vae_autocast_context: @@ -441,8 +424,7 @@ def __sample_inpainting( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/WuerstchenSampler.py b/modules/modelSampler/WuerstchenSampler.py index d679757c2..ac18ae6bc 100644 --- a/modules/modelSampler/WuerstchenSampler.py +++ b/modules/modelSampler/WuerstchenSampler.py @@ -11,7 +11,6 @@ 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 @@ -50,7 +49,7 @@ def __sample_prior( on_update_progress, ): # prepare prompt - self.model.prior_text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") prompt_embedding, pooled_prompt_embedding = self.model.encode_text( text=prompt, @@ -70,9 +69,6 @@ def __sample_prior( combined_pooled_prompt_embedding = torch.cat([pooled_negative_prompt_embedding, pooled_prompt_embedding]) \ .to(dtype=self.model.prior_train_dtype.torch_dtype()) - self.model.prior_text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps prior_noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device) timesteps = prior_noise_scheduler.timesteps @@ -95,7 +91,7 @@ def __sample_prior( clip_img = torch.zeros(size=(2, 1, 768), dtype=self.model.prior_train_dtype.torch_dtype(), device=combined_prompt_embedding.device) - self.model.prior_prior_to(self.train_device) + self.model.materialize_only("prior") for i, timestep in enumerate(tqdm(timesteps[:-1], desc="sampling")): timestep = torch.stack([timestep]).to(dtype=self.model.prior_train_dtype.torch_dtype()) @@ -134,9 +130,6 @@ def __sample_prior( on_update_progress(i + 1, len(timesteps)) - self.model.prior_prior_to(self.temp_device) - torch_gc() - if self.model_type.is_wuerstchen_v2(): latent_image = latent_image * 42.0 - 1.0 @@ -161,9 +154,9 @@ def __sample_decoder( ): # prepare prompt if self.model_type.is_wuerstchen_v2(): - self.model.decoder_text_encoder_to(self.train_device) + self.model.materialize_only("decoder_text_encoder") elif self.model_type.is_stable_cascade(): - self.model.prior_text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") tokenizer_output = decoder_tokenizer( prompt, padding='max_length', @@ -188,12 +181,6 @@ def __sample_decoder( if self.model_type.is_stable_cascade(): prompt_embedding = text_encoder_output.text_embeds.unsqueeze(1) - if self.model_type.is_wuerstchen_v2(): - self.model.decoder_text_encoder_to(self.temp_device) - elif self.model_type.is_stable_cascade(): - self.model.prior_text_encoder_to(self.temp_device) - torch_gc() - # prepare timesteps decoder_noise_scheduler.set_timesteps(10, device=self.train_device) timesteps = decoder_noise_scheduler.timesteps @@ -214,7 +201,7 @@ def __sample_decoder( if "generator" in set(inspect.signature(decoder_noise_scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = generator - self.model.decoder_decoder_to(self.train_device) + self.model.materialize_only("decoder") for i, timestep in enumerate(tqdm(timesteps[:-1], desc="sampling")): timestep = torch.stack([timestep]).to(dtype=self.model.prior_train_dtype.torch_dtype()) @@ -248,9 +235,6 @@ def __sample_decoder( on_update_progress(i + 1, len(timesteps)) - self.model.decoder_decoder_to(self.temp_device) - torch_gc() - return latent_image @torch.no_grad() @@ -322,15 +306,14 @@ def __sample_base( ) # decode vqgan - self.model.decoder_vqgan_to(self.train_device) + self.model.materialize_only("decoder_vqgan") latents = decoder_vqgan.config.scale_factor * latent_image image_tensor = decoder_vqgan.decode(latents).sample.clamp(0, 1) image_array = image_tensor.permute(0, 2, 3, 1).cpu().squeeze().float().numpy() image_array = (image_array * 255).round().astype("uint8") - self.model.decoder_vqgan_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSampler/ZImageSampler.py b/modules/modelSampler/ZImageSampler.py index 0cfcb46d6..74cf22bb4 100644 --- a/modules/modelSampler/ZImageSampler.py +++ b/modules/modelSampler/ZImageSampler.py @@ -12,7 +12,6 @@ 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 @@ -65,7 +64,7 @@ def __sample_base( #patch_size = 2 # prepare prompt - self.model.text_encoder_to(self.train_device) + self.model.materialize_only("text_encoder") batch_size = 2 if cfg_scale > 1.0 else 1 prompt_embedding = self.model.encode_text( @@ -74,9 +73,6 @@ def __sample_base( train_device=self.train_device, ) - self.model.text_encoder_to(self.temp_device) - torch_gc() - # prepare latent image latent_image = torch.randn( size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), @@ -94,7 +90,7 @@ def __sample_base( if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = generator - self.model.transformer_to(self.train_device) + self.model.materialize_only("transformer") for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): latent_model_input = latent_image.unsqueeze(2).to(dtype=self.model.train_dtype.torch_dtype()) latent_model_input = torch.cat([latent_model_input] * batch_size) @@ -118,17 +114,14 @@ def __sample_base( on_update_progress(i + 1, len(timesteps)) - self.model.transformer_to(self.temp_device) - torch_gc() - self.model.vae_to(self.train_device) + self.model.materialize_only("vae") latents = self.model.unscale_latents(latent_image) image = vae.decode(latents, return_dict=False)[0] image = image_processor.postprocess(image, output_type='pil') - self.model.vae_to(self.temp_device) - torch_gc() + self.model.evict() return ModelSamplerOutput( file_type=FileType.IMAGE, diff --git a/modules/modelSetup/AnimaFineTuneSetup.py b/modules/modelSetup/AnimaFineTuneSetup.py index 2825d7cf7..3c62b006f 100644 --- a/modules/modelSetup/AnimaFineTuneSetup.py +++ b/modules/modelSetup/AnimaFineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.ANIMA, TrainingMethod.FINE_TUNE) class AnimaFineTuneSetup( BaseAnimaSetup, ): - 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: AnimaModel, @@ -70,10 +56,14 @@ def setup_train_device( config: TrainConfig, ): vae_on_train_device = not config.latent_caching - - model.text_encoder_to(self.temp_device if config.latent_caching else self.train_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + text_encoder_on_train_device = not config.latent_caching + + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.text_conditioner.eval() diff --git a/modules/modelSetup/AnimaLoRASetup.py b/modules/modelSetup/AnimaLoRASetup.py index ef3086a10..b5edd80b7 100644 --- a/modules/modelSetup/AnimaLoRASetup.py +++ b/modules/modelSetup/AnimaLoRASetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.ANIMA, TrainingMethod.LORA) class AnimaLoRASetup( BaseAnimaSetup, ): - 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: AnimaModel, @@ -82,10 +68,14 @@ def setup_train_device( config: TrainConfig, ): vae_on_train_device = not config.latent_caching - - model.text_encoder_to(self.temp_device if config.latent_caching else self.train_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + text_encoder_on_train_device = not config.latent_caching + + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.text_conditioner.eval() diff --git a/modules/modelSetup/BaseAnimaSetup.py b/modules/modelSetup/BaseAnimaSetup.py index 0ed9bbb19..8c0b9e0aa 100644 --- a/modules/modelSetup/BaseAnimaSetup.py +++ b/modules/modelSetup/BaseAnimaSetup.py @@ -166,9 +166,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: AnimaModel, config: TrainConfig): - model.to(self.temp_device) - - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") model.eval() torch_gc() diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 01dcfab0c..11ab09c82 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -18,7 +18,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -265,10 +264,9 @@ def calculate_loss( def prepare_text_caching(self, model: ChromaModel, config: TrainConfig): - model.to(self.temp_device) - if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") + else: + model.evict() model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c412073e8..d5050c49f 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -16,7 +16,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -160,7 +159,5 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: ErnieModel, config: TrainConfig): - model.to(self.temp_device) - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index 6b2c1bb3e..63e8dfc70 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -18,7 +18,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -182,7 +181,5 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: FluxModel, config: TrainConfig): - model.to(self.temp_device) - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 55999fb20..398d4a1b4 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -326,13 +326,12 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: FluxModel, config: TrainConfig): - model.to(self.temp_device) - + parts = [] if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) - + parts.append("text_encoder") if not config.train_text_encoder_2_or_embedding(): - model.text_encoder_2_to(self.train_device) + parts.append("text_encoder_2") + model.materialize_only(*parts) model.eval() torch_gc() diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 1a06f961b..301574535 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -19,7 +19,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -411,19 +410,15 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: HiDreamModel, config: TrainConfig): - model.to(self.temp_device) - + parts = [] if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) - + parts.append("text_encoder") if not config.train_text_encoder_2_or_embedding(): - model.text_encoder_2_to(self.train_device) - + parts.append("text_encoder_2") if not config.train_text_encoder_3_or_embedding(): - model.text_encoder_3_to(self.train_device) - + parts.append("text_encoder_3") if not config.train_text_encoder_4_or_embedding(): - model.text_encoder_4_to(self.train_device) + parts.append("text_encoder_4") + model.materialize_only(*parts) model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index 8e99dce84..4f5713008 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -19,7 +19,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -294,13 +293,11 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: HunyuanVideoModel, config: TrainConfig): - model.to(self.temp_device) - + parts = [] if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) - + parts.append("text_encoder") if not config.train_text_encoder_2_or_embedding(): - model.text_encoder_2_to(self.train_device) + parts.append("text_encoder_2") + model.materialize_only(*parts) model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseIdeogramSetup.py b/modules/modelSetup/BaseIdeogramSetup.py index f067e8c2e..7fb1a0bff 100644 --- a/modules/modelSetup/BaseIdeogramSetup.py +++ b/modules/modelSetup/BaseIdeogramSetup.py @@ -16,7 +16,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -204,7 +203,5 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: IdeogramModel, config: TrainConfig): - model.to(self.temp_device) - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseKrea2Setup.py b/modules/modelSetup/BaseKrea2Setup.py index 52e9b5959..e4249e897 100644 --- a/modules/modelSetup/BaseKrea2Setup.py +++ b/modules/modelSetup/BaseKrea2Setup.py @@ -16,7 +16,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -178,10 +177,9 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: Krea2Model, config: TrainConfig): - model.to(self.temp_device) - if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") + else: + model.evict() model.eval() - torch_gc() diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 57bf3a40b..008e9a0c0 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -18,7 +18,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -326,10 +325,9 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: PixArtAlphaModel, config: TrainConfig): - model.to(self.temp_device) - if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") + else: + model.evict() model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index a618dc28f..0615a3c1e 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -16,7 +16,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -177,10 +176,9 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: QwenModel, config: TrainConfig): - model.to(self.temp_device) - if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") + else: + model.evict() model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index a96afd770..8718142a7 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -18,7 +18,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -246,10 +245,9 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: SanaModel, config: TrainConfig): - model.to(self.temp_device) - if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") + else: + model.evict() model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index 678d727ff..dd7cf0841 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -19,7 +19,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -346,16 +345,13 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusion3Model, config: TrainConfig): - model.to(self.temp_device) - + parts = [] if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) - + parts.append("text_encoder") if not config.train_text_encoder_2_or_embedding(): - model.text_encoder_2_to(self.train_device) - + parts.append("text_encoder_2") if not config.train_text_encoder_3_or_embedding(): - model.text_encoder_3_to(self.train_device) + parts.append("text_encoder_3") + model.materialize_only(*parts) model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 5be9fc97e..7a3f910c8 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -19,7 +19,6 @@ from modules.util.conv_util import apply_circular_padding_to_conv2d from modules.util.dtype_util import create_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -334,10 +333,11 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusionModel, config: TrainConfig): - model.to(self.temp_device) - if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") + else: + model.evict() + if model.depth_estimator is not None: + model.depth_estimator.to(self.temp_device) model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index e5c7bf0c3..54a080579 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -19,7 +19,6 @@ from modules.util.conv_util import apply_circular_padding_to_conv2d from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -379,13 +378,11 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusionXLModel, config: TrainConfig): - model.to(self.temp_device) - + parts = [] if not config.train_text_encoder_or_embedding(): - model.text_encoder_to(self.train_device) - + parts.append("text_encoder") if not config.train_text_encoder_2_or_embedding(): - model.text_encoder_2_to(self.train_device) + parts.append("text_encoder_2") + model.materialize_only(*parts) model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index 078e14fff..d73052f83 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -20,7 +20,6 @@ disable_fp16_autocast_context, ) from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -345,10 +344,9 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: WuerstchenModel, config: TrainConfig): - model.to(self.temp_device) - if not config.train_text_encoder_or_embedding(): - model.prior_text_encoder_to(self.train_device) + model.materialize_only("text_encoder") + else: + model.evict() model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index a180f57d2..b77f9bb72 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -17,7 +17,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -160,8 +159,6 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: ZImageModel, config: TrainConfig): - model.to(self.temp_device) - model.text_encoder_to(self.train_device) + model.materialize_only("text_encoder") model.eval() - torch_gc() diff --git a/modules/modelSetup/ChromaEmbeddingSetup.py b/modules/modelSetup/ChromaEmbeddingSetup.py index 88aef8db4..ce49b9db2 100644 --- a/modules/modelSetup/ChromaEmbeddingSetup.py +++ b/modules/modelSetup/ChromaEmbeddingSetup.py @@ -9,25 +9,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.CHROMA_1, TrainingMethod.EMBEDDING) class ChromaEmbeddingSetup( BaseChromaSetup, ): - 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: ChromaModel, @@ -74,9 +60,12 @@ def setup_train_device( ): vae_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if config.text_encoder.train_embedding: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/ChromaFineTuneSetup.py b/modules/modelSetup/ChromaFineTuneSetup.py index 95bd8bac6..64c2aa67f 100644 --- a/modules/modelSetup/ChromaFineTuneSetup.py +++ b/modules/modelSetup/ChromaFineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.CHROMA_1, TrainingMethod.FINE_TUNE) class ChromaFineTuneSetup( BaseChromaSetup, ): - 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: ChromaModel, @@ -87,9 +73,12 @@ def setup_train_device( config.train_text_encoder_or_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/ChromaLoRASetup.py b/modules/modelSetup/ChromaLoRASetup.py index 43e9c6572..63e3ef964 100644 --- a/modules/modelSetup/ChromaLoRASetup.py +++ b/modules/modelSetup/ChromaLoRASetup.py @@ -11,25 +11,11 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.CHROMA_1, TrainingMethod.LORA) class ChromaLoRASetup( BaseChromaSetup, ): - 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: ChromaModel, @@ -114,9 +100,12 @@ def setup_train_device( config.train_text_encoder_or_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/ErnieFineTuneSetup.py b/modules/modelSetup/ErnieFineTuneSetup.py index feadd4fca..dc7f2f770 100644 --- a/modules/modelSetup/ErnieFineTuneSetup.py +++ b/modules/modelSetup/ErnieFineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.ERNIE, TrainingMethod.FINE_TUNE) class ErnieFineTuneSetup( BaseErnieSetup, ): - 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: ErnieModel, @@ -67,9 +53,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/ErnieLoRASetup.py b/modules/modelSetup/ErnieLoRASetup.py index 1a70ef269..70b449dd0 100644 --- a/modules/modelSetup/ErnieLoRASetup.py +++ b/modules/modelSetup/ErnieLoRASetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.ERNIE, TrainingMethod.LORA) class ErnieLoRASetup( BaseErnieSetup, ): - 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: ErnieModel, @@ -77,9 +63,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/Flux2FineTuneSetup.py b/modules/modelSetup/Flux2FineTuneSetup.py index 7ca128a09..be310d512 100644 --- a/modules/modelSetup/Flux2FineTuneSetup.py +++ b/modules/modelSetup/Flux2FineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.FLUX_2, TrainingMethod.FINE_TUNE) class Flux2FineTuneSetup( BaseFlux2Setup, ): - 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: Flux2Model, @@ -67,9 +53,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/Flux2LoRASetup.py b/modules/modelSetup/Flux2LoRASetup.py index fe1750528..0358e62c1 100644 --- a/modules/modelSetup/Flux2LoRASetup.py +++ b/modules/modelSetup/Flux2LoRASetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.FLUX_2, TrainingMethod.LORA) class Flux2LoRASetup( BaseFlux2Setup, ): - 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: Flux2Model, @@ -80,9 +66,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/FluxEmbeddingSetup.py b/modules/modelSetup/FluxEmbeddingSetup.py index f31c5e922..4751d5e37 100644 --- a/modules/modelSetup/FluxEmbeddingSetup.py +++ b/modules/modelSetup/FluxEmbeddingSetup.py @@ -9,26 +9,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.FLUX_DEV_1, TrainingMethod.EMBEDDING) @factory.register(BaseModelSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.EMBEDDING) class FluxEmbeddingSetup( BaseFluxSetup, ): - 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: FluxModel, @@ -87,10 +73,14 @@ def setup_train_device( ): vae_on_train_device = not config.latent_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if config.text_encoder.train_embedding: + parts.append("text_encoder") + if config.text_encoder_2.train_embedding: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1 is not None: model.text_encoder_1.eval() diff --git a/modules/modelSetup/FluxFineTuneSetup.py b/modules/modelSetup/FluxFineTuneSetup.py index 5d2aa61cd..4dc713d97 100644 --- a/modules/modelSetup/FluxFineTuneSetup.py +++ b/modules/modelSetup/FluxFineTuneSetup.py @@ -10,26 +10,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.FLUX_DEV_1, TrainingMethod.FINE_TUNE) @factory.register(BaseModelSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.FINE_TUNE) class FluxFineTuneSetup( BaseFluxSetup, ): - 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: FluxModel, @@ -103,10 +89,14 @@ def setup_train_device( config.train_text_encoder_2_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/FluxLoRASetup.py b/modules/modelSetup/FluxLoRASetup.py index b883dc7b3..557710459 100644 --- a/modules/modelSetup/FluxLoRASetup.py +++ b/modules/modelSetup/FluxLoRASetup.py @@ -11,26 +11,12 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.FLUX_DEV_1, TrainingMethod.LORA) @factory.register(BaseModelSetup, ModelType.FLUX_FILL_DEV_1, TrainingMethod.LORA) class FluxLoRASetup( BaseFluxSetup, ): - 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: FluxModel, @@ -146,10 +132,14 @@ def setup_train_device( config.train_text_encoder_2_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/HiDreamEmbeddingSetup.py b/modules/modelSetup/HiDreamEmbeddingSetup.py index 659f3834e..b2e80df44 100644 --- a/modules/modelSetup/HiDreamEmbeddingSetup.py +++ b/modules/modelSetup/HiDreamEmbeddingSetup.py @@ -9,25 +9,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.HI_DREAM_FULL, TrainingMethod.EMBEDDING) class HiDreamEmbeddingSetup( BaseHiDreamSetup, ): - 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: HiDreamModel, @@ -106,12 +92,18 @@ def setup_train_device( ): vae_on_train_device = not config.latent_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) - model.text_encoder_3_to(self.train_device if config.text_encoder_3.train_embedding else self.temp_device) - model.text_encoder_4_to(self.train_device if config.text_encoder_4.train_embedding else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if config.text_encoder.train_embedding: + parts.append("text_encoder") + if config.text_encoder_2.train_embedding: + parts.append("text_encoder_2") + if config.text_encoder_3.train_embedding: + parts.append("text_encoder_3") + if config.text_encoder_4.train_embedding: + parts.append("text_encoder_4") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1 is not None: model.text_encoder_1.eval() diff --git a/modules/modelSetup/HiDreamFineTuneSetup.py b/modules/modelSetup/HiDreamFineTuneSetup.py index 2361e69ff..c1bad4337 100644 --- a/modules/modelSetup/HiDreamFineTuneSetup.py +++ b/modules/modelSetup/HiDreamFineTuneSetup.py @@ -12,25 +12,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.HI_DREAM_FULL, TrainingMethod.FINE_TUNE) class HiDreamFineTuneSetup( BaseHiDreamSetup, ): - 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: HiDreamModel, @@ -137,12 +123,18 @@ def setup_train_device( config.train_text_encoder_4_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.text_encoder_3_to(self.train_device if text_encoder_3_on_train_device else self.temp_device) - model.text_encoder_4_to(self.train_device if text_encoder_4_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if text_encoder_3_on_train_device: + parts.append("text_encoder_3") + if text_encoder_4_on_train_device: + parts.append("text_encoder_4") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/HiDreamLoRASetup.py b/modules/modelSetup/HiDreamLoRASetup.py index a63a30ab7..6bb90324d 100644 --- a/modules/modelSetup/HiDreamLoRASetup.py +++ b/modules/modelSetup/HiDreamLoRASetup.py @@ -13,25 +13,11 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.HI_DREAM_FULL, TrainingMethod.LORA) class HiDreamLoRASetup( BaseHiDreamSetup, ): - 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: HiDreamModel, @@ -209,12 +195,18 @@ def setup_train_device( config.train_text_encoder_4_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.text_encoder_3_to(self.train_device if text_encoder_3_on_train_device else self.temp_device) - model.text_encoder_4_to(self.train_device if text_encoder_4_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if text_encoder_3_on_train_device: + parts.append("text_encoder_3") + if text_encoder_4_on_train_device: + parts.append("text_encoder_4") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/HunyuanVideoEmbeddingSetup.py b/modules/modelSetup/HunyuanVideoEmbeddingSetup.py index 91022c9f6..4b1168e4b 100644 --- a/modules/modelSetup/HunyuanVideoEmbeddingSetup.py +++ b/modules/modelSetup/HunyuanVideoEmbeddingSetup.py @@ -9,25 +9,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.EMBEDDING) class HunyuanVideoEmbeddingSetup( BaseHunyuanVideoSetup, ): - 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: HunyuanVideoModel, @@ -86,10 +72,14 @@ def setup_train_device( ): vae_on_train_device = not config.latent_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if config.text_encoder.train_embedding: + parts.append("text_encoder") + if config.text_encoder_2.train_embedding: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1 is not None: model.text_encoder_1.eval() diff --git a/modules/modelSetup/HunyuanVideoFineTuneSetup.py b/modules/modelSetup/HunyuanVideoFineTuneSetup.py index 9bb26e5d0..66786ef19 100644 --- a/modules/modelSetup/HunyuanVideoFineTuneSetup.py +++ b/modules/modelSetup/HunyuanVideoFineTuneSetup.py @@ -17,18 +17,6 @@ class HunyuanVideoFineTuneSetup( BaseHunyuanVideoSetup, ): - 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: HunyuanVideoModel, @@ -103,10 +91,14 @@ def setup_train_device( config.train_text_encoder_2_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/HunyuanVideoLoRASetup.py b/modules/modelSetup/HunyuanVideoLoRASetup.py index 45a5f619e..987df90ee 100644 --- a/modules/modelSetup/HunyuanVideoLoRASetup.py +++ b/modules/modelSetup/HunyuanVideoLoRASetup.py @@ -13,25 +13,11 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.HUNYUAN_VIDEO, TrainingMethod.LORA) class HunyuanVideoLoRASetup( BaseHunyuanVideoSetup, ): - 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: HunyuanVideoModel, @@ -150,10 +136,14 @@ def setup_train_device( config.train_text_encoder_2_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/IdeogramFineTuneSetup.py b/modules/modelSetup/IdeogramFineTuneSetup.py index a259af03e..79dd5a05c 100644 --- a/modules/modelSetup/IdeogramFineTuneSetup.py +++ b/modules/modelSetup/IdeogramFineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.IDEOGRAM_4, TrainingMethod.FINE_TUNE) class IdeogramFineTuneSetup( BaseIdeogramSetup, ): - 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: IdeogramModel, @@ -71,11 +57,14 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) - # the unconditional transformer is only needed for sampling; keep it off the train device during training - model.unconditional_transformer_to(self.temp_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + # the unconditional transformer is only needed for sampling; materialize_only() evicts it as it's + # not in parts, keeping it off the train device during training + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/IdeogramLoRASetup.py b/modules/modelSetup/IdeogramLoRASetup.py index e3f83d8bf..fbdcec35e 100644 --- a/modules/modelSetup/IdeogramLoRASetup.py +++ b/modules/modelSetup/IdeogramLoRASetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.IDEOGRAM_4, TrainingMethod.LORA) class IdeogramLoRASetup( BaseIdeogramSetup, ): - 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: IdeogramModel, @@ -80,11 +66,14 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) - # the unconditional transformer is only needed for sampling; keep it off the train device during training - model.unconditional_transformer_to(self.temp_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + # the unconditional transformer is only needed for sampling; materialize_only() evicts it as it's + # not in parts, keeping it off the train device during training + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/Krea2FineTuneSetup.py b/modules/modelSetup/Krea2FineTuneSetup.py index c2ff053f1..07cc43944 100644 --- a/modules/modelSetup/Krea2FineTuneSetup.py +++ b/modules/modelSetup/Krea2FineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.KREA_2, TrainingMethod.FINE_TUNE) class Krea2FineTuneSetup( BaseKrea2Setup, ): - 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: Krea2Model, @@ -73,9 +59,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/Krea2LoRASetup.py b/modules/modelSetup/Krea2LoRASetup.py index 7310317bc..507ec8f59 100644 --- a/modules/modelSetup/Krea2LoRASetup.py +++ b/modules/modelSetup/Krea2LoRASetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.KREA_2, TrainingMethod.LORA) class Krea2LoRASetup( BaseKrea2Setup, ): - 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: Krea2Model, @@ -83,9 +69,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py index 232bf2e01..1c8413f3f 100644 --- a/modules/modelSetup/PixArtAlphaEmbeddingSetup.py +++ b/modules/modelSetup/PixArtAlphaEmbeddingSetup.py @@ -9,26 +9,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.PIXART_ALPHA, TrainingMethod.EMBEDDING) @factory.register(BaseModelSetup, ModelType.PIXART_SIGMA, TrainingMethod.EMBEDDING) class PixArtAlphaEmbeddingSetup( BasePixArtAlphaSetup, ): - 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: PixArtAlphaModel, @@ -74,9 +60,10 @@ def setup_train_device( ): vae_on_train_device = self.debug_mode - model.text_encoder_to(self.train_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer", "text_encoder"] + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/PixArtAlphaFineTuneSetup.py b/modules/modelSetup/PixArtAlphaFineTuneSetup.py index e34db3ac5..13bec1203 100644 --- a/modules/modelSetup/PixArtAlphaFineTuneSetup.py +++ b/modules/modelSetup/PixArtAlphaFineTuneSetup.py @@ -10,26 +10,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.PIXART_ALPHA, TrainingMethod.FINE_TUNE) @factory.register(BaseModelSetup, ModelType.PIXART_SIGMA, TrainingMethod.FINE_TUNE) class PixArtAlphaFineTuneSetup( BasePixArtAlphaSetup, ): - 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: PixArtAlphaModel, @@ -94,9 +80,12 @@ def setup_train_device( or config.train_any_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/PixArtAlphaLoRASetup.py b/modules/modelSetup/PixArtAlphaLoRASetup.py index 0456b4165..8c7de7766 100644 --- a/modules/modelSetup/PixArtAlphaLoRASetup.py +++ b/modules/modelSetup/PixArtAlphaLoRASetup.py @@ -11,26 +11,12 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.PIXART_ALPHA, TrainingMethod.LORA) @factory.register(BaseModelSetup, ModelType.PIXART_SIGMA, TrainingMethod.LORA) class PixArtAlphaLoRASetup( BasePixArtAlphaSetup, ): - 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: PixArtAlphaModel, @@ -115,9 +101,12 @@ def setup_train_device( or config.train_any_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/QwenFineTuneSetup.py b/modules/modelSetup/QwenFineTuneSetup.py index 7b9f5dc60..5c4296218 100644 --- a/modules/modelSetup/QwenFineTuneSetup.py +++ b/modules/modelSetup/QwenFineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.QWEN, TrainingMethod.FINE_TUNE) class QwenFineTuneSetup( BaseQwenSetup, ): - 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: QwenModel, @@ -74,9 +60,12 @@ def setup_train_device( config.train_text_encoder_or_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/QwenLoRASetup.py b/modules/modelSetup/QwenLoRASetup.py index bf114b0ea..15e637b73 100644 --- a/modules/modelSetup/QwenLoRASetup.py +++ b/modules/modelSetup/QwenLoRASetup.py @@ -11,25 +11,11 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.QWEN, TrainingMethod.LORA) class QwenLoRASetup( BaseQwenSetup, ): - 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: QwenModel, @@ -101,9 +87,12 @@ def setup_train_device( config.train_text_encoder_or_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/SanaEmbeddingSetup.py b/modules/modelSetup/SanaEmbeddingSetup.py index fad567b75..f3812a0f4 100644 --- a/modules/modelSetup/SanaEmbeddingSetup.py +++ b/modules/modelSetup/SanaEmbeddingSetup.py @@ -9,25 +9,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.SANA, TrainingMethod.EMBEDDING) class SanaEmbeddingSetup( BaseSanaSetup, ): - 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: SanaModel, @@ -73,9 +59,10 @@ def setup_train_device( ): vae_on_train_device = self.debug_mode - model.text_encoder_to(self.train_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer", "text_encoder"] + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/SanaFineTuneSetup.py b/modules/modelSetup/SanaFineTuneSetup.py index e782e856a..bf3f69740 100644 --- a/modules/modelSetup/SanaFineTuneSetup.py +++ b/modules/modelSetup/SanaFineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.SANA, TrainingMethod.FINE_TUNE) class SanaFineTuneSetup( BaseSanaSetup, ): - 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: SanaModel, @@ -87,9 +73,12 @@ def setup_train_device( or config.train_any_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/SanaLoRASetup.py b/modules/modelSetup/SanaLoRASetup.py index 70f378a9f..73ced8854 100644 --- a/modules/modelSetup/SanaLoRASetup.py +++ b/modules/modelSetup/SanaLoRASetup.py @@ -11,25 +11,11 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.SANA, TrainingMethod.LORA) class SanaLoRASetup( BaseSanaSetup, ): - 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: SanaModel, @@ -113,9 +99,12 @@ def setup_train_device( or config.train_any_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py index af857f996..437f17b38 100644 --- a/modules/modelSetup/StableDiffusion3EmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusion3EmbeddingSetup.py @@ -9,26 +9,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.EMBEDDING) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.EMBEDDING) class StableDiffusion3EmbeddingSetup( BaseStableDiffusion3Setup, ): - 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: StableDiffusion3Model, @@ -97,11 +83,16 @@ def setup_train_device( ): vae_on_train_device = not config.latent_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) - model.text_encoder_3_to(self.train_device if config.text_encoder_3.train_embedding else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if config.text_encoder.train_embedding: + parts.append("text_encoder") + if config.text_encoder_2.train_embedding: + parts.append("text_encoder_2") + if config.text_encoder_3.train_embedding: + parts.append("text_encoder_3") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1 is not None: model.text_encoder_1.eval() diff --git a/modules/modelSetup/StableDiffusion3FineTuneSetup.py b/modules/modelSetup/StableDiffusion3FineTuneSetup.py index 0c93025ed..e3e1d74f3 100644 --- a/modules/modelSetup/StableDiffusion3FineTuneSetup.py +++ b/modules/modelSetup/StableDiffusion3FineTuneSetup.py @@ -10,26 +10,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.FINE_TUNE) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.FINE_TUNE) class StableDiffusion3FineTuneSetup( BaseStableDiffusion3Setup, ): - 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: StableDiffusion3Model, @@ -117,11 +103,16 @@ def setup_train_device( config.train_text_encoder_3_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.text_encoder_3_to(self.train_device if text_encoder_3_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if text_encoder_3_on_train_device: + parts.append("text_encoder_3") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/StableDiffusion3LoRASetup.py b/modules/modelSetup/StableDiffusion3LoRASetup.py index db851027e..4dd2168bc 100644 --- a/modules/modelSetup/StableDiffusion3LoRASetup.py +++ b/modules/modelSetup/StableDiffusion3LoRASetup.py @@ -11,26 +11,12 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_3, TrainingMethod.LORA) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_35, TrainingMethod.LORA) class StableDiffusion3LoRASetup( BaseStableDiffusion3Setup, ): - 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: StableDiffusion3Model, @@ -176,11 +162,16 @@ def setup_train_device( config.train_text_encoder_3_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.text_encoder_3_to(self.train_device if text_encoder_3_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if text_encoder_3_on_train_device: + parts.append("text_encoder_3") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if model.text_encoder_1: if config.text_encoder.train: diff --git a/modules/modelSetup/StableDiffusionEmbeddingSetup.py b/modules/modelSetup/StableDiffusionEmbeddingSetup.py index b2f339ef2..2e53f5c98 100644 --- a/modules/modelSetup/StableDiffusionEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionEmbeddingSetup.py @@ -9,8 +9,6 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.EMBEDDING) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.EMBEDDING) @@ -23,18 +21,6 @@ class StableDiffusionEmbeddingSetup( BaseStableDiffusionSetup, ): - 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: StableDiffusionModel, @@ -84,10 +70,12 @@ def setup_train_device( ): vae_on_train_device = self.debug_mode or not config.latent_caching - model.text_encoder_to(self.train_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.unet_to(self.train_device) - model.depth_estimator_to(self.temp_device) + parts = ["unet", "text_encoder"] + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) + if model.depth_estimator is not None: + model.depth_estimator.to(self.temp_device) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/StableDiffusionFineTuneSetup.py b/modules/modelSetup/StableDiffusionFineTuneSetup.py index e32008411..57e0dc632 100644 --- a/modules/modelSetup/StableDiffusionFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneSetup.py @@ -10,8 +10,6 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.FINE_TUNE) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.FINE_TUNE) @@ -24,18 +22,6 @@ class StableDiffusionFineTuneSetup( BaseStableDiffusionSetup, ): - 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: StableDiffusionModel, @@ -102,10 +88,14 @@ def setup_train_device( or config.train_any_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.unet_to(self.train_device) - model.depth_estimator_to(self.temp_device) + parts = ["unet"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) + if model.depth_estimator is not None: + model.depth_estimator.to(self.temp_device) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py b/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py index 168020b53..8f947a3d6 100644 --- a/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py +++ b/modules/modelSetup/StableDiffusionFineTuneVaeSetup.py @@ -23,18 +23,6 @@ class StableDiffusionFineTuneVaeSetup( BaseStableDiffusionSetup, ): - 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: StableDiffusionModel, @@ -67,9 +55,7 @@ def setup_train_device( model: StableDiffusionModel, config: TrainConfig, ): - model.text_encoder.to(self.temp_device) - model.vae.to(self.train_device) - model.unet.to(self.temp_device) + model.materialize_only("vae") if model.depth_estimator is not None: model.depth_estimator.to(self.temp_device) diff --git a/modules/modelSetup/StableDiffusionLoRASetup.py b/modules/modelSetup/StableDiffusionLoRASetup.py index b6c4e8746..2ecc660e9 100644 --- a/modules/modelSetup/StableDiffusionLoRASetup.py +++ b/modules/modelSetup/StableDiffusionLoRASetup.py @@ -11,8 +11,6 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15, TrainingMethod.LORA) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_15_INPAINTING, TrainingMethod.LORA) @@ -25,18 +23,6 @@ class StableDiffusionLoRASetup( BaseStableDiffusionSetup, ): - 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: StableDiffusionModel, @@ -125,10 +111,14 @@ def setup_train_device( or config.train_any_embedding() \ or not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.unet_to(self.train_device) - model.depth_estimator_to(self.temp_device) + parts = ["unet"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) + if model.depth_estimator is not None: + model.depth_estimator.to(self.temp_device) if config.text_encoder.train: model.text_encoder.train() diff --git a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py index a2fe88ac7..acf1cc5f1 100644 --- a/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py +++ b/modules/modelSetup/StableDiffusionXLEmbeddingSetup.py @@ -9,26 +9,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.EMBEDDING) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.EMBEDDING) class StableDiffusionXLEmbeddingSetup( BaseStableDiffusionXLSetup, ): - 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: StableDiffusionXLModel, @@ -87,10 +73,14 @@ def setup_train_device( ): vae_on_train_device = not config.latent_caching - model.text_encoder_1_to(self.train_device if config.text_encoder.train_embedding else self.temp_device) - model.text_encoder_2_to(self.train_device if config.text_encoder_2.train_embedding else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.unet_to(self.train_device) + parts = ["unet"] + if config.text_encoder.train_embedding: + parts.append("text_encoder") + if config.text_encoder_2.train_embedding: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder_1.eval() model.text_encoder_2.eval() diff --git a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py index 7ed64132c..0cfe94d22 100644 --- a/modules/modelSetup/StableDiffusionXLFineTuneSetup.py +++ b/modules/modelSetup/StableDiffusionXLFineTuneSetup.py @@ -10,26 +10,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.FINE_TUNE) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.FINE_TUNE) class StableDiffusionXLFineTuneSetup( BaseStableDiffusionXLSetup, ): - 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: StableDiffusionXLModel, @@ -111,10 +97,14 @@ def setup_train_device( or config.train_any_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.unet_to(self.train_device) + parts = ["unet"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder_1.train() diff --git a/modules/modelSetup/StableDiffusionXLLoRASetup.py b/modules/modelSetup/StableDiffusionXLLoRASetup.py index 78aeb7f63..afa1b0cb2 100644 --- a/modules/modelSetup/StableDiffusionXLLoRASetup.py +++ b/modules/modelSetup/StableDiffusionXLLoRASetup.py @@ -11,26 +11,12 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE, TrainingMethod.LORA) @factory.register(BaseModelSetup, ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING, TrainingMethod.LORA) class StableDiffusionXLLoRASetup( BaseStableDiffusionXLSetup, ): - 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: StableDiffusionXLModel, @@ -142,10 +128,14 @@ def setup_train_device( config.train_text_encoder_2_or_embedding() \ or not config.latent_caching - model.text_encoder_1_to(self.train_device if text_encoder_1_on_train_device else self.temp_device) - model.text_encoder_2_to(self.train_device if text_encoder_2_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.unet_to(self.train_device) + parts = ["unet"] + if text_encoder_1_on_train_device: + parts.append("text_encoder") + if text_encoder_2_on_train_device: + parts.append("text_encoder_2") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) if config.text_encoder.train: model.text_encoder_1.train() diff --git a/modules/modelSetup/WuerstchenEmbeddingSetup.py b/modules/modelSetup/WuerstchenEmbeddingSetup.py index 7d0d65506..7e6c8ee30 100644 --- a/modules/modelSetup/WuerstchenEmbeddingSetup.py +++ b/modules/modelSetup/WuerstchenEmbeddingSetup.py @@ -9,26 +9,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.WUERSTCHEN_2, TrainingMethod.EMBEDDING) @factory.register(BaseModelSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.EMBEDDING) class WuerstchenEmbeddingSetup( BaseWuerstchenSetup, ): - 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: WuerstchenModel, @@ -78,14 +64,12 @@ def setup_train_device( ): effnet_on_train_device = not config.latent_caching - if model.model_type.is_wuerstchen_v2(): - model.decoder_text_encoder_to(self.temp_device) - model.decoder_decoder_to(self.temp_device) - model.decoder_vqgan_to(self.temp_device) - model.effnet_encoder_to(self.train_device if effnet_on_train_device else self.temp_device) - - model.prior_text_encoder_to(self.train_device) - model.prior_prior_to(self.train_device) + # decoder/decoder_text_encoder/decoder_vqgan are never needed during prior training; materialize_only() + # evicts them (decoder_text_encoder only exists in model_parts() for Wuerstchen v2, not Stable Cascade) + parts = ["prior", "text_encoder"] + if effnet_on_train_device: + parts.append("effnet_encoder") + model.materialize_only(*parts) if model.model_type.is_wuerstchen_v2(): model.decoder_text_encoder.eval() diff --git a/modules/modelSetup/WuerstchenFineTuneSetup.py b/modules/modelSetup/WuerstchenFineTuneSetup.py index 4f6262d63..830709a68 100644 --- a/modules/modelSetup/WuerstchenFineTuneSetup.py +++ b/modules/modelSetup/WuerstchenFineTuneSetup.py @@ -10,26 +10,12 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.WUERSTCHEN_2, TrainingMethod.FINE_TUNE) @factory.register(BaseModelSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.FINE_TUNE) class WuerstchenFineTuneSetup( BaseWuerstchenSetup, ): - 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: WuerstchenModel, @@ -86,20 +72,19 @@ def setup_train_device( config: TrainConfig, ): effnet_on_train_device = not config.latent_caching - - if model.model_type.is_wuerstchen_v2(): - model.decoder_text_encoder_to(self.temp_device) - model.decoder_decoder_to(self.temp_device) - model.decoder_vqgan_to(self.temp_device) - model.effnet_encoder_to(self.train_device if effnet_on_train_device else self.temp_device) - text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ or not config.latent_caching - model.prior_text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.prior_prior_to(self.train_device) + # decoder/decoder_text_encoder/decoder_vqgan are never needed during prior training; materialize_only() + # evicts them (decoder_text_encoder only exists in model_parts() for Wuerstchen v2, not Stable Cascade) + parts = ["prior"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if effnet_on_train_device: + parts.append("effnet_encoder") + model.materialize_only(*parts) if model.model_type.is_wuerstchen_v2(): model.decoder_text_encoder.eval() diff --git a/modules/modelSetup/WuerstchenLoRASetup.py b/modules/modelSetup/WuerstchenLoRASetup.py index 1bdc15f3a..c9b780e9d 100644 --- a/modules/modelSetup/WuerstchenLoRASetup.py +++ b/modules/modelSetup/WuerstchenLoRASetup.py @@ -11,26 +11,12 @@ from modules.util.torch_util import state_dict_has_prefix from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.WUERSTCHEN_2, TrainingMethod.LORA) @factory.register(BaseModelSetup, ModelType.STABLE_CASCADE_1, TrainingMethod.LORA) class WuerstchenLoRASetup( BaseWuerstchenSetup, ): - 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: WuerstchenModel, @@ -113,20 +99,19 @@ def setup_train_device( config: TrainConfig, ): effnet_on_train_device = not config.latent_caching - - if model.model_type.is_wuerstchen_v2(): - model.decoder_text_encoder_to(self.temp_device) - model.decoder_decoder_to(self.temp_device) - model.decoder_vqgan_to(self.temp_device) - model.effnet_encoder_to(self.train_device if effnet_on_train_device else self.temp_device) - text_encoder_on_train_device = \ config.text_encoder.train \ or config.train_any_embedding() \ or not config.latent_caching - model.prior_text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.prior_prior_to(self.train_device) + # decoder/decoder_text_encoder/decoder_vqgan are never needed during prior training; materialize_only() + # evicts them (decoder_text_encoder only exists in model_parts() for Wuerstchen v2, not Stable Cascade) + parts = ["prior"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if effnet_on_train_device: + parts.append("effnet_encoder") + model.materialize_only(*parts) if model.model_type.is_wuerstchen_v2(): model.decoder_text_encoder.eval() diff --git a/modules/modelSetup/ZImageFineTuneSetup.py b/modules/modelSetup/ZImageFineTuneSetup.py index 6f2642de7..8fd711ae4 100644 --- a/modules/modelSetup/ZImageFineTuneSetup.py +++ b/modules/modelSetup/ZImageFineTuneSetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.Z_IMAGE, TrainingMethod.FINE_TUNE) class ZImageFineTuneSetup( BaseZImageSetup, ): - 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: ZImageModel, @@ -67,9 +53,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/modelSetup/ZImageLoRASetup.py b/modules/modelSetup/ZImageLoRASetup.py index 85cd789d6..7b360a995 100644 --- a/modules/modelSetup/ZImageLoRASetup.py +++ b/modules/modelSetup/ZImageLoRASetup.py @@ -10,25 +10,11 @@ from modules.util.optimizer_util import init_model_parameters from modules.util.TrainProgress import TrainProgress -import torch - @factory.register(BaseModelSetup, ModelType.Z_IMAGE, TrainingMethod.LORA) class ZImageLoRASetup( BaseZImageSetup, ): - 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: ZImageModel, @@ -80,9 +66,12 @@ def setup_train_device( vae_on_train_device = not config.latent_caching text_encoder_on_train_device = not config.latent_caching - model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) - model.vae_to(self.train_device if vae_on_train_device else self.temp_device) - model.transformer_to(self.train_device) + parts = ["transformer"] + if text_encoder_on_train_device: + parts.append("text_encoder") + if vae_on_train_device: + parts.append("vae") + model.materialize_only(*parts) model.text_encoder.eval() model.vae.eval() diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 57dc6baa5..380ba39d8 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -139,9 +139,8 @@ 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.evict() self.model.eval() - torch_gc() self.callbacks.on_update_status("creating the data loader/caching") @@ -254,7 +253,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.evict() self.model.eval() sample_config = copy.copy(sample_config) @@ -718,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.evict() if backup: self.__backup(train_progress, True) if save: @@ -843,7 +842,7 @@ def sample_commands_fun(): def end(self): if self.one_step_trained: - self.model.to(self.temp_device) + self.model.evict() if self.config.backup_before_save and multi.is_master(): self.__backup(self.model.train_progress) @@ -876,7 +875,7 @@ def end(self): ) if self.model is not None: - self.model.to(self.temp_device) + self.model.evict() if multi.is_master(): self.tensorboard.close() diff --git a/modules/ui/SampleWindowController.py b/modules/ui/SampleWindowController.py index d1d24d643..45a02212f 100644 --- a/modules/ui/SampleWindowController.py +++ b/modules/ui/SampleWindowController.py @@ -96,7 +96,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.evict() return model From 43824a523687fb7b283976ab57a7a67e9f77c324 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:24:51 +0200 Subject: [PATCH 02/10] Use materialize_only("vae") consistently in debug-save and VAE sampler --- modules/dataLoader/ChromaBaseDataLoader.py | 2 +- modules/dataLoader/ErnieBaseDataLoader.py | 2 +- modules/dataLoader/Flux2BaseDataLoader.py | 2 +- modules/dataLoader/FluxBaseDataLoader.py | 2 +- modules/dataLoader/HiDreamBaseDataLoader.py | 2 +- modules/dataLoader/HunyuanVideoBaseDataLoader.py | 2 +- modules/dataLoader/IdeogramBaseDataLoader.py | 2 +- modules/dataLoader/Krea2BaseDataLoader.py | 2 +- modules/dataLoader/PixArtAlphaBaseDataLoader.py | 2 +- modules/dataLoader/QwenBaseDataLoader.py | 2 +- modules/dataLoader/SanaBaseDataLoader.py | 2 +- modules/dataLoader/StableDiffusion3BaseDataLoader.py | 2 +- modules/dataLoader/StableDiffusionBaseDataLoader.py | 2 +- modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py | 2 +- modules/dataLoader/StableDiffusionXLBaseDataLoader.py | 2 +- modules/dataLoader/ZImageBaseDataLoader.py | 2 +- modules/modelSampler/StableDiffusionVaeSampler.py | 4 ++-- 17 files changed, 18 insertions(+), 18 deletions(-) diff --git a/modules/dataLoader/ChromaBaseDataLoader.py b/modules/dataLoader/ChromaBaseDataLoader.py index 660d9ecd0..b83beec1c 100644 --- a/modules/dataLoader/ChromaBaseDataLoader.py +++ b/modules/dataLoader/ChromaBaseDataLoader.py @@ -120,7 +120,7 @@ def _debug_modules(self, config: TrainConfig, model: ChromaModel): #TODO clean u debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index 3da261272..dacedd807 100644 --- a/modules/dataLoader/ErnieBaseDataLoader.py +++ b/modules/dataLoader/ErnieBaseDataLoader.py @@ -110,7 +110,7 @@ def _debug_modules(self, config: TrainConfig, model: ErnieModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index 11815e8df..bfe191457 100644 --- a/modules/dataLoader/Flux2BaseDataLoader.py +++ b/modules/dataLoader/Flux2BaseDataLoader.py @@ -117,7 +117,7 @@ def _debug_modules(self, config: TrainConfig, model: Flux2Model): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/FluxBaseDataLoader.py b/modules/dataLoader/FluxBaseDataLoader.py index 1184a3d17..e18663311 100644 --- a/modules/dataLoader/FluxBaseDataLoader.py +++ b/modules/dataLoader/FluxBaseDataLoader.py @@ -143,7 +143,7 @@ def _debug_modules(self, config: TrainConfig, model: FluxModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/HiDreamBaseDataLoader.py b/modules/dataLoader/HiDreamBaseDataLoader.py index 02a35899f..9ed754bec 100644 --- a/modules/dataLoader/HiDreamBaseDataLoader.py +++ b/modules/dataLoader/HiDreamBaseDataLoader.py @@ -180,7 +180,7 @@ def _debug_modules(self, config: TrainConfig, model: HiDreamModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/HunyuanVideoBaseDataLoader.py b/modules/dataLoader/HunyuanVideoBaseDataLoader.py index 5c0e5ece9..308864efa 100644 --- a/modules/dataLoader/HunyuanVideoBaseDataLoader.py +++ b/modules/dataLoader/HunyuanVideoBaseDataLoader.py @@ -136,7 +136,7 @@ def _debug_modules(self, config: TrainConfig, model: HunyuanVideoModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/IdeogramBaseDataLoader.py b/modules/dataLoader/IdeogramBaseDataLoader.py index f6b3c8f26..4f987dac4 100644 --- a/modules/dataLoader/IdeogramBaseDataLoader.py +++ b/modules/dataLoader/IdeogramBaseDataLoader.py @@ -122,7 +122,7 @@ def _debug_modules(self, config: TrainConfig, model: IdeogramModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/Krea2BaseDataLoader.py b/modules/dataLoader/Krea2BaseDataLoader.py index 40ee02e27..c1a3dd083 100644 --- a/modules/dataLoader/Krea2BaseDataLoader.py +++ b/modules/dataLoader/Krea2BaseDataLoader.py @@ -131,7 +131,7 @@ def _debug_modules(self, config: TrainConfig, model: Krea2Model): #TODO clean up debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/PixArtAlphaBaseDataLoader.py b/modules/dataLoader/PixArtAlphaBaseDataLoader.py index 1e7ff0e3f..5e56496d4 100644 --- a/modules/dataLoader/PixArtAlphaBaseDataLoader.py +++ b/modules/dataLoader/PixArtAlphaBaseDataLoader.py @@ -121,7 +121,7 @@ def _debug_modules(self, config: TrainConfig, model: PixArtAlphaModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/QwenBaseDataLoader.py b/modules/dataLoader/QwenBaseDataLoader.py index 944dd70a9..8e99f7cab 100644 --- a/modules/dataLoader/QwenBaseDataLoader.py +++ b/modules/dataLoader/QwenBaseDataLoader.py @@ -124,7 +124,7 @@ def _debug_modules(self, config: TrainConfig, model: QwenModel): #TODO clean up debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/dataLoader/SanaBaseDataLoader.py b/modules/dataLoader/SanaBaseDataLoader.py index 2ebad9c6f..983d8d439 100644 --- a/modules/dataLoader/SanaBaseDataLoader.py +++ b/modules/dataLoader/SanaBaseDataLoader.py @@ -113,7 +113,7 @@ def _debug_modules(self, config: TrainConfig, model: SanaModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.train_dtype.torch_dtype()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/StableDiffusion3BaseDataLoader.py b/modules/dataLoader/StableDiffusion3BaseDataLoader.py index fd929b672..b9f261ba3 100644 --- a/modules/dataLoader/StableDiffusion3BaseDataLoader.py +++ b/modules/dataLoader/StableDiffusion3BaseDataLoader.py @@ -160,7 +160,7 @@ def _debug_modules(self, config: TrainConfig, model: StableDiffusion3Model): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/StableDiffusionBaseDataLoader.py b/modules/dataLoader/StableDiffusionBaseDataLoader.py index ee6a246be..f5f45c274 100644 --- a/modules/dataLoader/StableDiffusionBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionBaseDataLoader.py @@ -130,7 +130,7 @@ def _debug_modules(self, config: TrainConfig, model: StableDiffusionModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) diff --git a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py index 7e1da8b10..a0eda2913 100644 --- a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py +++ b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py @@ -244,7 +244,7 @@ def __debug_modules(self, config: TrainConfig, model: StableDiffusionModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae) upscale_mask = ScaleImage(in_name='latent_mask', out_name='decoded_mask', factor=8) diff --git a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py index e481afba8..3a4fc9517 100644 --- a/modules/dataLoader/StableDiffusionXLBaseDataLoader.py +++ b/modules/dataLoader/StableDiffusionXLBaseDataLoader.py @@ -138,7 +138,7 @@ def _debug_modules(self, config: TrainConfig, model: StableDiffusionXLModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.vae_train_dtype.torch_dtype()) decode_conditioning_image = DecodeVAE(in_name='latent_conditioning_image', out_name='decoded_conditioning_image', vae=model.vae, autocast_contexts=[model.autocast_context, model.vae_autocast_context], dtype=model.vae_train_dtype.torch_dtype()) diff --git a/modules/dataLoader/ZImageBaseDataLoader.py b/modules/dataLoader/ZImageBaseDataLoader.py index 4ed9e1e54..2617e8561 100644 --- a/modules/dataLoader/ZImageBaseDataLoader.py +++ b/modules/dataLoader/ZImageBaseDataLoader.py @@ -116,7 +116,7 @@ def _debug_modules(self, config: TrainConfig, model: ZImageModel): debug_dir = os.path.join(config.debug_dir, "dataloader") def before_save_fun(): - model.materialize("vae") + model.materialize_only("vae") 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) diff --git a/modules/modelSampler/StableDiffusionVaeSampler.py b/modules/modelSampler/StableDiffusionVaeSampler.py index 325a1d667..8ce49dde7 100644 --- a/modules/modelSampler/StableDiffusionVaeSampler.py +++ b/modules/modelSampler/StableDiffusionVaeSampler.py @@ -63,13 +63,13 @@ def sample( image_tensor = t_in(image).to(device=self.train_device, dtype=self.model.vae.dtype) image_tensor = image_tensor * 2 - 1 - self.model.materialize("vae") + self.model.materialize_only("vae") with torch.no_grad(): latent_image_tensor = self.model.vae.encode(image_tensor.unsqueeze(0)).latent_dist.mean image_tensor = self.model.vae.decode(latent_image_tensor).sample.squeeze() - self.model.evict("vae") + self.model.evict() image_tensor = (image_tensor + 1) * 0.5 image_tensor = image_tensor.clamp(0, 1) From 35a5dfad66443b2441c588d943c619dbaec238a6 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:10:59 +0200 Subject: [PATCH 03/10] Extract shared transformer-loading helper into HFModelLoaderMixin Extracts the transformer / text-encoder / vae loading logic duplicated across the per-model loaders into shared helpers on HFModelLoaderMixin: - _load_transformer: the from_single_file(..., quantization_config= GGUFQuantizationConfig(...)) / else-load-from-repo pattern for loading an optionally-GGUF-quantized transformer checkpoint, previously duplicated across 9 loaders. - _load_text_encoder: thin wrapper over _load_transformers_sub_module giving every loader one call site for the load-on-demand streaming branch to hook. - _load_vae: collapses the duplicated "separate vae repo overrides the base vae subfolder" branch across 15 diffusers-path loaders. Flux's separate-transformer/pipeline path and HunyuanVideo's ckpt path use structurally different else-branches and were left untouched, as were the safetensors/pipeline-sourced vae branches. Also fixes Ideogram's VAE Override being silently ignored (model_names. vae_model was never read by the loader) and a HiDream text_encoder_4 override regression introduced while switching to _load_text_encoder (the override repo holds text_encoder_4 at its root with no subfolder). Pure refactor otherwise, no behavior change. Co-Authored-By: Claude Opus 4.8 --- modules/modelLoader/AnimaModelLoader.py | 54 ++++-------- modules/modelLoader/ErnieModelLoader.py | 55 ++++-------- modules/modelLoader/Flux2ModelLoader.py | 58 ++++--------- modules/modelLoader/IdeogramModelLoader.py | 24 ++---- modules/modelLoader/ZImageModelLoader.py | 53 ++++-------- .../modelLoader/chroma/ChromaModelLoader.py | 53 ++++-------- modules/modelLoader/flux/FluxModelLoader.py | 53 ++++-------- .../modelLoader/hiDream/HiDreamModelLoader.py | 32 +++---- .../hunyuanVideo/HunyuanVideoModelLoader.py | 55 ++++-------- modules/modelLoader/krea2/Krea2ModelLoader.py | 56 ++++-------- .../modelLoader/mixin/HFModelLoaderMixin.py | 86 +++++++++++++++++++ .../pixartAlpha/PixArtAlphaModelLoader.py | 24 ++---- modules/modelLoader/qwen/QwenModelLoader.py | 56 ++++-------- modules/modelLoader/sana/SanaModelLoader.py | 24 ++---- .../StableDiffusionModelLoader.py | 24 ++---- .../StableDiffusion3ModelLoader.py | 28 +++--- .../StableDiffusionXLModelLoader.py | 26 ++---- .../wuerstchen/WuerstchenModelLoader.py | 6 +- 18 files changed, 297 insertions(+), 470 deletions(-) diff --git a/modules/modelLoader/AnimaModelLoader.py b/modules/modelLoader/AnimaModelLoader.py index f4d3662bc..77f963d3c 100644 --- a/modules/modelLoader/AnimaModelLoader.py +++ b/modules/modelLoader/AnimaModelLoader.py @@ -18,7 +18,6 @@ AutoencoderKLQwenImage, CosmosTransformer3DModel, FlowMatchEulerDiscreteScheduler, - GGUFQuantizationConfig, ) from transformers import Qwen2Tokenizer, Qwen3Model, T5TokenizerFast @@ -71,7 +70,7 @@ def __load_diffusers( subfolder="scheduler", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Qwen3Model, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -86,43 +85,22 @@ def __load_diffusers( torch_dtype=torch.bfloat16, ) - if vae_model_name: #TODO simplify - vae = self._load_diffusers_sub_module( - AutoencoderKLQwenImage, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKLQwenImage, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - if transformer_model_name: - transformer = CosmosTransformer3DModel.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( - CosmosTransformer3DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + CosmosTransformer3DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + config=base_model_name, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/ErnieModelLoader.py b/modules/modelLoader/ErnieModelLoader.py index af268c25c..a9f0c9259 100644 --- a/modules/modelLoader/ErnieModelLoader.py +++ b/modules/modelLoader/ErnieModelLoader.py @@ -11,13 +11,10 @@ from modules.util.ModelNames import ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes -import torch - from diffusers import ( AutoencoderKLFlux2, ErnieImageTransformer2DModel, FlowMatchEulerDiscreteScheduler, - GGUFQuantizationConfig, ) from transformers import AutoTokenizer, Mistral3Model @@ -56,33 +53,21 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - if transformer_model_name: - transformer = ErnieImageTransformer2DModel.from_single_file( - transformer_model_name, - config=base_model_name, - subfolder="transformer", - 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( - ErnieImageTransformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + ErnieImageTransformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + config=base_model_name, + ) tokenizer = AutoTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Mistral3Model, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -95,21 +80,13 @@ def __load_diffusers( 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", - ) + vae = self._load_vae( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/Flux2ModelLoader.py b/modules/modelLoader/Flux2ModelLoader.py index 33f3fe518..3b502b8de 100644 --- a/modules/modelLoader/Flux2ModelLoader.py +++ b/modules/modelLoader/Flux2ModelLoader.py @@ -11,13 +11,10 @@ from modules.util.ModelNames import ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes -import torch - from diffusers import ( AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler, Flux2Transformer2DModel, - GGUFQuantizationConfig, ) from transformers import ( Mistral3ForConditionalGeneration, @@ -60,27 +57,14 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - if transformer_model_name: - transformer = Flux2Transformer2DModel.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( - Flux2Transformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + Flux2Transformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + config=base_model_name, + ) if transformer.config.num_attention_heads == 48: #Flux2.Dev tokenizer = PixtralProcessor.from_pretrained( @@ -88,7 +72,7 @@ def __load_diffusers( subfolder="tokenizer", ).tokenizer - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Mistral3ForConditionalGeneration, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -100,7 +84,7 @@ def __load_diffusers( base_model_name, subfolder="tokenizer", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Qwen3ForCausalLM, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -113,21 +97,13 @@ def __load_diffusers( 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", - ) + vae = self._load_vae( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/IdeogramModelLoader.py b/modules/modelLoader/IdeogramModelLoader.py index aef400f08..3858ce732 100644 --- a/modules/modelLoader/IdeogramModelLoader.py +++ b/modules/modelLoader/IdeogramModelLoader.py @@ -73,7 +73,7 @@ def __load_diffusers( else: unconditional_transformer = None - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Qwen3VLModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -91,21 +91,13 @@ def __load_diffusers( 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", - ) + vae = self._load_vae( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/ZImageModelLoader.py b/modules/modelLoader/ZImageModelLoader.py index 308232823..f4b0b8db7 100644 --- a/modules/modelLoader/ZImageModelLoader.py +++ b/modules/modelLoader/ZImageModelLoader.py @@ -11,12 +11,9 @@ from modules.util.ModelNames import ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes -import torch - from diffusers import ( AutoencoderKL, FlowMatchEulerDiscreteScheduler, - GGUFQuantizationConfig, ZImageTransformer2DModel, ) from transformers import ( @@ -68,7 +65,7 @@ def __load_diffusers( subfolder="scheduler", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Qwen3ForCausalLM, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -76,41 +73,21 @@ def __load_diffusers( "text_encoder", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - if transformer_model_name: - transformer = ZImageTransformer2DModel.from_single_file( - transformer_model_name, - #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( - ZImageTransformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + ZImageTransformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/chroma/ChromaModelLoader.py b/modules/modelLoader/chroma/ChromaModelLoader.py index 7dcbef794..4a93cc674 100644 --- a/modules/modelLoader/chroma/ChromaModelLoader.py +++ b/modules/modelLoader/chroma/ChromaModelLoader.py @@ -9,13 +9,10 @@ from modules.util.ModelNames import ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes -import torch - from diffusers import ( AutoencoderKL, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler, - GGUFQuantizationConfig, ) from transformers import T5EncoderModel, T5Tokenizer @@ -63,7 +60,7 @@ def __load_diffusers( subfolder="scheduler", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -71,41 +68,21 @@ def __load_diffusers( "text_encoder", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - if transformer_model_name: - transformer = ChromaTransformer2DModel.from_single_file( - transformer_model_name, - #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( - ChromaTransformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + ChromaTransformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/flux/FluxModelLoader.py b/modules/modelLoader/flux/FluxModelLoader.py index d4f21ea2a..02547a950 100644 --- a/modules/modelLoader/flux/FluxModelLoader.py +++ b/modules/modelLoader/flux/FluxModelLoader.py @@ -16,7 +16,6 @@ FlowMatchEulerDiscreteScheduler, FluxPipeline, FluxTransformer2DModel, - GGUFQuantizationConfig, ) from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer @@ -81,7 +80,7 @@ def __load_diffusers( ) if include_text_encoder_1: - text_encoder_1 = self._load_transformers_sub_module( + text_encoder_1 = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -92,7 +91,7 @@ def __load_diffusers( text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_transformers_sub_module( + text_encoder_2 = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder_2, weight_dtypes.fallback_train_dtype, @@ -102,41 +101,21 @@ def __load_diffusers( else: text_encoder_2 = None - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - if transformer_model_name: - transformer = FluxTransformer2DModel.from_single_file( - transformer_model_name, - #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( - FluxTransformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + FluxTransformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + ) model.model_type = model_type model.tokenizer_1 = tokenizer_1 diff --git a/modules/modelLoader/hiDream/HiDreamModelLoader.py b/modules/modelLoader/hiDream/HiDreamModelLoader.py index b3e20f23c..4c0f002bb 100644 --- a/modules/modelLoader/hiDream/HiDreamModelLoader.py +++ b/modules/modelLoader/hiDream/HiDreamModelLoader.py @@ -92,7 +92,7 @@ def __load_diffusers( ) if include_text_encoder_1: - text_encoder_1 = self._load_transformers_sub_module( + text_encoder_1 = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -103,7 +103,7 @@ def __load_diffusers( text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_transformers_sub_module( + text_encoder_2 = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder_2, weight_dtypes.train_dtype, @@ -114,7 +114,7 @@ def __load_diffusers( text_encoder_2 = None if include_text_encoder_3: - text_encoder_3 = self._load_transformers_sub_module( + text_encoder_3 = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder_3, weight_dtypes.fallback_train_dtype, @@ -126,6 +126,8 @@ def __load_diffusers( if include_text_encoder_4: if text_encoder_4_model_name: + # override repo holds text_encoder_4 at its root, not in a base-model subfolder, so it bypasses + # _load_text_encoder (which always loads from a base-repo subfolder) and loads directly text_encoder_4 = self._load_transformers_sub_module( LlamaForCausalLM, weight_dtypes.text_encoder_4, @@ -133,7 +135,7 @@ def __load_diffusers( text_encoder_4_model_name, ) else: - text_encoder_4 = self._load_transformers_sub_module( + text_encoder_4 = self._load_text_encoder( LlamaForCausalLM, weight_dtypes.text_encoder_4, weight_dtypes.train_dtype, @@ -144,21 +146,13 @@ def __load_diffusers( else: text_encoder_4 = None - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) transformer = self._load_diffusers_sub_module( HiDreamImageTransformer2DModel, diff --git a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py index 85c91699b..e1f306d58 100644 --- a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py +++ b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py @@ -81,7 +81,7 @@ def __load_diffusers( ) if include_text_encoder_1: - text_encoder_1 = self._load_transformers_sub_module( + text_encoder_1 = self._load_text_encoder( LlamaModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -92,7 +92,7 @@ def __load_diffusers( text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_transformers_sub_module( + text_encoder_2 = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder_2, weight_dtypes.fallback_train_dtype, @@ -102,43 +102,22 @@ def __load_diffusers( else: text_encoder_2 = None - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKLHunyuanVideo, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKLHunyuanVideo, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKLHunyuanVideo, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - if transformer_model_name: - transformer = HunyuanVideoTransformer3DModel.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( - HunyuanVideoTransformer3DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + HunyuanVideoTransformer3DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + config=base_model_name, + ) model.model_type = model_type model.tokenizer_1 = tokenizer_1 diff --git a/modules/modelLoader/krea2/Krea2ModelLoader.py b/modules/modelLoader/krea2/Krea2ModelLoader.py index c3987e97c..df51d2e5e 100644 --- a/modules/modelLoader/krea2/Krea2ModelLoader.py +++ b/modules/modelLoader/krea2/Krea2ModelLoader.py @@ -8,12 +8,9 @@ from modules.util.ModelNames import ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes -import torch - from diffusers import ( AutoencoderKLQwenImage, FlowMatchEulerDiscreteScheduler, - GGUFQuantizationConfig, Krea2Transformer2DModel, ) from transformers import Qwen2Tokenizer, Qwen3VLModel @@ -62,7 +59,7 @@ def __load_diffusers( subfolder="scheduler", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Qwen3VLModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -70,43 +67,22 @@ def __load_diffusers( "text_encoder", ) - if vae_model_name: #TODO simplify - vae = self._load_diffusers_sub_module( - AutoencoderKLQwenImage, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKLQwenImage, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - if transformer_model_name: - transformer = Krea2Transformer2DModel.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( - Krea2Transformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + Krea2Transformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + config=base_model_name, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/mixin/HFModelLoaderMixin.py b/modules/modelLoader/mixin/HFModelLoaderMixin.py index f2f196257..4cdc457bd 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -6,6 +6,7 @@ from modules.util.config.TrainConfig import QuantizationConfig from modules.util.enum.DataType import DataType +from modules.util.ModelWeightDtypes import ModelWeightDtypes from modules.util.quantization_util import ( is_quantized_parameter, replace_linear_with_quantized_layers, @@ -14,6 +15,7 @@ import torch from torch import nn +from diffusers import GGUFQuantizationConfig from transformers.conversion_mapping import get_checkpoint_conversion_mapping from transformers.core_model_loading import rename_source_key @@ -328,3 +330,87 @@ def _convert_diffusers_sub_module_to_dtype( None, quantization, ) + + def _load_transformer( + self, + module_type, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + transformer_model_name: str, + quantization: QuantizationConfig, + config: str | None = None, + ): + # a single-file (optionally GGUF-quantized) checkpoint is loaded directly, using + # a separate repo to source the model config if the checkpoint doesn't carry one; + # otherwise the transformer is loaded from its subfolder in the base model repo + if transformer_model_name: + single_file_kwargs = {} + if config is not None: + single_file_kwargs["config"] = config + single_file_kwargs["subfolder"] = "transformer" + + transformer = module_type.from_single_file( + transformer_model_name, + **single_file_kwargs, + #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, + ) + return self._convert_diffusers_sub_module_to_dtype( + transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization, + ) + else: + return self._load_diffusers_sub_module( + module_type, + weight_dtypes.transformer, + weight_dtypes.train_dtype, + base_model_name, + "transformer", + quantization, + ) + + def _load_text_encoder( + self, + module_type, + dtype: DataType, + train_dtype: DataType, + base_model_name: str, + subfolder: str, + ): + # text encoders have no single-file override and always load from their subfolder in the base model + # repo; kept as a per-model entry point alongside _load_transformer / _load_vae. dtype/train_dtype are + # explicit rather than a weight_dtypes bundle since a model can hold several encoders (text_encoder, + # text_encoder_2, ...) with differing dtypes + return self._load_transformers_sub_module( + module_type, + dtype, + train_dtype, + base_model_name, + subfolder, + ) + + def _load_vae( + self, + module_type, + dtype: DataType, + train_dtype: DataType, + base_model_name: str, + vae_model_name: str, + ): + # a separate vae repo overrides the base model's vae subfolder when given. train_dtype is explicit + # since some models (e.g. SDXL) upgrade the vae to fallback_train_dtype to avoid fp16 overflow + if vae_model_name: + return self._load_diffusers_sub_module( + module_type, + dtype, + train_dtype, + vae_model_name, + ) + else: + return self._load_diffusers_sub_module( + module_type, + dtype, + train_dtype, + base_model_name, + "vae", + ) diff --git a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py index 467c29c8a..69e280c9d 100644 --- a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py +++ b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py @@ -52,7 +52,7 @@ def __load_diffusers( subfolder="scheduler", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -60,21 +60,13 @@ def __load_diffusers( "text_encoder", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) transformer = self._load_diffusers_sub_module( Transformer2DModel, diff --git a/modules/modelLoader/qwen/QwenModelLoader.py b/modules/modelLoader/qwen/QwenModelLoader.py index 953f15bfb..77a12bed7 100644 --- a/modules/modelLoader/qwen/QwenModelLoader.py +++ b/modules/modelLoader/qwen/QwenModelLoader.py @@ -8,12 +8,9 @@ from modules.util.ModelNames import ModelNames from modules.util.ModelWeightDtypes import ModelWeightDtypes -import torch - from diffusers import ( AutoencoderKLQwenImage, FlowMatchEulerDiscreteScheduler, - GGUFQuantizationConfig, QwenImageTransformer2DModel, ) from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer @@ -62,7 +59,7 @@ def __load_diffusers( subfolder="scheduler", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Qwen2_5_VLForConditionalGeneration, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -70,43 +67,22 @@ def __load_diffusers( "text_encoder", ) - if vae_model_name: #TODO simplify - vae = self._load_diffusers_sub_module( - AutoencoderKLQwenImage, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKLQwenImage, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - if transformer_model_name: - transformer = QwenImageTransformer2DModel.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( - QwenImageTransformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, - base_model_name, - "transformer", - quantization, - ) + transformer = self._load_transformer( + QwenImageTransformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + config=base_model_name, + ) model.model_type = model_type model.tokenizer = tokenizer diff --git a/modules/modelLoader/sana/SanaModelLoader.py b/modules/modelLoader/sana/SanaModelLoader.py index a904e3996..ec4d31239 100644 --- a/modules/modelLoader/sana/SanaModelLoader.py +++ b/modules/modelLoader/sana/SanaModelLoader.py @@ -52,7 +52,7 @@ def __load_diffusers( subfolder="scheduler", ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( Gemma2Model, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, @@ -60,21 +60,13 @@ def __load_diffusers( "text_encoder", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderDC, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderDC, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderDC, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) transformer = self._load_diffusers_sub_module( SanaTransformer2DModel, diff --git a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py index aa610f485..223d1da61 100644 --- a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py +++ b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py @@ -87,7 +87,7 @@ def __load_diffusers( original_noise_scheduler=noise_scheduler, ) - text_encoder = self._load_transformers_sub_module( + text_encoder = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -95,21 +95,13 @@ def __load_diffusers( "text_encoder", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) unet = self._load_diffusers_sub_module( UNet2DConditionModel, diff --git a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py index 47d87da74..0c4f4348a 100644 --- a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py +++ b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py @@ -81,7 +81,7 @@ def __load_diffusers( ) if include_text_encoder_1: - text_encoder_1 = self._load_transformers_sub_module( + text_encoder_1 = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -92,7 +92,7 @@ def __load_diffusers( text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_transformers_sub_module( + text_encoder_2 = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder_2, weight_dtypes.train_dtype, @@ -103,7 +103,7 @@ def __load_diffusers( text_encoder_2 = None if include_text_encoder_3: - text_encoder_3 = self._load_transformers_sub_module( + text_encoder_3 = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder_3, weight_dtypes.fallback_train_dtype, @@ -113,21 +113,13 @@ def __load_diffusers( else: text_encoder_3 = None - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) transformer = self._load_diffusers_sub_module( SD3Transformer2DModel, diff --git a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py index afbab6581..ba340006a 100644 --- a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py +++ b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py @@ -83,7 +83,7 @@ def __load_diffusers( original_noise_scheduler=noise_scheduler, ) - text_encoder_1 = self._load_transformers_sub_module( + text_encoder_1 = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -91,7 +91,7 @@ def __load_diffusers( "text_encoder", ) - text_encoder_2 = self._load_transformers_sub_module( + text_encoder_2 = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder_2, weight_dtypes.train_dtype, @@ -99,21 +99,13 @@ def __load_diffusers( "text_encoder_2", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.fallback_train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.fallback_train_dtype, - base_model_name, - "vae", - ) + vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.fallback_train_dtype, + base_model_name, + vae_model_name, + ) unet = self._load_diffusers_sub_module( UNet2DConditionModel, diff --git a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py index 188107a2c..d8e86a19b 100644 --- a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py +++ b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py @@ -75,7 +75,7 @@ def __load_diffusers( ) if model_type.is_wuerstchen_v2(): - decoder_text_encoder = self._load_transformers_sub_module( + decoder_text_encoder = self._load_text_encoder( CLIPTextModel, weight_dtypes.decoder_text_encoder, weight_dtypes.train_dtype, @@ -164,7 +164,7 @@ def __load_diffusers( ) if model_type.is_wuerstchen_v2(): - prior_text_encoder = self._load_transformers_sub_module( + prior_text_encoder = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -172,7 +172,7 @@ def __load_diffusers( "text_encoder", ) elif model_type.is_stable_cascade(): - prior_text_encoder = self._load_transformers_sub_module( + prior_text_encoder = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder, weight_dtypes.train_dtype, From 90f1616af4300fcb3f54c937e9aaef41853c4d10 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:18:55 +0200 Subject: [PATCH 04/10] Arena MemPool + single-buffer offload cache On CUDA, the layer-offload cache now uses one cache tensor instead of the multi-chunk split: a large cuda allocation is page-mapped, so one buffer packs with no inter-chunk tail waste, and the arena is filled per-layer from the CPU so no full resident source coexists with it. The host/pinned cache keeps the lazy multi-chunk split (its peak-doubling justification is host-only). Each layer-offload cache tensor, and each BaseModel component move, gets its own dedicated torch.cuda.MemPool, so the churny small tensors in the default pool can't wedge into a freed cache/component segment and strand it across an evict/reload cycle -- the cross-cycle fragmentation OOM on a tight budget. Pools are released once their tensors are freed. Shared MemPool helpers (create_mem_pool, mem_pool_context, supports_mem_pool) live in torch_util so both BaseModel._move_part and the offload conductor use the same wrapper. The alignment budget for the offload cache is sized from the actual offload-tensor count (TENSOR_ALIGNMENT_BYTES per tensor) instead of a fixed 4KB, since the unguarded ring wrap would otherwise silently overwrite live weights once a cache tensor holds enough tensors. Co-Authored-By: Claude Opus 4.8 --- modules/model/BaseModel.py | 17 +++++- modules/util/LayerOffloadConductor.py | 87 ++++++++++++++++++++------- modules/util/torch_util.py | 14 +++++ 3 files changed, 94 insertions(+), 24 deletions(-) diff --git a/modules/model/BaseModel.py b/modules/model/BaseModel.py index d7227cced..2a76e576a 100644 --- a/modules/model/BaseModel.py +++ b/modules/model/BaseModel.py @@ -11,7 +11,7 @@ from modules.util.enum.ModelType import ModelType from modules.util.modelSpec.ModelSpec import ModelSpec from modules.util.NamedParameterGroup import NamedParameterGroupCollection -from modules.util.torch_util import torch_gc +from modules.util.torch_util import create_mem_pool, mem_pool_context, supports_mem_pool, torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -97,6 +97,8 @@ def __init__( self.autocast_context = nullcontext() self.train_dtype = DataType.FLOAT_32 + self._mem_pools = {} + @property def train_device(self) -> torch.device: return torch.device(self.train_config.train_device) @@ -144,7 +146,18 @@ def _move_part(self, part: str, device: torch.device): # None when the part is excluded from training (e.g. a text encoder with include_text_encoder off): # it stays in model_parts() but the loader never populated it, so there is nothing to move. if component is not None: - component.to(device=device) + # give each component its own MemPool to avoid fragmenting the next reload + if supports_mem_pool(device): + pool = self._mem_pools.get(stem) + if pool is None: + pool = self._mem_pools[stem] = create_mem_pool(device) + with mem_pool_context(pool): + component.to(device=device) + else: + # the target has no MemPool (CPU): move normally and drop this component's pool from the + # earlier GPU move, so evict()'s torch_gc can release its segments + component.to(device=device) + self._mem_pools.pop(stem, None) lora = getattr(self, f"{stem}_lora", None) if lora is not None: diff --git a/modules/util/LayerOffloadConductor.py b/modules/util/LayerOffloadConductor.py index a69094c75..54dcef3d6 100644 --- a/modules/util/LayerOffloadConductor.py +++ b/modules/util/LayerOffloadConductor.py @@ -3,11 +3,13 @@ from typing import Any from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig -from modules.util.quantization_util import get_offload_tensor_bytes, offload_quantized +from modules.util.quantization_util import get_offload_tensor_bytes, get_offload_tensors, offload_quantized from modules.util.torch_util import ( + create_mem_pool, create_stream_context, device_equals, get_tensor_data, + mem_pool_context, pin_tensor_, replace_tensors_, tensors_match_device, @@ -34,12 +36,16 @@ def clone_tensor_allocator(tensor: torch.Tensor) -> torch.Tensor: return tensor.clone() -def ceil_16(number: int) -> int: - return number + (16 - (number % 16)) % 16 +# allocate_like places each cached tensor at an aligned offset, wasting up to this many bytes per tensor +TENSOR_ALIGNMENT_BYTES = 16 -def floor_16(number: int) -> int: - return number - (number % 16) +def align_up(number: int) -> int: + return number + (TENSOR_ALIGNMENT_BYTES - (number % TENSOR_ALIGNMENT_BYTES)) % TENSOR_ALIGNMENT_BYTES + + +def align_down(number: int) -> int: + return number - (number % TENSOR_ALIGNMENT_BYTES) class StaticLayerTensorAllocator: @@ -69,7 +75,7 @@ def allocate_like(self, source_tensor: torch.Tensor) -> torch.Tensor: total_cache_bytes = cache_tensor_size * len(self.__layer_allocator.cache_tensors) if self.__allocate_forward: cache_tensor_index = self.__allocation_end // cache_tensor_size - cache_tensor_allocation_end = ceil_16(self.__allocation_end % cache_tensor_size) + cache_tensor_allocation_end = align_up(self.__allocation_end % cache_tensor_size) if cache_tensor_allocation_end + num_bytes > cache_tensor_size: # move to the start of the next cache tensor @@ -100,7 +106,7 @@ def allocate_like(self, source_tensor: torch.Tensor) -> torch.Tensor: cache_tensor_index = len(self.__layer_allocator.cache_tensors) - 1 cache_tensor_allocation_start = cache_tensor_size - new_allocation_start = floor_16(cache_tensor_allocation_start - num_bytes) + new_allocation_start = align_down(cache_tensor_allocation_start - num_bytes) self.__layer_allocator.ensure_allocation(cache_tensor_index) cache_tensor = self.__layer_allocator.cache_tensors[cache_tensor_index] allocated_tensor = cache_tensor[new_allocation_start:new_allocation_start + num_bytes] @@ -153,31 +159,55 @@ def __init__( self.__tensor_allocators = [] + self.__mem_pool = None + def allocate_cache(self, layers: list[nn.Module], target_bytes: int): if not self.__allocate_statically or any(x is not None for x in self.cache_tensors): return log(f"allocating cache on device {self.device}") + # keep the cache tensor in its own MemPool to avoid fragmenting the next cycle's allocation + if self.__mem_pool is None: + self.__mem_pool = create_mem_pool(self.device) + self.__max_tensor_bytes = 0 self.__layer_bytes = [] + total_tensors = 0 # count of individual offload tensors == number of allocate_like calls == alignment slots for layer in layers: layer_tensor_bytes = [get_offload_tensor_bytes(x) for x in layer.modules()] + total_tensors += sum(len(get_offload_tensors(x)) for x in layer.modules()) self.__max_tensor_bytes = max(self.__max_tensor_bytes, *layer_tensor_bytes) self.__layer_bytes.append(sum(layer_tensor_bytes)) cache_bytes = target_bytes - num_cache_tensors = min( - # no more than 10% overhead - math.ceil(int(cache_bytes * 0.10) / self.__max_tensor_bytes), - # at least twice self.__max_tensor_bytes for each tensor - math.ceil(cache_bytes / (self.__max_tensor_bytes * 2)), - # no more than 10 cache tensors - 10 - ) - # add self.__max_tensor_bytes to ensure even the largest tensors can be allocated in the remaining space - # add 4kb for the alignment overhead - self.cache_tensor_size = math.ceil(cache_bytes / num_cache_tensors) + self.__max_tensor_bytes + 4096 + if self.device.type == "cuda": + # single cache tensor on the GPU: a large cuda allocation is page-mapped (assembled from scattered + # physical pages), so one buffer allocates as readily as many and packs with no inter-chunk tail waste. + # The GPU cache is filled one layer at a time from the CPU, so the destination buffer and a full + # resident source never coexist on the device -- no peak-doubling to guard against here. + num_cache_tensors = 1 + else: + # host/pinned cache keeps the multi-chunk split: the chunks are allocated lazily (per ensure_allocation) + # to cap peak host RAM while the resident model is copied into the pinned cache, which a single eager + # buffer would roughly double. + # TODO once the disk-streaming (load_on_demand) path lands here: collapse to a single buffer too when + # the layers are streamed per-part from the checkpoint -- then no resident model copy coexists with the + # pinned cache, so the peak-doubling that justifies chunking here does not occur. + num_cache_tensors = min( + # no more than 10% overhead + math.ceil(int(cache_bytes * 0.10) / self.__max_tensor_bytes), + # at least twice self.__max_tensor_bytes for each tensor + math.ceil(cache_bytes / (self.__max_tensor_bytes * 2)), + # no more than 10 cache tensors + 10 + ) + # the alignment budget must cover EVERY tensor packed into a cache tensor: allocate_like wastes up to + # TENSOR_ALIGNMENT_BYTES per tensor and the ring wrap is unguarded, so a fixed total would silently + # overwrite live weights once a cache tensor holds enough tensors. Size it from the actual tensor count. + alignment_bytes = TENSOR_ALIGNMENT_BYTES * total_tensors + # add self.__max_tensor_bytes so even the largest tensor fits in the space left after a ring wrap + self.cache_tensor_size = math.ceil(cache_bytes / num_cache_tensors) + self.__max_tensor_bytes + alignment_bytes self.__tensor_allocators = [None] * len(layers) self.cache_tensors = [None] * num_cache_tensors @@ -188,8 +218,12 @@ def ensure_allocation(self, cache_tensor_index: int): if self.cache_tensors[cache_tensor_index] is None: torch_gc() - self.cache_tensors[cache_tensor_index] = \ - torch.zeros((self.cache_tensor_size,), dtype=torch.int8, device=self.device) + # create the cache tensor inside the MemPool so it lands in the pool's isolated segments. the buffers + # are allocated lazily here (allocate_cache only sizes them), so the pool context wraps this + # allocation rather than allocate_cache. + with mem_pool_context(self.__mem_pool): + self.cache_tensors[cache_tensor_index] = \ + torch.zeros((self.cache_tensor_size,), dtype=torch.int8, device=self.device) log(f"tensor {cache_tensor_index} not allocated, allocating {self.cache_tensor_size} bytes") @@ -206,6 +240,15 @@ def deallocate_cache(self): self.cache_tensors = [None] * len(self.cache_tensors) self.__tensor_allocators = [None] * len(self.__tensor_allocators) + # the loop above leaves `cache_tensor` bound to the last tensor; clear it so that stray reference can't + # keep the MemPool alive through the torch_gc below + cache_tensor = None + + # drop the MemPool once its tensors are freed so its now-empty segments return to the driver for the + # default pool; a fresh one is created on the next allocate_cache. + if self.__mem_pool is not None: + self.__mem_pool = None + torch_gc() def get_allocator(self, layer_index: int, allocate_forward: bool) -> StaticLayerTensorAllocator | None: if self.__allocate_statically: @@ -248,7 +291,7 @@ def __init__( def reserve_cache(self, tensors: list[torch.Tensor]): num_bytes = sum(tensor.element_size() * tensor.numel() for tensor in tensors) \ - + len(tensors) * 16 # add enough padding for alignment + + len(tensors) * TENSOR_ALIGNMENT_BYTES # add enough padding for alignment if num_bytes == 0: return @@ -284,7 +327,7 @@ def allocate_like(self, source_tensor: torch.Tensor) -> torch.Tensor: cache_tensor = self.__cache_tensors[self.__current_cache_tensor] allocated_tensor = \ cache_tensor[self.__current_cache_tensor_offset:self.__current_cache_tensor_offset + num_bytes] - self.__current_cache_tensor_offset += ceil_16(num_bytes) + self.__current_cache_tensor_offset += align_up(num_bytes) return allocated_tensor.view(dtype=source_tensor.dtype).view(size=source_tensor.shape) diff --git a/modules/util/torch_util.py b/modules/util/torch_util.py index 408100bf9..4d199f282 100644 --- a/modules/util/torch_util.py +++ b/modules/util/torch_util.py @@ -15,6 +15,20 @@ torch_version = packaging.version.parse(torch.__version__) +def supports_mem_pool(device: torch.device) -> bool: + return device.type == "cuda" + + +def create_mem_pool(device: torch.device): + # a dedicated MemPool the caller can allocate into; None on devices without MemPool support (cpu/mps) + return torch.cuda.MemPool() if supports_mem_pool(device) else None + + +def mem_pool_context(mem_pool): + # route allocations made in this context into the given MemPool; no-op when it is None + return torch.cuda.use_mem_pool(mem_pool) if mem_pool is not None else nullcontext() + + def state_dict_has_prefix(state_dict: dict | None, prefix: str): if not state_dict: return False From e27dc4c92f9ccb104df1eaeec9b35402d5e8e34a Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:19:16 +0200 Subject: [PATCH 05/10] Pool each component's LoRA in a per-stem MemPool Squashed history of the mempool branch on top of PR #1620 (Arena MemPool + single-buffer offload cache / materialize-evict API): eviction handling for multi-TE samplers, materialize_only_text_encoders() helper, generic BaseModel.eval()/adapters(), per-stem LoRA pooling, and review cleanups. --- modules/model/BaseModel.py | 44 ++++++++++++++++++++++---------------- 1 file changed, 25 insertions(+), 19 deletions(-) diff --git a/modules/model/BaseModel.py b/modules/model/BaseModel.py index 2a76e576a..3b172461e 100644 --- a/modules/model/BaseModel.py +++ b/modules/model/BaseModel.py @@ -139,29 +139,35 @@ def _move_part(self, part: str, device: torch.device): stem = f"{part}_1" if hasattr(self, f"{part}_1") else part conductor = getattr(self, f"{stem}_offload_conductor", None) + lora = getattr(self, f"{stem}_lora", None) + component = None if conductor is not None else getattr(self, stem) + if conductor is not None: conductor.to(device) + + if component is None and lora is None: + return + + if supports_mem_pool(device): + # The component (when not conductor-managed) and its LoRA share a per-stem MemPool so both land + # contiguously and release together on evict. A conductor keeps its own pool, so the stem pool then + # holds only the LoRA, keeping its small tensors out of the default pool across the part's evict/reload. + pool = self._mem_pools.get(stem) + if pool is None: + pool = self._mem_pools[stem] = create_mem_pool(device) + with mem_pool_context(pool): + if component is not None: + component.to(device=device) + if lora is not None: + lora.to(device=device) else: - component = getattr(self, stem) # raises if `part` doesn't name a real attribute - # None when the part is excluded from training (e.g. a text encoder with include_text_encoder off): - # it stays in model_parts() but the loader never populated it, so there is nothing to move. + # the target has no MemPool (CPU): move normally and drop this stem's pool from the earlier GPU move, + # so evict()'s torch_gc can release its segments if component is not None: - # give each component its own MemPool to avoid fragmenting the next reload - if supports_mem_pool(device): - pool = self._mem_pools.get(stem) - if pool is None: - pool = self._mem_pools[stem] = create_mem_pool(device) - with mem_pool_context(pool): - component.to(device=device) - else: - # the target has no MemPool (CPU): move normally and drop this component's pool from the - # earlier GPU move, so evict()'s torch_gc can release its segments - component.to(device=device) - self._mem_pools.pop(stem, None) - - lora = getattr(self, f"{stem}_lora", None) - if lora is not None: - lora.to(device) + component.to(device=device) + if lora is not None: + lora.to(device=device) + self._mem_pools.pop(stem, None) def eval(self): # Put every present component on eval(); driven by the same part registry as materialize()/evict(). From ed56354bd9fb7918ead1ade4e423b6a4fbbdeeda Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:33:36 +0200 Subject: [PATCH 06/10] [centralize-setup-optimizations] Extract shared setup_optimizations helper into BaseModelSetup --- modules/modelSetup/BaseAnimaSetup.py | 26 ++--------- modules/modelSetup/BaseChromaSetup.py | 26 ++--------- modules/modelSetup/BaseErnieSetup.py | 25 ++-------- modules/modelSetup/BaseFlux2Setup.py | 30 +++--------- modules/modelSetup/BaseFluxSetup.py | 30 ++---------- modules/modelSetup/BaseHiDreamSetup.py | 46 +++---------------- modules/modelSetup/BaseHunyuanVideoSetup.py | 29 ++---------- modules/modelSetup/BaseIdeogramSetup.py | 40 +++------------- modules/modelSetup/BaseKrea2Setup.py | 25 ++-------- modules/modelSetup/BaseModelSetup.py | 44 +++++++++++++++++- modules/modelSetup/BasePixArtAlphaSetup.py | 23 ++-------- modules/modelSetup/BaseQwenSetup.py | 25 ++-------- modules/modelSetup/BaseSanaSetup.py | 27 ++++------- .../modelSetup/BaseStableDiffusion3Setup.py | 33 +++---------- .../modelSetup/BaseStableDiffusionSetup.py | 8 ++-- .../modelSetup/BaseStableDiffusionXLSetup.py | 9 ++-- modules/modelSetup/BaseWuerstchenSetup.py | 8 ++-- modules/modelSetup/BaseZImageSetup.py | 25 ++-------- modules/util/checkpointing_util.py | 13 +++--- modules/util/dtype_util.py | 8 ++-- 20 files changed, 141 insertions(+), 359 deletions(-) diff --git a/modules/modelSetup/BaseAnimaSetup.py b/modules/modelSetup/BaseAnimaSetup.py index 8c0b9e0aa..763462696 100644 --- a/modules/modelSetup/BaseAnimaSetup.py +++ b/modules/modelSetup/BaseAnimaSetup.py @@ -14,8 +14,6 @@ enable_checkpointing_for_qwen_transformer, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -23,7 +21,6 @@ from torch import Tensor -#TODO share more code with other models class BaseAnimaSetup( BaseModelSetup, ModelSetupDiffusionLossMixin, @@ -46,25 +43,10 @@ def setup_optimizations( model: AnimaModel, config: TrainConfig, ): - model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer) - 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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=False) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_qwen_transformer, attention_mask=False) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_qwen3_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def predict( self, diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 11ab09c82..c4d6df734 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -16,15 +16,12 @@ enable_checkpointing_for_t5_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch from torch import Tensor -#TODO share more code with Flux and other models class BaseChromaSetup( BaseModelSetup, ModelSetupDiffusionLossMixin, @@ -47,25 +44,10 @@ def setup_optimizations( model: ChromaModel, config: TrainConfig, ): - model.transformer_offload_conductor = enable_checkpointing_for_chroma_transformer(model.transformer, config, config.transformer) - 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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_chroma_transformer, attention_mask=True) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_t5_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def _setup_embeddings( self, diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index d5050c49f..379747ba1 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -14,8 +14,6 @@ enable_checkpointing_for_mistral_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -43,25 +41,10 @@ def setup_optimizations( model: ErnieModel, config: TrainConfig, ): - model.transformer_offload_conductor = enable_checkpointing_for_ernie_transformer(model.transformer, config, config.transformer) - model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_ernie_transformer, attention_mask=True) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_mistral_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def predict( self, diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index 63e8dfc70..a7509e77c 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -16,8 +16,6 @@ enable_checkpointing_for_qwen3_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -43,28 +41,12 @@ def setup_optimizations( model: Flux2Model, config: TrainConfig, ): - model.transformer_offload_conductor = enable_checkpointing_for_flux2_transformer(model.transformer, config, config.transformer) - 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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=False) + text_encoder_checkpointing_fn = enable_checkpointing_for_mistral_encoder_layers if model.is_dev() \ + else enable_checkpointing_for_qwen3_encoder_layers + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_flux2_transformer, attention_mask=False) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, text_encoder_checkpointing_fn, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def predict( self, diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 398d4a1b4..2bb5619dd 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -17,8 +17,6 @@ enable_checkpointing_for_t5_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress @@ -48,29 +46,11 @@ def setup_optimizations( model: FluxModel, config: TrainConfig, ): - 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: - 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.enable_autocast_cache) - - model.text_encoder_2_autocast_context, model.text_encoder_2_train_dtype = \ - disable_fp16_autocast_context( - self.train_device, - config.train_dtype, - config.fallback_train_dtype, - config.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder_1, self.train_device, model.train_dtype, config) - quantize_layers(model.text_encoder_2, self.train_device, model.text_encoder_2_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=False) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_flux_transformer, attention_mask=False) + self._setup_model_part(model, config, "text_encoder_1", config.text_encoder, enable_checkpointing_for_clip_encoder_layers) + self._setup_model_part(model, config, "text_encoder_2", config.text_encoder_2, enable_checkpointing_for_t5_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def _setup_embeddings( self, diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 301574535..71cbe8b24 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -17,8 +17,6 @@ enable_checkpointing_for_t5_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -47,43 +45,13 @@ def setup_optimizations( model: HiDreamModel, config: TrainConfig, ): - 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: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) - 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: - 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.enable_autocast_cache) - - model.text_encoder_3_autocast_context, model.text_encoder_3_train_dtype = \ - disable_fp16_autocast_context( - self.train_device, - config.train_dtype, - config.fallback_train_dtype, - config.enable_autocast_cache, - ) - - model.transformer_autocast_context, model.transformer_train_dtype = \ - disable_fp16_autocast_context( - self.train_device, - config.train_dtype, - config.fallback_train_dtype, - config.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder_1, self.train_device, model.train_dtype, config) - quantize_layers(model.text_encoder_2, self.train_device, model.train_dtype, config) - quantize_layers(model.text_encoder_3, self.train_device, model.text_encoder_3_train_dtype, config) - quantize_layers(model.text_encoder_4, self.train_device, model.train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.transformer_train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_hi_dream_transformer, disable_fp16_autocast=True, attention_mask=True) + self._setup_model_part(model, config, "text_encoder_1", config.text_encoder, enable_checkpointing_for_clip_encoder_layers) + self._setup_model_part(model, config, "text_encoder_2", config.text_encoder_2, enable_checkpointing_for_clip_encoder_layers) + self._setup_model_part(model, config, "text_encoder_3", config.text_encoder_3, enable_checkpointing_for_t5_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "text_encoder_4", config.text_encoder_4, enable_checkpointing_for_llama_encoder_layers) + self._setup_model_part(model, config, "vae", config.vae) def _setup_embeddings( self, diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index 4f5713008..a2749fef7 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -17,8 +17,6 @@ enable_checkpointing_for_llama_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -47,30 +45,13 @@ def setup_optimizations( model: HunyuanVideoModel, config: TrainConfig, ): - 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: - 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.enable_autocast_cache) - - model.transformer_autocast_context, model.transformer_train_dtype = \ - disable_fp16_autocast_context( - self.train_device, - config.train_dtype, - config.fallback_train_dtype, - config.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder_1, self.train_device, model.train_dtype, config) - quantize_layers(model.text_encoder_2, self.train_device, model.train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.transformer_train_dtype, config) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_hunyuan_video_transformer, disable_fp16_autocast=True, attention_mask=True) + self._setup_model_part(model, config, "text_encoder_1", config.text_encoder, enable_checkpointing_for_llama_encoder_layers) + self._setup_model_part(model, config, "text_encoder_2", config.text_encoder_2, enable_checkpointing_for_clip_encoder_layers) + self._setup_model_part(model, config, "vae", config.vae) model.vae.enable_tiling() - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) def _setup_embeddings( self, diff --git a/modules/modelSetup/BaseIdeogramSetup.py b/modules/modelSetup/BaseIdeogramSetup.py index 7fb1a0bff..966a4eb2f 100644 --- a/modules/modelSetup/BaseIdeogramSetup.py +++ b/modules/modelSetup/BaseIdeogramSetup.py @@ -14,8 +14,6 @@ enable_checkpointing_for_qwen3vl_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -43,37 +41,13 @@ def setup_optimizations( model: IdeogramModel, config: TrainConfig, ): - # Only the conditional transformer is trained, so gradient checkpointing applies there. - model.transformer_offload_conductor = \ - enable_checkpointing_for_ideogram_transformer(model.transformer, config, config.transformer) - - # The unconditional transformer is frozen, but it still benefits from layer offloading - # since both transformers need to fit in VRAM during sampling. It is optional, so may be unloaded. - if model.unconditional_transformer is not None: - model.unconditional_transformer_offload_conductor = \ - enable_checkpointing_for_ideogram_transformer(model.unconditional_transformer, config, config.unconditional_transformer) - - model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3vl_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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - quantize_layers(model.unconditional_transformer, self.train_device, model.train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=False) - if model.unconditional_transformer is not None: - self._set_attention_backend(model.unconditional_transformer, config.attention_mechanism, mask=False) + super().setup_optimizations(model, config) + # The unconditional transformer is frozen but still layer-offloaded so both transformers fit in VRAM + # during sampling; it is optional, so _setup_model_part skips it when unloaded. + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_ideogram_transformer, attention_mask=False) + self._setup_model_part(model, config, "unconditional_transformer", config.unconditional_transformer, enable_checkpointing_for_ideogram_transformer, attention_mask=False) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_qwen3vl_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def predict( self, diff --git a/modules/modelSetup/BaseKrea2Setup.py b/modules/modelSetup/BaseKrea2Setup.py index e4249e897..9701a4eb0 100644 --- a/modules/modelSetup/BaseKrea2Setup.py +++ b/modules/modelSetup/BaseKrea2Setup.py @@ -14,8 +14,6 @@ enable_checkpointing_for_qwen3vl_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -45,25 +43,10 @@ def setup_optimizations( model: Krea2Model, config: TrainConfig, ): - model.transformer_offload_conductor = enable_checkpointing_for_krea2_transformer(model.transformer, config, config.transformer) - model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3vl_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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_krea2_transformer, attention_mask=True) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_qwen3vl_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def predict( self, diff --git a/modules/modelSetup/BaseModelSetup.py b/modules/modelSetup/BaseModelSetup.py index 8a6818908..fea9c94b9 100644 --- a/modules/modelSetup/BaseModelSetup.py +++ b/modules/modelSetup/BaseModelSetup.py @@ -3,10 +3,12 @@ from modules.model.BaseModel import BaseModel from modules.util.config.TrainConfig import TrainConfig, TrainEmbeddingConfig, TrainModelPartConfig +from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context from modules.util.enum.AttentionMechanism import AttentionMechanism from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.ModuleFilter import ModuleFilter from modules.util.NamedParameterGroup import NamedParameterGroup, NamedParameterGroupCollection +from modules.util.quantization_util import quantize_layers from modules.util.TimedActionMixin import TimedActionMixin from modules.util.TrainProgress import TrainProgress @@ -47,7 +49,10 @@ def setup_optimizations( model: BaseModel, config: TrainConfig, ): - pass + # Model-wide dtype/autocast, shared by every leaf. Leaves call super() first so model.train_dtype is + # set before their first _setup_model_part, which reads it for the non-fp16 quantize path. + model.train_dtype = config.train_dtype + model.autocast_context = create_autocast_context(self.train_device, config.train_dtype, config.enable_autocast_cache) @abstractmethod def setup_model( @@ -237,6 +242,43 @@ def _setup_model_part_requires_grad( for param in self.frozen_parameters[unique_name]: param.requires_grad_(False) + def _setup_model_part( + self, + model, + config: TrainConfig, + attr: str, + config_part: TrainModelPartConfig, + checkpointing_fn=None, + *, + disable_fp16_autocast: bool = False, + attention_mask: bool | None = None, + ): + # Per-part optimization wiring, called once per model part from each leaf. The optional + # disable_fp16_autocast context and its dtype are stored per-part and can differ per part + # (e.g. HiDream disables fp16 for both text_encoder_3 and the transformer). checkpointing_fn returns + # None for non-offloadable parts (SD/SDXL UNet), so no conductor is stored for those. + module = getattr(model, attr) + if module is None: + return + + if checkpointing_fn is not None: + conductor = checkpointing_fn(module, config, config_part) + if conductor is not None: + setattr(model, f"{attr}_offload_conductor", conductor) + + if disable_fp16_autocast: + autocast_context, train_dtype = disable_fp16_autocast_context( + self.train_device, config.train_dtype, config.fallback_train_dtype, config.enable_autocast_cache) + setattr(model, f"{attr}_autocast_context", autocast_context) + setattr(model, f"{attr}_train_dtype", train_dtype) + else: + train_dtype = model.train_dtype + + quantize_layers(module, self.train_device, train_dtype, config) + + if attention_mask is not None: + self._set_attention_backend(module, config.attention_mechanism, mask=attention_mask) + @staticmethod def _set_attention_backend(component, attn: AttentionMechanism, mask: bool): match attn: diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 008e9a0c0..875bc4a8e 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -16,8 +16,6 @@ enable_checkpointing_for_t5_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -49,23 +47,10 @@ def setup_optimizations( model: PixArtAlphaModel, config: TrainConfig, ): - 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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_basic_transformer_blocks, attention_mask=True) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_t5_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def _setup_embeddings( self, diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index 0615a3c1e..ee42de6b2 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -14,15 +14,12 @@ enable_checkpointing_for_qwen_transformer, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch from torch import Tensor -#TODO share more code with other models class BaseQwenSetup( BaseModelSetup, ModelSetupDiffusionLossMixin, @@ -44,24 +41,10 @@ def setup_optimizations( model: QwenModel, config: TrainConfig, ): - model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer) - 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.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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_qwen_transformer, attention_mask=True) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_qwen25vl_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def predict( self, diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 8718142a7..0c979aa32 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -16,8 +16,7 @@ enable_checkpointing_for_sana_transformer, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers +from modules.util.dtype_util import disable_fp16_autocast_context from modules.util.TrainProgress import TrainProgress import torch @@ -50,19 +49,12 @@ def setup_optimizations( config: TrainConfig, ): - 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.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.enable_autocast_cache, - ) + super().setup_optimizations(model, config) + # Sana's vae runs under its own fp16-disabled autocast in predict(), set here rather than via + # _setup_model_part's autocast handling. Note a preexisting inconsistency (predates this refactor, + # kept as-is since Sana is largely outdated): unlike SDXL, model.vae_train_dtype is computed but never + # read anywhere, and the vae below is quantized with model.train_dtype, not this fp16-disabled dtype. model.vae_autocast_context, model.vae_train_dtype = disable_fp16_autocast_context( self.train_device, config.train_dtype, @@ -70,10 +62,9 @@ def setup_optimizations( config.enable_autocast_cache, ) - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_sana_transformer, attention_mask=True) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_gemma_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def _setup_embeddings( self, diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index dd7cf0841..9fd69ccf6 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -17,8 +17,6 @@ enable_checkpointing_for_t5_encoder_layers, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -46,31 +44,12 @@ def setup_optimizations( model: StableDiffusion3Model, config: TrainConfig, ): - 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: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) - 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, config.enable_autocast_cache) - - model.text_encoder_3_autocast_context, model.text_encoder_3_train_dtype = \ - disable_fp16_autocast_context( - self.train_device, - config.train_dtype, - config.fallback_train_dtype, - config.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder_1, self.train_device, model.train_dtype, config) - quantize_layers(model.text_encoder_2, self.train_device, model.train_dtype, config) - quantize_layers(model.text_encoder_3, self.train_device, model.text_encoder_3_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=False) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_stable_diffusion_3_transformer, attention_mask=False) + self._setup_model_part(model, config, "text_encoder_1", config.text_encoder, enable_checkpointing_for_clip_encoder_layers) + self._setup_model_part(model, config, "text_encoder_2", config.text_encoder_2, enable_checkpointing_for_clip_encoder_layers) + self._setup_model_part(model, config, "text_encoder_3", config.text_encoder_3, enable_checkpointing_for_t5_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def _setup_embeddings( self, diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 7a3f910c8..56179aff1 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -17,7 +17,6 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.conv_util import apply_circular_padding_to_conv2d -from modules.util.dtype_util import create_autocast_context from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress @@ -49,9 +48,11 @@ def setup_optimizations( model: StableDiffusionModel, config: TrainConfig, ): + # Not routed through _setup_model_part: the UNet's checkpointing needs supports_offloading=False, which + # _setup_model_part's checkpointing_fn slot doesn't pass, so the parts are wired by hand here. if config.unet.checkpointing_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, supports_offloading=False) enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config, config.text_encoder) if config.force_circular_padding: @@ -60,8 +61,7 @@ def setup_optimizations( if model.unet_lora is not None: apply_circular_padding_to_conv2d(model.unet_lora) - model.autocast_context, model.train_dtype = create_autocast_context( - self.train_device, config.train_dtype, config.enable_autocast_cache) + super().setup_optimizations(model, config) quantize_layers(model.text_encoder, self.train_device, model.train_dtype, config) quantize_layers(model.vae, self.train_device, model.train_dtype, config) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 54a080579..dd7a4fe6f 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -17,7 +17,7 @@ ) from modules.util.config.TrainConfig import TrainConfig from modules.util.conv_util import apply_circular_padding_to_conv2d -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context +from modules.util.dtype_util import disable_fp16_autocast_context from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress @@ -46,9 +46,11 @@ def setup_optimizations( model: StableDiffusionXLModel, config: TrainConfig, ): + # Not routed through _setup_model_part: the UNet's checkpointing needs supports_offloading=False, which + # _setup_model_part's checkpointing_fn slot doesn't pass, so the parts are wired by hand here. if config.unet.checkpointing_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, supports_offloading=False) 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) @@ -58,8 +60,7 @@ def setup_optimizations( if model.unet_lora is not None: apply_circular_padding_to_conv2d(model.unet_lora) - model.autocast_context, model.train_dtype = create_autocast_context( - self.train_device, config.train_dtype, config.enable_autocast_cache) + super().setup_optimizations(model, config) model.vae_autocast_context, model.vae_train_dtype = disable_fp16_autocast_context( self.train_device, diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index d73052f83..d2bbb200a 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -15,7 +15,6 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.conv_util import apply_circular_padding_to_conv2d from modules.util.dtype_util import ( - create_autocast_context, disable_bf16_on_fp16_autocast_context, disable_fp16_autocast_context, ) @@ -52,6 +51,10 @@ def setup_optimizations( model: WuerstchenModel, config: TrainConfig, ): + # Not routed through _setup_model_part: Wuerstchen's parts (prior_prior, decoder_*, effnet_encoder, + # prior_text_encoder) don't match the transformer/text_encoder/vae shape _setup_model_part assumes and + # take bespoke per-part contexts (stable-cascade prior fp16-disable, effnet bf16-on-fp16), so this + # setup is fully hand-rolled. if config.prior.checkpointing_enabled(): model.prior_prior.enable_gradient_checkpointing() enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config, config.text_encoder) @@ -63,8 +66,7 @@ def setup_optimizations( if model.prior_prior_lora is not None: apply_circular_padding_to_conv2d(model.prior_prior_lora) - model.autocast_context, model.train_dtype = create_autocast_context( - self.train_device, config.train_dtype, config.enable_autocast_cache) + super().setup_optimizations(model, config) if model.model_type.is_stable_cascade(): model.prior_autocast_context, model.prior_train_dtype = disable_fp16_autocast_context( diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index b77f9bb72..df0f0912c 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -15,8 +15,6 @@ enable_checkpointing_for_z_image_transformer, ) from modules.util.config.TrainConfig import TrainConfig -from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context -from modules.util.quantization_util import quantize_layers from modules.util.TrainProgress import TrainProgress import torch @@ -45,25 +43,10 @@ def setup_optimizations( model: ZImageModel, config: TrainConfig, ): - model.transformer_offload_conductor = enable_checkpointing_for_z_image_transformer(model.transformer, config, config.transformer) - 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.enable_autocast_cache) - - #TODO necessary if we don't train it? - 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.enable_autocast_cache, - ) - - quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) - quantize_layers(model.vae, self.train_device, model.train_dtype, config) - quantize_layers(model.transformer, self.train_device, model.train_dtype, config) - self._set_attention_backend(model.transformer, config.attention_mechanism, mask=True) + super().setup_optimizations(model, config) + self._setup_model_part(model, config, "transformer", config.transformer, enable_checkpointing_for_z_image_transformer, attention_mask=True) + self._setup_model_part(model, config, "text_encoder", config.text_encoder, enable_checkpointing_for_qwen3_encoder_layers, disable_fp16_autocast=True) + self._setup_model_part(model, config, "vae", config.vae) def predict( self, diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index 1f669e9a5..2df40065e 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -243,13 +243,14 @@ def enable_checkpointing( 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, + supports_offloading: 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() + # a conductor exists iff this part actually offloads: the user enabled it (part.offloading_enabled()) and the + # architecture can be driven by the conductor (supports_offloading). + offload = supports_offloading and part.offloading_enabled() conductor = LayerOffloadConductor(model, config, part) if offload else None checkpointing = part.checkpointing_enabled() @@ -298,12 +299,12 @@ def enable_checkpointing_for_basic_transformer_blocks( model: nn.Module, config: TrainConfig, part: TrainModelPartConfig, - offload_enabled: bool, + supports_offloading: bool = True, ) -> LayerOffloadConductor | None: return enable_checkpointing(model, config, part, config.compile, [ (BasicTransformerBlock , []), ], - offload_enabled = offload_enabled, + supports_offloading = supports_offloading, ) def enable_checkpointing_for_clip_encoder_layers( @@ -313,7 +314,7 @@ def enable_checkpointing_for_clip_encoder_layers( ): 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 - ], offload_enabled=False) # CLIP is non-offloadable; keep it plain-checkpointed so a migrated offload_fraction can't build a self-activating conductor + ], supports_offloading=False) # CLIP is non-offloadable; keep it plain-checkpointed so a migrated offload_fraction can't build a self-activating conductor def enable_checkpointing_for_t5_encoder_layers( model: nn.Module, diff --git a/modules/util/dtype_util.py b/modules/util/dtype_util.py index b6d6b2643..7dde18d45 100644 --- a/modules/util/dtype_util.py +++ b/modules/util/dtype_util.py @@ -20,7 +20,7 @@ def create_autocast_context( device: torch.device, train_dtype: DataType | None, enable_autocast_cache: bool, -) -> tuple[torch.autocast | nullcontext, DataType]: +) -> torch.autocast | nullcontext: torch_train_dtype = train_dtype.torch_dtype() if torch_train_dtype in (torch.float16, torch.bfloat16): @@ -34,13 +34,13 @@ def create_autocast_context( # MPS additionally needs macOS >= 14. print(f"Warning: Mixed precision training is untested on device type '{device.type}'.") return torch.autocast(device_type=device.type, dtype=torch_train_dtype, - cache_enabled=enable_autocast_cache), train_dtype + cache_enabled=enable_autocast_cache) elif device.type == "cuda": # float32/tfloat32 on CUDA (and ROCm, which also reports device type "cuda"): # CUDA accepts float32 as an autocast dtype and upcasts lower-precision weights # on the fly (this is undocumented but works). return torch.autocast(device_type=device.type, dtype=torch_train_dtype, - cache_enabled=enable_autocast_cache), train_dtype + cache_enabled=enable_autocast_cache) else: # float32/tfloat32 on a non-CUDA backend (cpu, mps, xpu, ...): those backends # reject fp32 autocast, so disable autocast and let the model run at its weight @@ -49,7 +49,7 @@ def create_autocast_context( print("Warning: float32 training does not upcast lower-precision weights on this device " "(only CUDA can autocast to float32); the model runs at its weight dtype. " "Set the weight data types to float32 for full precision.") - return torch.autocast(device_type=device.type, enabled=False), train_dtype + return torch.autocast(device_type=device.type, enabled=False) def disable_fp16_autocast_context( From 44c80f7849e6b443c37e048e8e204890fb21a199 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:01:22 +0200 Subject: [PATCH 07/10] Add ModelType.trainable_parts() and TrainConfig.part_trained_in_place() A part's 'train' flag defaults True even for parts the architecture can't train (e.g. the frozen text encoder on the newer transformer models), so taking it at face value misclassifies those parts as trained. Record the architecture-trainable parts per model type and expose part_trained_in_place() -- train and architecture-trainable and FINE_TUNE -- for the memory-management modes that must not silently discard in-place weight updates. Co-Authored-By: Claude Opus 4.8 (1M context) --- modules/util/config/TrainConfig.py | 9 +++++++ modules/util/enum/ModelType.py | 38 ++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index deb632f25..9d350cc8f 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -889,6 +889,15 @@ def weight_dtypes(self) -> ModelWeightDtypes: self.embedding_weight_dtype, ) + def part_trained_in_place(self, part: TrainModelPartConfig) -> bool: + # True iff a FINE_TUNE run updates this part's base weights. 'train' defaults True even for parts the + # architecture can't train (e.g. a frozen text encoder), so also require the model type to list the part as + # trainable. Gates the offload/streaming modes that would silently discard in-place weight updates. + if self.training_method != TrainingMethod.FINE_TUNE or not part.train: + return False + name = next((p for p in self.model_type.model_parts() if getattr(self, p) is part), None) + return name in self.model_type.trainable_parts() + def model_names(self) -> ModelNames: return ModelNames( base_model=self.base_model_name, diff --git a/modules/util/enum/ModelType.py b/modules/util/enum/ModelType.py index 727d71cfd..a64f9cee7 100644 --- a/modules/util/enum/ModelType.py +++ b/modules/util/enum/ModelType.py @@ -216,6 +216,9 @@ def text_encoder_parts(self) -> tuple[str, ...]: # the text encoder components, named "text_encoder"/"text_encoder_2"/... by convention (see below). return tuple(part for part in _MODEL_PARTS[self] if part.startswith("text_encoder")) + def trainable_parts(self) -> tuple[str, ...]: + return _TRAINABLE_PARTS[self] + def supported_lora_formats(self) -> list[ModelFormat]: formats = [ ModelFormat.DIFFUSERS_LORA, @@ -313,6 +316,41 @@ def supported_output_formats(self, training_method: TrainingMethod) -> list[Mode ModelType.IDEOGRAM_4: ("transformer", "text_encoder", "unconditional_transformer", "vae"), } +# subset of _MODEL_PARTS the architecture allows a run to train, for both LoRA and fine-tuning -- the parts each setup +# routes through _setup_model_part_requires_grad. Parts omitted here (VAE everywhere; the text encoder on the newer +# transformer models; Ideogram's unconditional_transformer; Wuerstchen's decoder stack) are architecture-frozen. +_TRAINABLE_PARTS: dict[ModelType, tuple[str, ...]] = { + ModelType.STABLE_DIFFUSION_15: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_15_INPAINTING: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_20: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_20_BASE: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_20_INPAINTING: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_20_DEPTH: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_21: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_21_BASE: ("unet", "text_encoder"), + ModelType.STABLE_DIFFUSION_3: ("transformer", "text_encoder", "text_encoder_2", "text_encoder_3"), + ModelType.STABLE_DIFFUSION_35: ("transformer", "text_encoder", "text_encoder_2", "text_encoder_3"), + ModelType.STABLE_DIFFUSION_XL_10_BASE: ("unet", "text_encoder", "text_encoder_2"), + ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING: ("unet", "text_encoder", "text_encoder_2"), + ModelType.WUERSTCHEN_2: ("prior", "text_encoder"), + ModelType.STABLE_CASCADE_1: ("prior", "text_encoder"), + ModelType.PIXART_ALPHA: ("transformer", "text_encoder"), + ModelType.PIXART_SIGMA: ("transformer", "text_encoder"), + ModelType.FLUX_DEV_1: ("transformer", "text_encoder", "text_encoder_2"), + ModelType.FLUX_FILL_DEV_1: ("transformer", "text_encoder", "text_encoder_2"), + ModelType.FLUX_2: ("transformer",), + ModelType.ANIMA: ("transformer",), + ModelType.SANA: ("transformer", "text_encoder"), + ModelType.HUNYUAN_VIDEO: ("transformer", "text_encoder", "text_encoder_2"), + ModelType.HI_DREAM_FULL: ("transformer", "text_encoder", "text_encoder_2", "text_encoder_3", "text_encoder_4"), + ModelType.CHROMA_1: ("transformer", "text_encoder"), + ModelType.QWEN: ("transformer", "text_encoder"), + ModelType.KREA_2: ("transformer",), + ModelType.Z_IMAGE: ("transformer",), + ModelType.ERNIE: ("transformer",), + ModelType.IDEOGRAM_4: ("transformer",), +} + class PeftType(Enum): LORA = 'LORA' From 3d9dc782d9cbd41c96b129a1b8b7f5e316709dba Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:24:59 +0200 Subject: [PATCH 08/10] Stream layers through the offload conductor from checkpoint (disk-offload) Wires LayerOffloadConductor to materialize/evict layers directly from the checkpoint (set_disk_materialize, per-layer key_prefix) so disk-offload load-on-demand and the GPU/CPU layer split can be used together (issue #69). A streamed sub-module is loaded as a meta skeleton; its real weights are streamed straight from the checkpoint to the compute device and quantized there on first use, so the full unquantized module never lands in system RAM. With cache_in_ram off the weights are discarded back to meta after each use and re-streamed on the next; with it on they stay resident in pinned RAM. Includes the unpublished part of the materialize()/evict() base work this depends on (the conductor.to split and the move to evicting only at save/teardown boundaries). --- modules/model/AnimaModel.py | 5 - modules/model/BaseModel.py | 67 ++- modules/model/ChromaModel.py | 5 - modules/model/ErnieModel.py | 5 - modules/model/Flux2Model.py | 5 - modules/model/FluxModel.py | 5 - modules/model/HiDreamModel.py | 7 - modules/model/HunyuanVideoModel.py | 5 - modules/model/IdeogramModel.py | 7 - modules/model/Krea2Model.py | 5 - modules/model/PixArtAlphaModel.py | 5 - modules/model/QwenModel.py | 5 - modules/model/SanaModel.py | 5 - modules/model/StableDiffusion3Model.py | 5 - modules/model/ZImageModel.py | 5 - modules/modelLoader/AnimaModelLoader.py | 31 +- modules/modelLoader/BaseModelLoader.py | 2 + modules/modelLoader/ErnieModelLoader.py | 28 +- modules/modelLoader/Flux2ModelLoader.py | 58 ++- .../GenericEmbeddingModelLoader.py | 12 +- .../modelLoader/GenericFineTuneModelLoader.py | 12 +- modules/modelLoader/GenericLoRAModelLoader.py | 12 +- modules/modelLoader/IdeogramModelLoader.py | 58 ++- modules/modelLoader/ZImageModelLoader.py | 25 +- .../modelLoader/chroma/ChromaModelLoader.py | 27 +- modules/modelLoader/flux/FluxModelLoader.py | 47 +- .../modelLoader/hiDream/HiDreamModelLoader.py | 91 ++-- .../hunyuanVideo/HunyuanVideoModelLoader.py | 43 +- modules/modelLoader/krea2/Krea2ModelLoader.py | 31 +- .../modelLoader/mixin/HFModelLoaderMixin.py | 451 ++++++++++++++++-- .../pixartAlpha/PixArtAlphaModelLoader.py | 38 +- modules/modelLoader/qwen/QwenModelLoader.py | 31 +- modules/modelLoader/sana/SanaModelLoader.py | 38 +- .../StableDiffusionModelLoader.py | 33 +- .../StableDiffusion3ModelLoader.py | 69 +-- .../StableDiffusionXLModelLoader.py | 76 +-- .../wuerstchen/WuerstchenModelLoader.py | 54 +-- modules/modelSampler/AnimaSampler.py | 2 - modules/modelSampler/ChromaSampler.py | 2 - modules/modelSampler/ErnieSampler.py | 2 - modules/modelSampler/Flux2Sampler.py | 2 - modules/modelSampler/FluxSampler.py | 4 - modules/modelSampler/HiDreamSampler.py | 2 - modules/modelSampler/HunyuanVideoSampler.py | 2 - modules/modelSampler/IdeogramSampler.py | 2 - modules/modelSampler/Krea2Sampler.py | 2 - modules/modelSampler/PixArtAlphaSampler.py | 2 - modules/modelSampler/QwenSampler.py | 2 - modules/modelSampler/SanaSampler.py | 2 - .../modelSampler/StableDiffusion3Sampler.py | 2 - .../modelSampler/StableDiffusionSampler.py | 4 - .../modelSampler/StableDiffusionVaeSampler.py | 2 - .../modelSampler/StableDiffusionXLSampler.py | 4 - modules/modelSampler/WuerstchenSampler.py | 2 - modules/modelSampler/ZImageSampler.py | 2 - modules/modelSetup/BaseChromaSetup.py | 2 - modules/modelSetup/BaseKrea2Setup.py | 2 - modules/modelSetup/BaseModelSetup.py | 20 +- modules/modelSetup/BasePixArtAlphaSetup.py | 2 - modules/modelSetup/BaseQwenSetup.py | 2 - modules/modelSetup/BaseSanaSetup.py | 2 - .../modelSetup/BaseStableDiffusionSetup.py | 2 - modules/modelSetup/BaseWuerstchenSetup.py | 2 - modules/module/AdditionalEmbeddingWrapper.py | 7 +- modules/module/quantized/LinearFp8.py | 28 +- modules/module/quantized/LinearNf4.py | 16 +- modules/module/quantized/LinearSVD.py | 22 +- modules/module/quantized/LinearW8A8.py | 24 +- .../quantized/mixin/QuantizedLinearMixin.py | 12 + modules/trainer/GenericTrainer.py | 4 +- modules/ui/BaseModelTabView.py | 7 + modules/ui/BaseTrainingTabView.py | 5 + modules/ui/SampleWindowController.py | 7 +- modules/util/LayerOffloadConductor.py | 230 +++++++-- modules/util/checkpointing_util.py | 8 + modules/util/config/TrainConfig.py | 7 + modules/util/disk_stream.py | 109 +++++ modules/util/quantization_util.py | 20 +- modules/util/torch_util.py | 18 + 79 files changed, 1373 insertions(+), 635 deletions(-) create mode 100644 modules/util/disk_stream.py diff --git a/modules/model/AnimaModel.py b/modules/model/AnimaModel.py index 9c365273d..51fdd1a93 100644 --- a/modules/model/AnimaModel.py +++ b/modules/model/AnimaModel.py @@ -7,7 +7,6 @@ from modules.util.convert_util import add_prefix from modules.util.enum.DataType import DataType from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -39,8 +38,6 @@ class AnimaModel(BaseModel): text_encoder_train_dtype: DataType - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent lora training data transformer_lora: LoRAModuleWrapper | None @@ -66,8 +63,6 @@ def __init__( self.text_encoder_train_dtype = DataType.FLOAT_32 - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.transformer_lora = None self.lora_state_dict = None diff --git a/modules/model/BaseModel.py b/modules/model/BaseModel.py index 3b172461e..bb2614c46 100644 --- a/modules/model/BaseModel.py +++ b/modules/model/BaseModel.py @@ -1,4 +1,5 @@ from abc import ABCMeta +from collections.abc import Callable from contextlib import nullcontext from uuid import uuid4 @@ -6,12 +7,14 @@ from modules.module.LoRAModule import LoRAModuleWrapper from modules.util.config.TrainConfig import TrainConfig from modules.util.convert_util import qkv_fusion +from modules.util.disk_stream import stream_module_to from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType +from modules.util.LayerOffloadConductor import LayerOffloadConductor from modules.util.modelSpec.ModelSpec import ModelSpec from modules.util.NamedParameterGroup import NamedParameterGroupCollection -from modules.util.torch_util import create_mem_pool, mem_pool_context, supports_mem_pool, torch_gc +from modules.util.torch_util import create_mem_pool, device_equals, mem_pool_context, supports_mem_pool, torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -79,6 +82,9 @@ class BaseModel(metaclass=ABCMeta): embedding_state_dicts: dict[str, dict[str, Tensor]] | None autocast_context: torch.autocast | nullcontext train_dtype: DataType + cache_in_ram: dict[str, bool] + offload_conductor: dict[str, LayerOffloadConductor] + materialize_fn: dict[str, Callable] def __init__( self, @@ -86,6 +92,9 @@ def __init__( ): self.model_type = model_type self.parameters = None + self.cache_in_ram = {} + self.offload_conductor = {} + self.materialize_fn = {} self.optimizer = None self.optimizer_state_dict = None self.param_group_mapping = None @@ -134,39 +143,57 @@ def materialize_only_text_encoders(self): self.materialize_only(*self.model_type.text_encoder_parts()) def _move_part(self, part: str, device: torch.device): - # The generic per-component move: `part` (or `part_1` for the first of several split text encoders), - # its LoRA (`{part}_lora`), and its layer-offload conductor (`{part}_offload_conductor`), if present. + # Move a component (`part`, or `part_1` for the first of several split text encoders) and its LoRA. The + # dispatch below routes through an offload conductor and/or a disk-stream materialize closure if present. stem = f"{part}_1" if hasattr(self, f"{part}_1") else part - conductor = getattr(self, f"{stem}_offload_conductor", None) - lora = getattr(self, f"{stem}_lora", None) - component = None if conductor is not None else getattr(self, stem) + conductor = self.offload_conductor.get(stem) + materialize_fn = self.materialize_fn.get(stem) + cache_in_ram = self.cache_in_ram.get(stem, True) if conductor is not None: - conductor.to(device) - - if component is None and lora is None: + if device_equals(device, self.train_device): + train_dtype = getattr(self, f"{stem}_train_dtype", self.train_dtype) + conductor.materialize( + train_dtype, name=part, materialize_fn=materialize_fn, + cache_in_ram=cache_in_ram) + else: + to_meta = materialize_fn is not None and not cache_in_ram + conductor.evict(to_meta=to_meta) + elif materialize_fn is not None: + streamed_component = getattr(self, stem) + train_dtype = getattr(self, f"{stem}_train_dtype", self.train_dtype) + stream_module_to( + streamed_component, device, materialize_fn, train_dtype, + cache_in_ram=cache_in_ram, name=part, temp_device=self.temp_device) + + # move into the shared stem pool: the base component itself (unless a conductor or stream owns its move) plus + # the LoRA. getattr(self, stem) is None for a part in model_parts() that was never populated (e.g. an omitted + # text encoder), so it drops out below. + to_move = [] + if conductor is None and materialize_fn is None: + to_move.append(getattr(self, stem)) + lora = getattr(self, f"{stem}_lora", None) + to_move.append(lora) + to_move = [module for module in to_move if module is not None] + if not to_move: return if supports_mem_pool(device): - # The component (when not conductor-managed) and its LoRA share a per-stem MemPool so both land - # contiguously and release together on evict. A conductor keeps its own pool, so the stem pool then - # holds only the LoRA, keeping its small tensors out of the default pool across the part's evict/reload. + # The component (when not conductor/stream-managed) and its LoRA share a per-stem MemPool so both release + # together on evict, keeping the LoRA's small tensors from pinning freed default-pool segments across the + # part's evict/reload cycle. A conductor keeps its own pool, so the stem pool then holds only the LoRA. pool = self._mem_pools.get(stem) if pool is None: pool = self._mem_pools[stem] = create_mem_pool(device) with mem_pool_context(pool): - if component is not None: - component.to(device=device) - if lora is not None: - lora.to(device=device) + for module in to_move: + module.to(device=device) else: # the target has no MemPool (CPU): move normally and drop this stem's pool from the earlier GPU move, # so evict()'s torch_gc can release its segments - if component is not None: - component.to(device=device) - if lora is not None: - lora.to(device=device) + for module in to_move: + module.to(device=device) self._mem_pools.pop(stem, None) def eval(self): diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index 4fbcf430b..3f15219e2 100644 --- a/modules/model/ChromaModel.py +++ b/modules/model/ChromaModel.py @@ -8,7 +8,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -52,8 +51,6 @@ class ChromaModel(BaseModel): text_encoder_train_dtype: DataType - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent embedding training data embedding: ChromaModelEmbedding | None @@ -84,8 +81,6 @@ def __init__( self.text_encoder_train_dtype = DataType.FLOAT_32 - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.embedding = None self.additional_embeddings = [] diff --git a/modules/model/ErnieModel.py b/modules/model/ErnieModel.py index 42bb1be74..2e920a47a 100644 --- a/modules/model/ErnieModel.py +++ b/modules/model/ErnieModel.py @@ -8,7 +8,6 @@ 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 import torch from torch import Tensor @@ -34,8 +33,6 @@ class ErnieModel(BaseModel): # autocast context text_encoder_autocast_context: torch.autocast | nullcontext - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None transformer_lora: LoRAModuleWrapper | None lora_state_dict: dict | None @@ -56,8 +53,6 @@ def __init__( self.text_encoder_autocast_context = nullcontext() - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.transformer_lora = None self.lora_state_dict = None diff --git a/modules/model/Flux2Model.py b/modules/model/Flux2Model.py index 79eb02006..f7887e3f2 100644 --- a/modules/model/Flux2Model.py +++ b/modules/model/Flux2Model.py @@ -6,7 +6,6 @@ from modules.module.LoRAModule import LoRAModuleWrapper from modules.util.convert_util import chunk_swap from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -43,8 +42,6 @@ class Flux2Model(BaseModel): # autocast context text_encoder_autocast_context: torch.autocast | nullcontext - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None transformer_lora: LoRAModuleWrapper | None lora_state_dict: dict | None @@ -65,8 +62,6 @@ def __init__( self.text_encoder_autocast_context = nullcontext() - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.transformer_lora = None self.lora_state_dict = None diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index 5236b3498..e61d1a4f3 100644 --- a/modules/model/FluxModel.py +++ b/modules/model/FluxModel.py @@ -11,7 +11,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -66,8 +65,6 @@ class FluxModel(BaseModel): text_encoder_2_train_dtype: DataType - text_encoder_2_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent embedding training data embedding: FluxModelEmbedding | None @@ -103,8 +100,6 @@ def __init__( self.text_encoder_2_train_dtype = DataType.FLOAT_32 - self.text_encoder_2_offload_conductor = None - self.transformer_offload_conductor = None self.embedding = None self.additional_embeddings = [] diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index f99b5116b..8d39d36e8 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -9,7 +9,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -98,9 +97,6 @@ class HiDreamModel(BaseModel): text_encoder_3_train_dtype: DataType transformer_train_dtype: DataType - text_encoder_3_offload_conductor: LayerOffloadConductor | None - text_encoder_4_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent embedding training data embedding: HiDreamModelEmbedding | None @@ -149,9 +145,6 @@ def __init__( self.text_encoder_3_train_dtype = DataType.FLOAT_32 self.transformer_train_dtype = DataType.FLOAT_32 - self.text_encoder_3_offload_conductor = None - self.text_encoder_4_offload_conductor = None - self.transformer_offload_conductor = None self.embedding = None self.additional_embeddings = [] diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 9ec6b4516..15bf26b7b 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -10,7 +10,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -80,8 +79,6 @@ class HunyuanVideoModel(BaseModel): transformer_train_dtype: DataType - text_encoder_1_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent embedding training data embedding: HunyuanVideoModelEmbedding | None @@ -118,8 +115,6 @@ def __init__( self.transformer_train_dtype = DataType.FLOAT_32 - self.text_encoder_1_offload_conductor = None - self.transformer_offload_conductor = None self.embedding = None self.additional_embeddings = [] diff --git a/modules/model/IdeogramModel.py b/modules/model/IdeogramModel.py index 70fae8444..755b6cb3e 100644 --- a/modules/model/IdeogramModel.py +++ b/modules/model/IdeogramModel.py @@ -7,7 +7,6 @@ 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 import torch from torch import Tensor @@ -36,9 +35,6 @@ class IdeogramModel(BaseModel): # autocast context text_encoder_autocast_context: torch.autocast | nullcontext - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None - unconditional_transformer_offload_conductor: LayerOffloadConductor | None transformer_lora: LoRAModuleWrapper | None lora_state_dict: dict | None @@ -60,9 +56,6 @@ def __init__( self.text_encoder_autocast_context = nullcontext() - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None - self.unconditional_transformer_offload_conductor = None self.transformer_lora = None self.lora_state_dict = None diff --git a/modules/model/Krea2Model.py b/modules/model/Krea2Model.py index cf5ca0d69..c1c7a8976 100644 --- a/modules/model/Krea2Model.py +++ b/modules/model/Krea2Model.py @@ -5,7 +5,6 @@ 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 import torch import torch.nn.functional as F @@ -51,8 +50,6 @@ class Krea2Model(BaseModel): # autocast context text_encoder_autocast_context: torch.autocast | nullcontext - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent lora training data transformer_lora: LoRAModuleWrapper | None @@ -74,8 +71,6 @@ def __init__( self.text_encoder_autocast_context = nullcontext() - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.transformer_lora = None self.lora_state_dict = None diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index bf36d1a21..066413fac 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -8,7 +8,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -54,8 +53,6 @@ class PixArtAlphaModel(BaseModel): text_encoder_train_dtype: DataType - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent embedding training data embedding: PixArtAlphaModelEmbedding | None @@ -86,8 +83,6 @@ def __init__( self.text_encoder_train_dtype = DataType.FLOAT_32 - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.embedding = None self.additional_embeddings = [] diff --git a/modules/model/QwenModel.py b/modules/model/QwenModel.py index 3b69d8f8b..f1b5f06a0 100644 --- a/modules/model/QwenModel.py +++ b/modules/model/QwenModel.py @@ -7,7 +7,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -38,8 +37,6 @@ class QwenModel(BaseModel): text_encoder_train_dtype: DataType - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent lora training data text_encoder_lora: LoRAModuleWrapper | None @@ -64,8 +61,6 @@ def __init__( self.text_encoder_train_dtype = DataType.FLOAT_32 - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.text_encoder_lora = None self.transformer_lora = None diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index f43e98ed5..10b932f03 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -8,7 +8,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -55,8 +54,6 @@ class SanaModel(BaseModel): text_encoder_train_dtype: DataType vae_train_dtype: DataType - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent embedding training data embedding: SanaModelEmbedding | None @@ -88,8 +85,6 @@ def __init__( self.text_encoder_train_dtype = DataType.FLOAT_32 self.vae_train_dtype = DataType.FLOAT_32 - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.embedding = None self.additional_embeddings = [] diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index e9324bb1b..076e58aed 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -10,7 +10,6 @@ from modules.util.enum.DataType import DataType from modules.util.enum.ModelFormat import ModelFormat from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -77,8 +76,6 @@ class StableDiffusion3Model(BaseModel): text_encoder_3_train_dtype: DataType - text_encoder_3_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent embedding training data embedding: StableDiffusion3ModelEmbedding | None @@ -119,8 +116,6 @@ def __init__( self.text_encoder_3_train_dtype = DataType.FLOAT_32 - self.text_encoder_3_offload_conductor = None - self.transformer_offload_conductor = None self.embedding = None self.additional_embeddings = [] diff --git a/modules/model/ZImageModel.py b/modules/model/ZImageModel.py index 8b60f2c28..942a6ce08 100644 --- a/modules/model/ZImageModel.py +++ b/modules/model/ZImageModel.py @@ -7,7 +7,6 @@ from modules.util.convert_util import fuse from modules.util.enum.DataType import DataType from modules.util.enum.ModelType import ModelType -from modules.util.LayerOffloadConductor import LayerOffloadConductor import torch from torch import Tensor @@ -42,8 +41,6 @@ class ZImageModel(BaseModel): text_encoder_train_dtype: DataType - text_encoder_offload_conductor: LayerOffloadConductor | None - transformer_offload_conductor: LayerOffloadConductor | None # persistent lora training data text_encoder_lora: LoRAModuleWrapper | None @@ -68,8 +65,6 @@ def __init__( self.text_encoder_train_dtype = DataType.FLOAT_32 #TODO - self.text_encoder_offload_conductor = None - self.transformer_offload_conductor = None self.transformer_lora = None self.lora_state_dict = None diff --git a/modules/modelLoader/AnimaModelLoader.py b/modules/modelLoader/AnimaModelLoader.py index 77f963d3c..edc38eade 100644 --- a/modules/modelLoader/AnimaModelLoader.py +++ b/modules/modelLoader/AnimaModelLoader.py @@ -37,10 +37,12 @@ def __load_internal( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: 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, + stream_from_disk, ) else: raise Exception("not an internal model") @@ -54,38 +56,40 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer = Qwen2Tokenizer.from_pretrained( + model.tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - t5_tokenizer = T5TokenizerFast.from_pretrained( + model.t5_tokenizer = T5TokenizerFast.from_pretrained( base_model_name, subfolder="t5_tokenizer", ) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( Qwen3Model, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) # conditioner is always bfloat16 — small adapter, no user dtype control - text_conditioner = AnimaTextConditioner.from_pretrained( + model.text_conditioner = AnimaTextConditioner.from_pretrained( base_model_name, subfolder="text_conditioner", torch_dtype=torch.bfloat16, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKLQwenImage, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -93,24 +97,16 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( CosmosTransformer3DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, config=base_model_name, + stream_from_disk=stream_from_disk, ) - model.model_type = model_type - model.tokenizer = tokenizer - model.t5_tokenizer = t5_tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.text_conditioner = text_conditioner - model.vae = vae - model.transformer = transformer - def load( #TODO share code between models self, model: AnimaModel, @@ -118,12 +114,14 @@ def load( #TODO share code between models model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + stream_from_disk, ) return except Exception: @@ -132,6 +130,7 @@ def load( #TODO share code between models try: self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/BaseModelLoader.py b/modules/modelLoader/BaseModelLoader.py index 4a560c2f1..44d908e22 100644 --- a/modules/modelLoader/BaseModelLoader.py +++ b/modules/modelLoader/BaseModelLoader.py @@ -49,5 +49,7 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, + cache_in_ram: dict[str, bool] | None = None, ) -> BaseModel | None: pass diff --git a/modules/modelLoader/ErnieModelLoader.py b/modules/modelLoader/ErnieModelLoader.py index a9f0c9259..0c9e3b85b 100644 --- a/modules/modelLoader/ErnieModelLoader.py +++ b/modules/modelLoader/ErnieModelLoader.py @@ -34,11 +34,12 @@ def __load_internal( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: 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, + quantization, stream_from_disk, ) else: raise Exception("not an internal model") @@ -52,35 +53,38 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( ErnieImageTransformer2DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, config=base_model_name, + stream_from_disk=stream_from_disk, ) - tokenizer = AutoTokenizer.from_pretrained( + model.tokenizer = AutoTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( Mistral3Model, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKLFlux2, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -88,13 +92,6 @@ def __load_diffusers( vae_model_name, ) - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: ErnieModel, @@ -117,13 +114,14 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): 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.vae_model, quantization, stream_from_disk, ) return except Exception: @@ -132,7 +130,7 @@ def load( try: self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, - model_names.vae_model, quantization, + model_names.vae_model, quantization, stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/Flux2ModelLoader.py b/modules/modelLoader/Flux2ModelLoader.py index 3b502b8de..0f27c3cbc 100644 --- a/modules/modelLoader/Flux2ModelLoader.py +++ b/modules/modelLoader/Flux2ModelLoader.py @@ -39,10 +39,12 @@ def __load_internal( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: 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, + model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name, + quantization, stream_from_disk, ) else: raise Exception("not an internal model") @@ -56,48 +58,46 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( Flux2Transformer2DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, config=base_model_name, + stream_from_disk=stream_from_disk, ) - if transformer.config.num_attention_heads == 48: #Flux2.Dev - tokenizer = PixtralProcessor.from_pretrained( + if model.transformer.config.num_attention_heads == 48: #Flux2.Dev + model.tokenizer = PixtralProcessor.from_pretrained( base_model_name, subfolder="tokenizer", ).tokenizer - - text_encoder = self._load_text_encoder( - Mistral3ForConditionalGeneration, - weight_dtypes.text_encoder, - weight_dtypes.fallback_train_dtype, - base_model_name, - "text_encoder", - ) + text_encoder_class = Mistral3ForConditionalGeneration else: #Flux2.Klein - tokenizer = Qwen2Tokenizer.from_pretrained( + model.tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - text_encoder = self._load_text_encoder( - Qwen3ForCausalLM, - weight_dtypes.text_encoder, - weight_dtypes.fallback_train_dtype, - base_model_name, - "text_encoder", - ) + text_encoder_class = Qwen3ForCausalLM - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( + text_encoder_class, + weight_dtypes.text_encoder, + weight_dtypes.fallback_train_dtype, + base_model_name, + "text_encoder", + stream_from_disk=stream_from_disk, + ) + + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKLFlux2, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -105,13 +105,6 @@ def __load_diffusers( vae_model_name, ) - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: Flux2Model, @@ -132,12 +125,14 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( - model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, + model_names.vae_model, quantization, stream_from_disk, ) return except Exception: @@ -145,7 +140,8 @@ def load( try: self.__load_diffusers( - model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, + model_names.vae_model, quantization, stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/GenericEmbeddingModelLoader.py b/modules/modelLoader/GenericEmbeddingModelLoader.py index 019502fb4..7106cd1f0 100644 --- a/modules/modelLoader/GenericEmbeddingModelLoader.py +++ b/modules/modelLoader/GenericEmbeddingModelLoader.py @@ -36,16 +36,26 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, + cache_in_ram: dict[str, bool] | None = None, ) -> model_class | None: base_model_loader = model_loader_class() embedding_loader = embedding_loader_class() model = model_class(model_type=model_type) + cache_in_ram = cache_in_ram or {} + if not stream_from_disk and any(not cache_in_ram.get(part, False) for part in model_type.model_parts()): + print("Warning: 'stream from disk' is off, so every component stays fully in RAM; disabling " + "'cache in ram' cannot free it without streaming.") + model.cache_in_ram = { + part: not stream_from_disk or cache_in_ram.get(part, False) + for part in model_type.model_parts() + } self._load_internal_data(model, model_names.embedding.model_name) model.model_spec = self._load_default_model_spec(model_type) if model_names.base_model is not None: - base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization) + base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization, stream_from_disk) embedding_loader.load(model, model_names.embedding.model_name, model_names) return model diff --git a/modules/modelLoader/GenericFineTuneModelLoader.py b/modules/modelLoader/GenericFineTuneModelLoader.py index 09915388f..410c3efcb 100644 --- a/modules/modelLoader/GenericFineTuneModelLoader.py +++ b/modules/modelLoader/GenericFineTuneModelLoader.py @@ -40,17 +40,27 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, + cache_in_ram: dict[str, bool] | None = None, ) -> model_class | None: base_model_loader = model_loader_class() if embedding_loader_class is not None: embedding_loader = embedding_loader_class() model = model_class(model_type=model_type) + cache_in_ram = cache_in_ram or {} + if not stream_from_disk and any(not cache_in_ram.get(part, False) for part in model_type.model_parts()): + print("Warning: 'stream from disk' is off, so every component stays fully in RAM; disabling " + "'cache in ram' cannot free it without streaming.") + model.cache_in_ram = { + part: not stream_from_disk or cache_in_ram.get(part, False) + for part in model_type.model_parts() + } self._load_internal_data(model, model_names.base_model) model.model_spec = self._load_default_model_spec(model_type) - base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization) + base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization, stream_from_disk) if embedding_loader_class is not None: embedding_loader.load(model, model_names.base_model, model_names) diff --git a/modules/modelLoader/GenericLoRAModelLoader.py b/modules/modelLoader/GenericLoRAModelLoader.py index d120eb008..4f005123b 100644 --- a/modules/modelLoader/GenericLoRAModelLoader.py +++ b/modules/modelLoader/GenericLoRAModelLoader.py @@ -37,6 +37,8 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, + cache_in_ram: dict[str, bool] | None = None, ) -> model_class | None: base_model_loader = model_loader_class() lora_model_loader = lora_loader_class() @@ -44,11 +46,19 @@ def load( embedding_loader = embedding_loader_class() model = model_class(model_type=model_type) + cache_in_ram = cache_in_ram or {} + if not stream_from_disk and any(not cache_in_ram.get(part, False) for part in model_type.model_parts()): + print("Warning: 'stream from disk' is off, so every component stays fully in RAM; disabling " + "'cache in ram' cannot free it without streaming.") + model.cache_in_ram = { + part: not stream_from_disk or cache_in_ram.get(part, False) + for part in model_type.model_parts() + } self._load_internal_data(model, model_names.lora) model.model_spec = self._load_default_model_spec(model_type) if model_names.base_model is not None: - base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization) + base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization, stream_from_disk) lora_model_loader.load(model, model_names) if embedding_loader_class is not None: embedding_loader.load(model, model_names.lora, model_names) diff --git a/modules/modelLoader/IdeogramModelLoader.py b/modules/modelLoader/IdeogramModelLoader.py index 3858ce732..126e0c305 100644 --- a/modules/modelLoader/IdeogramModelLoader.py +++ b/modules/modelLoader/IdeogramModelLoader.py @@ -34,9 +34,12 @@ def __load_internal( vae_model_name: str, include_unconditional_transformer: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): if os.path.isfile(os.path.join(base_model_name, "meta.json")): - self.__load_diffusers(model, model_type, weight_dtypes, base_model_name, vae_model_name, include_unconditional_transformer, quantization) + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, vae_model_name, include_unconditional_transformer, + quantization, stream_from_disk) else: raise Exception("not an internal model") @@ -49,20 +52,33 @@ def __load_diffusers( vae_model_name: str, include_unconditional_transformer: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): - transformer = self._load_diffusers_sub_module( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( Ideogram4Transformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, + weight_dtypes, base_model_name, - "transformer", + "", quantization, + stream_from_disk=stream_from_disk, ) # the unconditional transformer is frozen and only used for the negative branch of the dual-network CFG at - # sampling, so it has its own weight dtype independent of the trainable transformer's. It is optional: if not - # loaded, only cfg_scale<=1 sampling is possible. - if include_unconditional_transformer: - unconditional_transformer = self._load_diffusers_sub_module( + # sampling, so it has its own weight dtype independent of the trainable transformer's. It is optional. It uses + # _load_diffusers_sub_module directly (not _load_transformer) because of its own subfolder and dtype; in + # streaming mode that returns a materialize closure, otherwise a plain module. + if include_unconditional_transformer and stream_from_disk: + model.unconditional_transformer, model.materialize_fn["unconditional_transformer"] = \ + self._load_diffusers_sub_module( + Ideogram4Transformer2DModel, + weight_dtypes.unconditional_transformer, + weight_dtypes.train_dtype, + base_model_name, + "unconditional_transformer", + quantization, + stream_from_disk=True, + ) + elif include_unconditional_transformer: + model.unconditional_transformer = self._load_diffusers_sub_module( Ideogram4Transformer2DModel, weight_dtypes.unconditional_transformer, weight_dtypes.train_dtype, @@ -71,27 +87,28 @@ def __load_diffusers( quantization, ) else: - unconditional_transformer = None + model.unconditional_transformer = None - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( Qwen3VLModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - tokenizer = AutoTokenizer.from_pretrained( + model.tokenizer = AutoTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKLFlux2, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -99,14 +116,6 @@ def __load_diffusers( vae_model_name, ) - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer - model.unconditional_transformer = unconditional_transformer - def __load_safetensors( self, model: IdeogramModel, @@ -126,13 +135,14 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, - model_names.include_unconditional_transformer, quantization, + model_names.include_unconditional_transformer, quantization, stream_from_disk, ) return except Exception: @@ -141,7 +151,7 @@ def load( try: self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, - model_names.include_unconditional_transformer, quantization, + model_names.include_unconditional_transformer, quantization, stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/ZImageModelLoader.py b/modules/modelLoader/ZImageModelLoader.py index f4b0b8db7..63df43787 100644 --- a/modules/modelLoader/ZImageModelLoader.py +++ b/modules/modelLoader/ZImageModelLoader.py @@ -37,10 +37,12 @@ def __load_internal( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: 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, + stream_from_disk, ) else: raise Exception("not an internal model") @@ -54,26 +56,28 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer = Qwen2Tokenizer.from_pretrained( + model.tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( Qwen3ForCausalLM, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -81,21 +85,15 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( ZImageTransformer2DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, + stream_from_disk=stream_from_disk, ) - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: ZImageModel, @@ -116,12 +114,14 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + stream_from_disk, ) return except Exception: @@ -130,6 +130,7 @@ def load( try: self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/chroma/ChromaModelLoader.py b/modules/modelLoader/chroma/ChromaModelLoader.py index 4a93cc674..59994818d 100644 --- a/modules/modelLoader/chroma/ChromaModelLoader.py +++ b/modules/modelLoader/chroma/ChromaModelLoader.py @@ -32,10 +32,12 @@ def __load_internal( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: 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, + stream_from_disk, ) else: raise Exception("not an internal model") @@ -49,26 +51,29 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer = T5Tokenizer.from_pretrained( + model.tokenizer = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) + model.orig_tokenizer = copy.deepcopy(model.tokenizer) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -76,22 +81,15 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( ChromaTransformer2DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, + stream_from_disk=stream_from_disk, ) - 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 - model.transformer = transformer - def __load_safetensors( self, model: ChromaModel, @@ -112,12 +110,14 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + stream_from_disk, ) return except Exception: @@ -126,6 +126,7 @@ def load( try: self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/flux/FluxModelLoader.py b/modules/modelLoader/flux/FluxModelLoader.py index 02547a950..0a2cf9546 100644 --- a/modules/modelLoader/flux/FluxModelLoader.py +++ b/modules/modelLoader/flux/FluxModelLoader.py @@ -37,11 +37,12 @@ def __load_internal( include_text_encoder_1: bool, include_text_encoder_2: bool, quantization: QuantizationConfig, + stream_from_disk: 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, - include_text_encoder_1, include_text_encoder_2, quantization, + include_text_encoder_1, include_text_encoder_2, quantization, stream_from_disk, ) else: raise Exception("not an internal model") @@ -57,51 +58,56 @@ def __load_diffusers( include_text_encoder_1: bool, include_text_encoder_2: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): if include_text_encoder_1: - tokenizer_1 = CLIPTokenizer.from_pretrained( + model.tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) else: - tokenizer_1 = None + model.tokenizer_1 = None + model.orig_tokenizer_1 = copy.deepcopy(model.tokenizer_1) if include_text_encoder_2: - tokenizer_2 = T5Tokenizer.from_pretrained( + model.tokenizer_2 = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer_2", ) else: - tokenizer_2 = None + model.tokenizer_2 = None + model.orig_tokenizer_2 = copy.deepcopy(model.tokenizer_2) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) if include_text_encoder_1: - text_encoder_1 = self._load_text_encoder( + model.text_encoder_1, model.materialize_fn["text_encoder_1"] = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) else: - text_encoder_1 = None + model.text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_text_encoder( + model.text_encoder_2, model.materialize_fn["text_encoder_2"] = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder_2, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder_2", + stream_from_disk=stream_from_disk, ) else: - text_encoder_2 = None + model.text_encoder_2 = None - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -109,25 +115,15 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( FluxTransformer2DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, + stream_from_disk=stream_from_disk, ) - 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 - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: FluxModel, @@ -212,13 +208,14 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, - model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, + model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, stream_from_disk, ) return except Exception: @@ -227,7 +224,7 @@ def load( try: self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, - model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, + model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/hiDream/HiDreamModelLoader.py b/modules/modelLoader/hiDream/HiDreamModelLoader.py index 4c0f002bb..7c542bd00 100644 --- a/modules/modelLoader/hiDream/HiDreamModelLoader.py +++ b/modules/modelLoader/hiDream/HiDreamModelLoader.py @@ -44,11 +44,13 @@ def __load_internal( include_text_encoder_3: bool, include_text_encoder_4: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): if os.path.isfile(os.path.join(base_model_name, "meta.json")): self.__load_diffusers( model, model_type, weight_dtypes, base_model_name, text_encoder_4_model_name, vae_model_name, - include_text_encoder_1, include_text_encoder_2, include_text_encoder_3, include_text_encoder_4, quantization, + include_text_encoder_1, include_text_encoder_2, include_text_encoder_3, include_text_encoder_4, + quantization, stream_from_disk, ) else: raise Exception("not an internal model") @@ -66,87 +68,103 @@ def __load_diffusers( include_text_encoder_3: bool, include_text_encoder_4: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer_1 = CLIPTokenizer.from_pretrained( + model.tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) if include_text_encoder_1 else None - tokenizer_2 = CLIPTokenizer.from_pretrained( + model.tokenizer_2 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer_2", ) if include_text_encoder_2 else None - tokenizer_3 = T5Tokenizer.from_pretrained( + model.tokenizer_3 = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer_3", ) if include_text_encoder_3 else None - tokenizer_4 = LlamaTokenizerFast.from_pretrained( + model.tokenizer_4 = LlamaTokenizerFast.from_pretrained( text_encoder_4_model_name, ) if include_text_encoder_4 else None - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) if include_text_encoder_1: - text_encoder_1 = self._load_text_encoder( + model.text_encoder_1, model.materialize_fn["text_encoder_1"] = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder, weight_dtypes.train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) else: - text_encoder_1 = None + model.text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_text_encoder( + model.text_encoder_2, model.materialize_fn["text_encoder_2"] = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder_2, weight_dtypes.train_dtype, base_model_name, "text_encoder_2", + stream_from_disk=stream_from_disk, ) else: - text_encoder_2 = None + model.text_encoder_2 = None if include_text_encoder_3: - text_encoder_3 = self._load_text_encoder( + model.text_encoder_3, model.materialize_fn["text_encoder_3"] = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder_3, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder_3", + stream_from_disk=stream_from_disk, ) else: - text_encoder_3 = None + model.text_encoder_3 = None if include_text_encoder_4: if text_encoder_4_model_name: # override repo holds text_encoder_4 at its root, not in a base-model subfolder, so it bypasses - # _load_text_encoder (which always loads from a base-repo subfolder) and loads directly - text_encoder_4 = self._load_transformers_sub_module( - LlamaForCausalLM, - weight_dtypes.text_encoder_4, - weight_dtypes.train_dtype, - text_encoder_4_model_name, - ) + # _load_text_encoder (which always loads from a base-repo subfolder) and loads directly. + # _load_transformers_sub_module returns a (module, materialize_fn) pair only when streaming; a bare + # module otherwise. + if stream_from_disk: + model.text_encoder_4, model.materialize_fn["text_encoder_4"] = self._load_transformers_sub_module( + LlamaForCausalLM, + weight_dtypes.text_encoder_4, + weight_dtypes.train_dtype, + text_encoder_4_model_name, + stream_from_disk=True, + ) + else: + model.text_encoder_4 = self._load_transformers_sub_module( + LlamaForCausalLM, + weight_dtypes.text_encoder_4, + weight_dtypes.train_dtype, + text_encoder_4_model_name, + ) else: - text_encoder_4 = self._load_text_encoder( + model.text_encoder_4, model.materialize_fn["text_encoder_4"] = self._load_text_encoder( LlamaForCausalLM, weight_dtypes.text_encoder_4, weight_dtypes.train_dtype, base_model_name, "text_encoder_4", + stream_from_disk=stream_from_disk, ) else: - text_encoder_4 = None + model.text_encoder_4 = None - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -154,28 +172,15 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_diffusers_sub_module( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( HiDreamImageTransformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, + weight_dtypes, base_model_name, - "transformer", + "", quantization, + stream_from_disk=stream_from_disk, ) - model.model_type = model_type - model.tokenizer_1 = tokenizer_1 - model.tokenizer_2 = tokenizer_2 - model.tokenizer_3 = tokenizer_3 - model.tokenizer_4 = tokenizer_4 - model.noise_scheduler = noise_scheduler - model.text_encoder_1 = text_encoder_1 - model.text_encoder_2 = text_encoder_2 - model.text_encoder_3 = text_encoder_3 - model.text_encoder_4 = text_encoder_4 - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: HiDreamModel, @@ -262,6 +267,7 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] @@ -271,6 +277,7 @@ def load( model_names.text_encoder_4, model_names.vae_model, model_names.include_text_encoder, model_names.include_text_encoder_2, model_names.include_text_encoder_3, model_names.include_text_encoder_4, quantization, + stream_from_disk, ) self.__after_load(model) return @@ -283,12 +290,18 @@ def load( model_names.text_encoder_4, model_names.vae_model, model_names.include_text_encoder, model_names.include_text_encoder_2, model_names.include_text_encoder_3, model_names.include_text_encoder_4, quantization, + stream_from_disk, ) self.__after_load(model) return except Exception: stacktraces.append(traceback.format_exc()) + if stream_from_disk: + # the single-file loader below builds a full pipeline via from_single_file, which can't stream; fall + # back to loading it fully into RAM. + print(f"Warning: 'stream from disk' is not supported for single-file {model_type}; loading fully into RAM.") + try: self.__load_safetensors( model, model_type, weight_dtypes, model_names.base_model, diff --git a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py index e1f306d58..f8a219d67 100644 --- a/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py +++ b/modules/modelLoader/hunyuanVideo/HunyuanVideoModelLoader.py @@ -38,11 +38,12 @@ def __load_internal( include_text_encoder_1: bool, include_text_encoder_2: bool, quantization: QuantizationConfig, + stream_from_disk: 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, - include_text_encoder_1, include_text_encoder_2, quantization, + include_text_encoder_1, include_text_encoder_2, quantization, stream_from_disk, ) else: raise Exception("not an internal model") @@ -58,51 +59,54 @@ def __load_diffusers( include_text_encoder_1: bool, include_text_encoder_2: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): if include_text_encoder_1: - tokenizer_1 = LlamaTokenizerFast.from_pretrained( + model.tokenizer_1 = LlamaTokenizerFast.from_pretrained( base_model_name, subfolder="tokenizer", ) else: - tokenizer_1 = None + model.tokenizer_1 = None if include_text_encoder_2: - tokenizer_2 = CLIPTokenizer.from_pretrained( + model.tokenizer_2 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer_2", ) else: - tokenizer_2 = None + model.tokenizer_2 = None - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) if include_text_encoder_1: - text_encoder_1 = self._load_text_encoder( + model.text_encoder_1, model.materialize_fn["text_encoder_1"] = self._load_text_encoder( LlamaModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) else: - text_encoder_1 = None + model.text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_text_encoder( + model.text_encoder_2, model.materialize_fn["text_encoder_2"] = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder_2, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder_2", + stream_from_disk=stream_from_disk, ) else: - text_encoder_2 = None + model.text_encoder_2 = None - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKLHunyuanVideo, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -110,24 +114,16 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( HunyuanVideoTransformer3DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, config=base_model_name, + stream_from_disk=stream_from_disk, ) - model.model_type = model_type - model.tokenizer_1 = tokenizer_1 - model.tokenizer_2 = tokenizer_2 - model.noise_scheduler = noise_scheduler - model.text_encoder_1 = text_encoder_1 - model.text_encoder_2 = text_encoder_2 - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: HunyuanVideoModel, @@ -212,13 +208,14 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, - model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, + model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, stream_from_disk, ) self.__after_load(model) return @@ -228,7 +225,7 @@ def load( try: self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, - model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, + model_names.include_text_encoder, model_names.include_text_encoder_2, quantization, stream_from_disk, ) self.__after_load(model) return diff --git a/modules/modelLoader/krea2/Krea2ModelLoader.py b/modules/modelLoader/krea2/Krea2ModelLoader.py index df51d2e5e..b4789e5d6 100644 --- a/modules/modelLoader/krea2/Krea2ModelLoader.py +++ b/modules/modelLoader/krea2/Krea2ModelLoader.py @@ -31,10 +31,12 @@ def __load_internal( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: 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, + model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name, + quantization, stream_from_disk, ) else: raise Exception("not an internal model") @@ -48,26 +50,28 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer = Qwen2Tokenizer.from_pretrained( + model.tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( Qwen3VLModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKLQwenImage, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -75,22 +79,16 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( Krea2Transformer2DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, config=base_model_name, + stream_from_disk=stream_from_disk, ) - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: Krea2Model, @@ -110,12 +108,14 @@ def load( #TODO share code between models model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( - model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, + model_names.vae_model, quantization, stream_from_disk, ) return except Exception: @@ -123,7 +123,8 @@ def load( #TODO share code between models try: self.__load_diffusers( - model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, + model_names.vae_model, quantization, stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/mixin/HFModelLoaderMixin.py b/modules/modelLoader/mixin/HFModelLoaderMixin.py index 4cdc457bd..64b15eecd 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -1,16 +1,21 @@ import json import logging import os +import queue +import threading from abc import ABCMeta from itertools import repeat +from modules.module.quantized.mixin.QuantizedModuleMixin import QuantizedModuleMixin from modules.util.config.TrainConfig import QuantizationConfig from modules.util.enum.DataType import DataType from modules.util.ModelWeightDtypes import ModelWeightDtypes from modules.util.quantization_util import ( + is_quantized_module, is_quantized_parameter, replace_linear_with_quantized_layers, ) +from modules.util.torch_util import mem_pool_context import torch from torch import nn @@ -21,18 +26,226 @@ import accelerate import huggingface_hub +from accelerate.utils import set_module_tensor_to_device from huggingface_hub.utils import EntryNotFoundError +from safetensors import safe_open from safetensors.torch import load_file +from tqdm import tqdm # huggingface_hub 1.16+ uses httpx, which logs every HTTP request/response at INFO level. logging.getLogger("httpx").setLevel(logging.WARNING) +# reader threads striping the checkpoint into host RAM while the main thread does H2D + inline quant +STREAM_READER_THREADS = 4 + + +def __stream_reader( + tid: int, + nthreads: int, + work: list[tuple], + key_to_file: dict[str, str], + source_key_map: dict[str, str] | None, + out_queue: queue.Queue, + done, +): + # prefetch reader thread: reads a stripe of the work list into host RAM and feeds the bounded queue. Each thread + # owns its safe_open handles (a handle is not safe for concurrent get_tensor). + thread_handles: dict[str, object] = {} + try: + for i in range(tid, len(work), nthreads): + item = work[i] + path = key_to_file[item[0]] + handle = thread_handles.get(path) + if handle is None: + handle = thread_handles[path] = safe_open(path, framework="pt", device="cpu") + # cache key is the renamed module-layout key; the file stores the original, so read by the original + # (identity when no rename map was built). + read_key = source_key_map.get(item[0], item[0]) if source_key_map else item[0] + # get_tensor returns a lazy mmap view; .clone() forces the read off disk into host RAM + out_queue.put((item, handle.get_tensor(read_key).clone())) + except Exception as e: + out_queue.put(e) + finally: + out_queue.put(done) + + +def _intended_float_dtype( + module: nn.Module, + module_name: str, + tensor_name: str, + dtype: DataType, + train_dtype: DataType, + keep_in_fp32_modules: list[str], +) -> torch.dtype | None: + # target dtype for a streamed float tensor, or None to leave it unchanged. A param the quantizer will pack keeps + # its dtype (the quantizer converts it); keep-in-fp32 modules and a quantized component's leftover params go to + # train_dtype; everything else to the weight dtype. + if is_quantized_parameter(module, tensor_name): + return None + if dtype.is_quantized() or module_name in keep_in_fp32_modules: + # a caller without a train_dtype yet (budget sizing) gets None -> the budget over-estimates these from the + # fp32 skeleton; the stream-time caller always passes a real train_dtype. + return train_dtype.torch_dtype() if train_dtype is not None else None + return dtype.torch_dtype() + + +def _stamp_skeleton_float_dtypes( + sub_module: nn.Module, + dtype: DataType, + train_dtype: DataType, + keep_in_fp32_modules: list[str], +): + # stamp each meta-skeleton float param with the dtype the stream will give it, so the offload VRAM budget (which + # sizes the still-meta skeleton) measures the real post-load footprint, not init_empty_weights' fp32 default. Free: + # a meta tensor holds no data, so .to() only rewrites its declared dtype. Uses the same _intended_float_dtype helper + # as the stream-time cast so the two agree; quantized weights are left alone (sized via predict_offload_bytes). + # Buffers are not stamped: they never enter the offload budget. + for name, module in sub_module.named_modules(): + module_name = name.split(".")[-1] + for tensor_name, param in module.named_parameters(recurse=False): + if not torch.is_floating_point(param): + continue + target = _intended_float_dtype(module, module_name, tensor_name, dtype, train_dtype, keep_in_fp32_modules) + if target is not None and param.dtype != target: + param.data = param.data.to(dtype=target) + + +def stream_module_from_checkpoint( + module: nn.Module, + device: torch.device, + key_to_file: dict[str, str], + dtype: DataType, + train_dtype: DataType, + keep_in_fp32_modules: list[str], + tied_weights_keys: dict[str, str] | None, + quantize: bool, + key_prefix: str = "", + source_key_map: dict[str, str] | None = None, + part_name: str | None = None, + dest_pool=None, +): + # Fill a meta skeleton by streaming its checkpoint weights one tensor at a time, so the full checkpoint never lands + # in RAM. key_prefix scopes the lookup to one sub-module; keys stay checkpoint-absolute. dest_pool routes + # non-quantized weights straight into a MemPool (quantized modules pack in the default pool). + def dest_pool_for(sub_module): + return dest_pool if (dest_pool is not None and not is_quantized_module(sub_module)) else None + + # flat work list of every checkpoint-backed skeleton tensor, so the reader threads below can drive the reads. + work = [] # (key, sub_module, tensor_name, is_buffer, module_name) + for name, sub_module in module.named_modules(): + module_name = name.split(".")[-1] + # gradient checkpointing in compile mode wraps each block in a CheckpointLayer, inserting a ".checkpoint." + # level into the live path; the checkpoint keys have none, so strip it before lookup (as LoRAModule does). + lookup_name = name.replace(".checkpoint.", ".") + for tensor_name, param in list(sub_module.named_parameters(recurse=False)): + key = ".".join(p for p in (key_prefix, lookup_name, tensor_name) if p) + if key in key_to_file and param.is_meta: + work.append((key, sub_module, tensor_name, False, module_name)) + for tensor_name, _buffer in list(sub_module.named_buffers(recurse=False)): + # non-persistent buffers (rotary inv_freq etc.) are config-derived, not stored in the checkpoint + if tensor_name in sub_module._non_persistent_buffers_set: + continue + # no is_meta guard (unlike params): init_empty_weights materializes persistent buffers as REAL init values, + # so is_meta can't mean "not yet filled" -- always stream, else the init value survives (mis-normalizing the VAE). + key = ".".join(p for p in (key_prefix, lookup_name, tensor_name) if p) + if key in key_to_file: + work.append((key, sub_module, tensor_name, True, module_name)) + + # place() lands one tensor: cast floats to their intended dtype (quantizer-packed params keep theirs), move to the + # compute device, quantize inline once a layer's weight arrives so VRAM never holds the whole unquantized module. + # bar: one tick per streamed tensor; only a whole-module stream (part_name set) shows it, per-layer conductor calls stay silent. + bar = tqdm(total=len(work), unit="tensor", desc=f"streaming {part_name}", leave=False, smoothing=0.05) \ + if part_name is not None else None + + def quantize_if_ready(sub_module): + # quantize a module whose weight has landed (no longer meta): quantize() self-guards against a second call, so + # firing it the moment the weight arrives (rather than in the batch pass quantize_layers() does) is always safe. + if isinstance(sub_module, QuantizedModuleMixin) and not sub_module.weight.is_meta: + sub_module.compute_dtype = train_dtype.torch_dtype() + sub_module.quantize(device=device) + + def place(item, value): + _key, sub_module, tensor_name, is_buffer, module_name = item + # tensors that will be quantized stay at their original dtype (the quantizer converts them); everything else is + # cast to its intended dtype here. + if torch.is_floating_point(value): + target = _intended_float_dtype(sub_module, module_name, tensor_name, dtype, train_dtype, keep_in_fp32_modules) + if target is not None: + value = value.to(dtype=target) + with mem_pool_context(dest_pool_for(sub_module)): + set_module_tensor_to_device(sub_module, tensor_name, device, value=value, dtype=value.dtype) + # quantize outside the pool context so a quantized module's dequant scratch stays in the default pool + if quantize: + quantize_if_ready(sub_module) + if bar is not None: + bar.update(1) + + # reader threads stripe the work list into host RAM and feed a bounded queue; the main thread drains it and does + # H2D + inline quantize on the default stream. Both the parallel reads and overlapping them with the GPU work are + # wins. Each reader clones the tensor off its mmap and drops its safetensors handles when it exits (right after its + # stripe), so the file mmaps are released early rather than pinned until first use -- keeps page-cache pressure + # down. Tensors may land out of order -- place() addresses each by name and inline quant is order-free. + nthreads = STREAM_READER_THREADS + out_queue: queue.Queue = queue.Queue(maxsize=2 * nthreads) + done = object() + + threads = [ + threading.Thread( + target=__stream_reader, + args=(tid, nthreads, work, key_to_file, source_key_map, out_queue, done), + name=f"stream-reader-{tid}", daemon=True, + ) + for tid in range(nthreads) + ] + for t in threads: + t.start() + finished = 0 + while finished < nthreads: + got = out_queue.get() + if got is done: + finished += 1 + elif isinstance(got, Exception): + raise got + else: + place(*got) + for t in threads: + t.join() + + # tied weights (e.g. Qwen3 lm_head <-> embed_tokens) are saved once, so the target stays meta; fill it with an + # independent clone of the source (not an alias -- in-place quantize would corrupt both), then quantize. Both keys + # are module-root-relative, so whole-module streams only (key_prefix == ""). + if not key_prefix: + for target_key, source_key in (tied_weights_keys or {}).items(): + parent_path, _, target_name = target_key.rpartition(".") + target_module = module.get_submodule(parent_path) + if target_module._parameters[target_name].is_meta: + source = module.get_parameter(source_key) + with mem_pool_context(dest_pool_for(target_module)): + set_module_tensor_to_device( + target_module, target_name, device, value=source.detach().clone(), dtype=source.dtype) + if quantize: + quantize_if_ready(target_module) + + # non-persistent buffers (rotary inv_freq etc.) are skipped above but materialized REAL on cpu by init_empty_weights; + # move them to the device so the forward doesn't see cpu buffers vs device activations. Whole-module streams only. + if not key_prefix and device.type != "meta": + for sub_module in module.modules(): + for buffer_name in sub_module._non_persistent_buffers_set: + buffer = sub_module._buffers.get(buffer_name) + if buffer is not None and not buffer.is_meta: + with mem_pool_context(dest_pool_for(sub_module)): + sub_module._buffers[buffer_name] = buffer.to(device) + + if bar is not None: + bar.close() + class HFModelLoaderMixin(metaclass=ABCMeta): def __init__(self): super().__init__() - def __load_sub_module( + # ===== LEGACY (non-streaming) load path -- used only when Stream From Disk is off ===== + def __load_sub_module_legacy( self, sub_module: nn.Module, dtype: DataType, @@ -191,6 +404,7 @@ def __load_sub_module( module._parameters[tensor_name] = type(module._parameters[tensor_name])(source) return sub_module + # ===== end LEGACY load path ===== def _load_transformers_sub_module( self, @@ -199,6 +413,7 @@ def _load_transformers_sub_module( train_dtype: DataType, pretrained_model_name_or_path: str, subfolder: str = "", + stream_from_disk: bool = False, ): user_agent = { "file_type": "model", @@ -215,19 +430,110 @@ def _load_transformers_sub_module( with accelerate.init_empty_weights(): sub_module = module_type(config) - return self.__load_sub_module( - sub_module=sub_module, - dtype=dtype, - train_dtype=train_dtype, - keep_in_fp32_modules=module_type._keep_in_fp32_modules, - quantization=None, - pretrained_model_name_or_path=pretrained_model_name_or_path, - subfolder=subfolder, + if not stream_from_disk: + # LEGACY fallback: whole-checkpoint-into-RAM load + return self.__load_sub_module_legacy( + sub_module=sub_module, + dtype=dtype, + train_dtype=train_dtype, + keep_in_fp32_modules=module_type._keep_in_fp32_modules, + quantization=None, + pretrained_model_name_or_path=pretrained_model_name_or_path, + subfolder=subfolder, + model_filename="model.safetensors", + pytorch_model_filename="pytorch_model.bin", + shard_index_filename="model.safetensors.index.json", + ) + + keep_in_fp32_modules = module_type._keep_in_fp32_modules or [] + replace_linear_with_quantized_layers(sub_module, dtype, keep_in_fp32_modules, None, copy_parameters=False) + # stamp with train_dtype=None: the per-part train_dtype is only known at materialize, so keep-in-fp32 and + # quantized-leftover params stay fp32 in the skeleton (a safe budget over-estimate); see _intended_float_dtype + _stamp_skeleton_float_dtypes(sub_module, dtype, None, keep_in_fp32_modules) + + key_to_file = self.__resolve_shard_key_to_file( + pretrained_model_name_or_path, subfolder, model_filename="model.safetensors", - pytorch_model_filename="pytorch_model.bin", shard_index_filename="model.safetensors.index.json", ) + # some checkpoints (e.g. Ernie's Mistral3, Qwen's Qwen2_5_VL text encoders) were saved with an older module + # layout than transformers builds from the config now. Reuse transformers' own checkpoint conversion registry + # to rename the checkpoint keys to the module's layout so the streamed lookup finds them. diffusers sub-modules + # have no such registry (plain FrozenDict config, no model_type) and never need this. + weight_renamings = get_checkpoint_conversion_mapping(sub_module.config.model_type) \ + if hasattr(sub_module.config, 'model_type') else None + source_key_map = None + if weight_renamings: + meta_state_dict = sub_module.state_dict() + renamed_key_to_file = {} + # the rename maps each checkpoint key to the module's layout so the streamed lookup and the offload cache + # find it; the file itself still stores the original key, so keep renamed->original to read the tensor. + source_key_map = {} + for key, file in key_to_file.items(): + renamed = rename_source_key( + key, weight_renamings, [], prefix=sub_module.base_model_prefix, meta_state_dict=meta_state_dict, + )[0] + renamed_key_to_file[renamed] = file + source_key_map[renamed] = key + key_to_file = renamed_key_to_file + + return self.__finish_sub_module_load( + sub_module, dtype, train_dtype, keep_in_fp32_modules, key_to_file, source_key_map=source_key_map) + + def __resolve_shard_key_to_file( + self, + pretrained_model_name_or_path: str, + subfolder: str, + model_filename: str, + shard_index_filename: str, + ) -> dict[str, str]: + # map every checkpoint tensor key to the local safetensors file that holds it (downloading shards from the + # hub if the source is a repo id), so the streaming fill can read each tensor on demand. + is_local = os.path.isdir(pretrained_model_name_or_path) + + def resolve(filename: str) -> str | None: + # return a local path to `filename` (downloading it from the hub if needed), or None if it is absent + if is_local: + if subfolder: + path = os.path.join(pretrained_model_name_or_path, subfolder, filename) + else: + path = os.path.join(pretrained_model_name_or_path, filename) + return path if os.path.isfile(path) else None + try: + return huggingface_hub.hf_hub_download( + repo_id=pretrained_model_name_or_path, subfolder=subfolder, filename=filename) + except EntryNotFoundError: + return None + + key_to_file = {} + + index_path = resolve(shard_index_filename) + if index_path is not None: + with open(index_path, "r") as f: + weight_map = json.loads(f.read())["weight_map"] + shard_paths = {shard: resolve(shard) for shard in set(weight_map.values())} + for key, shard in weight_map.items(): + key_to_file[key] = shard_paths[shard] + return key_to_file + + # non-sharded: prefer the full-precision safetensors, fall back to the fp16 variant (some older repos, e.g. + # stable-diffusion-inpainting, ship only *.fp16.safetensors next to legacy pickle .bin files). Pickle .bin + # weights are not supported -- safe_open needs safetensors for random per-tensor reads. + fp16_filename = model_filename.replace(".safetensors", ".fp16.safetensors") + full_filename = resolve(model_filename) or resolve(fp16_filename) + if full_filename is None: + location = f"{pretrained_model_name_or_path}/{subfolder}" if subfolder else pretrained_model_name_or_path + raise FileNotFoundError( + f"No safetensors weights found for '{location}' (looked for {model_filename} and {fp16_filename}). " + f"Only pickle .bin checkpoints are present, which are not supported; convert the model to " + f"safetensors.") + with safe_open(full_filename, framework="pt") as f: + for key in f.keys(): # noqa: SIM118 -- safe_open handle, not a dict + key_to_file[key] = full_filename + + return key_to_file + def _load_diffusers_sub_module( self, module_type, @@ -236,6 +542,7 @@ def _load_diffusers_sub_module( pretrained_model_name_or_path: str, subfolder: str | None = None, quantization: QuantizationConfig | None = None, + stream_from_disk: bool = False, ): user_agent = { "file_type": "model", @@ -252,19 +559,67 @@ def _load_diffusers_sub_module( with accelerate.init_empty_weights(): sub_module = module_type.from_config(config) - return self.__load_sub_module( - sub_module=sub_module, - dtype=dtype, - train_dtype=train_dtype, - keep_in_fp32_modules=module_type._keep_in_fp32_modules, - quantization=quantization, - pretrained_model_name_or_path=pretrained_model_name_or_path, - subfolder=subfolder, + if not stream_from_disk: + # LEGACY fallback: whole-checkpoint-into-RAM load + return self.__load_sub_module_legacy( + sub_module=sub_module, + dtype=dtype, + train_dtype=train_dtype, + keep_in_fp32_modules=module_type._keep_in_fp32_modules, + quantization=quantization, + pretrained_model_name_or_path=pretrained_model_name_or_path, + subfolder=subfolder, + model_filename="diffusion_pytorch_model.safetensors", + pytorch_model_filename="diffusion_pytorch_model.bin", + shard_index_filename="diffusion_pytorch_model.safetensors.index.json", + ) + + keep_in_fp32_modules = module_type._keep_in_fp32_modules or [] + replace_linear_with_quantized_layers(sub_module, dtype, keep_in_fp32_modules, quantization, copy_parameters=False) + # stamp with train_dtype=None: the per-part train_dtype is only known at materialize, so keep-in-fp32 and + # quantized-leftover params stay fp32 in the skeleton (a safe budget over-estimate); see _intended_float_dtype + _stamp_skeleton_float_dtypes(sub_module, dtype, None, keep_in_fp32_modules) + + key_to_file = self.__resolve_shard_key_to_file( + pretrained_model_name_or_path, subfolder, model_filename="diffusion_pytorch_model.safetensors", - pytorch_model_filename="diffusion_pytorch_model.bin", shard_index_filename="diffusion_pytorch_model.safetensors.index.json", ) + # diffusers renamed deprecated attention-block weights (query->to_q etc.); older single-file checkpoints still + # use the old names. _fix_state_dict_keys_on_load rewrites them to the current layout, and since it only + # renames dict keys, applying it to the key->file map matches applying it to a state_dict. No-op for modern + # architectures. + if hasattr(sub_module, '_fix_state_dict_keys_on_load'): + sub_module._fix_state_dict_keys_on_load(key_to_file) + + return self.__finish_sub_module_load( + sub_module, dtype, train_dtype, keep_in_fp32_modules, key_to_file) + + def __finish_sub_module_load( + self, + sub_module: nn.Module, + dtype: DataType, + train_dtype: DataType, + keep_in_fp32_modules: list[str], + key_to_file: dict[str, str], + source_key_map: dict[str, str] | None = None, + ): + tied_weights_keys = getattr(sub_module, "_tied_weights_keys", None) + + # module/key_prefix let the layer-offload conductor reuse this same closure to stream one layer at a time + # (module=that layer, key_prefix=its path in the checkpoint) as well as the non-layer remainder + # (module=the whole sub-module, key_prefix=""). Whole-module callers pass neither and stream everything. + def materialize_fn( + module: nn.Module, device: torch.device, train_dtype: DataType, key_prefix: str = "", + part_name: str | None = None, dest_pool=None): + stream_module_from_checkpoint( + module, device, key_to_file, dtype, train_dtype, + keep_in_fp32_modules, tied_weights_keys, quantize=True, key_prefix=key_prefix, + source_key_map=source_key_map, part_name=part_name, dest_pool=dest_pool) + + return sub_module, materialize_fn + def __convert_sub_module_to_dtype( self, sub_module: nn.Module, @@ -339,10 +694,13 @@ def _load_transformer( transformer_model_name: str, quantization: QuantizationConfig, config: str | None = None, + stream_from_disk: bool = False, ): # a single-file (optionally GGUF-quantized) checkpoint is loaded directly, using # a separate repo to source the model config if the checkpoint doesn't carry one; - # otherwise the transformer is loaded from its subfolder in the base model repo + # otherwise the transformer is loaded from its subfolder in the base model repo. + # Always returns a (transformer, materialize_fn) pair -- materialize_fn None when not streamed -- so callers + # pass stream_from_disk through. if transformer_model_name: single_file_kwargs = {} if config is not None: @@ -356,10 +714,14 @@ def _load_transformer( 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, ) - return self._convert_diffusers_sub_module_to_dtype( + transformer = self._convert_diffusers_sub_module_to_dtype( transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization, ) - else: + return transformer, None + elif stream_from_disk: + # stream from disk: meta skeleton + materialize closure; weights are streamed and quantized to the compute + # device on use and evicted back to meta afterwards, so the full unquantized module never lands in RAM. + # train_dtype is applied per-materialize, not here. return self._load_diffusers_sub_module( module_type, weight_dtypes.transformer, @@ -367,7 +729,18 @@ def _load_transformer( base_model_name, "transformer", quantization, + stream_from_disk=True, + ) + else: + transformer = self._load_diffusers_sub_module( + module_type, + weight_dtypes.transformer, + weight_dtypes.train_dtype, + base_model_name, + "transformer", + quantization, ) + return transformer, None def _load_text_encoder( self, @@ -376,18 +749,30 @@ def _load_text_encoder( train_dtype: DataType, base_model_name: str, subfolder: str, + stream_from_disk: bool = False, ): - # text encoders have no single-file override and always load from their subfolder in the base model - # repo; kept as a per-model entry point alongside _load_transformer / _load_vae. dtype/train_dtype are - # explicit rather than a weight_dtypes bundle since a model can hold several encoders (text_encoder, - # text_encoder_2, ...) with differing dtypes - return self._load_transformers_sub_module( - module_type, - dtype, - train_dtype, - base_model_name, - subfolder, - ) + # text encoders have no single-file override and always load from their subfolder. Always returns a + # (text_encoder, materialize_fn) pair -- materialize_fn None when not streamed -- mirroring _load_transformer. + # dtype/train_dtype are explicit rather than a weight_dtypes bundle since a model can hold several encoders + # (text_encoder, text_encoder_2, ...) with differing dtypes. + if stream_from_disk: + return self._load_transformers_sub_module( + module_type, + dtype, + train_dtype, + base_model_name, + subfolder, + stream_from_disk=True, + ) + else: + text_encoder = self._load_transformers_sub_module( + module_type, + dtype, + train_dtype, + base_model_name, + subfolder, + ) + return text_encoder, None def _load_vae( self, diff --git a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py index 69e280c9d..9c6d647e6 100644 --- a/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py +++ b/modules/modelLoader/pixartAlpha/PixArtAlphaModelLoader.py @@ -27,9 +27,11 @@ def __load_internal( base_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): if os.path.isfile(os.path.join(base_model_name, "meta.json")): - self.__load_diffusers(model, model_type, weight_dtypes, base_model_name, vae_model_name, quantization) + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, vae_model_name, quantization, stream_from_disk) else: raise Exception("not an internal model") @@ -41,26 +43,29 @@ def __load_diffusers( base_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer = T5Tokenizer.from_pretrained( + model.tokenizer = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) + model.orig_tokenizer = copy.deepcopy(model.tokenizer) - noise_scheduler = DDIMScheduler.from_pretrained( + model.noise_scheduler = DDIMScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -68,23 +73,15 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_diffusers_sub_module( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( Transformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, + weight_dtypes, base_model_name, - "transformer", + "", quantization, + stream_from_disk=stream_from_disk, ) - 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 - model.transformer = transformer - def load( self, model: PixArtAlphaModel, @@ -92,19 +89,22 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ) -> PixArtAlphaModel | None: stacktraces = [] base_model_name = model_names.base_model try: - self.__load_internal(model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization) + self.__load_internal( + model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization, stream_from_disk) return except Exception: stacktraces.append(traceback.format_exc()) try: - self.__load_diffusers(model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization) + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization, stream_from_disk) return except Exception: stacktraces.append(traceback.format_exc()) diff --git a/modules/modelLoader/qwen/QwenModelLoader.py b/modules/modelLoader/qwen/QwenModelLoader.py index 77a12bed7..21a4e76f5 100644 --- a/modules/modelLoader/qwen/QwenModelLoader.py +++ b/modules/modelLoader/qwen/QwenModelLoader.py @@ -31,10 +31,12 @@ def __load_internal( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: 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, + model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name, + quantization, stream_from_disk, ) else: raise Exception("not an internal model") @@ -48,26 +50,28 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer = Qwen2Tokenizer.from_pretrained( + model.tokenizer = Qwen2Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( Qwen2_5_VLForConditionalGeneration, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKLQwenImage, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -75,22 +79,16 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_transformer( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( QwenImageTransformer2DModel, weight_dtypes, base_model_name, transformer_model_name, quantization, config=base_model_name, + stream_from_disk=stream_from_disk, ) - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: QwenModel, @@ -111,12 +109,14 @@ def load( #TODO share code between models model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] try: self.__load_internal( - model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, + model_names.vae_model, quantization, stream_from_disk, ) return except Exception: @@ -124,7 +124,8 @@ def load( #TODO share code between models try: self.__load_diffusers( - model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, + model_names.vae_model, quantization, stream_from_disk, ) return except Exception: diff --git a/modules/modelLoader/sana/SanaModelLoader.py b/modules/modelLoader/sana/SanaModelLoader.py index ec4d31239..74700bf54 100644 --- a/modules/modelLoader/sana/SanaModelLoader.py +++ b/modules/modelLoader/sana/SanaModelLoader.py @@ -27,9 +27,11 @@ def __load_internal( base_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): if os.path.isfile(os.path.join(base_model_name, "meta.json")): - self.__load_diffusers(model, model_type, weight_dtypes, base_model_name, vae_model_name, quantization) + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, vae_model_name, quantization, stream_from_disk) else: raise Exception("not an internal model") @@ -41,26 +43,29 @@ def __load_diffusers( base_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer = GemmaTokenizer.from_pretrained( + model.tokenizer = GemmaTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) + model.orig_tokenizer = copy.deepcopy(model.tokenizer) - noise_scheduler = DPMSolverMultistepScheduler.from_pretrained( + model.noise_scheduler = DPMSolverMultistepScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - text_encoder = self._load_text_encoder( + model.text_encoder, model.materialize_fn["text_encoder"] = self._load_text_encoder( Gemma2Model, weight_dtypes.text_encoder, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderDC, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -68,23 +73,15 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_diffusers_sub_module( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( SanaTransformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, + weight_dtypes, base_model_name, - "transformer", + "", quantization, + stream_from_disk=stream_from_disk, ) - 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 - model.transformer = transformer - def load( self, model: SanaModel, @@ -92,19 +89,22 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ) -> SanaModel | None: stacktraces = [] base_model_name = model_names.base_model try: - self.__load_internal(model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization) + self.__load_internal( + model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization, stream_from_disk) return except Exception: stacktraces.append(traceback.format_exc()) try: - self.__load_diffusers(model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization) + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, model_names.vae_model, quantization, stream_from_disk) return except Exception: stacktraces.append(traceback.format_exc()) diff --git a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py index 223d1da61..50dd4d12b 100644 --- a/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py +++ b/modules/modelLoader/stableDiffusion/StableDiffusionModelLoader.py @@ -73,21 +73,22 @@ def __load_diffusers( vae_model_name: str, quantization: QuantizationConfig, ): - tokenizer = CLIPTokenizer.from_pretrained( + model.tokenizer = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) + model.orig_tokenizer = copy.deepcopy(model.tokenizer) noise_scheduler = DDIMScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - noise_scheduler = create.create_noise_scheduler( + model.noise_scheduler = create.create_noise_scheduler( noise_scheduler=NoiseScheduler.DDIM, original_noise_scheduler=noise_scheduler, ) - text_encoder = self._load_text_encoder( + model.text_encoder, _ = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -95,7 +96,7 @@ def __load_diffusers( "text_encoder", ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -103,7 +104,7 @@ def __load_diffusers( vae_model_name, ) - unet = self._load_diffusers_sub_module( + model.unet = self._load_diffusers_sub_module( UNet2DConditionModel, weight_dtypes.unet, weight_dtypes.train_dtype, @@ -112,27 +113,17 @@ def __load_diffusers( quantization, ) - image_depth_processor = DPTImageProcessor.from_pretrained( + model.image_depth_processor = DPTImageProcessor.from_pretrained( base_model_name, subfolder="feature_extractor", ) if model_type.has_depth_input() else None - depth_estimator = DPTForDepthEstimation.from_pretrained( + model.depth_estimator = DPTForDepthEstimation.from_pretrained( base_model_name, subfolder="depth_estimator", torch_dtype=weight_dtypes.unet.torch_dtype(), # TODO: use depth estimator dtype ) if model_type.has_depth_input() else None - 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 - model.unet = unet - model.image_depth_processor = image_depth_processor - model.depth_estimator = depth_estimator - def __fix_nai_model(self, state_dict: dict) -> dict: # fix for loading models with an empty state_dict key while 'state_dict' in state_dict and len(state_dict['state_dict']) > 0: @@ -272,9 +263,17 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] + if stream_from_disk: + # SD 1.5 / 2.x checkpoints are almost always single-file (.ckpt/.safetensors loaded via + # download_from_original_stable_diffusion_ckpt), which builds a full pipeline and can't stream from a meta + # skeleton. The diffusers-subfolder path could stream its unet/text encoder like SDXL does, but wasn't + # wired up, as this is legacy. So the toggle is ignored here. + print(f"Warning: 'stream from disk' is not supported for {model_type}; loading the model fully into RAM.") + model.sd_config = self._load_sd_config(model_type, model_names.base_model) model.sd_config_filename = self._get_sd_config_name(model_type, model_names.base_model) diff --git a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py index 0c4f4348a..633e0a72d 100644 --- a/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py +++ b/modules/modelLoader/stableDiffusion3/StableDiffusion3ModelLoader.py @@ -30,11 +30,13 @@ def __load_internal( include_text_encoder_2: bool, include_text_encoder_3: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): if os.path.isfile(os.path.join(base_model_name, "meta.json")): self.__load_diffusers( model, model_type, weight_dtypes, base_model_name, vae_model_name, include_text_encoder_1, include_text_encoder_2, include_text_encoder_3, quantization, + stream_from_disk, ) else: raise Exception("not an internal model") @@ -50,70 +52,77 @@ def __load_diffusers( include_text_encoder_2: bool, include_text_encoder_3: bool, quantization: QuantizationConfig, + stream_from_disk: bool, ): if include_text_encoder_1: - tokenizer_1 = CLIPTokenizer.from_pretrained( + model.tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) else: - tokenizer_1 = None + model.tokenizer_1 = None + model.orig_tokenizer_1 = copy.deepcopy(model.tokenizer_1) if include_text_encoder_2: - tokenizer_2 = CLIPTokenizer.from_pretrained( + model.tokenizer_2 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer_2", ) else: - tokenizer_2 = None + model.tokenizer_2 = None + model.orig_tokenizer_2 = copy.deepcopy(model.tokenizer_2) if include_text_encoder_3: - tokenizer_3 = T5Tokenizer.from_pretrained( + model.tokenizer_3 = T5Tokenizer.from_pretrained( base_model_name, subfolder="tokenizer_3", ) else: - tokenizer_3 = None + model.tokenizer_3 = None + model.orig_tokenizer_3 = copy.deepcopy(model.tokenizer_3) - noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) if include_text_encoder_1: - text_encoder_1 = self._load_text_encoder( + model.text_encoder_1, model.materialize_fn["text_encoder_1"] = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder, weight_dtypes.train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) else: - text_encoder_1 = None + model.text_encoder_1 = None if include_text_encoder_2: - text_encoder_2 = self._load_text_encoder( + model.text_encoder_2, model.materialize_fn["text_encoder_2"] = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder_2, weight_dtypes.train_dtype, base_model_name, "text_encoder_2", + stream_from_disk=stream_from_disk, ) else: - text_encoder_2 = None + model.text_encoder_2 = None if include_text_encoder_3: - text_encoder_3 = self._load_text_encoder( + model.text_encoder_3, model.materialize_fn["text_encoder_3"] = self._load_text_encoder( T5EncoderModel, weight_dtypes.text_encoder_3, weight_dtypes.fallback_train_dtype, base_model_name, "text_encoder_3", + stream_from_disk=stream_from_disk, ) else: - text_encoder_3 = None + model.text_encoder_3 = None - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.train_dtype, @@ -121,29 +130,15 @@ def __load_diffusers( vae_model_name, ) - transformer = self._load_diffusers_sub_module( + model.transformer, model.materialize_fn["transformer"] = self._load_transformer( SD3Transformer2DModel, - weight_dtypes.transformer, - weight_dtypes.train_dtype, + weight_dtypes, base_model_name, - "transformer", + "", quantization, + stream_from_disk=stream_from_disk, ) - 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 - model.text_encoder_3 = text_encoder_3 - model.vae = vae - model.transformer = transformer - def __load_safetensors( self, model: StableDiffusion3Model, @@ -243,6 +238,7 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] @@ -250,7 +246,7 @@ def load( self.__load_internal( model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, model_names.include_text_encoder, model_names.include_text_encoder_2, - model_names.include_text_encoder_3, quantization, + model_names.include_text_encoder_3, quantization, stream_from_disk, ) return except Exception: @@ -260,12 +256,17 @@ def load( self.__load_diffusers( model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, model_names.include_text_encoder, model_names.include_text_encoder_2, - model_names.include_text_encoder_3, quantization, + model_names.include_text_encoder_3, quantization, stream_from_disk, ) return except Exception: stacktraces.append(traceback.format_exc()) + if stream_from_disk: + # the single-file loader below builds a full pipeline via from_single_file, which can't stream; fall + # back to loading it fully into RAM. + print(f"Warning: 'stream from disk' is not supported for single-file {model_type}; loading fully into RAM.") + try: self.__load_safetensors( model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, diff --git a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py index ba340006a..216e0186f 100644 --- a/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py +++ b/modules/modelLoader/stableDiffusionXL/StableDiffusionXLModelLoader.py @@ -49,9 +49,11 @@ def __load_internal( base_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): if os.path.isfile(os.path.join(base_model_name, "meta.json")): - self.__load_diffusers(model, model_type, weight_dtypes, base_model_name, vae_model_name, quantization) + self.__load_diffusers( + model, model_type, weight_dtypes, base_model_name, vae_model_name, quantization, stream_from_disk) else: raise Exception("not an internal model") @@ -63,43 +65,48 @@ def __load_diffusers( base_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - tokenizer_1 = CLIPTokenizer.from_pretrained( + model.tokenizer_1 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer", ) + model.orig_tokenizer_1 = copy.deepcopy(model.tokenizer_1) - tokenizer_2 = CLIPTokenizer.from_pretrained( + model.tokenizer_2 = CLIPTokenizer.from_pretrained( base_model_name, subfolder="tokenizer_2", ) + model.orig_tokenizer_2 = copy.deepcopy(model.tokenizer_2) noise_scheduler = DDIMScheduler.from_pretrained( base_model_name, subfolder="scheduler", ) - noise_scheduler = create.create_noise_scheduler( + model.noise_scheduler = create.create_noise_scheduler( noise_scheduler=NoiseScheduler.DDIM, original_noise_scheduler=noise_scheduler, ) - text_encoder_1 = self._load_text_encoder( + model.text_encoder_1, model.materialize_fn["text_encoder_1"] = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, base_model_name, "text_encoder", + stream_from_disk=stream_from_disk, ) - text_encoder_2 = self._load_text_encoder( + model.text_encoder_2, model.materialize_fn["text_encoder_2"] = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder_2, weight_dtypes.train_dtype, base_model_name, "text_encoder_2", + stream_from_disk=stream_from_disk, ) - vae = self._load_vae( + model.vae = self._load_vae( AutoencoderKL, weight_dtypes.vae, weight_dtypes.fallback_train_dtype, @@ -107,25 +114,28 @@ def __load_diffusers( vae_model_name, ) - unet = self._load_diffusers_sub_module( - UNet2DConditionModel, - weight_dtypes.unet, - weight_dtypes.train_dtype, - base_model_name, - "unet", - quantization, - ) - - 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 - model.vae = vae - model.unet = unet + # the SDXL UNet has no single-file transformer helper and lives in the "unet" subfolder, so it streams via + # _load_diffusers_sub_module directly, which returns a (module, materialize_fn) pair only when streaming and a + # bare module otherwise (train_dtype is applied per-materialize when streaming, so pass None there). + if stream_from_disk: + model.unet, model.materialize_fn["unet"] = self._load_diffusers_sub_module( + UNet2DConditionModel, + weight_dtypes.unet, + None, + base_model_name, + "unet", + quantization, + stream_from_disk=True, + ) + else: + model.unet = self._load_diffusers_sub_module( + UNet2DConditionModel, + weight_dtypes.unet, + weight_dtypes.train_dtype, + base_model_name, + "unet", + quantization, + ) def __load_ckpt( self, @@ -240,6 +250,7 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] @@ -247,17 +258,26 @@ def load( model.sd_config_filename = self._get_sd_config_name(model_type, model_names.base_model) try: - self.__load_internal(model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, quantization) + self.__load_internal( + model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, quantization, + stream_from_disk) return except Exception: stacktraces.append(traceback.format_exc()) try: - self.__load_diffusers(model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, quantization) + self.__load_diffusers( + model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, quantization, + stream_from_disk) return except Exception: stacktraces.append(traceback.format_exc()) + if stream_from_disk: + # the single-file loaders below build a full pipeline via from_single_file, which can't stream; fall + # back to loading it fully into RAM. + print(f"Warning: 'stream from disk' is not supported for single-file {model_type}; loading fully into RAM.") + try: self.__load_safetensors(model, model_type, weight_dtypes, model_names.base_model, model_names.vae_model, quantization) return diff --git a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py index d8e86a19b..c0d948880 100644 --- a/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py +++ b/modules/modelLoader/wuerstchen/WuerstchenModelLoader.py @@ -62,20 +62,20 @@ def __load_diffusers( quantization: QuantizationConfig, ): if model_type.is_wuerstchen_v2(): - decoder_tokenizer = CLIPTokenizer.from_pretrained( + model.decoder_tokenizer = CLIPTokenizer.from_pretrained( decoder_model_name, subfolder="tokenizer", ) if model_type.is_stable_cascade(): - decoder_tokenizer = None + model.decoder_tokenizer = None - decoder_noise_scheduler = DDPMWuerstchenScheduler.from_pretrained( + model.decoder_noise_scheduler = DDPMWuerstchenScheduler.from_pretrained( decoder_model_name, subfolder="scheduler", ) if model_type.is_wuerstchen_v2(): - decoder_text_encoder = self._load_text_encoder( + model.decoder_text_encoder, _ = self._load_text_encoder( CLIPTextModel, weight_dtypes.decoder_text_encoder, weight_dtypes.train_dtype, @@ -83,10 +83,10 @@ def __load_diffusers( "text_encoder", ) if model_type.is_stable_cascade(): - decoder_text_encoder = None + model.decoder_text_encoder = None if model_type.is_wuerstchen_v2(): - decoder_decoder = self._load_diffusers_sub_module( + model.decoder_decoder = self._load_diffusers_sub_module( WuerstchenDiffNeXt, weight_dtypes.decoder, weight_dtypes.train_dtype, @@ -94,7 +94,7 @@ def __load_diffusers( "decoder", ) elif model_type.is_stable_cascade(): - decoder_decoder = self._load_diffusers_sub_module( + model.decoder_decoder = self._load_diffusers_sub_module( StableCascadeUNet, weight_dtypes.decoder, weight_dtypes.train_dtype, @@ -102,7 +102,7 @@ def __load_diffusers( "decoder", ) - decoder_vqgan = self._load_diffusers_sub_module( + model.decoder_vqgan = self._load_diffusers_sub_module( PaellaVQModel, weight_dtypes.decoder_vqgan, weight_dtypes.train_dtype, @@ -111,7 +111,7 @@ def __load_diffusers( ) if model_type.is_wuerstchen_v2(): - effnet_encoder = self._load_diffusers_sub_module( + model.effnet_encoder = self._load_diffusers_sub_module( WuerstchenEfficientNetEncoder, weight_dtypes.effnet_encoder, weight_dtypes.fallback_train_dtype, @@ -121,12 +121,12 @@ def __load_diffusers( # TODO: this is a temporary workaround until the effnet weights are available in diffusers format effnet_encoder = WuerstchenEfficientNetEncoder(affine_batch_norm=False) effnet_encoder.load_state_dict(load_file(effnet_encoder_model_name)) - effnet_encoder = self._convert_diffusers_sub_module_to_dtype( + model.effnet_encoder = self._convert_diffusers_sub_module_to_dtype( effnet_encoder, weight_dtypes.effnet_encoder, weight_dtypes.fallback_train_dtype ) if model_type.is_wuerstchen_v2(): - prior_prior = self._load_diffusers_sub_module( + model.prior_prior = self._load_diffusers_sub_module( WuerstchenPrior, weight_dtypes.prior, weight_dtypes.train_dtype, @@ -145,11 +145,11 @@ def __load_diffusers( prior_config = json.load(config_file) prior_prior = StableCascadeUNet(**prior_config) prior_prior.load_state_dict(convert_stable_cascade_ckpt_to_diffusers(load_file(prior_prior_model_name))) - prior_prior = self._convert_diffusers_sub_module_to_dtype( + model.prior_prior = self._convert_diffusers_sub_module_to_dtype( prior_prior, weight_dtypes.prior, weight_dtypes.fallback_train_dtype, quantization, ) else: - prior_prior = self._load_diffusers_sub_module( + model.prior_prior = self._load_diffusers_sub_module( StableCascadeUNet, weight_dtypes.prior, weight_dtypes.fallback_train_dtype, @@ -158,13 +158,14 @@ def __load_diffusers( quantization, ) - prior_tokenizer = CLIPTokenizer.from_pretrained( + model.prior_tokenizer = CLIPTokenizer.from_pretrained( prior_model_name, subfolder="tokenizer", ) + model.orig_prior_tokenizer = copy.deepcopy(model.prior_tokenizer) if model_type.is_wuerstchen_v2(): - prior_text_encoder = self._load_text_encoder( + model.prior_text_encoder, _ = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -172,7 +173,7 @@ def __load_diffusers( "text_encoder", ) elif model_type.is_stable_cascade(): - prior_text_encoder = self._load_text_encoder( + model.prior_text_encoder, _ = self._load_text_encoder( CLIPTextModelWithProjection, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -180,24 +181,11 @@ def __load_diffusers( "text_encoder", ) - prior_noise_scheduler = DDPMWuerstchenScheduler.from_pretrained( + model.prior_noise_scheduler = DDPMWuerstchenScheduler.from_pretrained( prior_model_name, subfolder="scheduler", ) - model.model_type = model_type - model.decoder_tokenizer = decoder_tokenizer - model.decoder_noise_scheduler = decoder_noise_scheduler - model.decoder_text_encoder = decoder_text_encoder - model.decoder_decoder = decoder_decoder - 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 - def load( self, model: WuerstchenModel, @@ -205,9 +193,15 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] + if stream_from_disk: + # not supported: Stable Cascade loads its prior (single-file override) and effnet encoder by + # constructing the module and calling load_state_dict directly, which can't stream from a meta skeleton. + print(f"Warning: 'stream from disk' is not supported for {model_type}; loading the model fully into RAM.") + prior_model_name = model_names.base_model prior_prior_model_name = model_names.prior_model effnet_encoder_model_name = model_names.effnet_encoder_model diff --git a/modules/modelSampler/AnimaSampler.py b/modules/modelSampler/AnimaSampler.py index 974d3482a..d5de18ac1 100644 --- a/modules/modelSampler/AnimaSampler.py +++ b/modules/modelSampler/AnimaSampler.py @@ -126,8 +126,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = self.image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/ChromaSampler.py b/modules/modelSampler/ChromaSampler.py index 15441b164..23ef4aee5 100644 --- a/modules/modelSampler/ChromaSampler.py +++ b/modules/modelSampler/ChromaSampler.py @@ -145,8 +145,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/ErnieSampler.py b/modules/modelSampler/ErnieSampler.py index 78aee2405..a9cb57e0a 100644 --- a/modules/modelSampler/ErnieSampler.py +++ b/modules/modelSampler/ErnieSampler.py @@ -120,8 +120,6 @@ def __sample_base( image = image.cpu().permute(0, 2, 3, 1).float().numpy() image = [PILImage.fromarray((img * 255).astype(np.uint8)) for img in image] - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/Flux2Sampler.py b/modules/modelSampler/Flux2Sampler.py index 6653840fd..7ecbd5c83 100644 --- a/modules/modelSampler/Flux2Sampler.py +++ b/modules/modelSampler/Flux2Sampler.py @@ -145,8 +145,6 @@ def __sample_base( image = image_processor.postprocess(image, output_type='pil') - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/FluxSampler.py b/modules/modelSampler/FluxSampler.py index a7fb6a34d..fc2b2e0c8 100644 --- a/modules/modelSampler/FluxSampler.py +++ b/modules/modelSampler/FluxSampler.py @@ -158,8 +158,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] #TODO remove and test, from Flux and other models. True is the default image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], @@ -377,8 +375,6 @@ def __sample_inpainting( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py index eefc3a358..c5b30723b 100644 --- a/modules/modelSampler/HiDreamSampler.py +++ b/modules/modelSampler/HiDreamSampler.py @@ -147,8 +147,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py index 54f508f1d..10b22bfc9 100644 --- a/modules/modelSampler/HunyuanVideoSampler.py +++ b/modules/modelSampler/HunyuanVideoSampler.py @@ -143,8 +143,6 @@ def __sample_base( image = video_processor.postprocess(image, output_type='pt') - self.model.evict() - is_image = image.shape[2] == 1 if is_image: image = image.view((image.shape[0], image.shape[1], image.shape[3], image.shape[4])) diff --git a/modules/modelSampler/IdeogramSampler.py b/modules/modelSampler/IdeogramSampler.py index cfc5862b0..cb253db0e 100644 --- a/modules/modelSampler/IdeogramSampler.py +++ b/modules/modelSampler/IdeogramSampler.py @@ -198,8 +198,6 @@ def pack_conditioning(text_features: torch.Tensor, text_lengths: torch.Tensor) - image = image.cpu().permute(0, 2, 3, 1).float().numpy() image = [PILImage.fromarray((img * 255).astype(np.uint8)) for img in image] - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/Krea2Sampler.py b/modules/modelSampler/Krea2Sampler.py index e5f2a68c4..b83205a93 100644 --- a/modules/modelSampler/Krea2Sampler.py +++ b/modules/modelSampler/Krea2Sampler.py @@ -135,8 +135,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/PixArtAlphaSampler.py b/modules/modelSampler/PixArtAlphaSampler.py index 4f21e5250..f8c38f135 100644 --- a/modules/modelSampler/PixArtAlphaSampler.py +++ b/modules/modelSampler/PixArtAlphaSampler.py @@ -148,8 +148,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/QwenSampler.py b/modules/modelSampler/QwenSampler.py index 4ac4c3c17..c18eece7e 100644 --- a/modules/modelSampler/QwenSampler.py +++ b/modules/modelSampler/QwenSampler.py @@ -145,8 +145,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/SanaSampler.py b/modules/modelSampler/SanaSampler.py index 17d0645b6..f4089a222 100644 --- a/modules/modelSampler/SanaSampler.py +++ b/modules/modelSampler/SanaSampler.py @@ -133,8 +133,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/StableDiffusion3Sampler.py b/modules/modelSampler/StableDiffusion3Sampler.py index eaea37f06..f21f34627 100644 --- a/modules/modelSampler/StableDiffusion3Sampler.py +++ b/modules/modelSampler/StableDiffusion3Sampler.py @@ -145,8 +145,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/StableDiffusionSampler.py b/modules/modelSampler/StableDiffusionSampler.py index 1881a15bf..791290edc 100644 --- a/modules/modelSampler/StableDiffusionSampler.py +++ b/modules/modelSampler/StableDiffusionSampler.py @@ -160,8 +160,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], @@ -360,8 +358,6 @@ def __sample_inpainting( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/StableDiffusionVaeSampler.py b/modules/modelSampler/StableDiffusionVaeSampler.py index 8ce49dde7..787b0cfdd 100644 --- a/modules/modelSampler/StableDiffusionVaeSampler.py +++ b/modules/modelSampler/StableDiffusionVaeSampler.py @@ -69,8 +69,6 @@ def sample( latent_image_tensor = self.model.vae.encode(image_tensor.unsqueeze(0)).latent_dist.mean image_tensor = self.model.vae.decode(latent_image_tensor).sample.squeeze() - self.model.evict() - image_tensor = (image_tensor + 1) * 0.5 image_tensor = image_tensor.clamp(0, 1) diff --git a/modules/modelSampler/StableDiffusionXLSampler.py b/modules/modelSampler/StableDiffusionXLSampler.py index e351dac78..93d9f23d3 100644 --- a/modules/modelSampler/StableDiffusionXLSampler.py +++ b/modules/modelSampler/StableDiffusionXLSampler.py @@ -183,8 +183,6 @@ def __sample_base( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], @@ -424,8 +422,6 @@ def __sample_inpainting( do_denormalize = [True] * image.shape[0] image = image_processor.postprocess(image, output_type='pil', do_denormalize=do_denormalize) - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/WuerstchenSampler.py b/modules/modelSampler/WuerstchenSampler.py index ac18ae6bc..a1e9d8ba0 100644 --- a/modules/modelSampler/WuerstchenSampler.py +++ b/modules/modelSampler/WuerstchenSampler.py @@ -313,8 +313,6 @@ def __sample_base( image_array = image_tensor.permute(0, 2, 3, 1).cpu().squeeze().float().numpy() image_array = (image_array * 255).round().astype("uint8") - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=Image.fromarray(image_array), diff --git a/modules/modelSampler/ZImageSampler.py b/modules/modelSampler/ZImageSampler.py index 74cf22bb4..0e001df2a 100644 --- a/modules/modelSampler/ZImageSampler.py +++ b/modules/modelSampler/ZImageSampler.py @@ -121,8 +121,6 @@ def __sample_base( image = image_processor.postprocess(image, output_type='pil') - self.model.evict() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index c4d6df734..adf449e49 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -248,7 +248,5 @@ def calculate_loss( def prepare_text_caching(self, model: ChromaModel, config: TrainConfig): if not config.train_text_encoder_or_embedding(): model.materialize_only("text_encoder") - else: - model.evict() model.eval() diff --git a/modules/modelSetup/BaseKrea2Setup.py b/modules/modelSetup/BaseKrea2Setup.py index 9701a4eb0..bacf50741 100644 --- a/modules/modelSetup/BaseKrea2Setup.py +++ b/modules/modelSetup/BaseKrea2Setup.py @@ -162,7 +162,5 @@ def calculate_loss( def prepare_text_caching(self, model: Krea2Model, config: TrainConfig): if not config.train_text_encoder_or_embedding(): model.materialize_only("text_encoder") - else: - model.evict() model.eval() diff --git a/modules/modelSetup/BaseModelSetup.py b/modules/modelSetup/BaseModelSetup.py index fea9c94b9..2a4b37f2b 100644 --- a/modules/modelSetup/BaseModelSetup.py +++ b/modules/modelSetup/BaseModelSetup.py @@ -236,6 +236,13 @@ def _setup_model_part_requires_grad( not self.__stop_model_part_training_elapsed(unique_name, config, train_progress) model.requires_grad_(train_model_part) + # a streamed part (loaded as a meta skeleton) with cache-in-ram off is dropped to meta and re-streamed from + # the checkpoint on every reload, so training it would discard the update. Refuse the combination early. + if train_model_part and not config.cache_in_ram and any(p.is_meta for p in model.parameters()): + raise ValueError( + f"'{unique_name}' is trained with 'stream from disk' on and 'cache in ram' off -- the trained " + f"weights would be re-streamed from the checkpoint and lost. Enable 'cache in ram' for this part.") + #even if frozen parameters are not passed to the optimizer, required_grad has to be False. #otherwise, gradients accumulate in param.grad and waste vram if unique_name in self.frozen_parameters: @@ -253,18 +260,16 @@ def _setup_model_part( disable_fp16_autocast: bool = False, attention_mask: bool | None = None, ): - # Per-part optimization wiring, called once per model part from each leaf. The optional - # disable_fp16_autocast context and its dtype are stored per-part and can differ per part - # (e.g. HiDream disables fp16 for both text_encoder_3 and the transformer). checkpointing_fn returns - # None for non-offloadable parts (SD/SDXL UNet), so no conductor is stored for those. module = getattr(model, attr) if module is None: return + materialize_fn = model.materialize_fn.get(attr) + if checkpointing_fn is not None: conductor = checkpointing_fn(module, config, config_part) if conductor is not None: - setattr(model, f"{attr}_offload_conductor", conductor) + model.offload_conductor[attr] = conductor if disable_fp16_autocast: autocast_context, train_dtype = disable_fp16_autocast_context( @@ -274,7 +279,10 @@ def _setup_model_part( else: train_dtype = model.train_dtype - quantize_layers(module, self.train_device, train_dtype, config) + # a streamed module (materialize_fn set) stays on meta until materialized and is quantized per-materialize, + # so there is nothing to quantize here; a non-streamed module is quantized now. + if materialize_fn is None: + quantize_layers(module, self.train_device, train_dtype, config) if attention_mask is not None: self._set_attention_backend(module, config.attention_mechanism, mask=attention_mask) diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 875bc4a8e..c169cafe3 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -312,7 +312,5 @@ def calculate_loss( def prepare_text_caching(self, model: PixArtAlphaModel, config: TrainConfig): if not config.train_text_encoder_or_embedding(): model.materialize_only("text_encoder") - else: - model.evict() model.eval() diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index ee42de6b2..01852ae42 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -161,7 +161,5 @@ def calculate_loss( def prepare_text_caching(self, model: QwenModel, config: TrainConfig): if not config.train_text_encoder_or_embedding(): model.materialize_only("text_encoder") - else: - model.evict() model.eval() diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 0c979aa32..b3d2137b0 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -238,7 +238,5 @@ def calculate_loss( def prepare_text_caching(self, model: SanaModel, config: TrainConfig): if not config.train_text_encoder_or_embedding(): model.materialize_only("text_encoder") - else: - model.evict() model.eval() diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 56179aff1..0e10939ef 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -335,8 +335,6 @@ def calculate_loss( def prepare_text_caching(self, model: StableDiffusionModel, config: TrainConfig): if not config.train_text_encoder_or_embedding(): model.materialize_only("text_encoder") - else: - model.evict() if model.depth_estimator is not None: model.depth_estimator.to(self.temp_device) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index d2bbb200a..c6b807127 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -348,7 +348,5 @@ def calculate_loss( def prepare_text_caching(self, model: WuerstchenModel, config: TrainConfig): if not config.train_text_encoder_or_embedding(): model.materialize_only("text_encoder") - else: - model.evict() model.eval() diff --git a/modules/module/AdditionalEmbeddingWrapper.py b/modules/module/AdditionalEmbeddingWrapper.py index 573bcbb96..48cf0fd73 100644 --- a/modules/module/AdditionalEmbeddingWrapper.py +++ b/modules/module/AdditionalEmbeddingWrapper.py @@ -30,7 +30,12 @@ def __init__( self.is_applied = False self.orig_forward = self.orig_module.forward - self.orig_median_norm = torch.norm(self.orig_module.weight, dim=1).median().item() + # orig_median_norm is only read by normalize_embeddings(), which only touches learned embeddings. A text + # encoder left on meta (streamed but not materialized, because none of its embeddings are trained) never + # reaches that path, so skip the norm read that would otherwise fail on a meta tensor (#69). + self.orig_median_norm = None + if not self.orig_module.weight.is_meta: + self.orig_median_norm = torch.norm(self.orig_module.weight, dim=1).median().item() def forward(self, x, *args, **kwargs): # ensure that the original weights only contain as many embeddings as the unmodified tokenizer can create diff --git a/modules/module/quantized/LinearFp8.py b/modules/module/quantized/LinearFp8.py index 823572425..230b8dbe6 100644 --- a/modules/module/quantized/LinearFp8.py +++ b/modules/module/quantized/LinearFp8.py @@ -17,16 +17,25 @@ def __init__(self, *args, **kwargs): self.is_quantized = False self.fp8_dtype = torch.float8_e4m3fn - self._scale = torch.tensor(1.0, dtype=torch.float) - self.register_buffer("scale", self._scale) + self.register_buffer("scale", torch.tensor(1.0, dtype=torch.float)) self.compute_dtype = None def original_weight_shape(self) -> tuple[int, ...]: return self.weight.shape + def mark_needs_requantization(self): + self.is_quantized = False + + def predict_offload_bytes(self) -> int: + # weight quantizes to float8_e4m3fn (1 byte/elem, same shape); bias is left unchanged. Matches + # get_offload_tensors (weight + optional bias); the scalar scale buffer is not offload-counted. + weight_bytes = self.weight.numel() + bias_bytes = self.bias.numel() * self.bias.element_size() if self.bias is not None else 0 + return weight_bytes + bias_bytes + def unquantized_weight(self, dtype: torch.dtype, device: torch.device) -> torch.Tensor: - if self._scale is not None: - return self.weight.detach().to(dtype) * self._scale.to(dtype=dtype) + if self.scale is not None: + return self.weight.detach().to(dtype) * self.scale.to(dtype=dtype) else: return self.weight.detach().to(dtype=dtype) @@ -42,19 +51,22 @@ def quantize(self, device: torch.device | None = None): weight = weight.to(device=device) abs_max = weight.abs().max() - self._scale.copy_(torch.clamp(abs_max, min=1e-12) / torch.finfo(self.fp8_dtype).max) - weight = weight.div_(self._scale).to(dtype=self.fp8_dtype) + scale = torch.clamp(abs_max, min=1e-12) / torch.finfo(self.fp8_dtype).max + weight = weight.div_(scale).to(dtype=self.fp8_dtype) if device is not None: weight = weight.to(device=orig_device) + + # keep the scale on the weight's device (see LinearW8A8.quantize) + self.scale = scale.detach().to(orig_device) self.weight.data = weight def forward(self, x: torch.Tensor) -> torch.Tensor: weight = self.weight.detach() weight = weight.to(dtype=self.compute_dtype if self.compute_dtype is not None else x.dtype) - if self._scale is not None: - weight = weight.mul_(self._scale) + if self.scale is not None: + weight = weight.mul_(self.scale) x = nn.functional.linear(x, weight, self.bias) return x diff --git a/modules/module/quantized/LinearNf4.py b/modules/module/quantized/LinearNf4.py index 2a4bfbf17..718b65856 100644 --- a/modules/module/quantized/LinearNf4.py +++ b/modules/module/quantized/LinearNf4.py @@ -38,7 +38,21 @@ def __init__(self, *args, **kwargs): self.quant_state = None def original_weight_shape(self) -> tuple[int, ...]: - return self.weight.shape + # self.weight is repacked to a flat [N, 1] uint8 layout once quantized; self.shape keeps the original. + return self.shape + + def mark_needs_requantization(self): + self.is_quantized = False + + def predict_offload_bytes(self) -> int: + # nf4 packs the weight to 4-bit (2 values per uint8), and with double quant (compress_statistics) stores + # quant_state.absmax as one uint8 per block_size elements. Matches get_offload_tensors (packed weight + + # quant_state.absmax + optional bias); the small code/offset/nested-absmax buffers are not offload-counted. + numel = self.shape.numel() + weight_bytes = (numel + 1) // 2 + absmax_bytes = (numel + self.block_size - 1) // self.block_size + bias_bytes = self.bias.numel() * self.bias.element_size() if self.bias is not None else 0 + return weight_bytes + absmax_bytes + bias_bytes def unquantized_weight(self, dtype: torch.dtype, device: torch.device) -> torch.Tensor: if self.is_quantized: diff --git a/modules/module/quantized/LinearSVD.py b/modules/module/quantized/LinearSVD.py index 16e2f2650..3fcb3b777 100644 --- a/modules/module/quantized/LinearSVD.py +++ b/modules/module/quantized/LinearSVD.py @@ -1,3 +1,4 @@ +import os from abc import abstractmethod from contextlib import suppress @@ -51,6 +52,19 @@ def unquantized_weight(self, dtype: torch.dtype, device: torch.device) -> torch. else: return super().unquantized_weight(dtype, device) + def mark_needs_requantization(self): + # reset both the SVD split flag and the parent's base-weight flag so the next quantize() re-runs fully. + self.__svd_is_quantized = False + super().mark_needs_requantization() + + def predict_offload_bytes(self) -> int: + # the residual quantized weight (base quant type) plus the low-rank factors svd_up (out x rank) and + # svd_down (rank x in), both in svd_dtype. Sized from the meta skeleton -- the factors don't exist yet. + out_features, in_features = self.original_weight_shape() + svd_bytes = (out_features * self.rank + self.rank * in_features) \ + * torch.empty((), dtype=self.svd_dtype).element_size() + return super().predict_offload_bytes() + svd_bytes + @torch.no_grad() def quantize(self, device: torch.device | None = None): if self.__svd_is_quantized: @@ -73,11 +87,17 @@ def quantize(self, device: torch.device | None = None): U, S, Vh = torch.linalg.svd(W, full_matrices=False) if self.cache_dir is not None: + # write to a per-process temp then atomically rename in: under multi-GPU every rank quantizes + # concurrently and writes the same hash-named file, so a plain torch.save races and a reader can + # pick up a half-written file. os.replace is atomic on the same filesystem, so a concurrent reader + # sees either no file or a complete one, and multiple writers just overwrite with identical content. + tmp_filename = filename + f".tmp.{os.getpid()}" torch.save(( U[:, :self.max_cache_rank].clone(), S[:self.max_cache_rank].clone(), Vh[:self.max_cache_rank, :].clone(), - ), filename) + ), tmp_filename) + os.replace(tmp_filename, filename) U_r = U[:, :self.rank] S_r = S[:self.rank] diff --git a/modules/module/quantized/LinearW8A8.py b/modules/module/quantized/LinearW8A8.py index babbd1a80..a8d0432dd 100644 --- a/modules/module/quantized/LinearW8A8.py +++ b/modules/module/quantized/LinearW8A8.py @@ -80,27 +80,40 @@ class LinearW8A8( QuantizedModuleMixin, QuantizedLinearMixin, ): + is_quantized: bool + def __init__(self, dtype, *args, **kwargs): super().__init__(*args, **kwargs) assert dtype in [torch.int8, torch.float8_e4m3fn] self._dtype = dtype - self.__is_quantized = False + self.is_quantized = False self.compute_dtype = None self.register_buffer("scale", torch.tensor(1.0, dtype=torch.float32)) def original_weight_shape(self) -> tuple[int, ...]: return self.weight.shape + def mark_needs_requantization(self): + self.is_quantized = False + + def predict_offload_bytes(self) -> int: + # weight quantizes tensorwise to int8/float8_e4m3fn (both 1 byte/elem, same shape); bias is left + # unchanged. Matches get_offload_tensors (weight + optional bias); the scalar scale buffer is not + # offload-counted. _dtype is asserted int8/float8_e4m3fn in __init__, so 1 byte/elem is exact. + weight_bytes = self.weight.numel() + bias_bytes = self.bias.numel() * self.bias.element_size() if self.bias is not None else 0 + return weight_bytes + bias_bytes + def unquantized_weight(self, dtype: torch.dtype, device: torch.device) -> torch.Tensor: return dequantize(self.weight.detach(), self.scale).to(dtype) @torch.no_grad() def quantize(self, device: torch.device | None = None): - if self.__is_quantized: + if self.is_quantized: return - self.__is_quantized = True + self.is_quantized = True weight = self.weight.detach() orig_device = weight.device @@ -117,11 +130,12 @@ def quantize(self, device: torch.device | None = None): self.requires_grad_(False) self.weight.data = weight - self.scale.copy_(scale) + # keep the scale on the weight's device so the batched int8/fp8 path finds it co-located there + self.scale = scale.detach().to(orig_device) def forward(self, x_orig: torch.Tensor) -> torch.Tensor: assert not self.weight.requires_grad - assert self.__is_quantized + assert self.is_quantized x = x_orig.reshape(-1, x_orig.shape[-1]) if x.shape[0] > 16: diff --git a/modules/module/quantized/mixin/QuantizedLinearMixin.py b/modules/module/quantized/mixin/QuantizedLinearMixin.py index a81ccefa3..0dde8f022 100644 --- a/modules/module/quantized/mixin/QuantizedLinearMixin.py +++ b/modules/module/quantized/mixin/QuantizedLinearMixin.py @@ -11,3 +11,15 @@ def original_weight_shape(self) -> tuple[int, ...]: @abstractmethod def unquantized_weight(self, dtype: torch.dtype, device: torch.device) -> torch.Tensor: pass + + @abstractmethod + def mark_needs_requantization(self): + # reset the concrete class's is-quantized flag so the next materialize re-quantizes. Called by streaming + # eviction, which discards the packed weights back to meta. + pass + + def predict_offload_bytes(self) -> int: + # post-quantization offload footprint, predicted from the unpacked skeleton shape while the module is still a + # meta skeleton (the real packed tensors don't exist yet). + raise NotImplementedError( + f"{type(self).__name__} does not implement predict_offload_bytes (disk-offload conductor sizing)") diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index 380ba39d8..c39ecb1b7 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -131,6 +131,8 @@ def start(self): model_names=model_names, weight_dtypes=self.config.weight_dtypes(), quantization=self.config.quantization, + stream_from_disk=self.config.stream_from_disk, + cache_in_ram=self.config.cache_in_ram(), ) self.model.train_config = self.config @@ -139,7 +141,6 @@ 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.evict() self.model.eval() self.callbacks.on_update_status("creating the data loader/caching") @@ -253,7 +254,6 @@ 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.evict() self.model.eval() sample_config = copy.copy(sample_config) diff --git a/modules/ui/BaseModelTabView.py b/modules/ui/BaseModelTabView.py index 34a9ea8dd..7ce8225ea 100644 --- a/modules/ui/BaseModelTabView.py +++ b/modules/ui/BaseModelTabView.py @@ -103,6 +103,13 @@ def __create_base_dtype_components(self, frame, row: int, ui_state) -> int: row += 1 + # stream from disk + self.components.label(frame, row, 0, "Stream From Disk", + tooltip="Uses the streaming model loader to stream frozen weights from disk to VRAM on demand, greatly reducing RAM usage. Only turn off if you hit compatibility issues.") + self.components.switch(frame, row, 1, ui_state, "stream_from_disk") + + row += 1 + return row def __create_base_components( diff --git a/modules/ui/BaseTrainingTabView.py b/modules/ui/BaseTrainingTabView.py index 29a54bb1e..6aae18012 100644 --- a/modules/ui/BaseTrainingTabView.py +++ b/modules/ui/BaseTrainingTabView.py @@ -461,6 +461,11 @@ def __create_offloading_widgets(self, frame, row, ui_state, part, supports_check self.components.switch(frame, row, 1, ui_state, f"{part}.activation_offloading") row += 1 + self.components.label(frame, row, 0, "Cache In RAM", + tooltip="Keeps this model part's streamed weights in RAM between uses instead of re-reading them from disk on every use, trading RAM for loading speed. Only has an effect when \"Stream From Disk\" (model page) is enabled.") + self.components.switch(frame, row, 1, ui_state, f"{part}.cache_in_ram") + row += 1 + return row def __create_text_encoder_frame(self, master, row, ui_state, supports_clip_skip=True, supports_training=True, diff --git a/modules/ui/SampleWindowController.py b/modules/ui/SampleWindowController.py index 45a02212f..11e47a389 100644 --- a/modules/ui/SampleWindowController.py +++ b/modules/ui/SampleWindowController.py @@ -90,13 +90,14 @@ def load_model(self) -> BaseModel: model_names=model_names, weight_dtypes=self.initial_train_config.weight_dtypes(), quantization=self.initial_train_config.quantization, + stream_from_disk=self.initial_train_config.stream_from_disk, + cache_in_ram=self.initial_train_config.cache_in_ram(), ) model.train_config = self.initial_train_config 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.evict() return model @@ -145,3 +146,7 @@ def do_sample(self, on_sample, on_update_progress): on_sample=on_sample, on_update_progress=on_update_progress, ) + + # the sampler materializes parts on demand and no longer self-evicts; + # release VRAM now that this standalone sample window is idle again + self.model.evict() diff --git a/modules/util/LayerOffloadConductor.py b/modules/util/LayerOffloadConductor.py index 54dcef3d6..663a49a88 100644 --- a/modules/util/LayerOffloadConductor.py +++ b/modules/util/LayerOffloadConductor.py @@ -1,9 +1,17 @@ import math import random +from collections.abc import Callable from typing import Any from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig -from modules.util.quantization_util import get_offload_tensor_bytes, get_offload_tensors, offload_quantized +from modules.util.disk_stream import _is_evicted, evict_to_meta +from modules.util.enum.DataType import DataType +from modules.util.quantization_util import ( + get_offload_tensor_bytes, + get_offload_tensors, + is_quantized_module, + offload_quantized, +) from modules.util.torch_util import ( create_mem_pool, create_stream_context, @@ -22,6 +30,8 @@ import torch from torch import nn +from tqdm import tqdm + MESSAGES = [] @@ -36,7 +46,9 @@ def clone_tensor_allocator(tensor: torch.Tensor) -> torch.Tensor: return tensor.clone() -# allocate_like places each cached tensor at an aligned offset, wasting up to this many bytes per tensor +# allocate_like places each cached tensor at an aligned offset, wasting up to this many bytes per tensor. +# also the reserved size at the start of each cache tensor (see allocate_like); must stay >= 2 so no view +# ever lands at storage_offset 0 or 1, the two values torch.compile bakes into separate specialized graphs TENSOR_ALIGNMENT_BYTES = 16 @@ -75,16 +87,22 @@ def allocate_like(self, source_tensor: torch.Tensor) -> torch.Tensor: total_cache_bytes = cache_tensor_size * len(self.__layer_allocator.cache_tensors) if self.__allocate_forward: cache_tensor_index = self.__allocation_end // cache_tensor_size - cache_tensor_allocation_end = align_up(self.__allocation_end % cache_tensor_size) + # never hand out views at storage_offset 0: torch.compile creates a 0/1-specialized + # symbol for the storage_offset of any tensor with a dynamic dim, so an offset-0 view + # needs its own graph while one "2 <= offset" guard covers all + # other placements. keeping every view past the first alignment slot avoids those + # recompiles, and costs each tensor at most its alignment budget (the first tensor in a + # cache tensor previously wasted 0 of it) + cache_tensor_allocation_end = max(align_up(self.__allocation_end % cache_tensor_size), TENSOR_ALIGNMENT_BYTES) if cache_tensor_allocation_end + num_bytes > cache_tensor_size: # move to the start of the next cache tensor cache_tensor_index += 1 - cache_tensor_allocation_end = 0 + cache_tensor_allocation_end = TENSOR_ALIGNMENT_BYTES if cache_tensor_index * cache_tensor_size + cache_tensor_allocation_end + num_bytes > total_cache_bytes: # move to the first cache tensor cache_tensor_index = 0 - cache_tensor_allocation_end = 0 + cache_tensor_allocation_end = TENSOR_ALIGNMENT_BYTES self.__allocation_end = cache_tensor_index * cache_tensor_size + cache_tensor_allocation_end self.__layer_allocator.ensure_allocation(cache_tensor_index) @@ -97,7 +115,10 @@ def allocate_like(self, source_tensor: torch.Tensor) -> torch.Tensor: cache_tensor_index = self.__allocation_start // cache_tensor_size cache_tensor_allocation_start = self.__allocation_start % cache_tensor_size - if cache_tensor_allocation_start - num_bytes < 0: + # "< TENSOR_ALIGNMENT_BYTES" instead of "< 0": the first alignment slot of every cache + # tensor is reserved so no view lands at storage_offset 0 (see the forward-direction + # comment above) + if cache_tensor_allocation_start - num_bytes < TENSOR_ALIGNMENT_BYTES: # move to the end of the previous cache tensor cache_tensor_index -= 1 cache_tensor_allocation_start = cache_tensor_size @@ -161,7 +182,7 @@ def __init__( self.__mem_pool = None - def allocate_cache(self, layers: list[nn.Module], target_bytes: int): + def allocate_cache(self, layers: list[nn.Module], target_bytes: int, streaming: bool, cache_in_ram: bool): if not self.__allocate_statically or any(x is not None for x in self.cache_tensors): return @@ -187,13 +208,16 @@ def allocate_cache(self, layers: list[nn.Module], target_bytes: int): # The GPU cache is filled one layer at a time from the CPU, so the destination buffer and a full # resident source never coexist on the device -- no peak-doubling to guard against here. num_cache_tensors = 1 + elif streaming and not cache_in_ram: + # host/pinned cache, disk-streaming with cache_in_ram off: layers stream+quantize straight from the + # checkpoint and evict back to meta, so no resident copy ever coexists with the pinned cache -- none of the + # peak-doubling that justifies chunking below. A single large pinned buffer is fine: pin_tensor_ page-locks + # the existing scattered pages in place, and the CPU allocator has no pool to fragment. Same as the GPU cache. + num_cache_tensors = 1 else: - # host/pinned cache keeps the multi-chunk split: the chunks are allocated lazily (per ensure_allocation) - # to cap peak host RAM while the resident model is copied into the pinned cache, which a single eager - # buffer would roughly double. - # TODO once the disk-streaming (load_on_demand) path lands here: collapse to a single buffer too when - # the layers are streamed per-part from the checkpoint -- then no resident model copy coexists with the - # pinned cache, so the peak-doubling that justifies chunking here does not occur. + # host/pinned cache, resident model (classic offload, or streaming with cache_in_ram on): the chunks are + # allocated lazily (per ensure_allocation) to cap peak host RAM while the resident model is copied into + # the pinned cache (and, on evict, cloned back out of it), which a single eager buffer would roughly double. num_cache_tensors = min( # no more than 10% overhead math.ceil(int(cache_bytes * 0.10) / self.__max_tensor_bytes), @@ -250,6 +274,15 @@ def deallocate_cache(self): self.__mem_pool = None torch_gc() + @property + def mem_pool(self): + # the MemPool holding this allocator's cache tensor(s); also used to keep the conductor's resident non-layer + # remainder out of the default pool. allocate_cache creates it before the materialize layer loop; create it + # here too in case a caller reaches for it first. deallocate_cache drops it (static allocators only). + if self.__mem_pool is None: + self.__mem_pool = create_mem_pool(self.device) + return self.__mem_pool + def get_allocator(self, layer_index: int, allocate_forward: bool) -> StaticLayerTensorAllocator | None: if self.__allocate_statically: allocator = StaticLayerTensorAllocator(self, allocate_forward, layer_index) @@ -599,12 +632,16 @@ class LayerOffloadConductor: __is_forward_pass: bool __keep_graph: bool - __is_active: bool + __materialized: bool __deferred_layers: list[int] __config: TrainConfig + __disk_remainder_materialized: bool # whether the non-layer remainder (embedders/norms/proj) has been streamed since the last evict + __disk_layer_key_prefixes: list[str] # per-layer (indexed like __layers) checkpoint-absolute path, so a single layer subtree can be streamed on its own + __disk_module_name_by_id: dict[int, str] # module-name snapshot taken pre-wrapping, used to build the key prefixes above + def __init__( self, module: nn.Module, @@ -652,54 +689,97 @@ def __init__( self.__is_forward_pass = False self.__keep_graph = False - self.__is_active = False + self.__materialized = False self.__deferred_layers = [] self.__config = config + self.__disk_remainder_materialized = False + self.__disk_layer_key_prefixes = [] + self.__disk_module_name_by_id = {id(m): name for name, m in module.named_modules()} + def offload_activated(self) -> bool: return self.__offload_activations or self.__offload_layers - def to(self, device: torch.device): + def evict(self, to_meta: bool = False): torch_gc() self.__wait_all_layer_transfers() self.__wait_all_activation_transfers() - if device_equals(device, self.__temp_device): - log("to temp device") + log("to temp device") - # deallocate the cache before to take advantage of the gc - self.__train_device_layer_allocator.deallocate_cache() - self.__temp_device_layer_allocator.deallocate_cache() - self.__temp_device_activations_allocator.deallocate_cache() + if to_meta: + if self.__materialized: + self.__evict_to_meta() + else: + self.__evict_to_temp() + + def __evict_to_temp(self): + # move every layer and the non-layer remainder back to the temp device and free the static caches (the + # non-disk eviction path). Also the rollback for a resident conductor whose materialize() raised partway. + # deallocate the cache before to take advantage of the gc + self.__train_device_layer_allocator.deallocate_cache() + self.__temp_device_layer_allocator.deallocate_cache() + self.__temp_device_activations_allocator.deallocate_cache() + + self.__module_to_device_except_layers(self.__temp_device) + for layer_index, layer in enumerate(self.__layers): + self.__layers[layer_index].to(self.__temp_device) + for module in layer.modules(): + offload_quantized(module, self.__temp_device, allocator=clone_tensor_allocator) + self.__layer_device_map[layer_index] = None - self.__module_to_device_except_layers(self.__temp_device) - for layer_index, layer in enumerate(self.__layers): - self.__layers[layer_index].to(self.__temp_device) - for module in layer.modules(): - offload_quantized(module, self.__temp_device, allocator=clone_tensor_allocator) - self.__layer_device_map[layer_index] = None + self.__materialized = False - self.__is_active = False + def materialize( + self, train_dtype: DataType | None = None, name: str | None = None, + materialize_fn: Callable | None = None, cache_in_ram: bool = True): + torch_gc() - elif device_equals(device, self.__train_device): - log("to train device") + self.__wait_all_layer_transfers() + self.__wait_all_activation_transfers() + + streaming = materialize_fn is not None + + log("to train device") + try: self.__offload_strategy = LayerOffloadStrategy(self.__layers, self.__layer_offload_fraction) self.__train_device_layer_allocator.allocate_cache( - self.__layers, self.__offload_strategy.max_loaded_bytes) + self.__layers, self.__offload_strategy.max_loaded_bytes, streaming=streaming, cache_in_ram=cache_in_ram) self.__temp_device_layer_allocator.allocate_cache( - self.__layers, self.__offload_strategy.max_offloaded_bytes) - self.__module_to_device_except_layers(self.__train_device) - - # move all layers to the train device, then move offloadable tensors back to the temp device + self.__layers, self.__offload_strategy.max_offloaded_bytes, streaming=streaming, cache_in_ram=cache_in_ram) + # place the resident non-layer remainder onto the train device. When streaming, route it into the conductor + # pool: on a warm cache_in_ram re-activate it comes from cpu/temp and lands there directly (no default-pool + # copy to relocate); on a cold stream it is still meta here and gets skipped, then streamed below. + self.__module_to_device_except_layers( + self.__train_device, + pool=self.__train_device_layer_allocator.mem_pool if streaming else None) + + cold_layers = sum(1 for i, layer in enumerate(self.__layers) + if self.__layer_device_map[i] is None and _is_evicted(layer)) if streaming else 0 + disk_bar = tqdm(total=cold_layers, unit="layer", desc=f"streaming {name}", leave=False) \ + if cold_layers > 0 else None + + # bring each layer to the train device (streaming it if cold), then place it in its cache slot -- the GPU + # cache for an initially-loaded layer, the temp-device cache for an offloaded one for layer_index, layer in enumerate(self.__layers): if self.__layer_device_map[layer_index] is None: log(f"layer {layer_index} to train device") - layer.to(self.__train_device) + if streaming and _is_evicted(layer): + # cold materialize from the checkpoint: stream+quantize this layer onto the train device + # (per tensor, quantized inline). The offload_quantized below then copies it into its + # static cache slot -- GPU cache for a loaded layer, pinned CPU cache for an offloaded one. + materialize_fn( + layer, self.__train_device, train_dtype, + self.__disk_layer_key_prefixes[layer_index]) + if disk_bar is not None: + disk_bar.update(1) + else: + layer.to(self.__train_device) if layer_index in self.__offload_strategy.initial_loaded_layers: allocator = self.__train_device_layer_allocator.get_allocator( @@ -717,9 +797,35 @@ def to(self, device: torch.device): event = SyncEvent(self.__train_stream.record_event(), f"train on {self.__train_device}") self.__layer_train_event_map[layer_index] = event - self.__is_active = True + if disk_bar is not None: + disk_bar.close() + + if streaming and not self.__disk_remainder_materialized: + # the non-layer remainder (embedders/norms/proj) is still meta the first time; stream it to the train + # device now, where it stays resident. dest_pool routes the non-quantized weights straight into the + # conductor pool so no model weight sits in the default pool (which the optimizer state and quantize + # transients draw from). Quantized remainder weights pack in the default pool -- their dequant scratch + # stays out of the pool -- and are relocated into it just below, once small. + materialize_fn(self.__module, self.__train_device, train_dtype, "", + dest_pool=self.__train_device_layer_allocator.mem_pool) + self.__disk_remainder_materialized = True + self.__relocate_quantized_remainder_to_pool() + except Exception: + # a materialize that fails partway (typically OOM) leaves layers/cache tensors resident while + # __materialized is still False, so a later evict() would skip them and strand that VRAM. Force the unit + # back to its pre-materialize state, keyed on the actual weight state: a parameter still on meta means a + # cold disk-stream was in flight, so meta is the only valid target (re-stream next time, lossless since + # frozen); otherwise roll back to the temp device and keep the resident quantized copy. + if any(parameter.is_meta for parameter in self.__module.parameters()): + self.__evict_to_meta() + else: + self.__evict_to_temp() + # the rollback helpers no longer gc, and no caller gc's a failed materialize -- reclaim the stranded VRAM + # here before re-raising (evict() instead relies on BaseModel.evict()'s trailing gc). + torch_gc() + raise - torch_gc() + self.__materialized = True def add_layer(self, layer: nn.Module, included_offload_param_indices: list[int] = None): if included_offload_param_indices is None: @@ -729,13 +835,15 @@ def add_layer(self, layer: nn.Module, included_offload_param_indices: list[int] self.__layer_device_map.append(None) self.__layer_train_event_map.append(SyncEvent()) self.__layer_transfer_event_map.append(SyncEvent()) + # checkpoint-absolute path of this layer, for the per-layer disk stream (empty for a layer built outside self.__module) + self.__disk_layer_key_prefixes.append(self.__disk_module_name_by_id.get(id(layer), "")) self.__layer_activations_included_offload_param_indices_map.append(included_offload_param_indices) def start_forward(self, keep_graph: bool): log("starting forward") - if not self.__is_active: + if not self.__materialized: return if self.__async_transfer: @@ -750,7 +858,7 @@ def before_layer(self, layer_index: int, call_index: int, activations: Any) -> A log() log(f"before layer {layer_index}, {call_index}") - if not self.__is_active: + if not self.__materialized: return activations self.__call_index_layer_index_map[call_index] = layer_index @@ -811,7 +919,7 @@ def before_layer(self, layer_index: int, call_index: int, activations: Any) -> A def after_layer(self, layer_index: int, call_index: int, activations: Any): log(f"after layer {layer_index}, {call_index}") - if not self.__is_active: + if not self.__materialized: return # record stream @@ -831,9 +939,20 @@ def after_layer(self, layer_index: int, call_index: int, activations: Any): def __get_loaded_layers(self) -> list[int]: return [i for i in range(len(self.__layers)) if device_equals(self.__layer_device_map[i], self.__train_device)] + def __evict_to_meta(self): + evict_to_meta(self.__module) + for layer_index in range(len(self.__layers)): + self.__layer_device_map[layer_index] = None + self.__disk_remainder_materialized = False + self.__train_device_layer_allocator.deallocate_cache() + self.__temp_device_layer_allocator.deallocate_cache() + self.__temp_device_activations_allocator.deallocate_cache() + self.__materialized = False + def __module_to_device_except_layers( self, device: torch.device, + pool=None, ): sub_module_parameters = set(sum([list(x.parameters()) for x in self.__layers], [])) @@ -841,10 +960,37 @@ def convert(t): if t in sub_module_parameters or t.is_meta: return t + if pool is not None: + # place the (already-final) non-layer remainder weight straight into the conductor's pool instead of + # the default pool, which the optimizer state and quantize transients allocate from -- a weight left + # there fragments it and strands the region when the remainder is evicted. A weight from cpu/temp (warm + # cache_in_ram re-activate) lands in the pool directly; one already on the train device is relocated + # with a clone. + with mem_pool_context(pool): + return t.clone() if device_equals(t.device, device) else t.to(device=device) + return t.to(device=device) self.__module._apply(convert) + def __relocate_quantized_remainder_to_pool(self): + # the cold remainder stream packs quantized non-layer weights (e.g. a tied lm_head) in the default pool so + # their dequant scratch never enters the conductor pool. Copy just the packed weights into the pool now, so + # no model weight is left in the default pool (where the optimizer state and quantize transients would + # fragment/strand it). Small: the packed weights are a fraction of their fp size. Non-quantized remainder + # weights were streamed straight into the pool (dest_pool) and are not touched here. Layer modules are + # excluded -- they own their static cache slots -- matching __module_to_device_except_layers' scope. + pool = self.__train_device_layer_allocator.mem_pool + + def pool_clone(tensor): + with mem_pool_context(pool): + return tensor.clone() + + layer_modules = {module for layer in self.__layers for module in layer.modules()} + for module in self.__module.modules(): + if module not in layer_modules and is_quantized_module(module): + offload_quantized(module, self.__train_device, allocator=pool_clone) + def __clear_activations(self): self.__activations_map.clear() self.__call_index_layer_index_map.clear() diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index 2df40065e..e975029c0 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -245,6 +245,14 @@ 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 supports_offloading: bool = True, ) -> LayerOffloadConductor | None: + # A full fine-tune updates the base weights, but meta-eviction (stream_from_disk + cache_in_ram off) re-streams + # them from the checkpoint on each use, discarding those updates. Reject that combo. + if config.stream_from_disk and config.part_trained_in_place(part) and not part.cache_in_ram: + raise NotImplementedError( + "a fully fine-tuned component cannot stream from disk without keeping it cached in RAM: it re-streams " + "weights from the checkpoint on each use, discarding training updates. Enable 'Cache In RAM' for this " + "component") + if not part.checkpointing_or_offloading_enabled() and not compile: return None diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 9d350cc8f..80372d10d 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -270,6 +270,7 @@ class TrainModelPartConfig(BaseConfig): gradient_checkpointing: bool offload_fraction: float activation_offloading: bool + cache_in_ram: bool def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) @@ -310,6 +311,7 @@ def default_values(): data.append(("gradient_checkpointing", True, bool, False)) data.append(("offload_fraction", 0.0, float, False)) data.append(("activation_offloading", False, bool, False)) + data.append(("cache_in_ram", True, bool, False)) return TrainModelPartConfig(data) @@ -400,6 +402,7 @@ class TrainConfig(BaseConfig): async_offloading: bool force_circular_padding: bool compile: bool + stream_from_disk: bool # data settings concept_file_name: str @@ -889,6 +892,9 @@ def weight_dtypes(self) -> ModelWeightDtypes: self.embedding_weight_dtype, ) + def cache_in_ram(self) -> dict[str, bool]: + return {part: getattr(self, part).cache_in_ram for part in self.model_type.model_parts()} + def part_trained_in_place(self, part: TrainModelPartConfig) -> bool: # True iff a FINE_TUNE run updates this part's base weights. 'train' defaults True even for parts the # architecture can't train (e.g. a frozen text encoder), so also require the model type to list the part as @@ -1051,6 +1057,7 @@ def default_values() -> 'TrainConfig': data.append(("async_offloading", True, bool, False)) data.append(("force_circular_padding", False, bool, False)) data.append(("compile", False, bool, False)) + data.append(("stream_from_disk", True, bool, False)) # data settings data.append(("concept_file_name", "training_concepts/concepts.json", str, False)) diff --git a/modules/util/disk_stream.py b/modules/util/disk_stream.py new file mode 100644 index 000000000..0e825d2f8 --- /dev/null +++ b/modules/util/disk_stream.py @@ -0,0 +1,109 @@ +from collections.abc import Callable + +from modules.module.quantized.mixin.QuantizedLinearMixin import QuantizedLinearMixin +from modules.util.enum.DataType import DataType +from modules.util.torch_util import torch_gc + +import torch +from torch import nn + +# A streamed sub-module keeps its base weights frozen (LoRA training streams too -- only the adapter trains, so the +# streamed base weights never diverge from disk; a fully fine-tuned part cannot stream, its in-place updates would be +# discarded). It is loaded as a meta skeleton and its real weights are streamed straight from the checkpoint to the +# compute device and quantized the first time it is used -- so the full unquantized module never lands in system RAM. +# Both load paths share +# this materialize step and differ only in how they evict the weights off the compute device afterwards, selected by +# cache_in_ram: +# - cache_in_ram off: discard the weights to meta; re-materialize by re-streaming from the checkpoint. Frees both +# VRAM and RAM. Lossless because the module is frozen -- its weights never diverge from disk. +# - cache_in_ram on: keep the streamed+quantized weights resident on the temp device; re-materialize by moving them +# back to the compute device. Frees VRAM only, but avoids re-reading the checkpoint on every use. + + +def _is_evicted(module: nn.Module) -> bool: + # the skeleton is fully on meta between uses; a single real parameter means it is currently materialized + for parameter in module.parameters(): + return parameter.is_meta + return True + + +def _current_device(module: nn.Module) -> torch.device: + for parameter in module.parameters(): + return parameter.device + for buffer in module.buffers(): + return buffer.device + return torch.device("meta") + + +def evict_to_meta(module: nn.Module): + for sub_module in module.modules(): + for name, parameter in list(sub_module.named_parameters(recurse=False)): + if parameter.is_meta: + continue + if name == "weight" and isinstance(sub_module, QuantizedLinearMixin): + # a quantized weight is stored in a packed layout (nf4 packs to a flat [N, 1] tensor); reset it to a + # meta tensor of the original unpacked shape so the next materialize can stream the checkpoint weight + # back into it and re-quantize. Its dtype is irrelevant (the stream overwrites it), so keep the current. + sub_module.register_parameter(name, nn.Parameter( + torch.empty(sub_module.original_weight_shape(), dtype=parameter.dtype, device="meta"), + requires_grad=False)) + else: + sub_module.register_parameter( + name, nn.Parameter(parameter.detach().to("meta"), requires_grad=False)) + for name, buffer in list(sub_module._buffers.items()): + # non-persistent buffers (e.g. rotary inv_freq) are config-derived constants, not disk weights; + # keep them resident rather than evict and re-derive them. + if name in sub_module._non_persistent_buffers_set: + continue + if buffer is not None and not buffer.is_meta: + sub_module._buffers[name] = buffer.to("meta") + # let the next materialize() re-quantize the freshly streamed weights + if isinstance(sub_module, QuantizedLinearMixin): + sub_module.mark_needs_requantization() + + +def stream_module_to( + module: nn.Module, + device: torch.device, + materialize_fn: Callable[[nn.Module, torch.device, DataType], None], + train_dtype: DataType, + cache_in_ram: bool, + name: str, + temp_device: torch.device, +): + # module.to()-style entry point for a materialize-on-demand component; see the module-level comment for the + # materialize/evict semantics. Idempotent; train_dtype is used only when materializing. + if device.type not in ("meta", temp_device.type): + # target is the compute device -> materialize the module onto it + current = _current_device(module) + try: + if current.type == "meta": + # cold: stream+quantize the weights from the checkpoint onto the compute device + materialize_fn(module, device, train_dtype, part_name=name) + elif current.type == temp_device.type: + # warm (cache_in_ram): the quantized weights are staged resident on the temp device, move them back to + # the compute device. Dispatch on device *type* (not equality) so a module already on the compute + # device isn't dragged through module.to(), which would raise on the non-persistent buffers left on meta. + module.to(device=device) + except Exception: + # a materialize that fails partway (typically OOM) leaves already-streamed weights resident on the compute + # device -- live model state torch_gc can't reclaim, which can cascade into a second OOM. Roll back along the + # inverse of the failed move: a meta origin re-streams next time (drop the partial fill back to meta), a cpu + # origin keeps its RAM copy (move back to the temp device). + if current.type == "meta": + evict_to_meta(module) + # reclaim the partial fill now: this rollback runs under BaseModel.materialize, which (unlike + # evict) has no trailing torch_gc, so the stranded VRAM would otherwise survive into the re-raise. + torch_gc() + else: + module.to(device=current) + raise + elif not cache_in_ram: + if not _is_evicted(module): + evict_to_meta(module) + else: + # cache_in_ram: stage the resident quantized weights on the temp device. Only when currently on the compute + # device -- a module still on meta (never materialized) has nothing resident to stage and .to() can't move meta, + # so it stays a no-op here and streams from the checkpoint on its first materialize. + if _current_device(module).type not in (device.type, "meta"): + module.to(device=device) diff --git a/modules/util/quantization_util.py b/modules/util/quantization_util.py index 3570f670f..2524ef3d2 100644 --- a/modules/util/quantization_util.py +++ b/modules/util/quantization_util.py @@ -1,7 +1,6 @@ from collections.abc import Callable from functools import partial -import modules.util.multi_gpu_util as multi from modules.module.quantized.mixin.QuantizedLinearMixin import QuantizedLinearMixin from modules.module.quantized.mixin.QuantizedModuleMixin import QuantizedModuleMixin from modules.util.config.TrainConfig import QuantizationConfig, TrainConfig @@ -260,16 +259,20 @@ def is_quantized_parameter( return False +def is_quantized_module(module: nn.Module) -> bool: + return any(is_quantized_parameter(module, name) + for name, _ in module.named_parameters(recurse=False)) + + def quantize_layers(module: nn.Module, device: torch.device, train_dtype: DataType, config: TrainConfig): if module is None: return child_modules = list(module.modules()) - for _ in multi.master_first(): #avoid cache writing conflicts - for child_module in tqdm(child_modules, desc="Quantizing model weights", total=len(child_modules), delay=5, smoothing=0.1): - if isinstance(child_module, (QuantizedModuleMixin, GGUFLinear)): - child_module.compute_dtype = train_dtype.torch_dtype() - if isinstance(child_module, QuantizedModuleMixin): - child_module.quantize(device=device) + for child_module in tqdm(child_modules, desc="Quantizing model weights", total=len(child_modules), delay=5, smoothing=0.1): + if isinstance(child_module, (QuantizedModuleMixin, GGUFLinear)): + child_module.compute_dtype = train_dtype.torch_dtype() + if isinstance(child_module, QuantizedModuleMixin): + child_module.quantize(device=device) def get_unquantized_weight(module: nn.Linear, dtype: torch.dtype, device: torch.device) -> Tensor: assert isinstance(module, nn.Linear) @@ -305,6 +308,9 @@ def get_offload_tensors(module: nn.Module) -> list[torch.Tensor]: def get_offload_tensor_bytes(module: nn.Module) -> int: + if isinstance(module, QuantizedLinearMixin) and module.weight.is_meta: + return module.predict_offload_bytes() + tensors = get_offload_tensors(module) return sum(t.element_size() * t.numel() for t in tensors) diff --git a/modules/util/torch_util.py b/modules/util/torch_util.py index 4d199f282..2d2a9c1c6 100644 --- a/modules/util/torch_util.py +++ b/modules/util/torch_util.py @@ -1,4 +1,6 @@ +import contextlib import gc +import time from collections.abc import Callable from contextlib import nullcontext from typing import Any @@ -14,6 +16,22 @@ torch_version = packaging.version.parse(torch.__version__) +@contextlib.contextmanager +def timed(label: str, enabled: bool = True): + # wall-clock timing around a block; sync the compute device before and after so the measurement includes the + # async device transfer + (re)quantization rather than just the launch overhead. Forces a cuda sync per block, + # so enable only for ad-hoc profiling, not on the hot per-step path. + if not enabled: + yield + return + if torch.cuda.is_available(): + torch.cuda.synchronize() + start = time.perf_counter() + yield + if torch.cuda.is_available(): + torch.cuda.synchronize() + print(f"[timing] {label}: {time.perf_counter() - start:.3f}s") + def supports_mem_pool(device: torch.device) -> bool: return device.type == "cuda" From bb4fbb6b86833cacfd11aa91e791a77848bef640 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 24 Jul 2026 17:27:23 +0200 Subject: [PATCH 09/10] Skip dequantize for unquantized W8A8 weights Co-Authored-By: Claude Opus 4.8 (1M context) --- modules/module/quantized/LinearW8A8.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modules/module/quantized/LinearW8A8.py b/modules/module/quantized/LinearW8A8.py index a8d0432dd..d42e05f66 100644 --- a/modules/module/quantized/LinearW8A8.py +++ b/modules/module/quantized/LinearW8A8.py @@ -107,6 +107,8 @@ def predict_offload_bytes(self) -> int: return weight_bytes + bias_bytes def unquantized_weight(self, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if not self.is_quantized: + return self.weight.detach().to(dtype) return dequantize(self.weight.detach(), self.scale).to(dtype) @torch.no_grad() From 417cb5d68f2ae22588b459b32867438edd3d4245 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:54:44 +0200 Subject: [PATCH 10/10] Join stream-reader threads on abort, not just on success stream_module_from_checkpoint only joined its reader threads after the drain loop completed cleanly. A mid-stream failure (place() OOM, a reader error) unwound past the join, leaving up to STREAM_READER_THREADS daemon threads still executing inside safe_open().get_tensor(). When the run then tore down, force-killing a reader mid-mmap-read segfaulted on Windows (0xC0000005 inside torch_cpu, called from _safetensors_rust). Wrap the drain in try/finally: a stop Event lets readers break their stripe early, and the finally drains remaining items so a reader blocked on a full queue can post its done sentinel and exit before the unconditional join. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../modelLoader/mixin/HFModelLoaderMixin.py | 39 +++++++++++++------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/modules/modelLoader/mixin/HFModelLoaderMixin.py b/modules/modelLoader/mixin/HFModelLoaderMixin.py index 64b15eecd..4feffb638 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -47,12 +47,16 @@ def __stream_reader( source_key_map: dict[str, str] | None, out_queue: queue.Queue, done, + stop: threading.Event, ): # prefetch reader thread: reads a stripe of the work list into host RAM and feeds the bounded queue. Each thread - # owns its safe_open handles (a handle is not safe for concurrent get_tensor). + # owns its safe_open handles (a handle is not safe for concurrent get_tensor). stop lets the main thread break the + # stripe early on abort/OOM, so no reader is left executing inside safetensors when the stream unwinds. thread_handles: dict[str, object] = {} try: for i in range(tid, len(work), nthreads): + if stop.is_set(): + break item = work[i] path = key_to_file[item[0]] handle = thread_handles.get(path) @@ -188,11 +192,12 @@ def place(item, value): nthreads = STREAM_READER_THREADS out_queue: queue.Queue = queue.Queue(maxsize=2 * nthreads) done = object() + stop = threading.Event() threads = [ threading.Thread( target=__stream_reader, - args=(tid, nthreads, work, key_to_file, source_key_map, out_queue, done), + args=(tid, nthreads, work, key_to_file, source_key_map, out_queue, done, stop), name=f"stream-reader-{tid}", daemon=True, ) for tid in range(nthreads) @@ -200,16 +205,26 @@ def place(item, value): for t in threads: t.start() finished = 0 - while finished < nthreads: - got = out_queue.get() - if got is done: - finished += 1 - elif isinstance(got, Exception): - raise got - else: - place(*got) - for t in threads: - t.join() + try: + while finished < nthreads: + got = out_queue.get() + if got is done: + finished += 1 + elif isinstance(got, Exception): + raise got + else: + place(*got) + finally: + # On the happy path this just joins the already-finished readers. On an exception (place() OOM, a reader + # error) it signals the readers to stop and keeps draining so any reader blocked on a full queue can post its + # done sentinel and exit -- so no daemon reader is ever left executing inside safetensors when the stream + # unwinds, which on Windows would segfault (0xC0000005) when the thread is force-killed at teardown. + stop.set() + while finished < nthreads: + if out_queue.get() is done: + finished += 1 + for t in threads: + t.join() # tied weights (e.g. Qwen3 lm_head <-> embed_tokens) are saved once, so the target stays meta; fill it with an # independent clone of the source (not an alias -- in-place quantize would corrupt both), then quantize. Both keys