Chroma1-HD-DF11

Update: I have uploaded an updated version of this model, that should further reduce disk size and VRAM usage by ~82 MB. This is because I missed out on compressing a small portion of the model (the distilled_guidance_layer.layers) in my original upload. There is <u>no need to download again</u> if you are not having any issues with the older version.

For more information (including how to compress models yourself), check out https://huggingface.co/DFloat11 and https://github.com/LeanModels/DFloat11

This is my first time using DF11 to compress a model outside the Flux architecture. The process for compressing Flux-based models is much more straightforward as compared to other architectures because the compression code requires a pattern_dict as input, but the original example code only provides it for Flux, which meant I had to learn the notation myself and modify it to fit other models. At least Chroma is just a pruned version of Flux, so it was relatively simple to derive the correct pattern_dict this time. Do let me know if you run into any problems.

This is the pattern_dict I used for compression:

pattern_dict = {
    r"distilled_guidance_layer\.layers\.\d+": (
        "linear_1",
        "linear_2"
    ),
    r"transformer_blocks\.\d+": (
        "attn.to_q",
        "attn.to_k",
        "attn.to_v",
        "attn.add_k_proj",
        "attn.add_v_proj",
        "attn.add_q_proj",
        "attn.to_out.0",
        "attn.to_add_out",
        "ff.net.0.proj",
        "ff.net.2",
        "ff_context.net.0.proj",
        "ff_context.net.2",
    ),
    r"single_transformer_blocks\.\d+": (
        "proj_mlp",
        "proj_out",
        "attn.to_q",
        "attn.to_k",
        "attn.to_v",
    ),
}

How to Use

diffusers

1. Install the DFloat11 pip package (installs the CUDA kernel automatically; requires a CUDA-compatible GPU and PyTorch installed):

   pip install dfloat11[cuda12]
   # or if you have CUDA version 11:
   # pip install dfloat11[cuda11]
   

2. To use the DFloat11 model, run the following example code in Python:

   import torch
   from diffusers import ChromaPipeline, ChromaTransformer2DModel
   from dfloat11 import DFloat11Model
   from transformers.modeling_utils import no_init_weights
   with no_init_weights():
       transformer = ChromaTransformer2DModel.from_config(
           ChromaTransformer2DModel.load_config(
               "lodestones/Chroma1-HD",
               subfolder="transformer"
           ),
           torch_dtype=torch.bfloat16
       ).to(torch.bfloat16)
      
   pipe = ChromaPipeline.from_pretrained(
       "lodestones/Chroma1-HD",
       transformer=transformer,
       torch_dtype=torch.bfloat16
   )
   DFloat11Model.from_pretrained("mingyi456/Chroma1-HD-DF11", device='cpu', bfloat16_model=pipe.transformer)
   pipe.enable_model_cpu_offload()
   prompt = "A high-fashion close-up portrait of a blonde woman in clear sunglasses. The image uses a bold teal and red color split for dramatic lighting. The background is a simple teal-green. The photo is sharp and well-composed, and is designed for viewing with anaglyph 3D glasses for optimal effect. It looks professionally done."
   negative_prompt = "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors"
   image = pipe(
       prompt,
       negative_prompt=negative_prompt,
       generator=torch.Generator("cpu").manual_seed(0)
   ).images[0]
   image.save("Chroma1-HD.png")
   

ComfyUI

Refer to this model instead.