





'''simple docstring'''



import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




def 				snake_case_      (		SCREAMING_SNAKE_CASE__       ):







  '''simple docstring'''



  _snake_case     		=				[tensor.shape for tensor in tensor_list]
  return all(shape == shapes[0] for shape in shapes[1:]       )


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase					,    __UpperCamelCase					,    __UpperCamelCase					,    unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Union[str, Any]      =   StableDiffusionLatentUpscalePipeline
     UpperCAmelCase__       :						List[str]      =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
         '''height''',
         '''width''',
         '''cross_attention_kwargs''',
         '''negative_prompt_embeds''',
         '''prompt_embeds''',
     }
     UpperCAmelCase__       :						Dict      =   PipelineTesterMixin.required_optional_params - {'''num_images_per_prompt'''}
     UpperCAmelCase__       :						Tuple      =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     UpperCAmelCase__       :						List[str]      =   frozenset(
         []			)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
     UpperCAmelCase__       :						Any      =   frozenset([]			)

     UpperCAmelCase__       :						Any      =   True

     @property
     def 	UpperCamelCase(					self	):
       _snake_case     		=				1
       _snake_case     		=				4
       _snake_case     		=				(16, 16)

       _snake_case     		=				floats_tensor((batch_size, num_channels) + sizes						,   rng=random.Random(0	)	).to(lowerCamelCase	)
       return image

     def 	UpperCamelCase(					self	):
       torch.manual_seed(0	)
       _snake_case     		=				UNetaDConditionModel(
           act_fn="gelu"						,   attention_head_dim=8						,   norm_num_groups=lowerCamelCase						,   block_out_channels=[32, 32, 64, 64]						,   time_cond_proj_dim=160						,   conv_in_kernel=1						,   conv_out_kernel=1						,   cross_attention_dim=32						,   down_block_types=(
               "KDownBlock2D",
               "KCrossAttnDownBlock2D",
               "KCrossAttnDownBlock2D",
               "KCrossAttnDownBlock2D",
           )						,   in_channels=8						,   mid_block_type=lowerCamelCase						,   only_cross_attention=lowerCamelCase						,   out_channels=5						,   resnet_time_scale_shift="scale_shift"						,   time_embedding_type="fourier"						,   timestep_post_act="gelu"						,   up_block_types=("KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KUpBlock2D")						,   )
       _snake_case     		=				AutoencoderKL(
           block_out_channels=[32, 32, 64, 64]						,   in_channels=3						,   out_channels=3						,   down_block_types=[
               "DownEncoderBlock2D",
               "DownEncoderBlock2D",
               "DownEncoderBlock2D",
               "DownEncoderBlock2D",
           ]						,   up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]						,   latent_channels=4						,   )
       _snake_case     		=				EulerDiscreteScheduler(prediction_type="sample"	)
       _snake_case     		=				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=1_000						,   hidden_act="quick_gelu"						,   projection_dim=512						,   )
       _snake_case     		=				CLIPTextModel(lowerCamelCase	)
       _snake_case     		=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"	)

       _snake_case     		=				{
           "unet": model.eval(),
           "vae": vae.eval(),
           "scheduler": scheduler,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
       }

       return components

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase=0	):
       if str(lowerCamelCase	).startswith("mps"	):
         _snake_case     		=				torch.manual_seed(lowerCamelCase	)
       else:
         _snake_case     		=				torch.Generator(device=lowerCamelCase	).manual_seed(lowerCamelCase	)
       _snake_case     		=				{
           "prompt": "A painting of a squirrel eating a burger",
           "image": self.dummy_image.cpu(),
           "generator": generator,
           "num_inference_steps": 2,
           "output_type": "numpy",
       }
       return inputs

     def 	UpperCamelCase(					self	):
       _snake_case     		=				"cpu"

       _snake_case     		=				self.get_dummy_components()
       _snake_case     		=				self.pipeline_class(**lowerCamelCase	)
       pipe.to(lowerCamelCase	)
       pipe.set_progress_bar_config(disable=lowerCamelCase	)

       _snake_case     		=				self.get_dummy_inputs(lowerCamelCase	)
       _snake_case     		=				pipe(**lowerCamelCase	).images
       _snake_case     		=				image[0, -3:, -3:, -1]

       self.assertEqual(image.shape						,   (1, 256, 256, 3)	)
       _snake_case     		=				np.array(
           [0.47222412, 0.41921633, 0.44717434, 0.46874192, 0.42588258, 0.46150726, 0.4677534, 0.45583832, 0.48579055]	)
       _snake_case     		=				np.abs(image_slice.flatten() - expected_slice	).max()
       self.assertLessEqual(lowerCamelCase						,   1e-3	)

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

     def 	UpperCamelCase(					self	):
       super().test_cpu_offload_forward_pass(expected_max_diff=3e-3	)

     def 	UpperCamelCase(					self	):
       super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3	)

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

     def 	UpperCamelCase(					self	):
       super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3	)

     def 	UpperCamelCase(					self	):
       super().test_save_load_local(expected_max_difference=3e-3	)

     def 	UpperCamelCase(					self	):
       super().test_save_load_optional_components(expected_max_difference=3e-3	)






     def 	UpperCamelCase(					self	):
       _snake_case     		=				[
           "DDIMScheduler",
           "DDPMScheduler",
           "PNDMScheduler",
           "HeunDiscreteScheduler",
           "EulerAncestralDiscreteScheduler",
           "KDPM2DiscreteScheduler",
           "KDPM2AncestralDiscreteScheduler",
           "DPMSolverSDEScheduler",
       ]
       _snake_case     		=				self.get_dummy_components()
       _snake_case     		=				self.pipeline_class(**lowerCamelCase	)

       # make sure that PNDM does not need warm-up
       pipe.scheduler.register_to_config(skip_prk_steps=lowerCamelCase	)

       pipe.to(lowerCamelCase	)
       pipe.set_progress_bar_config(disable=lowerCamelCase	)
       _snake_case     		=				self.get_dummy_inputs(lowerCamelCase	)
       _snake_case     		=				2

       _snake_case     		=				[]
       for scheduler_enum in KarrasDiffusionSchedulers:
         if scheduler_enum.name in skip_schedulers:
           # no sigma schedulers are not supported
           # no schedulers
           continue

         _snake_case     		=				getattr(lowerCamelCase						,   scheduler_enum.name	)
         _snake_case     		=				scheduler_cls.from_config(pipe.scheduler.config	)
         _snake_case     		=				pipe(**lowerCamelCase	)[0]
         outputs.append(lowerCamelCase	)

       assert check_same_shape(lowerCamelCase	)







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



     '''simple docstring'''






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

     def 	UpperCamelCase(					self	):
       _snake_case     		=				torch.manual_seed(33	)

       _snake_case     		=				StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"						,   torch_dtype=torch.floataa	)
       pipe.to("cuda"	)

       _snake_case     		=				StableDiffusionLatentUpscalePipeline.from_pretrained(
           "stabilityai/sd-x2-latent-upscaler"						,   torch_dtype=torch.floataa	)
       upscaler.to("cuda"	)

       _snake_case     		=				"a photo of an astronaut high resolution, unreal engine, ultra realistic"

       _snake_case     		=				pipe(lowerCamelCase						,   generator=lowerCamelCase						,   output_type="latent"	).images

       _snake_case     		=				upscaler(
           prompt=lowerCamelCase						,   image=lowerCamelCase						,   num_inference_steps=20						,   guidance_scale=0						,   generator=lowerCamelCase						,   output_type="np"						,   ).images[0]

       _snake_case     		=				load_numpy(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy"	)
       assert np.abs((expected_image - image).mean()	) < 5e-2






     def 	UpperCamelCase(					self	):
       _snake_case     		=				torch.manual_seed(33	)

       _snake_case     		=				StableDiffusionLatentUpscalePipeline.from_pretrained(
           "stabilityai/sd-x2-latent-upscaler"						,   torch_dtype=torch.floataa	)
       upscaler.to("cuda"	)

       _snake_case     		=				"the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas"

       _snake_case     		=				load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png"	)

       _snake_case     		=				upscaler(
           prompt=lowerCamelCase						,   image=lowerCamelCase						,   num_inference_steps=20						,   guidance_scale=0						,   generator=lowerCamelCase						,   output_type="np"						,   ).images[0]

       _snake_case     		=				load_numpy(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy"	)
       assert np.abs((expected_image - image).max()	) < 5e-2









'''simple docstring'''



import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     Adafactor,
     AdamW,
     get_constant_schedule,
     get_constant_schedule_with_warmup,
     get_cosine_schedule_with_warmup,
     get_cosine_with_hard_restarts_schedule_with_warmup,
     get_inverse_sqrt_schedule,
     get_linear_schedule_with_warmup,
     get_polynomial_decay_schedule_with_warmup,
 )




def 				snake_case_      (		SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__=10       ):







  '''simple docstring'''



  _snake_case     		=				[]
  for _ in range(SCREAMING_SNAKE_CASE__       ):
    lrs.append(scheduler.get_lr()[0]       )
    scheduler.step()
  return lrs




def 				snake_case_      (		SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__=10       ):







  '''simple docstring'''



  _snake_case     		=				[]
  for step in range(SCREAMING_SNAKE_CASE__       ):
    lrs.append(scheduler.get_lr()[0]       )
    scheduler.step()
    if step == num_steps // 2:
      with tempfile.TemporaryDirectory() as tmpdirname:
        _snake_case     		=				os.path.join(SCREAMING_SNAKE_CASE__					,	"schedule.bin"       )
        torch.save(scheduler.state_dict()					,	SCREAMING_SNAKE_CASE__       )

        _snake_case     		=				torch.load(SCREAMING_SNAKE_CASE__       )
        scheduler.load_state_dict(SCREAMING_SNAKE_CASE__       )
  return lrs


@require_torch
class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       self.assertEqual(len(lowerCamelCase	)						,   len(lowerCamelCase	)	)
       for a, b in zip(lowerCamelCase						,   lowerCamelCase	):
         self.assertAlmostEqual(lowerCamelCase						,   lowerCamelCase						,   delta=lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				torch.tensor([0.1, -0.2, -0.1]						,   requires_grad=lowerCamelCase	)
       _snake_case     		=				torch.tensor([0.4, 0.2, -0.5]	)
       _snake_case     		=				nn.MSELoss()
       # No warmup, constant schedule, no gradient clipping
       _snake_case     		=				AdamW(params=[w]						,   lr=2e-1						,   weight_decay=0.0	)
       for _ in range(100	):
         _snake_case     		=				criterion(lowerCamelCase						,   lowerCamelCase	)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
       self.assertListAlmostEqual(w.tolist()						,   [0.4, 0.2, -0.5]						,   tol=1e-2	)






     def 	UpperCamelCase(					self	):
       _snake_case     		=				torch.tensor([0.1, -0.2, -0.1]						,   requires_grad=lowerCamelCase	)
       _snake_case     		=				torch.tensor([0.4, 0.2, -0.5]	)
       _snake_case     		=				nn.MSELoss()
       # No warmup, constant schedule, no gradient clipping
       _snake_case     		=				Adafactor(
           params=[w]						,   lr=1e-2						,   eps=(1e-30, 1e-3)						,   clip_threshold=1.0						,   decay_rate=-0.8						,   betaa=lowerCamelCase						,   weight_decay=0.0						,   relative_step=lowerCamelCase						,   scale_parameter=lowerCamelCase						,   warmup_init=lowerCamelCase						,   )
       for _ in range(1_000	):
         _snake_case     		=				criterion(lowerCamelCase						,   lowerCamelCase	)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
       self.assertListAlmostEqual(w.tolist()						,   [0.4, 0.2, -0.5]						,   tol=1e-2	)







@require_torch
class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Any      =   nn.Linear(50					,    50			) if is_torch_available() else None
     UpperCAmelCase__       :						str      =   AdamW(m.parameters()					,    lr=10.0			) if is_torch_available() else None
     UpperCAmelCase__       :						Optional[Any]      =   10

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase=None	):
       self.assertEqual(len(lowerCamelCase	)						,   len(lowerCamelCase	)	)
       for a, b in zip(lowerCamelCase						,   lowerCamelCase	):
         self.assertAlmostEqual(lowerCamelCase						,   lowerCamelCase						,   delta=lowerCamelCase						,   msg=lowerCamelCase	)






     def 	UpperCamelCase(					self	):
       _snake_case     		=				{"num_warmup_steps": 2, "num_training_steps": 10}
       # schedulers doct format
       # function: (sched_args_dict, expected_learning_rates)
       _snake_case     		=				{
           get_constant_schedule: ({}, [10.0] * self.num_steps),
           get_constant_schedule_with_warmup: (
               {"num_warmup_steps": 4},
               [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
           ),
           get_linear_schedule_with_warmup: (
               {**common_kwargs},
               [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
           ),
           get_cosine_schedule_with_warmup: (
               {**common_kwargs},
               [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
           ),
           get_cosine_with_hard_restarts_schedule_with_warmup: (
               {**common_kwargs, "num_cycles": 2},
               [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
           ),
           get_polynomial_decay_schedule_with_warmup: (
               {**common_kwargs, "power": 2.0, "lr_end": 1e-7},
               [0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
           ),
           get_inverse_sqrt_schedule: (
               {"num_warmup_steps": 2},
               [0.0, 5.0, 10.0, 8.165, 7.071, 6.325, 5.774, 5.345, 5.0, 4.714],
           ),
       }

       for scheduler_func, data in scheds.items():
         _snake_case		,		_snake_case     		=				data

         _snake_case     		=				scheduler_func(self.optimizer						,   **lowerCamelCase	)
         self.assertEqual(len([scheduler.get_lr()[0]]	)						,   1	)
         _snake_case     		=				unwrap_schedule(lowerCamelCase						,   self.num_steps	)
         self.assertListAlmostEqual(
             lowerCamelCase						,   lowerCamelCase						,   tol=1e-2						,   msg=F'''failed for {scheduler_func} in normal scheduler'''						,   )

         _snake_case     		=				scheduler_func(self.optimizer						,   **lowerCamelCase	)
         if scheduler_func.__name__ != "get_constant_schedule":
           LambdaScheduleWrapper.wrap_scheduler(lowerCamelCase	)  # wrap to test picklability of the schedule
         _snake_case     		=				unwrap_and_save_reload_schedule(lowerCamelCase						,   self.num_steps	)
         self.assertListEqual(lowerCamelCase						,   lowerCamelCase						,   msg=F'''failed for {scheduler_func} in save and reload'''	)







class       __SCREAMING_SNAKE_CASE	:



     '''simple docstring'''






     def __init__(					self						,   lowerCamelCase	):
       _snake_case     		=				fn

     def __call__(					self						,   *lowerCamelCase						,   **lowerCamelCase	):
       return self.fn(*lowerCamelCase						,   **lowerCamelCase	)






     @classmethod
     def 	UpperCamelCase(					self						,   lowerCamelCase	):
       _snake_case     		=				list(map(self						,   scheduler.lr_lambdas	)	)


