





"""simple docstring"""


import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

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


enable_full_determinism()



class      __SCREAMING_SNAKE_CASE    (						_UpperCamelCase		,  unittest.TestCase   ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__							:str	     =					ConsistencyModelPipeline
  SCREAMING_SNAKE_CASE__							:Optional[Any]	     =					UNCONDITIONAL_IMAGE_GENERATION_PARAMS
  SCREAMING_SNAKE_CASE__							:Dict	     =					UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

  # Override required_optional_params to remove num_images_per_prompt
  SCREAMING_SNAKE_CASE__							:List[str]	     =					frozenset(
      [
          "num_inference_steps",
          "generator",
          "latents",
          "output_type",
          "return_dict",
          "callback",
          "callback_steps",
      ]   )

  @property
  def    __SCREAMING_SNAKE_CASE				(		self   :       List[Any]	)		->				Optional[int]:
        _UpperCamelCase      :	str						=       UNetaDModel.from_pretrained(
            "diffusers/consistency-models-test"					,   subfolder="test_unet"					,   )
        return unet

  @property
  def    __SCREAMING_SNAKE_CASE				(		self   :       int	)		->				Any:
        _UpperCamelCase      :	Dict						=       UNetaDModel.from_pretrained(
            "diffusers/consistency-models-test"					,   subfolder="test_unet_class_cond"					,   )
        return unet

  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple					,   __a   :       Optional[Any]=False	)		->				Tuple:
        if class_cond:
              _UpperCamelCase      :	Union[str, Any]						=       self.dummy_cond_unet
        else:
              _UpperCamelCase      :	Any						=       self.dummy_uncond_unet

        # Default to CM multistep sampler
        _UpperCamelCase      :	List[Any]						=       CMStochasticIterativeScheduler(
            num_train_timesteps=40					,   sigma_min=0.0_02					,   sigma_max=80.0					,   )

        _UpperCamelCase      :	Optional[int]						=       {
            "unet": unet,
            "scheduler": scheduler,
        }

        return components

  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple					,   __a   :       Any					,   __a   :       Any=0	)		->				Optional[int]:
        if str(__a	).startswith("mps"	):
              _UpperCamelCase      :	Optional[int]						=       torch.manual_seed(__a	)
        else:
              _UpperCamelCase      :	List[str]						=       torch.Generator(device=__a	).manual_seed(__a	)

        _UpperCamelCase      :	Any						=       {
            "batch_size": 1,
            "num_inference_steps": None,
            "timesteps": [22, 0],
            "generator": generator,
            "output_type": "np",
        }

        return inputs

  def    __SCREAMING_SNAKE_CASE				(		self   :       str	)		->				int:
        _UpperCamelCase      :	Union[str, Any]						=       "cpu"  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase      :	Tuple						=       self.get_dummy_components()
        _UpperCamelCase      :	Optional[int]						=       ConsistencyModelPipeline(**__a	)
        _UpperCamelCase      :	int						=       pipe.to(__a	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	Dict						=       self.get_dummy_inputs(__a	)
        _UpperCamelCase      :	Optional[Any]						=       pipe(**__a	).images
        assert image.shape == (1, 32, 32, 3)

        _UpperCamelCase      :	Optional[int]						=       image[0, -3:, -3:, -1]
        _UpperCamelCase      :	List[str]						=       np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]	)

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Any	)		->				Union[str, Any]:
        _UpperCamelCase      :	Optional[int]						=       "cpu"  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase      :	Union[str, Any]						=       self.get_dummy_components(class_cond=__a	)
        _UpperCamelCase      :	List[Any]						=       ConsistencyModelPipeline(**__a	)
        _UpperCamelCase      :	Tuple						=       pipe.to(__a	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	Tuple						=       self.get_dummy_inputs(__a	)
        _UpperCamelCase      :	Tuple						=       0
        _UpperCamelCase      :	List[str]						=       pipe(**__a	).images
        assert image.shape == (1, 32, 32, 3)

        _UpperCamelCase      :	Union[str, Any]						=       image[0, -3:, -3:, -1]
        _UpperCamelCase      :	Optional[int]						=       np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]	)

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple	)		->				List[Any]:
        _UpperCamelCase      :	Tuple						=       "cpu"  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase      :	Any						=       self.get_dummy_components()
        _UpperCamelCase      :	List[Any]						=       ConsistencyModelPipeline(**__a	)
        _UpperCamelCase      :	Optional[Any]						=       pipe.to(__a	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	Union[str, Any]						=       self.get_dummy_inputs(__a	)
        _UpperCamelCase      :	Union[str, Any]						=       1
        _UpperCamelCase      :	Union[str, Any]						=       None
        _UpperCamelCase      :	Optional[int]						=       pipe(**__a	).images
        assert image.shape == (1, 32, 32, 3)

        _UpperCamelCase      :	str						=       image[0, -3:, -3:, -1]
        _UpperCamelCase      :	Optional[Any]						=       np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]	)

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


  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]	)		->				Union[str, Any]:
        _UpperCamelCase      :	Tuple						=       "cpu"  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase      :	Optional[Any]						=       self.get_dummy_components(class_cond=__a	)
        _UpperCamelCase      :	List[str]						=       ConsistencyModelPipeline(**__a	)
        _UpperCamelCase      :	List[str]						=       pipe.to(__a	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	Union[str, Any]						=       self.get_dummy_inputs(__a	)
        _UpperCamelCase      :	List[Any]						=       1
        _UpperCamelCase      :	List[str]						=       None
        _UpperCamelCase      :	int						=       0
        _UpperCamelCase      :	int						=       pipe(**__a	).images
        assert image.shape == (1, 32, 32, 3)

        _UpperCamelCase      :	Tuple						=       image[0, -3:, -3:, -1]
        _UpperCamelCase      :	Any						=       np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]	)

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





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







  '''simple docstring'''


  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[Any]	)		->				int:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict					,   __a   :       List[str]=0					,   __a   :       Optional[Any]=False					,   __a   :       Optional[Any]="cpu"					,   __a   :       int=torch.floataa					,   __a   :       Union[str, Any]=(1, 3, 64, 64)	)		->				Tuple:
        _UpperCamelCase      :	Optional[int]						=       torch.manual_seed(__a	)

        _UpperCamelCase      :	Optional[Any]						=       {
            "num_inference_steps": None,
            "timesteps": [22, 0],
            "class_labels": 0,
            "generator": generator,
            "output_type": "np",
        }

        if get_fixed_latents:
              _UpperCamelCase      :	List[str]						=       self.get_fixed_latents(seed=__a					,   device=__a					,   dtype=__a					,   shape=__a	)
              _UpperCamelCase      :	List[Any]						=       latents

        return inputs

  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple					,   __a   :       str=0					,   __a   :       int="cpu"					,   __a   :       List[str]=torch.floataa					,   __a   :       Tuple=(1, 3, 64, 64)	)		->				Dict:
        if type(__a	) == str:
              _UpperCamelCase      :	int						=       torch.device(__a	)
        _UpperCamelCase      :	Optional[int]						=       torch.Generator(device=__a	).manual_seed(__a	)
        _UpperCamelCase      :	List[Any]						=       randn_tensor(__a					,   generator=__a					,   device=__a					,   dtype=__a	)
        return latents

  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]	)		->				Any:
        _UpperCamelCase      :	Optional[Any]						=       UNetaDModel.from_pretrained("diffusers/consistency_models"					,   subfolder="diffusers_cd_imagenet64_l2"	)
        _UpperCamelCase      :	List[Any]						=       CMStochasticIterativeScheduler(
            num_train_timesteps=40					,   sigma_min=0.0_02					,   sigma_max=80.0					,   )
        _UpperCamelCase      :	List[Any]						=       ConsistencyModelPipeline(unet=__a					,   scheduler=__a	)
        pipe.to(torch_device=__a	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	Any						=       self.get_inputs()
        _UpperCamelCase      :	List[str]						=       pipe(**__a	).images
        assert image.shape == (1, 64, 64, 3)

        _UpperCamelCase      :	int						=       image[0, -3:, -3:, -1]

        _UpperCamelCase      :	Tuple						=       np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]	)

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

  def    __SCREAMING_SNAKE_CASE				(		self   :       Union[str, Any]	)		->				Any:
        _UpperCamelCase      :	Optional[Any]						=       UNetaDModel.from_pretrained("diffusers/consistency_models"					,   subfolder="diffusers_cd_imagenet64_l2"	)
        _UpperCamelCase      :	Union[str, Any]						=       CMStochasticIterativeScheduler(
            num_train_timesteps=40					,   sigma_min=0.0_02					,   sigma_max=80.0					,   )
        _UpperCamelCase      :	List[Any]						=       ConsistencyModelPipeline(unet=__a					,   scheduler=__a	)
        pipe.to(torch_device=__a	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	List[str]						=       self.get_inputs()
        _UpperCamelCase      :	List[str]						=       1
        _UpperCamelCase      :	Optional[Any]						=       None
        _UpperCamelCase      :	Optional[int]						=       pipe(**__a	).images
        assert image.shape == (1, 64, 64, 3)

        _UpperCamelCase      :	Optional[Any]						=       image[0, -3:, -3:, -1]

        _UpperCamelCase      :	int						=       np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]	)

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

  @require_torch_a
  def    __SCREAMING_SNAKE_CASE				(		self   :       List[Any]	)		->				Any:
        _UpperCamelCase      :	List[str]						=       UNetaDModel.from_pretrained("diffusers/consistency_models"					,   subfolder="diffusers_cd_imagenet64_l2"	)
        _UpperCamelCase      :	Union[str, Any]						=       CMStochasticIterativeScheduler(
            num_train_timesteps=40					,   sigma_min=0.0_02					,   sigma_max=80.0					,   )
        _UpperCamelCase      :	Any						=       ConsistencyModelPipeline(unet=__a					,   scheduler=__a	)
        pipe.to(torch_device=__a					,   torch_dtype=torch.floataa	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	Optional[int]						=       self.get_inputs(get_fixed_latents=__a					,   device=__a	)
        # Ensure usage of flash attention in torch 2.0
        with sdp_kernel(enable_flash=__a					,   enable_math=__a					,   enable_mem_efficient=__a	):
              _UpperCamelCase      :	Union[str, Any]						=       pipe(**__a	).images
        assert image.shape == (1, 64, 64, 3)

        _UpperCamelCase      :	Any						=       image[0, -3:, -3:, -1]

        _UpperCamelCase      :	Tuple						=       np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]	)

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


  @require_torch_a
  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple	)		->				str:
        _UpperCamelCase      :	List[Any]						=       UNetaDModel.from_pretrained("diffusers/consistency_models"					,   subfolder="diffusers_cd_imagenet64_l2"	)
        _UpperCamelCase      :	Optional[Any]						=       CMStochasticIterativeScheduler(
            num_train_timesteps=40					,   sigma_min=0.0_02					,   sigma_max=80.0					,   )
        _UpperCamelCase      :	Optional[Any]						=       ConsistencyModelPipeline(unet=__a					,   scheduler=__a	)
        pipe.to(torch_device=__a					,   torch_dtype=torch.floataa	)
        pipe.set_progress_bar_config(disable=__a	)

        _UpperCamelCase      :	int						=       self.get_inputs(get_fixed_latents=__a					,   device=__a	)
        _UpperCamelCase      :	Any						=       1
        _UpperCamelCase      :	Any						=       None
        # Ensure usage of flash attention in torch 2.0
        with sdp_kernel(enable_flash=__a					,   enable_math=__a					,   enable_mem_efficient=__a	):
              _UpperCamelCase      :	int						=       pipe(**__a	).images
        assert image.shape == (1, 64, 64, 3)

        _UpperCamelCase      :	Union[str, Any]						=       image[0, -3:, -3:, -1]

        _UpperCamelCase      :	Any						=       np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]	)

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







