monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix hidden_states in single_transformer_blocks for controlnet transformer

Browse Source
This commit is contained in:
filipstrand 2024-09-21 14:36:09 +02:00
parent a8d4300a86
commit abd7893612
1 changed files with 3 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/mflux/controlnet/transformer_controlnet.py
View File

@ -1,5 +1,3 @@
from typing import Tuple
import mlx.core as mx import mlx.core as mx
from mlx import nn from mlx import nn
@ -77,12 +75,12 @@ class TransformerControlnet(nn.Module):
single_block_samples = () single_block_samples = ()
for block in self.single_transformer_blocks: for block in self.single_transformer_blocks:
ctrlnet_hidden_states = block.forward( hidden_states = block.forward(
hidden_states=ctrlnet_hidden_states, hidden_states=hidden_states,
text_embeddings=text_embeddings, text_embeddings=text_embeddings,
rotary_embeddings=image_rotary_emb rotary_embeddings=image_rotary_emb
) )
single_block_samples = single_block_samples + (ctrlnet_hidden_states[:, encoder_hidden_states.shape[1] :],) single_block_samples = single_block_samples + (hidden_states[:, encoder_hidden_states.shape[1]:],)
controlnet_single_block_samples = () controlnet_single_block_samples = ()
for single_block_sample, controlnet_block in zip(single_block_samples, self.controlnet_single_blocks): for single_block_sample, controlnet_block in zip(single_block_samples, self.controlnet_single_blocks):