


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 				snake_case__(							UpperCAmelCase__,						UpperCAmelCase__,						unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Any					=							StableDiffusionSAGPipeline
						SCREAMING_SNAKE_CASE__      :  Optional[Any]					=							TEXT_TO_IMAGE_PARAMS
						SCREAMING_SNAKE_CASE__      :  str					=							TEXT_TO_IMAGE_BATCH_PARAMS
						SCREAMING_SNAKE_CASE__      :  Dict					=							TEXT_TO_IMAGE_IMAGE_PARAMS
						SCREAMING_SNAKE_CASE__      :  Any					=							TEXT_TO_IMAGE_IMAGE_PARAMS
						SCREAMING_SNAKE_CASE__      :  Optional[int]					=							False




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				UNetaDConditionModel(
										    block_out_channels=(3_2, 6_4)   ,						layers_per_block=2   ,						sample_size=3_2   ,						in_channels=4   ,						out_channels=4   ,						down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,						up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,						cross_attention_dim=3_2   ,						)
										lowerCAmelCase_ :      List[Any]    				=				DDIMScheduler(
										    beta_start=0.0_00_85   ,						beta_end=0.0_12   ,						beta_schedule='''scaled_linear'''   ,						clip_sample=__lowercase   ,						set_alpha_to_one=__lowercase   ,						)
										torch.manual_seed(0							)
										lowerCAmelCase_ :      List[Any]    				=				AutoencoderKL(
										    block_out_channels=[3_2, 6_4]   ,						in_channels=3   ,						out_channels=3   ,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,						latent_channels=4   ,						)
										torch.manual_seed(0							)
										lowerCAmelCase_ :      int    				=				CLIPTextConfig(
										    bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=3_2   ,						intermediate_size=3_7   ,						layer_norm_eps=1e-05   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=1_0_0_0   ,						)
										lowerCAmelCase_ :      Tuple    				=				CLIPTextModel(__lowercase							)
										lowerCAmelCase_ :      Optional[int]    				=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

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




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase=0							)					->    Tuple:
										if str(__lowercase							).startswith('''mps'''							):
														lowerCAmelCase_ :      List[str]    				=				torch.manual_seed(__lowercase							)
										else:
														lowerCAmelCase_ :      List[str]    				=				torch.Generator(device=__lowercase							).manual_seed(__lowercase							)
										lowerCAmelCase_ :      Any    				=				{
										    '''prompt''': '''.''',
										    '''generator''': generator,
										    '''num_inference_steps''': 2,
										    '''guidance_scale''': 1.0,
										    '''sag_scale''': 1.0,
										    '''output_type''': '''numpy''',
										}
										return inputs




						def 						lowercase_ (				self							)					->    Union[str, Any]:
										super().test_inference_batch_single_identical(expected_max_diff=3e-3							)







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







						'''simple docstring'''




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




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      List[str]    				=				StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''							)
										lowerCAmelCase_ :      Optional[int]    				=				sag_pipe.to(__lowercase							)
										sag_pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      List[Any]    				=				'''.'''
										lowerCAmelCase_ :      Optional[Any]    				=				torch.manual_seed(0							)
										lowerCAmelCase_ :      List[Any]    				=				sag_pipe(
										    [prompt]   ,						generator=__lowercase   ,						guidance_scale=7.5   ,						sag_scale=1.0   ,						num_inference_steps=2_0   ,						output_type='''np'''							)

										lowerCAmelCase_ :      List[str]    				=				output.images

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

										assert image.shape == (1, 5_1_2, 5_1_2, 3)
										lowerCAmelCase_ :      Any    				=				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 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      List[str]    				=				StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''							)
										lowerCAmelCase_ :      int    				=				sag_pipe.to(__lowercase							)
										sag_pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''.'''
										lowerCAmelCase_ :      str    				=				torch.manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				sag_pipe(
										    [prompt]   ,						generator=__lowercase   ,						guidance_scale=7.5   ,						sag_scale=1.0   ,						num_inference_steps=2_0   ,						output_type='''np'''							)

										lowerCAmelCase_ :      Optional[int]    				=				output.images

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

										assert image.shape == (1, 5_1_2, 5_1_2, 3)
										lowerCAmelCase_ :      Tuple    				=				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 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      Tuple    				=				StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''							)
										lowerCAmelCase_ :      Tuple    				=				sag_pipe.to(__lowercase							)
										sag_pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''.'''
										lowerCAmelCase_ :      Dict    				=				torch.manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				sag_pipe(
										    [prompt]   ,						width=7_6_8   ,						height=5_1_2   ,						generator=__lowercase   ,						guidance_scale=7.5   ,						sag_scale=1.0   ,						num_inference_steps=2_0   ,						output_type='''np'''   ,						)

										lowerCAmelCase_ :      Tuple    				=				output.images

										assert image.shape == (1, 5_1_2, 7_6_8, 3)



import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


_UpperCAmelCase			:  Tuple     =logging.get_logger(__name__)







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  List[str]					=							"""linear"""
						SCREAMING_SNAKE_CASE__      :  Union[str, Any]					=							"""cosine"""
						SCREAMING_SNAKE_CASE__      :  Dict					=							"""cosine_with_restarts"""
						SCREAMING_SNAKE_CASE__      :  List[str]					=							"""polynomial"""
						SCREAMING_SNAKE_CASE__      :  Dict					=							"""constant"""
						SCREAMING_SNAKE_CASE__      :  List[str]					=							"""constant_with_warmup"""
						SCREAMING_SNAKE_CASE__      :  str					=							"""piecewise_constant"""






