Model conversion with `--quantize-nbits` fails

[SaladDays831]

Branch: main
Conversion command:

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-vae-encoder --model-version "runwayml/stable-diffusion-v1-5" --unet-support-controlnet --quantize-nbits 6 --attention-implementation SPLIT_EINSUM_V2 --convert-controlnet "lllyasviel/control_v11p_sd15_canny" "lllyasviel/control_v11p_sd15_inpaint" "lllyasviel/control_v11p_sd15_openpose" --bundle-resources-for-swift-cli -o "/path/to/save"

The conversion fails after starting the Quantizing text_encoder step.
The conversion using the same command without --quantize-nbits 6 completes successfully.

INFO:__main__:Quantizing text_encoder
Running compression pass palettize_weights:   5%|████▉                                                                                            | 10/197 [00:10<03:09,  1.02s/ ops]
Traceback (most recent call last):
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/Users/danylo/Home/Projects/AI_stuff/apple:ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 1409, in <module>
    main(args)
  File "/Users/danylo/Home/Projects/AI_stuff/apple:ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 1271, in main
    quantize_weights(args)
  File "/Users/danylo/Home/Projects/AI_stuff/apple:ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 148, in quantize_weights
    _quantize_weights(
  File "/Users/danylo/Home/Projects/AI_stuff/apple:ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 184, in _quantize_weights
    model = ct.optimize.coreml.palettize_weights(mlmodel, config=config).save(out_path)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_post_training_quantization.py", line 262, in palettize_weights
    return _apply_graph_pass(mlmodel, weight_palettizer)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_post_training_quantization.py", line 67, in _apply_graph_pass
    graph_pass.apply(prog)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 104, in apply
    apply_block(f)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/converters/mil/mil/passes/helper.py", line 60, in wrapper
    return func(*args)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 101, in apply_block
    self.transform_op(op)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 537, in transform_op
    lut_params = self.compress(
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 505, in compress
    lut, indices = compress_kmeans(val, nbits)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 426, in compress_kmeans
    lut, indices = _get_kmeans_lookup_table_and_weight(nbits, val)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/models/neural_network/quantization_utils.py", line 424, in _get_kmeans_lookup_table_and_weight
    kmeans = KMeans(
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/base.py", line 1151, in wrapper
    return fit_method(estimator, *args, **kwargs)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/cluster/_kmeans.py", line 1471, in fit
    X = self._validate_data(
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/base.py", line 604, in _validate_data
    out = check_array(X, input_name="X", **check_params)
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/utils/validation.py", line 959, in check_array
    _assert_all_finite(
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/utils/validation.py", line 124, in _assert_all_finite
    _assert_all_finite_element_wise(
  File "/Users/danylo/anaconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/utils/validation.py", line 173, in _assert_all_finite_element_wise
    raise ValueError(msg_err)
ValueError: Input X contains infinity or a value too large for dtype('float64').

[atiorh]

Breaking into coremltools/optimize/coreml/_quantization_passes.py:426 after catching the ValueError:

In [1]: val.shape
Out[1]: (1, 1, 77, 77)

In [2]: val
Out[2]: 
array([[[[  0., -inf, -inf, ..., -inf, -inf, -inf],
         [  0.,   0., -inf, ..., -inf, -inf, -inf],
         [  0.,   0.,   0., ..., -inf, -inf, -inf],
         ...,
         [  0.,   0.,   0., ...,   0., -inf, -inf],
         [  0.,   0.,   0., ...,   0.,   0., -inf],
         [  0.,   0.,   0., ...,   0.,   0.,   0.]]]], dtype=float16)

[atiorh]

This seems to be caused by #241 which removed the attention mask override that replace inf with -1e4. So the offending "parameter" is the causal attention mask which is not an actual parameter and shouldn't be palletized anyways. We should exclude this type of tensor from being palletized and restore the -1e4 override.

[atiorh]

I fixed it by partially bringing back the patch that was removed in #241. Thanks again for the report @SaladDays831, please keep them coming!

[SaladDays831]

Thanks for your quick responses! Testing now

[indoflaven]

Hello - I seem to be getting the same or similar error.

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker --model-version stabilityai/sd-turbo -o output --quantize-nbits 6 --attention-implementation SPLIT_EINSUM_V2

Fails with this error:

INFO:__main__:Converted safety_checker
INFO:__main__:Quantizing weights to 6-bit precision
INFO:__main__:Quantizing text_encoder to 6-bit precision
INFO:__main__:Quantizing text_encoder
Running compression pass palettize_weights:   3%|██                                                                            | 10/373 [00:06<03:43,  1.63 ops/s]
Traceback (most recent call last):
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/Users/michaelhein/Documents/GitHub/ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 1524, in <module>
    main(args)
  File "/Users/michaelhein/Documents/GitHub/ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 1369, in main
    quantize_weights(args)
  File "/Users/michaelhein/Documents/GitHub/ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 147, in quantize_weights
    _quantize_weights(
  File "/Users/michaelhein/Documents/GitHub/ml-stable-diffusion/python_coreml_stable_diffusion/torch2coreml.py", line 183, in _quantize_weights
    model = ct.optimize.coreml.palettize_weights(mlmodel, config=config).save(out_path)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_post_training_quantization.py", line 268, in palettize_weights
    return _apply_graph_pass(mlmodel, weight_palettizer)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_post_training_quantization.py", line 72, in _apply_graph_pass
    graph_pass.apply(prog)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 117, in apply
    apply_block(f)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/converters/mil/mil/passes/helper.py", line 60, in wrapper
    return func(*args)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 114, in apply_block
    self.transform_op(op)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 591, in transform_op
    lut_params = self.compress(
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 559, in compress
    lut, indices = palettize_weights._get_lut_and_indices(val, mode, nbits, lut_function)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 516, in _get_lut_and_indices
    lut, indices = compress_kmeans(val, nbits)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/optimize/coreml/_quantization_passes.py", line 476, in compress_kmeans
    lut, indices = _get_kmeans_lookup_table_and_weight(nbits, val)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/coremltools/models/neural_network/quantization_utils.py", line 424, in _get_kmeans_lookup_table_and_weight
    kmeans = KMeans(
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/base.py", line 1152, in wrapper
    return fit_method(estimator, *args, **kwargs)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/cluster/_kmeans.py", line 1475, in fit
    X = self._validate_data(
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/base.py", line 605, in _validate_data
    out = check_array(X, input_name="X", **check_params)
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/utils/validation.py", line 957, in check_array
    _assert_all_finite(
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/utils/validation.py", line 122, in _assert_all_finite
    _assert_all_finite_element_wise(
  File "/opt/miniconda3/envs/coreml_stable_diffusion/lib/python3.8/site-packages/sklearn/utils/validation.py", line 171, in _assert_all_finite_element_wise
    raise ValueError(msg_err)
ValueError: Input X contains infinity or a value too large for dtype('float64').

[kacperd]

@indoflaven have you found a solution for this?

[indoflaven]

@indoflaven have you found a solution for this?

No - Gave up on this, sorry.