


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



     """simple docstring"""






     __magic_name__				=			'hf-internal-testing/tiny-stable-diffusion-pipe'
     __magic_name__				=			1
     __magic_name__				=			7_7
     __magic_name__				=			3_2
     __magic_name__				=			4
     __magic_name__				=			6_4
     __magic_name__				=			torch.manual_seed(0		)







     def a		(       self					,		snake_case__=0.9999      ):

            '''simple docstring'''

            _lowerCAmelCase			:							List[str]      	=				UNetaDConditionModel.from_pretrained(self.model_id					,		subfolder='unet'      )
            _lowerCAmelCase			:							Any      	=				unet.to(_SCREAMING_SNAKE_CASE      )
            _lowerCAmelCase			:							List[str]      	=				EMAModel(unet.parameters()					,		decay=_SCREAMING_SNAKE_CASE					,		model_cls=_SCREAMING_SNAKE_CASE					,		model_config=unet.config      )
            return unet, ema_unet







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							Any      	=				torch.randn(
                self.batch_size					,		self.num_in_channels					,		self.latent_height					,		self.latent_width					,		generator=self.generator      ).to(_SCREAMING_SNAKE_CASE      )
            _lowerCAmelCase			:							str      	=				torch.randint(0					,		1000					,		size=(self.batch_size,)					,		generator=self.generator      ).to(_SCREAMING_SNAKE_CASE      )
            _lowerCAmelCase			:							int      	=				torch.randn(
                self.batch_size					,		self.prompt_length					,		self.text_encoder_hidden_dim					,		generator=self.generator      ).to(_SCREAMING_SNAKE_CASE      )
            return noisy_latents, timesteps, encoder_hidden_states







     def a		(       self					,		snake_case__      ):

            '''simple docstring'''

            _lowerCAmelCase			:							List[str]      	=				{}
            for k, param in unet.state_dict().items():
                   _lowerCAmelCase			:							List[Any]      	=				torch.randn_like(_SCREAMING_SNAKE_CASE      ) + (param * torch.randn_like(_SCREAMING_SNAKE_CASE      ))
                   updated_state_dict.update({k: updated_param}      )
            unet.load_state_dict(_SCREAMING_SNAKE_CASE      )
            return unet







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							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 a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							List[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()      )
            _lowerCAmelCase			:							List[str]      	=				list(unet.parameters()      )
            for s_param, param in zip(ema_unet.shadow_params					,		_SCREAMING_SNAKE_CASE      ):
                   assert torch.allclose(_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE      )

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







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							Optional[int]      	=				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			:							int      	=				self.simulate_backprop(_SCREAMING_SNAKE_CASE      )

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

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







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							str      	=				self.get_models()

            # First backprop + EMA
            _lowerCAmelCase			:							List[str]      	=				self.simulate_backprop(_SCREAMING_SNAKE_CASE      )
            ema_unet.step(unet_step_one.parameters()      )
            _lowerCAmelCase			:							Union[str, Any]      	=				ema_unet.shadow_params

            # Second backprop + EMA
            _lowerCAmelCase			:							Dict      	=				self.simulate_backprop(_SCREAMING_SNAKE_CASE      )
            ema_unet.step(unet_step_two.parameters()      )
            _lowerCAmelCase			:							str      	=				ema_unet.shadow_params

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







     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							Tuple      	=				self.get_models(decay=0.0      )
            _lowerCAmelCase			:							List[Any]      	=				self.simulate_backprop(_SCREAMING_SNAKE_CASE      )
            ema_unet.step(unet_step_one.parameters()      )
            _lowerCAmelCase			:							List[Any]      	=				ema_unet.shadow_params

            _lowerCAmelCase			:							Dict      	=				self.simulate_backprop(_SCREAMING_SNAKE_CASE      )
            ema_unet.step(unet_step_two.parameters()      )
            _lowerCAmelCase			:							int      	=				ema_unet.shadow_params

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







     @skip_mps
     def a		(       self      ):

            '''simple docstring'''

            _lowerCAmelCase			:							Dict      	=				self.get_models()
            _lowerCAmelCase			:							Union[str, Any]      	=				self.get_dummy_inputs()

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

            # Since no EMA step has been performed the outputs should match.
            _lowerCAmelCase			:							Dict      	=				unet(_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE      ).sample
            _lowerCAmelCase			:							Union[str, Any]      	=				loaded_unet(_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE      ).sample

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









def    lowerCAmelCase__		(			_a				:	str     ,					_a				:	int = 13			):
		snake_case_ :							List[Any]         =				""
		for c in s:
				if "A" <= c <= "Z":
						out += chr(ord("A"			) + (ord(_a			) - ord("A"			) + n) % 26			)
				elif "a" <= c <= "z":
						out += chr(ord("a"			) + (ord(_a			) - ord("a"			) + n) % 26			)
				else:
						out += c
		return out

def    lowerCAmelCase__		(			):
		snake_case_ :							int         =				input("Enter message: "			)

		snake_case_ :							List[Any]         =				dencrypt(_a     ,					13			)
		print("Encryption:"     ,					_a			)

		snake_case_ :							int         =				dencrypt(_a     ,					13			)
		print("Decryption: "     ,					_a			)


if __name__ == "__main__":
							import doctest

							doctest.testmod()
							main()
