Move the offload and quantization to the base model

2026-08-28 08:37:56 -06:00
parent 702254688d
commit 85a6880643
45 changed files with 725 additions and 1140 deletions
--- a/toolkit/models/v2/diffusion_models/z_image.py
+++ b/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