"""simple docstring"""






# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()


class     UpperCAmelCase						(_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,unittest.TestCase  ):



    """simple docstring"""

    _UpperCAmelCase	:Optional[int]						    = StableDiffusionControlNetImgaImgPipeline
    _UpperCAmelCase	:Tuple						    = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
    _UpperCAmelCase	:Any						    = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    _UpperCAmelCase	:Union[str, Any]						    = IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}  )
    _UpperCAmelCase	:Dict						    = IMAGE_TO_IMAGE_IMAGE_PARAMS






    def _snake_case			(   self							):
       torch.manual_seed(0							)
       lowercase__:   List[Any]					=      UNetaDConditionModel(
           block_out_channels=(32, 64)    ,			layers_per_block=2    ,			sample_size=32    ,			in_channels=4    ,			out_channels=4    ,			down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,			up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')    ,			cross_attention_dim=32    ,			)
       torch.manual_seed(0							)
       lowercase__:   Dict					=      ControlNetModel(
           block_out_channels=(32, 64)    ,			layers_per_block=2    ,			in_channels=4    ,			down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,			cross_attention_dim=32    ,			conditioning_embedding_out_channels=(16, 32)    ,			)
       torch.manual_seed(0							)
       lowercase__:   Tuple					=      DDIMScheduler(
           beta_start=0.00_085    ,			beta_end=0.012    ,			beta_schedule='''scaled_linear'''    ,			clip_sample=_UpperCAmelCase    ,			set_alpha_to_one=_UpperCAmelCase    ,			)
       torch.manual_seed(0							)
       lowercase__:   Optional[int]					=      AutoencoderKL(
           block_out_channels=[32, 64]    ,			in_channels=3    ,			out_channels=3    ,			down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']    ,			up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']    ,			latent_channels=4    ,			)
       torch.manual_seed(0							)
       lowercase__:   Optional[int]					=      CLIPTextConfig(
           bos_token_id=0    ,			eos_token_id=2    ,			hidden_size=32    ,			intermediate_size=37    ,			layer_norm_eps=1e-0_5    ,			num_attention_heads=4    ,			num_hidden_layers=5    ,			pad_token_id=1    ,			vocab_size=1000    ,			)
       lowercase__:   Tuple					=      CLIPTextModel(_UpperCAmelCase							)
       lowercase__:   List[str]					=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

       lowercase__:   Tuple					=      {
           '''unet''': unet,
           '''controlnet''': controlnet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=0							):
       if str(_UpperCAmelCase							).startswith('''mps'''							):
          lowercase__:   List[Any]					=      torch.manual_seed(_UpperCAmelCase							)
       else:
          lowercase__:   List[Any]					=      torch.Generator(device=_UpperCAmelCase							).manual_seed(_UpperCAmelCase							)

       lowercase__:   Tuple					=      2
       lowercase__:   Any					=      randn_tensor(
           (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)    ,			generator=_UpperCAmelCase    ,			device=torch.device(_UpperCAmelCase							)    ,			)
       lowercase__:   Tuple					=      floats_tensor(control_image.shape    ,			rng=random.Random(_UpperCAmelCase							)							).to(_UpperCAmelCase							)
       lowercase__:   Any					=      image.cpu().permute(0    ,			2    ,			3    ,			1							)[0]
       lowercase__:   Tuple					=      Image.fromarray(np.uinta(_UpperCAmelCase							)							).convert('''RGB'''							).resize((64, 64)							)
       lowercase__:   List[Any]					=      {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
           '''image''': image,
           '''control_image''': control_image,
       }

       return inputs






    def _snake_case			(   self							):
       return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3							)






    @unittest.skipIf(
        torch_device != '''cuda''' or not is_xformers_available()    ,			reason='''XFormers attention is only available with CUDA and `xformers` installed'''    ,			)
    def _snake_case			(   self							):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3							)






    def _snake_case			(   self							):
       self._test_inference_batch_single_identical(expected_max_diff=2e-3							)



