import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class        SCREAMING_SNAKE_CASE__					(     A__							,unittest.TestCase ):





	'''simple docstring'''



	_UpperCAmelCase	:		List[Any]          = ConsistencyModelPipeline
	_UpperCAmelCase	:		Tuple          = UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	_UpperCAmelCase	:		Union[str, Any]          = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



	@property
	def 			A							(  self     :   Any     ):


				'''simple docstring'''


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



	@property
	def 			A							(  self     :   int     ):


				'''simple docstring'''


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



	def 			A							(  self     :   int      ,  lowercase     :   str=False     ):


				'''simple docstring'''


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

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

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

				return components



	def 			A							(  self     :   Union[str, Any]      ,  lowercase     :   int      ,  lowercase     :   Union[str, Any]=0     ):


				'''simple docstring'''


				if str(__A     ).startswith('mps'     ):
							_snake_case          =      torch.manual_seed(__A     )
				else:
							_snake_case          =      torch.Generator(device=__A     ).manual_seed(__A     )

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

				return inputs



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


				'''simple docstring'''


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

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

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


				'''simple docstring'''


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

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

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


				'''simple docstring'''


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

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

				_snake_case          =      image[0, -3:, -3:, -1]
				_snake_case          =      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     :   Optional[int]     ):


				'''simple docstring'''


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

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

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





	'''simple docstring'''



	def 			A							(  self     :   Tuple     ):


				'''simple docstring'''


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



	def 			A							(  self     :   List[Any]      ,  lowercase     :   Any=0      ,  lowercase     :   List[str]=False      ,  lowercase     :   Dict="cpu"      ,  lowercase     :   str=torch.floataa      ,  lowercase     :   Union[str, Any]=(1, 3, 64, 64)     ):


				'''simple docstring'''


				_snake_case          =      torch.manual_seed(__A     )

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

				if get_fixed_latents:
							_snake_case          =      self.get_fixed_latents(seed=__A      ,  device=__A      ,  dtype=__A      ,  shape=__A     )
							_snake_case          =      latents

				return inputs



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


				'''simple docstring'''


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



	def 			A							(  self     :   List[Any]     ):


				'''simple docstring'''


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

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

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

				_snake_case          =      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     :   Dict     ):


				'''simple docstring'''


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

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

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

				_snake_case          =      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     :   int     ):


				'''simple docstring'''


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

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

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

				_snake_case          =      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     :   Union[str, Any]     ):


				'''simple docstring'''


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

				_snake_case          =      self.get_inputs(get_fixed_latents=__A      ,  device=__A     )
				_snake_case          =      1
				_snake_case          =      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          =      pipe(**__A     ).images
				assert image.shape == (1, 64, 64, 3)

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

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

import argparse
import glob
import logging
import os
import time
from argparse import Namespace

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from torch.utils.data import DataLoader, TensorDataset

from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import glue_output_modes, glue_tasks_num_labels
from transformers import glue_processors as processors


_lowerCamelCase	:    str          =      logging.getLogger(__name__)




