
'''simple docstring'''



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

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





       """simple docstring"""




       SCREAMING_SNAKE_CASE					=     StableDiffusionImageVariationPipeline
       SCREAMING_SNAKE_CASE					=     IMAGE_VARIATION_PARAMS
       SCREAMING_SNAKE_CASE					=     IMAGE_VARIATION_BATCH_PARAMS
       SCREAMING_SNAKE_CASE					=     frozenset([])
       # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
       SCREAMING_SNAKE_CASE					=     frozenset([])






       def 				UpperCamelCase (							self					)			->	Optional[Any]:
          torch.manual_seed(0					)
          _SCREAMING_SNAKE_CASE          = UNetaDConditionModel(
              block_out_channels=(32, 64)				,							layers_per_block=2				,							sample_size=32				,							in_channels=4				,							out_channels=4				,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")				,							up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")				,							cross_attention_dim=32				,							)
          _SCREAMING_SNAKE_CASE          = PNDMScheduler(skip_prk_steps=A__					)
          torch.manual_seed(0					)
          _SCREAMING_SNAKE_CASE          = AutoencoderKL(
              block_out_channels=[32, 64]				,							in_channels=3				,							out_channels=3				,							down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]				,							up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]				,							latent_channels=4				,							)
          torch.manual_seed(0					)
          _SCREAMING_SNAKE_CASE          = CLIPVisionConfig(
              hidden_size=32				,							projection_dim=32				,							intermediate_size=37				,							layer_norm_eps=1E-05				,							num_attention_heads=4				,							num_hidden_layers=5				,							image_size=32				,							patch_size=4				,							)
          _SCREAMING_SNAKE_CASE          = CLIPVisionModelWithProjection(A__					)
          _SCREAMING_SNAKE_CASE          = CLIPImageProcessor(crop_size=32				,							size=32					)

          _SCREAMING_SNAKE_CASE          = {
              """unet""": unet,
              """scheduler""": scheduler,
              """vae""": vae,
              """image_encoder""": image_encoder,
              """feature_extractor""": feature_extractor,
              """safety_checker""": None,
          }
          return components






       def 				UpperCamelCase (							self				,							A__				,							A__=0					)			->	int:
          _SCREAMING_SNAKE_CASE          = floats_tensor((1, 3, 32, 32)				,							rng=random.Random(A__					)					)
          _SCREAMING_SNAKE_CASE          = image.cpu().permute(0				,							2				,							3				,							1					)[0]
          _SCREAMING_SNAKE_CASE          = Image.fromarray(np.uinta(A__					)					).convert("""RGB"""					).resize((32, 32)					)
          if str(A__					).startswith("""mps"""					):
             _SCREAMING_SNAKE_CASE          = torch.manual_seed(A__					)
          else:
             _SCREAMING_SNAKE_CASE          = torch.Generator(device=A__					).manual_seed(A__					)
          _SCREAMING_SNAKE_CASE          = {
              """image""": image,
              """generator""": generator,
              """num_inference_steps""": 2,
              """guidance_scale""": 6.0,
              """output_type""": """numpy""",
          }
          return inputs






       def 				UpperCamelCase (							self					)			->	List[str]:
          _SCREAMING_SNAKE_CASE          = """cpu"""  # ensure determinism for the device-dependent torch.Generator
          _SCREAMING_SNAKE_CASE          = self.get_dummy_components()
          _SCREAMING_SNAKE_CASE          = StableDiffusionImageVariationPipeline(**A__					)
          _SCREAMING_SNAKE_CASE          = sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

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

          assert image.shape == (1, 64, 64, 3)
          _SCREAMING_SNAKE_CASE          = np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]					)

          assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3






       def 				UpperCamelCase (							self					)			->	List[str]:
          _SCREAMING_SNAKE_CASE          = """cpu"""  # ensure determinism for the device-dependent torch.Generator
          _SCREAMING_SNAKE_CASE          = self.get_dummy_components()
          _SCREAMING_SNAKE_CASE          = StableDiffusionImageVariationPipeline(**A__					)
          _SCREAMING_SNAKE_CASE          = sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          _SCREAMING_SNAKE_CASE          = self.get_dummy_inputs(A__					)
          _SCREAMING_SNAKE_CASE          = 2 * [inputs["""image"""]]
          _SCREAMING_SNAKE_CASE          = sd_pipe(**A__					)

          _SCREAMING_SNAKE_CASE          = output.images

          _SCREAMING_SNAKE_CASE          = image[-1, -3:, -3:, -1]

          assert image.shape == (2, 64, 64, 3)
          _SCREAMING_SNAKE_CASE          = np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]					)

          assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3






       def 				UpperCamelCase (							self					)			->	Optional[int]:
          super().test_inference_batch_single_identical(expected_max_diff=3E-3					)






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





       """simple docstring"""




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






       def 				UpperCamelCase (							self				,							A__				,							A__="cpu"				,							A__=torch.floataa				,							A__=0					)			->	Optional[Any]:
          _SCREAMING_SNAKE_CASE          = torch.Generator(device=A__					).manual_seed(A__					)
          _SCREAMING_SNAKE_CASE          = load_image(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/input_image_vermeer.png"""					)
          _SCREAMING_SNAKE_CASE          = np.random.RandomState(A__					).standard_normal((1, 4, 64, 64)					)
          _SCREAMING_SNAKE_CASE          = torch.from_numpy(A__					).to(device=A__				,							dtype=A__					)
          _SCREAMING_SNAKE_CASE          = {
              """image""": init_image,
              """latents""": latents,
              """generator""": generator,
              """num_inference_steps""": 3,
              """guidance_scale""": 7.5,
              """output_type""": """numpy""",
          }
          return inputs






       def 				UpperCamelCase (							self					)			->	str:
          _SCREAMING_SNAKE_CASE          = StableDiffusionImageVariationPipeline.from_pretrained(
              """lambdalabs/sd-image-variations-diffusers"""				,							safety_checker=A__					)
          _SCREAMING_SNAKE_CASE          = sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          _SCREAMING_SNAKE_CASE          = self.get_inputs(A__					)
          _SCREAMING_SNAKE_CASE          = sd_pipe(**A__					).images
          _SCREAMING_SNAKE_CASE          = image[0, -3:, -3:, -1].flatten()

          assert image.shape == (1, 5_12, 5_12, 3)
          _SCREAMING_SNAKE_CASE          = np.array([0.8_4491, 0.9_0789, 0.7_5708, 0.7_8734, 0.8_3485, 0.7_0099, 0.6_6938, 0.6_8727, 0.6_1379]					)
          assert np.abs(image_slice - expected_slice					).max() < 6E-3






       def 				UpperCamelCase (							self					)			->	Dict:
          _SCREAMING_SNAKE_CASE          = 0

          def callback_fn(A__				,							A__				,							A__					) -> None:
             _SCREAMING_SNAKE_CASE          = True
             nonlocal number_of_steps
             number_of_steps += 1
             if step == 1:
                _SCREAMING_SNAKE_CASE          = latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                _SCREAMING_SNAKE_CASE          = latents[0, -3:, -3:, -1]
                _SCREAMING_SNAKE_CASE          = np.array(
                    [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]					)

                assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2
             elif step == 2:
                _SCREAMING_SNAKE_CASE          = latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                _SCREAMING_SNAKE_CASE          = latents[0, -3:, -3:, -1]
                _SCREAMING_SNAKE_CASE          = np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]					)

                assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2

          _SCREAMING_SNAKE_CASE          = False

          _SCREAMING_SNAKE_CASE          = StableDiffusionImageVariationPipeline.from_pretrained(
              """fusing/sd-image-variations-diffusers"""				,							safety_checker=A__				,							torch_dtype=torch.floataa				,							)
          pipe.to(A__					)
          pipe.set_progress_bar_config(disable=A__					)
          pipe.enable_attention_slicing()

          _SCREAMING_SNAKE_CASE          = self.get_inputs(A__				,							dtype=torch.floataa					)
          pipe(**A__				,							callback=A__				,							callback_steps=1					)
          assert callback_fn.has_been_called
          assert number_of_steps == inputs["num_inference_steps"]






       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          torch.cuda.empty_cache()
          torch.cuda.reset_max_memory_allocated()
          torch.cuda.reset_peak_memory_stats()

          _SCREAMING_SNAKE_CASE          = """fusing/sd-image-variations-diffusers"""
          _SCREAMING_SNAKE_CASE          = StableDiffusionImageVariationPipeline.from_pretrained(
              A__				,							safety_checker=A__				,							torch_dtype=torch.floataa					)
          _SCREAMING_SNAKE_CASE          = pipe.to(A__					)
          pipe.set_progress_bar_config(disable=A__					)
          pipe.enable_attention_slicing(1					)
          pipe.enable_sequential_cpu_offload()

          _SCREAMING_SNAKE_CASE          = self.get_inputs(A__				,							dtype=torch.floataa					)
          _SCREAMING_SNAKE_CASE          = pipe(**A__					)

          _SCREAMING_SNAKE_CASE          = torch.cuda.max_memory_allocated()
          # make sure that less than 2.6 GB is allocated
          assert mem_bytes < 2.6 * 10**9






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





       """simple docstring"""




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






       def 				UpperCamelCase (							self				,							A__				,							A__="cpu"				,							A__=torch.floataa				,							A__=0					)			->	Dict:
          _SCREAMING_SNAKE_CASE          = torch.Generator(device=A__					).manual_seed(A__					)
          _SCREAMING_SNAKE_CASE          = load_image(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/input_image_vermeer.png"""					)
          _SCREAMING_SNAKE_CASE          = np.random.RandomState(A__					).standard_normal((1, 4, 64, 64)					)
          _SCREAMING_SNAKE_CASE          = torch.from_numpy(A__					).to(device=A__				,							dtype=A__					)
          _SCREAMING_SNAKE_CASE          = {
              """image""": init_image,
              """latents""": latents,
              """generator""": generator,
              """num_inference_steps""": 50,
              """guidance_scale""": 7.5,
              """output_type""": """numpy""",
          }
          return inputs






       def 				UpperCamelCase (							self					)			->	str:
          _SCREAMING_SNAKE_CASE          = StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""					)
          sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          _SCREAMING_SNAKE_CASE          = self.get_inputs(A__					)
          _SCREAMING_SNAKE_CASE          = sd_pipe(**A__					).images[0]

          _SCREAMING_SNAKE_CASE          = load_numpy(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy"""					)
          _SCREAMING_SNAKE_CASE          = np.abs(expected_image - image					).max()
          assert max_diff < 1E-3






       def 				UpperCamelCase (							self					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE          = StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""					)
          _SCREAMING_SNAKE_CASE          = DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config					)
          sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          _SCREAMING_SNAKE_CASE          = self.get_inputs(A__					)
          _SCREAMING_SNAKE_CASE          = 25
          _SCREAMING_SNAKE_CASE          = sd_pipe(**A__					).images[0]

          _SCREAMING_SNAKE_CASE          = load_numpy(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"""					)
          _SCREAMING_SNAKE_CASE          = np.abs(expected_image - image					).max()
          assert max_diff < 1E-3








