diff --git a/README.md b/README.md index b2c85dd..19f783e 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py index 7ed4bd7..31a307f 100644 --- a/src/mflux/ui/defaults.py +++ b/src/mflux/ui/defaults.py @@ -7,4 +7,4 @@ MODEL_INFERENCE_STEPS = { "dev": 14, "schnell": 4, } -QUANTIZE_CHOICES = [4, 8] +QUANTIZE_CHOICES = [3, 4, 6, 8] diff --git a/src/mflux/weights/quantization_util.py b/src/mflux/weights/quantization_util.py index 3592f69..0afa685 100644 --- a/src/mflux/weights/quantization_util.py +++ b/src/mflux/weights/quantization_util.py @@ -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)