[diffusion] feat: support nunchaku for Z-Image-Turbo and flux.1 (int4)

[mickqian]

Motivation

Modifications

Performance and Accuracy Tests

Flux.1-dev

original:

nunchaku:

Z-Image-Turbo

original:

nunchaku:

Model	Resolution	Original (s)	Nunchaku (s)	Speedup
FLUX.1-dev	256×256	61.79	8.21	~7.5×
Z-Image-Turbo	512×512	5.84	2.06	~2.8×

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ping1jing2]

According to DRY principle, i think i might be better to change code below:

def get_param_names_mapping():
    block_rules = {
        # MLP related
        r"mlp_fc1\.(.*)$": r"ff.net.0.proj.\1",
        r"mlp_fc2\.(.*)$": r"ff.net.2.\1",
        r"mlp_context_fc1\.(.*)$": r"ff_context.net.0.proj.\1",
        r"mlp_context_fc2\.(.*)$": r"ff_context.net.2.\1",

        # Attention related
        r"qkv_proj\.(.*)$": r"attn.to_qkv.\1",
        r"qkv_proj_context\.(.*)$": r"attn.to_added_qkv.\1",
        r"attn\.add_qkv_proj\.(.*)$": r"attn.to_added_qkv.\1",
        r"out_proj\.(.*)$": r"attn.to_out.0.\1",
        r"out_proj_context\.(.*)$": r"attn.to_add_out.\1",

        # Norm related
        r"norm_q\.(.*)$": r"attn.norm_q.\1",
        r"norm_k\.(.*)$": r"attn.norm_k.\1",
        r"norm_added_q\.(.*)$": r"attn.norm_added_q.\1",
        r"norm_added_k\.(.*)$": r"attn.norm_added_k.\1",
    }

    mapping = {
        # HF diffusers format
        r"^transformer\.(\w*)\.(.*)$": r"\1.\2",
    }

    prefixes = ["transformer_blocks", "single_transformer_blocks"]
    for prefix in prefixes:
        for pattern, replacement in block_rules.items():
            mapping[rf"^{prefix}\.(\d+)\.{pattern}"] = rf"{prefix}.\1.{replacement}"

    # nunchaku quantization parameter name conversions (apply to all blocks)
    suffix_rules = {
        r"smooth_orig$": r"smooth_factor_orig",
        r"smooth$": r"smooth_factor",
        r"lora_down$": r"proj_down",
        r"lora_up$": r"proj_up",
    }
    for old, new in suffix_rules.items():
        mapping[rf"^(.*)\.{old}"] = rf"\1.{new}"

    return mapping

# nunchaku checkpoint uses different weight names; map to sglang flux layout
param_names_mapping: dict = field(default_factory=get_param_names_mapping)

[ping1jing2]

NunchakuFeedForward = None if nunchaku is not avaiable, but i saw we have self.feed_forward = NunchakuFeedForward(ff, **nunchaku_kwargs) below, so maybe we can change code like this?

_nunchaku_enabled = is_nunchaku_available()
if _nunchaku_enabled:
    from nunchaku.models.attention import NunchakuFeedForward  # type: ignore[import]

[mickqian]

/tag-and-rerun-ci