Flux: callback to decode intermediate images

[Clement-Lelievre]

In this doc there is a callback snippet allowing to display the intermediate images, for SDXL inferences, as 128*128 images.
This comes in handy eg to follow the denoising process.

For FLUX, I was able to display the intermediate images (in the output image's resolution, so not as 128*128 images) by decoding the latents inside the inference code (specifically the dunder call of the FluxPipeline), at each step. I subclassed FluxPipeline and that did the trick well.

This question is maybe for @asomoza because he is mentioned as the contributor in the doc:

could you please explain the snippet given in the article referenced by the doc above (copied below):

This will hopefully help me adapt this to FLUX.
In particular how you came up with the values used like the weights and biases.

def latents_to_rgb(latents):
    weights = (
        (60, -60, 25, -70),
        (60,  -5, 15, -50),
        (60,  10, -5, -35)
    )

    weights_tensor = torch.t(torch.tensor(weights, dtype=latents.dtype).to(latents.device))
    biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(latents.device)
    rgb_tensor = torch.einsum("...lxy,lr -> ...rxy", latents, weights_tensor) + biases_tensor.unsqueeze(-1).unsqueeze(-1)
    image_array = rgb_tensor.clamp(0, 255)[0].byte().cpu().numpy()
    image_array = image_array.transpose(1, 2, 0)  # Change the order of dimensions

    return Image.fromarray(image_array)

thank you

cc @sayakpaul @TimothyAlexisVass

[sayakpaul]

Cc: @asomoza

[asomoza]

Probably better to be answered by @TimothyAlexisVass, but the explanation is in the same part that you linked:

0: Luminance
1: Cyan/Red => equivalent to rgb(0, 255, 255)/rgb(255, 0, 0)
2: Lime/Medium Purple => equivalent to rgb(127, 255, 0)/rgb(127, 0, 255)
3: Pattern/structure.

If each value can range between -4 and 4 at the point of decoding, then in a 16-bit floating point format with half precision, each latent pixel can contain 16,384 distinct values for each of the 4 channels.

I haven't really looked at it, but Flux uses a 16ch VAE right? It won't be as easy to do the same trick with it.

The Flux generations are so slow though, that using the teaf1 won't make it noticeable slower.

[Clement-Lelievre]

yes from my experiments, flux latents have 16 channels

[TimothyAlexisVass]

In particular how you came up with the values used like the weights and biases.

Here are the experiments I did

I started out with trying to use the conversion formula for YCbCr to RGB which I then tried to adapt to the unique color space of SDXL. I then continued with a different approach, trying to get the best result until I realized that you can't get a perfect result and then just settled on some "good enough" numbers.

[TimothyAlexisVass]

I should definitely also mention that @madebyollin sent me some links in an email related to others having made similar experiments in the past. I will not link his name here without his consent, but you can reach out to him if you want.
I'm supposing you already know who he is.

[asomoza]

Actually, I remember that I saw that comfyui had a file for some models, I found that it has the ones for Flux, so probably you can use those values but you'll need to adapt the preview function for them.

This is probably your answer.