

"""simple docstring"""

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  __a				(	unittest.TestCase					):





       """simple docstring"""


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






       def  SCREAMING_SNAKE_CASE_  (      self ,				snake_case=0.9_999      ):





        """simple docstring"""

        lowerCAmelCase__      :							Tuple    =							UNetaDConditionModel.from_pretrained(self.model_id ,				subfolder="unet"      )
        lowerCAmelCase__      :							int    =							unet.to(lowercase_      )
        lowerCAmelCase__      :							Optional[int]    =							EMAModel(unet.parameters() ,				decay=lowercase_ ,				model_cls=lowercase_ ,				model_config=unet.config      )
        return unet, ema_unet






       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        lowerCAmelCase__      :							Optional[Any]    =							torch.randn(
            self.batch_size ,				self.num_in_channels ,				self.latent_height ,				self.latent_width ,				generator=self.generator      ).to(lowercase_      )
        lowerCAmelCase__      :							str    =							torch.randint(0 ,				1_000 ,				size=(self.batch_size,) ,				generator=self.generator      ).to(lowercase_      )
        lowerCAmelCase__      :							int    =							torch.randn(
            self.batch_size ,				self.prompt_length ,				self.text_encoder_hidden_dim ,				generator=self.generator      ).to(lowercase_      )
        return noisy_latents, timesteps, encoder_hidden_states






       def  SCREAMING_SNAKE_CASE_  (      self ,				snake_case      ):





        """simple docstring"""

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






       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        lowerCAmelCase__      :							List[str]    =							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  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

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

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






       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        lowerCAmelCase__      :							Tuple    =							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.
        lowerCAmelCase__      :							Optional[Any]    =							self.simulate_backprop(lowercase_      )

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

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






       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        lowerCAmelCase__      :							Tuple    =							self.get_models()

        # First backprop + EMA
        lowerCAmelCase__      :							int    =							self.simulate_backprop(lowercase_      )
        ema_unet.step(unet_step_one.parameters()      )
        lowerCAmelCase__      :							str    =							ema_unet.shadow_params

        # Second backprop + EMA
        lowerCAmelCase__      :							Union[str, Any]    =							self.simulate_backprop(lowercase_      )
        ema_unet.step(unet_step_two.parameters()      )
        lowerCAmelCase__      :							Union[str, Any]    =							ema_unet.shadow_params

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






       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

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

        lowerCAmelCase__      :							Optional[Any]    =							self.simulate_backprop(lowercase_      )
        ema_unet.step(unet_step_two.parameters()      )
        lowerCAmelCase__      :							Dict    =							ema_unet.shadow_params

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






       @skip_mps
       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        lowerCAmelCase__      :							Tuple    =							self.get_models()
        lowerCAmelCase__      :							str    =							self.get_dummy_inputs()

        with tempfile.TemporaryDirectory() as tmpdir:
         ema_unet.save_pretrained(lowercase_      )
         lowerCAmelCase__      :							int    =							UNetaDConditionModel.from_pretrained(lowercase_ ,				model_cls=lowercase_      )
         lowerCAmelCase__      :							Tuple    =							loaded_unet.to(unet.device      )

        # Since no EMA step has been performed the outputs should match.
        lowerCAmelCase__      :							str    =							unet(lowercase_ ,				lowercase_ ,				lowercase_      ).sample
        lowerCAmelCase__      :							int    =							loaded_unet(lowercase_ ,				lowercase_ ,				lowercase_      ).sample

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









"""simple docstring"""

from itertools import product



def 			SCREAMING_SNAKE_CASE       (   lowercase__ = 1_5      ) ->	int:
 lowerCAmelCase__      :							float    =							0.0
 lowerCAmelCase__      :							float
 lowerCAmelCase__      :							int
 lowerCAmelCase__      :							int
 lowerCAmelCase__      :							int
 lowerCAmelCase__      :							int
 lowerCAmelCase__      :							tuple[int, ...]

 for series in product(range(2      )				,				repeat=lowercase__      ):
  lowerCAmelCase__      :							Optional[int]    =							series.count(1      )
  lowerCAmelCase__      :							Dict    =							num_turns - num_blue
  if num_red >= num_blue:
   continue
  lowerCAmelCase__      :							Dict    =							1.0
  for ind, col in enumerate(lowercase__				,				2      ):
   if col == 0:
    prob *= (ind - 1) / ind
   else:
    prob *= 1 / ind

  total_prob += prob

 return int(1 / total_prob      )


if __name__ == "__main__":
    print(f"{solution() = }")






