






'''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    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,unittest.TestCase	):


			'''simple docstring'''





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






			def 				A      (   self							:						Tuple			)							->    Dict:

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = 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=lowercase ,      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 ,      )
									UpperCamelCase__         = 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			)
									UpperCamelCase__         = 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			)
									UpperCamelCase__         = 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 ,      )
									UpperCamelCase__         = CLIPTextModel(lowercase			)
									UpperCamelCase__         = CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" ,      local_files_only=lowercase			)

									UpperCamelCase__         = CLIPTextModelWithProjection(lowercase			)
									UpperCamelCase__         = CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" ,      local_files_only=lowercase			)

									UpperCamelCase__         = {
									    """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 				A      (   self							:						Optional[Any] ,      lowercase							:						Dict ,      lowercase							:						List[str]=0			)							->    Any:

									'''simple docstring'''







									UpperCamelCase__         = floats_tensor((1, 3, 3_2, 3_2) ,      rng=random.Random(lowercase			)			).to(lowercase			)
									UpperCamelCase__         = image / 2 + 0.5
									if str(lowercase			).startswith("""mps"""			):
															UpperCamelCase__         = torch.manual_seed(lowercase			)
									else:
															UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(lowercase			)
									UpperCamelCase__         = {
									    """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 				A      (   self							:						str			)							->    Dict:

									'''simple docstring'''







									UpperCamelCase__         = """cpu"""  # ensure determinism for the device-dependent torch.Generator
									UpperCamelCase__         = self.get_dummy_components()
									UpperCamelCase__         = StableDiffusionXLImgaImgPipeline(**lowercase			)
									UpperCamelCase__         = sd_pipe.to(lowercase			)
									sd_pipe.set_progress_bar_config(disable=lowercase			)

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

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

									UpperCamelCase__         = 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 				A      (   self							:						List[Any]			)							->    Dict:

									'''simple docstring'''







									super().test_attention_slicing_forward_pass(expected_max_diff=3e-3			)






			def 				A      (   self							:						Union[str, Any]			)							->    int:

									'''simple docstring'''







									super().test_inference_batch_single_identical(expected_max_diff=3e-3			)






			def 				A      (   self							:						List[str]			)							->    Tuple:

									'''simple docstring'''







									pass







			def 				A      (   self							:						List[Any]			)							->    Optional[Any]:

									'''simple docstring'''







									UpperCamelCase__         = self.get_dummy_components()
									UpperCamelCase__         = StableDiffusionXLImgaImgPipeline(**lowercase			)
									UpperCamelCase__         = sd_pipe.to(lowercase			)
									UpperCamelCase__         = sd_pipe.to(lowercase			)
									sd_pipe.set_progress_bar_config(disable=lowercase			)

									# forward without prompt embeds
									UpperCamelCase__         = self.get_dummy_inputs(lowercase			)
									UpperCamelCase__         = 3 * ["""this is a negative prompt"""]
									UpperCamelCase__         = negative_prompt
									UpperCamelCase__         = 3 * [inputs["""prompt"""]]

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

									# forward with prompt embeds
									UpperCamelCase__         = self.get_dummy_inputs(lowercase			)
									UpperCamelCase__         = 3 * ["""this is a negative prompt"""]
									UpperCamelCase__         = 3 * [inputs.pop("""prompt"""			)]

									(
									    (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   
									)         = sd_pipe.encode_prompt(lowercase ,      negative_prompt=lowercase			)

									UpperCamelCase__         = sd_pipe(
									    **lowercase ,      prompt_embeds=lowercase ,      negative_prompt_embeds=lowercase ,      pooled_prompt_embeds=lowercase ,      negative_pooled_prompt_embeds=lowercase ,      )
									UpperCamelCase__         = 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    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			def 				A      (   self							:						Optional[Any]			)							->    int:

									'''simple docstring'''







									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






			def 				A      (   self							:						Optional[Any] ,      lowercase							:						Any ,      lowercase							:						List[Any]="cpu" ,      lowercase							:						Union[str, Any]=torch.floataa ,      lowercase							:						int=0			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(lowercase			)
									UpperCamelCase__         = np.random.RandomState(lowercase			).standard_normal((1, 4, 6_4, 6_4)			)
									UpperCamelCase__         = torch.from_numpy(lowercase			).to(device=lowercase ,      dtype=lowercase			)
									UpperCamelCase__         = {
									    """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 				A      (   self							:						Any			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""			)
									pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

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

									assert image.shape == (1, 5_1_2, 5_1_2, 3)
									UpperCamelCase__         = 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








'''simple docstring'''


# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings('''ignore''', category=UserWarning, module='''torch.optim.lr_scheduler''')
class    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





			def __init__(   self							:						Optional[int] ,      lowercase							:						Optional[int] ,      lowercase							:						Optional[Any] ,      lowercase							:						bool = True ,      lowercase							:						bool = False			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = scheduler
									UpperCamelCase__         = optimizers if isinstance(lowercase ,      (list, tuple)			) else [optimizers]
									UpperCamelCase__         = split_batches
									UpperCamelCase__         = step_with_optimizer
									UpperCamelCase__         = GradientState()






			def 				A      (   self							:						int ,      *lowercase							:						int ,      **lowercase							:						str			)							->    str:

									'''simple docstring'''







									if not self.step_with_optimizer:
															# No link between scheduler and optimizer -> just step
															self.scheduler.step(*lowercase ,      **lowercase			)
															return

									# Otherwise, first make sure the optimizer was stepped.
									if not self.gradient_state.sync_gradients:
															if self.gradient_state.adjust_scheduler:
																					self.scheduler._step_count += 1
															return

									for opt in self.optimizers:
															if opt.step_was_skipped:
																					return
									if self.split_batches:
															# Split batches -> the training dataloader batch size is not changed so one step per training step
															self.scheduler.step(*lowercase ,      **lowercase			)
									else:
															# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
															# num_processes steps per training step
															UpperCamelCase__         = AcceleratorState().num_processes
															for _ in range(lowercase			):
																					# Special case when using OneCycle and `drop_last` was not used
																					if hasattr(self.scheduler ,      """total_steps"""			):
																											if self.scheduler._step_count <= self.scheduler.total_steps:
																																	self.scheduler.step(*lowercase ,      **lowercase			)
																					else:
																											self.scheduler.step(*lowercase ,      **lowercase			)






			def 				A      (   self							:						List[str]			)							->    Optional[Any]:

									'''simple docstring'''







									return self.scheduler.get_last_lr()






			def 				A      (   self							:						Any			)							->    int:

									'''simple docstring'''







									return self.scheduler.state_dict()






			def 				A      (   self							:						Any ,      lowercase							:						int			)							->    Optional[int]:

									'''simple docstring'''







									self.scheduler.load_state_dict(lowercase			)






			def 				A      (   self							:						str			)							->    int:

									'''simple docstring'''







									return self.scheduler.get_lr()







			def 				A      (   self							:						Union[str, Any] ,      *lowercase							:						List[Any] ,      **lowercase							:						str			)							->    str:

									'''simple docstring'''







									return self.scheduler.print_lr(*lowercase ,      **lowercase			)
