



"""simple docstring"""
import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
			from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
			from tensorflow.keras.optimizers import Adam


class 							lowerCamelCase_(   tf.keras.optimizers.schedules.LearningRateSchedule    ):





					'''simple docstring'''



					def __init__(    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__ = 1.0				,       lowerCamelCase__ = None				,       ):
									super().__init__()
									_lowerCamelCase        =				initial_learning_rate
									_lowerCamelCase        =				warmup_steps
									_lowerCamelCase        =				power
									_lowerCamelCase        =				decay_schedule_fn
									_lowerCamelCase        =				name

					def __call__(    self				,       lowerCamelCase__   ):
									with tf.name_scope(self.name or '''WarmUp'''   ) as name:
													# Implements polynomial warmup. i.e., if global_step < warmup_steps, the
													# learning rate will be `global_step/num_warmup_steps * init_lr`.
													_lowerCamelCase        =				tf.cast(lowerCamelCase__				,       tf.floataa   )
													_lowerCamelCase        =				tf.cast(self.warmup_steps				,       tf.floataa   )
													_lowerCamelCase        =				global_step_float / warmup_steps_float
													_lowerCamelCase        =				self.initial_learning_rate * tf.math.pow(lowerCamelCase__				,       self.power   )
													return tf.cond(
													    global_step_float < warmup_steps_float				,       lambda: warmup_learning_rate				,       lambda: self.decay_schedule_fn(step - self.warmup_steps   )				,       name=lowerCamelCase__				,       )





					def 			snake_case__      (    self   ):
									return {
									    "initial_learning_rate": self.initial_learning_rate,
									    "decay_schedule_fn": self.decay_schedule_fn,
									    "warmup_steps": self.warmup_steps,
									    "power": self.power,
									    "name": self.name,
									}



