




import posixpath
from pathlib import Path
from unittest.mock import patch

import pytest
from fsspec.implementations.local import AbstractFileSystem, LocalFileSystem, stringify_path
from fsspec.registry import _registry as _fsspec_registry
class 				__snake_case     (_a						):
				lowerCAmelCase__			=   "mock"



				def __init__(					self					:						Union[str, Any]						,		*_UpperCAmelCase					:						Union[str, Any]						,		_UpperCAmelCase					:						Optional[int]						,		**_UpperCAmelCase					:						List[str]					)     ->			Optional[int]:





								'''simple docstring'''




								super().__init__()
								_lowerCAmelCase			:			Tuple        =					LocalFileSystem(*_UpperCAmelCase						,		**_UpperCAmelCase					)
								_lowerCAmelCase			:			List[str]        =					Path(_UpperCAmelCase					).resolve().as_posix() + """/"""



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						Optional[int]						,		*_UpperCAmelCase					:						Dict						,		**_UpperCAmelCase					:						int					)     ->			int:





								'''simple docstring'''




								_lowerCAmelCase			:			str        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.mkdir(_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)



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





								'''simple docstring'''




								_lowerCAmelCase			:			List[str]        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.makedirs(_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						str						,		_UpperCAmelCase					:						str					)     ->			Any:





								'''simple docstring'''




								_lowerCAmelCase			:			Dict        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.rmdir(_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Union[str, Any]						,		_UpperCAmelCase					:						str						,		_UpperCAmelCase					:						List[Any]=True						,		*_UpperCAmelCase					:						Optional[int]						,		**_UpperCAmelCase					:						Tuple					)     ->			Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			Union[str, Any]        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								_lowerCAmelCase			:			str        =					self._fs.ls(_UpperCAmelCase						,		detail=_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)
								if detail:
												return [{**info, "name": info["name"][len(self.local_root_dir					) :]} for info in out]
								else:
												return [name[len(self.local_root_dir					) :] for name in out]



				def 					SCREAMING_SNAKE_CASE				(					self					:						List[Any]						,		_UpperCAmelCase					:						Any						,		*_UpperCAmelCase					:						List[str]						,		**_UpperCAmelCase					:						Optional[int]					)     ->			Union[str, Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			int        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								_lowerCAmelCase			:			Dict        =					dict(self._fs.info(_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)					)
								_lowerCAmelCase			:			Union[str, Any]        =					out["""name"""][len(self.local_root_dir					) :]
								return out



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[Any]						,		_UpperCAmelCase					:						Dict						,		_UpperCAmelCase					:						List[str]						,		*_UpperCAmelCase					:						Optional[Any]						,		**_UpperCAmelCase					:						List[str]					)     ->			str:





								'''simple docstring'''




								_lowerCAmelCase			:			Tuple        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								_lowerCAmelCase			:			Union[str, Any]        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.cp_file(_UpperCAmelCase						,		_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						Any						,		*_UpperCAmelCase					:						int						,		**_UpperCAmelCase					:						List[str]					)     ->			Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			Any        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.rm_file(_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[Any]						,		_UpperCAmelCase					:						Optional[Any]						,		*_UpperCAmelCase					:						str						,		**_UpperCAmelCase					:						List[Any]					)     ->			Tuple:





								'''simple docstring'''




								_lowerCAmelCase			:			Optional[Any]        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.rm(_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[int]						,		_UpperCAmelCase					:						Union[str, Any]						,		*_UpperCAmelCase					:						Dict						,		**_UpperCAmelCase					:						Union[str, Any]					)     ->			Any:





								'''simple docstring'''




								_lowerCAmelCase			:			Union[str, Any]        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs._open(_UpperCAmelCase						,		*_UpperCAmelCase						,		**_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						Tuple					)     ->			Any:





								'''simple docstring'''




								_lowerCAmelCase			:			List[Any]        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.created(_UpperCAmelCase					)



				def 					SCREAMING_SNAKE_CASE				(					self					:						Any						,		_UpperCAmelCase					:						Dict					)     ->			Dict:





								'''simple docstring'''




								_lowerCAmelCase			:			Optional[int]        =					posixpath.join(self.local_root_dir						,		self._strip_protocol(_UpperCAmelCase					)					)
								return self._fs.modified(_UpperCAmelCase					)



				@classmethod
				def 					SCREAMING_SNAKE_CASE				(					cls					:						List[Any]						,		_UpperCAmelCase					:						Tuple					)     ->			str:





								'''simple docstring'''




								_lowerCAmelCase			:			Any        =					stringify_path(_UpperCAmelCase					)
								if path.startswith("""mock://"""					):
												_lowerCAmelCase			:			Union[str, Any]        =					path[7:]
								return path






