




"""simple docstring"""





from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import (
					    TFFunnelBaseModel,
					    TFFunnelForMaskedLM,
					    TFFunnelForMultipleChoice,
					    TFFunnelForPreTraining,
					    TFFunnelForQuestionAnswering,
					    TFFunnelForSequenceClassification,
					    TFFunnelForTokenClassification,
					    TFFunnelModel,
					)



class      a__				:


			def __init__(      self      :    Optional[int]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[Any]=13      ,  UpperCamelCase_      :    Any=7      ,  UpperCamelCase_      :    List[Any]=True      ,  UpperCamelCase_      :    Optional[Any]=True      ,  UpperCamelCase_      :    Optional[Any]=True      ,  UpperCamelCase_      :    Union[str, Any]=True      ,  UpperCamelCase_      :    Dict=99      ,  UpperCamelCase_      :    Union[str, Any]=[1, 1, 2]      ,  UpperCamelCase_      :    List[Any]=1      ,  UpperCamelCase_      :    str=32      ,  UpperCamelCase_      :    Any=4      ,  UpperCamelCase_      :    Optional[Any]=8      ,  UpperCamelCase_      :    Any=37      ,  UpperCamelCase_      :    Union[str, Any]="gelu_new"      ,  UpperCamelCase_      :    Tuple=0.1      ,  UpperCamelCase_      :    Dict=0.1      ,  UpperCamelCase_      :    Any=0.0      ,  UpperCamelCase_      :    str=512      ,  UpperCamelCase_      :    Tuple=3      ,  UpperCamelCase_      :    int=0.02      ,  UpperCamelCase_      :    Tuple=3      ,  UpperCamelCase_      :    Tuple=4      ,  UpperCamelCase_      :    Tuple=None      ,  UpperCamelCase_      :    int=False      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					parent
						__UpperCAmelCase	: List[str]		=					batch_size
						__UpperCAmelCase	: Union[str, Any]		=					seq_length
						__UpperCAmelCase	: Optional[int]		=					is_training
						__UpperCAmelCase	: int		=					use_input_mask
						__UpperCAmelCase	: Dict		=					use_token_type_ids
						__UpperCAmelCase	: int		=					use_labels
						__UpperCAmelCase	: List[str]		=					vocab_size
						__UpperCAmelCase	: Union[str, Any]		=					block_sizes
						__UpperCAmelCase	: str		=					num_decoder_layers
						__UpperCAmelCase	: Tuple		=					d_model
						__UpperCAmelCase	: Optional[int]		=					n_head
						__UpperCAmelCase	: Optional[int]		=					d_head
						__UpperCAmelCase	: List[str]		=					d_inner
						__UpperCAmelCase	: str		=					hidden_act
						__UpperCAmelCase	: str		=					hidden_dropout
						__UpperCAmelCase	: str		=					attention_dropout
						__UpperCAmelCase	: Optional[Any]		=					activation_dropout
						__UpperCAmelCase	: Tuple		=					max_position_embeddings
						__UpperCAmelCase	: Any		=					type_vocab_size
						__UpperCAmelCase	: str		=					2
						__UpperCAmelCase	: Dict		=					num_labels
						__UpperCAmelCase	: str		=					num_choices
						__UpperCAmelCase	: Optional[int]		=					scope
						__UpperCAmelCase	: Union[str, Any]		=					initializer_std

						# Used in the tests to check the size of the first attention layer
						__UpperCAmelCase	: List[str]		=					n_head
						# Used in the tests to check the size of the first hidden state
						__UpperCAmelCase	: int		=					self.d_model
						# Used in the tests to check the number of output hidden states/attentions
						__UpperCAmelCase	: Optional[Any]		=					sum(self.block_sizes) + (0 if base else self.num_decoder_layers)
						# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
						# the last hidden state of the first block (which is the first hidden state of the decoder).
						if not base:
									__UpperCAmelCase	: Dict		=					self.num_hidden_layers + 2


			def       a_     (      self      :    int):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size)

						__UpperCAmelCase	: int		=					None
						if self.use_input_mask:
									__UpperCAmelCase	: Any		=					random_attention_mask([self.batch_size, self.seq_length])

						__UpperCAmelCase	: int		=					None
						if self.use_token_type_ids:
									__UpperCAmelCase	: Union[str, Any]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.type_vocab_size)

						__UpperCAmelCase	: Any		=					None
						__UpperCAmelCase	: int		=					None
						__UpperCAmelCase	: Any		=					None
						if self.use_labels:
									__UpperCAmelCase	: str		=					ids_tensor([self.batch_size]      ,  self.type_sequence_label_size)
									__UpperCAmelCase	: Union[str, Any]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.num_labels)
									__UpperCAmelCase	: Optional[Any]		=					ids_tensor([self.batch_size]      ,  self.num_choices)

						__UpperCAmelCase	: Tuple		=					FunnelConfig(
						    vocab_size=self.vocab_size      ,  block_sizes=self.block_sizes      ,  num_decoder_layers=self.num_decoder_layers      ,  d_model=self.d_model      ,  n_head=self.n_head      ,  d_head=self.d_head      ,  d_inner=self.d_inner      ,  hidden_act=self.hidden_act      ,  hidden_dropout=self.hidden_dropout      ,  attention_dropout=self.attention_dropout      ,  activation_dropout=self.activation_dropout      ,  max_position_embeddings=self.max_position_embeddings      ,  type_vocab_size=self.type_vocab_size      ,  initializer_std=self.initializer_std      ,  )

						return (
						    config,
						    input_ids,
						    token_type_ids,
						    input_mask,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						)


			def       a_     (      self      :    List[Any]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Union[str, Any]      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					TFFunnelModel(config=UpperCamelCase_)
						__UpperCAmelCase	: Any		=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						__UpperCAmelCase	: Optional[Any]		=					model(UpperCamelCase_)

						__UpperCAmelCase	: str		=					[input_ids, input_mask]
						__UpperCAmelCase	: Any		=					model(UpperCamelCase_)

						__UpperCAmelCase	: Optional[Any]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, self.seq_length, self.d_model))

						__UpperCAmelCase	: Optional[Any]		=					False
						__UpperCAmelCase	: Dict		=					TFFunnelModel(config=UpperCamelCase_)
						__UpperCAmelCase	: List[Any]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, self.seq_length, self.d_model))

						__UpperCAmelCase	: str		=					False
						__UpperCAmelCase	: int		=					TFFunnelModel(config=UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, self.seq_length, self.d_model))


			def       a_     (      self      :    str      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Any      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					TFFunnelBaseModel(config=UpperCamelCase_)
						__UpperCAmelCase	: List[str]		=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						__UpperCAmelCase	: str		=					model(UpperCamelCase_)

						__UpperCAmelCase	: Union[str, Any]		=					[input_ids, input_mask]
						__UpperCAmelCase	: int		=					model(UpperCamelCase_)

						__UpperCAmelCase	: Union[str, Any]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, 2, self.d_model))

						__UpperCAmelCase	: Dict		=					False
						__UpperCAmelCase	: List[Any]		=					TFFunnelBaseModel(config=UpperCamelCase_)
						__UpperCAmelCase	: Union[str, Any]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, 3, self.d_model))

						__UpperCAmelCase	: Any		=					False
						__UpperCAmelCase	: int		=					TFFunnelBaseModel(config=UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, 2, self.d_model))


			def       a_     (      self      :    Tuple      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[int]      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					TFFunnelForPreTraining(config=UpperCamelCase_)
						__UpperCAmelCase	: Optional[Any]		=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						__UpperCAmelCase	: Optional[Any]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length))


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Tuple      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					TFFunnelForMaskedLM(config=UpperCamelCase_)
						__UpperCAmelCase	: List[str]		=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						__UpperCAmelCase	: Optional[Any]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size))


			def       a_     (      self      :    List[str]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Tuple      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					self.num_labels
						__UpperCAmelCase	: Optional[Any]		=					TFFunnelForSequenceClassification(config=UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						__UpperCAmelCase	: List[str]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_labels))


			def       a_     (      self      :    Dict      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Any      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					self.num_choices
						__UpperCAmelCase	: Tuple		=					TFFunnelForMultipleChoice(config=UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					tf.tile(tf.expand_dims(UpperCamelCase_      ,  1)      ,  (1, self.num_choices, 1))
						__UpperCAmelCase	: Optional[Any]		=					tf.tile(tf.expand_dims(UpperCamelCase_      ,  1)      ,  (1, self.num_choices, 1))
						__UpperCAmelCase	: str		=					tf.tile(tf.expand_dims(UpperCamelCase_      ,  1)      ,  (1, self.num_choices, 1))
						__UpperCAmelCase	: Union[str, Any]		=					{
						    "input_ids": multiple_choice_inputs_ids,
						    "attention_mask": multiple_choice_input_mask,
						    "token_type_ids": multiple_choice_token_type_ids,
						}
						__UpperCAmelCase	: Any		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_choices))


			def       a_     (      self      :    Optional[Any]      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Tuple      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.num_labels
						__UpperCAmelCase	: Dict		=					TFFunnelForTokenClassification(config=UpperCamelCase_)
						__UpperCAmelCase	: List[Any]		=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						__UpperCAmelCase	: Dict		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.num_labels))


			def       a_     (      self      :    Any      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Any      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					TFFunnelForQuestionAnswering(config=UpperCamelCase_)
						__UpperCAmelCase	: Any		=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
						__UpperCAmelCase	: List[str]		=					model(UpperCamelCase_)
						self.parent.assertEqual(result.start_logits.shape      ,  (self.batch_size, self.seq_length))
						self.parent.assertEqual(result.end_logits.shape      ,  (self.batch_size, self.seq_length))


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					self.prepare_config_and_inputs()
						(
						    (
						    __UpperCAmelCase
						)	,     (
						    __UpperCAmelCase
						)	,     (
						    __UpperCAmelCase
						)	,     (
						    __UpperCAmelCase
						)	,     (
						    __UpperCAmelCase
						)	,     (
						    __UpperCAmelCase
						)	,     (
						    __UpperCAmelCase
						)	,     
						)	: Tuple		=					config_and_inputs
						__UpperCAmelCase	: str		=					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
						return config, inputs_dict






@require_tf
class      a__				(	__magic_name__							,    __magic_name__							,    unittest.TestCase   ):
			lowercase_    =	(
			    (
			        TFFunnelModel,
			        TFFunnelForMaskedLM,
			        TFFunnelForPreTraining,
			        TFFunnelForQuestionAnswering,
			        TFFunnelForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			lowercase_    =	(
			    {
			        "feature-extraction": (TFFunnelBaseModel, TFFunnelModel),
			        "fill-mask": TFFunnelForMaskedLM,
			        "question-answering": TFFunnelForQuestionAnswering,
			        "text-classification": TFFunnelForSequenceClassification,
			        "token-classification": TFFunnelForTokenClassification,
			        "zero-shot": TFFunnelForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			lowercase_    =	False
			lowercase_    =	False


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					TFFunnelModelTester(self)
						__UpperCAmelCase	: List[Any]		=					ConfigTester(self      ,  config_class=UpperCamelCase_)


			def       a_     (      self      :    Union[str, Any]):


						"""simple docstring"""




						self.config_tester.run_common_tests()


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: Any		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCamelCase_)


			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*UpperCamelCase_)


			def       a_     (      self      :    Dict):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_)


			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_)


			def       a_     (      self      :    Union[str, Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*UpperCamelCase_)






@require_tf
class      a__				(	__magic_name__							,    unittest.TestCase   ):
			lowercase_    =	(
			    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
			)
			lowercase_    =	False
			lowercase_    =	False


			def       a_     (      self      :    Dict):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					TFFunnelModelTester(self      ,  base=UpperCamelCase_)
						__UpperCAmelCase	: Any		=					ConfigTester(self      ,  config_class=UpperCamelCase_)


			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						self.config_tester.run_common_tests()


			def       a_     (      self      :    int):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_base_model(*UpperCamelCase_)


			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase_)


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase_)








