Unverified Commit 9912f57e authored by Muyang Li's avatar Muyang Li Committed by GitHub
Browse files

[Major] Support deleting the model

Fix GPU memory allocation issues
parents 889efafd 3235dcac
......@@ -179,7 +179,8 @@ class NunchakuFluxTransformerBlocks(nn.Module):
encoder_hidden_states = encoder_hidden_states.to(original_dtype).to(original_device)
return encoder_hidden_states, hidden_states
def __del__(self):
self.m.reset()
## copied from diffusers 0.30.3
def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
......
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment