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): ...@@ -179,7 +179,8 @@ class NunchakuFluxTransformerBlocks(nn.Module):
encoder_hidden_states = encoder_hidden_states.to(original_dtype).to(original_device) encoder_hidden_states = encoder_hidden_states.to(original_dtype).to(original_device)
return encoder_hidden_states, hidden_states return encoder_hidden_states, hidden_states
def __del__(self):
self.m.reset()
## copied from diffusers 0.30.3 ## copied from diffusers 0.30.3
def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
......
...@@ -130,7 +130,8 @@ class NunchakuSanaTransformerBlocks(nn.Module): ...@@ -130,7 +130,8 @@ class NunchakuSanaTransformerBlocks(nn.Module):
.to(original_dtype) .to(original_dtype)
.to(original_device) .to(original_device)
) )
def __del__(self):
self.m.reset()
class NunchakuSanaTransformer2DModel(SanaTransformer2DModel, NunchakuModelLoaderMixin): class NunchakuSanaTransformer2DModel(SanaTransformer2DModel, NunchakuModelLoaderMixin):
@classmethod @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