def        lowerCAmelCase_(   lowercase_     :							float						,   lowercase_     :							int						,   lowercase_     :							int						,   lowercase_     :							float = 0.0						,   lowercase_     :							float = 0.9						,   lowercase_     :							float = 0.9_9_9						,   lowercase_     :							float = 1e-8						,   lowercase_     :							Optional[float] = None						,   lowercase_     :							Optional[float] = None						,   lowercase_     :							float = 0.0						,   lowercase_     :							float = 1.0						,   lowercase_     :							Optional[List[str]] = None						,   )	->       List[Any]:

				'''simple docstring'''





				_lowerCamelCase        =				tf.keras.optimizers.schedules.PolynomialDecay(
				    initial_learning_rate=lowercase_						,   decay_steps=num_train_steps - num_warmup_steps						,   end_learning_rate=init_lr * min_lr_ratio						,   power=lowercase_						,   )
				if num_warmup_steps:
								_lowerCamelCase        =				WarmUp(
								    initial_learning_rate=lowercase_						,   decay_schedule_fn=lowercase_						,   warmup_steps=lowercase_						,   )
				if weight_decay_rate > 0.0:
								_lowerCamelCase        =				AdamWeightDecay(
								    learning_rate=lowercase_						,   weight_decay_rate=lowercase_						,   beta_a=lowercase_						,   beta_a=lowercase_						,   epsilon=lowercase_						,   clipnorm=lowercase_						,   global_clipnorm=lowercase_						,   exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias''']						,   include_in_weight_decay=lowercase_						,   )
				else:
								_lowerCamelCase        =				tf.keras.optimizers.Adam(
								    learning_rate=lowercase_						,   beta_a=lowercase_						,   beta_a=lowercase_						,   epsilon=lowercase_						,   clipnorm=lowercase_						,   global_clipnorm=lowercase_						,   )
				# We return the optimizer and the LR scheduler in order to better track the
				# evolution of the LR independently of the optimizer.
				return optimizer, lr_schedule


class 							lowerCamelCase_(   A__    ):





					'''simple docstring'''



					def __init__(    self				,       lowerCamelCase__ = 0.0_0_1				,       lowerCamelCase__ = 0.9				,       lowerCamelCase__ = 0.9_9_9				,       lowerCamelCase__ = 1e-7				,       lowerCamelCase__ = False				,       lowerCamelCase__ = 0.0				,       lowerCamelCase__ = None				,       lowerCamelCase__ = None				,       lowerCamelCase__ = "AdamWeightDecay"				,       **lowerCamelCase__				,       ):
									super().__init__(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       **lowerCamelCase__   )
									_lowerCamelCase        =				weight_decay_rate
									_lowerCamelCase        =				include_in_weight_decay
									_lowerCamelCase        =				exclude_from_weight_decay

					@classmethod
					def 			snake_case__      (    cls				,       lowerCamelCase__   ):
									_lowerCamelCase        =				{'''WarmUp''': WarmUp}
									return super(lowerCamelCase__				,       cls   ).from_config(lowerCamelCase__				,       custom_objects=lowerCamelCase__   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									super(lowerCamelCase__				,       self   )._prepare_local(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )
									_lowerCamelCase        =				tf.constant(
									    self.weight_decay_rate				,       name='''adam_weight_decay_rate'''   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				self._do_use_weight_decay(var.name   )
									if do_decay:
													return var.assign_sub(
													    learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]['''weight_decay_rate''']				,       use_locking=self._use_locking				,       )
									return tf.no_op()

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__=None				,       **lowerCamelCase__   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				list(zip(*lowerCamelCase__   )   )
									return super(lowerCamelCase__				,       self   ).apply_gradients(zip(lowerCamelCase__				,       lowerCamelCase__   )				,       name=lowerCamelCase__				,       **lowerCamelCase__   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									if apply_state is None:
													return self._decayed_lr_t[var_dtype], {}

									_lowerCamelCase        =				apply_state or {}
									_lowerCamelCase        =				apply_state.get((var_device, var_dtype)   )
									if coefficients is None:
													_lowerCamelCase        =				self._fallback_apply_state(lowerCamelCase__				,       lowerCamelCase__   )
													_lowerCamelCase        =				coefficients

									return coefficients["lr_t"], {"apply_state": apply_state}

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__=None   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				self._get_lr(var.device				,       var.dtype.base_dtype				,       lowerCamelCase__   )
									_lowerCamelCase        =				self._decay_weights_op(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )
									with tf.control_dependencies([decay]   ):
													return super(lowerCamelCase__				,       self   )._resource_apply_dense(lowerCamelCase__				,       lowerCamelCase__				,       **lowerCamelCase__   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__=None   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				self._get_lr(var.device				,       var.dtype.base_dtype				,       lowerCamelCase__   )
									_lowerCamelCase        =				self._decay_weights_op(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )
									with tf.control_dependencies([decay]   ):
													return super(lowerCamelCase__				,       self   )._resource_apply_sparse(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       **lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				super().get_config()
									config.update({'''weight_decay_rate''': self.weight_decay_rate}   )
									return config





					def 			snake_case__      (    self				,       lowerCamelCase__   ):
									if self.weight_decay_rate == 0:
													return False

									if self._include_in_weight_decay:
													for r in self._include_in_weight_decay:
																	if re.search(lowerCamelCase__				,       lowerCamelCase__   ) is not None:
																					return True

									if self._exclude_from_weight_decay:
													for r in self._exclude_from_weight_decay:
																	if re.search(lowerCamelCase__				,       lowerCamelCase__   ) is not None:
																					return False
									return True


class 							lowerCamelCase_(   A__    ):





					'''simple docstring'''



					def __init__(    self   ):
									_lowerCamelCase        =				[]
									_lowerCamelCase        =				None

					@property
					def 			snake_case__      (    self   ):
									if self._accum_steps is None:
													_lowerCamelCase        =				tf.Variable(
													    tf.constant(0				,       dtype=tf.intaa   )				,       trainable=lowerCamelCase__				,       synchronization=tf.VariableSynchronization.ON_READ				,       aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA				,       )

									return self._accum_steps.value()

					@property
					def 			snake_case__      (    self   ):
									if not self._gradients:
													raise ValueError('''The accumulator should be called first to initialize the gradients'''   )
									return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]

					def __call__(    self				,       lowerCamelCase__   ):
									if not self._gradients:
													_lowerCamelCase        =				self.step  # Create the step variable.
													self._gradients.extend(
													    [
													        tf.Variable(
													            tf.zeros_like(lowerCamelCase__   )				,       trainable=lowerCamelCase__				,       synchronization=tf.VariableSynchronization.ON_READ				,       aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA				,       )
													        if gradient is not None
													        else gradient
													        for gradient in gradients
													    ]   )
									if len(lowerCamelCase__   ) != len(self._gradients   ):
													raise ValueError(F"""Expected {len(self._gradients   )} gradients, but got {len(lowerCamelCase__   )}"""   )

									for accum_gradient, gradient in zip(self._gradients				,       lowerCamelCase__   ):
													if accum_gradient is not None and gradient is not None:
																	accum_gradient.assign_add(lowerCamelCase__   )

									self._accum_steps.assign_add(1   )





					def 			snake_case__      (    self   ):
									if not self._gradients:
													return
									self._accum_steps.assign(0   )
									for gradient in self._gradients:
													if gradient is not None:
																	gradient.assign(tf.zeros_like(lowerCamelCase__   )   )





"""simple docstring"""
import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 							lowerCamelCase_(   A__,				A__,				unittest.TestCase    ):





					'''simple docstring'''



					lowercase__   :       Tuple										= StableDiffusionXLImgaImgPipeline
					lowercase__   :       List[str]										= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
					lowercase__   :       int										= PipelineTesterMixin.required_optional_params - {'latents'}
					lowercase__   :       Union[str, Any]										= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
					lowercase__   :       Dict										= IMAGE_TO_IMAGE_IMAGE_PARAMS
					lowercase__   :       List[str]										= IMAGE_TO_IMAGE_IMAGE_PARAMS

					def 			snake_case__      (    self   ):
									torch.manual_seed(0   )
									_lowerCamelCase        =				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''')				,       attention_head_dim=(2, 4)				,       use_linear_projection=lowerCamelCase__				,       addition_embed_type='''text_time'''				,       addition_time_embed_dim=8				,       transformer_layers_per_block=(1, 2)				,       projection_class_embeddings_input_dim=8_0				,       cross_attention_dim=6_4				,       )
									_lowerCamelCase        =				EulerDiscreteScheduler(
									    beta_start=0.0_0_0_8_5				,       beta_end=0.0_1_2				,       steps_offset=1				,       beta_schedule='''scaled_linear'''				,       timestep_spacing='''leading'''				,       )
									torch.manual_seed(0   )
									_lowerCamelCase        =				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				,       sample_size=1_2_8				,       )
									torch.manual_seed(0   )
									_lowerCamelCase        =				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				,       hidden_act='''gelu'''				,       projection_dim=3_2				,       )
									_lowerCamelCase        =				CLIPTextModel(lowerCamelCase__   )
									_lowerCamelCase        =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				,       local_files_only=lowerCamelCase__   )

									_lowerCamelCase        =				CLIPTextModelWithProjection(lowerCamelCase__   )
									_lowerCamelCase        =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				,       local_files_only=lowerCamelCase__   )

									_lowerCamelCase        =				{
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									    '''text_encoder''': text_encoder,
									    '''tokenizer''': tokenizer,
									    '''text_encoder_2''': text_encoder_a,
									    '''tokenizer_2''': tokenizer_a,
									    # "safety_checker": None,
									    # "feature_extractor": None,
									}
									return components

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__=0   ):
									_lowerCamelCase        =				floats_tensor((1, 3, 3_2, 3_2)				,       rng=random.Random(lowerCamelCase__   )   ).to(lowerCamelCase__   )
									_lowerCamelCase        =				image / 2 + 0.5
									if str(lowerCamelCase__   ).startswith('''mps'''   ):
													_lowerCamelCase        =				torch.manual_seed(lowerCamelCase__   )
									else:
													_lowerCamelCase        =				torch.Generator(device=lowerCamelCase__   ).manual_seed(lowerCamelCase__   )
									_lowerCamelCase        =				{
									    '''prompt''': '''A painting of a squirrel eating a burger''',
									    '''image''': image,
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''guidance_scale''': 5.0,
									    '''output_type''': '''numpy''',
									    '''strength''': 0.7_5,
									}
									return inputs

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				StableDiffusionXLImgaImgPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe.to(lowerCamelCase__   )
									sd_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe(**lowerCamelCase__   ).images
									_lowerCamelCase        =				image[0, -3:, -3:, -1]

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

									_lowerCamelCase        =				np.array([0.4_6_5_6, 0.4_8_4_0, 0.4_4_3_9, 0.6_6_9_8, 0.5_5_7_4, 0.4_5_2_4, 0.5_7_9_9, 0.5_9_4_3, 0.5_1_6_5]   )

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

					def 			snake_case__      (    self   ):
									super().test_attention_slicing_forward_pass(expected_max_diff=3e-3   )

					def 			snake_case__      (    self   ):
									super().test_inference_batch_single_identical(expected_max_diff=3e-3   )

					def 			snake_case__      (    self   ):
									pass





					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				StableDiffusionXLImgaImgPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe.to(lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe.to(lowerCamelCase__   )
									sd_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									# forward without prompt embeds
									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				3 * ['''this is a negative prompt''']
									_lowerCamelCase        =				negative_prompt
									_lowerCamelCase        =				3 * [inputs['''prompt''']]

									_lowerCamelCase        =				sd_pipe(**lowerCamelCase__   )
									_lowerCamelCase        =				output.images[0, -3:, -3:, -1]

									# forward with prompt embeds
									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				3 * ['''this is a negative prompt''']
									_lowerCamelCase        =				3 * [inputs.pop('''prompt'''   )]

									(
									    (
									    _lowerCamelCase
									)  ,     (
									    _lowerCamelCase
									)  ,     (
									    _lowerCamelCase
									)  ,     (
									    _lowerCamelCase
									)  ,     
									)        =				sd_pipe.encode_prompt(lowerCamelCase__				,       negative_prompt=lowerCamelCase__   )

									_lowerCamelCase        =				sd_pipe(
									    **lowerCamelCase__				,       prompt_embeds=lowerCamelCase__				,       negative_prompt_embeds=lowerCamelCase__				,       pooled_prompt_embeds=lowerCamelCase__				,       negative_pooled_prompt_embeds=lowerCamelCase__				,       )
									_lowerCamelCase        =				output.images[0, -3:, -3:, -1]

									# make sure that it's equal
									assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()   ).max() < 1e-4


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





					'''simple docstring'''



					def 			snake_case__      (    self   ):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__="cpu"				,       lowerCamelCase__=torch.floataa				,       lowerCamelCase__=0   ):
									_lowerCamelCase        =				torch.Generator(device=lowerCamelCase__   ).manual_seed(lowerCamelCase__   )
									_lowerCamelCase        =				np.random.RandomState(lowerCamelCase__   ).standard_normal((1, 4, 6_4, 6_4)   )
									_lowerCamelCase        =				torch.from_numpy(lowerCamelCase__   ).to(device=lowerCamelCase__				,       dtype=lowerCamelCase__   )
									_lowerCamelCase        =				{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''latents''': latents,
									    '''generator''': generator,
									    '''num_inference_steps''': 3,
									    '''guidance_scale''': 7.5,
									    '''output_type''': '''numpy''',
									}
									return inputs





					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''   )
									pipe.to(lowerCamelCase__   )
									pipe.set_progress_bar_config(disable=lowerCamelCase__   )

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

									assert image.shape == (1, 5_1_2, 5_1_2, 3)
									_lowerCamelCase        =				np.array([0.4_9_4_9_3, 0.4_7_8_9_6, 0.4_0_7_9_8, 0.5_4_2_1_4, 0.5_3_2_1_2, 0.4_8_2_0_2, 0.4_7_6_5_6, 0.4_6_3_2_9, 0.4_8_5_0_6]   )
									assert np.abs(image_slice - expected_slice   ).max() < 7e-3
