


import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
							import torch

if is_vision_available():
							from PIL import Image

							from transformers import CLIPImageProcessor



class    __lowercase  (       unittest.TestCase							):

							"""simple docstring"""


							def __init__(    self				, A_				, A_=7				, A_=3				, A_=18				, A_=30				, A_=400				, A_=True				, A_=None				, A_=True				, A_=None				, A_=True				, A_=[0.4814_5466, 0.457_8275, 0.4082_1073]				, A_=[0.2686_2954, 0.2613_0258, 0.2757_7711]				, A_=True				, )->	int:
											_SCREAMING_SNAKE_CASE      = size if size is not None else {'shortest_edge': 20}
											_SCREAMING_SNAKE_CASE      = crop_size if crop_size is not None else {'height': 18, 'width': 18}
											_SCREAMING_SNAKE_CASE      = parent
											_SCREAMING_SNAKE_CASE      = batch_size
											_SCREAMING_SNAKE_CASE      = num_channels
											_SCREAMING_SNAKE_CASE      = image_size
											_SCREAMING_SNAKE_CASE      = min_resolution
											_SCREAMING_SNAKE_CASE      = max_resolution
											_SCREAMING_SNAKE_CASE      = do_resize
											_SCREAMING_SNAKE_CASE      = size
											_SCREAMING_SNAKE_CASE      = do_center_crop
											_SCREAMING_SNAKE_CASE      = crop_size
											_SCREAMING_SNAKE_CASE      = do_normalize
											_SCREAMING_SNAKE_CASE      = image_mean
											_SCREAMING_SNAKE_CASE      = image_std
											_SCREAMING_SNAKE_CASE      = do_convert_rgb





							def 		__magic_name__			(    self	)->	Optional[Any]:
											return {
											    "do_resize": self.do_resize,
											    "size": self.size,
											    "do_center_crop": self.do_center_crop,
											    "crop_size": self.crop_size,
											    "do_normalize": self.do_normalize,
											    "image_mean": self.image_mean,
											    "image_std": self.image_std,
											    "do_convert_rgb": self.do_convert_rgb,
											}





							def 		__magic_name__			(    self				, A_=False				, A_=False				, A_=False	)->	Dict:
											assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

											if equal_resolution:
															_SCREAMING_SNAKE_CASE      = []
															for i in range(self.batch_size	):
																			image_inputs.append(
																			    np.random.randint(
																			        255				, size=(self.num_channels, self.max_resolution, self.max_resolution)				, dtype=np.uinta	)	)
											else:
															_SCREAMING_SNAKE_CASE      = []
															for i in range(self.batch_size	):
																			_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = np.random.choice(np.arange(self.min_resolution				, self.max_resolution	)				, 2	)
																			image_inputs.append(np.random.randint(255				, size=(self.num_channels, width, height)				, dtype=np.uinta	)	)

											if not numpify and not torchify:
															# PIL expects the channel dimension as last dimension
															_SCREAMING_SNAKE_CASE      = [Image.fromarray(np.moveaxis(A_				, 0				, -1	)	) for x in image_inputs]

											if torchify:
															_SCREAMING_SNAKE_CASE      = [torch.from_numpy(A_	) for x in image_inputs]

											return image_inputs







