Skip to content

Commit 9912f57

Browse files
authored
[Major] Support deleting the model
Fix GPU memory allocation issues
2 parents 889efaf + 3235dca commit 9912f57

File tree

2 files changed

+4
-2
lines changed

2 files changed

+4
-2
lines changed

nunchaku/models/transformers/transformer_flux.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,8 @@ def forward_layer_at(
179179
encoder_hidden_states = encoder_hidden_states.to(original_dtype).to(original_device)
180180

181181
return encoder_hidden_states, hidden_states
182-
182+
def __del__(self):
183+
self.m.reset()
183184

184185
## copied from diffusers 0.30.3
185186
def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:

nunchaku/models/transformers/transformer_sana.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,8 @@ def forward_layer_at(
130130
.to(original_dtype)
131131
.to(original_device)
132132
)
133-
133+
def __del__(self):
134+
self.m.reset()
134135

135136
class NunchakuSanaTransformer2DModel(SanaTransformer2DModel, NunchakuModelLoaderMixin):
136137
@classmethod

0 commit comments

Comments
 (0)