



'''simple docstring'''


from typing import List, Union

import numpy as np

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, logging
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, ChunkPipeline


_lowerCamelCase	:       Optional[int]      	=	logging.get_logger(__name__)

class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





    def  A       (  self      :					List[str]  ,						UpperCamelCase__      :					str ):


       """simple docstring"""







       if isinstance(__a  ,						__a ):
          UpperCamelCase					=  [label.strip() for label in labels.split(',' ) if label.strip()]
       return labels






    def __call__(  self      :					Optional[Any]  ,						UpperCamelCase__      :					int  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					Dict ):


       """simple docstring"""







       if len(__a ) == 0 or len(__a ) == 0:
          raise ValueError('You must include at least one label and at least one sequence.' )
       if hypothesis_template.format(labels[0] ) == hypothesis_template:
          raise ValueError(
              (
                  'The provided hypothesis_template "{}" was not able to be formatted with the target labels. '
                  'Make sure the passed template includes formatting syntax such as {{}} where the label should go.'
              ).format(__a ) )

       if isinstance(__a  ,						__a ):
          UpperCamelCase					=  [sequences]

       UpperCamelCase					=  []
       for sequence in sequences:
          sequence_pairs.extend([[sequence, hypothesis_template.format(__a )] for label in labels] )

       return sequence_pairs, sequences

