Move the offload and quantization to the base model

toolkit/models/v2/diffusion_models/ltx2.py

@@ -21,6 +21,18 @@ class LTX2VideoTransformer3DModel(
     def get_transformer_block_names(cls):
         return ["transformer_blocks"]
 
+    def get_offload_ignore_modules(self):
+        # fp32 scale/shift tables (bare tensors) must stay resident
+        mods = []
+        for block in self.transformer_blocks:
+            mods += [
+                block.scale_shift_table,
+                block.audio_scale_shift_table,
+                block.video_a2v_cross_attn_scale_shift_table,
+                block.audio_a2v_cross_attn_scale_shift_table,
+            ]
+        return mods + [self.scale_shift_table, self.audio_scale_shift_table]
+
 
 class LTX2TextConnectors(DiffusersLTX2TextConnectors, OstrisModelMixin):
     aitk_subfolder = "connectors"

toolkit/models/v2/diffusion_models/wan.py

@@ -80,6 +80,12 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
     def get_transformer_block_names(cls):
         return ["blocks"]
 
+    def get_offload_ignore_modules(self):
+        # fp32 modulation tables must stay resident on the compute device
+        return [self.scale_shift_table] + [
+            block.scale_shift_table for block in self.blocks
+        ]
+
     @classmethod
     def convert_state_dict_on_load(cls, state_dict):
         # original/comfy wan keys -> diffusers layout via diffusers' own

toolkit/models/v2/diffusion_models/z_image.py

@@ -26,6 +26,9 @@ class ZImageTransformer2DModel(DiffusersZImageTransformer2DModel, OstrisModelMix
     def get_transformer_block_names(cls):
         return ["layers"]
 
+    def get_offload_ignore_modules(self):
+        return [self.x_pad_token, self.cap_pad_token]
+
     @classmethod
     def get_quantization_exclude_modules(cls):
         # sensitive modules kept in full precision (fnmatch patterns on module