wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-26 06:10:15 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

wan: Implement block level prefetching

Browse Source
This commit is contained in:
Rattus 2025-11-04 08:34:28 +10:00
parent 0814c1f63f
commit 2e843f309d
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
comfy/ldm/wan/model.py
View File

@ -538,6 +538,8 @@ class WanModel(torch.nn.Module):
List[Tensor]: List[Tensor]:
List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
""" """
prefetch_queue = comfy.ops.make_prefetch_queue(list(self.blocks))
comfy.ops.prefetch_queue_pop(prefetch_queue, x.device, None)
# embeddings # embeddings
x = self.patch_embedding(x.float()).to(x.dtype) x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:] grid_sizes = x.shape[2:]
@ -569,6 +571,7 @@ class WanModel(torch.nn.Module):
patches_replace = transformer_options.get("patches_replace", {}) patches_replace = transformer_options.get("patches_replace", {})
blocks_replace = patches_replace.get("dit", {}) blocks_replace = patches_replace.get("dit", {})
for i, block in enumerate(self.blocks): for i, block in enumerate(self.blocks):
comfy.ops.prefetch_queue_pop(prefetch_queue, x.device, block)
if ("double_block", i) in blocks_replace: if ("double_block", i) in blocks_replace:
def block_wrap(args): def block_wrap(args):
out = {} out = {}
@ -578,6 +581,7 @@ class WanModel(torch.nn.Module):
x = out["img"] x = out["img"]
else: else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options) x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
comfy.ops.prefetch_queue_pop(prefetch_queue, x.device, block)
# head # head
x = self.head(x, e) x = self.head(x, e)