@require_torch
@require_vision
class    __lowercase  (       lowercase_      ,     unittest.TestCase							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					Optional[int]        =						CLIPImageProcessor if is_vision_available() else None





							def 		__magic_name__			(    self	)->	Tuple:
											_SCREAMING_SNAKE_CASE      = CLIPImageProcessingTester(self	)





							@property
							def 		__magic_name__			(    self	)->	Dict:
											return self.image_processor_tester.prepare_image_processor_dict()





							def 		__magic_name__			(    self	)->	int:
											_SCREAMING_SNAKE_CASE      = self.image_processing_class(**self.image_processor_dict	)
											self.assertTrue(hasattr(A_				, 'do_resize'	)	)
											self.assertTrue(hasattr(A_				, 'size'	)	)
											self.assertTrue(hasattr(A_				, 'do_center_crop'	)	)
											self.assertTrue(hasattr(A_				, 'center_crop'	)	)
											self.assertTrue(hasattr(A_				, 'do_normalize'	)	)
											self.assertTrue(hasattr(A_				, 'image_mean'	)	)
											self.assertTrue(hasattr(A_				, 'image_std'	)	)
											self.assertTrue(hasattr(A_				, 'do_convert_rgb'	)	)





							def 		__magic_name__			(    self	)->	Union[str, Any]:
											_SCREAMING_SNAKE_CASE      = self.image_processing_class.from_dict(self.image_processor_dict	)
											self.assertEqual(image_processor.size				, {'shortest_edge': 20}	)
											self.assertEqual(image_processor.crop_size				, {'height': 18, 'width': 18}	)

											_SCREAMING_SNAKE_CASE      = self.image_processing_class.from_dict(self.image_processor_dict				, size=42				, crop_size=84	)
											self.assertEqual(image_processor.size				, {'shortest_edge': 42}	)
											self.assertEqual(image_processor.crop_size				, {'height': 84, 'width': 84}	)





							def 		__magic_name__			(    self	)->	Tuple:
											pass





							def 		__magic_name__			(    self	)->	Any:
											_SCREAMING_SNAKE_CASE      = self.image_processing_class(**self.image_processor_dict	)
											# create random PIL images
											_SCREAMING_SNAKE_CASE      = self.image_processor_tester.prepare_inputs(equal_resolution=A_	)
											for image in image_inputs:
															self.assertIsInstance(A_				, Image.Image	)

											# Test not batched input
											_SCREAMING_SNAKE_CASE      = image_processing(image_inputs[0]				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        1,
											        self.image_processor_tester.num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )

											# Test batched
											_SCREAMING_SNAKE_CASE      = image_processing(A_				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        self.image_processor_tester.batch_size,
											        self.image_processor_tester.num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )





							def 		__magic_name__			(    self	)->	Optional[Any]:
											_SCREAMING_SNAKE_CASE      = self.image_processing_class(**self.image_processor_dict	)
											# create random numpy tensors
											_SCREAMING_SNAKE_CASE      = self.image_processor_tester.prepare_inputs(equal_resolution=A_				, numpify=A_	)
											for image in image_inputs:
															self.assertIsInstance(A_				, np.ndarray	)

											# Test not batched input
											_SCREAMING_SNAKE_CASE      = image_processing(image_inputs[0]				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        1,
											        self.image_processor_tester.num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )

											# Test batched
											_SCREAMING_SNAKE_CASE      = image_processing(A_				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        self.image_processor_tester.batch_size,
											        self.image_processor_tester.num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )





							def 		__magic_name__			(    self	)->	Optional[Any]:
											_SCREAMING_SNAKE_CASE      = self.image_processing_class(**self.image_processor_dict	)
											# create random PyTorch tensors
											_SCREAMING_SNAKE_CASE      = self.image_processor_tester.prepare_inputs(equal_resolution=A_				, torchify=A_	)
											for image in image_inputs:
															self.assertIsInstance(A_				, torch.Tensor	)

											# Test not batched input
											_SCREAMING_SNAKE_CASE      = image_processing(image_inputs[0]				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        1,
											        self.image_processor_tester.num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )

											# Test batched
											_SCREAMING_SNAKE_CASE      = image_processing(A_				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        self.image_processor_tester.batch_size,
											        self.image_processor_tester.num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )







@require_torch
@require_vision
class    __lowercase  (       lowercase_      ,     unittest.TestCase							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					str        =						CLIPImageProcessor if is_vision_available() else None





							def 		__magic_name__			(    self	)->	Tuple:
											_SCREAMING_SNAKE_CASE      = CLIPImageProcessingTester(self				, num_channels=4	)
											_SCREAMING_SNAKE_CASE      = 3





							@property
							def 		__magic_name__			(    self	)->	str:
											return self.image_processor_tester.prepare_image_processor_dict()





							def 		__magic_name__			(    self	)->	Optional[Any]:
											_SCREAMING_SNAKE_CASE      = self.image_processing_class(**self.image_processor_dict	)
											self.assertTrue(hasattr(A_				, 'do_resize'	)	)
											self.assertTrue(hasattr(A_				, 'size'	)	)
											self.assertTrue(hasattr(A_				, 'do_center_crop'	)	)
											self.assertTrue(hasattr(A_				, 'center_crop'	)	)
											self.assertTrue(hasattr(A_				, 'do_normalize'	)	)
											self.assertTrue(hasattr(A_				, 'image_mean'	)	)
											self.assertTrue(hasattr(A_				, 'image_std'	)	)
											self.assertTrue(hasattr(A_				, 'do_convert_rgb'	)	)





							def 		__magic_name__			(    self	)->	Dict:
											pass





							def 		__magic_name__			(    self	)->	Any:
											_SCREAMING_SNAKE_CASE      = self.image_processing_class(**self.image_processor_dict	)
											# create random PIL images
											_SCREAMING_SNAKE_CASE      = self.image_processor_tester.prepare_inputs(equal_resolution=A_	)
											for image in image_inputs:
															self.assertIsInstance(A_				, Image.Image	)

											# Test not batched input
											_SCREAMING_SNAKE_CASE      = image_processing(image_inputs[0]				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        1,
											        self.expected_encoded_image_num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )

											# Test batched
											_SCREAMING_SNAKE_CASE      = image_processing(A_				, return_tensors='pt'	).pixel_values
											self.assertEqual(
											    encoded_images.shape				, (
											        self.image_processor_tester.batch_size,
											        self.expected_encoded_image_num_channels,
											        self.image_processor_tester.crop_size['height'],
											        self.image_processor_tester.crop_size['width'],
											    )				, )




import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()



class    __lowercase  (       unittest.TestCase							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					str        =						"hf-internal-testing/tiny-stable-diffusion-pipe"
							SCREAMING_SNAKE_CASE   :					str        =						1
							SCREAMING_SNAKE_CASE   :					Optional[Any]        =						77
							SCREAMING_SNAKE_CASE   :					Dict        =						32
							SCREAMING_SNAKE_CASE   :					Optional[Any]        =						4
							SCREAMING_SNAKE_CASE   :					Optional[int]        =						64
							SCREAMING_SNAKE_CASE   :					Tuple        =						torch.manual_seed(0							)





							def 		__magic_name__			(    self				, A_=0.9999	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      = UNetaDConditionModel.from_pretrained(self.model_id				, subfolder='unet'	)
											_SCREAMING_SNAKE_CASE      = unet.to(A_	)
											_SCREAMING_SNAKE_CASE      = EMAModel(unet.parameters()				, decay=A_				, model_cls=A_				, model_config=unet.config	)
											return unet, ema_unet





							def 		__magic_name__			(    self	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      = torch.randn(
											    self.batch_size				, self.num_in_channels				, self.latent_height				, self.latent_width				, generator=self.generator	).to(A_	)
											_SCREAMING_SNAKE_CASE      = torch.randint(0				, 1000				, size=(self.batch_size,)				, generator=self.generator	).to(A_	)
											_SCREAMING_SNAKE_CASE      = torch.randn(
											    self.batch_size				, self.prompt_length				, self.text_encoder_hidden_dim				, generator=self.generator	).to(A_	)
											return noisy_latents, timesteps, encoder_hidden_states





							def 		__magic_name__			(    self				, A_	)->	str:
											_SCREAMING_SNAKE_CASE      = {}
											for k, param in unet.state_dict().items():
															_SCREAMING_SNAKE_CASE      = torch.randn_like(A_	) + (param * torch.randn_like(A_	))
															updated_state_dict.update({k: updated_param}	)
											unet.load_state_dict(A_	)
											return unet





							def 		__magic_name__			(    self	)->	Dict:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models()
											# Take the first (hypothetical) EMA step.
											ema_unet.step(unet.parameters()	)
											assert ema_unet.optimization_step == 1

											# Take two more.
											for _ in range(2	):
															ema_unet.step(unet.parameters()	)
											assert ema_unet.optimization_step == 3





							def 		__magic_name__			(    self	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models()
											# Since the `unet` is not being updated (i.e., backprop'd)
											# there won't be any difference between the `params` of `unet`
											# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
											ema_unet.step(unet.parameters()	)
											_SCREAMING_SNAKE_CASE      = list(unet.parameters()	)
											for s_param, param in zip(ema_unet.shadow_params				, A_	):
															assert torch.allclose(A_				, A_	)

											# The above holds true even if we call `ema.step()` multiple times since
											# `unet` params are still not being updated.
											for _ in range(4	):
															ema_unet.step(unet.parameters()	)
											for s_param, param in zip(ema_unet.shadow_params				, A_	):
															assert torch.allclose(A_				, A_	)





							def 		__magic_name__			(    self	)->	Union[str, Any]:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models()
											# Here we simulate the parameter updates for `unet`. Since there might
											# be some parameters which are initialized to zero we take extra care to
											# initialize their values to something non-zero before the multiplication.
											_SCREAMING_SNAKE_CASE      = self.simulate_backprop(A_	)

											# Take the EMA step.
											ema_unet.step(unet_pseudo_updated_step_one.parameters()	)

											# Now the EMA'd parameters won't be equal to the original model parameters.
											_SCREAMING_SNAKE_CASE      = list(unet_pseudo_updated_step_one.parameters()	)
											for s_param, param in zip(ema_unet.shadow_params				, A_	):
															assert ~torch.allclose(A_				, A_	)

											# Ensure this is the case when we take multiple EMA steps.
											for _ in range(4	):
															ema_unet.step(unet.parameters()	)
											for s_param, param in zip(ema_unet.shadow_params				, A_	):
															assert ~torch.allclose(A_				, A_	)





							def 		__magic_name__			(    self	)->	Optional[int]:
											# If we call EMA step after a backpropagation consecutively for two times,
											# the shadow params from those two steps should be different.
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models()

											# First backprop + EMA
											_SCREAMING_SNAKE_CASE      = self.simulate_backprop(A_	)
											ema_unet.step(unet_step_one.parameters()	)
											_SCREAMING_SNAKE_CASE      = ema_unet.shadow_params

											# Second backprop + EMA
											_SCREAMING_SNAKE_CASE      = self.simulate_backprop(A_	)
											ema_unet.step(unet_step_two.parameters()	)
											_SCREAMING_SNAKE_CASE      = ema_unet.shadow_params

											for step_one, step_two in zip(A_				, A_	):
															assert ~torch.allclose(A_				, A_	)





							def 		__magic_name__			(    self	)->	Dict:
											# If there's no decay even if there are backprops, EMA steps
											# won't take any effect i.e., the shadow params would remain the
											# same.
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models(decay=0.0	)
											_SCREAMING_SNAKE_CASE      = self.simulate_backprop(A_	)
											ema_unet.step(unet_step_one.parameters()	)
											_SCREAMING_SNAKE_CASE      = ema_unet.shadow_params

											_SCREAMING_SNAKE_CASE      = self.simulate_backprop(A_	)
											ema_unet.step(unet_step_two.parameters()	)
											_SCREAMING_SNAKE_CASE      = ema_unet.shadow_params

											for step_one, step_two in zip(A_				, A_	):
															assert torch.allclose(A_				, A_	)





							@skip_mps
							def 		__magic_name__			(    self	)->	Optional[Any]:
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_models()
											_SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.get_dummy_inputs()

											with tempfile.TemporaryDirectory() as tmpdir:
															ema_unet.save_pretrained(A_	)
															_SCREAMING_SNAKE_CASE      = UNetaDConditionModel.from_pretrained(A_				, model_cls=A_	)
															_SCREAMING_SNAKE_CASE      = loaded_unet.to(unet.device	)

											# Since no EMA step has been performed the outputs should match.
											_SCREAMING_SNAKE_CASE      = unet(A_				, A_				, A_	).sample
											_SCREAMING_SNAKE_CASE      = loaded_unet(A_				, A_				, A_	).sample

											assert torch.allclose(A_				, A_				, atol=1e-4	)
