'''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   A_      (       SCREAMING_SNAKE_CASE_						)			->Union[str, Any]:
  lowercase_     =       [tensor.shape for tensor in tensor_list]
  return all(shape == shapes[0] for shape in shapes[1:]						)





class       _a    (		__a ,				__a ,				__a ,				unittest.TestCase      ):


  """simple docstring"""

  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 		lowerCamelCase__  (		self      :       Optional[Any]      ):







    '''simple docstring'''


    lowercase_     =       1
    lowercase_     =       4
    lowercase_     =       (16, 16)

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




  def 		lowerCamelCase__  (		self      :       int      ):







    '''simple docstring'''


    torch.manual_seed(0      )
    lowercase_     =       UNetaDConditionModel(
        act_fn="""gelu"""   ,				attention_head_dim=8   ,				norm_num_groups=lowercase_   ,				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=lowercase_   ,				only_cross_attention=lowercase_   ,				out_channels=5   ,				resnet_time_scale_shift="""scale_shift"""   ,				time_embedding_type="""fourier"""   ,				timestep_post_act="""gelu"""   ,				up_block_types=("""KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KUpBlock2D""")   ,				)
    lowercase_     =       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   ,				)
    lowercase_     =       EulerDiscreteScheduler(prediction_type="""sample"""      )
    lowercase_     =       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   ,				)
    lowercase_     =       CLIPTextModel(lowercase_      )
    lowercase_     =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

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

    return components




  def 		lowerCamelCase__  (		self      :       Optional[Any]   ,				lowercase_      :       int   ,				lowercase_      :       Optional[int]=0      ):







    '''simple docstring'''


    if str(lowercase_      ).startswith("""mps"""      ):
      lowercase_     =       torch.manual_seed(lowercase_      )
    else:
      lowercase_     =       torch.Generator(device=lowercase_      ).manual_seed(lowercase_      )
    lowercase_     =       {
        """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 		lowerCamelCase__  (		self      :       int      ):







    '''simple docstring'''


    lowercase_     =       """cpu"""

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

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

    self.assertEqual(image.shape   ,				(1, 256, 256, 3)      )
    lowercase_     =       np.array(
        [0.4_7_2_2_2_4_1_2, 0.4_1_9_2_1_6_3_3, 0.4_4_7_1_7_4_3_4, 0.4_6_8_7_4_1_9_2, 0.4_2_5_8_8_2_5_8, 0.4_6_1_5_0_7_2_6, 0.4_6_7_7_5_3_4, 0.4_5_5_8_3_8_3_2, 0.4_8_5_7_9_0_5_5]      )
    lowercase_     =       np.abs(image_slice.flatten() - expected_slice      ).max()
    self.assertLessEqual(lowercase_   ,				1e-3      )




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       List[Any]      ):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       int      ):







    '''simple docstring'''


    super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3      )




  def 		lowerCamelCase__  (		self      :       Tuple      ):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       Dict      ):







    '''simple docstring'''


    super().test_save_load_local(expected_max_difference=3e-3      )




  def 		lowerCamelCase__  (		self      :       Tuple      ):







    '''simple docstring'''


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





  def 		lowerCamelCase__  (		self      :       List[str]      ):







    '''simple docstring'''


    lowercase_     =       [
        """DDIMScheduler""",
        """DDPMScheduler""",
        """PNDMScheduler""",
        """HeunDiscreteScheduler""",
        """EulerAncestralDiscreteScheduler""",
        """KDPM2DiscreteScheduler""",
        """KDPM2AncestralDiscreteScheduler""",
        """DPMSolverSDEScheduler""",
    ]
    lowercase_     =       self.get_dummy_components()
    lowercase_     =       self.pipeline_class(**lowercase_      )

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

    pipe.to(lowercase_      )
    pipe.set_progress_bar_config(disable=lowercase_      )
    lowercase_     =       self.get_dummy_inputs(lowercase_      )
    lowercase_     =       2

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

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

    assert check_same_shape(lowercase_      )






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


  """simple docstring"""

  def 		lowerCamelCase__  (		self      :       Optional[int]      ):







    '''simple docstring'''


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




  def 		lowerCamelCase__  (		self      :       List[Any]      ):







    '''simple docstring'''


    lowercase_     =       torch.manual_seed(33      )

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

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

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

    lowercase_     =       pipe(lowercase_   ,				generator=lowercase_   ,				output_type="""latent"""      ).images

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

    lowercase_     =       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 		lowerCamelCase__  (		self      :       List[Any]      ):







    '''simple docstring'''


    lowercase_     =       torch.manual_seed(33      )

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

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

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

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

    lowercase_     =       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'''


from typing import TYPE_CHECKING

from ...file_utils import _LazyModule, is_torch_available
from ...utils import OptionalDependencyNotAvailable


__snake_case						  =		{
    """configuration_gpt_neox_japanese""": ["""GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP""", """GPTNeoXJapaneseConfig"""],
    """tokenization_gpt_neox_japanese""": ["""GPTNeoXJapaneseTokenizer"""],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __snake_case						  =		[
           """GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """GPTNeoXJapaneseForCausalLM""",
           """GPTNeoXJapaneseLayer""",
           """GPTNeoXJapaneseModel""",
           """GPTNeoXJapanesePreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_gpt_neox_japanese import GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXJapaneseConfig
       from .tokenization_gpt_neox_japanese import GPTNeoXJapaneseTokenizer

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_gpt_neox_japanese import (
                  GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST,
                  GPTNeoXJapaneseForCausalLM,
                  GPTNeoXJapaneseLayer,
                  GPTNeoXJapaneseModel,
                  GPTNeoXJapanesePreTrainedModel,
              )


else:
       import sys

       __snake_case						  =		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

