






'''simple docstring'''


import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()
class    _SCREAMING_SNAKE_CASE							(		__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE       ,unittest.TestCase	):


			'''simple docstring'''





			__a							:							int					     =					StableUnCLIPPipeline
			__a							:							int					     =					TEXT_TO_IMAGE_PARAMS
			__a							:							int					     =					TEXT_TO_IMAGE_BATCH_PARAMS
			__a							:							Dict					     =					TEXT_TO_IMAGE_IMAGE_PARAMS
			__a							:							Any					     =					TEXT_TO_IMAGE_IMAGE_PARAMS

			# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
			__a							:							str					     =					False






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

									'''simple docstring'''







									UpperCamelCase__         = 3_2
									UpperCamelCase__         = embedder_hidden_size

									# prior components

									torch.manual_seed(0			)
									UpperCamelCase__         = CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

									torch.manual_seed(0			)
									UpperCamelCase__         = CLIPTextModelWithProjection(
									    CLIPTextConfig(
									        bos_token_id=0 ,      eos_token_id=2 ,      hidden_size=__snake_case ,      projection_dim=__snake_case ,      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 ,      )			)

									torch.manual_seed(0			)
									UpperCamelCase__         = PriorTransformer(
									    num_attention_heads=2 ,      attention_head_dim=1_2 ,      embedding_dim=__snake_case ,      num_layers=1 ,      )

									torch.manual_seed(0			)
									UpperCamelCase__         = DDPMScheduler(
									    variance_type="""fixed_small_log""" ,      prediction_type="""sample""" ,      num_train_timesteps=1_0_0_0 ,      clip_sample=__snake_case ,      clip_sample_range=5.0 ,      beta_schedule="""squaredcos_cap_v2""" ,      )

									# regular denoising components

									torch.manual_seed(0			)
									UpperCamelCase__         = StableUnCLIPImageNormalizer(embedding_dim=__snake_case			)
									UpperCamelCase__         = DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""			)

									torch.manual_seed(0			)
									UpperCamelCase__         = CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

									torch.manual_seed(0			)
									UpperCamelCase__         = CLIPTextModel(
									    CLIPTextConfig(
									        bos_token_id=0 ,      eos_token_id=2 ,      hidden_size=__snake_case ,      projection_dim=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 ,      )			)

									torch.manual_seed(0			)
									UpperCamelCase__         = UNetaDConditionModel(
									    sample_size=3_2 ,      in_channels=4 ,      out_channels=4 ,      down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""") ,      block_out_channels=(3_2, 6_4) ,      attention_head_dim=(2, 4) ,      class_embed_type="""projection""" ,      projection_class_embeddings_input_dim=embedder_projection_dim * 2 ,      cross_attention_dim=__snake_case ,      layers_per_block=1 ,      upcast_attention=__snake_case ,      use_linear_projection=__snake_case ,      )

									torch.manual_seed(0			)
									UpperCamelCase__         = DDIMScheduler(
									    beta_schedule="""scaled_linear""" ,      beta_start=0.0_0_0_8_5 ,      beta_end=0.0_1_2 ,      prediction_type="""v_prediction""" ,      set_alpha_to_one=__snake_case ,      steps_offset=1 ,      )

									torch.manual_seed(0			)
									UpperCamelCase__         = AutoencoderKL()

									UpperCamelCase__         = {
									    # prior components
									    '''prior_tokenizer''': prior_tokenizer,
									    '''prior_text_encoder''': prior_text_encoder,
									    '''prior''': prior,
									    '''prior_scheduler''': prior_scheduler,
									    # image noising components
									    '''image_normalizer''': image_normalizer,
									    '''image_noising_scheduler''': image_noising_scheduler,
									    # regular denoising components
									    '''tokenizer''': tokenizer,
									    '''text_encoder''': text_encoder,
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									}

									return components






			def 				A      (   self							:						int ,      lowercase							:						Tuple ,      lowercase							:						str=0			)							->    Optional[int]:

									'''simple docstring'''







									if str(__snake_case			).startswith("""mps"""			):
															UpperCamelCase__         = torch.manual_seed(__snake_case			)
									else:
															UpperCamelCase__         = torch.Generator(device=__snake_case			).manual_seed(__snake_case			)
									UpperCamelCase__         = {
									    '''prompt''': '''A painting of a squirrel eating a burger''',
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''prior_num_inference_steps''': 2,
									    '''output_type''': '''numpy''',
									}
									return inputs






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

									'''simple docstring'''







									UpperCamelCase__         = torch_device == '''cpu'''

									self._test_attention_slicing_forward_pass(test_max_difference=__snake_case			)







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

									'''simple docstring'''







									UpperCamelCase__         = torch_device in ['''cpu''', '''mps''']

									self._test_inference_batch_single_identical(test_max_difference=__snake_case			)
@slow
@require_torch_gpu
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			def 				A      (   self							:						Tuple			)							->    Any:

									'''simple docstring'''







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






			def 				A      (   self							:						Any			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""			)

									UpperCamelCase__         = StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l""" ,      torch_dtype=torch.floataa			)
									pipe.to(__snake_case			)
									pipe.set_progress_bar_config(disable=__snake_case			)
									# stable unclip will oom when integration tests are run on a V100,
									# so turn on memory savings
									pipe.enable_attention_slicing()
									pipe.enable_sequential_cpu_offload()

									UpperCamelCase__         = torch.Generator(device="""cpu"""			).manual_seed(0			)
									UpperCamelCase__         = pipe("""anime turle""" ,      generator=__snake_case ,      output_type="""np"""			)

									UpperCamelCase__         = output.images[0]

									assert image.shape == (7_6_8, 7_6_8, 3)

									assert_mean_pixel_difference(__snake_case ,      __snake_case			)







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

									'''simple docstring'''







									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									UpperCamelCase__         = StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l""" ,      torch_dtype=torch.floataa			)
									UpperCamelCase__         = pipe.to(__snake_case			)
									pipe.set_progress_bar_config(disable=__snake_case			)
									pipe.enable_attention_slicing()
									pipe.enable_sequential_cpu_offload()

									UpperCamelCase__         = pipe(
									    """anime turtle""" ,      prior_num_inference_steps=2 ,      num_inference_steps=2 ,      output_type="""np""" ,      )

									UpperCamelCase__         = torch.cuda.max_memory_allocated()
									# make sure that less than 7 GB is allocated
									assert mem_bytes < 7 * 1_0**9








'''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			)
