





'''simple docstring'''





import math
from timeit import timeit
def   _A (	SCREAMING_SNAKE_CASE  :     Tuple		):





  """simple docstring"""



  a__			:      Optional[int]            =0
  a__			:      Any            =abs(SCREAMING_SNAKE_CASE		)
  while True:
    a__			:      str            =n // 10
    digits += 1
    if n == 0:
      break
  return digits
def   _A (	SCREAMING_SNAKE_CASE  :     Optional[int]		):





  """simple docstring"""



  return 1 if n == 0 else math.floor(math.log(abs(SCREAMING_SNAKE_CASE		) , 10		) + 1		)
def   _A (	SCREAMING_SNAKE_CASE  :     int		):





  """simple docstring"""



  return len(str(abs(SCREAMING_SNAKE_CASE		)		)		)
def   _A (	):





  """simple docstring"""



  from collections.abc import Callable

  def benchmark_a_function(SCREAMING_SNAKE_CASE  :     Dict , SCREAMING_SNAKE_CASE  :     Tuple		) -> None:
    a__			:      Tuple            =f'''{func.__name__}({value})'''
    a__			:      Union[str, Any]            =timeit(f'''__main__.{call}''' , setup="import __main__"		)
    print(f'''{call}: {func(SCREAMING_SNAKE_CASE		)} -- {timing} seconds'''		)

  for value in (262_144, 1_125_899_906_842_624, 1_267_650_600_228_229_401_496_703_205_376):
    for func in (num_digits, num_digits_fast, num_digits_faster):
      benchmark_a_function(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
    print()


if __name__ == "__main__":
     import doctest

     doctest.testmod()
     benchmark()













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        __lowerCAmelCase (					unittest.TestCase):
      _lowercase       :      str        =    """hf-internal-testing/tiny-stable-diffusion-pipe"""
      _lowercase       :      Optional[int]        =    1
      _lowercase       :      List[str]        =    77
      _lowercase       :      int        =    32
      _lowercase       :      Dict        =    4
      _lowercase       :      str        =    64
      _lowercase       :      Optional[Any]        =    torch.manual_seed(0)







      def      _lowercase						(					self	,   lowerCAmelCase__=0.99_99							)     ->		Tuple:



        '''simple docstring'''




        a__			:      int            =UNetaDConditionModel.from_pretrained(self.model_id	,   subfolder="unet"							)
        a__			:      Any            =unet.to(lowerCAmelCase__							)
        a__			:      int            =EMAModel(unet.parameters()	,   decay=lowerCAmelCase__	,   model_cls=lowerCAmelCase__	,   model_config=unet.config							)
        return unet, ema_unet







      def      _lowercase						(					self							)     ->		Dict:



        '''simple docstring'''




        a__			:      str            =torch.randn(
            self.batch_size	,   self.num_in_channels	,   self.latent_height	,   self.latent_width	,   generator=self.generator							).to(lowerCAmelCase__							)
        a__			:      Any            =torch.randint(0	,   1_0_0_0	,   size=(self.batch_size,)	,   generator=self.generator							).to(lowerCAmelCase__							)
        a__			:      Tuple            =torch.randn(
            self.batch_size	,   self.prompt_length	,   self.text_encoder_hidden_dim	,   generator=self.generator							).to(lowerCAmelCase__							)
        return noisy_latents, timesteps, encoder_hidden_states







      def      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Any:



        '''simple docstring'''




        a__			:      Any            ={}
        for k, param in unet.state_dict().items():
          a__			:      Tuple            =torch.randn_like(lowerCAmelCase__							) + (param * torch.randn_like(lowerCAmelCase__							))
          updated_state_dict.update({k: updated_param}							)
        unet.load_state_dict(lowerCAmelCase__							)
        return unet







      def      _lowercase						(					self							)     ->		Tuple:



        '''simple docstring'''




        a__      ,							a__			:      List[Any]            =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      _lowercase						(					self							)     ->		List[str]:



        '''simple docstring'''




        a__      ,							a__			:      Any            =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()							)
        a__			:      str            =list(unet.parameters()							)
        for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
          assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)

        # 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	,   lowerCAmelCase__							):
          assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)







      def      _lowercase						(					self							)     ->		List[Any]:



        '''simple docstring'''




        a__      ,							a__			:      Optional[Any]            =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.
        a__			:      List[str]            =self.simulate_backprop(lowerCAmelCase__							)

        # 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.
        a__			:      Optional[int]            =list(unet_pseudo_updated_step_one.parameters()							)
        for s_param, param in zip(ema_unet.shadow_params	,   lowerCAmelCase__							):
          assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)

        # 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	,   lowerCAmelCase__							):
          assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)







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



        '''simple docstring'''




        a__      ,							a__			:      Dict            =self.get_models()

        # First backprop + EMA
        a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_one.parameters()							)
        a__			:      Dict            =ema_unet.shadow_params

        # Second backprop + EMA
        a__			:      str            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_two.parameters()							)
        a__			:      List[str]            =ema_unet.shadow_params

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







      def      _lowercase						(					self							)     ->		List[str]:



        '''simple docstring'''




        a__      ,							a__			:      int            =self.get_models(decay=0.0							)
        a__			:      Optional[int]            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_one.parameters()							)
        a__			:      List[Any]            =ema_unet.shadow_params

        a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
        ema_unet.step(unet_step_two.parameters()							)
        a__			:      List[Any]            =ema_unet.shadow_params

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







      @skip_mps
      def      _lowercase						(					self							)     ->		Optional[Any]:



        '''simple docstring'''




        a__      ,							a__			:      Union[str, Any]            =self.get_models()
        a__      ,							a__      ,							a__			:      Tuple            =self.get_dummy_inputs()

        with tempfile.TemporaryDirectory() as tmpdir:
          ema_unet.save_pretrained(lowerCAmelCase__							)
          a__			:      List[str]            =UNetaDConditionModel.from_pretrained(lowerCAmelCase__	,   model_cls=lowerCAmelCase__							)
          a__			:      Tuple            =loaded_unet.to(unet.device							)

        # Since no EMA step has been performed the outputs should match.
        a__			:      Optional[Any]            =unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample
        a__			:      Optional[int]            =loaded_unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample

        assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__	,   atol=1E-4							)






