monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow 3-bit and 6-bit quantization (mlx > v0.21.0)

Browse Source
This commit is contained in:
filipstrand 2025-02-09 20:50:47 +01:00
parent 09d75e46b8
commit e3b2a65aa6
3 changed files with 11 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
README.md
View File

@ -378,7 +378,7 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
### 🗜️ Quantization
MFLUX supports running FLUX in 4-bit or 8-bit quantized mode. Running a quantized version can greatly speed up the
MFLUX supports running FLUX in 3, 4, 6, or 8-bit quantized mode. Running a quantized version can greatly speed up the
generation process and reduce the memory consumption by several gigabytes. [Quantized models also take up less disk space](#-size-comparisons-for-quantized-models).
```sh
@ -404,11 +404,10 @@ running the 8-bit quantized version on this particular machine. Unlike the non-q
The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
| 4 bit | 8 bit | Original (16 bit) |
|--------|---------|-------------------|
| 9.85GB | 18.16GB | 33.73GB |
| 3 bit | 4 bit | 6 bit | 8 bit | Original (16 bit) |
|--------|--------|---------|---------|-------------------|
| 7.52GB | 9.61GB | 13.81GB | 18.01GB | 33.73GB |
The reason weights sizes are not fully cut in half is because a small number of weights are not quantized and kept at full precision.
#### 💾 Saving a quantized version to disk

2
src/mflux/ui/defaults.py
View File

@ -7,4 +7,4 @@ MODEL_INFERENCE_STEPS = {
"dev": 14,
"schnell": 4,
}
QUANTIZE_CHOICES = [4, 8]
QUANTIZE_CHOICES = [3, 4, 6, 8]

13
src/mflux/weights/quantization_util.py
View File

@ -21,12 +21,10 @@ class QuantizationUtil:
if quantize is not None or q_level is not None:
bits = int(q_level) if q_level is not None else quantize
# fmt: off
nn.quantize(vae, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
nn.quantize(transformer, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 64, group_size=64, bits=bits)
nn.quantize(t5_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
nn.quantize(clip_text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear), group_size=64, bits=bits)
# fmt: on
nn.quantize(vae, bits=bits)
nn.quantize(transformer, bits=bits)
nn.quantize(t5_text_encoder, bits=bits)
nn.quantize(clip_text_encoder, bits=bits)
@staticmethod
def quantize_controlnet(
@ -35,6 +33,7 @@ class QuantizationUtil:
transformer_controlnet: nn.Module,
):
q_level = weights.meta_data.quantization_level
if quantize is not None or q_level is not None:
bits = int(q_level) if q_level is not None else quantize
nn.quantize(transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=bits) # fmt: off
nn.quantize(transformer_controlnet, bits=bits)