



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       _snake_case				(     UpperCAmelCase__ ,							UpperCAmelCase__ ,							UpperCAmelCase__ ,							unittest.TestCase ):






   '''simple docstring'''






   A__      : Optional[int]       =						StableDiffusionImageVariationPipeline
   A__      : Dict       =						IMAGE_VARIATION_PARAMS
   A__      : str       =						IMAGE_VARIATION_BATCH_PARAMS
   A__      : Any       =						frozenset([] )
   # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
   A__      : List[Any]       =						frozenset([] )


   def       A__    (    self:     int							)	->      List[str]:
       torch.manual_seed(0							)
       UpperCAmelCase_     :       List[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=snake_case__							)
       torch.manual_seed(0							)
       UpperCAmelCase_     :       Union[str, 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_     :       List[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(snake_case__							)
       UpperCAmelCase_     :       Dict  		=					CLIPImageProcessor(crop_size=32  ,size=32							)

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


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


   def       A__    (    self:     Optional[int]							)	->      int:
       UpperCAmelCase_     :       Union[str, Any]  		=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase_     :       Tuple  		=					self.get_dummy_components()
       UpperCAmelCase_     :       Optional[int]  		=					StableDiffusionImageVariationPipeline(**snake_case__							)
       UpperCAmelCase_     :       Tuple  		=					sd_pipe.to(snake_case__							)
       sd_pipe.set_progress_bar_config(disable=snake_case__							)

       UpperCAmelCase_     :       List[str]  		=					self.get_dummy_inputs(snake_case__							)
       UpperCAmelCase_     :       Optional[Any]  		=					sd_pipe(**snake_case__							).images
       UpperCAmelCase_     :       Optional[int]  		=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       UpperCAmelCase_     :       Union[str, Any]  		=					np.array([0.5_2_3_9, 0.5_7_2_3, 0.4_7_9_6, 0.5_0_4_9, 0.5_5_5_0, 0.4_6_8_5, 0.5_3_2_9, 0.4_8_9_1, 0.4_9_2_1]							)

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


   def       A__    (    self:     Dict							)	->      Optional[Any]:
       UpperCAmelCase_     :       Optional[Any]  		=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
       UpperCAmelCase_     :       Tuple  		=					self.get_dummy_components()
       UpperCAmelCase_     :       Optional[int]  		=					StableDiffusionImageVariationPipeline(**snake_case__							)
       UpperCAmelCase_     :       Tuple  		=					sd_pipe.to(snake_case__							)
       sd_pipe.set_progress_bar_config(disable=snake_case__							)

       UpperCAmelCase_     :       Optional[int]  		=					self.get_dummy_inputs(snake_case__							)
       UpperCAmelCase_     :       Any  		=					2 * [inputs["""image"""]]
       UpperCAmelCase_     :       Any  		=					sd_pipe(**snake_case__							)

       UpperCAmelCase_     :       List[str]  		=					output.images

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

       assert image.shape == (2, 64, 64, 3)
       UpperCAmelCase_     :       Tuple  		=					np.array([0.6_8_9_2, 0.5_6_3_7, 0.5_8_3_6, 0.5_7_7_1, 0.6_2_5_4, 0.6_4_0_9, 0.5_5_8_0, 0.5_5_6_9, 0.5_2_8_9]							)

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






   def       A__    (    self:     int							)	->      str:
       super().test_inference_batch_single_identical(expected_max_diff=3e-3							)







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






   '''simple docstring'''






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


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


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

       UpperCAmelCase_     :       str  		=					self.get_inputs(snake_case__							)
       UpperCAmelCase_     :       str  		=					sd_pipe(**snake_case__							).images
       UpperCAmelCase_     :       List[Any]  		=					image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 512, 3)
       UpperCAmelCase_     :       Optional[Any]  		=					np.array([0.8_4_4_9_1, 0.9_0_7_8_9, 0.7_5_7_0_8, 0.7_8_7_3_4, 0.8_3_4_8_5, 0.7_0_0_9_9, 0.6_6_9_3_8, 0.6_8_7_2_7, 0.6_1_3_7_9]							)
       assert np.abs(image_slice - expected_slice							).max() < 6e-3


   def       A__    (    self:     str							)	->      str:
       UpperCAmelCase_     :       Any  		=					0

       def callback_fn(lowerCamelCase_:     int  ,lowerCamelCase_:     int  ,lowerCamelCase_:     torch.FloatTensor							) -> None:
           UpperCAmelCase_     :       Union[str, Any]  		=					True
           nonlocal number_of_steps
           number_of_steps += 1
           if step == 1:
               UpperCAmelCase_     :       int  		=					latents.detach().cpu().numpy()
               assert latents.shape == (1, 4, 64, 64)
               UpperCAmelCase_     :       Any  		=					latents[0, -3:, -3:, -1]
               UpperCAmelCase_     :       Tuple  		=					np.array(
                   [-0.1_6_2_1, 0.2_8_3_7, -0.7_9_7_9, -0.1_2_2_1, -1.3_0_5_7, 0.7_6_8_1, -2.1_1_9_1, 0.0_4_6_4, 1.6_3_0_9]							)

               assert np.abs(latents_slice.flatten() - expected_slice							).max() < 5e-2
           elif step == 2:
               UpperCAmelCase_     :       Dict  		=					latents.detach().cpu().numpy()
               assert latents.shape == (1, 4, 64, 64)
               UpperCAmelCase_     :       List[Any]  		=					latents[0, -3:, -3:, -1]
               UpperCAmelCase_     :       Tuple  		=					np.array([0.6_2_9_9, 1.7_5_0_0, 1.1_9_9_2, -2.1_5_8_2, -1.8_9_9_4, 0.7_3_3_4, -0.7_0_9_0, 1.0_1_3_7, 1.5_2_7_3]							)

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

       UpperCAmelCase_     :       List[Any]  		=					False

       UpperCAmelCase_     :       Union[str, Any]  		=					StableDiffusionImageVariationPipeline.from_pretrained(
           """fusing/sd-image-variations-diffusers"""  ,safety_checker=snake_case__  ,torch_dtype=torch.floataa  ,)
       pipe.to(snake_case__							)
       pipe.set_progress_bar_config(disable=snake_case__							)
       pipe.enable_attention_slicing()

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






   def       A__    (    self:     Any							)	->      List[Any]:
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

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

       UpperCAmelCase_     :       Tuple  		=					self.get_inputs(snake_case__  ,dtype=torch.floataa							)
       UpperCAmelCase_     :       Any  		=					pipe(**snake_case__							)

       UpperCAmelCase_     :       Tuple  		=					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       _snake_case				(     unittest.TestCase ):






   '''simple docstring'''






   def       A__    (    self:     List[str]							)	->      Tuple:
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()


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


   def       A__    (    self:     Dict							)	->      int:
       UpperCAmelCase_     :       Optional[int]  		=					StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""							)
       sd_pipe.to(snake_case__							)
       sd_pipe.set_progress_bar_config(disable=snake_case__							)

       UpperCAmelCase_     :       List[Any]  		=					self.get_inputs(snake_case__							)
       UpperCAmelCase_     :       Dict  		=					sd_pipe(**snake_case__							).images[0]

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






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

       UpperCAmelCase_     :       List[Any]  		=					self.get_inputs(snake_case__							)
       UpperCAmelCase_     :       Tuple  		=					25
       UpperCAmelCase_     :       Optional[int]  		=					sd_pipe(**snake_case__							).images[0]

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






