add reset on delete to sana

3235dcac · Edna · GitHub · 3be44cbf · 3235dcac
Unverified Commit 3235dcac authored Apr 13, 2025 by Edna Committed by GitHub Apr 13, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

nunchaku/models/transformers/transformer_sana.py nunchaku/models/transformers/transformer_sana.py +2 -1

No files found.
--- a/nunchaku/models/transformers/transformer_sana.py
+++ b/nunchaku/models/transformers/transformer_sana.py
@@ -130,7 +130,8 @@ class NunchakuSanaTransformerBlocks(nn.Module):
            .to(original_dtype)
            .to(original_device)
        )
-
+    def __del__(self):
+        self.m.reset()

 class NunchakuSanaTransformer2DModel(SanaTransformer2DModel, NunchakuModelLoaderMixin):
    @classmethod