

'''simple docstring'''

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class lowerCamelCase   (    __lowerCAmelCase				,				__lowerCAmelCase				,				unittest.TestCase				):
				snake_case_							     = StableDiffusionParadigmsPipeline
				snake_case_							     = TEXT_TO_IMAGE_PARAMS
				snake_case_							     = TEXT_TO_IMAGE_BATCH_PARAMS
				snake_case_							     = TEXT_TO_IMAGE_IMAGE_PARAMS
				snake_case_							     = TEXT_TO_IMAGE_IMAGE_PARAMS

				def 					_lowerCamelCase    (       self )    ->							int:
					torch.manual_seed(0 )
					snake_case							=					UNetaDConditionModel(
					    block_out_channels=(32, 64),						layers_per_block=2,						sample_size=32,						in_channels=4,						out_channels=4,						down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D'),						up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D'),						cross_attention_dim=32,						attention_head_dim=(2, 4),						use_linear_projection=lowercase_,						)
					snake_case							=					DDIMParallelScheduler(
					    beta_start=0.00_085,						beta_end=0.012,						beta_schedule='scaled_linear',						clip_sample=lowercase_,						set_alpha_to_one=lowercase_,						)
					torch.manual_seed(0 )
					snake_case							=					AutoencoderKL(
					    block_out_channels=[32, 64],						in_channels=3,						out_channels=3,						down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D'],						up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D'],						latent_channels=4,						sample_size=128,						)
					torch.manual_seed(0 )
					snake_case							=					CLIPTextConfig(
					    bos_token_id=0,						eos_token_id=2,						hidden_size=32,						intermediate_size=37,						layer_norm_eps=1E-05,						num_attention_heads=4,						num_hidden_layers=5,						pad_token_id=1,						vocab_size=1000,						hidden_act='gelu',						projection_dim=512,						)
					snake_case							=					CLIPTextModel(lowercase_ )
					snake_case							=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

					snake_case							=					{
					    'unet': unet,
					    'scheduler': scheduler,
					    'vae': vae,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'safety_checker': None,
					    'feature_extractor': None,
					}
					return components

				def 					_lowerCamelCase    (       self,						lowercase_,						lowercase_=0 )    ->							List[str]:
					if str(lowercase_ ).startswith('mps' ):
						snake_case							=					torch.manual_seed(lowercase_ )
					else:
						snake_case							=					torch.Generator(device=lowercase_ ).manual_seed(lowercase_ )
					snake_case							=					{
					    'prompt': 'a photograph of an astronaut riding a horse',
					    'generator': generator,
					    'num_inference_steps': 10,
					    'guidance_scale': 6.0,
					    'output_type': 'numpy',
					    'parallel': 3,
					    'debug': True,
					}
					return inputs

				def 					_lowerCamelCase    (       self )    ->							int:
					snake_case							=					'cpu'  # ensure determinism for the device-dependent torch.Generator
					snake_case							=					self.get_dummy_components()
					snake_case							=					StableDiffusionParadigmsPipeline(**lowercase_ )
					snake_case							=					sd_pipe.to(lowercase_ )
					sd_pipe.set_progress_bar_config(disable=lowercase_ )

					snake_case							=					self.get_dummy_inputs(lowercase_ )
					snake_case							=					sd_pipe(**lowercase_ ).images
					snake_case							=					image[0, -3:, -3:, -1]
					assert image.shape == (1, 64, 64, 3)

					snake_case							=					np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023] )

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

				def 					_lowerCamelCase    (       self )    ->							Optional[Any]:
					snake_case							=					'cpu'  # ensure determinism for the device-dependent torch.Generator
					snake_case							=					self.get_dummy_components()
					torch.manual_seed(0 )
					snake_case							=					DDPMParallelScheduler()
					torch.manual_seed(0 )
					snake_case							=					StableDiffusionParadigmsPipeline(**lowercase_ )
					snake_case							=					sd_pipe.to(lowercase_ )
					sd_pipe.set_progress_bar_config(disable=lowercase_ )

					snake_case							=					self.get_dummy_inputs(lowercase_ )
					snake_case							=					sd_pipe(**lowercase_ ).images
					snake_case							=					image[0, -3:, -3:, -1]
					assert image.shape == (1, 64, 64, 3)

					snake_case							=					np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468] )

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

				def 					_lowerCamelCase    (       self )    ->							List[Any]:
					super().test_inference_batch_consistent(batch_sizes=[1, 2] )

				def 					_lowerCamelCase    (       self )    ->							Optional[Any]:
					super().test_inference_batch_single_identical(batch_size=2,						expected_max_diff=3E-3 )






				def 					_lowerCamelCase    (       self )    ->							str:
					snake_case							=					'cpu'  # ensure determinism for the device-dependent torch.Generator
					snake_case							=					self.get_dummy_components()
					snake_case							=					StableDiffusionParadigmsPipeline(**lowercase_ )
					snake_case							=					sd_pipe.to(lowercase_ )
					sd_pipe.set_progress_bar_config(disable=lowercase_ )

					snake_case							=					self.get_dummy_inputs(lowercase_ )
					snake_case							=					'french fries'
					snake_case							=					sd_pipe(**lowercase_,						negative_prompt=lowercase_ )
					snake_case							=					output.images
					snake_case							=					image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)

					snake_case							=					np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015] )

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







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

				def 					_lowerCamelCase    (       self )    ->							Tuple:
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()

				def 					_lowerCamelCase    (       self,						lowercase_=0 )    ->							Optional[Any]:
					snake_case							=					torch.Generator(device=lowercase_ ).manual_seed(lowercase_ )
					snake_case							=					{
					    'prompt': 'a photograph of an astronaut riding a horse',
					    'generator': generator,
					    'num_inference_steps': 10,
					    'guidance_scale': 7.5,
					    'output_type': 'numpy',
					    'parallel': 3,
					    'debug': True,
					}
					return inputs






				def 					_lowerCamelCase    (       self )    ->							Optional[int]:
					snake_case							=					'stabilityai/stable-diffusion-2-base'
					snake_case							=					DDIMParallelScheduler.from_pretrained(lowercase_,						subfolder='scheduler' )
					snake_case							=					StableDiffusionParadigmsPipeline.from_pretrained(lowercase_,						scheduler=lowercase_,						safety_checker=lowercase_ )
					pipe.to(lowercase_ )
					pipe.set_progress_bar_config(disable=lowercase_ )
					pipe.enable_attention_slicing()

					snake_case							=					self.get_inputs()
					snake_case							=					pipe(**lowercase_ ).images
					snake_case							=					image[0, -3:, -3:, -1].flatten()

					assert image.shape == (1, 512, 512, 3)

					snake_case							=					np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741] )

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







