


import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
UpperCAmelCase   :    Union[str, Any]     =		logging.get_logger(__name__)

UpperCAmelCase   :    Optional[int]     =		{
    "post_extract_proj": "feature_projection.projection",
    "encoder.pos_conv.0": "encoder.pos_conv_embed.conv",
    "self_attn.k_proj": "encoder.layers.*.attention.k_proj",
    "self_attn.v_proj": "encoder.layers.*.attention.v_proj",
    "self_attn.q_proj": "encoder.layers.*.attention.q_proj",
    "self_attn.out_proj": "encoder.layers.*.attention.out_proj",
    "self_attn_layer_norm": "encoder.layers.*.layer_norm",
    "fc1": "encoder.layers.*.feed_forward.intermediate_dense",
    "fc2": "encoder.layers.*.feed_forward.output_dense",
    "final_layer_norm": "encoder.layers.*.final_layer_norm",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "feature_projection.layer_norm",
    "quantizer.weight_proj": "quantizer.weight_proj",
    "quantizer.vars": "quantizer.codevectors",
    "project_q": "project_q",
    "final_proj": "project_hid",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}
UpperCAmelCase   :    List[str]     =		[
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]






def 	__lowerCamelCase						(							lowerCamelCase__		:						str  ,							lowerCamelCase__		:						Tuple  ,							lowerCamelCase__		:						str  ,							lowerCamelCase__		:						List[str]  ,							lowerCamelCase__		:						Dict      ):





			'''simple docstring'''




			for attribute in key.split("""."""      ):
						lowerCamelCase    		=  getattr(lowerCamelCase__  ,							lowerCamelCase__      )

			if weight_type is not None:
						lowerCamelCase    		=  getattr(lowerCamelCase__  ,							lowerCamelCase__      ).shape
			else:
						lowerCamelCase    		=  hf_pointer.shape

			assert hf_shape == value.shape, (
			    f'Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'
			    f' {value.shape} for {full_name}'
			)

			if weight_type == "weight":
						lowerCamelCase    		=  value
			elif weight_type == "weight_g":
						lowerCamelCase    		=  value
			elif weight_type == "weight_v":
						lowerCamelCase    		=  value
			elif weight_type == "bias":
						lowerCamelCase    		=  value
			else:
						lowerCamelCase    		=  value

			logger.info(f'{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'      )






def 	__lowerCamelCase						(							lowerCamelCase__		:						Optional[Any]  ,							lowerCamelCase__		:						Dict      ):





			'''simple docstring'''




			lowerCamelCase    		=  []
			lowerCamelCase    		=  fairseq_model.state_dict()

			lowerCamelCase    		=  hf_model.feature_extractor
			lowerCamelCase    		=  hf_model.adapter

			for name, value in fairseq_dict.items():
						lowerCamelCase    		=  False
						if "conv_layers" in name:
									load_conv_layer(
									    lowerCamelCase__  ,							lowerCamelCase__  ,							lowerCamelCase__  ,							lowerCamelCase__  ,							hf_model.config.feat_extract_norm == """group"""  ,							)
									lowerCamelCase    		=  True
						elif any(x in name for x in ["""adaptor""", """w2v_encoder.proj.""", """w2v_proj_ln."""]      ):
									load_adapter(lowerCamelCase__  ,							lowerCamelCase__  ,							lowerCamelCase__  ,							lowerCamelCase__      )
									lowerCamelCase    		=  True
						else:
									for key, mapped_key in MAPPING.items():
												if key in name or key.split("""w2v_model."""      )[-1] == name.split("""."""      )[0]:
															lowerCamelCase    		=  True
															if "*" in mapped_key:
																		lowerCamelCase    		=  name.split(lowerCamelCase__      )[0].split("""."""      )[-2]
																		lowerCamelCase    		=  mapped_key.replace("""*"""  ,							lowerCamelCase__      )
															if "weight_g" in name:
																		lowerCamelCase    		=  """weight_g"""
															elif "weight_v" in name:
																		lowerCamelCase    		=  """weight_v"""
															elif "bias" in name:
																		lowerCamelCase    		=  """bias"""
															elif "weight" in name:
																		lowerCamelCase    		=  """weight"""
															else:
																		lowerCamelCase    		=  None
															set_recursively(lowerCamelCase__  ,							lowerCamelCase__  ,							lowerCamelCase__  ,							lowerCamelCase__  ,							lowerCamelCase__      )
												continue
						if not is_used:
									unused_weights.append(lowerCamelCase__      )

			logger.warning(f'Unused weights: {unused_weights}'      )






