Skip to content

GitLab

Unverified Commit 2eb3fe1e authored by Hyunsung Lee's avatar Hyunsung Lee Committed by GitHub
Browse files

fix: fix stack error to enable Controlnet-Pro-2 working (#349)

parent 57147358
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 7 deletions
nunchaku/caching/utils.py
View file @ 2eb3fe1e
...@@ -396,16 +396,13 @@ class FluxCachedTransformerBlocks(nn.Module): ...@@ -396,16 +396,13 @@ class FluxCachedTransformerBlocks(nn.Module):
image_rotary_emb = image_rotary_emb.to(self.device) image_rotary_emb = image_rotary_emb.to(self.device)
if controlnet_block_samples is not None: if controlnet_block_samples is not None:
controlnet_block_samples = ( controlnet_block_samples = (
torch.stack(controlnet_block_samples).to(self.device) if len(controlnet_block_samples) > 0 else None torch.stack(controlnet_block_samples).to(self.device) if len(controlnet_block_samples) > 0 else None
) )
if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0: if controlnet_single_block_samples:
controlnet_single_block_samples = ( controlnet_single_block_samples = torch.stack(controlnet_single_block_samples).to(self.device)
torch.stack(controlnet_single_block_samples).to(self.device) else:
if len(controlnet_single_block_samples) > 0 controlnet_single_block_samples = None
else None
)
assert image_rotary_emb.ndim == 6 assert image_rotary_emb.ndim == 6
assert image_rotary_emb.shape[0] == 1 assert image_rotary_emb.shape[0] == 1
......
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