|
| 1 | +# AutoRound for Diffusion Models (Experimental) |
| 2 | + |
| 3 | +This feature is experimental and may be subject to changes, including potential bug fixes, API modifications, or adjustments to default parameters. |
| 4 | + |
| 5 | +## Quantization |
| 6 | + |
| 7 | +### API Usage (CPU/GPU) Recommended |
| 8 | + |
| 9 | +By default, AutoRoundDiffusion only quantizes the transformer module of diffusion models and uses `COCO2014 captions` for calibration. |
| 10 | + |
| 11 | +```python |
| 12 | +import torch |
| 13 | +from auto_round import AutoRound |
| 14 | +from diffusers import AutoPipelineForText2Image |
| 15 | + |
| 16 | +# Load the model |
| 17 | +model_name = "black-forest-labs/FLUX.1-dev" |
| 18 | +pipe = AutoPipelineForText2Image.from_pretrained(model_name, torch_dtype=torch.bfloat16) |
| 19 | + |
| 20 | +# Quantize the model |
| 21 | +autoround = AutoRound( |
| 22 | + pipe, |
| 23 | + scheme="MXFP8", |
| 24 | + dataset="coco2014", |
| 25 | + num_inference_steps=10, |
| 26 | + guidance_scale=7.5, |
| 27 | + generator_seed=None, |
| 28 | + batch_size=1, |
| 29 | +) |
| 30 | +autoround.quantize() |
| 31 | + |
| 32 | +# Save the quantized model |
| 33 | +output_dir = "./tmp_autoround" |
| 34 | +# Currently loading the quantized diffusion model is not supported, so use fake format |
| 35 | +autoround.save_quantized(output_dir, format="fake", inplace=True) |
| 36 | +``` |
| 37 | + |
| 38 | +- `dataset`: the dataset for quantization training. Currently only support coco2014 and user customized .tsv file. |
| 39 | + |
| 40 | +- `num_inference_steps`: The reference number of denoising steps. |
| 41 | + |
| 42 | +- `guidance_scale`: Control how much the image generation process follows the text prompt. The more it is, the more closely it follows the prompt. |
| 43 | + |
| 44 | +- `generator_seed`: A seed that controls the initial noise from which an image is generated. |
| 45 | + |
| 46 | +for more hyperparameters introduction, please refer [Homepage Detailed Hyperparameters](../../README.md#api-usage-gaudi2cpugpu) |
| 47 | + |
| 48 | +### CLI Usage |
| 49 | + |
| 50 | +A user guide detailing the full list of supported arguments is provided by calling ```auto-round -h``` on the |
| 51 | +terminal. |
| 52 | + |
| 53 | +```bash |
| 54 | +auto-round \ |
| 55 | + --model black-forest-labs/FLUX.1-dev \ |
| 56 | + --scheme MXFP8 \ |
| 57 | + --format fake \ |
| 58 | + --batch_size 1 \ |
| 59 | + --output_dir ./tmp_autoround |
| 60 | +``` |
| 61 | + |
| 62 | +### Diffusion Support Matrix |
| 63 | + |
| 64 | +For diffusion models, currently we only validate quantizaion on the FLUX.1-dev, which involves quantizing the transformer component of the pipeline. |
| 65 | + |
| 66 | +| Model | calibration dataset | |
| 67 | +|--------------|--------------| |
| 68 | +| black-forest-labs/FLUX.1-dev | COCO2014 | |
| 69 | + |
| 70 | + |
| 71 | + |
| 72 | +<details> |
| 73 | +<summary style="font-size:17px;">Calibration Dataset</summary> |
| 74 | + |
| 75 | +For diffusion models, we used [**coco2014**]("https://github.com/mlcommons/inference/raw/refs/heads/master/text_to_image/coco2014/captions/captions_source.tsv") calibration dataset as our default. |
| 76 | + |
| 77 | +If users want to use their own dataset, please build the dataset file in ".tsv" format following below structure and use it through argument --dataset (tsv file): |
| 78 | +``` |
| 79 | +id caption |
| 80 | +0 YOUR_PROMPT |
| 81 | +1 YOUR_PROMPT |
| 82 | +... ... |
| 83 | +``` |
| 84 | +- `id`: The id used to map generated images and prompts. |
| 85 | +- `caption`: The text prompt used to generate the images. |
| 86 | + |
| 87 | + |
| 88 | +</details> |
0 commit comments