class     UpperCAmelCase						(_UpperCAmelCase    ,_UpperCAmelCase    ,unittest.TestCase  ):



    """simple docstring"""

    _UpperCAmelCase	:Optional[int]						    = StableDiffusionControlNetImgaImgPipeline
    _UpperCAmelCase	:Optional[int]						    = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
    _UpperCAmelCase	:Dict						    = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    _UpperCAmelCase	:Tuple						    = frozenset([]  )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess






    def _snake_case			(   self							):
       torch.manual_seed(0							)
       lowercase__:   Dict					=      UNetaDConditionModel(
           block_out_channels=(32, 64)    ,			layers_per_block=2    ,			sample_size=32    ,			in_channels=4    ,			out_channels=4    ,			down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,			up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')    ,			cross_attention_dim=32    ,			)
       torch.manual_seed(0							)

       def init_weights(_UpperCAmelCase							):
          if isinstance(_UpperCAmelCase    ,			torch.nn.Convad							):
             torch.nn.init.normal(m.weight							)
             m.bias.data.fill_(1.0							)

       lowercase__:   Any					=      ControlNetModel(
           block_out_channels=(32, 64)    ,			layers_per_block=2    ,			in_channels=4    ,			down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,			cross_attention_dim=32    ,			conditioning_embedding_out_channels=(16, 32)    ,			)
       controlneta.controlnet_down_blocks.apply(_UpperCAmelCase							)

       torch.manual_seed(0							)
       lowercase__:   List[str]					=      ControlNetModel(
           block_out_channels=(32, 64)    ,			layers_per_block=2    ,			in_channels=4    ,			down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,			cross_attention_dim=32    ,			conditioning_embedding_out_channels=(16, 32)    ,			)
       controlneta.controlnet_down_blocks.apply(_UpperCAmelCase							)

       torch.manual_seed(0							)
       lowercase__:   Any					=      DDIMScheduler(
           beta_start=0.00_085    ,			beta_end=0.012    ,			beta_schedule='''scaled_linear'''    ,			clip_sample=_UpperCAmelCase    ,			set_alpha_to_one=_UpperCAmelCase    ,			)
       torch.manual_seed(0							)
       lowercase__:   Optional[Any]					=      AutoencoderKL(
           block_out_channels=[32, 64]    ,			in_channels=3    ,			out_channels=3    ,			down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']    ,			up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']    ,			latent_channels=4    ,			)
       torch.manual_seed(0							)
       lowercase__:   Tuple					=      CLIPTextConfig(
           bos_token_id=0    ,			eos_token_id=2    ,			hidden_size=32    ,			intermediate_size=37    ,			layer_norm_eps=1e-0_5    ,			num_attention_heads=4    ,			num_hidden_layers=5    ,			pad_token_id=1    ,			vocab_size=1000    ,			)
       lowercase__:   str					=      CLIPTextModel(_UpperCAmelCase							)
       lowercase__:   str					=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

       lowercase__:   str					=      MultiControlNetModel([controlneta, controlneta]							)

       lowercase__:   Union[str, Any]					=      {
           '''unet''': unet,
           '''controlnet''': controlnet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=0							):
       if str(_UpperCAmelCase							).startswith('''mps'''							):
          lowercase__:   Union[str, Any]					=      torch.manual_seed(_UpperCAmelCase							)
       else:
          lowercase__:   Union[str, Any]					=      torch.Generator(device=_UpperCAmelCase							).manual_seed(_UpperCAmelCase							)

       lowercase__:   List[Any]					=      2

       lowercase__:   Tuple					=      [
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)    ,			generator=_UpperCAmelCase    ,			device=torch.device(_UpperCAmelCase							)    ,			),
           randn_tensor(
               (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)    ,			generator=_UpperCAmelCase    ,			device=torch.device(_UpperCAmelCase							)    ,			),
       ]

       lowercase__:   int					=      floats_tensor(control_image[0].shape    ,			rng=random.Random(_UpperCAmelCase							)							).to(_UpperCAmelCase							)
       lowercase__:   List[Any]					=      image.cpu().permute(0    ,			2    ,			3    ,			1							)[0]
       lowercase__:   Union[str, Any]					=      Image.fromarray(np.uinta(_UpperCAmelCase							)							).convert('''RGB'''							).resize((64, 64)							)
       lowercase__:   int					=      {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
           '''image''': image,
           '''control_image''': control_image,
       }

       return inputs






    def _snake_case			(   self							):
       lowercase__:   Optional[Any]					=      self.get_dummy_components()
       lowercase__:   str					=      self.pipeline_class(**_UpperCAmelCase							)
       pipe.to(_UpperCAmelCase							)

       lowercase__:   Optional[Any]					=      10.0
       lowercase__:   Dict					=      4

       lowercase__:   Optional[int]					=      self.get_dummy_inputs(_UpperCAmelCase							)
       lowercase__:   int					=      steps
       lowercase__:   Dict					=      scale
       lowercase__:   Tuple					=      pipe(**_UpperCAmelCase							)[0]

       lowercase__:   Dict					=      self.get_dummy_inputs(_UpperCAmelCase							)
       lowercase__:   Optional[Any]					=      steps
       lowercase__:   int					=      scale
       lowercase__:   Tuple					=      pipe(**_UpperCAmelCase    ,			control_guidance_start=0.1    ,			control_guidance_end=0.2							)[0]

       lowercase__:   Optional[int]					=      self.get_dummy_inputs(_UpperCAmelCase							)
       lowercase__:   Optional[int]					=      steps
       lowercase__:   List[Any]					=      scale
       lowercase__:   Tuple					=      pipe(**_UpperCAmelCase    ,			control_guidance_start=[0.1, 0.3]    ,			control_guidance_end=[0.2, 0.7]							)[0]

       lowercase__:   Tuple					=      self.get_dummy_inputs(_UpperCAmelCase							)
       lowercase__:   Tuple					=      steps
       lowercase__:   List[Any]					=      scale
       lowercase__:   Union[str, Any]					=      pipe(**_UpperCAmelCase    ,			control_guidance_start=0.4    ,			control_guidance_end=[0.5, 0.8]							)[0]

       # make sure that all outputs are different
       assert np.sum(np.abs(output_a - output_a							)							) > 1e-3
       assert np.sum(np.abs(output_a - output_a							)							) > 1e-3
       assert np.sum(np.abs(output_a - output_a							)							) > 1e-3






    def _snake_case			(   self							):
       return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3							)






    @unittest.skipIf(
        torch_device != '''cuda''' or not is_xformers_available()    ,			reason='''XFormers attention is only available with CUDA and `xformers` installed'''    ,			)
    def _snake_case			(   self							):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3							)






    def _snake_case			(   self							):
       self._test_inference_batch_single_identical(expected_max_diff=2e-3							)






    def _snake_case			(   self							):
       lowercase__:   Optional[int]					=      self.get_dummy_components()
       lowercase__:   Optional[Any]					=      self.pipeline_class(**_UpperCAmelCase							)
       pipe.to(_UpperCAmelCase							)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase							)
       with tempfile.TemporaryDirectory() as tmpdir:
          try:
             # save_pretrained is not implemented for Multi-ControlNet
             pipe.save_pretrained(_UpperCAmelCase							)
          except NotImplementedError:
             pass



