theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-24 10:18:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
calmhawk
2025-03-30 01:13:34 +08:00
parent c7035ad911
commit 52896fa8dd
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
diffsynth/distributed/xdit_context_parallel.py
View File

@@ -124,4 +124,6 @@ def usp_attn_forward(self, x, freqs):
) )
x = x.flatten(2) x = x.flatten(2)
del q, k, v
torch.cuda.empty_cache()
return self.o(x) return self.o(x)

8
diffsynth/pipelines/wan_video.py
View File

@@ -396,13 +396,13 @@ def model_fn_wan_video(
else: else:
tea_cache_update = False tea_cache_update = False
# blocks
if use_unified_sequence_parallel:
if dist.is_initialized() and dist.get_world_size() > 1:
x = torch.chunk(x, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]
if tea_cache_update: if tea_cache_update:
x = tea_cache.update(x) x = tea_cache.update(x)
else: else:
# blocks
if use_unified_sequence_parallel:
if dist.is_initialized() and dist.get_world_size() > 1:
x = torch.chunk(x, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]
for block in dit.blocks: for block in dit.blocks:
x = block(x, context, t_mod, freqs) x = block(x, context, t_mod, freqs)
if tea_cache is not None: if tea_cache is not None: