Move the offload and quantization to the base model

extensions_built_in/diffusion_models/example_model/example_model.py

@@ -32,7 +32,7 @@ from toolkit.models.base_model import BaseModel
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
+from toolkit.util.quantize import quantize, get_qtype
 
 from .src.model import ExampleTransformer2DModel
 from .src.pipeline import ExamplePipeline, pad_prompt_embeds
@@ -165,21 +165,14 @@ class ExampleModel(BaseModel):
         del state_dict
         flush()  # gc + empty cuda cache; call it after dropping anything big
 
-        if self.model_config.quantize:
-            # quantize_model handles qtype selection, exclusions and device
-            # juggling, and leaves the model on CPU
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            flush()
-
-        if self.model_config.low_vram:
-            # leave it on CPU; get_noise_prediction moves it over when needed
-            transformer.to("cpu")
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
+        # quantize + offload + placement, all driven by model_config:
+        # component_load_kwargs derives qtype (incl. an accuracy recovery
+        # adapter), the layer-offload fraction and the target device
+        # (low_vram parks on CPU); aitk_post_load applies them. Models whose
+        # checkpoint sourcing is standard can collapse the build + this into
+        # one call: ExampleTransformer2DModel.load(path, **kwargs).
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
         flush()
-        # For partial layer offloading support see MemoryManager.attach usage
-        # in ../ideogram4/ideogram4.py or ../z_image/z_image.py.
 
         # --- text encoder + tokenizer (stock transformers model) ---
         self.print_and_status_update("Loading text encoder")

extensions_built_in/diffusion_models/example_model/src/model.py

@@ -51,6 +51,8 @@ import torch.nn.functional as F
 from torch import nn
 from torch.utils.checkpoint import checkpoint
 
+from toolkit.models.v2._mixin import OstrisModelMixin
+
 
 def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
     """Standard sinusoidal embedding.
@@ -135,8 +137,16 @@ class ExampleTransformerBlock(nn.Module):
         return x
 
 
-class ExampleTransformer2DModel(nn.Module):
-    """The denoiser. Plain ``nn.Module`` on purpose.
+class ExampleTransformer2DModel(nn.Module, OstrisModelMixin):
+    """The denoiser. Plain ``nn.Module`` plus ``OstrisModelMixin``.
+
+    The mixin supplies the universal component API every v2 module shares:
+    ``load()`` / ``load_model()`` / ``aitk_post_load()`` (quantize, layer
+    offloading and device placement driven by the holder's model_config via
+    ``BaseModel.component_load_kwargs(role)``), plus comfy-format save/load.
+    Override its class hooks (``get_transformer_block_names``,
+    ``get_quantization_exclude_modules``, ``get_offload_ignore_modules``,
+    ``convert_state_dict_on_load/save``) as needed.
 
     You could also subclass ``diffusers.ModelMixin``/``ConfigMixin`` (see
     ../../ernie_image/transformer.py) to get ``save_pretrained``,
@@ -153,6 +163,12 @@ class ExampleTransformer2DModel(nn.Module):
     names when deciding where to attach LoRA layers.
     """
 
+    @classmethod
+    def get_transformer_block_names(cls):
+        # attribute name(s) of the repeated-block ModuleList(s); the quantizer
+        # streams these blocks through the GPU one at a time
+        return ["blocks"]
+
     def __init__(
         self,
         in_channels: int = 16,    # VAE latent channels
@@ -183,8 +199,8 @@ class ExampleTransformer2DModel(nn.Module):
         )
 
         # ``blocks`` is the repeated-layer ModuleList. The attribute name is
-        # what ExampleModel.get_transformer_block_names() returns, which the
-        # LoRA code uses for block targeting / "transformer only" training.
+        # what get_transformer_block_names() returns, which the LoRA code uses
+        # for block targeting and the quantizer uses for block streaming.
         self.blocks = nn.ModuleList(
             [
                 ExampleTransformerBlock(hidden_size, num_heads)