


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_			(a__						,       unittest.TestCase						):




       '''simple docstring'''



       _a   							=							ConsistencyModelPipeline
       _a   							=							UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       _a   							=							UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



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



       @property
       def 				_lowerCAmelCase							(						self				:				str							)					->str:
          lowerCamelCase_						:   Optional[Any]		= UNetaDModel.from_pretrained(
              """diffusers/consistency-models-test"""		,				subfolder="""test_unet_class_cond"""		,				)
          return unet



       def 				_lowerCAmelCase							(						self				:				Tuple		,				__a				:				Tuple=False							)					->Tuple:
          if class_cond:
             lowerCamelCase_						:   Any		= self.dummy_cond_unet
          else:
             lowerCamelCase_						:   List[Any]		= self.dummy_uncond_unet

          # Default to CM multistep sampler
          lowerCamelCase_						:   Dict		= CMStochasticIterativeScheduler(
              num_train_timesteps=40		,				sigma_min=0.002		,				sigma_max=80.0		,				)

          lowerCamelCase_						:   str		= {
              """unet""": unet,
              """scheduler""": scheduler,
          }

          return components



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

          lowerCamelCase_						:   List[str]		= {
              """batch_size""": 1,
              """num_inference_steps""": None,
              """timesteps""": [22, 0],
              """generator""": generator,
              """output_type""": """np""",
          }

          return inputs



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

          lowerCamelCase_						:   int		= self.get_dummy_inputs(__a							)
          lowerCamelCase_						:   Any		= pipe(**__a							).images
          assert image.shape == (1, 32, 32, 3)

          lowerCamelCase_						:   Tuple		= image[0, -3:, -3:, -1]
          lowerCamelCase_						:   str		= np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004]							)

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



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

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

          lowerCamelCase_						:   Tuple		= image[0, -3:, -3:, -1]
          lowerCamelCase_						:   Any		= np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004]							)

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



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

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

          lowerCamelCase_						:   Dict		= image[0, -3:, -3:, -1]
          lowerCamelCase_						:   str		= np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987]							)

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




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

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

          lowerCamelCase_						:   Optional[int]		= image[0, -3:, -3:, -1]
          lowerCamelCase_						:   int		= np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987]							)

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







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




       '''simple docstring'''



       def 				_lowerCAmelCase							(						self				:				Tuple							)					->str:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



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

          lowerCamelCase_						:   Tuple		= {
              """num_inference_steps""": None,
              """timesteps""": [22, 0],
              """class_labels""": 0,
              """generator""": generator,
              """output_type""": """np""",
          }

          if get_fixed_latents:
             lowerCamelCase_						:   Tuple		= self.get_fixed_latents(seed=__a		,				device=__a		,				dtype=__a		,				shape=__a							)
             lowerCamelCase_						:   str		= latents

          return inputs



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



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

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

          lowerCamelCase_						:   List[str]		= image[0, -3:, -3:, -1]

          lowerCamelCase_						:   Tuple		= np.array([0.0_888, 0.0_881, 0.0_666, 0.0_479, 0.0_292, 0.0_195, 0.0_201, 0.0_163, 0.0_254]							)

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



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

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

          lowerCamelCase_						:   Tuple		= image[0, -3:, -3:, -1]

          lowerCamelCase_						:   List[str]		= np.array([0.0_340, 0.0_152, 0.0_063, 0.0_267, 0.0_221, 0.0_107, 0.0_416, 0.0_186, 0.0_217]							)

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



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

          lowerCamelCase_						:   str		= 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							):
             lowerCamelCase_						:   Dict		= pipe(**__a							).images
          assert image.shape == (1, 64, 64, 3)

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

          lowerCamelCase_						:   List[Any]		= np.array([0.1_875, 0.1_428, 0.1_289, 0.2_151, 0.2_092, 0.1_477, 0.1_877, 0.1_641, 0.1_353]							)

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




       @require_torch_a
       def 				_lowerCAmelCase							(						self				:				Union[str, Any]							)					->int:
          lowerCamelCase_						:   Dict		= UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,				subfolder="""diffusers_cd_imagenet64_l2"""							)
          lowerCamelCase_						:   Any		= CMStochasticIterativeScheduler(
              num_train_timesteps=40		,				sigma_min=0.002		,				sigma_max=80.0		,				)
          lowerCamelCase_						:   Tuple		= ConsistencyModelPipeline(unet=__a		,				scheduler=__a							)
          pipe.to(torch_device=__a		,				torch_dtype=torch.floataa							)
          pipe.set_progress_bar_config(disable=__a							)

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

          lowerCamelCase_						:   Optional[int]		= image[0, -3:, -3:, -1]

          lowerCamelCase_						:   Any		= np.array([0.1_663, 0.1_948, 0.2_275, 0.1_680, 0.1_204, 0.1_245, 0.1_858, 0.1_338, 0.2_095]							)

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