'''simple docstring'''



from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


UpperCamelCase__					:						Any											=     logging.get_logger(__name__)

class       _a (_lowerCamelCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE					=     'mt5'
       SCREAMING_SNAKE_CASE					=     ['past_key_values']






       def __init__(							self				,							A__=25_01_12				,							A__=5_12				,							A__=64				,							A__=10_24				,							A__=8				,							A__=None				,							A__=6				,							A__=32				,							A__=1_28				,							A__=0.1				,							A__=1E-6				,							A__=1.0				,							A__="gated-gelu"				,							A__=True				,							A__=True				,							A__="T5Tokenizer"				,							A__=False				,							A__=0				,							A__=1				,							A__=0				,							**A__				,							)			->	Union[str, Any]:
          super().__init__(
              is_encoder_decoder=A__				,							tokenizer_class=A__				,							tie_word_embeddings=A__				,							pad_token_id=A__				,							eos_token_id=A__				,							decoder_start_token_id=A__				,							**A__				,							)
          _SCREAMING_SNAKE_CASE          = vocab_size
          _SCREAMING_SNAKE_CASE          = d_model
          _SCREAMING_SNAKE_CASE          = d_kv
          _SCREAMING_SNAKE_CASE          = d_ff
          _SCREAMING_SNAKE_CASE          = num_layers
          _SCREAMING_SNAKE_CASE          = (
              num_decoder_layers if num_decoder_layers is not None else self.num_layers
          )  # default = symmetry
          _SCREAMING_SNAKE_CASE          = num_heads
          _SCREAMING_SNAKE_CASE          = relative_attention_num_buckets
          _SCREAMING_SNAKE_CASE          = relative_attention_max_distance
          _SCREAMING_SNAKE_CASE          = dropout_rate
          _SCREAMING_SNAKE_CASE          = layer_norm_epsilon
          _SCREAMING_SNAKE_CASE          = initializer_factor
          _SCREAMING_SNAKE_CASE          = feed_forward_proj
          _SCREAMING_SNAKE_CASE          = use_cache

          _SCREAMING_SNAKE_CASE          = self.feed_forward_proj.split("""-"""					)
          _SCREAMING_SNAKE_CASE          = act_info[-1]
          _SCREAMING_SNAKE_CASE          = act_info[0] == """gated"""

          if len(A__					) > 1 and act_info[0] != "gated" or len(A__					) > 2:
             raise ValueError(
                 F"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
                 """Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. """
                 """'gated-gelu' or 'relu'"""					)

          # for backwards compatibility
          if feed_forward_proj == "gated-gelu":
             _SCREAMING_SNAKE_CASE          = """gelu_new"""






       @property
       def 				UpperCamelCase (							self					)			->	int:
          return self.d_model






       @property
       def 				UpperCamelCase (							self					)			->	Optional[Any]:
          return self.num_heads






       @property
       def 				UpperCamelCase (							self					)			->	Dict:
          return self.num_layers






class       _a (_lowerCamelCase):





       """simple docstring"""




       @property
       # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
       def 				UpperCamelCase (							self					)			->	Mapping[str, Mapping[int, str]]:
          _SCREAMING_SNAKE_CASE          = {
              """input_ids""": {0: """batch""", 1: """encoder_sequence"""},
              """attention_mask""": {0: """batch""", 1: """encoder_sequence"""},
          }
          if self.use_past:
             _SCREAMING_SNAKE_CASE          = """past_encoder_sequence + sequence"""
             _SCREAMING_SNAKE_CASE          = {0: """batch"""}
             _SCREAMING_SNAKE_CASE          = {0: """batch""", 1: """past_decoder_sequence + sequence"""}
          else:
             _SCREAMING_SNAKE_CASE          = {0: """batch""", 1: """decoder_sequence"""}
             _SCREAMING_SNAKE_CASE          = {0: """batch""", 1: """decoder_sequence"""}

          if self.use_past:
             self.fill_with_past_key_values_(A__				,							direction="""inputs"""					)

          return common_inputs






       @property
       # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
       def 				UpperCamelCase (							self					)			->	int:
          return 13






       @property
       def 				UpperCamelCase (							self					)			->	float:
          return 5E-4






