
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 				UpperCAmelCase						( __lowerCamelCase    ,			__lowerCamelCase    ,			unittest.TestCase       ):
					a__:			Tuple      	=					StableDiffusionSAGPipeline
					a__:			str      	=					TEXT_TO_IMAGE_PARAMS
					a__:			Optional[Any]      	=					TEXT_TO_IMAGE_BATCH_PARAMS
					a__:			Dict      	=					TEXT_TO_IMAGE_IMAGE_PARAMS
					a__:			Optional[int]      	=					TEXT_TO_IMAGE_IMAGE_PARAMS
					a__:			Optional[Any]      	=					False

					def 							_lowerCAmelCase						(      self   :       List[Any]       ):
											torch.manual_seed(0       )
											lowercase							:     Optional[Any]									=	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			,   )
											lowercase							:     Tuple									=	DDIMScheduler(
											    beta_start=0.0_00_85			,   beta_end=0.0_12			,   beta_schedule='''scaled_linear'''			,   clip_sample=lowerCAmelCase			,   set_alpha_to_one=lowerCAmelCase			,   )
											torch.manual_seed(0       )
											lowercase							:     str									=	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       )
											lowercase							:     Optional[Any]									=	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			,   )
											lowercase							:     Optional[int]									=	CLIPTextModel(lowerCAmelCase       )
											lowercase							:     Union[str, Any]									=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )

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

					def 							_lowerCAmelCase						(      self   :       Dict			,   lowerCAmelCase   :       str			,   lowerCAmelCase   :       Dict=0       ):
											if str(lowerCAmelCase       ).startswith('''mps'''       ):
																	lowercase							:     Optional[int]									=	torch.manual_seed(lowerCAmelCase       )
											else:
																	lowercase							:     Dict									=	torch.Generator(device=lowerCAmelCase       ).manual_seed(lowerCAmelCase       )
											lowercase							:     Optional[int]									=	{
											    'prompt': '.',
											    'generator': generator,
											    'num_inference_steps': 2,
											    'guidance_scale': 1.0,
											    'sag_scale': 1.0,
											    'output_type': 'numpy',
											}
											return inputs




					def 							_lowerCAmelCase						(      self   :       Optional[int]       ):
											super().test_inference_batch_single_identical(expected_max_diff=3E-3       )







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

					def 							_lowerCAmelCase						(      self   :       Union[str, Any]       ):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()

					def 							_lowerCAmelCase						(      self   :       List[Any]       ):
											lowercase							:     Union[str, Any]									=	StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''       )
											lowercase							:     Tuple									=	sag_pipe.to(lowerCAmelCase       )
											sag_pipe.set_progress_bar_config(disable=lowerCAmelCase       )

											lowercase							:     str									=	'.'
											lowercase							:     str									=	torch.manual_seed(0       )
											lowercase							:     Optional[Any]									=	sag_pipe(
											    [prompt]			,   generator=lowerCAmelCase			,   guidance_scale=7.5			,   sag_scale=1.0			,   num_inference_steps=20			,   output_type='''np'''       )

											lowercase							:     Any									=	output.images

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

											assert image.shape == (1, 512, 512, 3)
											lowercase							:     Optional[int]									=	np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49]       )

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

					def 							_lowerCAmelCase						(      self   :       str       ):
											lowercase							:     Optional[int]									=	StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''       )
											lowercase							:     int									=	sag_pipe.to(lowerCAmelCase       )
											sag_pipe.set_progress_bar_config(disable=lowerCAmelCase       )

											lowercase							:     Dict									=	'.'
											lowercase							:     Union[str, Any]									=	torch.manual_seed(0       )
											lowercase							:     Any									=	sag_pipe(
											    [prompt]			,   generator=lowerCAmelCase			,   guidance_scale=7.5			,   sag_scale=1.0			,   num_inference_steps=20			,   output_type='''np'''       )

											lowercase							:     List[str]									=	output.images

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

											assert image.shape == (1, 512, 512, 3)
											lowercase							:     Optional[int]									=	np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71]       )

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




					def 							_lowerCAmelCase						(      self   :       Any       ):
											lowercase							:     Dict									=	StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''       )
											lowercase							:     List[str]									=	sag_pipe.to(lowerCAmelCase       )
											sag_pipe.set_progress_bar_config(disable=lowerCAmelCase       )

											lowercase							:     str									=	'.'
											lowercase							:     int									=	torch.manual_seed(0       )
											lowercase							:     Any									=	sag_pipe(
											    [prompt]			,   width=768			,   height=512			,   generator=lowerCAmelCase			,   guidance_scale=7.5			,   sag_scale=1.0			,   num_inference_steps=20			,   output_type='''np'''			,   )

											lowercase							:     List[str]									=	output.images

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






'''simple docstring'''






import argparse


lowercase				='docs/source/_static/js/custom.js'





def lowerCamelCase__     (		__lowerCamelCase			:				List[Any]				):







   '''simple docstring'''




   with open(__lowerCamelCase			,       encoding='utf-8'			,       newline='\n'				) as f:
      _UpperCAmelCase  :  Optional[Any]				       =f.readlines()
   _UpperCAmelCase  :  Dict				       =0

   # First let's put the right version
   while not lines[index].startswith('const stableVersion ='				):
      index += 1
   _UpperCAmelCase  :  Tuple				       =f"const stableVersion = \"v{version}\"\n"

   # Then update the dictionary
   while not lines[index].startswith('const versionMapping = {'				):
      index += 1

   # We go until the end
   while not lines[index].startswith('}'				):
      index += 1
   # We add the new version at the end
   lines[index - 1] += f"    \"v{version}\": \"v{version}\",\n"

   with open(__lowerCamelCase			,       'w'			,       encoding='utf-8'			,       newline='\n'				) as f:
      f.writelines(__lowerCamelCase				)


if __name__ == "__main__":
 lowercase				=argparse.ArgumentParser()
 parser.add_argument('--version', help='Release version.')
 lowercase				=parser.parse_args()
 update_custom_js(args.version)

