
"""simple docstring"""


import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible



class 	__a	(   unittest.TestCase						):






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[Any]      )->	Tuple:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "safety_checker/pytorch_model.bin",
                  "safety_checker/model.safetensors",
                  "vae/diffusion_pytorch_model.bin",
                  "vae/diffusion_pytorch_model.safetensors",
                  "text_encoder/pytorch_model.bin",
                  "text_encoder/model.safetensors",
                  "unet/diffusion_pytorch_model.bin",
                  "unet/diffusion_pytorch_model.safetensors",
              ]
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	Union[str, Any]:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "unet/diffusion_pytorch_model.bin",
                  "unet/diffusion_pytorch_model.safetensors",
              ]
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	List[Any]:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "safety_checker/pytorch_model.bin",
                  "safety_checker/model.safetensors",
                  "vae/diffusion_pytorch_model.bin",
                  "vae/diffusion_pytorch_model.safetensors",
                  "text_encoder/pytorch_model.bin",
                  "text_encoder/model.safetensors",
                  "unet/diffusion_pytorch_model.bin",
                  # Removed: 'unet/diffusion_pytorch_model.safetensors',
              ]
              self.assertFalse(is_safetensors_compatible(UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	Optional[int]:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "text_encoder/pytorch_model.bin",
                  "text_encoder/model.safetensors",
              ]
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[Any]      )->	List[str]:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "safety_checker/pytorch_model.bin",
                  "safety_checker/model.safetensors",
                  "vae/diffusion_pytorch_model.bin",
                  "vae/diffusion_pytorch_model.safetensors",
                  "text_encoder/pytorch_model.bin",
                  # Removed: 'text_encoder/model.safetensors',
                  "unet/diffusion_pytorch_model.bin",
                  "unet/diffusion_pytorch_model.safetensors",
              ]
              self.assertFalse(is_safetensors_compatible(UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple      )->	int:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "safety_checker/pytorch_model.fp16.bin",
                  "safety_checker/model.fp16.safetensors",
                  "vae/diffusion_pytorch_model.fp16.bin",
                  "vae/diffusion_pytorch_model.fp16.safetensors",
                  "text_encoder/pytorch_model.fp16.bin",
                  "text_encoder/model.fp16.safetensors",
                  "unet/diffusion_pytorch_model.fp16.bin",
                  "unet/diffusion_pytorch_model.fp16.safetensors",
              ]
              UpperCamelCase       			=		"fp16"
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_							,   variant=UpperCAmelCase_      )      )






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







              """simple docstring"""
              UpperCamelCase       			=		[
                  "unet/diffusion_pytorch_model.fp16.bin",
                  "unet/diffusion_pytorch_model.fp16.safetensors",
              ]
              UpperCamelCase       			=		"fp16"
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_							,   variant=UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[str]      )->	str:







              """simple docstring"""
              # pass variant but use the non-variant filenames
              UpperCamelCase       			=		[
                  "unet/diffusion_pytorch_model.bin",
                  "unet/diffusion_pytorch_model.safetensors",
              ]
              UpperCamelCase       			=		"fp16"
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_							,   variant=UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple      )->	int:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "safety_checker/pytorch_model.fp16.bin",
                  "safety_checker/model.fp16.safetensors",
                  "vae/diffusion_pytorch_model.fp16.bin",
                  "vae/diffusion_pytorch_model.fp16.safetensors",
                  "text_encoder/pytorch_model.fp16.bin",
                  "text_encoder/model.fp16.safetensors",
                  "unet/diffusion_pytorch_model.fp16.bin",
                  # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
              ]
              UpperCamelCase       			=		"fp16"
              self.assertFalse(is_safetensors_compatible(UpperCAmelCase_							,   variant=UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	Optional[int]:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "text_encoder/pytorch_model.fp16.bin",
                  "text_encoder/model.fp16.safetensors",
              ]
              UpperCamelCase       			=		"fp16"
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_							,   variant=UpperCAmelCase_      )      )






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







              """simple docstring"""
              # pass variant but use the non-variant filenames
              UpperCamelCase       			=		[
                  "text_encoder/pytorch_model.bin",
                  "text_encoder/model.safetensors",
              ]
              UpperCamelCase       			=		"fp16"
              self.assertTrue(is_safetensors_compatible(UpperCAmelCase_							,   variant=UpperCAmelCase_      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[str]      )->	str:







              """simple docstring"""
              UpperCamelCase       			=		[
                  "safety_checker/pytorch_model.fp16.bin",
                  "safety_checker/model.fp16.safetensors",
                  "vae/diffusion_pytorch_model.fp16.bin",
                  "vae/diffusion_pytorch_model.fp16.safetensors",
                  "text_encoder/pytorch_model.fp16.bin",
                  # 'text_encoder/model.fp16.safetensors',
                  "unet/diffusion_pytorch_model.fp16.bin",
                  "unet/diffusion_pytorch_model.fp16.safetensors",
              ]
              UpperCamelCase       			=		"fp16"
              self.assertFalse(is_safetensors_compatible(UpperCAmelCase_							,   variant=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}





