Adapt New Diffusion Model

Adapting AutoRound for a New Diffusion Model Architecture

Overview

AutoRound's new diffusion path uses auto_round/compressors/diffusion_mixin.py, auto_round/calibration/diffusion.py, and the quantizer implementations under auto_round/algorithms/quantization/. This skill covers what code changes are needed when a new diffusion model doesn't work out-of-the-box. Common reasons for adaptation:

• Transformer block type not registered in DIFFUSION_OUTPUT_CONFIGS
• Non-standard pipeline API (not compatible with pipe(prompts, ...))
• Hybrid architecture with both AR and diffusion components
• Model not detected as a diffusion model

Step 0: Diagnose the Problem

from auto_round import AutoRound

ar = AutoRound(
    "your-org/your-diffusion-model",
    scheme="W4A16",
    iters=2,
    nsamples=2,
    num_inference_steps=5,
)
ar.quantize_and_save(output_dir="./test_output", format="fake")

Step 1: Ensure Model Detection

AutoRound detects diffusion models by checking for model_index.json in the model directory:

# auto_round/utils/model.py
def is_diffusion_model(model_or_path):
    # Checks for model_index.json presence

If your model doesn't have model_index.json, either create one in the model directory or pass diffusion-specific options through new-architecture AutoRound kwargs:

ar = AutoRound(
    model,
    num_inference_steps=5,
)

Pipeline Loading

diffusion_load_model() uses AutoPipelineForText2Image.from_pretrained() and extracts pipe.transformer as the quantizable model. If your model uses a different attribute (e.g., pipe.unet), this needs adjustment in auto_round/utils/model.py.

Step 2: Register Transformer Block Output Config

This is the most common adaptation needed. DIFFUSION_OUTPUT_CONFIGS maps transformer block class names to their output tensor names. Without this, calibration crashes because AutoRound doesn't know how to collect activations.

Find your block class name

import diffusers

pipe = diffusers.AutoPipelineForText2Image.from_pretrained("your-model")
for name, module in pipe.transformer.named_modules():
    if hasattr(module, "forward") and "block" in name.lower():
        print(f"{name}: {type(module).__name__}")

Register in DIFFUSION_OUTPUT_CONFIGS

Edit auto_round/algorithms/quantization/base.py:

class BaseQuantizers:
    DIFFUSION_OUTPUT_CONFIGS = {
        "FluxTransformerBlock": ["encoder_hidden_states", "hidden_states"],
        "FluxSingleTransformerBlock": ["encoder_hidden_states", "hidden_states"],
        # Add your block type:
        "YourTransformerBlock": ["hidden_states"],  # output tensor names in order
    }

The list must match the exact order of tensors returned by the block's forward() method.

How to determine output tensor names

1. Read the block's forward() method in diffusers source code
2. Identify what tensors it returns (usually hidden_states, sometimes also encoder_hidden_states)
3. List them in the order they're returned

Example: If forward() returns (hidden_states, encoder_hidden_states):

BaseQuantizers.DIFFUSION_OUTPUT_CONFIGS["YourBlock"] = ["hidden_states", "encoder_hidden_states"]

Example: If forward() returns just hidden_states:

BaseQuantizers.DIFFUSION_OUTPUT_CONFIGS["YourBlock"] = ["hidden_states"]

Step 3: Handle Non-Standard Pipeline API

If your model's inference API differs from the standard pipe(prompts, guidance_scale=..., num_inference_steps=...), provide a custom pipeline function.

Option A: Add a custom pipeline dispatch in DiffusionCalibrator

Update auto_round/calibration/diffusion.py so DiffusionCalibrator.calib() dispatches through a small helper instead of calling pipe(...) directly:

class DiffusionCalibrator(LLMCalibrator):
    ...

    def _run_pipeline(self, pipe, prompts, generator):
        if getattr(pipe, "_autoround_pipeline_fn", None) is not None:
            pipe._autoround_pipeline_fn(
                pipe,
                prompts,
                guidance_scale=self.compressor.guidance_scale,
                num_inference_steps=self.compressor.num_inference_steps,
                generator=generator,
            )
            return
        pipe(
            prompts,
            guidance_scale=self.compressor.guidance_scale,
            num_inference_steps=self.compressor.num_inference_steps,
            generator=generator,
        )

Option B: Attach a model-specific function during model loading

For a known model family, attach _autoround_pipeline_fn in auto_round/utils/model.py or auto_round/special_model_handler.py:

pipe._autoround_pipeline_fn = your_model_pipeline_fn

Option C: Add a dedicated branch in DiffusionCalibrator

For full control, update auto_round/calibration/diffusion.py so DiffusionCalibrator.calib() dispatches through your custom pipeline function:

class DiffusionCalibrator(LLMCalibrator):
    ...

    def _run_pipeline(self, pipe, prompts):
        c = self.compressor
        generator = (
            None if c.generator_seed is None else torch.Generator(device=pipe.device).manual_seed(c.generator_seed)
        )
        pipe.your_custom_generate(
            prompts,
            steps=c.num_inference_steps,
            cfg=c.guidance_scale,
            generator=generator,
        )

Step 4: Add Hybrid AR+DiT Support

For models with both autoregressive and diffusion components (e.g., GLM-Image).

4a. Register AR component

Add hybrid routing through the new architecture. Start with auto_round/autoround.py, auto_round/compressors/entry.py, and auto_round/compressors/diffusion_mixin.py. If a reusable AR-component registry is needed, place it near the new routing code:

HYBRID_AR_COMPONENTS = [
    "vision_language_encoder",  # GLM-Image
    "your_ar_component",  # Your model's AR attribute name
]

The attribute name must match what exists on the diffusers pipeline object (i.e., pipe.your_ar_component).