'''simple docstring'''





import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


lowerCAmelCase			:							Optional[Any]     =False




class 	__snake_case	( unittest.TestCase		):





      '''simple docstring'''






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


          """simple docstring"""


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



      @property
      def 	_SCREAMING_SNAKE_CASE    (					self     :				Dict)				->Any:


          """simple docstring"""


          return 12



      @property
      def 	_SCREAMING_SNAKE_CASE    (					self     :				Tuple)				->Dict:


          """simple docstring"""


          return 12



      @property
      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[int])				->List[Any]:


          """simple docstring"""


          return 32



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


          """simple docstring"""


          torch.manual_seed(0)
          _lowerCamelCase     :		List[Any] 						=						VQModel(
              block_out_channels=[32, 64]    ,  in_channels=3    ,  out_channels=3    ,  down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]    ,  up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]    ,  latent_channels=3    ,  num_vq_embeddings=self.num_embed    ,  vq_embed_dim=3    ,  )
          return model



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


          """simple docstring"""


          _lowerCamelCase     :		int 						=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")
          return tokenizer



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


          """simple docstring"""


          torch.manual_seed(0)
          _lowerCamelCase     :		Dict 						=						CLIPTextConfig(
              bos_token_id=0    ,  eos_token_id=2    ,  hidden_size=self.text_embedder_hidden_size    ,  intermediate_size=37    ,  layer_norm_eps=1E-0_5    ,  num_attention_heads=4    ,  num_hidden_layers=5    ,  pad_token_id=1    ,  vocab_size=1000    ,  )
          return CLIPTextModel(_lowercase)



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


          """simple docstring"""


          torch.manual_seed(0)

          _lowerCamelCase     :		Union[str, Any] 						=						12
          _lowerCamelCase     :		Tuple 						=						12

          _lowerCamelCase     :		Tuple 						=						{
              """attention_bias""": True,
              """cross_attention_dim""": 32,
              """attention_head_dim""": height * width,
              """num_attention_heads""": 1,
              """num_vector_embeds""": self.num_embed,
              """num_embeds_ada_norm""": self.num_embeds_ada_norm,
              """norm_num_groups""": 32,
              """sample_size""": width,
              """activation_fn""": """geglu-approximate""",
          }

          _lowerCamelCase     :		Optional[Any] 						=						TransformeraDModel(**_lowercase)
          return model



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


          """simple docstring"""


          _lowerCamelCase     :		str 						=						"""cpu"""

          _lowerCamelCase     :		List[str] 						=						self.dummy_vqvae
          _lowerCamelCase     :		Any 						=						self.dummy_text_encoder
          _lowerCamelCase     :		Tuple 						=						self.dummy_tokenizer
          _lowerCamelCase     :		int 						=						self.dummy_transformer
          _lowerCamelCase     :		int 						=						VQDiffusionScheduler(self.num_embed)
          _lowerCamelCase     :		Dict 						=						LearnedClassifierFreeSamplingEmbeddings(learnable=_lowercase)

          _lowerCamelCase     :		Optional[Any] 						=						VQDiffusionPipeline(
              vqvae=_lowercase    ,  text_encoder=_lowercase    ,  tokenizer=_lowercase    ,  transformer=_lowercase    ,  scheduler=_lowercase    ,  learned_classifier_free_sampling_embeddings=_lowercase    ,  )
          _lowerCamelCase     :		int 						=						pipe.to(_lowercase)
          pipe.set_progress_bar_config(disable=_lowercase)

          _lowerCamelCase     :		List[Any] 						=						"""teddy bear playing in the pool"""

          _lowerCamelCase     :		Dict 						=						torch.Generator(device=_lowercase).manual_seed(0)
          _lowerCamelCase     :		List[Any] 						=						pipe([prompt]    ,  generator=_lowercase    ,  num_inference_steps=2    ,  output_type="""np""")
          _lowerCamelCase     :		Optional[int] 						=						output.images

          _lowerCamelCase     :		List[Any] 						=						torch.Generator(device=_lowercase).manual_seed(0)
          _lowerCamelCase     :		Dict 						=						pipe(
              [prompt]    ,  generator=_lowercase    ,  output_type="""np"""    ,  return_dict=_lowercase    ,  num_inference_steps=2)[0]

          _lowerCamelCase     :		List[Any] 						=						image[0, -3:, -3:, -1]
          _lowerCamelCase     :		Any 						=						image_from_tuple[0, -3:, -3:, -1]

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

          _lowerCamelCase     :		Dict 						=						np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2])

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




      def 	_SCREAMING_SNAKE_CASE    (					self     :				Tuple)				->str:


          """simple docstring"""


          _lowerCamelCase     :		int 						=						"""cpu"""

          _lowerCamelCase     :		List[Any] 						=						self.dummy_vqvae
          _lowerCamelCase     :		Optional[int] 						=						self.dummy_text_encoder
          _lowerCamelCase     :		List[Any] 						=						self.dummy_tokenizer
          _lowerCamelCase     :		Union[str, Any] 						=						self.dummy_transformer
          _lowerCamelCase     :		str 						=						VQDiffusionScheduler(self.num_embed)
          _lowerCamelCase     :		List[Any] 						=						LearnedClassifierFreeSamplingEmbeddings(
              learnable=_lowercase    ,  hidden_size=self.text_embedder_hidden_size    ,  length=tokenizer.model_max_length)

          _lowerCamelCase     :		Union[str, Any] 						=						VQDiffusionPipeline(
              vqvae=_lowercase    ,  text_encoder=_lowercase    ,  tokenizer=_lowercase    ,  transformer=_lowercase    ,  scheduler=_lowercase    ,  learned_classifier_free_sampling_embeddings=_lowercase    ,  )
          _lowerCamelCase     :		Any 						=						pipe.to(_lowercase)
          pipe.set_progress_bar_config(disable=_lowercase)

          _lowerCamelCase     :		Tuple 						=						"""teddy bear playing in the pool"""

          _lowerCamelCase     :		str 						=						torch.Generator(device=_lowercase).manual_seed(0)
          _lowerCamelCase     :		Tuple 						=						pipe([prompt]    ,  generator=_lowercase    ,  num_inference_steps=2    ,  output_type="""np""")
          _lowerCamelCase     :		Dict 						=						output.images

          _lowerCamelCase     :		Union[str, Any] 						=						torch.Generator(device=_lowercase).manual_seed(0)
          _lowerCamelCase     :		Any 						=						pipe(
              [prompt]    ,  generator=_lowercase    ,  output_type="""np"""    ,  return_dict=_lowercase    ,  num_inference_steps=2)[0]

          _lowerCamelCase     :		Optional[Any] 						=						image[0, -3:, -3:, -1]
          _lowerCamelCase     :		int 						=						image_from_tuple[0, -3:, -3:, -1]

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

          _lowerCamelCase     :		int 						=						np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 2.0
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2





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





      '''simple docstring'''






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


          """simple docstring"""


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




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


          """simple docstring"""


          _lowerCamelCase     :		List[Any] 						=						load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
              """/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy""")

          _lowerCamelCase     :		str 						=						VQDiffusionPipeline.from_pretrained("""microsoft/vq-diffusion-ithq""")
          _lowerCamelCase     :		Optional[Any] 						=						pipeline.to(_lowercase)
          pipeline.set_progress_bar_config(disable=_lowercase)

          # requires GPU generator for gumbel softmax
          # don't use GPU generator in tests though
          _lowerCamelCase     :		Any 						=						torch.Generator(device=_lowercase).manual_seed(0)
          _lowerCamelCase     :		Optional[int] 						=						pipeline(
              """teddy bear playing in the pool"""    ,  num_images_per_prompt=1    ,  generator=_lowercase    ,  output_type="""np"""    ,  )

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

          assert image.shape == (256, 256, 3)
          assert np.abs(expected_image - image).max() < 2.0





