Move the offload and quantization to the base model

extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py

@@ -43,7 +43,7 @@ import torch
 import yaml
 
 from transformers import AutoTokenizer, Qwen3VLTextModel
-from optimum.quanto import freeze
+
 
 from toolkit.accelerator import unwrap_model
 from toolkit.basic import flush
@@ -56,7 +56,7 @@ from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
 )
 from toolkit.train_tools import apply_noise_offset
-from toolkit.util.quantize import quantize, get_qtype, quantize_model
+
 
 from .src.transformer_prx import PRXTransformer2DModel
 from .src.pipeline import PRXPixelPipeline
@@ -127,36 +127,22 @@ class PRXPixelT2IModel(BaseModel):
         self.print_and_status_update("Loading transformer")
         # from_pretrained reads config.json (bottleneck_size, resolution_embeds,
         # in_channels=3, ...) and the safetensors in one shot.
-        transformer = PRXTransformer2DModel.load_model(model_path, dtype=dtype)
-        transformer.to(dtype=dtype)
-        flush()
-
-        if self.model_config.quantize:
-            self.print_and_status_update("Quantizing transformer")
-            quantize_model(self, transformer)
-            flush()
-
-        if self.model_config.low_vram:
-            transformer.to("cpu")
-        else:
-            transformer.to(self.device_torch, dtype=dtype)
+        # load + quantize + offload + placement, all driven by model_config
+        transformer = PRXTransformer2DModel.load(
+            model_path, **self.component_load_kwargs("transformer")
+        )
         flush()
 
         # --- text encoder + tokenizer (Qwen3-VL text tower) ---
         self.print_and_status_update("Loading text encoder")
         tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer")
-        text_encoder = Qwen3VLTextOnlyEncoder.load_model(model_path, dtype=dtype)
-        text_encoder.to(self.te_device_torch)
+        text_encoder = Qwen3VLTextOnlyEncoder.load(
+            model_path, **self.component_load_kwargs("te")
+        )
         text_encoder.eval()
         text_encoder.requires_grad_(False)
         flush()
 
-        if self.model_config.quantize_te:
-            self.print_and_status_update("Quantizing text encoder")
-            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
-            freeze(text_encoder)
-            flush()
-
         # --- "VAE": identity, since PRXPixel is pixel space ---
         self.print_and_status_update("Preparing pixel-space VAE (identity)")
         vae = FakeVAE(scaling_factor=1.0)