class 				__snake_case     (_a						):
				lowerCAmelCase__			=   "tmp"
				lowerCAmelCase__			=   None



				def __init__(					self					:						Optional[Any]						,		*_UpperCAmelCase					:						Tuple						,		**_UpperCAmelCase					:						Tuple					)     ->			List[Any]:





								'''simple docstring'''




								assert self.tmp_dir is not None, "TmpDirFileSystem.tmp_dir is not set"
								super().__init__(*_UpperCAmelCase						,		**_UpperCAmelCase						,		local_root_dir=self.tmp_dir						,		auto_mkdir=_UpperCAmelCase					)



				@classmethod
				def 					SCREAMING_SNAKE_CASE				(					cls					:						int						,		_UpperCAmelCase					:						Tuple					)     ->			Dict:





								'''simple docstring'''




								_lowerCAmelCase			:			int        =					stringify_path(_UpperCAmelCase					)
								if path.startswith("""tmp://"""					):
												_lowerCAmelCase			:			str        =					path[6:]
								return path






@pytest.fixture
def 							_UpperCAmelCase     ():
				'''simple docstring'''
				_lowerCAmelCase			:			str        =					MockFileSystem
				_lowerCAmelCase			:			Any        =					TmpDirFileSystem
				yield
				del _fsspec_registry["mock"]
				del _fsspec_registry["tmp"]




@pytest.fixture
def 							_UpperCAmelCase     (UpperCamelCase_       :       Dict					,			UpperCamelCase_       :       List[Any]				):
				'''simple docstring'''
				_lowerCAmelCase			:			int        =					tmp_path_factory.mktemp("""mockfs"""				)
				return MockFileSystem(local_root_dir=UpperCamelCase_					,			auto_mkdir=UpperCamelCase_				)





@pytest.fixture
def 							_UpperCAmelCase     (UpperCamelCase_       :       Optional[int]					,			UpperCamelCase_       :       Union[str, Any]				):
				'''simple docstring'''
				_lowerCAmelCase			:			Any        =					tmp_path_factory.mktemp("""tmpfs"""				)
				with patch.object(UpperCamelCase_					,			"""tmp_dir"""					,			UpperCamelCase_				):
								yield TmpDirFileSystem()







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 				__snake_case     (unittest.TestCase						):
				lowerCAmelCase__			=   "hf-internal-testing/tiny-stable-diffusion-pipe"
				lowerCAmelCase__			=   1
				lowerCAmelCase__			=   7_7
				lowerCAmelCase__			=   3_2
				lowerCAmelCase__			=   4
				lowerCAmelCase__			=   lowerCAmelCase__			=   6_4
				lowerCAmelCase__			=   torch.manual_seed(0						)



				def 					SCREAMING_SNAKE_CASE				(					self					:						List[str]						,		_UpperCAmelCase					:						Optional[Any]=0.9999					)     ->			List[Any]:





								'''simple docstring'''




								_lowerCAmelCase			:			List[Any]        =					UNetaDConditionModel.from_pretrained(self.model_id						,		subfolder="""unet"""					)
								_lowerCAmelCase			:			Optional[Any]        =					unet.to(_UpperCAmelCase					)
								_lowerCAmelCase			:			Optional[int]        =					EMAModel(unet.parameters()						,		decay=_UpperCAmelCase						,		model_cls=_UpperCAmelCase						,		model_config=unet.config					)
								return unet, ema_unet



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





								'''simple docstring'''




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



				def 					SCREAMING_SNAKE_CASE				(					self					:						Dict						,		_UpperCAmelCase					:						Dict					)     ->			Dict:





								'''simple docstring'''




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



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[Any]					)     ->			Tuple:





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			Optional[int]        =					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					:						Optional[int]					)     ->			str:





								'''simple docstring'''




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

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



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





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			Union[str, 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.
								_lowerCAmelCase			:			str        =					self.simulate_backprop(_UpperCAmelCase					)

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

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



				def 					SCREAMING_SNAKE_CASE				(					self					:						Optional[int]					)     ->			Tuple:





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			str        =					self.get_models()

								# First backprop + EMA
								_lowerCAmelCase			:			Optional[int]        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_one.parameters()					)
								_lowerCAmelCase			:			List[str]        =					ema_unet.shadow_params

								# Second backprop + EMA
								_lowerCAmelCase			:			Optional[int]        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_two.parameters()					)
								_lowerCAmelCase			:			Tuple        =					ema_unet.shadow_params

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



				def 					SCREAMING_SNAKE_CASE				(					self					:						List[Any]					)     ->			str:





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			int        =					self.get_models(decay=0.0					)
								_lowerCAmelCase			:			Dict        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_one.parameters()					)
								_lowerCAmelCase			:			Any        =					ema_unet.shadow_params

								_lowerCAmelCase			:			List[Any]        =					self.simulate_backprop(_UpperCAmelCase					)
								ema_unet.step(unet_step_two.parameters()					)
								_lowerCAmelCase			:			int        =					ema_unet.shadow_params

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



				@skip_mps
				def 					SCREAMING_SNAKE_CASE				(					self					:						Union[str, Any]					)     ->			str:





								'''simple docstring'''




								_lowerCAmelCase    ,			_lowerCAmelCase			:			Tuple        =					self.get_models()
								_lowerCAmelCase    ,			_lowerCAmelCase    ,			_lowerCAmelCase			:			Tuple        =					self.get_dummy_inputs()

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

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

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

