





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

_lowerCamelCase   				=			logging.get_logger(__name__)




@add_end_docstrings(UpperCamelCase__    )
class   UpperCamelCase_       (     UpperCamelCase__    ):


							def __init__(							self    :List[Any]							,				**__A    :Optional[Any]		)   ->   int:



									"""simple docstring"""



									super().__init__(**__A		)
									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(__A		)


							def      _snake_case  (							self    :str							,				**__A    :Optional[Any]		)   ->   str:



									"""simple docstring"""



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


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



									"""simple docstring"""



									return super().__call__(__A							,				*__A							,				num_workers=__A							,				batch_size=__A							,				**__A		)


							def      _snake_case  (							self    :Any							,				__A    :int							,				__A    :Optional[int]=64							,				__A    :int = 0							,				__A    :float = 512 / 1500							,				__A    :Optional[int] = 32							,				__A    :Optional[int] = 1							,				)   ->   Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      load_image(__A		)
									SCREAMING_SNAKE_CASE__			=      self.image_processor.size["""longest_edge"""]
									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      self.image_processor.generate_crop_boxes(
									    __A							,				__A							,				__A							,				__A							,				__A							,				__A		)
									SCREAMING_SNAKE_CASE__			=      self.image_processor(images=__A							,				return_tensors="""pt"""		)

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

									SCREAMING_SNAKE_CASE__			=      grid_points.shape[1]
									SCREAMING_SNAKE_CASE__			=      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							,				__A							,				__A		):
											SCREAMING_SNAKE_CASE__			=      grid_points[:, i : i + points_per_batch, :, :]
											SCREAMING_SNAKE_CASE__			=      input_labels[:, i : i + points_per_batch]
											SCREAMING_SNAKE_CASE__			=      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      _snake_case  (							self    :Optional[Any]							,				__A    :str							,				__A    :List[Any]=0.8_8							,				__A    :Tuple=0.9_5							,				__A    :Optional[Any]=0							,				__A    :List[Any]=1							,				)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      model_inputs.pop("""input_boxes"""		)
									SCREAMING_SNAKE_CASE__			=      model_inputs.pop("""is_last"""		)
									SCREAMING_SNAKE_CASE__			=      model_inputs.pop("""original_sizes"""		).tolist()
									SCREAMING_SNAKE_CASE__			=      model_inputs.pop("""reshaped_input_sizes"""		).tolist()

									SCREAMING_SNAKE_CASE__			=      self.model(**__A		)

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


							def      _snake_case  (							self    :Union[str, Any]							,				__A    :Union[str, Any]							,				__A    :Any=False							,				__A    :Any=False							,				__A    :str=0.7							,				)   ->   List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      []
									SCREAMING_SNAKE_CASE__			=      []
									SCREAMING_SNAKE_CASE__			=      []
									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"""		)		)

									SCREAMING_SNAKE_CASE__			=      torch.cat(__A		)
									SCREAMING_SNAKE_CASE__			=      torch.cat(__A		)
									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      self.image_processor.post_process_for_mask_generation(
									    __A							,				__A							,				__A							,				__A		)

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

									SCREAMING_SNAKE_CASE__			=      {}
									if output_rle_mask:
											SCREAMING_SNAKE_CASE__			=      rle_mask

									if output_bboxes_mask:
											SCREAMING_SNAKE_CASE__			=      bounding_boxes

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






import inspect
import unittest




class   UpperCamelCase_       (     unittest.TestCase    ):


							def      _snake_case  (							self    :str		)   ->   Union[str, Any]:



									"""simple docstring"""



									try:
											import diffusers  # noqa: F401
									except ImportError:
											assert False


							def      _snake_case  (							self    :Any		)   ->   Any:



									"""simple docstring"""



									import diffusers
									from diffusers.dependency_versions_table import deps

									SCREAMING_SNAKE_CASE__			=      inspect.getmembers(__A							,				inspect.isclass		)

									for cls_name, cls_module in all_classes:
											if "dummy_" in cls_module.__module__:
													for backend in cls_module._backends:
															if backend == "k_diffusion":
																	SCREAMING_SNAKE_CASE__			=      """k-diffusion"""
															elif backend == "invisible_watermark":
																	SCREAMING_SNAKE_CASE__			=      """invisible-watermark"""
															assert backend in deps, f'''{backend} is not in the deps table!'''