
"""simple docstring"""





import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor


class A_	(unittest.TestCase      ):




  '''simple docstring'''







  def __init__(		self				,   lowercase_				,   lowercase_=3				,   lowercase_=32				,   lowercase_=3				,   lowercase_=10				,   lowercase_=[10, 20, 30, 40]				,   lowercase_=[1, 1, 2, 1]				,   lowercase_=True				,   lowercase_=True				,   lowercase_="relu"				,   lowercase_=3				,   lowercase_=None				,   ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    parent
     UpperCAmelCase_					:					Tuple			     =    batch_size
     UpperCAmelCase_					:					List[Any]			     =    image_size
     UpperCAmelCase_					:					Optional[int]			     =    num_channels
     UpperCAmelCase_					:					Union[str, Any]			     =    embeddings_size
     UpperCAmelCase_					:					str			     =    hidden_sizes
     UpperCAmelCase_					:					List[Any]			     =    depths
     UpperCAmelCase_					:					List[str]			     =    is_training
     UpperCAmelCase_					:					Optional[int]			     =    use_labels
     UpperCAmelCase_					:					List[str]			     =    hidden_act
     UpperCAmelCase_					:					Tuple			     =    num_labels
     UpperCAmelCase_					:					Dict			     =    scope
     UpperCAmelCase_					:					Tuple			     =    len(lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					int			     =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

     UpperCAmelCase_					:					List[str]			     =    self.get_config()

     return config, pixel_values







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return RegNetConfig(
         num_channels=self.num_channels				,   embeddings_size=self.embeddings_size				,   hidden_sizes=self.hidden_sizes				,   depths=self.depths				,   hidden_act=self.hidden_act				,   num_labels=self.num_labels				,   image_size=self.image_size				,   )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    FlaxRegNetModel(config=lowercase_  )
     UpperCAmelCase_					:					List[str]			     =    model(lowercase_  )

     # Output shape (b, c, h, w)
     self.parent.assertEqual(
         result.last_hidden_state.shape				,   (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)				,   )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					int			     =    self.num_labels
     UpperCAmelCase_					:					str			     =    FlaxRegNetForImageClassification(config=lowercase_  )
     UpperCAmelCase_					:					Tuple			     =    model(lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.num_labels)  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					str			     =    self.prepare_config_and_inputs()
     UpperCAmelCase_     ,   UpperCAmelCase_					:					List[Any]			     =    config_and_inputs
     UpperCAmelCase_					:					List[Any]			     =    {"pixel_values": pixel_values}
     return config, inputs_dict




@require_flax
class A_	(lowercase__				,unittest.TestCase      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Any       =		(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

  SCREAMING_SNAKE_CASE__			:						Any       =		False
  SCREAMING_SNAKE_CASE__			:						int       =		False
  SCREAMING_SNAKE_CASE__			:						Tuple       =		False







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    FlaxRegNetModelTester(self  )
     UpperCAmelCase_					:					str			     =    ConfigTester(self				,   config_class=lowercase_				,   has_text_modality=lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     self.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Dict			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_  )







  @unittest.skip(reason="RegNet does not use inputs_embeds"  )
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     pass







  @unittest.skip(reason="RegNet does not support input and output embeddings"  )
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     pass







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					List[str]			     =    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        UpperCAmelCase_					:					Union[str, Any]			     =    model_class(lowercase_  )
        UpperCAmelCase_					:					str			     =    inspect.signature(model.__call__  )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        UpperCAmelCase_					:					Tuple			     =    [*signature.parameters.keys()]

        UpperCAmelCase_					:					List[Any]			     =    ["pixel_values"]
        self.assertListEqual(arg_names[:1]				,   lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     def check_hidden_states_output(lowercase_				,   lowercase_				,   lowercase_  ):
        UpperCAmelCase_					:					Optional[int]			     =    model_class(lowercase_  )

        UpperCAmelCase_					:					List[str]			     =    model(**self._prepare_for_class(lowercase_				,   lowercase_  )  )

        UpperCAmelCase_					:					Tuple			     =    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

        UpperCAmelCase_					:					int			     =    self.model_tester.num_stages
        self.assertEqual(len(lowercase_  )				,   expected_num_stages + 1  )

     UpperCAmelCase_     ,   UpperCAmelCase_					:					Optional[int]			     =    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        UpperCAmelCase_					:					int			     =    True
        check_hidden_states_output(lowercase_				,   lowercase_				,   lowercase_  )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        UpperCAmelCase_					:					Optional[Any]			     =    True

        check_hidden_states_output(lowercase_				,   lowercase_				,   lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					List[Any]			     =    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__  ):
           UpperCAmelCase_					:					Optional[int]			     =    self._prepare_for_class(lowercase_				,   lowercase_  )
           UpperCAmelCase_					:					Dict			     =    model_class(lowercase_  )

           @jax.jit
           def model_jitted(lowercase_				,   **lowercase_  ):
              return model(pixel_values=lowercase_				,   **lowercase_  )

           with self.subTest("JIT Enabled"  ):
              UpperCAmelCase_					:					Dict			     =    model_jitted(**lowercase_  ).to_tuple()

           with self.subTest("JIT Disabled"  ):
              with jax.disable_jit():
                 UpperCAmelCase_					:					int			     =    model_jitted(**lowercase_  ).to_tuple()

           self.assertEqual(len(lowercase_  )				,   len(lowercase_  )  )
           for jitted_output, output in zip(lowercase_				,   lowercase_  ):
              self.assertEqual(jitted_output.shape				,   output.shape  )







def 	__a     (		):
   UpperCAmelCase_					:					Optional[int]			     =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
   return image


@require_flax
class A_	(unittest.TestCase      ):




  '''simple docstring'''







  @cached_property
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return AutoImageProcessor.from_pretrained("facebook/regnet-y-040"  ) if is_vision_available() else None







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					str			     =    FlaxRegNetForImageClassification.from_pretrained("facebook/regnet-y-040"  )

     UpperCAmelCase_					:					str			     =    self.default_image_processor
     UpperCAmelCase_					:					Optional[Any]			     =    prepare_img()
     UpperCAmelCase_					:					Union[str, Any]			     =    image_processor(images=lowercase_				,   return_tensors="np"  )

     UpperCAmelCase_					:					str			     =    model(**lowercase_  )

     # verify the logits
     UpperCAmelCase_					:					Tuple			     =    (1, 1000)
     self.assertEqual(outputs.logits.shape				,   lowercase_  )

     UpperCAmelCase_					:					List[str]			     =    jnp.array([-0.41_80, -1.50_51, -3.48_36]  )

     self.assertTrue(jnp.allclose(outputs.logits[0, :3]				,   lowercase_				,   atol=1E-4  )  )




"""simple docstring"""





import argparse
import io

import requests
import torch
from omegaconf import OmegaConf

from diffusers import AutoencoderKL
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    assign_to_checkpoint,
    conv_attn_to_linear,
    create_vae_diffusers_config,
    renew_vae_attention_paths,
    renew_vae_resnet_paths,
)







def 	__a     (		__lowerCamelCase,		__lowerCamelCase  ):
   UpperCAmelCase_					:					str			     =    checkpoint

   UpperCAmelCase_					:					int			     =    {}

   UpperCAmelCase_					:					Optional[int]			     =    vae_state_dict["encoder.conv_in.weight"]
   UpperCAmelCase_					:					List[str]			     =    vae_state_dict["encoder.conv_in.bias"]
   UpperCAmelCase_					:					Optional[Any]			     =    vae_state_dict["encoder.conv_out.weight"]
   UpperCAmelCase_					:					Optional[int]			     =    vae_state_dict["encoder.conv_out.bias"]
   UpperCAmelCase_					:					List[Any]			     =    vae_state_dict["encoder.norm_out.weight"]
   UpperCAmelCase_					:					Union[str, Any]			     =    vae_state_dict["encoder.norm_out.bias"]

   UpperCAmelCase_					:					Any			     =    vae_state_dict["decoder.conv_in.weight"]
   UpperCAmelCase_					:					int			     =    vae_state_dict["decoder.conv_in.bias"]
   UpperCAmelCase_					:					Any			     =    vae_state_dict["decoder.conv_out.weight"]
   UpperCAmelCase_					:					Tuple			     =    vae_state_dict["decoder.conv_out.bias"]
   UpperCAmelCase_					:					List[Any]			     =    vae_state_dict["decoder.norm_out.weight"]
   UpperCAmelCase_					:					Tuple			     =    vae_state_dict["decoder.norm_out.bias"]

   UpperCAmelCase_					:					str			     =    vae_state_dict["quant_conv.weight"]
   UpperCAmelCase_					:					Optional[Any]			     =    vae_state_dict["quant_conv.bias"]
   UpperCAmelCase_					:					List[str]			     =    vae_state_dict["post_quant_conv.weight"]
   UpperCAmelCase_					:					List[Any]			     =    vae_state_dict["post_quant_conv.bias"]

   # Retrieves the keys for the encoder down blocks only
   UpperCAmelCase_					:					Optional[Any]			     =    len({".".join(layer.split("."  )[:3]  ) for layer in vae_state_dict if "encoder.down" in layer}  )
   UpperCAmelCase_					:					Optional[Any]			     =    {
       layer_id: [key for key in vae_state_dict if f"""down.{layer_id}""" in key] for layer_id in range(__lowerCamelCase  )
   }

   # Retrieves the keys for the decoder up blocks only
   UpperCAmelCase_					:					Dict			     =    len({".".join(layer.split("."  )[:3]  ) for layer in vae_state_dict if "decoder.up" in layer}  )
   UpperCAmelCase_					:					Optional[int]			     =    {
       layer_id: [key for key in vae_state_dict if f"""up.{layer_id}""" in key] for layer_id in range(__lowerCamelCase  )
   }

   for i in range(__lowerCamelCase  ):
      UpperCAmelCase_					:					Any			     =    [key for key in down_blocks[i] if f"""down.{i}""" in key and f"""down.{i}.downsample""" not in key]

      if f"""encoder.down.{i}.downsample.conv.weight""" in vae_state_dict:
         UpperCAmelCase_					:					Dict			     =    vae_state_dict.pop(
             f"""encoder.down.{i}.downsample.conv.weight"""  )
         UpperCAmelCase_					:					Dict			     =    vae_state_dict.pop(
             f"""encoder.down.{i}.downsample.conv.bias"""  )

      UpperCAmelCase_					:					List[str]			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					Optional[Any]			     =    {"old": f"""down.{i}.block""", "new": f"""down_blocks.{i}.resnets"""}
      assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )

   UpperCAmelCase_					:					List[Any]			     =    [key for key in vae_state_dict if "encoder.mid.block" in key]
   UpperCAmelCase_					:					Tuple			     =    2
   for i in range(1,		num_mid_res_blocks + 1  ):
      UpperCAmelCase_					:					str			     =    [key for key in mid_resnets if f"""encoder.mid.block_{i}""" in key]

      UpperCAmelCase_					:					List[Any]			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					Optional[int]			     =    {"old": f"""mid.block_{i}""", "new": f"""mid_block.resnets.{i - 1}"""}
      assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )

   UpperCAmelCase_					:					Optional[int]			     =    [key for key in vae_state_dict if "encoder.mid.attn" in key]
   UpperCAmelCase_					:					Union[str, Any]			     =    renew_vae_attention_paths(__lowerCamelCase  )
   UpperCAmelCase_					:					int			     =    {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
   assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )
   conv_attn_to_linear(__lowerCamelCase  )

   for i in range(__lowerCamelCase  ):
      UpperCAmelCase_					:					Optional[Any]			     =    num_up_blocks - 1 - i
      UpperCAmelCase_					:					Any			     =    [
          key for key in up_blocks[block_id] if f"""up.{block_id}""" in key and f"""up.{block_id}.upsample""" not in key
      ]

      if f"""decoder.up.{block_id}.upsample.conv.weight""" in vae_state_dict:
         UpperCAmelCase_					:					str			     =    vae_state_dict[
             f"""decoder.up.{block_id}.upsample.conv.weight"""
         ]
         UpperCAmelCase_					:					Optional[Any]			     =    vae_state_dict[
             f"""decoder.up.{block_id}.upsample.conv.bias"""
         ]

      UpperCAmelCase_					:					Dict			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					List[str]			     =    {"old": f"""up.{block_id}.block""", "new": f"""up_blocks.{i}.resnets"""}
      assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )

   UpperCAmelCase_					:					Optional[Any]			     =    [key for key in vae_state_dict if "decoder.mid.block" in key]
   UpperCAmelCase_					:					List[Any]			     =    2
   for i in range(1,		num_mid_res_blocks + 1  ):
      UpperCAmelCase_					:					str			     =    [key for key in mid_resnets if f"""decoder.mid.block_{i}""" in key]

      UpperCAmelCase_					:					Tuple			     =    renew_vae_resnet_paths(__lowerCamelCase  )
      UpperCAmelCase_					:					Tuple			     =    {"old": f"""mid.block_{i}""", "new": f"""mid_block.resnets.{i - 1}"""}
      assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )

   UpperCAmelCase_					:					Optional[int]			     =    [key for key in vae_state_dict if "decoder.mid.attn" in key]
   UpperCAmelCase_					:					Any			     =    renew_vae_attention_paths(__lowerCamelCase  )
   UpperCAmelCase_					:					Union[str, Any]			     =    {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
   assign_to_checkpoint(__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		additional_replacements=[meta_path],		config=__lowerCamelCase  )
   conv_attn_to_linear(__lowerCamelCase  )
   return new_checkpoint







def 	__a     (		__lowerCamelCase,		__lowerCamelCase,		):
   # Only support V1
   UpperCAmelCase_					:					List[str]			     =    requests.get(
       " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"  )
   UpperCAmelCase_					:					List[Any]			     =    io.BytesIO(r.content  )

   UpperCAmelCase_					:					Any			     =    OmegaConf.load(__lowerCamelCase  )
   UpperCAmelCase_					:					Union[str, Any]			     =    512
   UpperCAmelCase_					:					Optional[Any]			     =    "cuda" if torch.cuda.is_available() else "cpu"
   if checkpoint_path.endswith("safetensors"  ):
      from safetensors import safe_open

      UpperCAmelCase_					:					int			     =    {}
      with safe_open(__lowerCamelCase,		framework="pt",		device="cpu"  ) as f:
         for key in f.keys():
            UpperCAmelCase_					:					Tuple			     =    f.get_tensor(__lowerCamelCase  )
   else:
      UpperCAmelCase_					:					Any			     =    torch.load(__lowerCamelCase,		map_location=__lowerCamelCase  )["state_dict"]

   # Convert the VAE model.
   UpperCAmelCase_					:					Dict			     =    create_vae_diffusers_config(__lowerCamelCase,		image_size=__lowerCamelCase  )
   UpperCAmelCase_					:					Union[str, Any]			     =    custom_convert_ldm_vae_checkpoint(__lowerCamelCase,		__lowerCamelCase  )

   UpperCAmelCase_					:					int			     =    AutoencoderKL(**__lowerCamelCase  )
   vae.load_state_dict(__lowerCamelCase  )
   vae.save_pretrained(__lowerCamelCase  )


if __name__ == "__main__":
    _a		       =							argparse.ArgumentParser()

    parser.add_argument('--vae_pt_path', default=None, type=str, required=True, help='Path to the VAE.pt to convert.')
    parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the VAE.pt to convert.')

    _a		       =							parser.parse_args()

    vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)


