diffusers 0.30.3__py3-none-any.whl → 0.32.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (268) hide show
  1. diffusers/__init__.py +97 -4
  2. diffusers/callbacks.py +56 -3
  3. diffusers/configuration_utils.py +13 -1
  4. diffusers/image_processor.py +282 -71
  5. diffusers/loaders/__init__.py +24 -3
  6. diffusers/loaders/ip_adapter.py +543 -16
  7. diffusers/loaders/lora_base.py +138 -125
  8. diffusers/loaders/lora_conversion_utils.py +647 -0
  9. diffusers/loaders/lora_pipeline.py +2216 -230
  10. diffusers/loaders/peft.py +380 -0
  11. diffusers/loaders/single_file_model.py +71 -4
  12. diffusers/loaders/single_file_utils.py +597 -10
  13. diffusers/loaders/textual_inversion.py +5 -3
  14. diffusers/loaders/transformer_flux.py +181 -0
  15. diffusers/loaders/transformer_sd3.py +89 -0
  16. diffusers/loaders/unet.py +56 -12
  17. diffusers/models/__init__.py +49 -12
  18. diffusers/models/activations.py +22 -9
  19. diffusers/models/adapter.py +53 -53
  20. diffusers/models/attention.py +98 -13
  21. diffusers/models/attention_flax.py +1 -1
  22. diffusers/models/attention_processor.py +2160 -346
  23. diffusers/models/autoencoders/__init__.py +5 -0
  24. diffusers/models/autoencoders/autoencoder_dc.py +620 -0
  25. diffusers/models/autoencoders/autoencoder_kl.py +73 -12
  26. diffusers/models/autoencoders/autoencoder_kl_allegro.py +1149 -0
  27. diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +213 -105
  28. diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py +1176 -0
  29. diffusers/models/autoencoders/autoencoder_kl_ltx.py +1338 -0
  30. diffusers/models/autoencoders/autoencoder_kl_mochi.py +1166 -0
  31. diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +3 -10
  32. diffusers/models/autoencoders/autoencoder_tiny.py +4 -2
  33. diffusers/models/autoencoders/vae.py +18 -5
  34. diffusers/models/controlnet.py +47 -802
  35. diffusers/models/controlnet_flux.py +70 -0
  36. diffusers/models/controlnet_sd3.py +26 -376
  37. diffusers/models/controlnet_sparsectrl.py +46 -719
  38. diffusers/models/controlnets/__init__.py +23 -0
  39. diffusers/models/controlnets/controlnet.py +872 -0
  40. diffusers/models/{controlnet_flax.py → controlnets/controlnet_flax.py} +5 -5
  41. diffusers/models/controlnets/controlnet_flux.py +536 -0
  42. diffusers/models/{controlnet_hunyuan.py → controlnets/controlnet_hunyuan.py} +7 -7
  43. diffusers/models/controlnets/controlnet_sd3.py +489 -0
  44. diffusers/models/controlnets/controlnet_sparsectrl.py +788 -0
  45. diffusers/models/controlnets/controlnet_union.py +832 -0
  46. diffusers/models/{controlnet_xs.py → controlnets/controlnet_xs.py} +14 -13
  47. diffusers/models/controlnets/multicontrolnet.py +183 -0
  48. diffusers/models/embeddings.py +996 -92
  49. diffusers/models/embeddings_flax.py +23 -9
  50. diffusers/models/model_loading_utils.py +264 -14
  51. diffusers/models/modeling_flax_utils.py +1 -1
  52. diffusers/models/modeling_utils.py +334 -51
  53. diffusers/models/normalization.py +157 -13
  54. diffusers/models/transformers/__init__.py +6 -0
  55. diffusers/models/transformers/auraflow_transformer_2d.py +3 -2
  56. diffusers/models/transformers/cogvideox_transformer_3d.py +69 -13
  57. diffusers/models/transformers/dit_transformer_2d.py +1 -1
  58. diffusers/models/transformers/latte_transformer_3d.py +4 -4
  59. diffusers/models/transformers/pixart_transformer_2d.py +10 -2
  60. diffusers/models/transformers/sana_transformer.py +488 -0
  61. diffusers/models/transformers/stable_audio_transformer.py +1 -1
  62. diffusers/models/transformers/transformer_2d.py +1 -1
  63. diffusers/models/transformers/transformer_allegro.py +422 -0
  64. diffusers/models/transformers/transformer_cogview3plus.py +386 -0
  65. diffusers/models/transformers/transformer_flux.py +189 -51
  66. diffusers/models/transformers/transformer_hunyuan_video.py +789 -0
  67. diffusers/models/transformers/transformer_ltx.py +469 -0
  68. diffusers/models/transformers/transformer_mochi.py +499 -0
  69. diffusers/models/transformers/transformer_sd3.py +112 -18
  70. diffusers/models/transformers/transformer_temporal.py +1 -1
  71. diffusers/models/unets/unet_1d_blocks.py +1 -1
  72. diffusers/models/unets/unet_2d.py +8 -1
  73. diffusers/models/unets/unet_2d_blocks.py +88 -21
  74. diffusers/models/unets/unet_2d_condition.py +9 -9
  75. diffusers/models/unets/unet_3d_blocks.py +9 -7
  76. diffusers/models/unets/unet_motion_model.py +46 -68
  77. diffusers/models/unets/unet_spatio_temporal_condition.py +23 -0
  78. diffusers/models/unets/unet_stable_cascade.py +2 -2
  79. diffusers/models/unets/uvit_2d.py +1 -1
  80. diffusers/models/upsampling.py +14 -6
  81. diffusers/pipelines/__init__.py +69 -6
  82. diffusers/pipelines/allegro/__init__.py +48 -0
  83. diffusers/pipelines/allegro/pipeline_allegro.py +938 -0
  84. diffusers/pipelines/allegro/pipeline_output.py +23 -0
  85. diffusers/pipelines/animatediff/__init__.py +2 -0
  86. diffusers/pipelines/animatediff/pipeline_animatediff.py +45 -21
  87. diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py +52 -22
  88. diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +18 -4
  89. diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py +3 -1
  90. diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +104 -72
  91. diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py +1341 -0
  92. diffusers/pipelines/audioldm2/modeling_audioldm2.py +3 -3
  93. diffusers/pipelines/aura_flow/pipeline_aura_flow.py +2 -9
  94. diffusers/pipelines/auto_pipeline.py +88 -10
  95. diffusers/pipelines/blip_diffusion/modeling_blip2.py +1 -1
  96. diffusers/pipelines/cogvideo/__init__.py +2 -0
  97. diffusers/pipelines/cogvideo/pipeline_cogvideox.py +80 -39
  98. diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py +825 -0
  99. diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py +108 -50
  100. diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py +89 -50
  101. diffusers/pipelines/cogview3/__init__.py +47 -0
  102. diffusers/pipelines/cogview3/pipeline_cogview3plus.py +674 -0
  103. diffusers/pipelines/cogview3/pipeline_output.py +21 -0
  104. diffusers/pipelines/controlnet/__init__.py +86 -80
  105. diffusers/pipelines/controlnet/multicontrolnet.py +7 -178
  106. diffusers/pipelines/controlnet/pipeline_controlnet.py +20 -3
  107. diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +9 -2
  108. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +9 -2
  109. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +37 -15
  110. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +12 -4
  111. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +9 -4
  112. diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py +1790 -0
  113. diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py +1501 -0
  114. diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py +1627 -0
  115. diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py +22 -4
  116. diffusers/pipelines/controlnet_sd3/__init__.py +4 -0
  117. diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +56 -20
  118. diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py +1153 -0
  119. diffusers/pipelines/ddpm/pipeline_ddpm.py +2 -2
  120. diffusers/pipelines/deepfloyd_if/pipeline_output.py +6 -5
  121. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +16 -4
  122. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +1 -1
  123. diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +32 -9
  124. diffusers/pipelines/flux/__init__.py +23 -1
  125. diffusers/pipelines/flux/modeling_flux.py +47 -0
  126. diffusers/pipelines/flux/pipeline_flux.py +256 -48
  127. diffusers/pipelines/flux/pipeline_flux_control.py +889 -0
  128. diffusers/pipelines/flux/pipeline_flux_control_img2img.py +945 -0
  129. diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +1141 -0
  130. diffusers/pipelines/flux/pipeline_flux_controlnet.py +1006 -0
  131. diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py +998 -0
  132. diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +1204 -0
  133. diffusers/pipelines/flux/pipeline_flux_fill.py +969 -0
  134. diffusers/pipelines/flux/pipeline_flux_img2img.py +856 -0
  135. diffusers/pipelines/flux/pipeline_flux_inpaint.py +1022 -0
  136. diffusers/pipelines/flux/pipeline_flux_prior_redux.py +492 -0
  137. diffusers/pipelines/flux/pipeline_output.py +16 -0
  138. diffusers/pipelines/free_noise_utils.py +365 -5
  139. diffusers/pipelines/hunyuan_video/__init__.py +48 -0
  140. diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py +687 -0
  141. diffusers/pipelines/hunyuan_video/pipeline_output.py +20 -0
  142. diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py +20 -4
  143. diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py +9 -9
  144. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +2 -2
  145. diffusers/pipelines/kolors/pipeline_kolors.py +1 -1
  146. diffusers/pipelines/kolors/pipeline_kolors_img2img.py +14 -11
  147. diffusers/pipelines/kolors/text_encoder.py +2 -2
  148. diffusers/pipelines/kolors/tokenizer.py +4 -0
  149. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +1 -1
  150. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +1 -1
  151. diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +1 -1
  152. diffusers/pipelines/latte/pipeline_latte.py +2 -2
  153. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +15 -3
  154. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +15 -3
  155. diffusers/pipelines/ltx/__init__.py +50 -0
  156. diffusers/pipelines/ltx/pipeline_ltx.py +789 -0
  157. diffusers/pipelines/ltx/pipeline_ltx_image2video.py +885 -0
  158. diffusers/pipelines/ltx/pipeline_output.py +20 -0
  159. diffusers/pipelines/lumina/pipeline_lumina.py +3 -10
  160. diffusers/pipelines/mochi/__init__.py +48 -0
  161. diffusers/pipelines/mochi/pipeline_mochi.py +748 -0
  162. diffusers/pipelines/mochi/pipeline_output.py +20 -0
  163. diffusers/pipelines/pag/__init__.py +13 -0
  164. diffusers/pipelines/pag/pag_utils.py +8 -2
  165. diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py +2 -3
  166. diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py +1543 -0
  167. diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py +3 -5
  168. diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py +1683 -0
  169. diffusers/pipelines/pag/pipeline_pag_hunyuandit.py +22 -6
  170. diffusers/pipelines/pag/pipeline_pag_kolors.py +1 -1
  171. diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py +7 -14
  172. diffusers/pipelines/pag/pipeline_pag_sana.py +886 -0
  173. diffusers/pipelines/pag/pipeline_pag_sd.py +18 -6
  174. diffusers/pipelines/pag/pipeline_pag_sd_3.py +18 -9
  175. diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py +1058 -0
  176. diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py +5 -1
  177. diffusers/pipelines/pag/pipeline_pag_sd_img2img.py +1094 -0
  178. diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py +1356 -0
  179. diffusers/pipelines/pag/pipeline_pag_sd_xl.py +18 -6
  180. diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py +31 -16
  181. diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py +42 -19
  182. diffusers/pipelines/pia/pipeline_pia.py +2 -0
  183. diffusers/pipelines/pipeline_flax_utils.py +1 -1
  184. diffusers/pipelines/pipeline_loading_utils.py +250 -31
  185. diffusers/pipelines/pipeline_utils.py +158 -186
  186. diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +7 -14
  187. diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py +7 -14
  188. diffusers/pipelines/sana/__init__.py +47 -0
  189. diffusers/pipelines/sana/pipeline_output.py +21 -0
  190. diffusers/pipelines/sana/pipeline_sana.py +884 -0
  191. diffusers/pipelines/stable_audio/pipeline_stable_audio.py +12 -1
  192. diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +35 -3
  193. diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +2 -2
  194. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +46 -9
  195. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +1 -1
  196. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +1 -1
  197. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +241 -81
  198. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +228 -23
  199. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py +82 -13
  200. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +60 -11
  201. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +11 -1
  202. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +1 -1
  203. diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +16 -4
  204. diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +16 -4
  205. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +16 -12
  206. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +29 -22
  207. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +29 -22
  208. diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +1 -1
  209. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +1 -1
  210. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +16 -4
  211. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +15 -3
  212. diffusers/pipelines/unidiffuser/modeling_uvit.py +2 -2
  213. diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +1 -1
  214. diffusers/quantizers/__init__.py +16 -0
  215. diffusers/quantizers/auto.py +139 -0
  216. diffusers/quantizers/base.py +233 -0
  217. diffusers/quantizers/bitsandbytes/__init__.py +2 -0
  218. diffusers/quantizers/bitsandbytes/bnb_quantizer.py +561 -0
  219. diffusers/quantizers/bitsandbytes/utils.py +306 -0
  220. diffusers/quantizers/gguf/__init__.py +1 -0
  221. diffusers/quantizers/gguf/gguf_quantizer.py +159 -0
  222. diffusers/quantizers/gguf/utils.py +456 -0
  223. diffusers/quantizers/quantization_config.py +669 -0
  224. diffusers/quantizers/torchao/__init__.py +15 -0
  225. diffusers/quantizers/torchao/torchao_quantizer.py +285 -0
  226. diffusers/schedulers/scheduling_ddim.py +4 -1
  227. diffusers/schedulers/scheduling_ddim_cogvideox.py +4 -1
  228. diffusers/schedulers/scheduling_ddim_parallel.py +4 -1
  229. diffusers/schedulers/scheduling_ddpm.py +6 -7
  230. diffusers/schedulers/scheduling_ddpm_parallel.py +6 -7
  231. diffusers/schedulers/scheduling_deis_multistep.py +102 -6
  232. diffusers/schedulers/scheduling_dpmsolver_multistep.py +113 -6
  233. diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +111 -5
  234. diffusers/schedulers/scheduling_dpmsolver_sde.py +125 -10
  235. diffusers/schedulers/scheduling_dpmsolver_singlestep.py +126 -7
  236. diffusers/schedulers/scheduling_edm_euler.py +8 -6
  237. diffusers/schedulers/scheduling_euler_ancestral_discrete.py +4 -1
  238. diffusers/schedulers/scheduling_euler_discrete.py +92 -7
  239. diffusers/schedulers/scheduling_flow_match_euler_discrete.py +153 -6
  240. diffusers/schedulers/scheduling_flow_match_heun_discrete.py +4 -5
  241. diffusers/schedulers/scheduling_heun_discrete.py +114 -8
  242. diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +116 -11
  243. diffusers/schedulers/scheduling_k_dpm_2_discrete.py +110 -8
  244. diffusers/schedulers/scheduling_lcm.py +2 -6
  245. diffusers/schedulers/scheduling_lms_discrete.py +76 -1
  246. diffusers/schedulers/scheduling_repaint.py +1 -1
  247. diffusers/schedulers/scheduling_sasolver.py +102 -6
  248. diffusers/schedulers/scheduling_tcd.py +2 -6
  249. diffusers/schedulers/scheduling_unclip.py +4 -1
  250. diffusers/schedulers/scheduling_unipc_multistep.py +127 -5
  251. diffusers/training_utils.py +63 -19
  252. diffusers/utils/__init__.py +7 -1
  253. diffusers/utils/constants.py +1 -0
  254. diffusers/utils/dummy_pt_objects.py +240 -0
  255. diffusers/utils/dummy_torch_and_transformers_objects.py +435 -0
  256. diffusers/utils/dynamic_modules_utils.py +3 -3
  257. diffusers/utils/hub_utils.py +44 -40
  258. diffusers/utils/import_utils.py +98 -8
  259. diffusers/utils/loading_utils.py +28 -4
  260. diffusers/utils/peft_utils.py +6 -3
  261. diffusers/utils/testing_utils.py +115 -1
  262. diffusers/utils/torch_utils.py +3 -0
  263. {diffusers-0.30.3.dist-info → diffusers-0.32.0.dist-info}/METADATA +73 -72
  264. {diffusers-0.30.3.dist-info → diffusers-0.32.0.dist-info}/RECORD +268 -193
  265. {diffusers-0.30.3.dist-info → diffusers-0.32.0.dist-info}/WHEEL +1 -1
  266. {diffusers-0.30.3.dist-info → diffusers-0.32.0.dist-info}/LICENSE +0 -0
  267. {diffusers-0.30.3.dist-info → diffusers-0.32.0.dist-info}/entry_points.txt +0 -0
  268. {diffusers-0.30.3.dist-info → diffusers-0.32.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,669 @@
1
+ #!/usr/bin/env python
2
+ # coding=utf-8
3
+
4
+ # Copyright 2023 The HuggingFace Inc. team. All rights reserved.
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ """
19
+ Adapted from
20
+ https://github.com/huggingface/transformers/blob/52cb4034ada381fe1ffe8d428a1076e5411a8026/src/transformers/utils/quantization_config.py
21
+ """
22
+
23
+ import copy
24
+ import importlib.metadata
25
+ import inspect
26
+ import json
27
+ import os
28
+ from dataclasses import dataclass
29
+ from enum import Enum
30
+ from functools import partial
31
+ from typing import Any, Dict, List, Optional, Union
32
+
33
+ from packaging import version
34
+
35
+ from ..utils import is_torch_available, is_torchao_available, logging
36
+
37
+
38
+ if is_torch_available():
39
+ import torch
40
+
41
+ logger = logging.get_logger(__name__)
42
+
43
+
44
+ class QuantizationMethod(str, Enum):
45
+ BITS_AND_BYTES = "bitsandbytes"
46
+ GGUF = "gguf"
47
+ TORCHAO = "torchao"
48
+
49
+
50
+ @dataclass
51
+ class QuantizationConfigMixin:
52
+ """
53
+ Mixin class for quantization config
54
+ """
55
+
56
+ quant_method: QuantizationMethod
57
+ _exclude_attributes_at_init = []
58
+
59
+ @classmethod
60
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
61
+ """
62
+ Instantiates a [`QuantizationConfigMixin`] from a Python dictionary of parameters.
63
+
64
+ Args:
65
+ config_dict (`Dict[str, Any]`):
66
+ Dictionary that will be used to instantiate the configuration object.
67
+ return_unused_kwargs (`bool`,*optional*, defaults to `False`):
68
+ Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in
69
+ `PreTrainedModel`.
70
+ kwargs (`Dict[str, Any]`):
71
+ Additional parameters from which to initialize the configuration object.
72
+
73
+ Returns:
74
+ [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters.
75
+ """
76
+
77
+ config = cls(**config_dict)
78
+
79
+ to_remove = []
80
+ for key, value in kwargs.items():
81
+ if hasattr(config, key):
82
+ setattr(config, key, value)
83
+ to_remove.append(key)
84
+ for key in to_remove:
85
+ kwargs.pop(key, None)
86
+
87
+ if return_unused_kwargs:
88
+ return config, kwargs
89
+ else:
90
+ return config
91
+
92
+ def to_json_file(self, json_file_path: Union[str, os.PathLike]):
93
+ """
94
+ Save this instance to a JSON file.
95
+
96
+ Args:
97
+ json_file_path (`str` or `os.PathLike`):
98
+ Path to the JSON file in which this configuration instance's parameters will be saved.
99
+ use_diff (`bool`, *optional*, defaults to `True`):
100
+ If set to `True`, only the difference between the config instance and the default
101
+ `QuantizationConfig()` is serialized to JSON file.
102
+ """
103
+ with open(json_file_path, "w", encoding="utf-8") as writer:
104
+ config_dict = self.to_dict()
105
+ json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
106
+
107
+ writer.write(json_string)
108
+
109
+ def to_dict(self) -> Dict[str, Any]:
110
+ """
111
+ Serializes this instance to a Python dictionary. Returns:
112
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
113
+ """
114
+ return copy.deepcopy(self.__dict__)
115
+
116
+ def __iter__(self):
117
+ """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin"""
118
+ for attr, value in copy.deepcopy(self.__dict__).items():
119
+ yield attr, value
120
+
121
+ def __repr__(self):
122
+ return f"{self.__class__.__name__} {self.to_json_string()}"
123
+
124
+ def to_json_string(self, use_diff: bool = True) -> str:
125
+ """
126
+ Serializes this instance to a JSON string.
127
+
128
+ Args:
129
+ use_diff (`bool`, *optional*, defaults to `True`):
130
+ If set to `True`, only the difference between the config instance and the default `PretrainedConfig()`
131
+ is serialized to JSON string.
132
+
133
+ Returns:
134
+ `str`: String containing all the attributes that make up this configuration instance in JSON format.
135
+ """
136
+ if use_diff is True:
137
+ config_dict = self.to_diff_dict()
138
+ else:
139
+ config_dict = self.to_dict()
140
+ return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
141
+
142
+ def update(self, **kwargs):
143
+ """
144
+ Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes,
145
+ returning all the unused kwargs.
146
+
147
+ Args:
148
+ kwargs (`Dict[str, Any]`):
149
+ Dictionary of attributes to tentatively update this class.
150
+
151
+ Returns:
152
+ `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance.
153
+ """
154
+ to_remove = []
155
+ for key, value in kwargs.items():
156
+ if hasattr(self, key):
157
+ setattr(self, key, value)
158
+ to_remove.append(key)
159
+
160
+ # Remove all the attributes that were updated, without modifying the input dict
161
+ unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove}
162
+ return unused_kwargs
163
+
164
+
165
+ @dataclass
166
+ class BitsAndBytesConfig(QuantizationConfigMixin):
167
+ """
168
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
169
+ loaded using `bitsandbytes`.
170
+
171
+ This replaces `load_in_8bit` or `load_in_4bit`therefore both options are mutually exclusive.
172
+
173
+ Currently only supports `LLM.int8()`, `FP4`, and `NF4` quantization. If more methods are added to `bitsandbytes`,
174
+ then more arguments will be added to this class.
175
+
176
+ Args:
177
+ load_in_8bit (`bool`, *optional*, defaults to `False`):
178
+ This flag is used to enable 8-bit quantization with LLM.int8().
179
+ load_in_4bit (`bool`, *optional*, defaults to `False`):
180
+ This flag is used to enable 4-bit quantization by replacing the Linear layers with FP4/NF4 layers from
181
+ `bitsandbytes`.
182
+ llm_int8_threshold (`float`, *optional*, defaults to 6.0):
183
+ This corresponds to the outlier threshold for outlier detection as described in `LLM.int8() : 8-bit Matrix
184
+ Multiplication for Transformers at Scale` paper: https://arxiv.org/abs/2208.07339 Any hidden states value
185
+ that is above this threshold will be considered an outlier and the operation on those values will be done
186
+ in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], but
187
+ there are some exceptional systematic outliers that are very differently distributed for large models.
188
+ These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of
189
+ magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6,
190
+ but a lower threshold might be needed for more unstable models (small models, fine-tuning).
191
+ llm_int8_skip_modules (`List[str]`, *optional*):
192
+ An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as
193
+ Jukebox that has several heads in different places and not necessarily at the last position. For example
194
+ for `CausalLM` models, the last `lm_head` is typically kept in its original `dtype`.
195
+ llm_int8_enable_fp32_cpu_offload (`bool`, *optional*, defaults to `False`):
196
+ This flag is used for advanced use cases and users that are aware of this feature. If you want to split
197
+ your model in different parts and run some parts in int8 on GPU and some parts in fp32 on CPU, you can use
198
+ this flag. This is useful for offloading large models such as `google/flan-t5-xxl`. Note that the int8
199
+ operations will not be run on CPU.
200
+ llm_int8_has_fp16_weight (`bool`, *optional*, defaults to `False`):
201
+ This flag runs LLM.int8() with 16-bit main weights. This is useful for fine-tuning as the weights do not
202
+ have to be converted back and forth for the backward pass.
203
+ bnb_4bit_compute_dtype (`torch.dtype` or str, *optional*, defaults to `torch.float32`):
204
+ This sets the computational type which might be different than the input type. For example, inputs might be
205
+ fp32, but computation can be set to bf16 for speedups.
206
+ bnb_4bit_quant_type (`str`, *optional*, defaults to `"fp4"`):
207
+ This sets the quantization data type in the bnb.nn.Linear4Bit layers. Options are FP4 and NF4 data types
208
+ which are specified by `fp4` or `nf4`.
209
+ bnb_4bit_use_double_quant (`bool`, *optional*, defaults to `False`):
210
+ This flag is used for nested quantization where the quantization constants from the first quantization are
211
+ quantized again.
212
+ bnb_4bit_quant_storage (`torch.dtype` or str, *optional*, defaults to `torch.uint8`):
213
+ This sets the storage type to pack the quanitzed 4-bit prarams.
214
+ kwargs (`Dict[str, Any]`, *optional*):
215
+ Additional parameters from which to initialize the configuration object.
216
+ """
217
+
218
+ _exclude_attributes_at_init = ["_load_in_4bit", "_load_in_8bit", "quant_method"]
219
+
220
+ def __init__(
221
+ self,
222
+ load_in_8bit=False,
223
+ load_in_4bit=False,
224
+ llm_int8_threshold=6.0,
225
+ llm_int8_skip_modules=None,
226
+ llm_int8_enable_fp32_cpu_offload=False,
227
+ llm_int8_has_fp16_weight=False,
228
+ bnb_4bit_compute_dtype=None,
229
+ bnb_4bit_quant_type="fp4",
230
+ bnb_4bit_use_double_quant=False,
231
+ bnb_4bit_quant_storage=None,
232
+ **kwargs,
233
+ ):
234
+ self.quant_method = QuantizationMethod.BITS_AND_BYTES
235
+
236
+ if load_in_4bit and load_in_8bit:
237
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
238
+
239
+ self._load_in_8bit = load_in_8bit
240
+ self._load_in_4bit = load_in_4bit
241
+ self.llm_int8_threshold = llm_int8_threshold
242
+ self.llm_int8_skip_modules = llm_int8_skip_modules
243
+ self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload
244
+ self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight
245
+ self.bnb_4bit_quant_type = bnb_4bit_quant_type
246
+ self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
247
+
248
+ if bnb_4bit_compute_dtype is None:
249
+ self.bnb_4bit_compute_dtype = torch.float32
250
+ elif isinstance(bnb_4bit_compute_dtype, str):
251
+ self.bnb_4bit_compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
252
+ elif isinstance(bnb_4bit_compute_dtype, torch.dtype):
253
+ self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
254
+ else:
255
+ raise ValueError("bnb_4bit_compute_dtype must be a string or a torch.dtype")
256
+
257
+ if bnb_4bit_quant_storage is None:
258
+ self.bnb_4bit_quant_storage = torch.uint8
259
+ elif isinstance(bnb_4bit_quant_storage, str):
260
+ if bnb_4bit_quant_storage not in ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]:
261
+ raise ValueError(
262
+ "`bnb_4bit_quant_storage` must be a valid string (one of 'float16', 'float32', 'int8', 'uint8', 'float64', 'bfloat16') "
263
+ )
264
+ self.bnb_4bit_quant_storage = getattr(torch, bnb_4bit_quant_storage)
265
+ elif isinstance(bnb_4bit_quant_storage, torch.dtype):
266
+ self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
267
+ else:
268
+ raise ValueError("bnb_4bit_quant_storage must be a string or a torch.dtype")
269
+
270
+ if kwargs and not all(k in self._exclude_attributes_at_init for k in kwargs):
271
+ logger.warning(f"Unused kwargs: {list(kwargs.keys())}. These kwargs are not used in {self.__class__}.")
272
+
273
+ self.post_init()
274
+
275
+ @property
276
+ def load_in_4bit(self):
277
+ return self._load_in_4bit
278
+
279
+ @load_in_4bit.setter
280
+ def load_in_4bit(self, value: bool):
281
+ if not isinstance(value, bool):
282
+ raise TypeError("load_in_4bit must be a boolean")
283
+
284
+ if self.load_in_8bit and value:
285
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
286
+ self._load_in_4bit = value
287
+
288
+ @property
289
+ def load_in_8bit(self):
290
+ return self._load_in_8bit
291
+
292
+ @load_in_8bit.setter
293
+ def load_in_8bit(self, value: bool):
294
+ if not isinstance(value, bool):
295
+ raise TypeError("load_in_8bit must be a boolean")
296
+
297
+ if self.load_in_4bit and value:
298
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
299
+ self._load_in_8bit = value
300
+
301
+ def post_init(self):
302
+ r"""
303
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
304
+ """
305
+ if not isinstance(self.load_in_4bit, bool):
306
+ raise TypeError("load_in_4bit must be a boolean")
307
+
308
+ if not isinstance(self.load_in_8bit, bool):
309
+ raise TypeError("load_in_8bit must be a boolean")
310
+
311
+ if not isinstance(self.llm_int8_threshold, float):
312
+ raise TypeError("llm_int8_threshold must be a float")
313
+
314
+ if self.llm_int8_skip_modules is not None and not isinstance(self.llm_int8_skip_modules, list):
315
+ raise TypeError("llm_int8_skip_modules must be a list of strings")
316
+ if not isinstance(self.llm_int8_enable_fp32_cpu_offload, bool):
317
+ raise TypeError("llm_int8_enable_fp32_cpu_offload must be a boolean")
318
+
319
+ if not isinstance(self.llm_int8_has_fp16_weight, bool):
320
+ raise TypeError("llm_int8_has_fp16_weight must be a boolean")
321
+
322
+ if self.bnb_4bit_compute_dtype is not None and not isinstance(self.bnb_4bit_compute_dtype, torch.dtype):
323
+ raise TypeError("bnb_4bit_compute_dtype must be torch.dtype")
324
+
325
+ if not isinstance(self.bnb_4bit_quant_type, str):
326
+ raise TypeError("bnb_4bit_quant_type must be a string")
327
+
328
+ if not isinstance(self.bnb_4bit_use_double_quant, bool):
329
+ raise TypeError("bnb_4bit_use_double_quant must be a boolean")
330
+
331
+ if self.load_in_4bit and not version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse(
332
+ "0.39.0"
333
+ ):
334
+ raise ValueError(
335
+ "4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version"
336
+ )
337
+
338
+ def is_quantizable(self):
339
+ r"""
340
+ Returns `True` if the model is quantizable, `False` otherwise.
341
+ """
342
+ return self.load_in_8bit or self.load_in_4bit
343
+
344
+ def quantization_method(self):
345
+ r"""
346
+ This method returns the quantization method used for the model. If the model is not quantizable, it returns
347
+ `None`.
348
+ """
349
+ if self.load_in_8bit:
350
+ return "llm_int8"
351
+ elif self.load_in_4bit and self.bnb_4bit_quant_type == "fp4":
352
+ return "fp4"
353
+ elif self.load_in_4bit and self.bnb_4bit_quant_type == "nf4":
354
+ return "nf4"
355
+ else:
356
+ return None
357
+
358
+ def to_dict(self) -> Dict[str, Any]:
359
+ """
360
+ Serializes this instance to a Python dictionary. Returns:
361
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
362
+ """
363
+ output = copy.deepcopy(self.__dict__)
364
+ output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1]
365
+ output["bnb_4bit_quant_storage"] = str(output["bnb_4bit_quant_storage"]).split(".")[1]
366
+ output["load_in_4bit"] = self.load_in_4bit
367
+ output["load_in_8bit"] = self.load_in_8bit
368
+
369
+ return output
370
+
371
+ def __repr__(self):
372
+ config_dict = self.to_dict()
373
+ return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
374
+
375
+ def to_diff_dict(self) -> Dict[str, Any]:
376
+ """
377
+ Removes all attributes from config which correspond to the default config attributes for better readability and
378
+ serializes to a Python dictionary.
379
+
380
+ Returns:
381
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
382
+ """
383
+ config_dict = self.to_dict()
384
+
385
+ # get the default config dict
386
+ default_config_dict = BitsAndBytesConfig().to_dict()
387
+
388
+ serializable_config_dict = {}
389
+
390
+ # only serialize values that differ from the default config
391
+ for key, value in config_dict.items():
392
+ if value != default_config_dict[key]:
393
+ serializable_config_dict[key] = value
394
+
395
+ return serializable_config_dict
396
+
397
+
398
+ @dataclass
399
+ class GGUFQuantizationConfig(QuantizationConfigMixin):
400
+ """This is a config class for GGUF Quantization techniques.
401
+
402
+ Args:
403
+ compute_dtype: (`torch.dtype`, defaults to `torch.float32`):
404
+ This sets the computational type which might be different than the input type. For example, inputs might be
405
+ fp32, but computation can be set to bf16 for speedups.
406
+
407
+ """
408
+
409
+ def __init__(self, compute_dtype: Optional["torch.dtype"] = None):
410
+ self.quant_method = QuantizationMethod.GGUF
411
+ self.compute_dtype = compute_dtype
412
+ self.pre_quantized = True
413
+
414
+ # TODO: (Dhruv) Add this as an init argument when we can support loading unquantized checkpoints.
415
+ self.modules_to_not_convert = None
416
+
417
+ if self.compute_dtype is None:
418
+ self.compute_dtype = torch.float32
419
+
420
+
421
+ @dataclass
422
+ class TorchAoConfig(QuantizationConfigMixin):
423
+ """This is a config class for torchao quantization/sparsity techniques.
424
+
425
+ Args:
426
+ quant_type (`str`):
427
+ The type of quantization we want to use, currently supporting:
428
+ - **Integer quantization:**
429
+ - Full function names: `int4_weight_only`, `int8_dynamic_activation_int4_weight`,
430
+ `int8_weight_only`, `int8_dynamic_activation_int8_weight`
431
+ - Shorthands: `int4wo`, `int4dq`, `int8wo`, `int8dq`
432
+
433
+ - **Floating point 8-bit quantization:**
434
+ - Full function names: `float8_weight_only`, `float8_dynamic_activation_float8_weight`,
435
+ `float8_static_activation_float8_weight`
436
+ - Shorthands: `float8wo`, `float8wo_e5m2`, `float8wo_e4m3`, `float8dq`, `float8dq_e4m3`,
437
+ `float8_e4m3_tensor`, `float8_e4m3_row`,
438
+
439
+ - **Floating point X-bit quantization:**
440
+ - Full function names: `fpx_weight_only`
441
+ - Shorthands: `fpX_eAwB`, where `X` is the number of bits (between `1` to `7`), `A` is the number
442
+ of exponent bits and `B` is the number of mantissa bits. The constraint of `X == A + B + 1` must
443
+ be satisfied for a given shorthand notation.
444
+
445
+ - **Unsigned Integer quantization:**
446
+ - Full function names: `uintx_weight_only`
447
+ - Shorthands: `uint1wo`, `uint2wo`, `uint3wo`, `uint4wo`, `uint5wo`, `uint6wo`, `uint7wo`
448
+ modules_to_not_convert (`List[str]`, *optional*, default to `None`):
449
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have some
450
+ modules left in their original precision.
451
+ kwargs (`Dict[str, Any]`, *optional*):
452
+ The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization
453
+ supports two keyword arguments `group_size` and `inner_k_tiles` currently. More API examples and
454
+ documentation of arguments can be found in
455
+ https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques
456
+
457
+ Example:
458
+ ```python
459
+ from diffusers import FluxTransformer2DModel, TorchAoConfig
460
+
461
+ quantization_config = TorchAoConfig("int8wo")
462
+ transformer = FluxTransformer2DModel.from_pretrained(
463
+ "black-forest-labs/Flux.1-Dev",
464
+ subfolder="transformer",
465
+ quantization_config=quantization_config,
466
+ torch_dtype=torch.bfloat16,
467
+ )
468
+ ```
469
+ """
470
+
471
+ def __init__(self, quant_type: str, modules_to_not_convert: Optional[List[str]] = None, **kwargs) -> None:
472
+ self.quant_method = QuantizationMethod.TORCHAO
473
+ self.quant_type = quant_type
474
+ self.modules_to_not_convert = modules_to_not_convert
475
+
476
+ # When we load from serialized config, "quant_type_kwargs" will be the key
477
+ if "quant_type_kwargs" in kwargs:
478
+ self.quant_type_kwargs = kwargs["quant_type_kwargs"]
479
+ else:
480
+ self.quant_type_kwargs = kwargs
481
+
482
+ TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method()
483
+ if self.quant_type not in TORCHAO_QUANT_TYPE_METHODS.keys():
484
+ raise ValueError(
485
+ f"Requested quantization type: {self.quant_type} is not supported yet or is incorrect. If you think the "
486
+ f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues."
487
+ )
488
+
489
+ method = TORCHAO_QUANT_TYPE_METHODS[self.quant_type]
490
+ signature = inspect.signature(method)
491
+ all_kwargs = {
492
+ param.name
493
+ for param in signature.parameters.values()
494
+ if param.kind in [inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD]
495
+ }
496
+ unsupported_kwargs = list(self.quant_type_kwargs.keys() - all_kwargs)
497
+
498
+ if len(unsupported_kwargs) > 0:
499
+ raise ValueError(
500
+ f'The quantization method "{quant_type}" does not support the following keyword arguments: '
501
+ f"{unsupported_kwargs}. The following keywords arguments are supported: {all_kwargs}."
502
+ )
503
+
504
+ @classmethod
505
+ def _get_torchao_quant_type_to_method(cls):
506
+ r"""
507
+ Returns supported torchao quantization types with all commonly used notations.
508
+ """
509
+
510
+ if is_torchao_available():
511
+ # TODO(aryan): Support autoquant and sparsify
512
+ from torchao.quantization import (
513
+ float8_dynamic_activation_float8_weight,
514
+ float8_static_activation_float8_weight,
515
+ float8_weight_only,
516
+ fpx_weight_only,
517
+ int4_weight_only,
518
+ int8_dynamic_activation_int4_weight,
519
+ int8_dynamic_activation_int8_weight,
520
+ int8_weight_only,
521
+ uintx_weight_only,
522
+ )
523
+
524
+ # TODO(aryan): Add a note on how to use PerAxis and PerGroup observers
525
+ from torchao.quantization.observer import PerRow, PerTensor
526
+
527
+ def generate_float8dq_types(dtype: torch.dtype):
528
+ name = "e5m2" if dtype == torch.float8_e5m2 else "e4m3"
529
+ types = {}
530
+
531
+ for granularity_cls in [PerTensor, PerRow]:
532
+ # Note: Activation and Weights cannot have different granularities
533
+ granularity_name = "tensor" if granularity_cls is PerTensor else "row"
534
+ types[f"float8dq_{name}_{granularity_name}"] = partial(
535
+ float8_dynamic_activation_float8_weight,
536
+ activation_dtype=dtype,
537
+ weight_dtype=dtype,
538
+ granularity=(granularity_cls(), granularity_cls()),
539
+ )
540
+
541
+ return types
542
+
543
+ def generate_fpx_quantization_types(bits: int):
544
+ types = {}
545
+
546
+ for ebits in range(1, bits):
547
+ mbits = bits - ebits - 1
548
+ types[f"fp{bits}_e{ebits}m{mbits}"] = partial(fpx_weight_only, ebits=ebits, mbits=mbits)
549
+
550
+ non_sign_bits = bits - 1
551
+ default_ebits = (non_sign_bits + 1) // 2
552
+ default_mbits = non_sign_bits - default_ebits
553
+ types[f"fp{bits}"] = partial(fpx_weight_only, ebits=default_ebits, mbits=default_mbits)
554
+
555
+ return types
556
+
557
+ INT4_QUANTIZATION_TYPES = {
558
+ # int4 weight + bfloat16/float16 activation
559
+ "int4wo": int4_weight_only,
560
+ "int4_weight_only": int4_weight_only,
561
+ # int4 weight + int8 activation
562
+ "int4dq": int8_dynamic_activation_int4_weight,
563
+ "int8_dynamic_activation_int4_weight": int8_dynamic_activation_int4_weight,
564
+ }
565
+
566
+ INT8_QUANTIZATION_TYPES = {
567
+ # int8 weight + bfloat16/float16 activation
568
+ "int8wo": int8_weight_only,
569
+ "int8_weight_only": int8_weight_only,
570
+ # int8 weight + int8 activation
571
+ "int8dq": int8_dynamic_activation_int8_weight,
572
+ "int8_dynamic_activation_int8_weight": int8_dynamic_activation_int8_weight,
573
+ }
574
+
575
+ # TODO(aryan): handle torch 2.2/2.3
576
+ FLOATX_QUANTIZATION_TYPES = {
577
+ # float8_e5m2 weight + bfloat16/float16 activation
578
+ "float8wo": partial(float8_weight_only, weight_dtype=torch.float8_e5m2),
579
+ "float8_weight_only": float8_weight_only,
580
+ "float8wo_e5m2": partial(float8_weight_only, weight_dtype=torch.float8_e5m2),
581
+ # float8_e4m3 weight + bfloat16/float16 activation
582
+ "float8wo_e4m3": partial(float8_weight_only, weight_dtype=torch.float8_e4m3fn),
583
+ # float8_e5m2 weight + float8 activation (dynamic)
584
+ "float8dq": float8_dynamic_activation_float8_weight,
585
+ "float8_dynamic_activation_float8_weight": float8_dynamic_activation_float8_weight,
586
+ # ===== Matrix multiplication is not supported in float8_e5m2 so the following errors out.
587
+ # However, changing activation_dtype=torch.float8_e4m3 might work here =====
588
+ # "float8dq_e5m2": partial(
589
+ # float8_dynamic_activation_float8_weight,
590
+ # activation_dtype=torch.float8_e5m2,
591
+ # weight_dtype=torch.float8_e5m2,
592
+ # ),
593
+ # **generate_float8dq_types(torch.float8_e5m2),
594
+ # ===== =====
595
+ # float8_e4m3 weight + float8 activation (dynamic)
596
+ "float8dq_e4m3": partial(
597
+ float8_dynamic_activation_float8_weight,
598
+ activation_dtype=torch.float8_e4m3fn,
599
+ weight_dtype=torch.float8_e4m3fn,
600
+ ),
601
+ **generate_float8dq_types(torch.float8_e4m3fn),
602
+ # float8 weight + float8 activation (static)
603
+ "float8_static_activation_float8_weight": float8_static_activation_float8_weight,
604
+ # For fpx, only x <= 8 is supported by default. Other dtypes can be explored by users directly
605
+ # fpx weight + bfloat16/float16 activation
606
+ **generate_fpx_quantization_types(3),
607
+ **generate_fpx_quantization_types(4),
608
+ **generate_fpx_quantization_types(5),
609
+ **generate_fpx_quantization_types(6),
610
+ **generate_fpx_quantization_types(7),
611
+ }
612
+
613
+ UINTX_QUANTIZATION_DTYPES = {
614
+ "uintx_weight_only": uintx_weight_only,
615
+ "uint1wo": partial(uintx_weight_only, dtype=torch.uint1),
616
+ "uint2wo": partial(uintx_weight_only, dtype=torch.uint2),
617
+ "uint3wo": partial(uintx_weight_only, dtype=torch.uint3),
618
+ "uint4wo": partial(uintx_weight_only, dtype=torch.uint4),
619
+ "uint5wo": partial(uintx_weight_only, dtype=torch.uint5),
620
+ "uint6wo": partial(uintx_weight_only, dtype=torch.uint6),
621
+ "uint7wo": partial(uintx_weight_only, dtype=torch.uint7),
622
+ # "uint8wo": partial(uintx_weight_only, dtype=torch.uint8), # uint8 quantization is not supported
623
+ }
624
+
625
+ QUANTIZATION_TYPES = {}
626
+ QUANTIZATION_TYPES.update(INT4_QUANTIZATION_TYPES)
627
+ QUANTIZATION_TYPES.update(INT8_QUANTIZATION_TYPES)
628
+ QUANTIZATION_TYPES.update(UINTX_QUANTIZATION_DTYPES)
629
+
630
+ if cls._is_cuda_capability_atleast_8_9():
631
+ QUANTIZATION_TYPES.update(FLOATX_QUANTIZATION_TYPES)
632
+
633
+ return QUANTIZATION_TYPES
634
+ else:
635
+ raise ValueError(
636
+ "TorchAoConfig requires torchao to be installed, please install with `pip install torchao`"
637
+ )
638
+
639
+ @staticmethod
640
+ def _is_cuda_capability_atleast_8_9() -> bool:
641
+ if not torch.cuda.is_available():
642
+ raise RuntimeError("TorchAO requires a CUDA compatible GPU and installation of PyTorch.")
643
+
644
+ major, minor = torch.cuda.get_device_capability()
645
+ if major == 8:
646
+ return minor >= 9
647
+ return major >= 9
648
+
649
+ def get_apply_tensor_subclass(self):
650
+ TORCHAO_QUANT_TYPE_METHODS = self._get_torchao_quant_type_to_method()
651
+ return TORCHAO_QUANT_TYPE_METHODS[self.quant_type](**self.quant_type_kwargs)
652
+
653
+ def __repr__(self):
654
+ r"""
655
+ Example of how this looks for `TorchAoConfig("uint_a16w4", group_size=32)`:
656
+
657
+ ```
658
+ TorchAoConfig {
659
+ "modules_to_not_convert": null,
660
+ "quant_method": "torchao",
661
+ "quant_type": "uint_a16w4",
662
+ "quant_type_kwargs": {
663
+ "group_size": 32
664
+ }
665
+ }
666
+ ```
667
+ """
668
+ config_dict = self.to_dict()
669
+ return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
@@ -0,0 +1,15 @@
1
+ # Copyright 2024 The HuggingFace Inc. team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ from .torchao_quantizer import TorchAoHfQuantizer