Please support Chroma is a 8.9 billion parameter rectified flow transformer capable of generating images from text descriptions. Based on FLUX.1 [schnell] with heavy architectural modifications.

[nitinmukesh]

https://huggingface.co/lodestones/Chroma

[pcprinciple]

yeah, it's a fantastic model.
I‘m also hope the community can support this model. However, the model is still in training, so it may be a good idea to wait until the model is officially released before quantizing it.

[josephrocca]

it may be a good idea to wait until the model is officially released before quantizing it

Note that this issue is about supporting the model within the Nunchaku engine (and DeepCompressor), rather than the quantization itself. There's no need to wait for training to finish - Chroma is already usable. Once Nunchaku/DeepCompressor supports it every version of Chroma will be quantized by someone in the community. SVDQuant is more expensive than the average quantization method, but it's still "pennies" basically, and it's relatively easy to do once the engine supports it.

Also, RE the title of this issue "with heavy architectural modifications" - the modifications aren't too heavy. Just replacing the per-transformer-block guidance layers with a much smaller distilled version, and masking the encoder_hidden_states (leaving one unmasked padding token at the end of the prompt), and removing CLIP - see this comment for more info:

• Support Chroma - Flux based model with architecture changes huggingface/diffusers#11010 (comment)

[josephrocca]

@lmxyy It's possible that Nunchaku and DeepCompressor already "support" Chroma, if it's possible to pass an attention_mask to NunchakuFluxTransformerBlocks. Is that possible? If it is, then Chroma can be converted into a Schnell-compatible form like this:

# docker run --rm -it --gpus all -v $(pwd):/workspace -w /workspace pytorch/pytorch:2.6.0-cuda12.6-cudnn9-devel bash

# Create Schnell-compatible variant of Chroma by downloading both Chroma and Schnell safetensor files, and copying Chroma's matching weights over to Schnell. This works because lodestone *distilled* the guidance layers instead of completely pruning them, so we can actually just use Schnell's guidance stuff. This comes at the cost of bloating the model back to Schnell's original size (which is fine for now).
CHROMA_VERSION="15"
apt-get update && apt-get install aria2 -y && pip3 install safetensors
cd /workspace
aria2c -x 16 -s 16 -o chroma.safetensors "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v${CHROMA_VERSION}.safetensors?download=true"
aria2c -x 16 -s 16 -o flux1-schnell.safetensors "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors?download=true"
python3 -c '
from safetensors import safe_open
from safetensors.torch import save_file
with safe_open("/workspace/chroma.safetensors", framework="pt", device="cpu") as chroma, safe_open("/workspace/flux1-schnell.safetensors", framework="pt", device="cpu") as schnell:
    chroma_tensors = {key: chroma.get_tensor(key) for key in chroma.keys()}
    schnell_tensors = {key: schnell.get_tensor(key) for key in schnell.keys()}
matching_keys = set(chroma_tensors).intersection(schnell_tensors)
for key in matching_keys:
    schnell_tensors[key] = chroma_tensors[key]
save_file(schnell_tensors, "/workspace/chroma-schnell-compat.safetensors")
'

So then we only need to:

1. Mask the padding tokens of the T5 prompt (except the last padding token), construct the attention_mask (for the joint txt+img transformer blocks only) like this
2. Feed in torch.zeros for the CLIP embeddings (Chroma doesn't need/use the CLIP embeddings)

[lmxyy]

@lmxyy It's possible that Nunchaku and DeepCompressor already "support" Chroma, if it's possible to pass an attention_mask to NunchakuFluxTransformerBlocks. Is that possible? If it is, then Chroma can be converted into a Schnell-compatible form like this:

docker run --rm -it --gpus all -v $(pwd):/workspace -w /workspace pytorch/pytorch:2.6.0-cuda12.6-cudnn9-devel bash

Create Schnell-compatible variant of Chroma by downloading both Chroma and Schnell safetensor files, and copying Chroma's matching weights over to Schnell. This works because lodestone distilled the guidance layers instead of completely pruning them, so we can actually just use Schnell's guidance stuff. This comes at the cost of bloating the model back to Schnell's original size (which is fine for now).

CHROMA_VERSION="15"
apt-get update && apt-get install aria2 -y && pip3 install safetensors
cd /workspace
aria2c -x 16 -s 16 -o chroma.safetensors "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v${CHROMA_VERSION}.safetensors?download=true"
aria2c -x 16 -s 16 -o flux1-schnell.safetensors "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors?download=true"
python3 -c '
from safetensors import safe_open
from safetensors.torch import save_file
with safe_open("/workspace/chroma.safetensors", framework="pt", device="cpu") as chroma, safe_open("/workspace/flux1-schnell.safetensors", framework="pt", device="cpu") as schnell:
chroma_tensors = {key: chroma.get_tensor(key) for key in chroma.keys()}
schnell_tensors = {key: schnell.get_tensor(key) for key in schnell.keys()}
matching_keys = set(chroma_tensors).intersection(schnell_tensors)
for key in matching_keys:
schnell_tensors[key] = chroma_tensors[key]
save_file(schnell_tensors, "/workspace/chroma-schnell-compat.safetensors")
'
So then we only need to:

1. Mask the padding tokens of the T5 prompt (except the last padding token), construct the attention_mask (for the joint txt+img transformer blocks only) like this
2. Feed in torch.zeros for the CLIP embeddings (Chroma doesn't need/use the CLIP embeddings)

Should be possible

[josephrocca]

I kinda sorta got it working here in a very hacky manner:

• https://gist.github.com/josephrocca/dd339fc20c18dd48524c57b6e4005486

The only weird thing is that when using the quants, I need to set CFG to 1 for the last couple of steps or it burns the image. So I'm probably doing something wrong there. (Edit: See below)

[josephrocca]

Actually, the "set CFG to 1 for the last couple of steps" hack isn't needed anymore. I think it was due to unconditional generation being unstable in earlier versions of Chroma (due to the change in masking, which the model hadn't adapted to yet). SVDQuant of Chroma v27 seems to work well in my initial testing.

Note though that the above script creates a "schnell-compatible" version of Chroma, so it's not taking advantage of smaller 8.9B size. I.e. it has a bunch of useless params in the checkpoint that are wasting flops. But it might be a starting point for official Chroma support.

[lmxyy]

@Aprilhuu Could you take a look? This model seems quite similar to FLUX and should be straightforward to support in both DeepCompressor and Nunchaku.