'''simple docstring'''







import logging

# standard libraries
import sys

import numpy as np
import pytest  # type: ignore

# Custom/local libraries
from matrix import matrix_operation as matop

_snake_case					: List[str]     			=					[[12, 10], [3, 9]]
_snake_case					: Any     			=					[[3, 4], [7, 4]]
_snake_case					: List[str]     			=					[[3, 0, 2], [2, 0, -2], [0, 1, 1]]
_snake_case					: Any     			=					[[3, 0, -2], [2, 0, 2], [0, 1, 1]]
_snake_case					: Any     			=					[[3, 0, 2], [2, 0, -2], [0, 1, 1], [2, 0, -2]]
_snake_case					: Optional[Any]     			=					[1]
_snake_case					: str     			=					[2]

_snake_case					: Optional[int]     			=					logging.getLogger()
_snake_case					: Optional[Any]     			=					logging.DEBUG
_snake_case					: Tuple     			=					logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)



@pytest.mark.mat_ops
@pytest.mark.parametrize(
    ("mat1", "mat2")						,      [(mat_a, mat_b), (mat_c, mat_d), (mat_d, mat_e), (mat_f, mat_h)]	)
def       _a		(						_SCREAMING_SNAKE_CASE		: List[Any]						,      _SCREAMING_SNAKE_CASE		: Tuple	):
						if (np.array(_SCREAMING_SNAKE_CASE	)).shape < (2, 2) or (np.array(_SCREAMING_SNAKE_CASE	)).shape < (2, 2):
												with pytest.raises(_SCREAMING_SNAKE_CASE	):
																		logger.info(F'\n\t{test_addition.__name__} returned integer'	)
																		matop.add(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)
						elif (np.array(_SCREAMING_SNAKE_CASE	)).shape == (np.array(_SCREAMING_SNAKE_CASE	)).shape:
												logger.info(F'\n\t{test_addition.__name__} with same matrix dims'	)
												_SCREAMING_SNAKE_CASE					=			(np.array(_SCREAMING_SNAKE_CASE	) + np.array(_SCREAMING_SNAKE_CASE	)).tolist()
												_SCREAMING_SNAKE_CASE					=			matop.add(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)
												assert theo == act
						else:
												with pytest.raises(_SCREAMING_SNAKE_CASE	):
																		logger.info(F'\n\t{test_addition.__name__} with different matrix dims'	)
																		matop.add(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)



