Tensor size mismatch error

[AndreiZoltan]

If one run main.py he will get the following error:

hidden_states.shape=torch.Size([2, 2304, 3072])
encoder_hidden_states.shape=torch.Size([2, 512, 3072])
image_rotary_emb.shape=torch.Size([2, 1, 2816, 64, 2, 2])
  0%|                                                                | 0/28 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/root/inpainting/FLUX-Controlnet-Inpainting/main.py", line 36, in <module>
    result = pipe(
             ^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/root/inpainting/FLUX-Controlnet-Inpainting/pipeline_flux_controlnet_inpaint.py", line 959, in __call__
    ) = self.controlnet(
        ^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/inpainting/FLUX-Controlnet-Inpainting/controlnet_flux.py", line 329, in forward
    encoder_hidden_states, hidden_states = block(
                                           ^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/inpainting/FLUX-Controlnet-Inpainting/transformer_flux.py", line 214, in forward
    attn_output, context_attn_output = self.attn(
                                       ^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/diffusers/models/attention_processor.py", line 490, in forward
    return self.processor(
           ^^^^^^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/diffusers/models/attention_processor.py", line 1762, in __call__
    query = apply_rotary_emb(query, image_rotary_emb)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/flux/lib/python3.11/site-packages/diffusers/models/embeddings.py", line 680, in apply_rotary_emb
    out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
           ~~~~~~~~~~^~~~~
RuntimeError: The size of tensor a (128) must match the size of tensor b (2) at non-singleton dimension 6

[everythoughthelps]

have you solved this problem?

[JPlin]

Same problem as #7