






import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()
class 		SCREAMING_SNAKE_CASE__	(  UpperCAmelCase					,      UpperCAmelCase					,      unittest.TestCase     ):

  '''simple docstring'''







  _lowerCamelCase    =  AutoencoderKL
  _lowerCamelCase    =  '''sample'''
  _lowerCamelCase    =  1E-2


  @property
  def  lowerCamelCase       (      self		:			str							)						->	Tuple:


         """simple docstring"""
         _UpperCAmelCase							=    4
         _UpperCAmelCase							=    3
         _UpperCAmelCase							=    (32, 32)

         _UpperCAmelCase							=    floats_tensor((batch_size, num_channels) + sizes							).to(lowerCamelCase							)

         return {"sample": image}


  @property
  def  lowerCamelCase       (      self		:			Optional[Any]							)						->	List[str]:


         """simple docstring"""
         return (3, 32, 32)


  @property
  def  lowerCamelCase       (      self		:			List[Any]							)						->	Tuple:


         """simple docstring"""
         return (3, 32, 32)


  def  lowerCamelCase       (      self		:			Optional[Any]							)						->	Any:


         """simple docstring"""
         _UpperCAmelCase							=    {
             """block_out_channels""": [32, 64],
             """in_channels""": 3,
             """out_channels""": 3,
             """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
             """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
             """latent_channels""": 4,
         }
         _UpperCAmelCase							=    self.dummy_input
         return init_dict, inputs_dict


  def  lowerCamelCase       (      self		:			Union[str, Any]							)						->	Optional[int]:


         """simple docstring"""
         pass


  def  lowerCamelCase       (      self		:			Union[str, Any]							)						->	Tuple:


         """simple docstring"""
         pass


  @unittest.skipIf(torch_device == """mps"""       ,				"""Gradient checkpointing skipped on MPS"""							)
  def  lowerCamelCase       (      self		:			Optional[Any]							)						->	Any:


         """simple docstring"""
         # enable deterministic behavior for gradient checkpointing
         _UpperCAmelCase	,			_UpperCAmelCase							=    self.prepare_init_args_and_inputs_for_common()
         _UpperCAmelCase							=    self.model_class(**lowerCamelCase							)
         model.to(lowerCamelCase							)

         assert not model.is_gradient_checkpointing and model.training

         _UpperCAmelCase							=    model(**lowerCamelCase							).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model.zero_grad()

         _UpperCAmelCase							=    torch.randn_like(lowerCamelCase							)
         _UpperCAmelCase							=    (out - labels).mean()
         loss.backward()

         # re-instantiate the model now enabling gradient checkpointing
         _UpperCAmelCase							=    self.model_class(**lowerCamelCase							)
         # clone model
         model_a.load_state_dict(model.state_dict()							)
         model_a.to(lowerCamelCase							)
         model_a.enable_gradient_checkpointing()

         assert model_a.is_gradient_checkpointing and model_a.training

         _UpperCAmelCase							=    model_a(**lowerCamelCase							).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model_a.zero_grad()
         _UpperCAmelCase							=    (out_a - labels).mean()
         loss_a.backward()

         # compare the output and parameters gradients
         self.assertTrue((loss - loss_a).abs() < 1E-5							)
         _UpperCAmelCase							=    dict(model.named_parameters()							)
         _UpperCAmelCase							=    dict(model_a.named_parameters()							)
         for name, param in named_params.items():
                self.assertTrue(torch_all_close(param.grad.data       ,				named_params_a[name].grad.data       ,				atol=5E-5							)							)


  def  lowerCamelCase       (      self		:			int							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase	,			_UpperCAmelCase							=    AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""       ,				output_loading_info=lowerCamelCase							)
         self.assertIsNotNone(lowerCamelCase							)
         self.assertEqual(len(loading_info["""missing_keys"""]							)       ,				0							)

         model.to(lowerCamelCase							)
         _UpperCAmelCase							=    model(**self.dummy_input							)

         assert image is not None, "Make sure output is not None"





  def  lowerCamelCase       (      self		:			Optional[int]							)						->	Dict:


         """simple docstring"""
         _UpperCAmelCase							=    AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""							)
         _UpperCAmelCase							=    model.to(lowerCamelCase							)
         model.eval()

         if torch_device == "mps":
                _UpperCAmelCase							=    torch.manual_seed(0							)
         else:
                _UpperCAmelCase							=    torch.Generator(device=lowerCamelCase							).manual_seed(0							)

         _UpperCAmelCase							=    torch.randn(
             1       ,				model.config.in_channels       ,				model.config.sample_size       ,				model.config.sample_size       ,				generator=torch.manual_seed(0							)       ,				)
         _UpperCAmelCase							=    image.to(lowerCamelCase							)
         with torch.no_grad():
                _UpperCAmelCase							=    model(lowerCamelCase       ,				sample_posterior=lowerCamelCase       ,				generator=lowerCamelCase							).sample

         _UpperCAmelCase							=    output[0, -1, -3:, -3:].flatten().cpu()

         # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
         # the expected output slices are not the same for CPU and GPU.
         if torch_device == "mps":
                _UpperCAmelCase							=    torch.tensor(
                    [
                        -4.0078E-01,
                        -3.8323E-04,
                        -1.2681E-01,
                        -1.1462E-01,
                        2.0095E-01,
                        1.0893E-01,
                        -8.8247E-02,
                        -3.0361E-01,
                        -9.8644E-03,
                    ]							)
         elif torch_device == "cpu":
                _UpperCAmelCase							=    torch.tensor(
                    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]							)
         else:
                _UpperCAmelCase							=    torch.tensor(
                    [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]							)

         self.assertTrue(torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				rtol=1E-2							)							)



