
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    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 				lowercase					(		_SCREAMING_SNAKE_CASE  ,   _SCREAMING_SNAKE_CASE  ,   unittest.TestCase   ):
		__lowercase       :       List[str]	     =				StableDiffusionSAGPipeline
		__lowercase       :       Any	     =				TEXT_TO_IMAGE_PARAMS
		__lowercase       :       Tuple	     =				TEXT_TO_IMAGE_BATCH_PARAMS
		__lowercase       :       Any	     =				TEXT_TO_IMAGE_IMAGE_PARAMS
		__lowercase       :       Any	     =				TEXT_TO_IMAGE_IMAGE_PARAMS
		__lowercase       :       Union[str, Any]	     =				False



		def __UpperCamelCase ( self       )  ->					List[str]:
				"""simple docstring"""
				torch.manual_seed(0       )
				UpperCamelCase		=     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				,      )
				UpperCamelCase		=     DDIMScheduler(
				    beta_start=0.0_0085				,      beta_end=0.012				,      beta_schedule='scaled_linear'				,      clip_sample=A_				,      set_alpha_to_one=A_				,      )
				torch.manual_seed(0       )
				UpperCamelCase		=     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				,      )
				torch.manual_seed(0       )
				UpperCamelCase		=     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=1_000				,      )
				UpperCamelCase		=     CLIPTextModel(A_       )
				UpperCamelCase		=     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

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



		def __UpperCamelCase ( self				,      A_				,      A_=0       )  ->					Optional[int]:
				"""simple docstring"""
				if str(A_       ).startswith('mps'       ):
						UpperCamelCase		=     torch.manual_seed(A_       )
				else:
						UpperCamelCase		=     torch.Generator(device=A_       ).manual_seed(A_       )
				UpperCamelCase		=     {
				    'prompt': '.',
				    'generator': generator,
				    'num_inference_steps': 2,
				    'guidance_scale': 1.0,
				    'sag_scale': 1.0,
				    'output_type': 'numpy',
				}
				return inputs



		def __UpperCamelCase ( self       )  ->					Tuple:
				"""simple docstring"""
				super().test_inference_batch_single_identical(expected_max_diff=3e-3       )






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



		def __UpperCamelCase ( self       )  ->					Any:
				"""simple docstring"""
				# clean up the VRAM after each test
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()



		def __UpperCamelCase ( self       )  ->					Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase		=     StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'       )
				UpperCamelCase		=     sag_pipe.to(A_       )
				sag_pipe.set_progress_bar_config(disable=A_       )

				UpperCamelCase		=     '.'
				UpperCamelCase		=     torch.manual_seed(0       )
				UpperCamelCase		=     sag_pipe(
				    [prompt]				,      generator=A_				,      guidance_scale=7.5				,      sag_scale=1.0				,      num_inference_steps=20				,      output_type='np'       )

				UpperCamelCase		=     output.images

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

				assert image.shape == (1, 512, 512, 3)
				UpperCamelCase		=     np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]       )

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



		def __UpperCamelCase ( self       )  ->					Optional[Any]:
				"""simple docstring"""
				UpperCamelCase		=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'       )
				UpperCamelCase		=     sag_pipe.to(A_       )
				sag_pipe.set_progress_bar_config(disable=A_       )

				UpperCamelCase		=     '.'
				UpperCamelCase		=     torch.manual_seed(0       )
				UpperCamelCase		=     sag_pipe(
				    [prompt]				,      generator=A_				,      guidance_scale=7.5				,      sag_scale=1.0				,      num_inference_steps=20				,      output_type='np'       )

				UpperCamelCase		=     output.images

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

				assert image.shape == (1, 512, 512, 3)
				UpperCamelCase		=     np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]       )

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



		def __UpperCamelCase ( self       )  ->					List[str]:
				"""simple docstring"""
				UpperCamelCase		=     StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'       )
				UpperCamelCase		=     sag_pipe.to(A_       )
				sag_pipe.set_progress_bar_config(disable=A_       )

				UpperCamelCase		=     '.'
				UpperCamelCase		=     torch.manual_seed(0       )
				UpperCamelCase		=     sag_pipe(
				    [prompt]				,      width=768				,      height=512				,      generator=A_				,      guidance_scale=7.5				,      sag_scale=1.0				,      num_inference_steps=20				,      output_type='np'				,      )

				UpperCamelCase		=     output.images

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


from ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCAmelCase					: Union[str, Any]      =       logging.get_logger(__name__)

_UpperCAmelCase					: List[str]      =       {
    "EleutherAI/gpt-neox-20b": "https://huggingface.co/EleutherAI/gpt-neox-20b/resolve/main/config.json",
    # See all GPTNeoX models at https://huggingface.co/models?filter=gpt_neox
}


class 				lowercase					(		_SCREAMING_SNAKE_CASE   ):
		__lowercase       :       List[str]	     =				"gpt_neox"



		def __init__( self				,      A_=50_432				,      A_=6_144				,      A_=44				,      A_=64				,      A_=24_576				,      A_="gelu"				,      A_=0.25				,      A_=10_000				,      A_=0.0				,      A_=0.0				,      A_=0.1				,      A_=2_048				,      A_=0.02				,      A_=1e-5				,      A_=True				,      A_=0				,      A_=2				,      A_=False				,      A_=True				,      A_=None				,      **A_				,      )  ->					Tuple:
				"""simple docstring"""
				super().__init__(bos_token_id=A_				,      eos_token_id=A_				,      **A_       )
				UpperCamelCase		=     vocab_size
				UpperCamelCase		=     max_position_embeddings
				UpperCamelCase		=     hidden_size
				UpperCamelCase		=     num_hidden_layers
				UpperCamelCase		=     num_attention_heads
				UpperCamelCase		=     intermediate_size
				UpperCamelCase		=     hidden_act
				UpperCamelCase		=     rotary_pct
				UpperCamelCase		=     rotary_emb_base
				UpperCamelCase		=     attention_dropout
				UpperCamelCase		=     hidden_dropout
				UpperCamelCase		=     classifier_dropout
				UpperCamelCase		=     initializer_range
				UpperCamelCase		=     layer_norm_eps
				UpperCamelCase		=     use_cache
				UpperCamelCase		=     tie_word_embeddings
				UpperCamelCase		=     use_parallel_residual
				UpperCamelCase		=     rope_scaling
				self._rope_scaling_validation()

				if self.hidden_size % self.num_attention_heads != 0:
						raise ValueError(
						    'The hidden size is not divisble by the number of attention heads! Make sure to update them!'       )



		def __UpperCamelCase ( self       )  ->					Tuple:
				"""simple docstring"""
				if self.rope_scaling is None:
						return

				if not isinstance(self.rope_scaling				,      A_       ) or len(self.rope_scaling       ) != 2:
						raise ValueError(
						    '`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, '
						    F'''got {self.rope_scaling}'''       )
				UpperCamelCase		=     self.rope_scaling.get('type'				,      A_       )
				UpperCamelCase		=     self.rope_scaling.get('factor'				,      A_       )
				if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
						raise ValueError(
						    F'''`rope_scaling`\'s name field must be one of [\'linear\', \'dynamic\'], got {rope_scaling_type}'''       )
				if rope_scaling_factor is None or not isinstance(A_				,      A_       ) or rope_scaling_factor <= 1.0:
						raise ValueError(F'''`rope_scaling`\'s factor field must be an float > 1, got {rope_scaling_factor}'''       )
