
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     lowerCamelCase      (       __lowerCamelCase ,		unittest.TestCase					):
	UpperCamelCase_  :					Any    		=	ConsistencyModelPipeline
	UpperCamelCase_  :					Optional[Any]    		=	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	UpperCamelCase_  :					Union[str, Any]    		=	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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







	@property
	def  snake_case__     (     self      :List[str]			)		->      Optional[Any]:







				"""simple docstring"""






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







	@property
	def  snake_case__     (     self      :Union[str, Any]			)		->      List[str]:







				"""simple docstring"""






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







	def  snake_case__     (     self      :List[Any]   ,     lowercase      :Any=False			)		->      Any:







				"""simple docstring"""






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

				# Default to CM multistep sampler
				SCREAMING_SNAKE_CASE     						= CMStochasticIterativeScheduler(
				    num_train_timesteps=4_0   ,     sigma_min=0.0_02   ,     sigma_max=80.0   ,     )

				SCREAMING_SNAKE_CASE     						= {
				    '''unet''': unet,
				    '''scheduler''': scheduler,
				}

				return components







	def  snake_case__     (     self      :Union[str, Any]   ,     lowercase      :Any   ,     lowercase      :Dict=0			)		->      Any:







				"""simple docstring"""






				if str(lowercase			).startswith('''mps'''			):
							SCREAMING_SNAKE_CASE     						= torch.manual_seed(lowercase			)
				else:
							SCREAMING_SNAKE_CASE     						= torch.Generator(device=lowercase			).manual_seed(lowercase			)

				SCREAMING_SNAKE_CASE     						= {
				    '''batch_size''': 1,
				    '''num_inference_steps''': None,
				    '''timesteps''': [2_2, 0],
				    '''generator''': generator,
				    '''output_type''': '''np''',
				}

				return inputs







	def  snake_case__     (     self      :Optional[int]			)		->      str:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
				SCREAMING_SNAKE_CASE     						= self.get_dummy_components()
				SCREAMING_SNAKE_CASE     						= ConsistencyModelPipeline(**lowercase			)
				SCREAMING_SNAKE_CASE     						= pipe.to(lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= self.get_dummy_inputs(lowercase			)
				SCREAMING_SNAKE_CASE     						= pipe(**lowercase			).images
				assert image.shape == (1, 3_2, 3_2, 3)

				SCREAMING_SNAKE_CASE     						= image[0, -3:, -3:, -1]
				SCREAMING_SNAKE_CASE     						= 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  snake_case__     (     self      :Optional[Any]			)		->      Dict:







				"""simple docstring"""






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

				SCREAMING_SNAKE_CASE     						= self.get_dummy_inputs(lowercase			)
				SCREAMING_SNAKE_CASE     						= 0
				SCREAMING_SNAKE_CASE     						= pipe(**lowercase			).images
				assert image.shape == (1, 3_2, 3_2, 3)

				SCREAMING_SNAKE_CASE     						= image[0, -3:, -3:, -1]
				SCREAMING_SNAKE_CASE     						= 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  snake_case__     (     self      :Union[str, Any]			)		->      str:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
				SCREAMING_SNAKE_CASE     						= self.get_dummy_components()
				SCREAMING_SNAKE_CASE     						= ConsistencyModelPipeline(**lowercase			)
				SCREAMING_SNAKE_CASE     						= pipe.to(lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= self.get_dummy_inputs(lowercase			)
				SCREAMING_SNAKE_CASE     						= 1
				SCREAMING_SNAKE_CASE     						= None
				SCREAMING_SNAKE_CASE     						= pipe(**lowercase			).images
				assert image.shape == (1, 3_2, 3_2, 3)

				SCREAMING_SNAKE_CASE     						= image[0, -3:, -3:, -1]
				SCREAMING_SNAKE_CASE     						= 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  snake_case__     (     self      :Optional[int]			)		->      Optional[Any]:







				"""simple docstring"""






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

				SCREAMING_SNAKE_CASE     						= self.get_dummy_inputs(lowercase			)
				SCREAMING_SNAKE_CASE     						= 1
				SCREAMING_SNAKE_CASE     						= None
				SCREAMING_SNAKE_CASE     						= 0
				SCREAMING_SNAKE_CASE     						= pipe(**lowercase			).images
				assert image.shape == (1, 3_2, 3_2, 3)

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







	def  snake_case__     (     self      :Any			)		->      Optional[int]:







				"""simple docstring"""






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







	def  snake_case__     (     self      :int   ,     lowercase      :List[Any]=0   ,     lowercase      :Optional[Any]=False   ,     lowercase      :Any="cpu"   ,     lowercase      :Any=torch.floataa   ,     lowercase      :List[str]=(1, 3, 6_4, 6_4)			)		->      List[str]:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= torch.manual_seed(lowercase			)

				SCREAMING_SNAKE_CASE     						= {
				    '''num_inference_steps''': None,
				    '''timesteps''': [2_2, 0],
				    '''class_labels''': 0,
				    '''generator''': generator,
				    '''output_type''': '''np''',
				}

				if get_fixed_latents:
							SCREAMING_SNAKE_CASE     						= self.get_fixed_latents(seed=lowercase   ,     device=lowercase   ,     dtype=lowercase   ,     shape=lowercase			)
							SCREAMING_SNAKE_CASE     						= latents

				return inputs







	def  snake_case__     (     self      :str   ,     lowercase      :Optional[int]=0   ,     lowercase      :Dict="cpu"   ,     lowercase      :Optional[int]=torch.floataa   ,     lowercase      :str=(1, 3, 6_4, 6_4)			)		->      List[Any]:







				"""simple docstring"""






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







	def  snake_case__     (     self      :Tuple			)		->      Optional[Any]:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= UNetaDModel.from_pretrained('''diffusers/consistency_models'''   ,     subfolder='''diffusers_cd_imagenet64_l2'''			)
				SCREAMING_SNAKE_CASE     						= CMStochasticIterativeScheduler(
				    num_train_timesteps=4_0   ,     sigma_min=0.0_02   ,     sigma_max=80.0   ,     )
				SCREAMING_SNAKE_CASE     						= ConsistencyModelPipeline(unet=lowercase   ,     scheduler=lowercase			)
				pipe.to(torch_device=lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= self.get_inputs()
				SCREAMING_SNAKE_CASE     						= pipe(**lowercase			).images
				assert image.shape == (1, 6_4, 6_4, 3)

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

				SCREAMING_SNAKE_CASE     						= 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  snake_case__     (     self      :List[str]			)		->      Dict:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= UNetaDModel.from_pretrained('''diffusers/consistency_models'''   ,     subfolder='''diffusers_cd_imagenet64_l2'''			)
				SCREAMING_SNAKE_CASE     						= CMStochasticIterativeScheduler(
				    num_train_timesteps=4_0   ,     sigma_min=0.0_02   ,     sigma_max=80.0   ,     )
				SCREAMING_SNAKE_CASE     						= ConsistencyModelPipeline(unet=lowercase   ,     scheduler=lowercase			)
				pipe.to(torch_device=lowercase			)
				pipe.set_progress_bar_config(disable=lowercase			)

				SCREAMING_SNAKE_CASE     						= self.get_inputs()
				SCREAMING_SNAKE_CASE     						= 1
				SCREAMING_SNAKE_CASE     						= None
				SCREAMING_SNAKE_CASE     						= pipe(**lowercase			).images
				assert image.shape == (1, 6_4, 6_4, 3)

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

				SCREAMING_SNAKE_CASE     						= 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  snake_case__     (     self      :Optional[Any]			)		->      int:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= UNetaDModel.from_pretrained('''diffusers/consistency_models'''   ,     subfolder='''diffusers_cd_imagenet64_l2'''			)
				SCREAMING_SNAKE_CASE     						= CMStochasticIterativeScheduler(
				    num_train_timesteps=4_0   ,     sigma_min=0.0_02   ,     sigma_max=80.0   ,     )
				SCREAMING_SNAKE_CASE     						= ConsistencyModelPipeline(unet=lowercase   ,     scheduler=lowercase			)
				pipe.to(torch_device=lowercase   ,     torch_dtype=torch.floataa			)
				pipe.set_progress_bar_config(disable=lowercase			)

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

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

				SCREAMING_SNAKE_CASE     						= 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  snake_case__     (     self      :Tuple			)		->      Any:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= UNetaDModel.from_pretrained('''diffusers/consistency_models'''   ,     subfolder='''diffusers_cd_imagenet64_l2'''			)
				SCREAMING_SNAKE_CASE     						= CMStochasticIterativeScheduler(
				    num_train_timesteps=4_0   ,     sigma_min=0.0_02   ,     sigma_max=80.0   ,     )
				SCREAMING_SNAKE_CASE     						= ConsistencyModelPipeline(unet=lowercase   ,     scheduler=lowercase			)
				pipe.to(torch_device=lowercase   ,     torch_dtype=torch.floataa			)
				pipe.set_progress_bar_config(disable=lowercase			)

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

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

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

import os
from argparse import ArgumentParser
from typing import List

import torch.utils.data

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node


__lowerCAmelCase					   =       4
__lowerCAmelCase					   =       3





class     lowerCamelCase      (       __lowerCamelCase					):
	pass



def    a      ( a	)     ->Union[str, Any]:




			'''simple docstring'''


			for shard in shards:
						for i in range(a	):
									yield {"i": i, "shard": shard}


def    a      ( )     ->List[str]:




			'''simple docstring'''


			SCREAMING_SNAKE_CASE     						= int(os.environ['''RANK''']	)
			SCREAMING_SNAKE_CASE     						= int(os.environ['''WORLD_SIZE''']	)

			SCREAMING_SNAKE_CASE     						= ArgumentParser()
			parser.add_argument('''--streaming'''		,						type=a	)
			parser.add_argument('''--local_rank'''		,						type=a	)
			parser.add_argument('''--num_workers'''		,						type=a		,						default=0	)
			SCREAMING_SNAKE_CASE     						= parser.parse_args()
			SCREAMING_SNAKE_CASE     						= args.streaming
			SCREAMING_SNAKE_CASE     						= args.num_workers

			SCREAMING_SNAKE_CASE     						= {'''shards''': [F"""shard_{shard_idx}""" for shard_idx in range(a	)]}
			SCREAMING_SNAKE_CASE     						= IterableDataset.from_generator(a		,						gen_kwargs=a	)
			if not streaming:
						SCREAMING_SNAKE_CASE     						= Dataset.from_list(list(a	)	)

			SCREAMING_SNAKE_CASE     						= split_dataset_by_node(a		,						rank=a		,						world_size=a	)
			SCREAMING_SNAKE_CASE     						= torch.utils.data.DataLoader(a		,						num_workers=a	)

			SCREAMING_SNAKE_CASE     						= NUM_SHARDS * NUM_ITEMS_PER_SHARD
			SCREAMING_SNAKE_CASE     						= full_size // world_size
			expected_local_size += int(rank < (full_size % world_size)	)

			SCREAMING_SNAKE_CASE     						= sum(1 for _ in dataloader	)
			if local_size != expected_local_size:
						raise FailedTestError(F"""local_size {local_size} != expected_local_size {expected_local_size}"""	)


if __name__ == "__main__":
							main()