Move the offload and quantization to the base model

extensions_built_in/diffusion_models/mageflow/mageflow.py

@@ -36,7 +36,6 @@ from safetensors.torch import load_file, save_file
 import huggingface_hub
 from transformers import AutoProcessor, AutoTokenizer
 from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
-from optimum.quanto import freeze
 
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.models.base_model import BaseModel
@@ -47,8 +46,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.transformer import MageFlow, MageFlowParams
 from .src.vae import MageVAE
@@ -262,50 +259,13 @@ class MageFlowModel(BaseModel):
 
         transformer = self._load_transformer()
 
-        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,
-            )
-
-        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, 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)
+        # quantize + offload + placement, all driven by model_config
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
         flush()
 
         vae = self._load_vae()

extensions_built_in/diffusion_models/mageflow/src/transformer.py

@@ -684,6 +684,26 @@ class MageFlowParams:
 
 
 class MageFlow(nn.Module, OstrisModelMixin):
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["transformer_blocks"]
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        # sensitive modules kept in full precision (fnmatch patterns):
+        #   img_in / txt_in / txt_norm - input projections
+        #   time_text_embed*           - timestep embedder feeding every
+        #                                block's modulation
+        #   norm_out* / proj_out       - final adaptive norm / output projection
+        return [
+            "img_in",
+            "txt_in",
+            "txt_norm",
+            "time_text_embed*",
+            "norm_out*",
+            "proj_out",
+        ]
+
     def __init__(self, params: MageFlowParams):
         super().__init__()
         self.params = params