
"""simple docstring"""
def    _lowerCAmelCase    (     lowerCamelCase__							:							str,						lowerCamelCase__							:							Any				)  ->    Optional[int]:
  _SCREAMING_SNAKE_CASE   :				Dict							    = [1]
  for i in range(2,						lowerCamelCase__				):
    factorials.append(factorials[-1] * i				)
  assert 0 <= k < factorials[-1] * n, "k out of bounds"

  _SCREAMING_SNAKE_CASE   :				Any							    = []
  _SCREAMING_SNAKE_CASE   :				str							    = list(range(lowerCamelCase__				)				)

  # Find permutation
  while factorials:
    _SCREAMING_SNAKE_CASE   :				Any							    = factorials.pop()
    _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Optional[Any]							    = divmod(lowerCamelCase__,						lowerCamelCase__				)
    permutation.append(elements[number]				)
    elements.remove(elements[number]				)
  permutation.append(elements[0]				)

  return permutation


if __name__ == "__main__":
      import doctest

      doctest.testmod()





"""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    _lowerCAmelCase    (     lowerCamelCase__							:							Optional[Any]				)  ->    Optional[Any]:
  _SCREAMING_SNAKE_CASE   :				Any							    = [tensor.shape for tensor in tensor_list]
  return all(shape == shapes[0] for shape in shapes[1:]				)
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
     A__      =			StableDiffusionLatentUpscalePipeline
     A__      =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
         """height""",
         """width""",
         """cross_attention_kwargs""",
         """negative_prompt_embeds""",
         """prompt_embeds""",
     }
     A__      =			PipelineTesterMixin.required_optional_params - {"""num_images_per_prompt"""}
     A__      =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     A__      =			frozenset(
         []  )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
     A__      =			frozenset([]  )

     A__      =			True
     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = 1
       _SCREAMING_SNAKE_CASE   :				int							    = 4
       _SCREAMING_SNAKE_CASE   :				str							    = (16, 16)

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = floats_tensor((batch_size, num_channels) + sizes  ,     rng=random.Random(0				)				).to(snake_case__				)
       return image
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = UNetaDConditionModel(
           act_fn="gelu"  ,     attention_head_dim=8  ,     norm_num_groups=snake_case__  ,     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=snake_case__  ,     only_cross_attention=snake_case__  ,     out_channels=5  ,     resnet_time_scale_shift="scale_shift"  ,     time_embedding_type="fourier"  ,     timestep_post_act="gelu"  ,     up_block_types=("KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KUpBlock2D")  ,     )
       _SCREAMING_SNAKE_CASE   :				Any							    = 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  ,     )
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = EulerDiscreteScheduler(prediction_type="sample"				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = 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=1000  ,     hidden_act="quick_gelu"  ,     projection_dim=512  ,     )
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = CLIPTextModel(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				str							    = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)

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

       return components
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__=0				):




       """simple docstring"""


       if str(snake_case__				).startswith("mps"				):
         _SCREAMING_SNAKE_CASE   :				Optional[int]							    = torch.manual_seed(snake_case__				)
       else:
         _SCREAMING_SNAKE_CASE   :				Dict							    = torch.Generator(device=snake_case__				).manual_seed(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				int							    = {
           "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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = "cpu"

       _SCREAMING_SNAKE_CASE   :				Any							    = self.get_dummy_components()
       _SCREAMING_SNAKE_CASE   :				str							    = self.pipeline_class(**snake_case__				)
       pipe.to(snake_case__				)
       pipe.set_progress_bar_config(disable=snake_case__				)

       _SCREAMING_SNAKE_CASE   :				str							    = self.get_dummy_inputs(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = pipe(**snake_case__				).images
       _SCREAMING_SNAKE_CASE   :				str							    = image[0, -3:, -3:, -1]

       self.assertEqual(image.shape  ,     (1, 256, 256, 3)				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = np.array(
           [0.47_222_412, 0.41_921_633, 0.44_717_434, 0.46_874_192, 0.42_588_258, 0.46_150_726, 0.4_677_534, 0.45_583_832, 0.48_579_055]				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = np.abs(image_slice.flatten() - expected_slice				).max()
       self.assertLessEqual(snake_case__  ,     1E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().test_attention_slicing_forward_pass(expected_max_diff=7E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().test_cpu_offload_forward_pass(expected_max_diff=3E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().test_inference_batch_single_identical(expected_max_diff=7E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().test_save_load_local(expected_max_difference=3E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().test_save_load_optional_components(expected_max_difference=3E-3				)




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = [
           "DDIMScheduler",
           "DDPMScheduler",
           "PNDMScheduler",
           "HeunDiscreteScheduler",
           "EulerAncestralDiscreteScheduler",
           "KDPM2DiscreteScheduler",
           "KDPM2AncestralDiscreteScheduler",
           "DPMSolverSDEScheduler",
       ]
       _SCREAMING_SNAKE_CASE   :				Any							    = self.get_dummy_components()
       _SCREAMING_SNAKE_CASE   :				Dict							    = self.pipeline_class(**snake_case__				)

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

       pipe.to(snake_case__				)
       pipe.set_progress_bar_config(disable=snake_case__				)
       _SCREAMING_SNAKE_CASE   :				str							    = self.get_dummy_inputs(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = 2

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

         _SCREAMING_SNAKE_CASE   :				List[str]							    = getattr(snake_case__  ,     scheduler_enum.name				)
         _SCREAMING_SNAKE_CASE   :				Dict							    = scheduler_cls.from_config(pipe.scheduler.config				)
         _SCREAMING_SNAKE_CASE   :				int							    = pipe(**snake_case__				)[0]
         outputs.append(snake_case__				)

       assert check_same_shape(snake_case__				)
@require_torch_gpu
@slow
class 			UpperCamelCase      (							unittest.TestCase  ):
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.manual_seed(33				)

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

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

       _SCREAMING_SNAKE_CASE   :				str							    = "a photo of an astronaut high resolution, unreal engine, ultra realistic"

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = pipe(snake_case__  ,     generator=snake_case__  ,     output_type="latent"				).images

       _SCREAMING_SNAKE_CASE   :				Tuple							    = upscaler(
           prompt=snake_case__  ,     image=snake_case__  ,     num_inference_steps=20  ,     guidance_scale=0  ,     generator=snake_case__  ,     output_type="np"  ,     ).images[0]

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = 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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = torch.manual_seed(33				)

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

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = "the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas"

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

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = upscaler(
           prompt=snake_case__  ,     image=snake_case__  ,     num_inference_steps=20  ,     guidance_scale=0  ,     generator=snake_case__  ,     output_type="np"  ,     ).images[0]

       _SCREAMING_SNAKE_CASE   :				int							    = 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