def 	__lowerCamelCase						(							lowerCamelCase__		:						List[Any]  ,							lowerCamelCase__		:						Tuple  ,							lowerCamelCase__		:						Tuple  ,							lowerCamelCase__		:						Union[str, Any]  ,							lowerCamelCase__		:						Any      ):





			'''simple docstring'''




			lowerCamelCase    		=  full_name.split("""conv_layers."""      )[-1]
			lowerCamelCase    		=  name.split("""."""      )
			lowerCamelCase    		=  int(items[0]      )
			lowerCamelCase    		=  int(items[1]      )

			if type_id == 0:
						if "bias" in name:
									assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
									    f'{full_name} has size {value.shape}, but'
									    f' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'
									)
									lowerCamelCase    		=  value
									logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.'      )
						elif "weight" in name:
									assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
									    f'{full_name} has size {value.shape}, but'
									    f' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'
									)
									lowerCamelCase    		=  value
									logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.'      )
			elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
						if "bias" in name:
									assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
									    f'{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'
									    " found."
									)
									lowerCamelCase    		=  value
									logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'      )
						elif "weight" in name:
									assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
									    f'{full_name} has size {value.shape}, but'
									    f' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'
									)
									lowerCamelCase    		=  value
									logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'      )
			else:
						unused_weights.append(lowerCamelCase__      )






def 	__lowerCamelCase						(							lowerCamelCase__		:						Tuple  ,							lowerCamelCase__		:						Any  ,							lowerCamelCase__		:						Any  ,							lowerCamelCase__		:						List[str]      ):





			'''simple docstring'''




			lowerCamelCase    		=  full_name.split("""adaptor."""      )[-1]
			lowerCamelCase    		=  name.split("""."""      )

			if items[1].isdigit():
						lowerCamelCase    		=  int(items[1]      )
			else:
						lowerCamelCase    		=  None

			if "adaptor" not in full_name:
						if "proj_ln" in full_name:
									# has to be layer norm
									if "bias" in name:
												assert (
												    value.shape == adapter.proj_layer_norm.bias.data.shape
												), f'{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'
												lowerCamelCase    		=  value
												logger.info(f'Adapter proj layer norm bias was initialized from {full_name}.'      )
									if "weight" in name:
												assert (
												    value.shape == adapter.proj_layer_norm.weight.data.shape
												), f'{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'
												lowerCamelCase    		=  value
						else:
									# has to be projection layer
									if "bias" in name:
												assert (
												    value.shape == adapter.proj.bias.data.shape
												), f'{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'
												lowerCamelCase    		=  value
												logger.info(f'Adapter proj layer bias was initialized from {full_name}.'      )
									if "weight" in name:
												assert (
												    value.shape == adapter.proj.weight.data.shape
												), f'{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'
												lowerCamelCase    		=  value
												logger.info(f'Adapter proj layer weight was initialized from {full_name}.'      )
			elif isinstance(lowerCamelCase__  ,							lowerCamelCase__      ):
						if "bias" in name:
									assert (
									    value.shape == adapter.layers[layer_id].conv.bias.data.shape
									), f'{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'
									lowerCamelCase    		=  value
									logger.info(f'Adapter layer {layer_id} bias was initialized from {full_name}.'      )
						elif "weight" in name:
									assert (
									    value.shape == adapter.layers[layer_id].conv.weight.data.shape
									), f'{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'
									lowerCamelCase    		=  value
									logger.info(f'Adapter layer {layer_id} bias was initialized from {full_name}.'      )
			else:
						unused_weights.append(lowerCamelCase__      )






def 	__lowerCamelCase						(							lowerCamelCase__		:						str      ):





			'''simple docstring'''




			lowerCamelCase		,		lowerCamelCase    		=  emb.weight.shape
			lowerCamelCase    		=  nn.Linear(lowerCamelCase__  ,							lowerCamelCase__  ,							bias=lowerCamelCase__      )
			lowerCamelCase    		=  emb.weight.data
			return lin_layer