class        SCREAMING_SNAKE_CASE__					(     UpperCAmelCase ):





	'''simple docstring'''



	_UpperCAmelCase	:		Any          = "sequence-classification"



	def __init__(  self     :   Optional[int]      ,  lowercase     :   Optional[Any]     ):


				'''simple docstring'''


				if type(lowercase     ) == dict:
							_snake_case          =      Namespace(**lowercase     )
				_snake_case          =      glue_output_modes[hparams.task]
				_snake_case          =      glue_tasks_num_labels[hparams.task]

				super().__init__(lowercase      ,  lowercase      ,  self.mode     )



	def 			A							(  self     :   List[str]      ,  **lowercase     :   Optional[Any]     ):


				'''simple docstring'''


				return self.model(**lowercase     )



	def 			A							(  self     :   Tuple      ,  lowercase     :   Optional[Any]      ,  lowercase     :   int     ):


				'''simple docstring'''


				_snake_case          =      {'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}

				if self.config.model_type not in ["distilbert", "bart"]:
							_snake_case          =      batch[2] if self.config.model_type in ['bert', 'xlnet', 'albert'] else None

				_snake_case          =      self(**lowercase     )
				_snake_case          =      outputs[0]

				_snake_case          =      self.trainer.lr_schedulers[0]['scheduler']
				_snake_case          =      {'loss': loss, 'rate': lr_scheduler.get_last_lr()[-1]}
				return {"loss": loss, "log": tensorboard_logs}



	def 			A							(  self     :   str     ):


				'''simple docstring'''


				_snake_case          =      self.hparams
				_snake_case          =      processors[args.task]()
				_snake_case          =      processor.get_labels()

				for mode in ["train", "dev"]:
							_snake_case          =      self._feature_file(lowercase     )
							if os.path.exists(lowercase     ) and not args.overwrite_cache:
										logger.info('Loading features from cached file %s'      ,  lowercase     )
							else:
										logger.info('Creating features from dataset file at %s'      ,  args.data_dir     )
										_snake_case          =      (
										    processor.get_dev_examples(args.data_dir     )
										    if mode == 'dev'
										    else processor.get_train_examples(args.data_dir     )
										)
										_snake_case          =      convert_examples_to_features(
										    lowercase      ,  self.tokenizer      ,  max_length=args.max_seq_length      ,  label_list=self.labels      ,  output_mode=args.glue_output_mode      ,  )
										logger.info('Saving features into cached file %s'      ,  lowercase     )
										torch.save(lowercase      ,  lowercase     )



	def 			A							(  self     :   List[str]      ,  lowercase     :   str      ,  lowercase     :   int      ,  lowercase     :   bool = False     ):


				'''simple docstring'''


				_snake_case          =      'dev' if mode == 'test' else mode

				_snake_case          =      self._feature_file(lowercase     )
				logger.info('Loading features from cached file %s'      ,  lowercase     )
				_snake_case          =      torch.load(lowercase     )
				_snake_case          =      torch.tensor([f.input_ids for f in features]      ,  dtype=torch.long     )
				_snake_case          =      torch.tensor([f.attention_mask for f in features]      ,  dtype=torch.long     )
				_snake_case          =      torch.tensor([f.token_type_ids for f in features]      ,  dtype=torch.long     )
				if self.hparams.glue_output_mode == "classification":
							_snake_case          =      torch.tensor([f.label for f in features]      ,  dtype=torch.long     )
				elif self.hparams.glue_output_mode == "regression":
							_snake_case          =      torch.tensor([f.label for f in features]      ,  dtype=torch.float     )

				return DataLoader(
				    TensorDataset(lowercase      ,  lowercase      ,  lowercase      ,  lowercase     )      ,  batch_size=lowercase      ,  shuffle=lowercase      ,  )



	def 			A							(  self     :   str      ,  lowercase     :   Dict      ,  lowercase     :   Optional[Any]     ):


				'''simple docstring'''


				_snake_case          =      {'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}

				if self.config.model_type not in ["distilbert", "bart"]:
							_snake_case          =      batch[2] if self.config.model_type in ['bert', 'xlnet', 'albert'] else None

				_snake_case          =      self(**lowercase     )
				_snake_case     , _snake_case          =      outputs[:2]
				_snake_case          =      logits.detach().cpu().numpy()
				_snake_case          =      inputs['labels'].detach().cpu().numpy()

				return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}



	def 			A							(  self     :   Optional[Any]      ,  lowercase     :   Optional[Any]     ):


				'''simple docstring'''


				_snake_case          =      torch.stack([x['val_loss'] for x in outputs]     ).mean().detach().cpu().item()
				_snake_case          =      np.concatenate([x['pred'] for x in outputs]      ,  axis=0     )

				if self.hparams.glue_output_mode == "classification":
							_snake_case          =      np.argmax(lowercase      ,  axis=1     )
				elif self.hparams.glue_output_mode == "regression":
							_snake_case          =      np.squeeze(lowercase     )

				_snake_case          =      np.concatenate([x['target'] for x in outputs]      ,  axis=0     )
				_snake_case          =      [[] for _ in range(out_label_ids.shape[0]     )]
				_snake_case          =      [[] for _ in range(out_label_ids.shape[0]     )]

				_snake_case          =      {**{'val_loss': val_loss_mean}, **compute_metrics(self.hparams.task      ,  lowercase      ,  lowercase     )}

				_snake_case          =      dict(results.items()     )
				_snake_case          =      results
				return ret, preds_list, out_label_list



	def 			A							(  self     :   List[str]      ,  lowercase     :   list     ):


				'''simple docstring'''


				_snake_case     , _snake_case     , _snake_case          =      self._eval_end(lowercase     )
				_snake_case          =      ret['log']
				return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



	def 			A							(  self     :   Tuple      ,  lowercase     :   Tuple     ):


				'''simple docstring'''


				_snake_case     , _snake_case     , _snake_case          =      self._eval_end(lowercase     )
				_snake_case          =      ret['log']
				# `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
				return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}





	@staticmethod
	def 			A							(  lowercase     :   Optional[int]      ,  lowercase     :   Optional[int]     ):


				'''simple docstring'''


				BaseTransformer.add_model_specific_args(lowercase      ,  lowercase     )
				parser.add_argument(
				    '--max_seq_length'      ,  default=128      ,  type=lowercase      ,  help=(
				        'The maximum total input sequence length after tokenization. Sequences longer '
				        'than this will be truncated, sequences shorter will be padded.'
				    )      ,  )

				parser.add_argument(
				    '--task'      ,  default=''      ,  type=lowercase      ,  required=lowercase      ,  help='The GLUE task to run'      ,  )
				parser.add_argument(
				    '--gpus'      ,  default=0      ,  type=lowercase      ,  help='The number of GPUs allocated for this, it is by default 0 meaning none'      ,  )

				parser.add_argument(
				    '--overwrite_cache'      ,  action='store_true'      ,  help='Overwrite the cached training and evaluation sets'     )

				return parser






def 	a_		(	)			->  List[str]:
			_snake_case          =      argparse.ArgumentParser()
			add_generic_args(__lowercase	, os.getcwd()			)
			_snake_case          =      GLUETransformer.add_model_specific_args(__lowercase	, os.getcwd()			)
			_snake_case          =      parser.parse_args()

			# If output_dir not provided, a folder will be generated in pwd
			if args.output_dir is None:
						_snake_case          =      os.path.join(
						    './results'	, f'''{args.task}_{time.strftime('%Y%m%d_%H%M%S'			)}'''	, )
						os.makedirs(args.output_dir			)

			_snake_case          =      GLUETransformer(__lowercase			)
			_snake_case          =      generic_train(__lowercase	, __lowercase			)

			# Optionally, predict on dev set and write to output_dir
			if args.do_predict:
						_snake_case          =      sorted(glob.glob(os.path.join(args.output_dir	, 'checkpoint-epoch=*.ckpt'			)	, recursive=__lowercase			)			)
						_snake_case          =      model.load_from_checkpoint(checkpoints[-1]			)
						return trainer.test(__lowercase			)


if __name__ == "__main__":
			main()