
"""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 	_lowercase    (							lowerCAmelCase_						, lowerCAmelCase_						, lowerCAmelCase_						, unittest.TestCase					):


     '''simple docstring'''





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



     def 			lowerCAmelCase__							(   self	)->       Union[str, Any]:
          torch.manual_seed(0	)
          UpperCAmelCase__				:			Any							=    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			,      )
          UpperCAmelCase__				:			Dict							=    PNDMScheduler(skip_prk_steps=__UpperCamelCase	)
          torch.manual_seed(0	)
          UpperCAmelCase__				:			Optional[Any]							=    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	)
          UpperCAmelCase__				:			Any							=    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			,      )
          UpperCAmelCase__				:			List[Any]							=    CLIPVisionModelWithProjection(__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[int]							=    CLIPImageProcessor(crop_size=32			,      size=32	)

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



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



     def 			lowerCAmelCase__							(   self	)->       str:
          UpperCAmelCase__				:			List[Any]							=    "cpu"  # ensure determinism for the device-dependent torch.Generator
          UpperCAmelCase__				:			int							=    self.get_dummy_components()
          UpperCAmelCase__				:			int							=    StableDiffusionImageVariationPipeline(**__UpperCamelCase	)
          UpperCAmelCase__				:			Dict							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				:			Union[str, Any]							=    self.get_dummy_inputs(__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[int]							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			str							=    image[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          UpperCAmelCase__				:			Optional[int]							=    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 			lowerCAmelCase__							(   self	)->       List[Any]:
          UpperCAmelCase__				:			List[str]							=    "cpu"  # ensure determinism for the device-dependent torch.Generator
          UpperCAmelCase__				:			List[Any]							=    self.get_dummy_components()
          UpperCAmelCase__				:			Dict							=    StableDiffusionImageVariationPipeline(**__UpperCamelCase	)
          UpperCAmelCase__				:			int							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				:			int							=    self.get_dummy_inputs(__UpperCamelCase	)
          UpperCAmelCase__				:			Dict							=    2 * [inputs["image"]]
          UpperCAmelCase__				:			Optional[int]							=    sd_pipe(**__UpperCamelCase	)

          UpperCAmelCase__				:			Union[str, Any]							=    output.images

          UpperCAmelCase__				:			int							=    image[-1, -3:, -3:, -1]

          assert image.shape == (2, 64, 64, 3)
          UpperCAmelCase__				:			Dict							=    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 			lowerCAmelCase__							(   self	)->       str:
          super().test_inference_batch_single_identical(expected_max_diff=3E-3	)







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


     '''simple docstring'''





     def 			lowerCAmelCase__							(   self	)->       Dict:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



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



     def 			lowerCAmelCase__							(   self	)->       Dict:
          UpperCAmelCase__				:			Union[str, Any]							=    StableDiffusionImageVariationPipeline.from_pretrained(
              "lambdalabs/sd-image-variations-diffusers"			,      safety_checker=__UpperCamelCase	)
          UpperCAmelCase__				:			Tuple							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				:			Any							=    self.get_inputs(__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[int]							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			Optional[int]							=    image[0, -3:, -3:, -1].flatten()

          assert image.shape == (1, 5_12, 5_12, 3)
          UpperCAmelCase__				:			List[str]							=    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 			lowerCAmelCase__							(   self	)->       Tuple:
          UpperCAmelCase__				:			Optional[Any]							=    0

          def callback_fn(__UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	) -> None:
               UpperCAmelCase__				:			Optional[Any]							=    True
               nonlocal number_of_steps
               number_of_steps += 1
               if step == 1:
                    UpperCAmelCase__				:			Dict							=    latents.detach().cpu().numpy()
                    assert latents.shape == (1, 4, 64, 64)
                    UpperCAmelCase__				:			Any							=    latents[0, -3:, -3:, -1]
                    UpperCAmelCase__				:			List[str]							=    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:
                    UpperCAmelCase__				:			int							=    latents.detach().cpu().numpy()
                    assert latents.shape == (1, 4, 64, 64)
                    UpperCAmelCase__				:			Optional[int]							=    latents[0, -3:, -3:, -1]
                    UpperCAmelCase__				:			int							=    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

          UpperCAmelCase__				:			List[Any]							=    False

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

          UpperCAmelCase__				:			Union[str, Any]							=    self.get_inputs(__UpperCamelCase			,      dtype=torch.floataa	)
          pipe(**__UpperCamelCase			,      callback=__UpperCamelCase			,      callback_steps=1	)
          assert callback_fn.has_been_called
          assert number_of_steps == inputs["num_inference_steps"]



     def 			lowerCAmelCase__							(   self	)->       Dict:
          torch.cuda.empty_cache()
          torch.cuda.reset_max_memory_allocated()
          torch.cuda.reset_peak_memory_stats()

          UpperCAmelCase__				:			List[Any]							=    "fusing/sd-image-variations-diffusers"
          UpperCAmelCase__				:			int							=    StableDiffusionImageVariationPipeline.from_pretrained(
              __UpperCamelCase			,      safety_checker=__UpperCamelCase			,      torch_dtype=torch.floataa	)
          UpperCAmelCase__				:			Optional[int]							=    pipe.to(__UpperCamelCase	)
          pipe.set_progress_bar_config(disable=__UpperCamelCase	)
          pipe.enable_attention_slicing(1	)
          pipe.enable_sequential_cpu_offload()

          UpperCAmelCase__				:			int							=    self.get_inputs(__UpperCamelCase			,      dtype=torch.floataa	)
          UpperCAmelCase__				:			List[Any]							=    pipe(**__UpperCamelCase	)

          UpperCAmelCase__				:			Optional[Any]							=    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 	_lowercase    (							unittest.TestCase					):


     '''simple docstring'''





     def 			lowerCAmelCase__							(   self	)->       Optional[Any]:
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



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



     def 			lowerCAmelCase__							(   self	)->       Optional[int]:
          UpperCAmelCase__				:			Any							=    StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"	)
          sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				:			List[Any]							=    self.get_inputs(__UpperCamelCase	)
          UpperCAmelCase__				:			str							=    sd_pipe(**__UpperCamelCase	).images[0]

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



     def 			lowerCAmelCase__							(   self	)->       int:
          UpperCAmelCase__				:			List[str]							=    StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"	)
          UpperCAmelCase__				:			Union[str, Any]							=    DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config	)
          sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				:			int							=    self.get_inputs(__UpperCamelCase	)
          UpperCAmelCase__				:			Dict							=    25
          UpperCAmelCase__				:			Any							=    sd_pipe(**__UpperCamelCase	).images[0]

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





"""simple docstring"""




from typing import TYPE_CHECKING

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


A__   :						Optional[Any] 							=     {"""configuration_vilt""": ["""VILT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViltConfig"""]}

try:
    if not is_vision_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    A__   :						List[Any] 							=     ["""ViltFeatureExtractor"""]
    A__   :						Any 							=     ["""ViltImageProcessor"""]
    A__   :						int 							=     ["""ViltProcessor"""]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    A__   :						str 							=     [
        """VILT_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """ViltForImageAndTextRetrieval""",
        """ViltForImagesAndTextClassification""",
        """ViltForTokenClassification""",
        """ViltForMaskedLM""",
        """ViltForQuestionAnswering""",
        """ViltLayer""",
        """ViltModel""",
        """ViltPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_vilt import VILT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViltConfig

    try:
        if not is_vision_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .feature_extraction_vilt import ViltFeatureExtractor
        from .image_processing_vilt import ViltImageProcessor
        from .processing_vilt import ViltProcessor

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_vilt import (
            VILT_PRETRAINED_MODEL_ARCHIVE_LIST,
            ViltForImageAndTextRetrieval,
            ViltForImagesAndTextClassification,
            ViltForMaskedLM,
            ViltForQuestionAnswering,
            ViltForTokenClassification,
            ViltLayer,
            ViltModel,
            ViltPreTrainedModel,
        )


else:
    import sys

    A__   :						Any 							=     _LazyModule(__name__, globals()["""__file__"""], _import_structure)