"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


lowerCamelCase__				  =   {
    "configuration_owlvit": [
        "OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "OwlViTConfig",
        "OwlViTOnnxConfig",
        "OwlViTTextConfig",
        "OwlViTVisionConfig",
    ],
    "processing_owlvit": ["OwlViTProcessor"],
}


try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__				  =   ["OwlViTFeatureExtractor"]
  lowerCamelCase__				  =   ["OwlViTImageProcessor"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__				  =   [
      "OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "OwlViTModel",
      "OwlViTPreTrainedModel",
      "OwlViTTextModel",
      "OwlViTVisionModel",
      "OwlViTForObjectDetection",
  ]

if TYPE_CHECKING:
  from .configuration_owlvit import (
      OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      OwlViTConfig,
      OwlViTOnnxConfig,
      OwlViTTextConfig,
      OwlViTVisionConfig,
  )
  from .processing_owlvit import OwlViTProcessor

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_owlvit import OwlViTFeatureExtractor
    from .image_processing_owlvit import OwlViTImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_owlvit import (
        OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        OwlViTForObjectDetection,
        OwlViTModel,
        OwlViTPreTrainedModel,
        OwlViTTextModel,
        OwlViTVisionModel,
    )

else:
  import sys

  lowerCamelCase__				  =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
