


'''simple docstring'''



from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
      from PIL import Image

      from ..image_utils import load_image

if is_torch_available():
      import torch

      from transformers.modeling_outputs import BaseModelOutput

      from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

__UpperCAmelCase      :Dict			    =       logging.get_logger(__name__)



@add_end_docstrings(_a	)
class      a    (		_a	):


     """simple docstring"""






     def __init__(    self					:		List[str]							,			**snake_case					:		Dict      )						->							Optional[int]:
          super().__init__(**snake_case      )

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

          requires_backends(self							,			'''vision'''      )
          self.check_model_type(snake_case      )

     def __call__(    self					:		List[str]							,			snake_case					:		Union[str, "Image.Image", List[Dict[str, Any]]]							,			snake_case					:		Union[str, List[str]] = None							,			**snake_case					:		List[str]							,			)						->							Optional[Any]:
          if "text_queries" in kwargs:
               __UpperCAmelCase     :			str  							=      kwargs.pop('''text_queries'''      )

          if isinstance(snake_case							,			(str, Image.Image)      ):
               __UpperCAmelCase     :			Optional[Any]  							=      {'''image''': image, '''candidate_labels''': candidate_labels}
          else:
               __UpperCAmelCase     :			List[str]  							=      image
          __UpperCAmelCase     :			Optional[Any]  							=      super().__call__(snake_case							,			**snake_case      )
          return results

     def 			lowerCamelCase__			(    self					:		Optional[int]							,			**snake_case					:		Optional[int]      )						->							int:
          __UpperCAmelCase     :			Tuple  							=      {}
          if "threshold" in kwargs:
               __UpperCAmelCase     :			Optional[int]  							=      kwargs['''threshold''']
          if "top_k" in kwargs:
               __UpperCAmelCase     :			Optional[int]  							=      kwargs['''top_k''']
          return {}, {}, postprocess_params

     def 			lowerCamelCase__			(    self					:		Optional[int]							,			snake_case					:		Union[str, Any]      )						->							int:
          __UpperCAmelCase     :			Any  							=      load_image(inputs['''image''']      )
          __UpperCAmelCase     :			int  							=      inputs['''candidate_labels''']
          if isinstance(snake_case							,			snake_case      ):
               __UpperCAmelCase     :			str  							=      candidate_labels.split(''','''      )

          __UpperCAmelCase     :			Union[str, Any]  							=      torch.tensor([[image.height, image.width]]							,			dtype=torch.intaa      )
          for i, candidate_label in enumerate(snake_case      ):
               __UpperCAmelCase     :			Union[str, Any]  							=      self.tokenizer(snake_case							,			return_tensors=self.framework      )
               __UpperCAmelCase     :			int  							=      self.image_processor(snake_case							,			return_tensors=self.framework      )
               yield {
                   "is_last": i == len(snake_case      ) - 1,
                   "target_size": target_size,
                   "candidate_label": candidate_label,
                   **text_inputs,
                   **image_features,
               }

     def 			lowerCamelCase__			(    self					:		Union[str, Any]							,			snake_case					:		List[str]      )						->							List[str]:
          __UpperCAmelCase     :			Optional[int]  							=      model_inputs.pop('''target_size'''      )
          __UpperCAmelCase     :			Any  							=      model_inputs.pop('''candidate_label'''      )
          __UpperCAmelCase     :			List[str]  							=      model_inputs.pop('''is_last'''      )

          __UpperCAmelCase     :			Optional[Any]  							=      self.model(**snake_case      )

          __UpperCAmelCase     :			str  							=      {'''target_size''': target_size, '''candidate_label''': candidate_label, '''is_last''': is_last, **outputs}
          return model_outputs

     def 			lowerCamelCase__			(    self					:		Any							,			snake_case					:		Optional[Any]							,			snake_case					:		List[Any]=0.1							,			snake_case					:		str=None      )						->							Dict:
          __UpperCAmelCase     :			Any  							=      []
          for model_output in model_outputs:
               __UpperCAmelCase     :			List[Any]  							=      model_output['''candidate_label''']
               __UpperCAmelCase     :			Optional[Any]  							=      BaseModelOutput(snake_case      )
               __UpperCAmelCase     :			Optional[Any]  							=      self.image_processor.post_process_object_detection(
                   outputs=snake_case							,			threshold=snake_case							,			target_sizes=model_output['''target_size''']      )[0]

               for index in outputs["scores"].nonzero():
                    __UpperCAmelCase     :			Optional[int]  							=      outputs['''scores'''][index].item()
                    __UpperCAmelCase     :			Optional[int]  							=      self._get_bounding_box(outputs['''boxes'''][index][0]      )

                    __UpperCAmelCase     :			List[str]  							=      {'''score''': score, '''label''': label, '''box''': box}
                    results.append(snake_case      )

          __UpperCAmelCase     :			Any  							=      sorted(snake_case							,			key=lambda snake_case      : x["score"]							,			reverse=snake_case      )
          if top_k:
               __UpperCAmelCase     :			str  							=      results[:top_k]

          return results






     def 			lowerCamelCase__			(    self					:		int							,			snake_case					:		"torch.Tensor"      )						->							Dict[str, int]:
          if self.framework != "pt":
               raise ValueError('''The ZeroShotObjectDetectionPipeline is only available in PyTorch.'''      )
          __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase							,       __UpperCAmelCase     :			Dict  							=      box.int().tolist()
          __UpperCAmelCase     :			List[Any]  							=      {
              '''xmin''': xmin,
              '''ymin''': ymin,
              '''xmax''': xmax,
              '''ymax''': ymax,
          }
          return bbox



