Streaming model loader

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)

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)

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)

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)

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())

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())

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)

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)

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)

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())

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)

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())

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())

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())

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)

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())

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,

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)

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():

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()