


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'''




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






   def _snake_case   (		self						,    __A=0.9999			):






     """simple docstring"""





     lowerCamelCase			: str   =    UNetaDConditionModel.from_pretrained(self.model_id						,    subfolder="unet"			)
     lowerCamelCase			: Tuple   =    unet.to(__A			)
     lowerCamelCase			: int   =    EMAModel(unet.parameters()						,    decay=__A						,    model_cls=__A						,    model_config=unet.config			)
     return unet, ema_unet






   def _snake_case   (		self			):






     """simple docstring"""





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






   def _snake_case   (		self						,    __A			):






     """simple docstring"""





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






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase   ,     lowerCamelCase			: Union[str, 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 _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase   ,     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			: int   =    list(unet.parameters()			)
     for s_param, param in zip(ema_unet.shadow_params						,    __A			):
       assert torch.allclose(__A						,    __A			)

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






   def _snake_case   (		self			):






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

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

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






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase   ,     lowerCamelCase			: List[str]   =    self.get_models()

     # First backprop + EMA
     lowerCamelCase			: Optional[Any]   =    self.simulate_backprop(__A			)
     ema_unet.step(unet_step_one.parameters()			)
     lowerCamelCase			: Union[str, Any]   =    ema_unet.shadow_params

     # Second backprop + EMA
     lowerCamelCase			: Tuple   =    self.simulate_backprop(__A			)
     ema_unet.step(unet_step_two.parameters()			)
     lowerCamelCase			: Tuple   =    ema_unet.shadow_params

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






   def _snake_case   (		self			):






     """simple docstring"""





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

     lowerCamelCase			: int   =    self.simulate_backprop(__A			)
     ema_unet.step(unet_step_two.parameters()			)
     lowerCamelCase			: Optional[int]   =    ema_unet.shadow_params

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






   @skip_mps
   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase   ,     lowerCamelCase			: Optional[Any]   =    self.get_models()
     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase			: str   =    self.get_dummy_inputs()

     with tempfile.TemporaryDirectory() as tmpdir:
       ema_unet.save_pretrained(__A			)
       lowerCamelCase			: Any   =    UNetaDConditionModel.from_pretrained(__A						,    model_cls=__A			)
       lowerCamelCase			: Dict   =    loaded_unet.to(unet.device			)

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

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







import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
   from PIL import Image

   from transformers import AutoProcessor, BertTokenizer, CLIPImageProcessor, GitProcessor, PreTrainedTokenizerFast
@require_vision
class      UpperCAmelCase_ ( unittest.TestCase   ):

   '''simple docstring'''




   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    tempfile.mkdtemp()

     lowerCamelCase			: List[Any]   =    CLIPImageProcessor()
     lowerCamelCase			: str   =    BertTokenizer.from_pretrained(
         "hf-internal-testing/tiny-random-BertModel"						,    model_input_names=["input_ids", "attention_mask"]			)

     lowerCamelCase			: List[str]   =    GitProcessor(__A						,    __A			)

     processor.save_pretrained(self.tmpdirname			)






   def _snake_case   (		self						,    **__A			):






     """simple docstring"""





     return AutoProcessor.from_pretrained(self.tmpdirname						,    **__A			).tokenizer






   def _snake_case   (		self						,    **__A			):






     """simple docstring"""





     return AutoProcessor.from_pretrained(self.tmpdirname						,    **__A			).image_processor






   def _snake_case   (		self			):






     """simple docstring"""





     shutil.rmtree(self.tmpdirname			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: str   =    [np.random.randint(255						,    size=(3, 30, 400)						,    dtype=np.uinta			)]

     lowerCamelCase			: Any   =    [Image.fromarray(np.moveaxis(__A						,    0						,    -1			)			) for x in image_inputs]

     return image_inputs






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    GitProcessor(tokenizer=self.get_tokenizer()						,    image_processor=self.get_image_processor()			)
     processor.save_pretrained(self.tmpdirname			)

     lowerCamelCase			: Optional[Any]   =    self.get_tokenizer(bos_token="(BOS)"						,    eos_token="(EOS)"			)
     lowerCamelCase			: List[str]   =    self.get_image_processor(do_normalize=__A						,    padding_value=1.0			)

     lowerCamelCase			: Dict   =    GitProcessor.from_pretrained(
         self.tmpdirname						,    bos_token="(BOS)"						,    eos_token="(EOS)"						,    do_normalize=__A						,    padding_value=1.0			)

     self.assertEqual(processor.tokenizer.get_vocab()						,    tokenizer_add_kwargs.get_vocab()			)
     self.assertIsInstance(processor.tokenizer						,    __A			)

     self.assertEqual(processor.image_processor.to_json_string()						,    image_processor_add_kwargs.to_json_string()			)
     self.assertIsInstance(processor.image_processor						,    __A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Any   =    self.get_image_processor()
     lowerCamelCase			: Optional[int]   =    self.get_tokenizer()

     lowerCamelCase			: Optional[Any]   =    GitProcessor(tokenizer=__A						,    image_processor=__A			)

     lowerCamelCase			: List[str]   =    self.prepare_image_inputs()

     lowerCamelCase			: str   =    image_processor(__A						,    return_tensors="np"			)
     lowerCamelCase			: Union[str, Any]   =    processor(images=__A						,    return_tensors="np"			)

     for key in input_feat_extract.keys():
       self.assertAlmostEqual(input_feat_extract[key].sum()						,    input_processor[key].sum()						,    delta=1e-2			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: int   =    self.get_image_processor()
     lowerCamelCase			: Dict   =    self.get_tokenizer()

     lowerCamelCase			: Optional[int]   =    GitProcessor(tokenizer=__A						,    image_processor=__A			)

     lowerCamelCase			: int   =    "lower newer"

     lowerCamelCase			: str   =    processor(text=__A			)

     lowerCamelCase			: Union[str, Any]   =    tokenizer(__A						,    return_token_type_ids=__A			)

     for key in encoded_tok.keys():
       self.assertListEqual(encoded_tok[key]						,    encoded_processor[key]			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Dict   =    self.get_image_processor()
     lowerCamelCase			: Optional[int]   =    self.get_tokenizer()

     lowerCamelCase			: List[str]   =    GitProcessor(tokenizer=__A						,    image_processor=__A			)

     lowerCamelCase			: str   =    "lower newer"
     lowerCamelCase			: Tuple   =    self.prepare_image_inputs()

     lowerCamelCase			: str   =    processor(text=__A						,    images=__A			)

     self.assertListEqual(list(inputs.keys()			)						,    ["input_ids", "attention_mask", "pixel_values"]			)

     # test if it raises when no input is passed
     with pytest.raises(__A			):
       processor()






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Union[str, Any]   =    self.get_image_processor()
     lowerCamelCase			: Tuple   =    self.get_tokenizer()

     lowerCamelCase			: Tuple   =    GitProcessor(tokenizer=__A						,    image_processor=__A			)

     lowerCamelCase			: Optional[Any]   =    [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

     lowerCamelCase			: str   =    processor.batch_decode(__A			)
     lowerCamelCase			: int   =    tokenizer.batch_decode(__A			)

     self.assertListEqual(__A						,    __A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: List[Any]   =    self.get_image_processor()
     lowerCamelCase			: int   =    self.get_tokenizer()

     lowerCamelCase			: Optional[int]   =    GitProcessor(tokenizer=__A						,    image_processor=__A			)

     lowerCamelCase			: Optional[Any]   =    "lower newer"
     lowerCamelCase			: List[Any]   =    self.prepare_image_inputs()

     lowerCamelCase			: Optional[int]   =    processor(text=__A						,    images=__A			)

     # For now the processor supports only ['input_ids', 'attention_mask', 'pixel_values']
     self.assertListEqual(list(inputs.keys()			)						,    ["input_ids", "attention_mask", "pixel_values"]			)



