

'''simple docstring'''






import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np






class     lowercase_			:




  """simple docstring"""




  def       SCREAMING_SNAKE_CASE   (   self     :  Union[str, Any]    ):
   torch.manual_seed(0    )
   __lowercase													=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''    )

   torch.manual_seed(0    )
   __lowercase													=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''    )

   torch.manual_seed(0    )
   __lowercase													=				UNetaDConditionModel(
       sample_size=3_2      ,layers_per_block=1      ,block_out_channels=[3_2, 6_4]      ,down_block_types=[
           '''ResnetDownsampleBlock2D''',
           '''SimpleCrossAttnDownBlock2D''',
       ]      ,mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''      ,up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']      ,in_channels=3      ,out_channels=6      ,cross_attention_dim=3_2      ,encoder_hid_dim=3_2      ,attention_head_dim=8      ,addition_embed_type='''text'''      ,addition_embed_type_num_heads=2      ,cross_attention_norm='''group_norm'''      ,resnet_time_scale_shift='''scale_shift'''      ,act_fn='''gelu'''      ,)
   unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

   torch.manual_seed(0    )
   __lowercase													=				DDPMScheduler(
       num_train_timesteps=1_0_0_0      ,beta_schedule='''squaredcos_cap_v2'''      ,beta_start=0.0_0_0_1      ,beta_end=0.0_2      ,thresholding=_UpperCAmelCase      ,dynamic_thresholding_ratio=0.9_5      ,sample_max_value=1.0      ,prediction_type='''epsilon'''      ,variance_type='''learned_range'''      ,)

   torch.manual_seed(0    )
   __lowercase													=				IFWatermarker()

   return {
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "unet": unet,
       "scheduler": scheduler,
       "watermarker": watermarker,
       "safety_checker": None,
       "feature_extractor": None,
   }






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   torch.manual_seed(0    )
   __lowercase													=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''    )

   torch.manual_seed(0    )
   __lowercase													=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''    )

   torch.manual_seed(0    )
   __lowercase													=				UNetaDConditionModel(
       sample_size=3_2      ,layers_per_block=[1, 2]      ,block_out_channels=[3_2, 6_4]      ,down_block_types=[
           '''ResnetDownsampleBlock2D''',
           '''SimpleCrossAttnDownBlock2D''',
       ]      ,mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''      ,up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']      ,in_channels=6      ,out_channels=6      ,cross_attention_dim=3_2      ,encoder_hid_dim=3_2      ,attention_head_dim=8      ,addition_embed_type='''text'''      ,addition_embed_type_num_heads=2      ,cross_attention_norm='''group_norm'''      ,resnet_time_scale_shift='''scale_shift'''      ,act_fn='''gelu'''      ,class_embed_type='''timestep'''      ,mid_block_scale_factor=1.4_1_4      ,time_embedding_act_fn='''gelu'''      ,time_embedding_dim=3_2      ,)
   unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

   torch.manual_seed(0    )
   __lowercase													=				DDPMScheduler(
       num_train_timesteps=1_0_0_0      ,beta_schedule='''squaredcos_cap_v2'''      ,beta_start=0.0_0_0_1      ,beta_end=0.0_2      ,thresholding=_UpperCAmelCase      ,dynamic_thresholding_ratio=0.9_5      ,sample_max_value=1.0      ,prediction_type='''epsilon'''      ,variance_type='''learned_range'''      ,)

   torch.manual_seed(0    )
   __lowercase													=				DDPMScheduler(
       num_train_timesteps=1_0_0_0      ,beta_schedule='''squaredcos_cap_v2'''      ,beta_start=0.0_0_0_1      ,beta_end=0.0_2      ,)

   torch.manual_seed(0    )
   __lowercase													=				IFWatermarker()

   return {
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "unet": unet,
       "scheduler": scheduler,
       "image_noising_scheduler": image_noising_scheduler,
       "watermarker": watermarker,
       "safety_checker": None,
       "feature_extractor": None,
   }






  def       SCREAMING_SNAKE_CASE   (   self     :  Union[str, Any]    ):
   __lowercase													=				self.get_dummy_components()
   __lowercase													=				self.pipeline_class(**_UpperCAmelCase    )
   pipe.to(_UpperCAmelCase    )
   pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

   __lowercase													=				self.get_dummy_inputs(_UpperCAmelCase    )

   __lowercase													=				inputs['''prompt''']
   __lowercase													=				inputs['''generator''']
   __lowercase													=				inputs['''num_inference_steps''']
   __lowercase													=				inputs['''output_type''']

   if "image" in inputs:
    __lowercase													=				inputs['''image''']
   else:
    __lowercase													=				None

   if "mask_image" in inputs:
    __lowercase													=				inputs['''mask_image''']
   else:
    __lowercase													=				None

   if "original_image" in inputs:
    __lowercase													=				inputs['''original_image''']
   else:
    __lowercase													=				None

   __lowercase													=				pipe.encode_prompt(_UpperCAmelCase    )

   # inputs with prompt converted to embeddings
   __lowercase													=				{
       '''prompt_embeds''': prompt_embeds,
       '''negative_prompt_embeds''': negative_prompt_embeds,
       '''generator''': generator,
       '''num_inference_steps''': num_inference_steps,
       '''output_type''': output_type,
   }

   if image is not None:
    __lowercase													=				image

   if mask_image is not None:
    __lowercase													=				mask_image

   if original_image is not None:
    __lowercase													=				original_image

   # set all optional components to None
   for optional_component in pipe._optional_components:
    setattr(_UpperCAmelCase      ,_UpperCAmelCase      ,_UpperCAmelCase    )

   __lowercase													=				pipe(**_UpperCAmelCase    )[0]

   with tempfile.TemporaryDirectory() as tmpdir:
    pipe.save_pretrained(_UpperCAmelCase    )
    __lowercase													=				self.pipeline_class.from_pretrained(_UpperCAmelCase    )
    pipe_loaded.to(_UpperCAmelCase    )
    pipe_loaded.set_progress_bar_config(disable=_UpperCAmelCase    )

   pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

   for optional_component in pipe._optional_components:
    self.assertTrue(
        getattr(_UpperCAmelCase      ,_UpperCAmelCase    ) is None      ,F"`{optional_component}` did not stay set to None after loading."      ,)

   __lowercase													=				self.get_dummy_inputs(_UpperCAmelCase    )

   __lowercase													=				inputs['''generator''']
   __lowercase													=				inputs['''num_inference_steps''']
   __lowercase													=				inputs['''output_type''']

   # inputs with prompt converted to embeddings
   __lowercase													=				{
       '''prompt_embeds''': prompt_embeds,
       '''negative_prompt_embeds''': negative_prompt_embeds,
       '''generator''': generator,
       '''num_inference_steps''': num_inference_steps,
       '''output_type''': output_type,
   }

   if image is not None:
    __lowercase													=				image

   if mask_image is not None:
    __lowercase													=				mask_image

   if original_image is not None:
    __lowercase													=				original_image

   __lowercase													=				pipe_loaded(**_UpperCAmelCase    )[0]

   __lowercase													=				np.abs(to_np(_UpperCAmelCase    ) - to_np(_UpperCAmelCase    )    ).max()
   self.assertLess(_UpperCAmelCase      ,1e-4    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				self.get_dummy_components()
   __lowercase													=				self.pipeline_class(**_UpperCAmelCase    )
   pipe.to(_UpperCAmelCase    )
   pipe.set_progress_bar_config(disable=_UpperCAmelCase    )

   __lowercase													=				self.get_dummy_inputs(_UpperCAmelCase    )
   __lowercase													=				pipe(**_UpperCAmelCase    )[0]

   with tempfile.TemporaryDirectory() as tmpdir:
    pipe.save_pretrained(_UpperCAmelCase    )
    __lowercase													=				self.pipeline_class.from_pretrained(_UpperCAmelCase    )
    pipe_loaded.to(_UpperCAmelCase    )
    pipe_loaded.set_progress_bar_config(disable=_UpperCAmelCase    )

   pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

   __lowercase													=				self.get_dummy_inputs(_UpperCAmelCase    )
   __lowercase													=				pipe_loaded(**_UpperCAmelCase    )[0]

   __lowercase													=				np.abs(to_np(_UpperCAmelCase    ) - to_np(_UpperCAmelCase    )    ).max()
   self.assertLess(_UpperCAmelCase      ,1e-4    )




'''simple docstring'''






import copy
import json
import os
import tempfile

from transformers import is_torch_available

from .test_configuration_utils import config_common_kwargs






class     lowercase_			(lowerCamelCase__   ):




  """simple docstring"""




  def __init__(   self     :  Optional[Any]      ,lowercase__     :  int      ,lowercase__     :  List[str]=None      ,lowercase__     :  Any=True      ,lowercase__     :  Union[str, Any]=None      ,**lowercase__     :  Dict    ):
   __lowercase													=				parent
   __lowercase													=				config_class
   __lowercase													=				has_text_modality
   __lowercase													=				kwargs
   __lowercase													=				common_properties






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )
   __lowercase													=				(
       ['''hidden_size''', '''num_attention_heads''', '''num_hidden_layers''']
       if self.common_properties is None
       else self.common_properties
   )

   # Add common fields for text models
   if self.has_text_modality:
    common_properties.extend(['''vocab_size''']    )

   # Test that config has the common properties as getters
   for prop in common_properties:
    self.parent.assertTrue(hasattr(lowercase__      ,lowercase__    )      ,msg=F"`{prop}` does not exist"    )

   # Test that config has the common properties as setter
   for idx, name in enumerate(lowercase__    ):
    try:
     setattr(lowercase__      ,lowercase__      ,lowercase__    )
     self.parent.assertEqual(
         getattr(lowercase__      ,lowercase__    )      ,lowercase__      ,msg=F"`{name} value {idx} expected, but was {getattr(lowercase__      ,lowercase__    )}"    )
    except NotImplementedError:
     # Some models might not be able to implement setters for common_properties
     # In that case, a NotImplementedError is raised
     pass

        # Test if config class can be called with Config(prop_name=..)
   for idx, name in enumerate(lowercase__    ):
    try:
     __lowercase													=				self.config_class(**{name: idx}    )
     self.parent.assertEqual(
         getattr(lowercase__      ,lowercase__    )      ,lowercase__      ,msg=F"`{name} value {idx} expected, but was {getattr(lowercase__      ,lowercase__    )}"    )
    except NotImplementedError:
     # Some models might not be able to implement setters for common_properties
     # In that case, a NotImplementedError is raised
     pass






  def       SCREAMING_SNAKE_CASE   (   self     :  str    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )
   __lowercase													=				json.loads(config.to_json_string()    )
   for key, value in self.inputs_dict.items():
    self.parent.assertEqual(obj[key]      ,lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  int    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   with tempfile.TemporaryDirectory() as tmpdirname:
    __lowercase													=				os.path.join(lowercase__      ,'''config.json'''    )
    config_first.to_json_file(lowercase__    )
    __lowercase													=				self.config_class.from_json_file(lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   with tempfile.TemporaryDirectory() as tmpdirname:
    config_first.save_pretrained(lowercase__    )
    __lowercase													=				self.config_class.from_pretrained(lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   __lowercase													=				'''test'''
   with tempfile.TemporaryDirectory() as tmpdirname:
    __lowercase													=				os.path.join(lowercase__      ,lowercase__    )
    config_first.save_pretrained(lowercase__    )
    __lowercase													=				self.config_class.from_pretrained(lowercase__      ,subfolder=lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				self.config_class(**self.inputs_dict      ,num_labels=5    )
   self.parent.assertEqual(len(config.idalabel    )      ,5    )
   self.parent.assertEqual(len(config.labelaid    )      ,5    )

   __lowercase													=				3
   self.parent.assertEqual(len(config.idalabel    )      ,3    )
   self.parent.assertEqual(len(config.labelaid    )      ,3    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Tuple    ):
   if self.config_class.is_composition:
    return
   __lowercase													=				self.config_class()
   self.parent.assertIsNotNone(lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				copy.deepcopy(lowercase__    )
   __lowercase													=				self.config_class(**lowercase__    )
   __lowercase													=				[]
   for key, value in config_common_kwargs.items():
    if key == "torch_dtype":
     if not is_torch_available():
      continue
     else:
      import torch

      if config.torch_dtype != torch.floataa:
       wrong_values.append(('''torch_dtype''', config.torch_dtype, torch.floataa)    )
    elif getattr(lowercase__      ,lowercase__    ) != value:
     wrong_values.append((key, getattr(lowercase__      ,lowercase__    ), value)    )

   if len(lowercase__    ) > 0:
    __lowercase													=				'''\n'''.join([F"- {v[0]}: got {v[1]} instead of {v[2]}" for v in wrong_values]    )
    raise ValueError(F"The following keys were not properly set in the config:\n{errors}"    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[int]    ):
   self.create_and_test_config_common_properties()
   self.create_and_test_config_to_json_string()
   self.create_and_test_config_to_json_file()
   self.create_and_test_config_from_and_save_pretrained()
   self.create_and_test_config_from_and_save_pretrained_subfolder()
   self.create_and_test_config_with_num_labels()
   self.check_config_can_be_init_without_params()
   self.check_config_arguments_init()