@slow
class 		SCREAMING_SNAKE_CASE__	(  unittest.TestCase     ):

  '''simple docstring'''







  def  lowerCamelCase       (      self		:			Union[str, Any]       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Tuple							)						->	Optional[int]:


         """simple docstring"""
         return f"""gaussian_noise_s={seed}_shape={"_".join([str(lowerCamelCase							) for s in shape]							)}.npy"""


  def  lowerCamelCase       (      self		:			Dict							)						->	Dict:


         """simple docstring"""
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()


  def  lowerCamelCase       (      self		:			str       ,				lowerCamelCase		:			Dict=0       ,				lowerCamelCase		:			Optional[Any]=(4, 3, 512, 512)       ,				lowerCamelCase		:			str=False							)						->	Any:


         """simple docstring"""
         _UpperCAmelCase							=    torch.floataa if fpaa else torch.floataa
         _UpperCAmelCase							=    torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCamelCase       ,				lowerCamelCase							)							)							).to(lowerCamelCase							).to(lowerCamelCase							)
         return image


  def  lowerCamelCase       (      self		:			Dict       ,				lowerCamelCase		:			Optional[int]="CompVis/stable-diffusion-v1-4"       ,				lowerCamelCase		:			Any=False							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    """fp16""" if fpaa else None
         _UpperCAmelCase							=    torch.floataa if fpaa else torch.floataa

         _UpperCAmelCase							=    AutoencoderKL.from_pretrained(
             lowerCamelCase       ,				subfolder="""vae"""       ,				torch_dtype=lowerCamelCase       ,				revision=lowerCamelCase       ,				)
         model.to(lowerCamelCase							).eval()

         return model


  def  lowerCamelCase       (      self		:			Dict       ,				lowerCamelCase		:			Union[str, Any]=0							)						->	str:


         """simple docstring"""
         if torch_device == "mps":
                return torch.manual_seed(lowerCamelCase							)
         return torch.Generator(device=lowerCamelCase							).manual_seed(lowerCamelCase							)


  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
          [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
          # fmt: on
      ]							)
  def  lowerCamelCase       (      self		:			List[str]       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Optional[Any]							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model()
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase							)
         _UpperCAmelCase							=    self.get_generator(lowerCamelCase							)

         with torch.no_grad():
                _UpperCAmelCase							=    model(lowerCamelCase       ,				generator=lowerCamelCase       ,				sample_posterior=lowerCamelCase							).sample

         assert sample.shape == image.shape

         _UpperCAmelCase							=    sample[-1, -2:, -2:, :2].flatten().float().cpu()
         _UpperCAmelCase							=    torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice							)

         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=3E-3							)


  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
          [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
          # fmt: on
      ]							)
  @require_torch_gpu
  def  lowerCamelCase       (      self		:			int       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			List[str]							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model(fpaa=lowerCamelCase							)
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase       ,				fpaa=lowerCamelCase							)
         _UpperCAmelCase							=    self.get_generator(lowerCamelCase							)

         with torch.no_grad():
                _UpperCAmelCase							=    model(lowerCamelCase       ,				generator=lowerCamelCase       ,				sample_posterior=lowerCamelCase							).sample

         assert sample.shape == image.shape

         _UpperCAmelCase							=    sample[-1, -2:, :2, -2:].flatten().float().cpu()
         _UpperCAmelCase							=    torch.tensor(lowerCamelCase							)

         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=1E-2							)


  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
          [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
          # fmt: on
      ]							)
  def  lowerCamelCase       (      self		:			Tuple       ,				lowerCamelCase		:			Tuple       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			Optional[Any]							)						->	List[str]:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model()
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase							)

         with torch.no_grad():
                _UpperCAmelCase							=    model(lowerCamelCase							).sample

         assert sample.shape == image.shape

         _UpperCAmelCase							=    sample[-1, -2:, -2:, :2].flatten().float().cpu()
         _UpperCAmelCase							=    torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice							)

         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=3E-3							)


  @parameterized.expand(
      [
          # fmt: off
          [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
          [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
          # fmt: on
      ]							)
  @require_torch_gpu
  def  lowerCamelCase       (      self		:			Optional[Any]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			List[Any]							)						->	int:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model()
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase       ,				shape=(3, 4, 64, 64)							)

         with torch.no_grad():
                _UpperCAmelCase							=    model.decode(lowerCamelCase							).sample

         assert list(sample.shape							) == [3, 3, 512, 512]

         _UpperCAmelCase							=    sample[-1, -2:, :2, -2:].flatten().cpu()
         _UpperCAmelCase							=    torch.tensor(lowerCamelCase							)

         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=1E-3							)


  @parameterized.expand(
      [
          # fmt: off
          [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
          [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
          # fmt: on
      ]							)
  @require_torch_gpu
  def  lowerCamelCase       (      self		:			Tuple       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			int							)						->	Dict:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model(fpaa=lowerCamelCase							)
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase       ,				shape=(3, 4, 64, 64)       ,				fpaa=lowerCamelCase							)

         with torch.no_grad():
                _UpperCAmelCase							=    model.decode(lowerCamelCase							).sample

         assert list(sample.shape							) == [3, 3, 512, 512]

         _UpperCAmelCase							=    sample[-1, -2:, :2, -2:].flatten().float().cpu()
         _UpperCAmelCase							=    torch.tensor(lowerCamelCase							)

         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=5E-3							)


  @parameterized.expand([(13,), (16,), (27,)]							)
  @require_torch_gpu
  @unittest.skipIf(not is_xformers_available()       ,				reason="""xformers is not required when using PyTorch 2.0."""							)
  def  lowerCamelCase       (      self		:			Union[str, Any]       ,				lowerCamelCase		:			str							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model(fpaa=lowerCamelCase							)
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase       ,				shape=(3, 4, 64, 64)       ,				fpaa=lowerCamelCase							)

         with torch.no_grad():
                _UpperCAmelCase							=    model.decode(lowerCamelCase							).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
                _UpperCAmelCase							=    model.decode(lowerCamelCase							).sample

         assert list(sample.shape							) == [3, 3, 512, 512]

         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=1E-1							)


  @parameterized.expand([(13,), (16,), (37,)]							)
  @require_torch_gpu
  @unittest.skipIf(not is_xformers_available()       ,				reason="""xformers is not required when using PyTorch 2.0."""							)
  def  lowerCamelCase       (      self		:			Tuple       ,				lowerCamelCase		:			Dict							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model()
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase       ,				shape=(3, 4, 64, 64)							)

         with torch.no_grad():
                _UpperCAmelCase							=    model.decode(lowerCamelCase							).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
                _UpperCAmelCase							=    model.decode(lowerCamelCase							).sample

         assert list(sample.shape							) == [3, 3, 512, 512]

         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=1E-2							)





  @parameterized.expand(
      [
          # fmt: off
          [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
          [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
          # fmt: on
      ]							)
  def  lowerCamelCase       (      self		:			List[str]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Any							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_sd_vae_model()
         _UpperCAmelCase							=    self.get_sd_image(lowerCamelCase							)
         _UpperCAmelCase							=    self.get_generator(lowerCamelCase							)

         with torch.no_grad():
                _UpperCAmelCase							=    model.encode(lowerCamelCase							).latent_dist
                _UpperCAmelCase							=    dist.sample(generator=lowerCamelCase							)

         assert list(sample.shape							) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

         _UpperCAmelCase							=    sample[0, -1, -3:, -3:].flatten().cpu()
         _UpperCAmelCase							=    torch.tensor(lowerCamelCase							)

         _UpperCAmelCase							=    3E-3 if torch_device != """mps""" else 1E-2
         assert torch_all_close(lowerCamelCase       ,				lowerCamelCase       ,				atol=lowerCamelCase							)






from __future__ import annotations

from math import gcd





def  __magic_name__  (				lowercase				,       lowercase = 2				,       lowercase = 1				,       lowercase = 3				,       )				-> int | None:

							"""simple docstring"""



							if num < 2:
														raise ValueError("""The input value cannot be less than 2"""				)

							# Because of the relationship between ``f(f(x))`` and ``f(x)``, this
							# algorithm struggles to find factors that are divisible by two.
							# As a workaround, we specifically check for two and even inputs.
							#   See: https://math.stackexchange.com/a/2856214/165820
							if num > 2 and num % 2 == 0:
														return 2

							# Pollard's Rho algorithm requires a function that returns pseudorandom
							# values between 0 <= X < ``num``.  It doesn't need to be random in the
							# sense that the output value is cryptographically secure or difficult
							# to calculate, it only needs to be random in the sense that all output
							# values should be equally likely to appear.
							# For this reason, Pollard suggested using ``f(x) = (x**2 - 1) % num``
							# However, the success of Pollard's algorithm isn't guaranteed and is
							# determined in part by the initial seed and the chosen random function.
							# To make retries easier, we will instead use ``f(x) = (x**2 + C) % num``
							# where ``C`` is a value that we can modify between each attempt.
							def rand_fn(lowercase				,       lowercase				,       lowercase				) -> int:
														return (pow(lowercase				,       2				) + step) % modulus

							for _ in range(lowercase				):
														# These track the position within the cycle detection logic.
														lowercase_	:   List[Any]  				=     seed
														lowercase_	:   Any  				=     seed

														while True:
																					# At each iteration, the tortoise moves one step and the hare moves two.
																					lowercase_	:   List[str]  				=     rand_fn(lowercase				,       lowercase				,       lowercase				)
																					lowercase_	:   str  				=     rand_fn(lowercase				,       lowercase				,       lowercase				)
																					lowercase_	:   Any  				=     rand_fn(lowercase				,       lowercase				,       lowercase				)

																					# At some point both the tortoise and the hare will enter a cycle whose
																					# length ``p`` is a divisor of ``num``.  Once in that cycle, at some point
																					# the tortoise and hare will end up on the same value modulo ``p``.
																					# We can detect when this happens because the position difference between
																					# the tortoise and the hare will share a common divisor with ``num``.
																					lowercase_	:   Tuple  				=     gcd(hare - tortoise				,       lowercase				)

																					if divisor == 1:
																												# No common divisor yet, just keep searching.
																												continue
																					else:
																												# We found a common divisor!
																												if divisor == num:
																																			# Unfortunately, the divisor is ``num`` itself and is useless.
																																			break
																												else:
																																			# The divisor is a nontrivial factor of ``num``!
																																			return divisor

        # If we made it here, then this attempt failed.
        # We need to pick a new starting seed for the tortoise and hare
        # in addition to a new step value for the random function.
        # To keep this example implementation deterministic, the
        # new values will be generated based on currently available
        # values instead of using something like ``random.randint``.

        # We can use the hare's position as the new seed.
        # This is actually what Richard Brent's the "optimized" variant does.
														lowercase_	:   int  				=     hare

														# The new step value for the random function can just be incremented.
														# At first the results will be similar to what the old function would
														# have produced, but the value will quickly diverge after a bit.
														step += 1

							# We haven't found a divisor within the requested number of attempts.
							# We were unlucky or ``num`` itself is actually prime.
							return None


if __name__ == "__main__":
							import argparse

							UpperCAmelCase_							=				argparse.ArgumentParser()
							parser.add_argument(
							    """num""",
							    type=int,
							    help="""The value to find a divisor of""",
							)
							parser.add_argument(
							    """--attempts""",
							    type=int,
							    default=3,
							    help="""The number of attempts before giving up""",
							)
							UpperCAmelCase_							=				parser.parse_args()

							UpperCAmelCase_							=				pollard_rho(args.num, attempts=args.attempts)
							if divisor is None:
														print(F'''{args.num} is probably prime''')
							else:
														UpperCAmelCase_							=				args.num // divisor
														print(F'''{args.num} = {divisor} * {quotient}''')