


from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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 (
		    TFRoFormerForCausalLM,
		    TFRoFormerForMaskedLM,
		    TFRoFormerForMultipleChoice,
		    TFRoFormerForQuestionAnswering,
		    TFRoFormerForSequenceClassification,
		    TFRoFormerForTokenClassification,
		    TFRoFormerModel,
		)
		from transformers.models.roformer.modeling_tf_roformer import (
		    TFRoFormerSelfAttention,
		    TFRoFormerSinusoidalPositionalEmbedding,
		)







class        __lowercase       :



		"""simple docstring"""

		def __init__( self		,      A		,      A=13		,      A=7		,      A=True		,      A=True		,      A=True		,      A=True		,      A=99		,      A=32		,      A=2		,      A=4		,      A=37		,      A="gelu"		,      A=0.1		,      A=0.1		,      A=5_12		,      A=16		,      A=2		,      A=0.02		,      A=3		,      A=4		,      A=None		,      )  -> Tuple:







					'''simple docstring'''







					lowerCamelCase    		=  parent
					lowerCamelCase    		=  13
					lowerCamelCase    		=  7
					lowerCamelCase    		=  True
					lowerCamelCase    		=  True
					lowerCamelCase    		=  True
					lowerCamelCase    		=  True
					lowerCamelCase    		=  99
					lowerCamelCase    		=  32
					lowerCamelCase    		=  2
					lowerCamelCase    		=  4
					lowerCamelCase    		=  37
					lowerCamelCase    		=  """gelu"""
					lowerCamelCase    		=  0.1
					lowerCamelCase    		=  0.1
					lowerCamelCase    		=  5_12
					lowerCamelCase    		=  16
					lowerCamelCase    		=  2
					lowerCamelCase    		=  0.02
					lowerCamelCase    		=  3
					lowerCamelCase    		=  4
					lowerCamelCase    		=  None







		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







					lowerCamelCase    		=  ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size	)

					lowerCamelCase    		=  None
					if self.use_input_mask:
								lowerCamelCase    		=  random_attention_mask([self.batch_size, self.seq_length]	)

					lowerCamelCase    		=  None
					if self.use_token_type_ids:
								lowerCamelCase    		=  ids_tensor([self.batch_size, self.seq_length]		,      self.type_vocab_size	)

					lowerCamelCase    		=  None
					lowerCamelCase    		=  None
					lowerCamelCase    		=  None
					if self.use_labels:
								lowerCamelCase    		=  ids_tensor([self.batch_size]		,      self.type_sequence_label_size	)
								lowerCamelCase    		=  ids_tensor([self.batch_size, self.seq_length]		,      self.num_labels	)
								lowerCamelCase    		=  ids_tensor([self.batch_size]		,      self.num_choices	)

					lowerCamelCase    		=  RoFormerConfig(
					    vocab_size=self.vocab_size		,      hidden_size=self.hidden_size		,      num_hidden_layers=self.num_hidden_layers		,      num_attention_heads=self.num_attention_heads		,      intermediate_size=self.intermediate_size		,      hidden_act=self.hidden_act		,      hidden_dropout_prob=self.hidden_dropout_prob		,      attention_probs_dropout_prob=self.attention_probs_dropout_prob		,      max_position_embeddings=self.max_position_embeddings		,      type_vocab_size=self.type_vocab_size		,      initializer_range=self.initializer_range		,      return_dict=A		,      )

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







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Tuple:







					'''simple docstring'''







					lowerCamelCase    		=  TFRoFormerModel(config=A	)
					lowerCamelCase    		=  {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

					lowerCamelCase    		=  [input_ids, input_mask]
					lowerCamelCase    		=  model(A	)

					lowerCamelCase    		=  model(A	)

					self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, self.seq_length, self.hidden_size)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  True
					lowerCamelCase    		=  TFRoFormerForCausalLM(config=A	)
					lowerCamelCase    		=  {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}
					lowerCamelCase    		=  model(A	)["""logits"""]
					self.parent.assertListEqual(
					    list(prediction_scores.numpy().shape	)		,      [self.batch_size, self.seq_length, self.vocab_size]	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  TFRoFormerForMaskedLM(config=A	)
					lowerCamelCase    		=  {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}
					lowerCamelCase    		=  model(A	)
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  self.num_labels
					lowerCamelCase    		=  TFRoFormerForSequenceClassification(config=A	)
					lowerCamelCase    		=  {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}

					lowerCamelCase    		=  model(A	)
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> Dict:







					'''simple docstring'''







					lowerCamelCase    		=  self.num_choices
					lowerCamelCase    		=  TFRoFormerForMultipleChoice(config=A	)
					lowerCamelCase    		=  tf.tile(tf.expand_dims(A		,      1	)		,      (1, self.num_choices, 1)	)
					lowerCamelCase    		=  tf.tile(tf.expand_dims(A		,      1	)		,      (1, self.num_choices, 1)	)
					lowerCamelCase    		=  tf.tile(tf.expand_dims(A		,      1	)		,      (1, self.num_choices, 1)	)
					lowerCamelCase    		=  {
					    """input_ids""": multiple_choice_inputs_ids,
					    """attention_mask""": multiple_choice_input_mask,
					    """token_type_ids""": multiple_choice_token_type_ids,
					}
					lowerCamelCase    		=  model(A	)
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  self.num_labels
					lowerCamelCase    		=  TFRoFormerForTokenClassification(config=A	)
					lowerCamelCase    		=  {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}
					lowerCamelCase    		=  model(A	)
					self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.num_labels)	)







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A		,      A		,      A	)  -> List[str]:







					'''simple docstring'''







					lowerCamelCase    		=  TFRoFormerForQuestionAnswering(config=A	)
					lowerCamelCase    		=  {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}

					lowerCamelCase    		=  model(A	)
					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	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  self.prepare_config_and_inputs()
					(
					    (
					    lowerCamelCase
					)		,		(
					    lowerCamelCase
					)		,		(
					    lowerCamelCase
					)		,		(
					    lowerCamelCase
					)		,		(
					    lowerCamelCase
					)		,		(
					    lowerCamelCase
					)		,		(
					    lowerCamelCase
					)		,		
					)    		=  config_and_inputs
					lowerCamelCase    		=  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
					return config, inputs_dict







