Move the offload and quantization to the base model

extensions_built_in/diffusion_models/krea2/krea2.py

@@ -47,7 +47,6 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 )
 from toolkit.accelerator import unwrap_model
 from toolkit.metadata import get_meta_for_safetensors
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
 from toolkit.memory_management import MemoryManager
 
 from .src.mmdit import (
@@ -383,55 +382,12 @@ class Krea2Model(QwenImageVAEHolderMixin, BaseModel):
             if self.model_config.qtype == "qfloat8":
                 self.model_config.qtype = "float8"
 
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            flush()
-
-        if (
-            self.model_config.layer_offloading
-            and self.model_config.layer_offloading_transformer_percent > 0
-        ):
-            MemoryManager.attach(
-                transformer,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_transformer_percent,
-                ignore_modules=[
-                    module
-                    for module in transformer.modules()
-                    if isinstance(module, (SimpleModulation, DoubleSharedModulation))
-                ],
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving transformer to CPU")
-            transformer.to("cpu")
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
 
         tokenizer, processor, vl_processor, text_encoder = self._load_text_encoder()
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing text encoder")
-            text_encoder.to(self.device_torch)
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-        if (
-            self.model_config.layer_offloading
-            and self.model_config.layer_offloading_text_encoder_percent > 0
-        ):
-            MemoryManager.attach(
-                text_encoder,
-                self.device_torch,
-                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
-            )
-
-        if self.model_config.low_vram:
-            self.print_and_status_update("Moving text encoder to CPU")
-            text_encoder.to("cpu")
-        else:
-            text_encoder.to(self.device_torch)
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         flush()
 
         vae = self._load_vae()

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -411,6 +411,14 @@ class SingleStreamBlock(nn.Module):
 
 
 class SingleStreamDiT(nn.Module, OstrisModelMixin):
+    def get_offload_ignore_modules(self):
+        # modulation modules hold tiny live state the offloader must not page
+        return [
+            module
+            for module in self.modules()
+            if isinstance(module, (SimpleModulation, DoubleSharedModulation))
+        ]
+
     def __init__(self, config: SingleMMDiTConfig):
         super().__init__()
         self.config = config