

'''simple docstring'''



import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 				_snake_case					(   a_     , unittest.TestCase			):
  SCREAMING_SNAKE_CASE			:       Optional[int]				=     DDIMPipeline
  SCREAMING_SNAKE_CASE			:       Dict				=     UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  SCREAMING_SNAKE_CASE			:       List[Any]				=     PipelineTesterMixin.required_optional_params - {
      '''num_images_per_prompt''',
      '''latents''',
      '''callback''',
      '''callback_steps''',
  }
  SCREAMING_SNAKE_CASE			:       Optional[int]				=     UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
  SCREAMING_SNAKE_CASE			:       Any				=     False

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      torch.manual_seed(0	)
      lowerCAmelCase 						=							UNetaDModel(
          block_out_channels=(32, 64)      , layers_per_block=2      , sample_size=32      , in_channels=3      , out_channels=3      , down_block_types=('DownBlock2D', 'AttnDownBlock2D')      , up_block_types=('AttnUpBlock2D', 'UpBlock2D')      , )
      lowerCAmelCase 						=							DDIMScheduler()
      lowerCAmelCase 						=							{'unet': unet, 'scheduler': scheduler}
      return components

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE=0	):





      '''simple docstring'''

      if str(_SCREAMING_SNAKE_CASE	).startswith('mps'	):
          lowerCAmelCase 						=							torch.manual_seed(_SCREAMING_SNAKE_CASE	)
      else:
          lowerCAmelCase 						=							torch.Generator(device=_SCREAMING_SNAKE_CASE	).manual_seed(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							{
          'batch_size': 1,
          'generator': generator,
          'num_inference_steps': 2,
          'output_type': 'numpy',
      }
      return inputs

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							'cpu'

      lowerCAmelCase 						=							self.get_dummy_components()
      lowerCAmelCase 						=							self.pipeline_class(**_SCREAMING_SNAKE_CASE	)
      pipe.to(_SCREAMING_SNAKE_CASE	)
      pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							self.get_dummy_inputs(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							pipe(**_SCREAMING_SNAKE_CASE	).images
      lowerCAmelCase 						=							image[0, -3:, -3:, -1]

      self.assertEqual(image.shape      , (1, 32, 32, 3)	)
      lowerCAmelCase 						=							np.array(
          [1.000e00, 5.717e-01, 4.717e-01, 1.000e00, 0.000e00, 1.000e00, 3.000e-04, 0.000e00, 9.000e-04]	)
      lowerCAmelCase 						=							np.abs(image_slice.flatten() - expected_slice	).max()
      self.assertLessEqual(_SCREAMING_SNAKE_CASE      , 1e-3	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      super().test_save_load_local(expected_max_difference=3e-3	)

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      super().test_save_load_optional_components(expected_max_difference=3e-3	)


  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      super().test_inference_batch_single_identical(expected_max_diff=3e-3	)



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

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							'google/ddpm-cifar10-32'

      lowerCAmelCase 						=							UNetaDModel.from_pretrained(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							DDIMScheduler()

      lowerCAmelCase 						=							DDIMPipeline(unet=_SCREAMING_SNAKE_CASE      , scheduler=_SCREAMING_SNAKE_CASE	)
      ddim.to(_SCREAMING_SNAKE_CASE	)
      ddim.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							torch.manual_seed(0	)
      lowerCAmelCase 						=							ddim(generator=_SCREAMING_SNAKE_CASE      , eta=0.0      , output_type='numpy'	).images

      lowerCAmelCase 						=							image[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      lowerCAmelCase 						=							np.array([0.1_723, 0.1_617, 0.1_600, 0.1_626, 0.1_497, 0.1_513, 0.1_505, 0.1_442, 0.1_453]	)

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


  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase 						=							'google/ddpm-ema-bedroom-256'

      lowerCAmelCase 						=							UNetaDModel.from_pretrained(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							DDIMScheduler.from_pretrained(_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							DDIMPipeline(unet=_SCREAMING_SNAKE_CASE      , scheduler=_SCREAMING_SNAKE_CASE	)
      ddpm.to(_SCREAMING_SNAKE_CASE	)
      ddpm.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE	)

      lowerCAmelCase 						=							torch.manual_seed(0	)
      lowerCAmelCase 						=							ddpm(generator=_SCREAMING_SNAKE_CASE      , output_type='numpy'	).images

      lowerCAmelCase 						=							image[0, -3:, -3:, -1]

      assert image.shape == (1, 2_56, 2_56, 3)
      lowerCAmelCase 						=							np.array([0.0_060, 0.0_201, 0.0_344, 0.0_024, 0.0_018, 0.0_002, 0.0_022, 0.0_000, 0.0_069]	)

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








'''simple docstring'''



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

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



@add_end_docstrings(a_			)
class 				_snake_case					(   a_			):

  def __init__(				self      , *_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      super().__init__(*_SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	)
      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      , _SCREAMING_SNAKE_CASE=None	):





      '''simple docstring'''

      lowerCAmelCase 						=							{}
      if top_k is not None:
          lowerCAmelCase 						=							top_k
      return {}, {}, postprocess_params

  def __call__(				self      , _SCREAMING_SNAKE_CASE      , **_SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

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

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							load_image(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							self.image_processor(images=_SCREAMING_SNAKE_CASE      , return_tensors=self.framework	)
      return model_inputs

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      lowerCAmelCase 						=							self.model(**_SCREAMING_SNAKE_CASE	)
      return model_outputs


  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE=5	):





      '''simple docstring'''

      if top_k > self.model.config.num_labels:
          lowerCAmelCase 						=							self.model.config.num_labels

      if self.framework == "pt":
          lowerCAmelCase 						=							model_outputs.logits.softmax(-1	)[0]
          lowerCAmelCase						,   lowerCAmelCase 						=							probs.topk(_SCREAMING_SNAKE_CASE	)
      elif self.framework == "tf":
          lowerCAmelCase 						=							stable_softmax(model_outputs.logits      , axis=-1	)[0]
          lowerCAmelCase 						=							tf.math.top_k(_SCREAMING_SNAKE_CASE      , k=_SCREAMING_SNAKE_CASE	)
          lowerCAmelCase						,   lowerCAmelCase 						=							topk.values.numpy(), topk.indices.numpy()
      else:
          raise ValueError(F'Unsupported framework: {self.framework}'	)

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





