Skip to content

GitLab

Commit 48b2dc3c authored by Hyunsung Lee's avatar Hyunsung Lee Committed by muyangli
Browse files

Merge pull request #85 from mit-han-lab/ita9naiwa/controlnet-fix

parent 16979769
Hide whitespace changes
Inline Side-by-side
Showing with 10 additions and 9 deletions
nunchaku/models/transformers/transformer_flux.py
View file @ 48b2dc3c
...@@ -82,15 +82,16 @@ class NunchakuFluxTransformerBlocks(nn.Module): ...@@ -82,15 +82,16 @@ class NunchakuFluxTransformerBlocks(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 = ( if len(controlnet_block_samples) > 0:
torch.stack(controlnet_block_samples).to(self.device) if len(controlnet_block_samples) > 0 else None controlnet_block_samples = torch.stack(controlnet_block_samples).to(self.device)
) else:
if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0: controlnet_block_samples = None
controlnet_single_block_samples = (
torch.stack(controlnet_single_block_samples).to(self.device) if controlnet_single_block_samples is not None:
if len(controlnet_single_block_samples) > 0 if len(controlnet_single_block_samples) > 0:
else None controlnet_single_block_samples = torch.stack(controlnet_single_block_samples).to(self.device)
) else:
controlnet_single_block_samples = 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