






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 	_UpperCAmelCase  :




      """simple docstring"""



      def 			snake_case_      (   self		):



          '''simple docstring'''
          torch.manual_seed(0		)
          lowerCAmelCase__   :List[str]      							=				TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"		)

          torch.manual_seed(0		)
          lowerCAmelCase__   :List[Any]      							=				AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"		)

          torch.manual_seed(0		)
          lowerCAmelCase__   :Tuple      							=				UNetaDConditionModel(
              sample_size=32    ,	layers_per_block=1    ,	block_out_channels=[32, 64]    ,	down_block_types=[
                  "ResnetDownsampleBlock2D",
                  "SimpleCrossAttnDownBlock2D",
              ]    ,	mid_block_type="UNetMidBlock2DSimpleCrossAttn"    ,	up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]    ,	in_channels=3    ,	out_channels=6    ,	cross_attention_dim=32    ,	encoder_hid_dim=32    ,	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		)
          lowerCAmelCase__   :Dict      							=				DDPMScheduler(
              num_train_timesteps=1_000    ,	beta_schedule="squaredcos_cap_v2"    ,	beta_start=0.0001    ,	beta_end=0.02    ,	thresholding=_lowerCAmelCase    ,	dynamic_thresholding_ratio=0.95    ,	sample_max_value=1.0    ,	prediction_type="epsilon"    ,	variance_type="learned_range"    ,	)

          torch.manual_seed(0		)
          lowerCAmelCase__   :List[str]      							=				IFWatermarker()

          return {
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "unet": unet,
              "scheduler": scheduler,
              "watermarker": watermarker,
              "safety_checker": None,
              "feature_extractor": None,
          }
      def 			snake_case_      (   self		):



          '''simple docstring'''
          torch.manual_seed(0		)
          lowerCAmelCase__   :int      							=				TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"		)

          torch.manual_seed(0		)
          lowerCAmelCase__   :Dict      							=				AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"		)

          torch.manual_seed(0		)
          lowerCAmelCase__   :Union[str, Any]      							=				UNetaDConditionModel(
              sample_size=32    ,	layers_per_block=[1, 2]    ,	block_out_channels=[32, 64]    ,	down_block_types=[
                  "ResnetDownsampleBlock2D",
                  "SimpleCrossAttnDownBlock2D",
              ]    ,	mid_block_type="UNetMidBlock2DSimpleCrossAttn"    ,	up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]    ,	in_channels=6    ,	out_channels=6    ,	cross_attention_dim=32    ,	encoder_hid_dim=32    ,	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.414    ,	time_embedding_act_fn="gelu"    ,	time_embedding_dim=32    ,	)
          unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

          torch.manual_seed(0		)
          lowerCAmelCase__   :Tuple      							=				DDPMScheduler(
              num_train_timesteps=1_000    ,	beta_schedule="squaredcos_cap_v2"    ,	beta_start=0.0001    ,	beta_end=0.02    ,	thresholding=_lowerCAmelCase    ,	dynamic_thresholding_ratio=0.95    ,	sample_max_value=1.0    ,	prediction_type="epsilon"    ,	variance_type="learned_range"    ,	)

          torch.manual_seed(0		)
          lowerCAmelCase__   :str      							=				DDPMScheduler(
              num_train_timesteps=1_000    ,	beta_schedule="squaredcos_cap_v2"    ,	beta_start=0.0001    ,	beta_end=0.02    ,	)

          torch.manual_seed(0		)
          lowerCAmelCase__   :List[str]      							=				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 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				self.get_dummy_components()
          lowerCAmelCase__   :Any      							=				self.pipeline_class(**_lowerCAmelCase		)
          pipe.to(_lowerCAmelCase		)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

          lowerCAmelCase__   :int      							=				self.get_dummy_inputs(_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[int]      							=				inputs["prompt"]
          lowerCAmelCase__   :Tuple      							=				inputs["generator"]
          lowerCAmelCase__   :List[Any]      							=				inputs["num_inference_steps"]
          lowerCAmelCase__   :Tuple      							=				inputs["output_type"]

          if "image" in inputs:
              lowerCAmelCase__   :Dict      							=				inputs["image"]
          else:
              lowerCAmelCase__   :Union[str, Any]      							=				None

          if "mask_image" in inputs:
              lowerCAmelCase__   :int      							=				inputs["mask_image"]
          else:
              lowerCAmelCase__   :str      							=				None

          if "original_image" in inputs:
              lowerCAmelCase__   :List[Any]      							=				inputs["original_image"]
          else:
              lowerCAmelCase__   :Dict      							=				None

          lowerCAmelCase__   :int      							=				pipe.encode_prompt(_lowerCAmelCase		)

          # inputs with prompt converted to embeddings
          lowerCAmelCase__   :Optional[Any]      							=				{
              "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:
              lowerCAmelCase__   :List[Any]      							=				image

          if mask_image is not None:
              lowerCAmelCase__   :Any      							=				mask_image

          if original_image is not None:
              lowerCAmelCase__   :Tuple      							=				original_image

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

          lowerCAmelCase__   :Optional[int]      							=				pipe(**_lowerCAmelCase		)[0]

          with tempfile.TemporaryDirectory() as tmpdir:
              pipe.save_pretrained(_lowerCAmelCase		)
              lowerCAmelCase__   :str      							=				self.pipeline_class.from_pretrained(_lowerCAmelCase		)
              pipe_loaded.to(_lowerCAmelCase		)
              pipe_loaded.set_progress_bar_config(disable=_lowerCAmelCase		)

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

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

          lowerCAmelCase__   :Any      							=				self.get_dummy_inputs(_lowerCAmelCase		)

          lowerCAmelCase__   :Tuple      							=				inputs["generator"]
          lowerCAmelCase__   :Union[str, Any]      							=				inputs["num_inference_steps"]
          lowerCAmelCase__   :str      							=				inputs["output_type"]

          # inputs with prompt converted to embeddings
          lowerCAmelCase__   :Union[str, Any]      							=				{
              "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:
              lowerCAmelCase__   :Tuple      							=				image

          if mask_image is not None:
              lowerCAmelCase__   :str      							=				mask_image

          if original_image is not None:
              lowerCAmelCase__   :Union[str, Any]      							=				original_image

          lowerCAmelCase__   :List[Any]      							=				pipe_loaded(**_lowerCAmelCase		)[0]

          lowerCAmelCase__   :Dict      							=				np.abs(to_np(_lowerCAmelCase		) - to_np(_lowerCAmelCase		)		).max()
          self.assertLess(_lowerCAmelCase    ,	1e-4		)



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				self.get_dummy_components()
          lowerCAmelCase__   :Optional[int]      							=				self.pipeline_class(**_lowerCAmelCase		)
          pipe.to(_lowerCAmelCase		)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[int]      							=				self.get_dummy_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				pipe(**_lowerCAmelCase		)[0]

          with tempfile.TemporaryDirectory() as tmpdir:
              pipe.save_pretrained(_lowerCAmelCase		)
              lowerCAmelCase__   :Dict      							=				self.pipeline_class.from_pretrained(_lowerCAmelCase		)
              pipe_loaded.to(_lowerCAmelCase		)
              pipe_loaded.set_progress_bar_config(disable=_lowerCAmelCase		)

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

          lowerCAmelCase__   :Optional[int]      							=				self.get_dummy_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				pipe_loaded(**_lowerCAmelCase		)[0]

          lowerCAmelCase__   :str      							=				np.abs(to_np(_lowerCAmelCase		) - to_np(_lowerCAmelCase		)		).max()
          self.assertLess(_lowerCAmelCase    ,	1e-4		)











from ...configuration_utils import PretrainedConfig
from ...utils import logging


_a    :   Dict        =			logging.get_logger(__name__)

_a    :   Optional[Any]        =			{
    """caidas/swin2sr-classicalsr-x2-64""": (
        """https://huggingface.co/caidas/swin2sr-classicalsr-x2-64/resolve/main/config.json"""
    ),
}







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    '''swin2sr'''

      A          =    {
          '''hidden_size''': '''embed_dim''',
          '''num_attention_heads''': '''num_heads''',
          '''num_hidden_layers''': '''num_layers''',
      }
      def __init__(   self    ,	_lowerCAmelCase=64    ,	_lowerCAmelCase=1    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=180    ,	_lowerCAmelCase=[6, 6, 6, 6, 6, 6]    ,	_lowerCAmelCase=[6, 6, 6, 6, 6, 6]    ,	_lowerCAmelCase=8    ,	_lowerCAmelCase=2.0    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=0.0    ,	_lowerCAmelCase=0.0    ,	_lowerCAmelCase=0.1    ,	_lowerCAmelCase="gelu"    ,	_lowerCAmelCase=False    ,	_lowerCAmelCase=0.02    ,	_lowerCAmelCase=1e-5    ,	_lowerCAmelCase=2    ,	_lowerCAmelCase=1.0    ,	_lowerCAmelCase="1conv"    ,	_lowerCAmelCase="pixelshuffle"    ,	**_lowerCAmelCase    ,	):



          '''simple docstring'''
          super().__init__(**_lowerCAmelCase		)

          lowerCAmelCase__   :Union[str, Any]      							=				image_size
          lowerCAmelCase__   :List[Any]      							=				patch_size
          lowerCAmelCase__   :Any      							=				num_channels
          lowerCAmelCase__   :Union[str, Any]      							=				embed_dim
          lowerCAmelCase__   :Optional[int]      							=				depths
          lowerCAmelCase__   :int      							=				len(_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				num_heads
          lowerCAmelCase__   :str      							=				window_size
          lowerCAmelCase__   :List[str]      							=				mlp_ratio
          lowerCAmelCase__   :List[Any]      							=				qkv_bias
          lowerCAmelCase__   :Union[str, Any]      							=				hidden_dropout_prob
          lowerCAmelCase__   :Optional[int]      							=				attention_probs_dropout_prob
          lowerCAmelCase__   :Dict      							=				drop_path_rate
          lowerCAmelCase__   :Tuple      							=				hidden_act
          lowerCAmelCase__   :Dict      							=				use_absolute_embeddings
          lowerCAmelCase__   :Tuple      							=				layer_norm_eps
          lowerCAmelCase__   :Dict      							=				initializer_range
          lowerCAmelCase__   :Optional[int]      							=				upscale
          lowerCAmelCase__   :Optional[Any]      							=				img_range
          lowerCAmelCase__   :List[str]      							=				resi_connection
          lowerCAmelCase__   :Union[str, Any]      							=				upsampler



