
"""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


SCREAMING_SNAKE_CASE			  =  logging.getLogger(__name__)

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



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







       """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__      (				UpperCAmelCase_					)->				Union[str, Any]:







       """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__      (				UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_=False		,   UpperCAmelCase_=False					)->				Tuple:







       """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__      (				UpperCAmelCase_					)->				List[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__      (				UpperCAmelCase_		,   UpperCAmelCase_					)->				str:







       """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__      (				UpperCAmelCase_		,   UpperCAmelCase_					)->				int:







       """simple docstring"""

       def fusea(UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_					):
              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__      (				UpperCAmelCase_		,   UpperCAmelCase_					)->				Dict:







       """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__      (				UpperCAmelCase_					)->				int:







       """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__      (				UpperCAmelCase_					)->				Dict:







       """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__      (				UpperCAmelCase_		,   UpperCAmelCase_=25		,   UpperCAmelCase_=1_80		,   UpperCAmelCase_=None					)->				str:







       """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__      (				UpperCAmelCase_					)->				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__      (				UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_					)->				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__      (				UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_="both"		,   **UpperCAmelCase_					)->				Tuple:







       """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__      (				UpperCAmelCase_		,   UpperCAmelCase_		,   **UpperCAmelCase_					)->				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"""


from collections import defaultdict
from typing import Optional

from ..image_utils import load_image
from ..utils import (
    add_end_docstrings,
    is_torch_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_torch_available():
      import torch

      from ..models.auto.modeling_auto import MODEL_FOR_MASK_GENERATION_MAPPING

SCREAMING_SNAKE_CASE			  =  logging.get_logger(__name__)



@add_end_docstrings(_lowerCAmelCase						)
class 	__a	(   _lowerCAmelCase						):






       def __init__(	self			:				str							,   **UpperCAmelCase_			:				List[str]      )->	Tuple:







              """simple docstring"""
              super().__init__(**UpperCAmelCase_      )
              requires_backends(self							,   "vision"      )
              requires_backends(self							,   "torch"      )

              if self.framework != "pt":
                     raise ValueError(f"The {self.__class__} is only available in PyTorch."      )

              self.check_model_type(UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple							,   **UpperCAmelCase_			:				int      )->	List[str]:







              """simple docstring"""
              UpperCamelCase       			=		{}
              UpperCamelCase       			=		{}
              UpperCamelCase       			=		{}
              # preprocess args
              if "points_per_batch" in kwargs:
                     UpperCamelCase       			=		kwargs["points_per_batch"]
              if "points_per_crop" in kwargs:
                     UpperCamelCase       			=		kwargs["points_per_crop"]
              if "crops_n_layers" in kwargs:
                     UpperCamelCase       			=		kwargs["crops_n_layers"]
              if "crop_overlap_ratio" in kwargs:
                     UpperCamelCase       			=		kwargs["crop_overlap_ratio"]
              if "crop_n_points_downscale_factor" in kwargs:
                     UpperCamelCase       			=		kwargs["crop_n_points_downscale_factor"]
              # postprocess args
              if "pred_iou_thresh" in kwargs:
                     UpperCamelCase       			=		kwargs["pred_iou_thresh"]
              if "stability_score_offset" in kwargs:
                     UpperCamelCase       			=		kwargs["stability_score_offset"]
              if "mask_threshold" in kwargs:
                     UpperCamelCase       			=		kwargs["mask_threshold"]
              if "stability_score_thresh" in kwargs:
                     UpperCamelCase       			=		kwargs["stability_score_thresh"]
              if "crops_nms_thresh" in kwargs:
                     UpperCamelCase       			=		kwargs["crops_nms_thresh"]
              if "output_rle_mask" in kwargs:
                     UpperCamelCase       			=		kwargs["output_rle_mask"]
              if "output_bboxes_mask" in kwargs:
                     UpperCamelCase       			=		kwargs["output_bboxes_mask"]
              return preprocess_kwargs, forward_params, postprocess_kwargs






       def __call__(	self			:				Optional[int]							,   UpperCAmelCase_			:				Optional[int]							,   *UpperCAmelCase_			:				Tuple							,   UpperCAmelCase_			:				int=None							,   UpperCAmelCase_			:				Tuple=None							,   **UpperCAmelCase_			:				str      )->	Any:







              """simple docstring"""
              return super().__call__(UpperCAmelCase_							,   *UpperCAmelCase_							,   num_workers=UpperCAmelCase_							,   batch_size=UpperCAmelCase_							,   **UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]							,   UpperCAmelCase_			:				Optional[int]							,   UpperCAmelCase_			:				str=64							,   UpperCAmelCase_			:				int = 0							,   UpperCAmelCase_			:				float = 512 / 1_500							,   UpperCAmelCase_			:				Optional[int] = 32							,   UpperCAmelCase_			:				Optional[int] = 1							,   )->	List[Any]:







              """simple docstring"""
              UpperCamelCase       			=		load_image(UpperCAmelCase_      )
              UpperCamelCase       			=		self.image_processor.size["longest_edge"]
              UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase       			=		self.image_processor.generate_crop_boxes(
                  UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_      )
              UpperCamelCase       			=		self.image_processor(images=UpperCAmelCase_							,   return_tensors="pt"      )

              with self.device_placement():
                     if self.framework == "pt":
                            UpperCamelCase       			=		self.get_inference_context()
                            with inference_context():
                                   UpperCamelCase       			=		self._ensure_tensor_on_device(UpperCAmelCase_							,   device=self.device      )
                                   UpperCamelCase       			=		self.model.get_image_embeddings(model_inputs.pop("pixel_values"      )      )
                                   UpperCamelCase       			=		image_embeddings

              UpperCamelCase       			=		grid_points.shape[1]
              UpperCamelCase       			=		points_per_batch if points_per_batch is not None else n_points

              if points_per_batch <= 0:
                     raise ValueError(
                         "Cannot have points_per_batch<=0. Must be >=1 to returned batched outputs. "
                         "To return all points at once, set points_per_batch to None"      )

              for i in range(0							,   UpperCAmelCase_							,   UpperCAmelCase_      ):
                     UpperCamelCase       			=		grid_points[:, i : i + points_per_batch, :, :]
                     UpperCamelCase       			=		input_labels[:, i : i + points_per_batch]
                     UpperCamelCase       			=		i == n_points - points_per_batch
                     yield {
                         "input_points": batched_points,
                         "input_labels": labels,
                         "input_boxes": crop_boxes,
                         "is_last": is_last,
                         **model_inputs,
                     }






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Union[str, Any]							,   UpperCAmelCase_			:				Union[str, Any]							,   UpperCAmelCase_			:				int=0.88							,   UpperCAmelCase_			:				Optional[int]=0.95							,   UpperCAmelCase_			:				Dict=0							,   UpperCAmelCase_			:				int=1							,   )->	Any:







              """simple docstring"""
              UpperCamelCase       			=		model_inputs.pop("input_boxes"      )
              UpperCamelCase       			=		model_inputs.pop("is_last"      )
              UpperCamelCase       			=		model_inputs.pop("original_sizes"      ).tolist()
              UpperCamelCase       			=		model_inputs.pop("reshaped_input_sizes"      ).tolist()

              UpperCamelCase       			=		self.model(**UpperCAmelCase_      )

              # post processing happens here in order to avoid CPU GPU copies of ALL the masks
              UpperCamelCase       			=		model_outputs["pred_masks"]
              UpperCamelCase       			=		self.image_processor.post_process_masks(
                  UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   binarize=UpperCAmelCase_      )
              UpperCamelCase       			=		model_outputs["iou_scores"]
              UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase       			=		self.image_processor.filter_masks(
                  masks[0]							,   iou_scores[0]							,   original_sizes[0]							,   input_boxes[0]							,   UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   )
              return {
                  "masks": masks,
                  "is_last": is_last,
                  "boxes": boxes,
                  "iou_scores": iou_scores,
              }






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any							,   UpperCAmelCase_			:				Tuple							,   UpperCAmelCase_			:				List[str]=False							,   UpperCAmelCase_			:				str=False							,   UpperCAmelCase_			:				Optional[int]=0.7							,   )->	Dict:







              """simple docstring"""
              UpperCamelCase       			=		[]
              UpperCamelCase       			=		[]
              UpperCamelCase       			=		[]
              for model_output in model_outputs:
                     all_scores.append(model_output.pop("iou_scores"      )      )
                     all_masks.extend(model_output.pop("masks"      )      )
                     all_boxes.append(model_output.pop("boxes"      )      )

              UpperCamelCase       			=		torch.cat(UpperCAmelCase_      )
              UpperCamelCase       			=		torch.cat(UpperCAmelCase_      )
              UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase       			=		self.image_processor.post_process_for_mask_generation(
                  UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_							,   UpperCAmelCase_      )

              UpperCamelCase       			=		defaultdict(UpperCAmelCase_      )
              for output in model_outputs:
                     for k, v in output.items():
                            extra[k].append(UpperCAmelCase_      )

              UpperCamelCase       			=		{}
              if output_rle_mask:
                     UpperCamelCase       			=		rle_mask

              if output_bboxes_mask:
                     UpperCamelCase       			=		bounding_boxes

              return {"masks": output_masks, "scores": iou_scores, **optional, **extra}





