






"""simple docstring"""
import gc
import unittest

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

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

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


enable_full_determinism()







class      UpperCamelCase_			(__A							,		unittest.TestCase	):
	__magic_name__       =							ConsistencyModelPipeline
	__magic_name__       =							UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	__magic_name__       =							UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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





	@property
	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     str:
				UpperCAmelCase_	:					Union[str, Any]     	=   UNetaDModel.from_pretrained(
				    "diffusers/consistency-models-test"      ,      subfolder="test_unet"      ,      )
				return unet





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





	def 	_SCREAMING_SNAKE_CASE ( self   :					int      ,      lowerCAmelCase_   :					Optional[Any]=False		)		->     List[Any]:
				if class_cond:
							UpperCAmelCase_	:					Union[str, Any]     	=   self.dummy_cond_unet
				else:
							UpperCAmelCase_	:					Any     	=   self.dummy_uncond_unet

				# Default to CM multistep sampler
				UpperCAmelCase_	:					int     	=   CMStochasticIterativeScheduler(
				    num_train_timesteps=40      ,      sigma_min=0.0_0_2      ,      sigma_max=8_0.0      ,      )

				UpperCAmelCase_	:					Any     	=   {
				    "unet": unet,
				    "scheduler": scheduler,
				}

				return components





	def 	_SCREAMING_SNAKE_CASE ( self   :					int      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Any=0		)		->     List[Any]:
				if str(lowerCAmelCase_		).startswith("mps"		):
							UpperCAmelCase_	:					Any     	=   torch.manual_seed(lowerCAmelCase_		)
				else:
							UpperCAmelCase_	:					List[str]     	=   torch.Generator(device=lowerCAmelCase_		).manual_seed(lowerCAmelCase_		)

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

				return inputs





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

				UpperCAmelCase_	:					List[Any]     	=   self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_	:					Any     	=   pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_	:					List[str]     	=   image[0, -3:, -3:, -1]
				UpperCAmelCase_	:					List[Any]     	=   np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]		)

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





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

				UpperCAmelCase_	:					Optional[Any]     	=   self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_	:					str     	=   0
				UpperCAmelCase_	:					int     	=   pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_	:					Tuple     	=   image[0, -3:, -3:, -1]
				UpperCAmelCase_	:					Dict     	=   np.array([0.3_5_7_2, 0.6_2_7_3, 0.4_0_3_1, 0.3_9_6_1, 0.4_3_2_1, 0.5_7_3_0, 0.5_2_6_6, 0.4_7_8_0, 0.5_0_0_4]		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     List[str]:
				UpperCAmelCase_	:					Optional[int]     	=   "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_	:					Dict     	=   self.get_dummy_components()
				UpperCAmelCase_	:					Tuple     	=   ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_	:					Optional[Any]     	=   pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_	:					Tuple     	=   self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_	:					Optional[int]     	=   1
				UpperCAmelCase_	:					List[Any]     	=   None
				UpperCAmelCase_	:					Dict     	=   pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 32, 32, 3)

				UpperCAmelCase_	:					Dict     	=   image[0, -3:, -3:, -1]
				UpperCAmelCase_	:					Optional[int]     	=   np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     Dict:
				UpperCAmelCase_	:					int     	=   "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_	:					Tuple     	=   self.get_dummy_components(class_cond=lowerCAmelCase_		)
				UpperCAmelCase_	:					List[Any]     	=   ConsistencyModelPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_	:					Tuple     	=   pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

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

				UpperCAmelCase_	:					List[str]     	=   image[0, -3:, -3:, -1]
				UpperCAmelCase_	:					Optional[int]     	=   np.array([0.5_0_0_4, 0.5_0_0_4, 0.4_9_9_4, 0.5_0_0_8, 0.4_9_7_6, 0.5_0_1_8, 0.4_9_9_0, 0.4_9_8_2, 0.4_9_8_7]		)

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







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





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     int:
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





	def 	_SCREAMING_SNAKE_CASE ( self   :					Any      ,      lowerCAmelCase_   :					Any=0      ,      lowerCAmelCase_   :					Tuple=False      ,      lowerCAmelCase_   :					Optional[Any]="cpu"      ,      lowerCAmelCase_   :					Dict=torch.floataa      ,      lowerCAmelCase_   :					Optional[Any]=(1, 3, 64, 64)		)		->     str:
				UpperCAmelCase_	:					str     	=   torch.manual_seed(lowerCAmelCase_		)

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

				if get_fixed_latents:
							UpperCAmelCase_	:					str     	=   self.get_fixed_latents(seed=lowerCAmelCase_      ,      device=lowerCAmelCase_      ,      dtype=lowerCAmelCase_      ,      shape=lowerCAmelCase_		)
							UpperCAmelCase_	:					List[str]     	=   latents

				return inputs





	def 	_SCREAMING_SNAKE_CASE ( self   :					Any      ,      lowerCAmelCase_   :					List[Any]=0      ,      lowerCAmelCase_   :					Any="cpu"      ,      lowerCAmelCase_   :					Optional[Any]=torch.floataa      ,      lowerCAmelCase_   :					int=(1, 3, 64, 64)		)		->     Optional[Any]:
				if type(lowerCAmelCase_		) == str:
							UpperCAmelCase_	:					Tuple     	=   torch.device(lowerCAmelCase_		)
				UpperCAmelCase_	:					Tuple     	=   torch.Generator(device=lowerCAmelCase_		).manual_seed(lowerCAmelCase_		)
				UpperCAmelCase_	:					List[Any]     	=   randn_tensor(lowerCAmelCase_      ,      generator=lowerCAmelCase_      ,      device=lowerCAmelCase_      ,      dtype=lowerCAmelCase_		)
				return latents





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     List[str]:
				UpperCAmelCase_	:					List[str]     	=   UNetaDModel.from_pretrained("diffusers/consistency_models"      ,      subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_	:					str     	=   CMStochasticIterativeScheduler(
				    num_train_timesteps=40      ,      sigma_min=0.0_0_2      ,      sigma_max=8_0.0      ,      )
				UpperCAmelCase_	:					int     	=   ConsistencyModelPipeline(unet=lowerCAmelCase_      ,      scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_	:					Union[str, Any]     	=   self.get_inputs()
				UpperCAmelCase_	:					int     	=   pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

				UpperCAmelCase_	:					Any     	=   image[0, -3:, -3:, -1]

				UpperCAmelCase_	:					Union[str, Any]     	=   np.array([0.0_8_8_8, 0.0_8_8_1, 0.0_6_6_6, 0.0_4_7_9, 0.0_2_9_2, 0.0_1_9_5, 0.0_2_0_1, 0.0_1_6_3, 0.0_2_5_4]		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     List[Any]:
				UpperCAmelCase_	:					List[str]     	=   UNetaDModel.from_pretrained("diffusers/consistency_models"      ,      subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_	:					Any     	=   CMStochasticIterativeScheduler(
				    num_train_timesteps=40      ,      sigma_min=0.0_0_2      ,      sigma_max=8_0.0      ,      )
				UpperCAmelCase_	:					str     	=   ConsistencyModelPipeline(unet=lowerCAmelCase_      ,      scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_	:					str     	=   self.get_inputs()
				UpperCAmelCase_	:					Union[str, Any]     	=   1
				UpperCAmelCase_	:					Union[str, Any]     	=   None
				UpperCAmelCase_	:					List[str]     	=   pipe(**lowerCAmelCase_		).images
				assert image.shape == (1, 64, 64, 3)

				UpperCAmelCase_	:					Dict     	=   image[0, -3:, -3:, -1]

				UpperCAmelCase_	:					Union[str, Any]     	=   np.array([0.0_3_4_0, 0.0_1_5_2, 0.0_0_6_3, 0.0_2_6_7, 0.0_2_2_1, 0.0_1_0_7, 0.0_4_1_6, 0.0_1_8_6, 0.0_2_1_7]		)

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





	@require_torch_a
	def 	_SCREAMING_SNAKE_CASE ( self   :					Any		)		->     List[str]:
				UpperCAmelCase_	:					Optional[int]     	=   UNetaDModel.from_pretrained("diffusers/consistency_models"      ,      subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_	:					int     	=   CMStochasticIterativeScheduler(
				    num_train_timesteps=40      ,      sigma_min=0.0_0_2      ,      sigma_max=8_0.0      ,      )
				UpperCAmelCase_	:					Optional[Any]     	=   ConsistencyModelPipeline(unet=lowerCAmelCase_      ,      scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_      ,      torch_dtype=torch.floataa		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

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

				UpperCAmelCase_	:					Dict     	=   image[0, -3:, -3:, -1]

				UpperCAmelCase_	:					int     	=   np.array([0.1_8_7_5, 0.1_4_2_8, 0.1_2_8_9, 0.2_1_5_1, 0.2_0_9_2, 0.1_4_7_7, 0.1_8_7_7, 0.1_6_4_1, 0.1_3_5_3]		)

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





	@require_torch_a
	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     int:
				UpperCAmelCase_	:					Tuple     	=   UNetaDModel.from_pretrained("diffusers/consistency_models"      ,      subfolder="diffusers_cd_imagenet64_l2"		)
				UpperCAmelCase_	:					Any     	=   CMStochasticIterativeScheduler(
				    num_train_timesteps=40      ,      sigma_min=0.0_0_2      ,      sigma_max=8_0.0      ,      )
				UpperCAmelCase_	:					Optional[Any]     	=   ConsistencyModelPipeline(unet=lowerCAmelCase_      ,      scheduler=lowerCAmelCase_		)
				pipe.to(torch_device=lowerCAmelCase_      ,      torch_dtype=torch.floataa		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

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

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

				UpperCAmelCase_	:					Dict     	=   np.array([0.1_6_6_3, 0.1_9_4_8, 0.2_2_7_5, 0.1_6_8_0, 0.1_2_0_4, 0.1_2_4_5, 0.1_8_5_8, 0.1_3_3_8, 0.2_0_9_5]		)

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









import argparse
import os

import torch

from transformers.utils import WEIGHTS_NAME


a__		     =     ['''small''', '''medium''', '''large''']

a__		     =     '''lm_head.decoder.weight'''
a__		     =     '''lm_head.weight'''





def 				A__ (snake_case   :		str		,	snake_case   :		str				)     ->						Union[str, Any]:
					__UpperCamelCase					:	Union[str, Any]             =      torch.load(snake_case				)
					__UpperCamelCase					:	int             =      d.pop(snake_case				)
					os.makedirs(snake_case		,	exist_ok=snake_case				)
					torch.save(snake_case		,	os.path.join(snake_case		,	snake_case				)				)


if __name__ == "__main__":
	a__		     =     argparse.ArgumentParser()
	parser.add_argument('''--dialogpt_path''', default='''.''', type=str)
	a__		     =     parser.parse_args()
	for MODEL in DIALOGPT_MODELS:
		a__		     =     os.path.join(args.dialogpt_path, f"{MODEL}_ft.pkl")
		a__		     =     f"./DialoGPT-{MODEL}"
		convert_dialogpt_checkpoint(
		    checkpoint_path,
		    pytorch_dump_folder_path,
		)