@pytest.mark.mat_ops
@pytest.mark.parametrize(
    ("mat1", "mat2")						,      [(mat_a, mat_b), (mat_c, mat_d), (mat_d, mat_e), (mat_f, mat_h)]	)
def       _a		(						_SCREAMING_SNAKE_CASE		: Optional[Any]						,      _SCREAMING_SNAKE_CASE		: List[str]	):
						if (np.array(_SCREAMING_SNAKE_CASE	)).shape < (2, 2) or (np.array(_SCREAMING_SNAKE_CASE	)).shape < (2, 2):
												with pytest.raises(_SCREAMING_SNAKE_CASE	):
																		logger.info(F'\n\t{test_subtraction.__name__} returned integer'	)
																		matop.subtract(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)
						elif (np.array(_SCREAMING_SNAKE_CASE	)).shape == (np.array(_SCREAMING_SNAKE_CASE	)).shape:
												logger.info(F'\n\t{test_subtraction.__name__} with same matrix dims'	)
												_SCREAMING_SNAKE_CASE					=			(np.array(_SCREAMING_SNAKE_CASE	) - np.array(_SCREAMING_SNAKE_CASE	)).tolist()
												_SCREAMING_SNAKE_CASE					=			matop.subtract(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)
												assert theo == act
						else:
												with pytest.raises(_SCREAMING_SNAKE_CASE	):
																		logger.info(F'\n\t{test_subtraction.__name__} with different matrix dims'	)
																		assert matop.subtract(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)



@pytest.mark.mat_ops
@pytest.mark.parametrize(
    ("mat1", "mat2")						,      [(mat_a, mat_b), (mat_c, mat_d), (mat_d, mat_e), (mat_f, mat_h)]	)
def       _a		(						_SCREAMING_SNAKE_CASE		: Any						,      _SCREAMING_SNAKE_CASE		: List[str]	):
						if (np.array(_SCREAMING_SNAKE_CASE	)).shape < (2, 2) or (np.array(_SCREAMING_SNAKE_CASE	)).shape < (2, 2):
												logger.info(F'\n\t{test_multiplication.__name__} returned integer'	)
												with pytest.raises(_SCREAMING_SNAKE_CASE	):
																		matop.add(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)
						elif (np.array(_SCREAMING_SNAKE_CASE	)).shape == (np.array(_SCREAMING_SNAKE_CASE	)).shape:
												logger.info(F'\n\t{test_multiplication.__name__} meets dim requirements'	)
												_SCREAMING_SNAKE_CASE					=			(np.matmul(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)).tolist()
												_SCREAMING_SNAKE_CASE					=			matop.multiply(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)
												assert theo == act
						else:
												with pytest.raises(_SCREAMING_SNAKE_CASE	):
																		logger.info(
																		    F'\n\t{test_multiplication.__name__} does not meet dim requirements'	)
																		assert matop.subtract(_SCREAMING_SNAKE_CASE						,      _SCREAMING_SNAKE_CASE	)



@pytest.mark.mat_ops
def       _a		(						):
						_SCREAMING_SNAKE_CASE					=			(3.5 * np.array(_SCREAMING_SNAKE_CASE	)).tolist()
						_SCREAMING_SNAKE_CASE					=			matop.scalar_multiply(_SCREAMING_SNAKE_CASE						,      3.5	)
						assert theo == act



@pytest.mark.mat_ops
def       _a		(						):
						_SCREAMING_SNAKE_CASE					=			(np.identity(5	)).tolist()
						_SCREAMING_SNAKE_CASE					=			matop.identity(5	)
						assert theo == act




@pytest.mark.mat_ops
@pytest.mark.parametrize("mat"						,      [mat_a, mat_b, mat_c, mat_d, mat_e, mat_f]	)
def       _a		(						_SCREAMING_SNAKE_CASE		: Tuple	):
						if (np.array(_SCREAMING_SNAKE_CASE	)).shape < (2, 2):
												with pytest.raises(_SCREAMING_SNAKE_CASE	):
																		logger.info(F'\n\t{test_transpose.__name__} returned integer'	)
																		matop.transpose(_SCREAMING_SNAKE_CASE	)
						else:
												_SCREAMING_SNAKE_CASE					=			(np.transpose(_SCREAMING_SNAKE_CASE	)).tolist()
												_SCREAMING_SNAKE_CASE					=			matop.transpose(_SCREAMING_SNAKE_CASE						,      return_map=_SCREAMING_SNAKE_CASE	)
												assert theo == act