import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer




class 						__lowercase		:




      """simple docstring"""





      def __init__(				self    ,				A    ,				A=1_3    ,				A=7    ,				A=True    ,				A=True    ,				A=False    ,				A=True    ,				A=9_9    ,				A=3_2    ,				A=5    ,				A=4    ,				A=3_7    ,				A="gelu"    ,				A=0.1    ,				A=0.1    ,				A=5_1_2    ,				A=1_6    ,				A=2    ,				A=0.02    ,				A=3    ,				A=4    ,				A=None    ,				)		->    Tuple:
         snake_case				:    Optional[int]           =	parent
         snake_case				:    Tuple           =	batch_size
         snake_case				:    int           =	seq_length
         snake_case				:    int           =	is_training
         snake_case				:    str           =	use_input_mask
         snake_case				:    Tuple           =	use_token_type_ids
         snake_case				:    Any           =	use_labels
         snake_case				:    List[Any]           =	vocab_size
         snake_case				:    Tuple           =	hidden_size
         snake_case				:    Any           =	num_hidden_layers
         snake_case				:    Any           =	num_attention_heads
         snake_case				:    Optional[Any]           =	intermediate_size
         snake_case				:    Any           =	hidden_act
         snake_case				:    Optional[Any]           =	hidden_dropout_prob
         snake_case				:    str           =	attention_probs_dropout_prob
         snake_case				:    Optional[int]           =	max_position_embeddings
         snake_case				:    Dict           =	type_vocab_size
         snake_case				:    str           =	type_sequence_label_size
         snake_case				:    List[Any]           =	initializer_range
         snake_case				:    List[Any]           =	num_labels
         snake_case				:    Union[str, Any]           =	num_choices
         snake_case				:    Dict           =	scope

      def      UpperCAmelCase				(				self     )		->    str:
         snake_case				:    List[Any]           =	ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size     )

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

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

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

         snake_case				:    Dict           =	self.get_config()

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

      def      UpperCAmelCase				(				self     )		->    Tuple:
         return LlamaConfig(
             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    ,				is_decoder=lowercase_    ,				initializer_range=self.initializer_range    ,				)

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A     )		->    List[str]:
         snake_case				:    Any           =	LlamaModel(config=lowercase_     )
         model.to(lowercase_     )
         model.eval()
         snake_case				:    Tuple           =	model(lowercase_    ,				attention_mask=lowercase_     )
         snake_case				:    Dict           =	model(lowercase_     )
         self.parent.assertEqual(result.last_hidden_state.shape    ,				(self.batch_size, self.seq_length, self.hidden_size)     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				)		->    List[str]:
         snake_case				:    Dict           =	True
         snake_case				:    Dict           =	LlamaModel(lowercase_     )
         model.to(lowercase_     )
         model.eval()
         snake_case				:    List[str]           =	model(
             lowercase_    ,				attention_mask=lowercase_    ,				encoder_hidden_states=lowercase_    ,				encoder_attention_mask=lowercase_    ,				)
         snake_case				:    str           =	model(
             lowercase_    ,				attention_mask=lowercase_    ,				encoder_hidden_states=lowercase_    ,				)
         snake_case				:    Optional[int]           =	model(lowercase_    ,				attention_mask=lowercase_     )
         self.parent.assertEqual(result.last_hidden_state.shape    ,				(self.batch_size, self.seq_length, self.hidden_size)     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				)		->    List[str]:
         snake_case				:    str           =	LlamaForCausalLM(config=lowercase_     )
         model.to(lowercase_     )
         model.eval()
         snake_case				:    Optional[Any]           =	model(lowercase_    ,				attention_mask=lowercase_    ,				labels=lowercase_     )
         self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size)     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				A    ,				)		->    int:
         snake_case				:    Dict           =	True
         snake_case				:    str           =	True
         snake_case				:    Dict           =	LlamaForCausalLM(config=lowercase_     )
         model.to(lowercase_     )
         model.eval()

         # first forward pass
         snake_case				:    Optional[Any]           =	model(
             lowercase_    ,				attention_mask=lowercase_    ,				encoder_hidden_states=lowercase_    ,				encoder_attention_mask=lowercase_    ,				use_cache=lowercase_    ,				)
         snake_case				:    Optional[Any]           =	outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         snake_case				:    Any           =	ids_tensor((self.batch_size, 3)    ,				config.vocab_size     )
         snake_case				:    Dict           =	ids_tensor((self.batch_size, 3)    ,				vocab_size=2     )

         # append to next input_ids and
         snake_case				:    List[str]           =	torch.cat([input_ids, next_tokens]    ,				dim=-1     )
         snake_case				:    List[Any]           =	torch.cat([input_mask, next_mask]    ,				dim=-1     )

         snake_case				:    List[str]           =	model(
             lowercase_    ,				attention_mask=lowercase_    ,				encoder_hidden_states=lowercase_    ,				encoder_attention_mask=lowercase_    ,				output_hidden_states=lowercase_    ,				)["""hidden_states"""][0]
         snake_case				:    int           =	model(
             lowercase_    ,				attention_mask=lowercase_    ,				encoder_hidden_states=lowercase_    ,				encoder_attention_mask=lowercase_    ,				past_key_values=lowercase_    ,				output_hidden_states=lowercase_    ,				)["""hidden_states"""][0]

         # select random slice
         snake_case				:    Any           =	ids_tensor((1,)    ,				output_from_past.shape[-1]     ).item()
         snake_case				:    int           =	output_from_no_past[:, -3:, random_slice_idx].detach()
         snake_case				:    List[Any]           =	output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]     )

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(lowercase_    ,				lowercase_    ,				atol=1e-3     )     )

      def      UpperCAmelCase				(				self     )		->    str:
         snake_case				:    Tuple           =	self.prepare_config_and_inputs()
         (
             (
             snake_case       
         )     , (
             snake_case       
         )     , (
             snake_case       
         )     , (
             snake_case       
         )     , (
             snake_case       
         )     , (
             snake_case       
         )     , (
             snake_case       
         )     ,        
         )				:    int           =	config_and_inputs
         snake_case				:    Optional[Any]           =	{"""input_ids""": input_ids, """attention_mask""": input_mask}
         return config, inputs_dict




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




      """simple docstring"""





      _snake_case												=						(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
      _snake_case												=						(LlamaForCausalLM,) if is_torch_available() else ()
      _snake_case												=						(
          {
              '''feature-extraction''': LlamaModel,
              '''text-classification''': LlamaForSequenceClassification,
              '''text-generation''': LlamaForCausalLM,
              '''zero-shot''': LlamaForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _snake_case												=						False
      _snake_case												=						False

      def      UpperCAmelCase				(				self     )		->    List[str]:
         snake_case				:    str           =	LlamaModelTester(self     )
         snake_case				:    List[str]           =	ConfigTester(self    ,				config_class=lowercase_    ,				hidden_size=3_7     )

      def      UpperCAmelCase				(				self     )		->    Any:
         self.config_tester.run_common_tests()

      def      UpperCAmelCase				(				self     )		->    int:
         snake_case				:    Any           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowercase_     )

      def      UpperCAmelCase				(				self     )		->    Optional[int]:
         snake_case				:    Optional[Any]           =	self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
            snake_case				:    List[str]           =	type
            self.model_tester.create_and_check_model(*lowercase_     )

      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case     , snake_case				:    Dict           =	self.model_tester.prepare_config_and_inputs_for_common()
         snake_case				:    Optional[int]           =	3
         snake_case				:    List[Any]           =	input_dict["""input_ids"""]
         snake_case				:    Optional[Any]           =	input_ids.ne(1     ).to(lowercase_     )
         snake_case				:    int           =	ids_tensor([self.model_tester.batch_size]    ,				self.model_tester.type_sequence_label_size     )
         snake_case				:    Any           =	LlamaForSequenceClassification(lowercase_     )
         model.to(lowercase_     )
         model.eval()
         snake_case				:    Optional[int]           =	model(lowercase_    ,				attention_mask=lowercase_    ,				labels=lowercase_     )
         self.assertEqual(result.logits.shape    ,				(self.model_tester.batch_size, self.model_tester.num_labels)     )

      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         snake_case     , snake_case				:    Tuple           =	self.model_tester.prepare_config_and_inputs_for_common()
         snake_case				:    Dict           =	3
         snake_case				:    Optional[Any]           =	"""single_label_classification"""
         snake_case				:    Any           =	input_dict["""input_ids"""]
         snake_case				:    Any           =	input_ids.ne(1     ).to(lowercase_     )
         snake_case				:    Dict           =	ids_tensor([self.model_tester.batch_size]    ,				self.model_tester.type_sequence_label_size     )
         snake_case				:    Tuple           =	LlamaForSequenceClassification(lowercase_     )
         model.to(lowercase_     )
         model.eval()
         snake_case				:    Tuple           =	model(lowercase_    ,				attention_mask=lowercase_    ,				labels=lowercase_     )
         self.assertEqual(result.logits.shape    ,				(self.model_tester.batch_size, self.model_tester.num_labels)     )

      def      UpperCAmelCase				(				self     )		->    List[str]:
         snake_case     , snake_case				:    Optional[int]           =	self.model_tester.prepare_config_and_inputs_for_common()
         snake_case				:    List[Any]           =	3
         snake_case				:    Optional[Any]           =	"""multi_label_classification"""
         snake_case				:    Union[str, Any]           =	input_dict["""input_ids"""]
         snake_case				:    List[str]           =	input_ids.ne(1     ).to(lowercase_     )
         snake_case				:    List[Any]           =	ids_tensor(
             [self.model_tester.batch_size, config.num_labels]    ,				self.model_tester.type_sequence_label_size     ).to(torch.float     )
         snake_case				:    Optional[int]           =	LlamaForSequenceClassification(lowercase_     )
         model.to(lowercase_     )
         model.eval()
         snake_case				:    List[Any]           =	model(lowercase_    ,				attention_mask=lowercase_    ,				labels=lowercase_     )
         self.assertEqual(result.logits.shape    ,				(self.model_tester.batch_size, self.model_tester.num_labels)     )

      @unittest.skip("""LLaMA buffers include complex numbers, which breaks this test"""     )
      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         pass

      @parameterized.expand([("""linear""",), ("""dynamic""",)]     )
      def      UpperCAmelCase				(				self    ,				A     )		->    int:
         snake_case     , snake_case				:    Optional[Any]           =	self.model_tester.prepare_config_and_inputs_for_common()
         snake_case				:    Tuple           =	ids_tensor([1, 1_0]    ,				config.vocab_size     )
         snake_case				:    str           =	ids_tensor([1, int(config.max_position_embeddings * 1.5     )]    ,				config.vocab_size     )

         set_seed(4_2     )  # Fixed seed at init time so the two models get the same random weights
         snake_case				:    Tuple           =	LlamaModel(lowercase_     )
         original_model.to(lowercase_     )
         original_model.eval()
         snake_case				:    Tuple           =	original_model(lowercase_     ).last_hidden_state
         snake_case				:    Optional[int]           =	original_model(lowercase_     ).last_hidden_state

         set_seed(4_2     )  # Fixed seed at init time so the two models get the same random weights
         snake_case				:    Dict           =	{"""type""": scaling_type, """factor""": 10.0}
         snake_case				:    Optional[int]           =	LlamaModel(lowercase_     )
         scaled_model.to(lowercase_     )
         scaled_model.eval()
         snake_case				:    List[Any]           =	scaled_model(lowercase_     ).last_hidden_state
         snake_case				:    Optional[int]           =	scaled_model(lowercase_     ).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
            self.assertTrue(torch.allclose(lowercase_    ,				lowercase_    ,				atol=1e-5     )     )
         else:
            self.assertFalse(torch.allclose(lowercase_    ,				lowercase_    ,				atol=1e-5     )     )

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(lowercase_    ,				lowercase_    ,				atol=1e-5     )     )




@require_torch
class 						__lowercase		(unittest.TestCase							):




      """simple docstring"""





      @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""     )
      @slow
      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case				:    Dict           =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
         snake_case				:    Any           =	LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-7b-hf"""    ,				device_map="""auto"""     )
         snake_case				:    Tuple           =	model(torch.tensor([input_ids]     )     )
         # Expected mean on dim = -1
         snake_case				:    Any           =	torch.tensor([[-6.65_50, -4.12_27, -4.98_59, -3.24_06, 0.82_62, -3.00_33, 1.29_64, -3.36_99]]     )
         torch.testing.assert_close(out.mean(-1     )    ,				lowercase_    ,				atol=1e-2    ,				rtol=1e-2     )
         # slicing logits[0, 0, 0:30]
         # fmt: off
         snake_case				:    List[str]           =	torch.tensor([-12.82_81, -7.44_53, -0.46_39, -8.06_25, -7.25_00, -8.00_00, -6.48_83, -7.76_95, -7.84_38, -7.03_12, -6.21_88, -7.13_28, -1.84_96, 1.99_61, -8.62_50, -6.72_27, -12.82_81, -6.94_92, -7.07_42, -7.78_52, -7.58_20, -7.90_62, -6.93_75, -7.98_05, -8.34_38, -8.15_62, -8.04_69, -7.62_50, -7.74_22, -7.33_98,]     )
         # fmt: on
         torch.testing.assert_close(out[0, 0, :3_0]    ,				lowercase_    ,				atol=1e-5    ,				rtol=1e-5     )

      @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""     )
      @slow
      def      UpperCAmelCase				(				self     )		->    Optional[int]:
         snake_case				:    List[str]           =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
         snake_case				:    Optional[int]           =	LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-13b-hf"""    ,				device_map="""auto"""     )
         snake_case				:    Union[str, Any]           =	model(torch.tensor(lowercase_     )     )
         # Expected mean on dim = -1
         snake_case				:    int           =	torch.tensor([[-2.06_22, -1.27_94, -1.16_38, -0.97_88, -1.46_03, -1.02_38, -1.78_93, -1.44_11]]     )
         torch.testing.assert_close(out.mean(-1     )    ,				lowercase_    ,				atol=1e-2    ,				rtol=1e-2     )
         # slicing logits[0, 0, 0:30]
         # fmt: off
         snake_case				:    Optional[int]           =	torch.tensor([-8.14_06, -8.05_47, 2.74_61, -1.23_44, -0.14_48, -1.82_62, -1.00_20, -1.81_54, -1.68_95, -1.85_16, -2.35_74, -0.92_77, 3.75_98, 6.57_42, -1.29_98, -0.11_77, -8.14_06, -2.96_88, -2.91_99, -3.16_99, -3.52_54, -2.35_55, -2.79_88, -3.41_41, -2.82_62, -4.51_95, -3.33_79, -3.31_64, -2.78_32, -3.02_73]     )
         # fmt: on
         torch.testing.assert_close(out[0, 0, :3_0]    ,				lowercase_    ,				atol=1e-5    ,				rtol=1e-5     )

      @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""     )
      @slow
      def      UpperCAmelCase				(				self     )		->    str:
         snake_case				:    int           =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
         snake_case				:    Any           =	LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-13b-chat-hf"""    ,				device_map="""auto"""     )
         snake_case				:    List[Any]           =	model(torch.tensor(lowercase_     )     )
         # Expected mean on dim = -1
         snake_case				:    List[str]           =	torch.tensor([[-0.85_62, -1.85_20, -0.75_51, -0.41_62, -1.51_61, -1.20_38, -2.48_23, -2.32_54]]     )
         torch.testing.assert_close(out.mean(-1     )    ,				lowercase_    ,				atol=1e-2    ,				rtol=1e-2     )
         # slicing logits[0, 0, 0:30]
         # fmt: off
         snake_case				:    Optional[Any]           =	torch.tensor([-2.22_27, 4.88_28, 0.90_23, -0.45_78, -0.78_71, -0.10_33, -0.62_21, -0.57_86, -0.78_03, -1.06_74, -1.29_20, -0.15_70, 0.80_08, 2.07_23, -0.94_97, 0.27_71, -2.22_27, -0.76_12, -1.43_46, -1.20_61, -1.64_26, -0.30_00, -0.71_39, -1.19_34, -1.86_91, -1.69_73, -1.59_47, -1.27_05, -0.35_23, -0.55_13]     )
         # fmt: on
         torch.testing.assert_close(out.mean(-1     )    ,				lowercase_    ,				atol=1e-2    ,				rtol=1e-2     )

      @unittest.skip(
          """Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test"""     )
      @slow
      def      UpperCAmelCase				(				self     )		->    Any:
         snake_case				:    List[Any]           =	[1, 3_0_6, 4_6_5_8, 2_7_8, 6_5_9_3, 3_1_0, 2_8_3_4, 3_3_8]
         snake_case				:    Any           =	LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-70b-hf"""    ,				device_map="""auto"""     )
         snake_case				:    Any           =	model(torch.tensor(lowercase_     )     )

         snake_case				:    Any           =	torch.tensor(
             [[-4.23_27, -3.33_60, -4.66_65, -4.76_31, -1.81_80, -3.41_70, -1.42_11, -3.18_10]]    ,				dtype=torch.floataa     )
         torch.testing.assert_close(out.mean(-1     )    ,				lowercase_    ,				atol=1e-2    ,				rtol=1e-2     )
         # fmt: off
         snake_case				:    Any           =	torch.tensor([-9.49_22, -3.95_51, 1.79_98, -5.67_58, -5.10_55, -5.89_84, -4.83_20, -6.80_86, -6.53_91, -5.61_72, -5.58_20, -5.53_52, 1.78_81, 3.62_89, -6.51_17, -3.47_85, -9.50_00, -6.03_52, -6.81_25, -6.01_95, -6.68_36, -5.47_27, -6.28_12, -6.03_91, -7.33_98, -7.42_97, -7.48_44, -6.58_20, -5.87_89, -5.53_12]     )
         # fmt: on
         torch.testing.assert_close(out[0, 0, :3_0]    ,				lowercase_    ,				atol=1e-5    ,				rtol=1e-5     )

      @unittest.skip("""Model is curently gated"""     )
      @slow
      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case				:    int           =	"""Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the \"princi"""
         snake_case				:    List[str]           =	"""Simply put, the theory of relativity states that """
         snake_case				:    List[str]           =	LlamaTokenizer.from_pretrained("""meta-llama/Llama-2-13b-chat-hf"""     )
         snake_case				:    str           =	tokenizer.encode(lowercase_    ,				return_tensors="""pt"""     )
         snake_case				:    List[Any]           =	LlamaForCausalLM.from_pretrained(
             """meta-llama/Llama-2-13b-chat-hf"""    ,				device_map="""sequential"""    ,				use_safetensors=lowercase_     )

         # greedy generation outputs
         snake_case				:    Optional[int]           =	model.generate(lowercase_    ,				max_new_tokens=6_4    ,				top_p=lowercase_    ,				temperature=1    ,				do_sample=lowercase_     )
         snake_case				:    Tuple           =	tokenizer.decode(generated_ids[0]    ,				skip_special_tokens=lowercase_     )
         self.assertEqual(lowercase_    ,				lowercase_     )











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




      """simple docstring"""





      _snake_case												=						ConsistencyModelPipeline
      _snake_case												=						UNCONDITIONAL_IMAGE_GENERATION_PARAMS
      _snake_case												=						UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

      @property
      def      UpperCAmelCase				(				self     )		->    List[str]:
         snake_case				:    Dict           =	UNetaDModel.from_pretrained(
             """diffusers/consistency-models-test"""    ,				subfolder="""test_unet"""    ,				)
         return unet

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

      def      UpperCAmelCase				(				self    ,				A=False     )		->    Optional[Any]:
         if class_cond:
            snake_case				:    List[str]           =	self.dummy_cond_unet
         else:
            snake_case				:    Optional[Any]           =	self.dummy_uncond_unet

         # Default to CM multistep sampler
         snake_case				:    Optional[Any]           =	CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    ,				sigma_min=0.0_02    ,				sigma_max=80.0    ,				)

         snake_case				:    Dict           =	{
             """unet""": unet,
             """scheduler""": scheduler,
         }

         return components

      def      UpperCAmelCase				(				self    ,				A    ,				A=0     )		->    Optional[int]:
         if str(A     ).startswith("""mps"""     ):
            snake_case				:    Union[str, Any]           =	torch.manual_seed(A     )
         else:
            snake_case				:    Dict           =	torch.Generator(device=A     ).manual_seed(A     )

         snake_case				:    Tuple           =	{
             """batch_size""": 1,
             """num_inference_steps""": None,
             """timesteps""": [2_2, 0],
             """generator""": generator,
             """output_type""": """np""",
         }

         return inputs

      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case				:    Union[str, Any]           =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         snake_case				:    List[Any]           =	self.get_dummy_components()
         snake_case				:    Dict           =	ConsistencyModelPipeline(**A     )
         snake_case				:    int           =	pipe.to(A     )
         pipe.set_progress_bar_config(disable=A     )

         snake_case				:    Any           =	self.get_dummy_inputs(A     )
         snake_case				:    Tuple           =	pipe(**A     ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         snake_case				:    Dict           =	image[0, -3:, -3:, -1]
         snake_case				:    Tuple           =	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      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case				:    Optional[int]           =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         snake_case				:    Tuple           =	self.get_dummy_components(class_cond=A     )
         snake_case				:    Dict           =	ConsistencyModelPipeline(**A     )
         snake_case				:    List[Any]           =	pipe.to(A     )
         pipe.set_progress_bar_config(disable=A     )

         snake_case				:    str           =	self.get_dummy_inputs(A     )
         snake_case				:    Optional[int]           =	0
         snake_case				:    List[Any]           =	pipe(**A     ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         snake_case				:    Optional[int]           =	image[0, -3:, -3:, -1]
         snake_case				:    Dict           =	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      UpperCAmelCase				(				self     )		->    Dict:
         snake_case				:    str           =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         snake_case				:    Optional[int]           =	self.get_dummy_components()
         snake_case				:    List[str]           =	ConsistencyModelPipeline(**A     )
         snake_case				:    Dict           =	pipe.to(A     )
         pipe.set_progress_bar_config(disable=A     )

         snake_case				:    Dict           =	self.get_dummy_inputs(A     )
         snake_case				:    Tuple           =	1
         snake_case				:    Optional[Any]           =	None
         snake_case				:    List[Any]           =	pipe(**A     ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         snake_case				:    Optional[Any]           =	image[0, -3:, -3:, -1]
         snake_case				:    List[str]           =	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      UpperCAmelCase				(				self     )		->    Optional[Any]:
         snake_case				:    Optional[int]           =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         snake_case				:    Dict           =	self.get_dummy_components(class_cond=A     )
         snake_case				:    List[Any]           =	ConsistencyModelPipeline(**A     )
         snake_case				:    List[Any]           =	pipe.to(A     )
         pipe.set_progress_bar_config(disable=A     )

         snake_case				:    Optional[Any]           =	self.get_dummy_inputs(A     )
         snake_case				:    Optional[Any]           =	1
         snake_case				:    Any           =	None
         snake_case				:    Optional[int]           =	0
         snake_case				:    List[Any]           =	pipe(**A     ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         snake_case				:    Optional[int]           =	image[0, -3:, -3:, -1]
         snake_case				:    List[str]           =	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 						__lowercase		(unittest.TestCase							):




      """simple docstring"""





      def      UpperCAmelCase				(				self     )		->    Dict:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()

      def      UpperCAmelCase				(				self    ,				A=0    ,				A=False    ,				A="cpu"    ,				A=torch.floataa    ,				A=(1, 3, 6_4, 6_4)     )		->    int:
         snake_case				:    Union[str, Any]           =	torch.manual_seed(A     )

         snake_case				:    Tuple           =	{
             """num_inference_steps""": None,
             """timesteps""": [2_2, 0],
             """class_labels""": 0,
             """generator""": generator,
             """output_type""": """np""",
         }

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

         return inputs

      def      UpperCAmelCase				(				self    ,				A=0    ,				A="cpu"    ,				A=torch.floataa    ,				A=(1, 3, 6_4, 6_4)     )		->    Any:
         if type(A     ) == str:
            snake_case				:    List[str]           =	torch.device(A     )
         snake_case				:    Any           =	torch.Generator(device=A     ).manual_seed(A     )
         snake_case				:    Dict           =	randn_tensor(A    ,				generator=A    ,				device=A    ,				dtype=A     )
         return latents

      def      UpperCAmelCase				(				self     )		->    Dict:
         snake_case				:    List[Any]           =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""    ,				subfolder="""diffusers_cd_imagenet64_l2"""     )
         snake_case				:    Optional[int]           =	CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    ,				sigma_min=0.0_02    ,				sigma_max=80.0    ,				)
         snake_case				:    Any           =	ConsistencyModelPipeline(unet=A    ,				scheduler=A     )
         pipe.to(torch_device=A     )
         pipe.set_progress_bar_config(disable=A     )

         snake_case				:    Tuple           =	self.get_inputs()
         snake_case				:    List[str]           =	pipe(**A     ).images
         assert image.shape == (1, 6_4, 6_4, 3)

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

         snake_case				:    Union[str, Any]           =	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      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    Tuple           =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""    ,				subfolder="""diffusers_cd_imagenet64_l2"""     )
         snake_case				:    int           =	CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    ,				sigma_min=0.0_02    ,				sigma_max=80.0    ,				)
         snake_case				:    List[str]           =	ConsistencyModelPipeline(unet=A    ,				scheduler=A     )
         pipe.to(torch_device=A     )
         pipe.set_progress_bar_config(disable=A     )

         snake_case				:    Union[str, Any]           =	self.get_inputs()
         snake_case				:    Tuple           =	1
         snake_case				:    List[str]           =	None
         snake_case				:    Any           =	pipe(**A     ).images
         assert image.shape == (1, 6_4, 6_4, 3)

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

         snake_case				:    Union[str, Any]           =	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      UpperCAmelCase				(				self     )		->    Optional[int]:
         snake_case				:    List[Any]           =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""    ,				subfolder="""diffusers_cd_imagenet64_l2"""     )
         snake_case				:    Optional[Any]           =	CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    ,				sigma_min=0.0_02    ,				sigma_max=80.0    ,				)
         snake_case				:    List[Any]           =	ConsistencyModelPipeline(unet=A    ,				scheduler=A     )
         pipe.to(torch_device=A    ,				torch_dtype=torch.floataa     )
         pipe.set_progress_bar_config(disable=A     )

         snake_case				:    List[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     ):
            snake_case				:    Tuple           =	pipe(**A     ).images
         assert image.shape == (1, 6_4, 6_4, 3)

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

         snake_case				:    int           =	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      UpperCAmelCase				(				self     )		->    Optional[Any]:
         snake_case				:    Optional[Any]           =	UNetaDModel.from_pretrained("""diffusers/consistency_models"""    ,				subfolder="""diffusers_cd_imagenet64_l2"""     )
         snake_case				:    Any           =	CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    ,				sigma_min=0.0_02    ,				sigma_max=80.0    ,				)
         snake_case				:    Dict           =	ConsistencyModelPipeline(unet=A    ,				scheduler=A     )
         pipe.to(torch_device=A    ,				torch_dtype=torch.floataa     )
         pipe.set_progress_bar_config(disable=A     )

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

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

         snake_case				:    Dict           =	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





