Backwards Compatibility Issue: MFlux-ComfyUI plugin, ComfyMFlux branch, Error: AttributeError: 'T5Encoder' object has no attribute 'forward' in prompt_embeds = flux.t5_text_encoder.forward(t5_tokens)

[InformEthics]

Hi,

I am developing and maintaining the ComfyMFlux branch of the MFlux-ComfyUI plugin from @raysers.

https://github.com/raysers/Mflux-ComfyUI/blob/ComfyMFlux/ComfyMflux.py

After updating from MFlux 0.4.1 to 0.5.1 I get the error

!!! Exception during processing !!! 'T5Encoder' object has no attribute 'forward'
Traceback (most recent call last):
File "/Users/igar/ComfyUI/execution.py", line 327, in execute
output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/execution.py", line 202, in get_output_data
return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/execution.py", line 174, in _map_node_over_list
process_inputs(input_dict, i)
File "/Users/igar/ComfyUI/execution.py", line 163, in process_inputs
results.append(getattr(obj, func)(**inputs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/custom_nodes/Mflux-ComfyUI/ComfyMflux.py", line 377, in generate
generated_images = generate_image(
^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/custom_nodes/Mflux-ComfyUI/ComfyMflux.py", line 108, in generate_image
prompt_embeds = flux.t5_text_encoder.forward(t5_tokens)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Caskroom/miniconda/base/lib/python3.12/site-packages/mlx/nn/layers/base.py", line 103, in getattr
super(Module, self).getattribute(key)
AttributeError: 'T5Encoder' object has no attribute 'forward'

After adopting the script by importing T5FeedForward:

from mflux.models.text_encoder.t5_encoder.t5_feed_forward import T5FeedForward

forward = T5FeedForward()
prompt_embeds = forward(t5_tokens)

I get the error: ValueError: [broadcast_shapes] Shapes (4096) and (1,512) cannot be broadcast.

!!! Exception during processing !!! [broadcast_shapes] Shapes (4096) and (1,512) cannot be broadcast.
Traceback (most recent call last):
File "/Users/igar/ComfyUI/execution.py", line 327, in execute
output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/execution.py", line 202, in get_output_data
return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/execution.py", line 174, in _map_node_over_list
process_inputs(input_dict, i)
File "/Users/igar/ComfyUI/execution.py", line 163, in process_inputs
results.append(getattr(obj, func)(**inputs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/custom_nodes/Mflux-ComfyUI/ComfyMflux.py", line 381, in generate
generated_images = generate_image(
^^^^^^^^^^^^^^^
File "/Users/igar/ComfyUI/custom_nodes/Mflux-ComfyUI/ComfyMflux.py", line 111, in generate_image
prompt_embeds = forward(t5_tokens)
^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Caskroom/miniconda/base/lib/python3.12/site-packages/mflux/models/text_encoder/t5_encoder/t5_feed_forward.py", line 17, in call
forwarded_states = self.layer_norm(hidden_states)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Caskroom/miniconda/base/lib/python3.12/site-packages/mflux/models/text_encoder/t5_encoder/t5_layer_norm.py", line 18, in call
return self.weight * hidden_states
~~~~~~~~~~~~^~~~~~~~~~~~~~~
ValueError: [broadcast_shapes] Shapes (4096) and (1,512) cannot be broadcast.

Can you explain me how to adopt our Comfy UI plugin to the new version of MFlux?

Thanks

[filipstrand]

@InformEthics Hi, yes this happens because we switched how we call nn modules from:

        prompt_embeds = self.t5_text_encoder.forward(t5_tokens)

to

        prompt_embeds = self.t5_text_encoder(t5_tokens)

Skipping the .forward part. This new way is more standard, but will have to be updated for newer releases of MFLUX.

[InformEthics]

It perfectly solved the issues, I had to delete all the forward attributes from other modules as well and now it works... thanks a lot for quick help. Feel free to try the updated version of ComfyMFlux that runs with MFlux 0.5.1. !