Add offloading support for krea2

extensions_built_in/diffusion_models/krea2/krea2.py

@@ -39,8 +39,14 @@ from toolkit.samplers.custom_flowmatch_sampler import (
 from toolkit.accelerator import unwrap_model
 from toolkit.metadata import get_meta_for_safetensors
 from toolkit.util.quantize import quantize, get_qtype, quantize_model
+from toolkit.memory_management import MemoryManager
 
-from .src.mmdit import SingleStreamDiT, SingleMMDiTConfig
+from .src.mmdit import (
+    DoubleSharedModulation,
+    SimpleModulation,
+    SingleMMDiTConfig,
+    SingleStreamDiT,
+)
 from .src.text_encoder import encode_krea_prompt, SELECT_LAYERS
 from .src.pipeline import Krea2Pipeline, pad_text_features, predict_velocity
 
@@ -234,6 +240,21 @@ class Krea2Model(BaseModel):
             quantize_model(self, transformer)
             flush()
 
+        if (
+            self.model_config.layer_offloading
+            and self.model_config.layer_offloading_transformer_percent > 0
+        ):
+            MemoryManager.attach(
+                transformer,
+                self.device_torch,
+                offload_percent=self.model_config.layer_offloading_transformer_percent,
+                ignore_modules=[
+                    module
+                    for module in transformer.modules()
+                    if isinstance(module, (SimpleModulation, DoubleSharedModulation))
+                ],
+            )
+
         if self.model_config.low_vram:
             self.print_and_status_update("Moving transformer to CPU")
             transformer.to("cpu")
@@ -248,6 +269,16 @@ class Krea2Model(BaseModel):
             quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
             freeze(text_encoder)
             flush()
+        if (
+            self.model_config.layer_offloading
+            and self.model_config.layer_offloading_text_encoder_percent > 0
+        ):
+            MemoryManager.attach(
+                text_encoder,
+                self.device_torch,
+                offload_percent=self.model_config.layer_offloading_text_encoder_percent,
+            )
+
         if self.model_config.low_vram:
             self.print_and_status_update("Moving text encoder to CPU")
             text_encoder.to("cpu")