@require_tf
class        __lowercase       (							a_			,							a_			,							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					int			       =		(
		    (
		        TFRoFormerModel,
		        TFRoFormerForCausalLM,
		        TFRoFormerForMaskedLM,
		        TFRoFormerForQuestionAnswering,
		        TFRoFormerForSequenceClassification,
		        TFRoFormerForTokenClassification,
		        TFRoFormerForMultipleChoice,
		    )
		    if is_tf_available()
		    else ()
		)
		UpperCamelCase     :					int			       =		(
		    {
		        "feature-extraction": TFRoFormerModel,
		        "fill-mask": TFRoFormerForMaskedLM,
		        "question-answering": TFRoFormerForQuestionAnswering,
		        "text-classification": TFRoFormerForSequenceClassification,
		        "text-generation": TFRoFormerForCausalLM,
		        "token-classification": TFRoFormerForTokenClassification,
		        "zero-shot": TFRoFormerForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)

		UpperCamelCase     :					Optional[int]			       =		False
		UpperCamelCase     :					Optional[Any]			       =		False







		def 					__A   ( self		,      A		,      A		,      A		,      A		,      A	)  -> Tuple:







					'''simple docstring'''







					if pipeline_test_casse_name == "TextGenerationPipelineTests":
								return True

					return False







		def 					__A   ( self	)  -> Dict:







					'''simple docstring'''







					lowerCamelCase    		=  TFRoFormerModelTester(self	)
					lowerCamelCase    		=  ConfigTester(self		,      config_class=A		,      hidden_size=37	)







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					self.config_tester.run_common_tests()







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*A	)







		def 					__A   ( self	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*A	)







		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_lm_head(*A	)







		def 					__A   ( self	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_multiple_choice(*A	)







		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_question_answering(*A	)







		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_sequence_classification(*A	)







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*A	)







		@slow
		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  TFRoFormerModel.from_pretrained("""junnyu/roformer_chinese_base"""	)
					self.assertIsNotNone(A	)







@require_tf
class        __lowercase       (							unittest.TestCase  ):



		"""simple docstring"""

		@slow
		def 					__A   ( self	)  -> Tuple:







					'''simple docstring'''







					lowerCamelCase    		=  TFRoFormerForMaskedLM.from_pretrained("""junnyu/roformer_chinese_base"""	)
					lowerCamelCase    		=  tf.constant([[0, 1, 2, 3, 4, 5]]	)
					lowerCamelCase    		=  model(A	)[0]

					# TODO Replace vocab size
					lowerCamelCase    		=  5_00_00

					lowerCamelCase    		=  [1, 6, vocab_size]
					self.assertEqual(output.shape		,      A	)

					print(output[:, :3, :3]	)

					# TODO Replace values below with what was printed above.
					lowerCamelCase    		=  tf.constant(
					    [
					        [
					            [-0.12053341, -1.0264901, 0.29221946],
					            [-1.5133783, 0.197433, 0.15190607],
					            [-5.0135403, -3.900256, -0.84038764],
					        ]
					    ]	)
					tf.debugging.assert_near(output[:, :3, :3]		,      A		,      atol=1e-4	)







@require_tf
class        __lowercase       (							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					int			       =		1e-4







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  tf.constant([[4, 10]]	)
					lowerCamelCase    		=  TFRoFormerSinusoidalPositionalEmbedding(num_positions=6		,      embedding_dim=6	)

					lowerCamelCase    		=  emba(input_ids.shape	)
					lowerCamelCase    		=  tf.constant(
					    [[0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 1.0000], [0.8415, 0.0464, 0.0022, 0.5403, 0.9989, 1.0000]]	)

					tf.debugging.assert_near(A		,      A		,      atol=self.tolerance	)







		def 					__A   ( self	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  tf.constant(
					    [
					        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
					        [0.8415, 0.8219, 0.8020, 0.7819, 0.7617],
					        [0.9093, 0.9364, 0.9581, 0.9749, 0.9870],
					    ]	)
					lowerCamelCase    		=  TFRoFormerSinusoidalPositionalEmbedding(num_positions=5_12		,      embedding_dim=5_12	)
					emba([2, 16, 5_12]	)
					lowerCamelCase    		=  emba.weight[:3, :5]

					tf.debugging.assert_near(A		,      A		,      atol=self.tolerance	)







@require_tf
class        __lowercase       (							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					List[Any]			       =		1e-4







		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







					lowerCamelCase    		=  tf.reshape(tf.range(2 * 12 * 16 * 64		,      dtype=tf.floataa	)		,      shape=(2, 12, 16, 64)	) / 1_00

					lowerCamelCase    		=  -tf.reshape(tf.range(2 * 12 * 16 * 64		,      dtype=tf.floataa	)		,      shape=(2, 12, 16, 64)	) / 1_00

					lowerCamelCase    		=  TFRoFormerSinusoidalPositionalEmbedding(num_positions=32		,      embedding_dim=64	)
					lowerCamelCase    		=  embed_positions([2, 16, 7_68]	)[None, None, :, :]

					lowerCamelCase		,		lowerCamelCase    		=  TFRoFormerSelfAttention.apply_rotary_position_embeddings(
					    A		,      A		,      A	)

					lowerCamelCase    		=  tf.constant(
					    [
					        [0.0000, 0.0100, 0.0200, 0.0300, 0.0400, 0.0500, 0.0600, 0.0700],
					        [-0.2012, 0.8897, 0.0263, 0.9401, 0.2074, 0.9463, 0.3481, 0.9343],
					        [-1.7057, 0.6271, -1.2145, 1.3897, -0.6303, 1.7647, -0.1173, 1.8985],
					        [-2.1731, -1.6397, -2.7358, 0.2854, -2.1840, 1.7183, -1.3018, 2.4871],
					        [0.2717, -3.6173, -2.9206, -2.1988, -3.6638, 0.3858, -2.9155, 2.2980],
					        [3.9859, -2.1580, -0.7984, -4.4904, -4.1181, -2.0252, -4.4782, 1.1253],
					    ]	)
					lowerCamelCase    		=  tf.constant(
					    [
					        [0.0000, -0.0100, -0.0200, -0.0300, -0.0400, -0.0500, -0.0600, -0.0700],
					        [0.2012, -0.8897, -0.0263, -0.9401, -0.2074, -0.9463, -0.3481, -0.9343],
					        [1.7057, -0.6271, 1.2145, -1.3897, 0.6303, -1.7647, 0.1173, -1.8985],
					        [2.1731, 1.6397, 2.7358, -0.2854, 2.1840, -1.7183, 1.3018, -2.4871],
					        [-0.2717, 3.6173, 2.9206, 2.1988, 3.6638, -0.3858, 2.9155, -2.2980],
					        [-3.9859, 2.1580, 0.7984, 4.4904, 4.1181, 2.0252, 4.4782, -1.1253],
					    ]	)

					tf.debugging.assert_near(query_layer[0, 0, :6, :8]		,      A		,      atol=self.tolerance	)
					tf.debugging.assert_near(key_layer[0, 0, :6, :8]		,      A		,      atol=self.tolerance	)





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        __lowercase       (							a_			,							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					List[str]			       =		ConsistencyModelPipeline
		UpperCamelCase     :					int			       =		UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		UpperCamelCase     :					Any			       =		UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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







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







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained(
					    """diffusers/consistency-models-test"""		,      subfolder="""test_unet"""		,      )
					return unet







		@property
		def 					__A   ( self	)  -> str:







					'''simple docstring'''







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







		def 					__A   ( self		,      A=False	)  -> Tuple:







					'''simple docstring'''







					if class_cond:
								lowerCamelCase    		=  self.dummy_cond_unet
					else:
								lowerCamelCase    		=  self.dummy_uncond_unet

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

					lowerCamelCase    		=  {
					    """unet""": unet,
					    """scheduler""": scheduler,
					}

					return components







		def 					__A   ( self		,      A		,      A=0	)  -> List[str]:







					'''simple docstring'''







					if str(A	).startswith("""mps"""	):
								lowerCamelCase    		=  torch.manual_seed(A	)
					else:
								lowerCamelCase    		=  torch.Generator(device=A	).manual_seed(A	)

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

					return inputs







		def 					__A   ( self	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components()
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  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	)  -> Dict:







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components(class_cond=A	)
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  0
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  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	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components()
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  1
					lowerCamelCase    		=  None
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  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'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components(class_cond=A	)
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  1
					lowerCamelCase    		=  None
					lowerCamelCase    		=  0
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  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        __lowercase       (							unittest.TestCase  ):



		"""simple docstring"""

		def 					__A   ( self	)  -> Optional[Any]:







					'''simple docstring'''







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







		def 					__A   ( self		,      A=0		,      A=False		,      A="cpu"		,      A=torch.floataa		,      A=(1, 3, 64, 64)	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  torch.manual_seed(A	)

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

					if get_fixed_latents:
								lowerCamelCase    		=  self.get_fixed_latents(seed=A		,      device=A		,      dtype=A		,      shape=A	)
								lowerCamelCase    		=  latents

					return inputs







		def 					__A   ( self		,      A=0		,      A="cpu"		,      A=torch.floataa		,      A=(1, 3, 64, 64)	)  -> Union[str, Any]:







					'''simple docstring'''







					if type(A	) == str:
								lowerCamelCase    		=  torch.device(A	)
					lowerCamelCase    		=  torch.Generator(device=A	).manual_seed(A	)
					lowerCamelCase    		=  randn_tensor(A		,      generator=A		,      device=A		,      dtype=A	)
					return latents







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_inputs()
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 64, 64, 3)

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

					lowerCamelCase    		=  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	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_inputs()
					lowerCamelCase    		=  1
					lowerCamelCase    		=  None
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 64, 64, 3)

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

					lowerCamelCase    		=  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	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A		,      torch_dtype=torch.floataa	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  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    		=  pipe(**A	).images
					assert image.shape == (1, 64, 64, 3)

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

					lowerCamelCase    		=  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	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A		,      torch_dtype=torch.floataa	)
					pipe.set_progress_bar_config(disable=A	)

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

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

					lowerCamelCase    		=  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