from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


snake_case__				:	Tuple					       =					logging.get_logger(__name__)


if is_vision_available():
      import PIL





class   SCREAMING_SNAKE_CASE_			(a__						):




       '''simple docstring'''



       _a   							=							["pixel_values"]



       def __init__(						self				:				List[str]		,				__a				:				bool = True		,				__a				:				Dict[str, int] = None		,				__a				:				PILImageResampling = PILImageResampling.BICUBIC		,				__a				:				bool = True		,				__a				:				Dict[str, int] = None		,				__a				:				bool = True		,				__a				:				Union[int, float] = 1 / 255		,				__a				:				bool = True		,				__a				:				Optional[Union[float, List[float]]] = None		,				__a				:				Optional[Union[float, List[float]]] = None		,				__a				:				bool = True		,				**__a				:				str		,				)					->None:
          super().__init__(**__a							)
          lowerCamelCase_						:   Any		= size if size is not None else {"""shortest_edge""": 224}
          lowerCamelCase_						:   List[Any]		= get_size_dict(__a		,				default_to_square=__a							)
          lowerCamelCase_						:   List[str]		= crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
          lowerCamelCase_						:   Any		= get_size_dict(__a		,				default_to_square=__a		,				param_name="""crop_size"""							)

          lowerCamelCase_						:   List[Any]		= do_resize
          lowerCamelCase_						:   str		= size
          lowerCamelCase_						:   Any		= resample
          lowerCamelCase_						:   List[Any]		= do_center_crop
          lowerCamelCase_						:   Any		= crop_size
          lowerCamelCase_						:   List[Any]		= do_rescale
          lowerCamelCase_						:   Tuple		= rescale_factor
          lowerCamelCase_						:   Dict		= do_normalize
          lowerCamelCase_						:   Any		= image_mean if image_mean is not None else OPENAI_CLIP_MEAN
          lowerCamelCase_						:   int		= image_std if image_std is not None else OPENAI_CLIP_STD
          lowerCamelCase_						:   Tuple		= do_convert_rgb



       def 				_lowerCAmelCase							(						self				:				int		,				__a				:				np.ndarray		,				__a				:				Dict[str, int]		,				__a				:				PILImageResampling = PILImageResampling.BICUBIC		,				__a				:				Optional[Union[str, ChannelDimension]] = None		,				**__a				:				Dict		,				)					->np.ndarray:
          lowerCamelCase_						:   Optional[Any]		= get_size_dict(__a		,				default_to_square=__a							)
          if "shortest_edge" not in size:
             raise ValueError(F'''The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'''							)
          lowerCamelCase_						:   Dict		= get_resize_output_image_size(__a		,				size=size["""shortest_edge"""]		,				default_to_square=__a							)
          return resize(__a		,				size=__a		,				resample=__a		,				data_format=__a		,				**__a							)



       def 				_lowerCAmelCase							(						self				:				Union[str, Any]		,				__a				:				np.ndarray		,				__a				:				Dict[str, int]		,				__a				:				Optional[Union[str, ChannelDimension]] = None		,				**__a				:				List[Any]		,				)					->np.ndarray:
          lowerCamelCase_						:   List[Any]		= get_size_dict(__a							)
          if "height" not in size or "width" not in size:
             raise ValueError(F'''The `size` parameter must contain the keys (height, width). Got {size.keys()}'''							)
          return center_crop(__a		,				size=(size["""height"""], size["""width"""])		,				data_format=__a		,				**__a							)



       def 				_lowerCAmelCase							(						self				:				Optional[Any]		,				__a				:				np.ndarray		,				__a				:				Union[int, float]		,				__a				:				Optional[Union[str, ChannelDimension]] = None		,				**__a				:				Dict		,				)					->Dict:
          return rescale(__a		,				scale=__a		,				data_format=__a		,				**__a							)



       def 				_lowerCAmelCase							(						self				:				List[Any]		,				__a				:				np.ndarray		,				__a				:				Union[float, List[float]]		,				__a				:				Union[float, List[float]]		,				__a				:				Optional[Union[str, ChannelDimension]] = None		,				**__a				:				List[str]		,				)					->np.ndarray:
          return normalize(__a		,				mean=__a		,				std=__a		,				data_format=__a		,				**__a							)




       def 				_lowerCAmelCase							(						self				:				str		,				__a				:				ImageInput		,				__a				:				bool = None		,				__a				:				Dict[str, int] = None		,				__a				:				PILImageResampling = None		,				__a				:				bool = None		,				__a				:				int = None		,				__a				:				bool = None		,				__a				:				float = None		,				__a				:				bool = None		,				__a				:				Optional[Union[float, List[float]]] = None		,				__a				:				Optional[Union[float, List[float]]] = None		,				__a				:				bool = None		,				__a				:				Optional[Union[str, TensorType]] = None		,				__a				:				Optional[ChannelDimension] = ChannelDimension.FIRST		,				**__a				:				List[str]		,				)					->PIL.Image.Image:
          lowerCamelCase_						:   List[Any]		= do_resize if do_resize is not None else self.do_resize
          lowerCamelCase_						:   Any		= size if size is not None else self.size
          lowerCamelCase_						:   str		= get_size_dict(__a		,				param_name="""size"""		,				default_to_square=__a							)
          lowerCamelCase_						:   List[str]		= resample if resample is not None else self.resample
          lowerCamelCase_						:   List[str]		= do_center_crop if do_center_crop is not None else self.do_center_crop
          lowerCamelCase_						:   List[Any]		= crop_size if crop_size is not None else self.crop_size
          lowerCamelCase_						:   Optional[int]		= get_size_dict(__a		,				param_name="""crop_size"""		,				default_to_square=__a							)
          lowerCamelCase_						:   List[Any]		= do_rescale if do_rescale is not None else self.do_rescale
          lowerCamelCase_						:   List[str]		= rescale_factor if rescale_factor is not None else self.rescale_factor
          lowerCamelCase_						:   Dict		= do_normalize if do_normalize is not None else self.do_normalize
          lowerCamelCase_						:   Dict		= image_mean if image_mean is not None else self.image_mean
          lowerCamelCase_						:   Tuple		= image_std if image_std is not None else self.image_std
          lowerCamelCase_						:   Union[str, Any]		= do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

          lowerCamelCase_						:   Union[str, Any]		= make_list_of_images(__a							)

          if not valid_images(__a							):
             raise ValueError(
                 """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
                 """torch.Tensor, tf.Tensor or jax.ndarray."""							)

          if do_resize and size is None:
             raise ValueError("""Size must be specified if do_resize is True."""							)

          if do_center_crop and crop_size is None:
             raise ValueError("""Crop size must be specified if do_center_crop is True."""							)

          if do_rescale and rescale_factor is None:
             raise ValueError("""Rescale factor must be specified if do_rescale is True."""							)

          if do_normalize and (image_mean is None or image_std is None):
             raise ValueError("""Image mean and std must be specified if do_normalize is True."""							)

          # PIL RGBA images are converted to RGB
          if do_convert_rgb:
             lowerCamelCase_						:   str		= [convert_to_rgb(__a							) for image in images]

          # All transformations expect numpy arrays.
          lowerCamelCase_						:   str		= [to_numpy_array(__a							) for image in images]

          if do_resize:
             lowerCamelCase_						:   int		= [self.resize(image=__a		,				size=__a		,				resample=__a							) for image in images]

          if do_center_crop:
             lowerCamelCase_						:   Dict		= [self.center_crop(image=__a		,				size=__a							) for image in images]

          if do_rescale:
             lowerCamelCase_						:   List[str]		= [self.rescale(image=__a		,				scale=__a							) for image in images]

          if do_normalize:
             lowerCamelCase_						:   Tuple		= [self.normalize(image=__a		,				mean=__a		,				std=__a							) for image in images]

          lowerCamelCase_						:   Tuple		= [to_channel_dimension_format(__a		,				__a							) for image in images]

          lowerCamelCase_						:   Tuple		= {"""pixel_values""": images}
          return BatchFeature(data=__a		,				tensor_type=__a							)