"""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      a__				(	__magic_name__							,    unittest.TestCase   ):
			lowercase_    =	ConsistencyModelPipeline
			lowercase_    =	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			lowercase_    =	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


			@property
			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					UNetaDModel.from_pretrained(
						    "diffusers/consistency-models-test"      ,  subfolder="test_unet"      ,  )
						return unet


			@property
			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					UNetaDModel.from_pretrained(
						    "diffusers/consistency-models-test"      ,  subfolder="test_unet_class_cond"      ,  )
						return unet


			def       a_     (      self      :    Any      ,  UpperCamelCase_      :    int=False):


						"""simple docstring"""




						if class_cond:
									__UpperCAmelCase	: List[Any]		=					self.dummy_cond_unet
						else:
									__UpperCAmelCase	: Optional[int]		=					self.dummy_uncond_unet

						# Default to CM multistep sampler
						__UpperCAmelCase	: List[str]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )

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

						return components


			def       a_     (      self      :    List[str]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[Any]=0):


						"""simple docstring"""




						if str(UpperCamelCase_).startswith("mps"):
									__UpperCAmelCase	: str		=					torch.manual_seed(UpperCamelCase_)
						else:
									__UpperCAmelCase	: Optional[Any]		=					torch.Generator(device=UpperCamelCase_).manual_seed(UpperCamelCase_)

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

						return inputs


			def       a_     (      self      :    int):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Union[str, Any]		=					self.get_dummy_components()
						__UpperCAmelCase	: str		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: Any		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Any		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: str		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Dict		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: List[Any]		=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004])

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


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Optional[int]		=					self.get_dummy_components(class_cond=UpperCamelCase_)
						__UpperCAmelCase	: Optional[Any]		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: str		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					0
						__UpperCAmelCase	: int		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Any		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: List[Any]		=					np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004])

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


			def       a_     (      self      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: Optional[int]		=					self.get_dummy_components()
						__UpperCAmelCase	: Optional[Any]		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: Optional[Any]		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: List[str]		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Dict		=					1
						__UpperCAmelCase	: int		=					None
						__UpperCAmelCase	: List[Any]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: Tuple		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987])

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


			def       a_     (      self      :    List[str]):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					"cpu"  # ensure determinism for the device-dependent torch.Generator
						__UpperCAmelCase	: List[str]		=					self.get_dummy_components(class_cond=UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					ConsistencyModelPipeline(**UpperCamelCase_)
						__UpperCAmelCase	: int		=					pipe.to(UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Tuple		=					self.get_dummy_inputs(UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					1
						__UpperCAmelCase	: Optional[int]		=					None
						__UpperCAmelCase	: Tuple		=					0
						__UpperCAmelCase	: Tuple		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 32, 32, 3)

						__UpperCAmelCase	: List[str]		=					image[0, -3:, -3:, -1]
						__UpperCAmelCase	: Dict		=					np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987])

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






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


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




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


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    Optional[int]=0      ,  UpperCamelCase_      :    Tuple=False      ,  UpperCamelCase_      :    int="cpu"      ,  UpperCamelCase_      :    Any=torch.floataa      ,  UpperCamelCase_      :    List[str]=(1, 3, 64, 64)):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					torch.manual_seed(UpperCamelCase_)

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

						if get_fixed_latents:
									__UpperCAmelCase	: int		=					self.get_fixed_latents(seed=UpperCamelCase_      ,  device=UpperCamelCase_      ,  dtype=UpperCamelCase_      ,  shape=UpperCamelCase_)
									__UpperCAmelCase	: Optional[int]		=					latents

						return inputs


			def       a_     (      self      :    Union[str, Any]      ,  UpperCamelCase_      :    int=0      ,  UpperCamelCase_      :    Tuple="cpu"      ,  UpperCamelCase_      :    Tuple=torch.floataa      ,  UpperCamelCase_      :    Optional[Any]=(1, 3, 64, 64)):


						"""simple docstring"""




						if type(UpperCamelCase_) == str:
									__UpperCAmelCase	: Union[str, Any]		=					torch.device(UpperCamelCase_)
						__UpperCAmelCase	: Union[str, Any]		=					torch.Generator(device=UpperCamelCase_).manual_seed(UpperCamelCase_)
						__UpperCAmelCase	: Optional[int]		=					randn_tensor(UpperCamelCase_      ,  generator=UpperCamelCase_      ,  device=UpperCamelCase_      ,  dtype=UpperCamelCase_)
						return latents


			def       a_     (      self      :    List[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[int]		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Dict		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: Dict		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: Dict		=					self.get_inputs()
						__UpperCAmelCase	: List[str]		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

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

						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254])

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


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Optional[Any]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: Any		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

						__UpperCAmelCase	: int		=					self.get_inputs()
						__UpperCAmelCase	: str		=					1
						__UpperCAmelCase	: Union[str, Any]		=					None
						__UpperCAmelCase	: Tuple		=					pipe(**UpperCamelCase_).images
						assert image.shape == (1, 64, 64, 3)

						__UpperCAmelCase	: str		=					image[0, -3:, -3:, -1]

						__UpperCAmelCase	: Union[str, Any]		=					np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217])

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


			@require_torch_a
			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Optional[int]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: int		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_      ,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

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

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

						__UpperCAmelCase	: Optional[int]		=					np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353])

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


			@require_torch_a
			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					UNetaDModel.from_pretrained("diffusers/consistency_models"      ,  subfolder="diffusers_cd_imagenet64_l2")
						__UpperCAmelCase	: Union[str, Any]		=					CMStochasticIterativeScheduler(
						    num_train_timesteps=40      ,  sigma_min=0.002      ,  sigma_max=80.0      ,  )
						__UpperCAmelCase	: List[Any]		=					ConsistencyModelPipeline(unet=UpperCamelCase_      ,  scheduler=UpperCamelCase_)
						pipe.to(torch_device=UpperCamelCase_      ,  torch_dtype=torch.floataa)
						pipe.set_progress_bar_config(disable=UpperCamelCase_)

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

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

						__UpperCAmelCase	: List[str]		=					np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095])

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


