Fix TeaCache bug with usp support integration and optimize memory usage by clearing attn cache

2025-03-30 01:13:34 +08:00
parent c7035ad911
commit 52896fa8dd
2 changed files with 6 additions and 4 deletions
--- a/diffsynth/distributed/xdit_context_parallel.py
+++ b/diffsynth/distributed/xdit_context_parallel.py
@@ -124,4 +124,6 @@ def usp_attn_forward(self, x, freqs):
    )
    x = x.flatten(2)

+    del q, k, v
+    torch.cuda.empty_cache()
    return self.o(x)