Lots of rendering artifacts, doesn't follow prompts well, error messages

[Adreitz]

I am using a 64GB M3 Max MBP on 15.0.1. Just got MLX Flux Dev working in comfyui this morning, but it's not working well. While it is somewhat faster than regular PyTorch, it uses much more memory (swaps a lot of memory even at 1024x1024, especially during VAE decoding), produces a lot of poor-quality images, doesn't follow prompts well if they contain more than one concept (e.g. "a cat"), and gives error messages if the cfg control (is this meant to be real cfg or flux guidance?) is anything other than 0.

This is an image generated using the rather complex prompt, "Drone photograph of a small flat rock island in the middle of the ocean with three neatly-dressed people standing on it, looking at a small fishing boat. The boat is motoring away from the island and the pilot waves farewell. On the horizon to the right, the sun is rising. On the horizon to the left, massive dark storm clouds are gathering." As you can see, it understood "fishing boat", "island", and "ocean" and not much more, which is uncharacteristic for Flux. The image also has crosshatch artifacts throughout. I've also seen dotted vertical lines and repeated noise patterns in other images.

If I take this same workflow and set cfg to 1, I get the following error message:

[concatenate] All the input array dimensions must match exactly except for the concatenation axis. However, the provided shapes are (1,24,512,128), (2,24,4032,128), and the concatenation axis is 2.

Any further attempts at MLX rendering (whether just reattempting the above with cfg 1, another cfg value, or even setting cfg back to 0) result in this other error message:

[addmm] Last dimension of first input with shape (1,1,1,3072) must match second to last dimension of second input with shape (0,1).

The only solution I've found to regain rendering is to restart the back end.

So, unfortunately, this doesn't look ready for more than experimental use at this point. I'll be happy to retry when these issues start to be fixed.

[thoddnn]

Thanks for your feedback ! I will check those issues.

Feel free to contribute as it's a new open-source project.

[Adreitz]

Thanks. I am unfortunately not a programmer and am not really in a position where I'm able to learn right now. But if you have anything you'd like me to test, feel free to send it.

[Adreitz]

I retested your MLX nodes on OS 15.1.1, since I saw you had done an update some weeks ago. I haven't done extensive testing, but I retested the above prompt and seed and things are looking much better, more how I'd expect Flux to look. Thanks for the fix!

Setting cfg to anything other than 0 still produces the error messages as I indicated previously, including the behavior that the backend remains broken after returning to cfg 0 until it is restarted.

Since you have a latent decoder node, would it be possible to have a latent encoder as well? This would give more flexibility for image to image workflows. Or is the latent format the same so that the normal Pytorch VAE Encode node would still work?