'''simple docstring'''



import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin



@skip_mps
class      a    (		_a     ,  _a     ,  unittest.TestCase	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			Dict							       =       IFPipeline
     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       TEXT_TO_IMAGE_BATCH_PARAMS
     SCREAMING_SNAKE_CASE	:			str							       =       PipelineTesterMixin.required_optional_params - {"latents"}

     def 			lowerCamelCase__			(    self					:		Tuple      )						->							Optional[Any]:
          return self._get_dummy_components()

     def 			lowerCamelCase__			(    self					:		Optional[Any]							,			snake_case					:		Optional[int]							,			snake_case					:		List[Any]=0      )						->							Optional[Any]:
          if str(snake_case      ).startswith('''mps'''      ):
               __UpperCAmelCase     :			Optional[Any]  							=      torch.manual_seed(snake_case      )
          else:
               __UpperCAmelCase     :			Dict  							=      torch.Generator(device=snake_case      ).manual_seed(snake_case      )

          __UpperCAmelCase     :			str  							=      {
              '''prompt''': '''A painting of a squirrel eating a burger''',
              '''generator''': generator,
              '''num_inference_steps''': 2,
              '''output_type''': '''numpy''',
          }

          return inputs

     def 			lowerCamelCase__			(    self					:		List[str]      )						->							Optional[int]:
          self._test_save_load_optional_components()

     @unittest.skipIf(torch_device != '''cuda'''							,			reason='''float16 requires CUDA'''      )
     def 			lowerCamelCase__			(    self					:		Dict      )						->							Union[str, 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					:		List[str]      )						->							Tuple:
          self._test_attention_slicing_forward_pass(expected_max_diff=1E-2      )

     def 			lowerCamelCase__			(    self					:		Any      )						->							List[str]:
          self._test_save_load_local()

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






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



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


     """simple docstring"""






     def 			lowerCamelCase__			(    self					:		Dict      )						->							Dict:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()

     def 			lowerCamelCase__			(    self					:		Any      )						->							Tuple:
          # if

          __UpperCAmelCase     :			Union[str, Any]  							=      IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''							,			variant='''fp16'''							,			torch_dtype=torch.floataa      )

          __UpperCAmelCase     :			List[str]  							=      IFSuperResolutionPipeline.from_pretrained(
              '''DeepFloyd/IF-II-L-v1.0'''							,			variant='''fp16'''							,			torch_dtype=torch.floataa							,			text_encoder=snake_case							,			tokenizer=snake_case      )

          # pre compute text embeddings and remove T5 to save memory

          pipe_a.text_encoder.to('''cuda'''      )

          __UpperCAmelCase							,       __UpperCAmelCase     :			Tuple  							=      pipe_a.encode_prompt('''anime turtle'''							,			device='''cuda'''      )

          del pipe_a.tokenizer
          del pipe_a.text_encoder
          gc.collect()

          __UpperCAmelCase     :			Any  							=      None
          __UpperCAmelCase     :			Optional[Any]  							=      None

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

          self._test_if(snake_case							,			snake_case							,			snake_case							,			snake_case      )

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # img2img

          __UpperCAmelCase     :			Union[str, Any]  							=      IFImgaImgPipeline(**pipe_a.components      )
          __UpperCAmelCase     :			int  							=      IFImgaImgSuperResolutionPipeline(**pipe_a.components      )

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

          self._test_if_imgaimg(snake_case							,			snake_case							,			snake_case							,			snake_case      )

          pipe_a.remove_all_hooks()
          pipe_a.remove_all_hooks()

          # inpainting

          __UpperCAmelCase     :			List[str]  							=      IFInpaintingPipeline(**pipe_a.components      )
          __UpperCAmelCase     :			List[str]  							=      IFInpaintingSuperResolutionPipeline(**pipe_a.components      )

          pipe_a.enable_model_cpu_offload()
          pipe_a.enable_model_cpu_offload()

          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )
          pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()      )

          self._test_if_inpainting(snake_case							,			snake_case							,			snake_case							,			snake_case      )

     def 			lowerCamelCase__			(    self					:		List[str]							,			snake_case					:		Any							,			snake_case					:		Union[str, Any]							,			snake_case					:		Optional[int]							,			snake_case					:		str      )						->							Optional[int]:
          # pipeline 1

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			Dict  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )
          __UpperCAmelCase     :			List[str]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			num_inference_steps=2							,			generator=snake_case							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			List[Any]  							=      output.images[0]

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

          __UpperCAmelCase     :			List[str]  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 13 * 10**9

          __UpperCAmelCase     :			Optional[Any]  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

          # pipeline 2

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			int  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

          __UpperCAmelCase     :			Optional[int]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )

          __UpperCAmelCase     :			List[Any]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			generator=snake_case							,			num_inference_steps=2							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			Optional[Any]  							=      output.images[0]

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

          __UpperCAmelCase     :			List[str]  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          __UpperCAmelCase     :			Tuple  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

     def 			lowerCamelCase__			(    self					:		Optional[int]							,			snake_case					:		Optional[int]							,			snake_case					:		Union[str, Any]							,			snake_case					:		str							,			snake_case					:		Dict      )						->							str:
          # pipeline 1

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			Optional[Any]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )

          __UpperCAmelCase     :			Dict  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

          __UpperCAmelCase     :			Dict  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			num_inference_steps=2							,			generator=snake_case							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			int  							=      output.images[0]

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

          __UpperCAmelCase     :			Any  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 10 * 10**9

          __UpperCAmelCase     :			Optional[Any]  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

          # pipeline 2

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			Optional[int]  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

          __UpperCAmelCase     :			Optional[Any]  							=      floats_tensor((1, 3, 256, 256)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			List[str]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )

          __UpperCAmelCase     :			List[str]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			original_image=snake_case							,			generator=snake_case							,			num_inference_steps=2							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			int  							=      output.images[0]

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

          __UpperCAmelCase     :			int  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          __UpperCAmelCase     :			Any  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )






     def 			lowerCamelCase__			(    self					:		str							,			snake_case					:		Union[str, Any]							,			snake_case					:		Optional[int]							,			snake_case					:		Optional[Any]							,			snake_case					:		Optional[Any]      )						->							Optional[int]:
          # pipeline 1

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			List[Any]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			List[str]  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(1      )      ).to(snake_case      )

          __UpperCAmelCase     :			Optional[Any]  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )
          __UpperCAmelCase     :			Dict  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			mask_image=snake_case							,			num_inference_steps=2							,			generator=snake_case							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			Union[str, Any]  							=      output.images[0]

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

          __UpperCAmelCase     :			Tuple  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 10 * 10**9

          __UpperCAmelCase     :			Tuple  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )

          # pipeline 2

          _start_torch_memory_measurement()

          __UpperCAmelCase     :			int  							=      torch.Generator(device='''cpu'''      ).manual_seed(0      )

          __UpperCAmelCase     :			Tuple  							=      floats_tensor((1, 3, 64, 64)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			Any  							=      floats_tensor((1, 3, 256, 256)							,			rng=random.Random(0      )      ).to(snake_case      )
          __UpperCAmelCase     :			Tuple  							=      floats_tensor((1, 3, 256, 256)							,			rng=random.Random(1      )      ).to(snake_case      )

          __UpperCAmelCase     :			Union[str, Any]  							=      pipe_a(
              prompt_embeds=snake_case							,			negative_prompt_embeds=snake_case							,			image=snake_case							,			mask_image=snake_case							,			original_image=snake_case							,			generator=snake_case							,			num_inference_steps=2							,			output_type='''np'''							,			)

          __UpperCAmelCase     :			List[Any]  							=      output.images[0]

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

          __UpperCAmelCase     :			Any  							=      torch.cuda.max_memory_allocated()
          assert mem_bytes < 4 * 10**9

          __UpperCAmelCase     :			str  							=      load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'''      )
          assert_mean_pixel_difference(snake_case							,			snake_case      )




def      _a     (						):





     '''simple docstring'''






     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()