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..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.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/ErnieBaseDataLoader.py b/modules/dataLoader/ErnieBaseDataLoader.py index 032180bc3..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.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/Flux2BaseDataLoader.py b/modules/dataLoader/Flux2BaseDataLoader.py index a6bc3a05d..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.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/FluxBaseDataLoader.py b/modules/dataLoader/FluxBaseDataLoader.py index d23dee3a8..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.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()) 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..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.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()) 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..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.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/IdeogramBaseDataLoader.py b/modules/dataLoader/IdeogramBaseDataLoader.py index 28de23e40..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.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/Krea2BaseDataLoader.py b/modules/dataLoader/Krea2BaseDataLoader.py index 5d243f637..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.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/PixArtAlphaBaseDataLoader.py b/modules/dataLoader/PixArtAlphaBaseDataLoader.py index bb110dc2e..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.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()) 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..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.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/SanaBaseDataLoader.py b/modules/dataLoader/SanaBaseDataLoader.py index a44ff8130..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.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, 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..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.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()) 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..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.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()) 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..a0eda2913 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_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 12739adf8..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.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, 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..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.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/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..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,17 +63,10 @@ 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 - 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 +120,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..bb2614c46 100644 --- a/modules/model/BaseModel.py +++ b/modules/model/BaseModel.py @@ -1,4 +1,5 @@ -from abc import ABCMeta, abstractmethod +from abc import ABCMeta +from collections.abc import Callable from contextlib import nullcontext from uuid import uuid4 @@ -6,11 +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, device_equals, mem_pool_context, supports_mem_pool, torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -78,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, @@ -85,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 @@ -96,17 +106,116 @@ def __init__( self.autocast_context = nullcontext() self.train_dtype = DataType.FLOAT_32 - @abstractmethod - def to(self, device: torch.device): - pass + self._mem_pools = {} + + @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): + # 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 = 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: + 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/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): + 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 + for module in to_move: + module.to(device=device) + self._mem_pools.pop(stem, None) - @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..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 = [] @@ -95,12 +90,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 +153,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..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,47 +53,10 @@ 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 - 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..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,17 +62,10 @@ 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 - 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 +116,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..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 = [] @@ -116,13 +111,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 +194,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..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 = [] @@ -167,15 +160,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 +210,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..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 = [] @@ -131,13 +126,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 +218,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..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,19 +56,10 @@ 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 - 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 +78,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..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,17 +71,10 @@ 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 - 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 +117,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..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 = [] @@ -97,12 +92,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 +147,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..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,19 +61,11 @@ 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 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 +76,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..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 = [] @@ -99,12 +94,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 +112,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..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 = [] @@ -134,14 +129,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 +218,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..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,17 +65,10 @@ 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 - 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 +87,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/modelLoader/AnimaModelLoader.py b/modules/modelLoader/AnimaModelLoader.py index f4d3662bc..edc38eade 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 @@ -38,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") @@ -55,83 +56,56 @@ 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_transformers_sub_module( + 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, ) - 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", - ) - - 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, - ) + model.vae = self._load_vae( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - 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 + 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, + ) def load( #TODO share code between models self, @@ -140,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: @@ -154,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 af268c25c..0c9e3b85b 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 @@ -37,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") @@ -55,68 +53,44 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - 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, - ) + 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_transformers_sub_module( + 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", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKLFlux2, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKLFlux2, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) - - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer + model.vae = self._load_vae( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) def __load_safetensors( self, @@ -140,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: @@ -155,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 33f3fe518..0f27c3cbc 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, @@ -42,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") @@ -59,82 +58,52 @@ def __load_diffusers( transformer_model_name: str, vae_model_name: str, quantization: QuantizationConfig, + stream_from_disk: bool, ): - 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, - ) + 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_transformers_sub_module( - 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_transformers_sub_module( - 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, - subfolder="scheduler", + "text_encoder", + stream_from_disk=stream_from_disk, ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKLFlux2, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKLFlux2, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) + model.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + base_model_name, + subfolder="scheduler", + ) - 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.vae = self._load_vae( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) def __load_safetensors( self, @@ -156,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: @@ -169,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 aef400f08..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,49 +87,34 @@ def __load_diffusers( quantization, ) else: - unconditional_transformer = None + model.unconditional_transformer = None - text_encoder = self._load_transformers_sub_module( + 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", ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKLFlux2, - weight_dtypes.vae, - weight_dtypes.train_dtype, - vae_model_name, - ) - else: - vae = self._load_diffusers_sub_module( - AutoencoderKLFlux2, - weight_dtypes.vae, - weight_dtypes.train_dtype, - base_model_name, - "vae", - ) - - model.model_type = model_type - model.tokenizer = tokenizer - model.noise_scheduler = noise_scheduler - model.text_encoder = text_encoder - model.vae = vae - model.transformer = transformer - model.unconditional_transformer = unconditional_transformer + model.vae = self._load_vae( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) def __load_safetensors( self, @@ -134,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: @@ -149,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 308232823..63df43787 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 ( @@ -40,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") @@ -57,67 +56,43 @@ 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_transformers_sub_module( + 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, ) - 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", - ) - - 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, - ) + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + 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.transformer, model.materialize_fn["transformer"] = self._load_transformer( + ZImageTransformer2DModel, + weight_dtypes, + base_model_name, + transformer_model_name, + quantization, + stream_from_disk=stream_from_disk, + ) def __load_safetensors( self, @@ -139,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: @@ -153,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 7dcbef794..59994818d 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 @@ -35,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") @@ -52,68 +51,44 @@ 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_transformers_sub_module( + 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, ) - 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", - ) - - 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, - ) + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - 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 + 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, + ) def __load_safetensors( self, @@ -135,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: @@ -149,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 d4f21ea2a..0a2cf9546 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 @@ -38,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") @@ -58,96 +58,71 @@ 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_transformers_sub_module( + 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_transformers_sub_module( + 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 - - 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", - ) + model.text_encoder_2 = None - 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, - ) + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - 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 + 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, + ) def __load_safetensors( self, @@ -233,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: @@ -248,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 b3e20f23c..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,122 +68,119 @@ 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_transformers_sub_module( + 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_transformers_sub_module( + 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_transformers_sub_module( + 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: - text_encoder_4 = self._load_transformers_sub_module( - LlamaForCausalLM, - weight_dtypes.text_encoder_4, - weight_dtypes.train_dtype, - 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. + # _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_transformers_sub_module( + 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 - 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", - ) + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + 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, @@ -268,6 +267,7 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] @@ -277,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 @@ -289,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 85c91699b..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,96 +59,70 @@ 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_transformers_sub_module( + 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_transformers_sub_module( + 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 - 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", - ) - - 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, - ) + model.vae = self._load_vae( + AutoencoderKLHunyuanVideo, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - 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 + 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, + ) def __load_safetensors( self, @@ -233,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 @@ -249,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 c3987e97c..b4789e5d6 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 @@ -34,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") @@ -51,69 +50,44 @@ 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_transformers_sub_module( + 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, ) - 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", - ) - - 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, - ) + model.vae = self._load_vae( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + 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.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, + ) def __load_safetensors( self, @@ -134,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: @@ -147,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 f2f196257..4feffb638 100644 --- a/modules/modelLoader/mixin/HFModelLoaderMixin.py +++ b/modules/modelLoader/mixin/HFModelLoaderMixin.py @@ -1,36 +1,266 @@ 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 +from diffusers import GGUFQuantizationConfig from transformers.conversion_mapping import get_checkpoint_conversion_mapping from transformers.core_model_loading import rename_source_key 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, + 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). 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) + 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() + stop = threading.Event() + + threads = [ + threading.Thread( + target=__stream_reader, + 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) + ] + for t in threads: + t.start() + finished = 0 + 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 + # 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, @@ -189,6 +419,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, @@ -197,6 +428,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", @@ -213,19 +445,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, @@ -234,6 +557,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", @@ -250,19 +574,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, @@ -328,3 +700,117 @@ 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, + 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. + # 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: + 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, + ) + transformer = self._convert_diffusers_sub_module_to_dtype( + transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization, + ) + 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, + weight_dtypes.train_dtype, + 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, + module_type, + dtype: DataType, + 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. 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, + 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..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,58 +43,45 @@ 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_transformers_sub_module( + 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, ) - 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", - ) + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + 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, @@ -100,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 953f15bfb..21a4e76f5 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 @@ -34,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") @@ -51,69 +50,44 @@ 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_transformers_sub_module( + 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, ) - 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", - ) - - 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, - ) + model.vae = self._load_vae( + AutoencoderKLQwenImage, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + 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.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, + ) def __load_safetensors( self, @@ -135,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: @@ -148,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 a904e3996..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,58 +43,45 @@ 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_transformers_sub_module( + 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, ) - 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", - ) + model.vae = self._load_vae( + AutoencoderDC, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + 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, @@ -100,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 aa610f485..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_transformers_sub_module( + model.text_encoder, _ = self._load_text_encoder( CLIPTextModel, weight_dtypes.text_encoder, weight_dtypes.train_dtype, @@ -95,23 +96,15 @@ 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", - ) + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + vae_model_name, + ) - unet = self._load_diffusers_sub_module( + model.unet = self._load_diffusers_sub_module( UNet2DConditionModel, weight_dtypes.unet, weight_dtypes.train_dtype, @@ -120,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: @@ -280,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 47d87da74..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,108 +52,93 @@ 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_transformers_sub_module( + 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_transformers_sub_module( + 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_transformers_sub_module( + 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 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", - ) + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + 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, @@ -251,6 +238,7 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] @@ -258,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: @@ -268,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 afbab6581..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,78 +65,78 @@ 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_transformers_sub_module( + 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_transformers_sub_module( + 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, ) - if vae_model_name: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.fallback_train_dtype, - vae_model_name, + model.vae = self._load_vae( + AutoencoderKL, + weight_dtypes.vae, + weight_dtypes.fallback_train_dtype, + base_model_name, + vae_model_name, + ) + + # 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: - vae = self._load_diffusers_sub_module( - AutoencoderKL, - weight_dtypes.vae, - weight_dtypes.fallback_train_dtype, + model.unet = self._load_diffusers_sub_module( + UNet2DConditionModel, + weight_dtypes.unet, + weight_dtypes.train_dtype, base_model_name, - "vae", + "unet", + quantization, ) - 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 - def __load_ckpt( self, model: StableDiffusionXLModel, @@ -248,6 +250,7 @@ def load( model_names: ModelNames, weight_dtypes: ModelWeightDtypes, quantization: QuantizationConfig, + stream_from_disk: bool = False, ): stacktraces = [] @@ -255,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 188107a2c..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_transformers_sub_module( + 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_transformers_sub_module( + 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_transformers_sub_module( + 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 2c7322816..d5de18ac1 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/ChromaSampler.py b/modules/modelSampler/ChromaSampler.py index bc89e5cab..23ef4aee5 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/ErnieSampler.py b/modules/modelSampler/ErnieSampler.py index 2122e5343..a9cb57e0a 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/Flux2Sampler.py b/modules/modelSampler/Flux2Sampler.py index 0a4cca9b9..7ecbd5c83 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,9 +145,6 @@ def __sample_base( image = image_processor.postprocess(image, output_type='pil') - self.model.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/FluxSampler.py b/modules/modelSampler/FluxSampler.py index fbe532053..fc2b2e0c8 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], @@ -222,7 +213,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 +287,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 +298,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 +325,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 +360,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 +367,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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/HiDreamSampler.py b/modules/modelSampler/HiDreamSampler.py index ee7569ccf..c5b30723b 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/HunyuanVideoSampler.py b/modules/modelSampler/HunyuanVideoSampler.py index c12056c92..10b22bfc9 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,20 +135,14 @@ 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() - 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 0a5dd420b..cb253db0e 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/Krea2Sampler.py b/modules/modelSampler/Krea2Sampler.py index 20aecd5b1..b83205a93 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/PixArtAlphaSampler.py b/modules/modelSampler/PixArtAlphaSampler.py index 9adddf3f2..f8c38f135 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/QwenSampler.py b/modules/modelSampler/QwenSampler.py index 4ca604102..c18eece7e 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/SanaSampler.py b/modules/modelSampler/SanaSampler.py index 6251ab87e..f4089a222 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/StableDiffusion3Sampler.py b/modules/modelSampler/StableDiffusion3Sampler.py index fe159e914..f21f34627 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/StableDiffusionSampler.py b/modules/modelSampler/StableDiffusionSampler.py index 92bfcd759..791290edc 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], @@ -223,7 +213,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 +267,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 +284,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 +312,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 +349,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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/StableDiffusionVaeSampler.py b/modules/modelSampler/StableDiffusionVaeSampler.py index 254c73fdd..787b0cfdd 100644 --- a/modules/modelSampler/StableDiffusionVaeSampler.py +++ b/modules/modelSampler/StableDiffusionVaeSampler.py @@ -63,14 +63,12 @@ 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_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.vae_to(self.temp_device) - 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..93d9f23d3 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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], @@ -247,7 +237,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 +293,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 +315,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 +376,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 +412,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,9 +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.vae_to(self.temp_device) - torch_gc() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], diff --git a/modules/modelSampler/WuerstchenSampler.py b/modules/modelSampler/WuerstchenSampler.py index d679757c2..a1e9d8ba0 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,16 +306,13 @@ 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() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=Image.fromarray(image_array), diff --git a/modules/modelSampler/ZImageSampler.py b/modules/modelSampler/ZImageSampler.py index 0cfcb46d6..0e001df2a 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,18 +114,13 @@ 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() - return ModelSamplerOutput( file_type=FileType.IMAGE, data=image[0], 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..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, @@ -166,9 +148,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..adf449e49 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -16,16 +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.torch_util import torch_gc 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, @@ -48,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, @@ -265,10 +246,7 @@ 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") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c412073e8..379747ba1 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -14,9 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -44,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, @@ -160,7 +142,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..a7509e77c 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -16,9 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -44,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, @@ -182,7 +163,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..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, @@ -326,13 +306,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..71cbe8b24 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -17,9 +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 import torch @@ -48,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, @@ -411,19 +378,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..a2749fef7 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -17,9 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -48,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, @@ -294,13 +274,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..966a4eb2f 100644 --- a/modules/modelSetup/BaseIdeogramSetup.py +++ b/modules/modelSetup/BaseIdeogramSetup.py @@ -14,9 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -44,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, @@ -204,7 +177,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..bacf50741 100644 --- a/modules/modelSetup/BaseKrea2Setup.py +++ b/modules/modelSetup/BaseKrea2Setup.py @@ -14,9 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -46,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, @@ -178,10 +160,7 @@ 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") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseModelSetup.py b/modules/modelSetup/BaseModelSetup.py index 8a6818908..2a4b37f2b 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( @@ -231,12 +236,57 @@ 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: 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, + ): + 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: + model.offload_conductor[attr] = 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 + + # 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) + @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 57bf3a40b..c169cafe3 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -16,9 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -50,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, @@ -326,10 +310,7 @@ 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") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index a618dc28f..01852ae42 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -14,16 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch from torch import Tensor -#TODO share more code with other models class BaseQwenSetup( BaseModelSetup, ModelSetupDiffusionLossMixin, @@ -45,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, @@ -177,10 +159,7 @@ 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") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index a96afd770..b3d2137b0 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -16,9 +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.torch_util import torch_gc +from modules.util.dtype_util import disable_fp16_autocast_context from modules.util.TrainProgress import TrainProgress import torch @@ -51,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, @@ -71,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, @@ -246,10 +236,7 @@ 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") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index 678d727ff..9fd69ccf6 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -17,9 +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 import torch @@ -47,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, @@ -346,16 +324,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..0e10939ef 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -17,9 +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 from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -50,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: @@ -61,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) @@ -334,10 +333,9 @@ 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") + 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..dd7a4fe6f 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -17,9 +17,8 @@ ) 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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -47,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) @@ -59,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, @@ -379,13 +379,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..c6b807127 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -15,12 +15,10 @@ 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, ) from modules.util.quantization_util import quantize_layers -from modules.util.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -53,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) @@ -64,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( @@ -345,10 +346,7 @@ 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") model.eval() - torch_gc() diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index a180f57d2..df0f0912c 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -15,9 +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.torch_util import torch_gc from modules.util.TrainProgress import TrainProgress import torch @@ -46,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, @@ -160,8 +142,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/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..d42e05f66 100644 --- a/modules/module/quantized/LinearW8A8.py +++ b/modules/module/quantized/LinearW8A8.py @@ -80,27 +80,42 @@ 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: + if not self.is_quantized: + return self.weight.detach().to(dtype) 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 +132,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 57dc6baa5..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,9 +141,7 @@ def start(self): self.model_setup.setup_optimizations(self.model, self.config) self.model_setup.setup_train_device(self.model, self.config) self.model_setup.setup_model(self.model, self.config) - self.model.to(self.temp_device) self.model.eval() - torch_gc() self.callbacks.on_update_status("creating the data loader/caching") @@ -254,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.to(self.temp_device) 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/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 d1d24d643..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.to(torch.device(self.initial_train_config.temp_device)) 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 a69094c75..663a49a88 100644 --- a/modules/util/LayerOffloadConductor.py +++ b/modules/util/LayerOffloadConductor.py @@ -1,13 +1,23 @@ 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, 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, device_equals, get_tensor_data, + mem_pool_context, pin_tensor_, replace_tensors_, tensors_match_device, @@ -20,6 +30,8 @@ import torch from torch import nn +from tqdm import tqdm + MESSAGES = [] @@ -34,12 +46,18 @@ 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. +# 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 + + +def align_up(number: int) -> int: + return number + (TENSOR_ALIGNMENT_BYTES - (number % TENSOR_ALIGNMENT_BYTES)) % TENSOR_ALIGNMENT_BYTES -def floor_16(number: int) -> int: - return number - (number % 16) +def align_down(number: int) -> int: + return number - (number % TENSOR_ALIGNMENT_BYTES) class StaticLayerTensorAllocator: @@ -69,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 = ceil_16(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) @@ -91,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 @@ -100,7 +127,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 +180,58 @@ def __init__( self.__tensor_allocators = [] - def allocate_cache(self, layers: list[nn.Module], target_bytes: int): + self.__mem_pool = None + + 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 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 + 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, 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), + # 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 +242,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 +264,24 @@ 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() + + @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: @@ -248,7 +324,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 +360,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) @@ -556,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, @@ -609,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() + + self.__wait_all_layer_transfers() + self.__wait_all_activation_transfers() - elif device_equals(device, self.__train_device): - log("to train device") + 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( @@ -674,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: @@ -686,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: @@ -707,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 @@ -768,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 @@ -788,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], [])) @@ -798,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 1f669e9a5..e975029c0 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -243,13 +243,22 @@ 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: + # 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 - # 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 +307,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 +322,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/config/TrainConfig.py b/modules/util/config/TrainConfig.py index deb632f25..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,18 @@ 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 + # 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, @@ -1042,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/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( 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' 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 408100bf9..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,36 @@ 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" + + +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: