
"""simple docstring"""





import logging
import re

import pytorch_quantization
import pytorch_quantization.nn as quant_nn
import torch
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor


_a														=		logging.getLogger(__name__)

_a														=		50  # max width of layer names
_a														=		70  # max width of quantizer names



def       lowerCamelCase__		(			__snake_case	)							->		Dict:


					"""simple docstring"""





					_UpperCamelCase							 =  parser.add_argument_group('''quant_trainer arguments'''	)
					group.add_argument('''--wprec''',     type=UpperCamelCase__,     default=8,     help='''weight precision'''	)
					group.add_argument('''--aprec''',     type=UpperCamelCase__,     default=8,     help='''activation precision'''	)
					group.add_argument('''--quant-per-tensor''',     action='''store_true''',     help='''per tensor weight scaling'''	)
					group.add_argument('''--quant-disable''',     action='''store_true''',     help='''disable all quantizers'''	)
					group.add_argument('''--quant-disable-embeddings''',     action='''store_true''',     help='''disable all embeddings quantizers'''	)
					group.add_argument('''--quant-disable-keyword''',     type=UpperCamelCase__,     nargs='''+''',     help='''disable quantizers by keyword'''	)
					group.add_argument('''--quant-disable-layer-module''',     type=UpperCamelCase__,     help='''disable quantizers by keyword under layer.'''	)
					group.add_argument('''--quant-enable-layer-module''',     type=UpperCamelCase__,     help='''enable quantizers by keyword under layer'''	)
					group.add_argument('''--calibrator''',     default='''max''',     help='''which quantization range calibrator to use'''	)
					group.add_argument('''--percentile''',     default=UpperCamelCase__,     type=UpperCamelCase__,     help='''percentile for PercentileCalibrator'''	)
					group.add_argument('''--fuse-qkv''',     action='''store_true''',     help='''use the same scale factor for qkv'''	)
					group.add_argument('''--clip-gelu''',     metavar='''N''',     type=UpperCamelCase__,     help='''clip gelu output maximum value to N'''	)
					group.add_argument(
					    '''--recalibrate-weights''',     action='''store_true''',     help=(
					        '''recalibrate weight amaxes by taking the max of the weights.'''
					        ''' amaxes will be computed with the current quantization granularity (axis).'''
					    ),     )



def       lowerCamelCase__		(			__snake_case	)							->		int:


					"""simple docstring"""





					if args.calibrator == "max":
										_UpperCamelCase							 =  '''max'''
					elif args.calibrator == "percentile":
										if args.percentile is None:
															raise ValueError('''Specify --percentile when using percentile calibrator'''	)
										_UpperCamelCase							 =  '''histogram'''
					elif args.calibrator == "mse":
										_UpperCamelCase							 =  '''histogram'''
					else:
										raise ValueError(F'''Invalid calibrator {args.calibrator}'''	)

					_UpperCamelCase							 =  QuantDescriptor(num_bits=args.aprec,     calib_method=UpperCamelCase__	)
					_UpperCamelCase							 =  QuantDescriptor(num_bits=args.wprec,     axis=(None if args.quant_per_tensor else (0,))	)
					quant_nn.QuantLinear.set_default_quant_desc_input(UpperCamelCase__	)
					quant_nn.QuantLinear.set_default_quant_desc_weight(UpperCamelCase__	)



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case=False,     __snake_case=False	)							->		List[str]:


					"""simple docstring"""





					logger.info('''Configuring Model for Quantization'''	)
					logger.info(F'''using quantization package {pytorch_quantization.__file__}'''	)

					if not calib:
										if args.quant_disable_embeddings:
															set_quantizer_by_name(UpperCamelCase__,     ['''embeddings'''],     which='''weight''',     _disabled=UpperCamelCase__	)

										if args.quant_disable:
															set_quantizer_by_name(UpperCamelCase__,     [''''''],     _disabled=UpperCamelCase__	)

										if args.quant_disable_keyword:
															set_quantizer_by_name(UpperCamelCase__,     args.quant_disable_keyword,     _disabled=UpperCamelCase__	)

										if args.quant_disable_layer_module:
															set_quantizer_by_name(UpperCamelCase__,     [r'''layer.\d+.''' + args.quant_disable_layer_module],     _disabled=UpperCamelCase__	)

										if args.quant_enable_layer_module:
															set_quantizer_by_name(UpperCamelCase__,     [r'''layer.\d+.''' + args.quant_enable_layer_module],     _disabled=UpperCamelCase__	)

										if args.recalibrate_weights:
															recalibrate_weights(UpperCamelCase__	)

										if args.fuse_qkv:
															fuse_qkv(UpperCamelCase__,     UpperCamelCase__	)

					if args.clip_gelu:
										clip_gelu(UpperCamelCase__,     args.clip_gelu	)

					# if args.local_rank in [-1, 0] and not calib:
					print_quant_summary(UpperCamelCase__	)



def       lowerCamelCase__		(			__snake_case	)							->		Optional[Any]:


					"""simple docstring"""





					logger.info('''Enabling Calibration'''	)
					for name, module in model.named_modules():
										if name.endswith('''_quantizer'''	):
															if module._calibrator is not None:
																				module.disable_quant()
																				module.enable_calib()
															else:
																				module.disable()
															logger.info(F'''{name:80}: {module}'''	)



def       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		Optional[Any]:


					"""simple docstring"""





					logger.info('''Loading calibrated amax'''	)
					for name, module in model.named_modules():
										if name.endswith('''_quantizer'''	):
															if module._calibrator is not None:
																				if isinstance(module._calibrator,     calib.MaxCalibrator	):
																									module.load_calib_amax()
																				else:
																									module.load_calib_amax('''percentile''',     percentile=args.percentile	)
																				module.enable_quant()
																				module.disable_calib()
															else:
																				module.enable()
					model.cuda()
					print_quant_summary(UpperCamelCase__	)



def       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		Tuple:


					"""simple docstring"""





					def fusea(__snake_case,     __snake_case,     __snake_case	):
										for mod in [qq, qk, qv]:
															if not hasattr(UpperCamelCase__,     '''_amax'''	):
																				print('''          WARNING: NO AMAX BUFFER'''	)
																				return
										_UpperCamelCase							 =  qq._amax.detach().item()
										_UpperCamelCase							 =  qk._amax.detach().item()
										_UpperCamelCase							 =  qv._amax.detach().item()

										_UpperCamelCase							 =  max(UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__	)
										qq._amax.fill_(UpperCamelCase__	)
										qk._amax.fill_(UpperCamelCase__	)
										qv._amax.fill_(UpperCamelCase__	)
										logger.info(F'''          q={q:5.2f} k={k:5.2f} v={v:5.2f} -> {amax:5.2f}'''	)

					for name, mod in model.named_modules():
										if name.endswith('''.attention.self'''	):
															logger.info(F'''FUSE_QKV: {name:{name_width}}'''	)
															fusea(mod.matmul_q_input_quantizer,     mod.matmul_k_input_quantizer,     mod.matmul_v_input_quantizer	)
															if args.quant_per_tensor:
																				fusea(mod.query._weight_quantizer,     mod.key._weight_quantizer,     mod.value._weight_quantizer	)



def       lowerCamelCase__		(			__snake_case,     __snake_case	)							->		Union[str, Any]:


					"""simple docstring"""





					for name, mod in model.named_modules():
										if name.endswith('''.output.dense'''	) and not name.endswith('''attention.output.dense'''	):
															_UpperCamelCase							 =  mod._input_quantizer._amax.data.detach().item()
															mod._input_quantizer._amax.data.detach().clamp_(max=UpperCamelCase__	)
															_UpperCamelCase							 =  mod._input_quantizer._amax.data.detach().item()
															logger.info(F'''CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}'''	)



def       lowerCamelCase__		(			__snake_case	)							->		Tuple:


					"""simple docstring"""





					for name, mod in model.named_modules():
										if hasattr(UpperCamelCase__,     '''_weight_quantizer'''	) and mod._weight_quantizer.axis is not None:
															_UpperCamelCase							 =  mod.weight.shape[0]
															_UpperCamelCase							 =  mod._weight_quantizer._amax.detach()
															_UpperCamelCase							 =  torch.ones(UpperCamelCase__,     dtype=amax.dtype,     device=amax.device	) * amax
															print(F'''expanding {name} {amax} -> {mod._weight_quantizer._amax}'''	)



def       lowerCamelCase__		(			__snake_case	)							->		Tuple:


					"""simple docstring"""





					for name, mod in model.named_modules():
										if hasattr(UpperCamelCase__,     '''_weight_quantizer'''	):
															if not hasattr(mod.weight_quantizer,     '''_amax'''	):
																				print('''RECALIB: {name:{name_width}} WARNING: NO AMAX BUFFER'''	)
																				continue

															# determine which axes to reduce across
															# e.g. a 4D tensor quantized per axis 0 should reduce over (1,2,3)
															_UpperCamelCase							 =  set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis	)
															_UpperCamelCase							 =  set(range(len(mod.weight.size()	)	)	) - axis_set
															_UpperCamelCase							 =  pytorch_quantization.utils.reduce_amax(mod.weight,     axis=UpperCamelCase__,     keepdims=UpperCamelCase__	).detach()
															logger.info(F'''RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}'''	)
															_UpperCamelCase							 =  amax



def       lowerCamelCase__		(			__snake_case,     __snake_case=25,     __snake_case=1_80,     __snake_case=None	)							->		Optional[Any]:


					"""simple docstring"""





					if ignore is None:
										_UpperCamelCase							 =  []
					elif not isinstance(UpperCamelCase__,     UpperCamelCase__	):
										_UpperCamelCase							 =  [ignore]

					_UpperCamelCase							 =  0
					for name, mod in model.named_modules():
										if not hasattr(UpperCamelCase__,     '''weight'''	):
															continue
										_UpperCamelCase							 =  max(UpperCamelCase__,     len(UpperCamelCase__	)	)

					for name, mod in model.named_modules():
										_UpperCamelCase							 =  getattr(UpperCamelCase__,     '''_input_quantizer''',     UpperCamelCase__	)
										_UpperCamelCase							 =  getattr(UpperCamelCase__,     '''_weight_quantizer''',     UpperCamelCase__	)
										if not hasattr(UpperCamelCase__,     '''weight'''	):
															continue
										if type(UpperCamelCase__	) in ignore:
															continue
										if [True for s in ignore if type(UpperCamelCase__	) is str and s in name]:
															continue
										_UpperCamelCase							 =  F'''Act:{input_q.extra_repr()}'''
										_UpperCamelCase							 =  F'''Wgt:{weight_q.extra_repr()}'''
										_UpperCamelCase							 =  F'''{name:{name_width}} {act_str} {wgt_str}'''
										if len(UpperCamelCase__	) <= line_width:
															logger.info(UpperCamelCase__	)
										else:
															logger.info(F'''{name:{name_width}} {act_str}'''	)
															logger.info(F'''{"  ":{name_width}} {wgt_str}'''	)



def       lowerCamelCase__		(			__snake_case	)							->		Tuple:


					"""simple docstring"""





					_UpperCamelCase							 =  0
					for name, mod in model.named_modules():
										if isinstance(UpperCamelCase__,     pytorch_quantization.nn.TensorQuantizer	):
															print(F'''{name:80} {mod}'''	)
															count += 1
					print(F'''{count} TensorQuantizers found in model'''	)



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case,     __snake_case,     __snake_case	)							->		List[Any]:


					"""simple docstring"""





					_UpperCamelCase							 =  getattr(UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__	)
					if quantizer_mod is not None:
										assert hasattr(UpperCamelCase__,     UpperCamelCase__	)
										setattr(UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__	)
					else:
										logger.warning(F'''{name} has no {quantizer}'''	)



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case="both",     **__snake_case	)							->		str:


					"""simple docstring"""





					_UpperCamelCase							 =  F'''Warning: changing {which} quantizers of {name:{qname_width}}'''
					for k, v in kwargs.items():
										s += F''' {k}={v}'''
										if which in ["input", "both"]:
															set_quantizer(UpperCamelCase__,     UpperCamelCase__,     '''_input_quantizer''',     UpperCamelCase__,     UpperCamelCase__	)
										if which in ["weight", "both"]:
															set_quantizer(UpperCamelCase__,     UpperCamelCase__,     '''_weight_quantizer''',     UpperCamelCase__,     UpperCamelCase__	)
					logger.info(UpperCamelCase__	)



def       lowerCamelCase__		(			__snake_case,     __snake_case,     **__snake_case	)							->		Union[str, Any]:


					"""simple docstring"""





					for name, mod in model.named_modules():
										if hasattr(UpperCamelCase__,     '''_input_quantizer'''	) or hasattr(UpperCamelCase__,     '''_weight_quantizer'''	):
															for n in names:
																				if re.search(UpperCamelCase__,     UpperCamelCase__	):
																									set_quantizers(UpperCamelCase__,     UpperCamelCase__,     **UpperCamelCase__	)
										elif name.endswith('''_quantizer'''	):
															for n in names:
																				if re.search(UpperCamelCase__,     UpperCamelCase__	):
																									_UpperCamelCase							 =  F'''Warning: changing {name:{name_width}}'''
																									for k, v in kwargs.items():
																														s += F''' {k}={v}'''
																														setattr(UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__	)
																									logger.info(UpperCamelCase__	)











"""simple docstring"""
import random
import unittest

import torch

from diffusers import IFImgaImgSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin
@skip_mps
class 							_UpperCAmelCase							( a		,a		,unittest.TestCase      ):





     '''simple docstring'''



     a__									=IFImgaImgSuperResolutionPipeline
     a__									=TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''width''', '''height'''}
     a__									=TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''original_image'''}      )
     a__									=PipelineTesterMixin.required_optional_params - {'''latents'''}





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         return self._get_superresolution_dummy_components()





     def  __lowerCAmelCase	(					self       ,	A       ,	A=0		)							->					Union[str, Any]:
         if str(A		).startswith('''mps'''		):
             _UpperCAmelCase							:  Any											=    torch.manual_seed(A		)
         else:
             _UpperCAmelCase							:  int											=    torch.Generator(device=A		).manual_seed(A		)

         _UpperCAmelCase							:  str											=    floats_tensor((1, 3, 3_2, 3_2)       ,	rng=random.Random(A		)		).to(A		)
         _UpperCAmelCase							:  Dict											=    floats_tensor((1, 3, 1_6, 1_6)       ,	rng=random.Random(A		)		).to(A		)

         _UpperCAmelCase							:  List[Any]											=    {
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''image''': image,
             '''original_image''': original_image,
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''output_type''': '''numpy''',
         }

         return inputs





     @unittest.skipIf(
         torch_device != '''cuda''' or not is_xformers_available()       ,	reason='''XFormers attention is only available with CUDA and `xformers` installed'''       ,	)
     def  __lowerCAmelCase	(					self		)							->					List[Any]:
         self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3		)





     def  __lowerCAmelCase	(					self		)							->					List[str]:
         self._test_save_load_optional_components()





     @unittest.skipIf(torch_device != '''cuda'''       ,	reason='''float16 requires CUDA'''		)
     def  __lowerCAmelCase	(					self		)							->					Optional[Any]:
         # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
         super().test_save_load_floataa(expected_max_diff=1E-1		)





     def  __lowerCAmelCase	(					self		)							->					int:
         self._test_attention_slicing_forward_pass(expected_max_diff=1E-2		)





     def  __lowerCAmelCase	(					self		)							->					Union[str, Any]:
         self._test_save_load_local()





     def  __lowerCAmelCase	(					self		)							->					Union[str, Any]:
         self._test_inference_batch_single_identical(
             expected_max_diff=1E-2       ,	)


