Move the offload and quantization to the base model

extensions_built_in/diffusion_models/krea2/src/mmdit.py

@@ -411,6 +411,14 @@ class SingleStreamBlock(nn.Module):
 
 
 class SingleStreamDiT(nn.Module, OstrisModelMixin):
+    def get_offload_ignore_modules(self):
+        # modulation modules hold tiny live state the offloader must not page
+        return [
+            module
+            for module in self.modules()
+            if isinstance(module, (SimpleModulation, DoubleSharedModulation))
+        ]
+
     def __init__(self, config: SingleMMDiTConfig):
         super().__init__()
         self.config = config