"""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    (					lowercase   ,			unittest.TestCase		):
				UpperCAmelCase       :   str					=	ConsistencyModelPipeline
				UpperCAmelCase       :   Optional[Any]					=	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
				UpperCAmelCase       :   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       _lowercase  (self       :   Tuple)					->			Union[str, Any]:
					__snake_case    :							int   	=				UNetaDModel.from_pretrained(
					    'diffusers/consistency-models-test'			,						subfolder='test_unet'			,						)
					return unet



				@property
				def       _lowercase  (self       :   Union[str, Any])					->			Union[str, Any]:
					__snake_case    :							int   	=				UNetaDModel.from_pretrained(
					    'diffusers/consistency-models-test'			,						subfolder='test_unet_class_cond'			,						)
					return unet



				def       _lowercase  (self       :   Dict			,						_A       :   str=False)					->			List[Any]:
					if class_cond:
						__snake_case    :							int   	=				self.dummy_cond_unet
					else:
						__snake_case    :							Any   	=				self.dummy_uncond_unet

					# Default to CM multistep sampler
					__snake_case    :							List[str]   	=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40			,						sigma_min=0.002			,						sigma_max=80.0			,						)

					__snake_case    :							Optional[Any]   	=				{
					    'unet': unet,
					    'scheduler': scheduler,
					}

					return components



				def       _lowercase  (self       :   Dict			,						_A       :   Optional[Any]			,						_A       :   Union[str, Any]=0)					->			List[Any]:
					if str(_A).startswith('mps'):
						__snake_case    :							Tuple   	=				torch.manual_seed(_A)
					else:
						__snake_case    :							List[str]   	=				torch.Generator(device=_A).manual_seed(_A)

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

					return inputs



				def       _lowercase  (self       :   List[Any])					->			str:
					__snake_case    :							int   	=				'cpu'  # ensure determinism for the device-dependent torch.Generator
					__snake_case    :							Union[str, Any]   	=				self.get_dummy_components()
					__snake_case    :							Tuple   	=				ConsistencyModelPipeline(**_A)
					__snake_case    :							Optional[Any]   	=				pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case    :							Dict   	=				self.get_dummy_inputs(_A)
					__snake_case    :							Optional[Any]   	=				pipe(**_A).images
					assert image.shape == (1, 32, 32, 3)

					__snake_case    :							Any   	=				image[0, -3:, -3:, -1]
					__snake_case    :							str   	=				np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004])

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



				def       _lowercase  (self       :   List[Any])					->			str:
					__snake_case    :							Dict   	=				'cpu'  # ensure determinism for the device-dependent torch.Generator
					__snake_case    :							List[Any]   	=				self.get_dummy_components(class_cond=_A)
					__snake_case    :							Optional[int]   	=				ConsistencyModelPipeline(**_A)
					__snake_case    :							Optional[int]   	=				pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case    :							Tuple   	=				self.get_dummy_inputs(_A)
					__snake_case    :							Tuple   	=				0
					__snake_case    :							Optional[Any]   	=				pipe(**_A).images
					assert image.shape == (1, 32, 32, 3)

					__snake_case    :							List[str]   	=				image[0, -3:, -3:, -1]
					__snake_case    :							Dict   	=				np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004])

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



				def       _lowercase  (self       :   Union[str, Any])					->			Optional[Any]:
					__snake_case    :							str   	=				'cpu'  # ensure determinism for the device-dependent torch.Generator
					__snake_case    :							Any   	=				self.get_dummy_components()
					__snake_case    :							Optional[Any]   	=				ConsistencyModelPipeline(**_A)
					__snake_case    :							str   	=				pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case    :							Tuple   	=				self.get_dummy_inputs(_A)
					__snake_case    :							Union[str, Any]   	=				1
					__snake_case    :							Union[str, Any]   	=				None
					__snake_case    :							str   	=				pipe(**_A).images
					assert image.shape == (1, 32, 32, 3)

					__snake_case    :							str   	=				image[0, -3:, -3:, -1]
					__snake_case    :							Dict   	=				np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987])

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






				def       _lowercase  (self       :   List[Any])					->			Optional[Any]:
					__snake_case    :							Any   	=				'cpu'  # ensure determinism for the device-dependent torch.Generator
					__snake_case    :							Dict   	=				self.get_dummy_components(class_cond=_A)
					__snake_case    :							str   	=				ConsistencyModelPipeline(**_A)
					__snake_case    :							Any   	=				pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

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

					__snake_case    :							Dict   	=				image[0, -3:, -3:, -1]
					__snake_case    :							Optional[int]   	=				np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987])

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







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



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



				def       _lowercase  (self       :   Union[str, Any]			,						_A       :   Any=0			,						_A       :   List[Any]=False			,						_A       :   str="cpu"			,						_A       :   Union[str, Any]=torch.floataa			,						_A       :   Any=(1, 3, 64, 64))					->			Tuple:
					__snake_case    :							Union[str, Any]   	=				torch.manual_seed(_A)

					__snake_case    :							Any   	=				{
					    'num_inference_steps': None,
					    'timesteps': [22, 0],
					    'class_labels': 0,
					    'generator': generator,
					    'output_type': 'np',
					}

					if get_fixed_latents:
						__snake_case    :							str   	=				self.get_fixed_latents(seed=_A			,						device=_A			,						dtype=_A			,						shape=_A)
						__snake_case    :							int   	=				latents

					return inputs



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



				def       _lowercase  (self       :   Tuple)					->			int:
					__snake_case    :							Tuple   	=				UNetaDModel.from_pretrained('diffusers/consistency_models'			,						subfolder='diffusers_cd_imagenet64_l2')
					__snake_case    :							str   	=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40			,						sigma_min=0.002			,						sigma_max=80.0			,						)
					__snake_case    :							Union[str, Any]   	=				ConsistencyModelPipeline(unet=_A			,						scheduler=_A)
					pipe.to(torch_device=_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case    :							List[Any]   	=				self.get_inputs()
					__snake_case    :							Tuple   	=				pipe(**_A).images
					assert image.shape == (1, 64, 64, 3)

					__snake_case    :							str   	=				image[0, -3:, -3:, -1]

					__snake_case    :							Optional[int]   	=				np.array([0.0_888, 0.0_881, 0.0_666, 0.0_479, 0.0_292, 0.0_195, 0.0_201, 0.0_163, 0.0_254])

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



				def       _lowercase  (self       :   str)					->			Dict:
					__snake_case    :							Any   	=				UNetaDModel.from_pretrained('diffusers/consistency_models'			,						subfolder='diffusers_cd_imagenet64_l2')
					__snake_case    :							Dict   	=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40			,						sigma_min=0.002			,						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    :							Optional[Any]   	=				self.get_inputs()
					__snake_case    :							str   	=				1
					__snake_case    :							str   	=				None
					__snake_case    :							Optional[Any]   	=				pipe(**_A).images
					assert image.shape == (1, 64, 64, 3)

					__snake_case    :							Optional[int]   	=				image[0, -3:, -3:, -1]

					__snake_case    :							Optional[int]   	=				np.array([0.0_340, 0.0_152, 0.0_063, 0.0_267, 0.0_221, 0.0_107, 0.0_416, 0.0_186, 0.0_217])

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



				@require_torch_a
				def       _lowercase  (self       :   Any)					->			List[Any]:
					__snake_case    :							Tuple   	=				UNetaDModel.from_pretrained('diffusers/consistency_models'			,						subfolder='diffusers_cd_imagenet64_l2')
					__snake_case    :							str   	=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40			,						sigma_min=0.002			,						sigma_max=80.0			,						)
					__snake_case    :							Tuple   	=				ConsistencyModelPipeline(unet=_A			,						scheduler=_A)
					pipe.to(torch_device=_A			,						torch_dtype=torch.floataa)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case    :							Tuple   	=				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    :							Union[str, Any]   	=				pipe(**_A).images
					assert image.shape == (1, 64, 64, 3)

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

					__snake_case    :							Tuple   	=				np.array([0.1_875, 0.1_428, 0.1_289, 0.2_151, 0.2_092, 0.1_477, 0.1_877, 0.1_641, 0.1_353])

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






				@require_torch_a
				def       _lowercase  (self       :   Dict)					->			Tuple:
					__snake_case    :							Any   	=				UNetaDModel.from_pretrained('diffusers/consistency_models'			,						subfolder='diffusers_cd_imagenet64_l2')
					__snake_case    :							Dict   	=				CMStochasticIterativeScheduler(
					    num_train_timesteps=40			,						sigma_min=0.002			,						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    :							Tuple   	=				self.get_inputs(get_fixed_latents=_A			,						device=_A)
					__snake_case    :							Dict   	=				1
					__snake_case    :							Optional[int]   	=				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, 64, 64, 3)

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

					__snake_case    :							Union[str, Any]   	=				np.array([0.1_663, 0.1_948, 0.2_275, 0.1_680, 0.1_204, 0.1_245, 0.1_858, 0.1_338, 0.2_095])

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






"""simple docstring"""







import unittest
from queue import Empty
from threading import Thread

from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
from transformers.testing_utils import CaptureStdout, require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
	import torch

	from transformers import AutoModelForCausalLM

@require_torch
class 	UpperCamelCase    (					unittest.TestCase		):



				def       _lowercase  (self       :   List[Any])					->			Dict:
					__snake_case    :							int   	=				AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2')
					__snake_case    :							Optional[Any]   	=				AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2').to(_A)
					__snake_case    :							Union[str, Any]   	=				-1

					__snake_case    :							Optional[Any]   	=				ids_tensor((1, 5)			,						vocab_size=model.config.vocab_size).to(_A)
					__snake_case    :							Any   	=				model.generate(_A			,						max_new_tokens=10			,						do_sample=_A)
					__snake_case    :							str   	=				tokenizer.decode(greedy_ids[0])

					with CaptureStdout() as cs:
						__snake_case    :							List[str]   	=				TextStreamer(_A)
						model.generate(_A			,						max_new_tokens=10			,						do_sample=_A			,						streamer=_A)
					# The greedy text should be printed to stdout, except for the final "\n" in the streamer
					__snake_case    :							Optional[Any]   	=				cs.out[:-1]

					self.assertEqual(_A			,						_A)



				def       _lowercase  (self       :   int)					->			Optional[Any]:
					__snake_case    :							Any   	=				AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2')
					__snake_case    :							List[Any]   	=				AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2').to(_A)
					__snake_case    :							Optional[int]   	=				-1

					__snake_case    :							str   	=				ids_tensor((1, 5)			,						vocab_size=model.config.vocab_size).to(_A)
					__snake_case    :							int   	=				model.generate(_A			,						max_new_tokens=10			,						do_sample=_A)
					__snake_case    :							int   	=				tokenizer.decode(greedy_ids[0])

					__snake_case    :							int   	=				TextIteratorStreamer(_A)
					__snake_case    :							List[Any]   	=				{'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
					__snake_case    :							Any   	=				Thread(target=model.generate			,						kwargs=_A)
					thread.start()
					__snake_case    :							Dict   	=				''
					for new_text in streamer:
						streamer_text += new_text

					self.assertEqual(_A			,						_A)



				def       _lowercase  (self       :   Any)					->			List[str]:
					__snake_case    :							List[str]   	=				AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2')
					__snake_case    :							Any   	=				AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2').to(_A)
					__snake_case    :							Union[str, Any]   	=				-1

					__snake_case    :							List[Any]   	=				ids_tensor((1, 5)			,						vocab_size=model.config.vocab_size).to(_A)
					__snake_case    :							Optional[int]   	=				model.generate(_A			,						max_new_tokens=10			,						do_sample=_A)
					__snake_case    :							Union[str, Any]   	=				greedy_ids[:, input_ids.shape[1] :]
					__snake_case    :							Optional[Any]   	=				tokenizer.decode(new_greedy_ids[0])

					with CaptureStdout() as cs:
						__snake_case    :							str   	=				TextStreamer(_A			,						skip_prompt=_A)
						model.generate(_A			,						max_new_tokens=10			,						do_sample=_A			,						streamer=_A)
					# The greedy text should be printed to stdout, except for the final "\n" in the streamer
					__snake_case    :							List[str]   	=				cs.out[:-1]

					self.assertEqual(_A			,						_A)



				def       _lowercase  (self       :   List[str])					->			Dict:
					# Tests that we can pass `decode_kwargs` to the streamer to control how the tokens are decoded. Must be tested
					# with actual models -- the dummy models' tokenizers are not aligned with their models, and
					# `skip_special_tokens=True` has no effect on them
					__snake_case    :							int   	=				AutoTokenizer.from_pretrained('distilgpt2')
					__snake_case    :							Dict   	=				AutoModelForCausalLM.from_pretrained('distilgpt2').to(_A)
					__snake_case    :							Union[str, Any]   	=				-1

					__snake_case    :							List[str]   	=				torch.ones((1, 5)			,						device=_A).long() * model.config.bos_token_id
					with CaptureStdout() as cs:
						__snake_case    :							Tuple   	=				TextStreamer(_A			,						skip_special_tokens=_A)
						model.generate(_A			,						max_new_tokens=1			,						do_sample=_A			,						streamer=_A)

					# The prompt contains a special token, so the streamer should not print it. As such, the output text, when
					# re-tokenized, must only contain one token
					__snake_case    :							Tuple   	=				cs.out[:-1]  # Remove the final "\n"
					__snake_case    :							int   	=				tokenizer(_A			,						return_tensors='pt')
					self.assertEqual(streamer_text_tokenized.input_ids.shape			,						(1, 1))






				def       _lowercase  (self       :   Optional[int])					->			List[str]:
					__snake_case    :							Optional[int]   	=				AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2')
					__snake_case    :							Optional[int]   	=				AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2').to(_A)
					__snake_case    :							int   	=				-1

					__snake_case    :							List[Any]   	=				ids_tensor((1, 5)			,						vocab_size=model.config.vocab_size).to(_A)
					__snake_case    :							List[Any]   	=				TextIteratorStreamer(_A			,						timeout=0.001)
					__snake_case    :							str   	=				{'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
					__snake_case    :							str   	=				Thread(target=model.generate			,						kwargs=_A)
					thread.start()

					# The streamer will timeout after 0.001 seconds, so an exception will be raised
					with self.assertRaises(_A):
						__snake_case    :							Any   	=				''
						for new_text in streamer:
							streamer_text += new_text