@add_end_docstrings(_a		)
class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





    def __init__(  self      :					Optional[int]  ,						UpperCamelCase__      :					Union[str, Any]=ZeroShotClassificationArgumentHandler()  ,						*UpperCamelCase__      :					Optional[Any]  ,						**UpperCamelCase__      :					Dict ):


       """simple docstring"""







       UpperCamelCase					=  args_parser
       super().__init__(*__a  ,						**__a )
       if self.entailment_id == -1:
          logger.warning(
              'Failed to determine \'entailment\' label id from the label2id mapping in the model config. Setting to '
              '-1. Define a descriptive label2id mapping in the model config to ensure correct outputs.' )






    @property
    def  A       (  self      :					Dict ):


       """simple docstring"""







       for label, ind in self.model.config.labelaid.items():
          if label.lower().startswith('entail' ):
             return ind
       return -1






    def  A       (  self      :					int  ,						UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					Optional[Any]=True  ,						UpperCamelCase__      :					Union[str, Any]=True  ,						UpperCamelCase__      :					Optional[Any]=TruncationStrategy.ONLY_FIRST  ,						**UpperCamelCase__      :					str ):


       """simple docstring"""







       UpperCamelCase					=  self.framework
       if self.tokenizer.pad_token is None:
          # Override for tokenizers not supporting padding
          logger.error(
              'Tokenizer was not supporting padding necessary for zero-shot, attempting to use '
              ' `pad_token=eos_token`' )
          UpperCamelCase					=  self.tokenizer.eos_token
       try:
          UpperCamelCase					=  self.tokenizer(
              __a  ,						add_special_tokens=__a  ,						return_tensors=__a  ,						padding=__a  ,						truncation=__a  ,						)
       except Exception as e:
          if "too short" in str(__a ):
             # tokenizers might yell that we want to truncate
             # to a value that is not even reached by the input.
             # In that case we don't want to truncate.
             # It seems there's not a really better way to catch that
             # exception.

             UpperCamelCase					=  self.tokenizer(
                 __a  ,						add_special_tokens=__a  ,						return_tensors=__a  ,						padding=__a  ,						truncation=TruncationStrategy.DO_NOT_TRUNCATE  ,						)
          else:
             raise e

       return inputs






    def  A       (  self      :					int  ,						**UpperCamelCase__      :					Optional[int] ):


       """simple docstring"""







       if kwargs.get('multi_class'  ,						__a ) is not None:
          UpperCamelCase					=  kwargs['multi_class']
          logger.warning(
              'The `multi_class` argument has been deprecated and renamed to `multi_label`. '
              '`multi_class` will be removed in a future version of Transformers.' )
       UpperCamelCase					=  {}
       if "candidate_labels" in kwargs:
          UpperCamelCase					=  self._args_parser._parse_labels(kwargs['candidate_labels'] )
       if "hypothesis_template" in kwargs:
          UpperCamelCase					=  kwargs['hypothesis_template']

       UpperCamelCase					=  {}
       if "multi_label" in kwargs:
          UpperCamelCase					=  kwargs['multi_label']
       return preprocess_params, {}, postprocess_params






    def __call__(  self      :					Optional[Any]  ,						UpperCamelCase__      :					Union[str, Any]  ,						*UpperCamelCase__      :					Any  ,						**UpperCamelCase__      :					Optional[int]  ,						):


       """simple docstring"""







       if len(__a ) == 0:
          pass
       elif len(__a ) == 1 and "candidate_labels" not in kwargs:
          UpperCamelCase					=  args[0]
       else:
          raise ValueError(f"""Unable to understand extra arguments {args}""" )

       return super().__call__(__a  ,						**__a )






    def  A       (  self      :					Any  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					Union[str, Any]=None  ,						UpperCamelCase__      :					Dict="This example is {}." ):


       """simple docstring"""







       UpperCamelCase				,      UpperCamelCase					=  self._args_parser(__a  ,						__a  ,						__a )

       for i, (candidate_label, sequence_pair) in enumerate(zip(__a  ,						__a ) ):
          UpperCamelCase					=  self._parse_and_tokenize([sequence_pair] )

          yield {
              "candidate_label": candidate_label,
              "sequence": sequences[0],
              "is_last": i == len(__a ) - 1,
              **model_input,
          }






    def  A       (  self      :					int  ,						UpperCamelCase__      :					Optional[Any] ):


       """simple docstring"""







       UpperCamelCase					=  inputs['candidate_label']
       UpperCamelCase					=  inputs['sequence']
       UpperCamelCase					=  {k: inputs[k] for k in self.tokenizer.model_input_names}
       UpperCamelCase					=  self.model(**__a )

       UpperCamelCase					=  {
           'candidate_label': candidate_label,
           'sequence': sequence,
           'is_last': inputs['is_last'],
           **outputs,
       }
       return model_outputs






    def  A       (  self      :					Tuple  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					Optional[int]=False ):


       """simple docstring"""







       UpperCamelCase					=  [outputs['candidate_label'] for outputs in model_outputs]
       UpperCamelCase					=  [outputs['sequence'] for outputs in model_outputs]
       UpperCamelCase					=  np.concatenate([output['logits'].numpy() for output in model_outputs] )
       UpperCamelCase					=  logits.shape[0]
       UpperCamelCase					=  len(__a )
       UpperCamelCase					=  N // n
       UpperCamelCase					=  logits.reshape((num_sequences, n, -1) )

       if multi_label or len(__a ) == 1:
          # softmax over the entailment vs. contradiction dim for each label independently
          UpperCamelCase					=  self.entailment_id
          UpperCamelCase					=  -1 if entailment_id == 0 else 0
          UpperCamelCase					=  reshaped_outputs[..., [contradiction_id, entailment_id]]
          UpperCamelCase					=  np.exp(__a ) / np.exp(__a ).sum(-1  ,						keepdims=__a )
          UpperCamelCase					=  scores[..., 1]
       else:
          # softmax the "entailment" logits over all candidate labels
          UpperCamelCase					=  reshaped_outputs[..., self.entailment_id]
          UpperCamelCase					=  np.exp(__a ) / np.exp(__a ).sum(-1  ,						keepdims=__a )

       UpperCamelCase					=  list(reversed(scores[0].argsort() ) )
       return {
           "sequence": sequences[0],
           "labels": [candidate_labels[i] for i in top_inds],
           "scores": scores[0, top_inds].tolist(),
       }



"""simple docstring"""





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






