Chroma fails to mask attention in the transformer

[AmericanPresidentJimmyCarter]

Describe the bug

Chroma has attention masks on both the transformer and T5. Right now only T5 computes with the proper mask. The attention masks from T5 should carry forward into the transformer and mask padding tokens except for the last padding token. See the reference implementation.

https://github.com/lodestone-rock/flow/blob/master/src/models/chroma/model.py#L222-L257

Reproduction

import torch
from diffusers import ChromaTransformer2DModel
from transformers import T5EncoderModel, T5Tokenizer

from optimum.quanto import freeze, qint8, qfloat8, quantize

from chroma_pipeline import ChromaPipeline

bfl_repo = "black-forest-labs/FLUX.1-schnell"
dtype = torch.bfloat16

transformer = ChromaTransformer2DModel.from_single_file(
    "chroma-unlocked-v37-detail-calibrated.safetensors",
    torch_dtype=dtype,
).to('cuda')
quantize(transformer, weights=qint8, exclude=["proj_out", "x_embedder", "norm_out",
    "context_embedder", "distilled_guidance_layer"])
freeze(transformer)

text_encoder = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype).to('cuda')
quantize(text_encoder, weights=qint8)
freeze(text_encoder)
tokenizer = T5Tokenizer.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype)

pipe = ChromaPipeline.from_pretrained(bfl_repo, transformer=transformer, text_encoder=text_encoder,
    tokenizer=tokenizer, torch_dtype=dtype).to('cuda')
generator = torch.Generator().manual_seed(0)
image = pipe(
    prompt='pikachu playing a violin on mars, sign in the background says, "welcome to mars!!',
    negative_prompt="",
    width=1024,
    height=1024,
    num_inference_steps=35, 
    generator=generator,
    guidance_scale=5,
).images[0]
image.save('chroma/test_chroma0.png')

Logs

System Info

linux 24.04, 3090, diffusers main

Who can help?

@DN6 @Ednaordinary

[AmericanPresidentJimmyCarter]

Refer also to your own code: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/chroma/pipeline_chroma.py#L791-L820

The masks are computed when you prepare the text embeddings and then thrown away and not returned.

[AmericanPresidentJimmyCarter]

Comfy outputs look like this.

[AmericanPresidentJimmyCarter]

The 8bit quant of T5 should not be an issue either, as the model itself is trained with T5 quantized to 8 bits.

[Ednaordinary]

Hm, good catch. I believe this happened when we refactored for batch as I'm getting worse output now than before then but regardless I'll see about fixing it

[AmericanPresidentJimmyCarter]

No worries. Trying to read through the reference code is tricky too. We have def modify_mask_to_attend_padding(mask, max_seq_length, num_extra_padding=8): but then attn_padding: int = 1, in forward seems to overwrite it, so it appears that the only unmasked padding token is the last one (index 511).

We should be getting output similar to the HF space: https://huggingface.co/spaces/gokaygokay/Chroma

[AmericanPresidentJimmyCarter]

Gave it a quick try with unquantized T5 too. The results are different but not better.

[AmericanPresidentJimmyCarter]

Another interesting finding. Using GGUF T5 gives wildly different results than bf16 or qint8 quanto.

bfl_repo = "black-forest-labs/FLUX.1-schnell"
dtype = torch.bfloat16

transformer = ChromaTransformer2DModel.from_single_file(
    "/home/user/Programs/playing-with-diffusers/chroma-unlocked-v37-detail-calibrated.safetensors",
    torch_dtype=dtype,
).to('cuda')
quantize(transformer, weights=qint8, exclude=["proj_out", "x_embedder", "norm_out",
    "context_embedder", "distilled_guidance_layer"])
freeze(transformer)

text_encoder = T5EncoderModel.from_pretrained(
    'city96/t5-v1_1-xxl-encoder-gguf',
    gguf_file='t5-v1_1-xxl-encoder-Q8_0.gguf',
    torch_dtype=dtype,
).to('cuda')
tokenizer = T5Tokenizer.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype)