@slow
@require_torch_gpu
class     UpperCAmelCase						(unittest.TestCase  ):



    """simple docstring"""

    def _snake_case			(   self							):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






    def _snake_case			(   self							):
       lowercase__:   List[str]					=      ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''							)

       lowercase__:   List[Any]					=      StableDiffusionControlNetImgaImgPipeline.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''    ,			safety_checker=_UpperCAmelCase    ,			controlnet=_UpperCAmelCase							)
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=_UpperCAmelCase							)

       lowercase__:   Union[str, Any]					=      torch.Generator(device='''cpu'''							).manual_seed(0							)
       lowercase__:   str					=      '''evil space-punk bird'''
       lowercase__:   List[Any]					=      load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''							).resize((512, 512)							)
       lowercase__:   str					=      load_image(
           '''https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'''							).resize((512, 512)							)

       lowercase__:   Any					=      pipe(
           _UpperCAmelCase    ,			_UpperCAmelCase    ,			control_image=_UpperCAmelCase    ,			generator=_UpperCAmelCase    ,			output_type='''np'''    ,			num_inference_steps=50    ,			strength=0.6    ,			)

       lowercase__:   int					=      output.images[0]

       assert image.shape == (512, 512, 3)

       lowercase__:   int					=      load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy'''							)

       assert np.abs(expected_image - image							).max() < 9e-2

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

__A			   = logging.get_logger(__name__)


@add_end_docstrings(_UpperCAmelCase  )
class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    def __init__(   self    ,			**_UpperCAmelCase							):
       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 _snake_case			(   self    ,			**_UpperCAmelCase							):
       lowercase__:   List[Any]					=      {}
       lowercase__:   List[Any]					=      {}
       lowercase__:   Dict					=      {}
       # preprocess args
       if "points_per_batch" in kwargs:
          lowercase__:   Dict					=      kwargs['''points_per_batch''']
       if "points_per_crop" in kwargs:
          lowercase__:   Any					=      kwargs['''points_per_crop''']
       if "crops_n_layers" in kwargs:
          lowercase__:   Union[str, Any]					=      kwargs['''crops_n_layers''']
       if "crop_overlap_ratio" in kwargs:
          lowercase__:   Optional[Any]					=      kwargs['''crop_overlap_ratio''']
       if "crop_n_points_downscale_factor" in kwargs:
          lowercase__:   Union[str, Any]					=      kwargs['''crop_n_points_downscale_factor''']
       # postprocess args
       if "pred_iou_thresh" in kwargs:
          lowercase__:   Any					=      kwargs['''pred_iou_thresh''']
       if "stability_score_offset" in kwargs:
          lowercase__:   Tuple					=      kwargs['''stability_score_offset''']
       if "mask_threshold" in kwargs:
          lowercase__:   List[str]					=      kwargs['''mask_threshold''']
       if "stability_score_thresh" in kwargs:
          lowercase__:   str					=      kwargs['''stability_score_thresh''']
       if "crops_nms_thresh" in kwargs:
          lowercase__:   List[str]					=      kwargs['''crops_nms_thresh''']
       if "output_rle_mask" in kwargs:
          lowercase__:   Dict					=      kwargs['''output_rle_mask''']
       if "output_bboxes_mask" in kwargs:
          lowercase__:   int					=      kwargs['''output_bboxes_mask''']
       return preprocess_kwargs, forward_params, postprocess_kwargs






    def __call__(   self    ,			_UpperCAmelCase    ,			*_UpperCAmelCase    ,			_UpperCAmelCase=None    ,			_UpperCAmelCase=None    ,			**_UpperCAmelCase							):
       return super().__call__(_UpperCAmelCase    ,			*_UpperCAmelCase    ,			num_workers=_UpperCAmelCase    ,			batch_size=_UpperCAmelCase    ,			**_UpperCAmelCase							)






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=64    ,			_UpperCAmelCase = 0    ,			_UpperCAmelCase = 512 / 1500    ,			_UpperCAmelCase = 32    ,			_UpperCAmelCase = 1    ,			):
       lowercase__:   Union[str, Any]					=      load_image(_UpperCAmelCase							)
       lowercase__:   Dict					=      self.image_processor.size['''longest_edge''']
       lowercase__, lowercase__, lowercase__, lowercase__:   Optional[Any]					=      self.image_processor.generate_crop_boxes(
           _UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							)
       lowercase__:   List[Any]					=      self.image_processor(images=_UpperCAmelCase    ,			return_tensors='''pt'''							)

       with self.device_placement():
          if self.framework == "pt":
             lowercase__:   Tuple					=      self.get_inference_context()
             with inference_context():
                lowercase__:   Optional[Any]					=      self._ensure_tensor_on_device(_UpperCAmelCase    ,			device=self.device							)
                lowercase__:   Any					=      self.model.get_image_embeddings(model_inputs.pop('''pixel_values'''							)							)
                lowercase__:   Tuple					=      image_embeddings

       lowercase__:   Optional[Any]					=      grid_points.shape[1]
       lowercase__:   Tuple					=      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							):
          lowercase__:   Dict					=      grid_points[:, i : i + points_per_batch, :, :]
          lowercase__:   int					=      input_labels[:, i : i + points_per_batch]
          lowercase__:   Any					=      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    ,			_UpperCAmelCase    ,			_UpperCAmelCase=0.88    ,			_UpperCAmelCase=0.95    ,			_UpperCAmelCase=0    ,			_UpperCAmelCase=1    ,			):
       lowercase__:   List[Any]					=      model_inputs.pop('''input_boxes'''							)
       lowercase__:   List[Any]					=      model_inputs.pop('''is_last'''							)
       lowercase__:   Any					=      model_inputs.pop('''original_sizes'''							).tolist()
       lowercase__:   Union[str, Any]					=      model_inputs.pop('''reshaped_input_sizes'''							).tolist()

       lowercase__:   List[Any]					=      self.model(**_UpperCAmelCase							)

       # post processing happens here in order to avoid CPU GPU copies of ALL the masks
       lowercase__:   int					=      model_outputs['''pred_masks''']
       lowercase__:   str					=      self.image_processor.post_process_masks(
           _UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			binarize=_UpperCAmelCase							)
       lowercase__:   str					=      model_outputs['''iou_scores''']
       lowercase__, lowercase__, lowercase__:   Optional[int]					=      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 _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=False    ,			_UpperCAmelCase=False    ,			_UpperCAmelCase=0.7    ,			):
       lowercase__:   int					=      []
       lowercase__:   str					=      []
       lowercase__:   List[Any]					=      []
       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'''							)							)

       lowercase__:   Any					=      torch.cat(_UpperCAmelCase							)
       lowercase__:   Dict					=      torch.cat(_UpperCAmelCase							)
       lowercase__, lowercase__, lowercase__, lowercase__:   Any					=      self.image_processor.post_process_for_mask_generation(
           _UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							)

       lowercase__:   Union[str, Any]					=      defaultdict(_UpperCAmelCase							)
       for output in model_outputs:
          for k, v in output.items():
             extra[k].append(_UpperCAmelCase							)

       lowercase__:   Any					=      {}
       if output_rle_mask:
          lowercase__:   Optional[Any]					=      rle_mask

       if output_bboxes_mask:
          lowercase__:   Optional[int]					=      bounding_boxes

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