@skip_mps
class     _UpperCAmelCase(    lowerCamelCase      ,     lowerCamelCase      ,     unittest.TestCase	):
		lowercase__						=      StableDiffusionPanoramaPipeline
		lowercase__						=      TEXT_TO_IMAGE_PARAMS
		lowercase__						=      TEXT_TO_IMAGE_BATCH_PARAMS
		lowercase__						=      TEXT_TO_IMAGE_IMAGE_PARAMS
		lowercase__						=      TEXT_TO_IMAGE_IMAGE_PARAMS



		def 						UpperCAmelCase							(					self)       ->       Dict:






							'''simple docstring'''
							torch.manual_seed(0)
							_UpperCamelCase							 =  UNetaDConditionModel(
							    block_out_channels=(32, 64)				,		layers_per_block=1				,		sample_size=32				,		in_channels=4				,		out_channels=4				,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,		cross_attention_dim=32				,		)
							_UpperCamelCase							 =  DDIMScheduler()
							torch.manual_seed(0)
							_UpperCamelCase							 =  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)
							_UpperCamelCase							 =  CLIPTextConfig(
							    bos_token_id=0				,		eos_token_id=2				,		hidden_size=32				,		intermediate_size=37				,		layer_norm_eps=1e-05				,		num_attention_heads=4				,		num_hidden_layers=5				,		pad_token_id=1				,		vocab_size=10_00				,		)
							_UpperCamelCase							 =  CLIPTextModel(__a)
							_UpperCamelCase							 =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

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



		def 						UpperCAmelCase							(					self				,		__a				,		__a=0)       ->       int:






							'''simple docstring'''
							_UpperCamelCase							 =  torch.manual_seed(__a)
							_UpperCamelCase							 =  {
							    '''prompt''': '''a photo of the dolomites''',
							    '''generator''': generator,
							    # Setting height and width to None to prevent OOMs on CPU.
							    '''height''': None,
							    '''width''': None,
							    '''num_inference_steps''': 1,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}
							return inputs



		def 						UpperCAmelCase							(					self)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  sd_pipe(**__a).images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							_UpperCamelCase							 =  np.array([0.6186, 0.5374, 0.4915, 0.4135, 0.4114, 0.4563, 0.5128, 0.4977, 0.4757])

							assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2



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






							'''simple docstring'''
							super().test_inference_batch_consistent(batch_sizes=[1, 2])



		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							super().test_inference_batch_single_identical(batch_size=2				,		expected_max_diff=3.25e-3)



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






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  '''french fries'''
							_UpperCamelCase							 =  sd_pipe(**__a				,		negative_prompt=__a)
							_UpperCamelCase							 =  output.images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]

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

							_UpperCamelCase							 =  np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757])

							assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2



		def 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  sd_pipe(**__a				,		view_batch_size=2)
							_UpperCamelCase							 =  output.images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]

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

							_UpperCamelCase							 =  np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757])

							assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2



		def 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  EulerAncestralDiscreteScheduler(
							    beta_start=0.0_0085				,		beta_end=0.012				,		beta_schedule='''scaled_linear''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  sd_pipe(**__a).images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]

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

							_UpperCamelCase							 =  np.array([0.4024, 0.6510, 0.4901, 0.5378, 0.5813, 0.5622, 0.4795, 0.4467, 0.4952])

							assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2




		def 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  PNDMScheduler(
							    beta_start=0.0_0085				,		beta_end=0.012				,		beta_schedule='''scaled_linear'''				,		skip_prk_steps=__a)
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline(**__a)
							_UpperCamelCase							 =  sd_pipe.to(__a)
							sd_pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  sd_pipe(**__a).images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1]

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

							_UpperCamelCase							 =  np.array([0.6391, 0.6291, 0.4861, 0.5134, 0.5552, 0.4578, 0.5032, 0.5023, 0.4539])

							assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2






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



		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



		def 						UpperCAmelCase							(					self				,		__a=0)       ->       List[str]:






							'''simple docstring'''
							_UpperCamelCase							 =  torch.manual_seed(__a)
							_UpperCamelCase							 =  {
							    '''prompt''': '''a photo of the dolomites''',
							    '''generator''': generator,
							    '''num_inference_steps''': 3,
							    '''guidance_scale''': 7.5,
							    '''output_type''': '''numpy''',
							}
							return inputs



		def 						UpperCAmelCase							(					self)       ->       Any:






							'''simple docstring'''
							_UpperCamelCase							 =  '''stabilityai/stable-diffusion-2-base'''
							_UpperCamelCase							 =  DDIMScheduler.from_pretrained(__a				,		subfolder='''scheduler''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(__a				,		scheduler=__a				,		safety_checker=__a)
							pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing()

							_UpperCamelCase							 =  self.get_inputs()
							_UpperCamelCase							 =  pipe(**__a).images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 5_12, 20_48, 3)

							_UpperCamelCase							 =  np.array(
							    [
							        0.3696_8392,
							        0.2702_5372,
							        0.3244_6766,
							        0.2837_9387,
							        0.3636_3274,
							        0.3073_3347,
							        0.2710_0027,
							        0.2705_4125,
							        0.2553_6096,
							    ])

							assert np.abs(expected_slice - image_slice).max() < 1e-2



		def 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-base'''				,		safety_checker=__a)
							_UpperCamelCase							 =  LMSDiscreteScheduler.from_config(pipe.scheduler.config)
							pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing()

							_UpperCamelCase							 =  self.get_inputs()
							_UpperCamelCase							 =  pipe(**__a).images
							_UpperCamelCase							 =  image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 5_12, 20_48, 3)

							_UpperCamelCase							 =  np.array(
							    [
							        [
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							            0.0,
							        ]
							    ])

							assert np.abs(expected_slice - image_slice).max() < 1e-3



		def 						UpperCAmelCase							(					self)       ->       Dict:






							'''simple docstring'''
							_UpperCamelCase							 =  0

							def callback_fn(__a				,		__a				,		__a) -> None:
												_UpperCamelCase							 =  True
												nonlocal number_of_steps
												number_of_steps += 1
												if step == 1:
																	_UpperCamelCase							 =  latents.detach().cpu().numpy()
																	assert latents.shape == (1, 4, 64, 2_56)
																	_UpperCamelCase							 =  latents[0, -3:, -3:, -1]

																	_UpperCamelCase							 =  np.array(
																	    [
																	        0.1868_1869,
																	        0.3390_7816,
																	        0.536_1276,
																	        0.1443_2865,
																	        -0.0285_6611,
																	        -0.7394_1123,
																	        0.2339_7987,
																	        0.4732_2682,
																	        -0.3782_3164,
																	    ])
																	assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2
												elif step == 2:
																	_UpperCamelCase							 =  latents.detach().cpu().numpy()
																	assert latents.shape == (1, 4, 64, 2_56)
																	_UpperCamelCase							 =  latents[0, -3:, -3:, -1]

																	_UpperCamelCase							 =  np.array(
																	    [
																	        0.1853_9645,
																	        0.3398_7248,
																	        0.537_8559,
																	        0.1443_7142,
																	        -0.0245_5261,
																	        -0.733_8317,
																	        0.2399_0755,
																	        0.4735_6272,
																	        -0.378_6505,
																	    ])

																	assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2

							_UpperCamelCase							 =  False

							_UpperCamelCase							 =  '''stabilityai/stable-diffusion-2-base'''
							_UpperCamelCase							 =  DDIMScheduler.from_pretrained(__a				,		subfolder='''scheduler''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(__a				,		scheduler=__a				,		safety_checker=__a)
							_UpperCamelCase							 =  pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing()

							_UpperCamelCase							 =  self.get_inputs()
							pipe(**__a				,		callback=__a				,		callback_steps=1)
							assert callback_fn.has_been_called
							assert number_of_steps == 3




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






							'''simple docstring'''
							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							_UpperCamelCase							 =  '''stabilityai/stable-diffusion-2-base'''
							_UpperCamelCase							 =  DDIMScheduler.from_pretrained(__a				,		subfolder='''scheduler''')
							_UpperCamelCase							 =  StableDiffusionPanoramaPipeline.from_pretrained(__a				,		scheduler=__a				,		safety_checker=__a)
							_UpperCamelCase							 =  pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)
							pipe.enable_attention_slicing(1)
							pipe.enable_sequential_cpu_offload()

							_UpperCamelCase							 =  self.get_inputs()
							_UpperCamelCase							 =  pipe(**__a)

							_UpperCamelCase							 =  torch.cuda.max_memory_allocated()
							# make sure that less than 5.2 GB is allocated
							assert mem_bytes < 5.5 * 10**9