import enum
import inspect
import warnings
from functools import wraps
from typing import Callable, Optional

from .logging import get_logger


lowercase_      :							Optional[int]    							=					set()
lowercase_      :							List[Any]    							=					get_logger(__name__)


def 				__SCREAMING_SNAKE_CASE    (     snake_case_ = None		):

						'''simple docstring'''


						def decorator(snake_case_		):
												global _emitted_deprecation_warnings

												if inspect.isclass(snake_case_		):
																		_UpperCAmelCase						=	deprecated_class_or_function.__init__
																		_UpperCAmelCase						=	deprecated_class_or_function.__name__
												else:
																		_UpperCAmelCase						=	deprecated_class_or_function
																		_UpperCAmelCase						=	deprecated_function.__name__
																		# Support deprecating __init__ class method: class name instead
																		_UpperCAmelCase						=	name if name != "__init__" else deprecated_function.__qualname__.split("."		)[-2]

												_UpperCAmelCase						=	(
												    f"""{name} is deprecated and will be removed in the next major version of datasets.""" + f""" {help_message}"""
												    if help_message
												    else ""
												)

												@wraps(snake_case_		)
												def wrapper(*snake_case_      ,       **snake_case_		):
																		_UpperCAmelCase						=	hash(snake_case_		)
																		if func_hash not in _emitted_deprecation_warnings:
																								warnings.warn(snake_case_      ,       category=snake_case_      ,       stacklevel=2		)
																								_emitted_deprecation_warnings.add(snake_case_		)
																		return deprecated_function(*snake_case_      ,       **snake_case_		)

												_UpperCAmelCase						=	"deprecated"

												if inspect.isclass(snake_case_		):
																		_UpperCAmelCase						=	wrapper
																		return deprecated_class_or_function
												else:
																		return wrapper

						return decorator






class        __lowerCAmelCase					(			enum.EnumMeta					):







			def __getattribute__(			cls       :     str       ,       snake_case__       :     Optional[Any] ):




									"""simple docstring"""


									_UpperCAmelCase						=	super().__getattribute__(snake_case__ )
									if isinstance(snake_case__       ,       enum.Enum ) and obj._on_access:
															obj._on_access()
									return obj







			def __getitem__(			cls       :     str       ,       snake_case__       :     Union[str, Any] ):




									"""simple docstring"""


									_UpperCAmelCase						=	super().__getitem__(snake_case__ )
									if member._on_access:
															member._on_access()
									return member







			def __call__(			cls       :     Dict       ,       snake_case__       :     Any       ,       snake_case__       :     List[str]=None       ,       *, snake_case__       :     Optional[int]=None       ,       snake_case__       :     List[Any]=None       ,       snake_case__       :     Tuple=None       ,       snake_case__       :     List[Any]=1 ):




									"""simple docstring"""


									_UpperCAmelCase						=	super().__call__(snake_case__       ,       snake_case__       ,       module=snake_case__       ,       qualname=snake_case__       ,       type=snake_case__       ,       start=snake_case__ )
									if isinstance(snake_case__       ,       enum.Enum ) and obj._on_access:
															obj._on_access()
									return obj






class        __lowerCAmelCase					(			enum.Enum   ,						metaclass=UpperCAmelCase__					):







			def __new__(			cls       :     Any       ,       snake_case__       :     Dict ):




									"""simple docstring"""


									_UpperCAmelCase						=	object.__new__(cls )
									_UpperCAmelCase						=	value
									_UpperCAmelCase						=	member.deprecate
									return member







			@property
			def UpperCamelCase     (			self       :     str ):




									"""simple docstring"""


									return ""







			def UpperCamelCase     (			self       :     Optional[int] ):




									"""simple docstring"""


									_UpperCAmelCase						=	F""" {self.help_message}""" if self.help_message else ""
									warnings.warn(
									    F"""'{self.__objclass__.__name__}' is deprecated and will be removed in the next major version of datasets."""
									    + help_message       ,       snake_case__       ,       stacklevel=3       ,       )