@torch.no_grad()
def 	__lowerCamelCase						(							lowerCamelCase__		:						str  ,							lowerCamelCase__		:						List[str]  ,							lowerCamelCase__		:						List[str]  ,							lowerCamelCase__		:						Any  ,							lowerCamelCase__		:						List[Any]  ,							lowerCamelCase__		:						Dict  ,							lowerCamelCase__		:						List[str]  ,							lowerCamelCase__		:						Optional[int]  ,							lowerCamelCase__		:						Optional[Any]  ,							lowerCamelCase__		:						List[Any]  ,							lowerCamelCase__		:						Optional[Any]  ,							):





			'''simple docstring'''




			lowerCamelCase    		=  WavaVecaConfig.from_pretrained(
			    lowerCamelCase__  ,							add_adapter=lowerCamelCase__  ,							adapter_stride=lowerCamelCase__  ,							adapter_kernel_size=lowerCamelCase__  ,							use_auth_token=lowerCamelCase__  ,							output_hidden_size=lowerCamelCase__  ,							)
			lowerCamelCase    		=  MBartConfig.from_pretrained(lowerCamelCase__      )

			# load model
			lowerCamelCase		,		lowerCamelCase		,		lowerCamelCase    		=  fairseq.checkpoint_utils.load_model_ensemble_and_task(
			    [checkpoint_path]  ,							arg_overrides={
			        """config_yaml""": config_yaml_path,
			        """data""": """/""".join(dict_path.split("""/"""      )[:-1]      ),
			        """w2v_path""": checkpoint_path,
			        """load_pretrained_decoder_from""": None,
			    }  ,							)
			lowerCamelCase    		=  model[0].eval()

			# load feature extractor
			lowerCamelCase    		=  WavaVecaFeatureExtractor.from_pretrained(lowerCamelCase__  ,							use_auth_token=lowerCamelCase__      )

			# set weights for wav2vec2 encoder
			lowerCamelCase    		=  WavaVecaModel(lowerCamelCase__      )

			recursively_load_weights_wavaveca(model.encoder  ,							lowerCamelCase__      )

			# load decoder weights
			lowerCamelCase    		=  MBartForCausalLM(lowerCamelCase__      )
			lowerCamelCase		,		lowerCamelCase    		=  hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()  ,							strict=lowerCamelCase__      )
			logger.warning(f'The following keys are missing when loading the decoder weights: {missing_keys}'      )
			logger.warning(f'The following keys are unexpected when loading the decoder weights: {unexpected_keys}'      )

			lowerCamelCase    		=  SpeechEncoderDecoderModel(encoder=lowerCamelCase__  ,							decoder=lowerCamelCase__      )
			lowerCamelCase    		=  False

			lowerCamelCase    		=  MBartaaTokenizer(lowerCamelCase__      )
			tokenizer.save_pretrained(lowerCamelCase__      )

			lowerCamelCase    		=  hf_wavavec.config.to_dict()
			lowerCamelCase    		=  tokenizer.pad_token_id
			lowerCamelCase    		=  tokenizer.bos_token_id
			lowerCamelCase    		=  tokenizer.eos_token_id
			lowerCamelCase    		=  """mbart50"""
			lowerCamelCase    		=  """wav2vec2"""

			lowerCamelCase    		=  tokenizer.eos_token_id
			lowerCamelCase    		=  250004
			lowerCamelCase    		=  tokenizer.eos_token_id

			lowerCamelCase    		=  SpeechEncoderDecoderConfig.from_dict(lowerCamelCase__      )

			hf_wavavec.save_pretrained(lowerCamelCase__      )
			feature_extractor.save_pretrained(lowerCamelCase__      )


if __name__ == "__main__":
		UpperCAmelCase   :    Dict     =		argparse.ArgumentParser()
		parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
		parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
		parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model")
		parser.add_argument("--config_yaml_path", default=None, type=str, help="Path to yaml file of fine-tuned model")
		parser.add_argument(
		    "--encoder_config_path",
		    default="facebook/wav2vec2-xls-r-1b",
		    type=str,
		    help="Path to hf encoder wav2vec2 checkpoint config",
		)
		parser.add_argument(
		    "--decoder_config_path",
		    default="facebook/mbart-large-50-one-to-many-mmt",
		    type=str,
		    help="Path to hf decoder checkpoint config",
		)
		parser.add_argument("--add_adapter", default=True, type=bool, help="whethere to add model adapter layers")
		parser.add_argument("--adapter_stride", default=2, type=int, help="stride of adapter layers")
		parser.add_argument("--adapter_kernel_size", default=3, type=int, help="kernel size of adapter layers")
		parser.add_argument("--encoder_output_dim", default=10_24, type=int, help="encoder output dim")
		parser.add_argument("--start_token_id", default=25_00_04, type=int, help="`decoder_start_token_id` of model config")

		UpperCAmelCase   :    Dict     =		parser.parse_args()
		convert_wavaveca_checkpoint(
		    args.checkpoint_path,
		    args.pytorch_dump_folder_path,
		    args.dict_path,
		    args.config_yaml_path,
		    encoder_config_path=args.encoder_config_path,
		    decoder_config_path=args.decoder_config_path,
		    add_adapter=args.add_adapter,
		    adapter_kernel_size=args.adapter_kernel_size,
		    adapter_stride=args.adapter_stride,
		    decoder_start_token_id=args.start_token_id,
		    encoder_output_dim=args.encoder_output_dim,
		)





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