It looks like it's directly casting T5 to 8-bit in the training code, which I believe can significantly degrade T5 performance: https://github.com/lodestone-rock/flow/blob/master/src/trainer/train_chroma.py#L494-L504

edit: Also tried this

def cast_linear(module, dtype):
    """
    Recursively cast all nn.Linear layers in the model to bfloat16.
    """
    for name, child in module.named_children():
        # If the child module is nn.Linear, cast it to bf16
        if isinstance(child, torch.nn.Linear):
            child.to(dtype)
        else:
            # Recursively apply to child modules
            cast_linear(child, dtype)

bfl_repo = "black-forest-labs/FLUX.1-schnell"
dtype = torch.bfloat16

transformer = ChromaTransformer2DModel.from_single_file(
    "/home/user/Programs/playing-with-diffusers/chroma-unlocked-v37-detail-calibrated.safetensors",
    torch_dtype=dtype,
).to('cuda')
quantize(transformer, weights=qint8, exclude=["proj_out", "x_embedder", "norm_out",
    "context_embedder", "distilled_guidance_layer"])
freeze(transformer)

text_encoder = T5EncoderModel.from_pretrained(
    bfl_repo,
    subfolder="text_encoder_2",
    torch_dtype=dtype,
).to('cuda')
cast_linear(text_encoder, torch.float8_e4m3fn)
tokenizer = T5Tokenizer.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype)

then

        with torch.autocast('cuda', dtype=torch.float32):
            (
                prompt_embeds,
                text_ids,
                negative_prompt_embeds,
                negative_text_ids,
            ) = self.encode_prompt(
                prompt=prompt,
                negative_prompt=negative_prompt,
                prompt_embeds=prompt_embeds,
                negative_prompt_embeds=negative_prompt_embeds,
                do_classifier_free_guidance=self.do_classifier_free_guidance,
                device=device,
                num_images_per_prompt=num_images_per_prompt,
                max_sequence_length=max_sequence_length,
                lora_scale=lora_scale,
            )

This should be close to 1:1 with the training code.

It returns a black image with RuntimeWarning: invalid value encountered in cast images = (images * 255).round().astype("uint8"). I think in order to match the reference implementation's training we need to figure out how fp8 T5 is even being inferenced.

Their T5 implementation itself seems to be modified:

https://github.com/lodestone-rock/flow/blob/master/src/models/chroma/module/t5.py#L188-L205

[Ednaordinary]

@AmericanPresidentJimmyCarter Could you share your comfy workflow? I'm actually unable to reproduce getting a good image of pikachu on mars there either (regardless of negative prompt, etc). With #11725 we should be on par I believe

[AmericanPresidentJimmyCarter]

They were from a reddit user, so I don't know what they actually use. Maybe it's a matter of negative conditioning.

[AmericanPresidentJimmyCarter]

So there's a few other issues I identified with Chroma that maybe warrant separate issues:

1. With higher CFG, about 50% of the time it outputs blurry images. To mitigate this you need to set CFG to 2-3 on the first 8% of timesteps or else the model doesn't work and just yields very blurry images.
2. Chroma is very sensitive to the quantization method of T5 and T5 is also sensitive to quant, and we're still not sure if how we inference it is similar to how it was trained.

The attention mask issue is solved with @DN6 's PR.

[DN6]

With higher CFG, about 50% of the time it outputs blurry images. To mitigate this you need to set CFG to 2-3 on the first 8% of timesteps or else the model doesn't work and just yields very blurry images.

We could handle this via callback?

Chroma is very sensitive to the quantization method of T5 and T5 is also sensitive to quant, and we're still not sure if how we inference it is similar to how it was trained.

This one is tricker since it looks like T5 was modified in the original repo?

[AmericanPresidentJimmyCarter]

We could handle this via callback?

Yes, I think so. Comfy uses CFG scheduler nodes that let you modify per step CFG.

Chroma is very sensitive to the quantization method of T5 and T5 is also sensitive to quant, and we're still not sure if how we inference it is similar to how it was trained.

Yeah -- I tried briefly to get that running but it was non-trivial.