from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
    from PIL import Image

    from ..image_utils import load_image

if is_tf_available():
    import tensorflow as tf

    from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
    from ..tf_utils import stable_softmax

if is_torch_available():
    from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

lowerCAmelCase			:							Optional[Any]     =logging.get_logger(__name__)




@add_end_docstrings(__lowerCAmelCase		)
class 	__snake_case	( __lowerCAmelCase		):





      '''simple docstring'''






      def __init__(					self     :				str    ,  *_UpperCamelCase     :				int    ,  **_UpperCamelCase     :				List[str])				->Tuple:


          """simple docstring"""


          super().__init__(*_UpperCamelCase    ,  **_UpperCamelCase)
          requires_backends(self    ,  """vision""")
          self.check_model_type(
              TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
              if self.framework == """tf"""
              else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Dict    ,  _UpperCamelCase     :				List[str]=None)				->Optional[int]:


          """simple docstring"""


          _lowerCamelCase     :		Optional[int] 						=						{}
          if top_k is not None:
              _lowerCamelCase     :		str 						=						top_k
          return {}, {}, postprocess_params



      def __call__(					self     :				Optional[int]    ,  _UpperCamelCase     :				Union[str, List[str], "Image.Image", List["Image.Image"]]    ,  **_UpperCamelCase     :				Optional[int])				->Dict:


          """simple docstring"""


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



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


          """simple docstring"""


          _lowerCamelCase     :		Tuple 						=						load_image(_UpperCamelCase)
          _lowerCamelCase     :		Any 						=						self.image_processor(images=_UpperCamelCase    ,  return_tensors=self.framework)
          return model_inputs



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


          """simple docstring"""


          _lowerCamelCase     :		Any 						=						self.model(**_UpperCamelCase)
          return model_outputs




      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any]    ,  _UpperCamelCase     :				List[str]    ,  _UpperCamelCase     :				List[str]=5)				->str:


          """simple docstring"""


          if top_k > self.model.config.num_labels:
              _lowerCamelCase     :		Union[str, Any] 						=						self.model.config.num_labels

          if self.framework == "pt":
              _lowerCamelCase     :		Optional[Any] 						=						model_outputs.logits.softmax(-1)[0]
              _lowerCamelCase						,							_lowerCamelCase     :		Dict 						=						probs.topk(_UpperCamelCase)
          elif self.framework == "tf":
              _lowerCamelCase     :		List[Any] 						=						stable_softmax(model_outputs.logits    ,  axis=-1)[0]
              _lowerCamelCase     :		List[Any] 						=						tf.math.top_k(_UpperCamelCase    ,  k=_UpperCamelCase)
              _lowerCamelCase						,							_lowerCamelCase     :		str 						=						topk.values.numpy(), topk.indices.numpy()
          else:
              raise ValueError(F"""Unsupported framework: {self.framework}""")

          _lowerCamelCase     :		str 						=						scores.tolist()
          _lowerCamelCase     :		str 						=						ids.tolist()
          return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(_UpperCamelCase    ,  _UpperCamelCase)]