def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ = -1 )-> Tuple:
				return LambdaLR(lowerCAmelCase_ ,   lambda lowerCAmelCase_ : 1 ,   last_epoch=lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = -1 )-> List[Any]:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1.0 ,   lowerCAmelCase_ ) )
								return 1.0

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = -1 )-> int:
				lowerCAmelCase_ :      Optional[int]    				=				{}
				lowerCAmelCase_ :      Union[str, Any]    				=				step_rules.split(''',''' )
				for rule_str in rule_list[:-1]:
								lowerCAmelCase_				,  lowerCAmelCase_ :      str    				=				rule_str.split(''':''' )
								lowerCAmelCase_ :      int    				=				int(lowerCAmelCase_ )
								lowerCAmelCase_ :      str    				=				float(lowerCAmelCase_ )
								lowerCAmelCase_ :      List[Any]    				=				value
				lowerCAmelCase_ :      int    				=				float(rule_list[-1] )

				def create_rules_function(lowerCAmelCase_ ,   lowerCAmelCase_ ):
								def rule_func(lowerCAmelCase_ ) -> float:
												lowerCAmelCase_ :      Tuple    				=				sorted(rules_dict.keys() )
												for i, sorted_step in enumerate(lowerCAmelCase_ ):
																if steps < sorted_step:
																				return rules_dict[sorted_steps[i]]
												return last_lr_multiple

								return rule_func

				lowerCAmelCase_ :      Tuple    				=				create_rules_function(lowerCAmelCase_ ,   lowerCAmelCase_ )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_=-1 )-> Optional[int]:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								return max(
								    0.0 ,   float(num_training_steps - current_step ) / float(max(1 ,   num_training_steps - num_warmup_steps ) ) )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = 0.5 ,   lowerCAmelCase_ = -1 )-> List[Any]:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								lowerCAmelCase_ :      Tuple    				=				float(current_step - num_warmup_steps ) / float(max(1 ,   num_training_steps - num_warmup_steps ) )
								return max(0.0 ,   0.5 * (1.0 + math.cos(math.pi * float(lowerCAmelCase_ ) * 2.0 * progress )) )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = 1 ,   lowerCAmelCase_ = -1 )-> Dict:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								lowerCAmelCase_ :      List[Any]    				=				float(current_step - num_warmup_steps ) / float(max(1 ,   num_training_steps - num_warmup_steps ) )
								if progress >= 1.0:
												return 0.0
								return max(0.0 ,   0.5 * (1.0 + math.cos(math.pi * ((float(lowerCAmelCase_ ) * progress) % 1.0) )) )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_=1e-7 ,   lowerCAmelCase_=1.0 ,   lowerCAmelCase_=-1 )-> Any:
				lowerCAmelCase_ :      Dict    				=				optimizer.defaults['''lr''']
				if not (lr_init > lr_end):
								raise ValueError(f"""lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})""" )

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								elif current_step > num_training_steps:
												return lr_end / lr_init  # as LambdaLR multiplies by lr_init
								else:
												lowerCAmelCase_ :      List[Any]    				=				lr_init - lr_end
												lowerCAmelCase_ :      Optional[Any]    				=				num_training_steps - num_warmup_steps
												lowerCAmelCase_ :      Any    				=				1 - (current_step - num_warmup_steps) / decay_steps
												lowerCAmelCase_ :      List[Any]    				=				lr_range * pct_remaining**power + lr_end
												return decay / lr_init  # as LambdaLR multiplies by lr_init

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )


_UpperCAmelCase			:  Union[str, Any]     ={
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}




def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = None ,   lowerCAmelCase_ = None ,   lowerCAmelCase_ = None ,   lowerCAmelCase_ = 1 ,   lowerCAmelCase_ = 1.0 ,   lowerCAmelCase_ = -1 ,   )-> Optional[int]:
				lowerCAmelCase_ :      Union[str, Any]    				=				SchedulerType(lowerCAmelCase_ )
				lowerCAmelCase_ :      Optional[int]    				=				TYPE_TO_SCHEDULER_FUNCTION[name]
				if name == SchedulerType.CONSTANT:
								return schedule_func(lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )

				if name == SchedulerType.PIECEWISE_CONSTANT:
								return schedule_func(lowerCAmelCase_ ,   step_rules=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )

				# All other schedulers require `num_warmup_steps`
				if num_warmup_steps is None:
								raise ValueError(f"""{name} requires `num_warmup_steps`, please provide that argument.""" )

				if name == SchedulerType.CONSTANT_WITH_WARMUP:
								return schedule_func(lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )

				# All other schedulers require `num_training_steps`
				if num_training_steps is None:
								raise ValueError(f"""{name} requires `num_training_steps`, please provide that argument.""" )

				if name == SchedulerType.COSINE_WITH_RESTARTS:
								return schedule_func(
								    lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   num_training_steps=lowerCAmelCase_ ,   num_cycles=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ ,   )

				if name == SchedulerType.POLYNOMIAL:
								return schedule_func(
								    lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   num_training_steps=lowerCAmelCase_ ,   power=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ ,   )

				return schedule_func(
				    lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   num_training_steps=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )