




'''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        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					str:
     """simple docstring"""




     __a       =   checkpoint

     __a       =   {}

     __a       =   vae_state_dict["""encoder.conv_in.weight"""]
     __a       =   vae_state_dict["""encoder.conv_in.bias"""]
     __a       =   vae_state_dict["""encoder.conv_out.weight"""]
     __a       =   vae_state_dict["""encoder.conv_out.bias"""]
     __a       =   vae_state_dict["""encoder.norm_out.weight"""]
     __a       =   vae_state_dict["""encoder.norm_out.bias"""]

     __a       =   vae_state_dict["""decoder.conv_in.weight"""]
     __a       =   vae_state_dict["""decoder.conv_in.bias"""]
     __a       =   vae_state_dict["""decoder.conv_out.weight"""]
     __a       =   vae_state_dict["""decoder.conv_out.bias"""]
     __a       =   vae_state_dict["""decoder.norm_out.weight"""]
     __a       =   vae_state_dict["""decoder.norm_out.bias"""]

     __a       =   vae_state_dict["""quant_conv.weight"""]
     __a       =   vae_state_dict["""quant_conv.bias"""]
     __a       =   vae_state_dict["""post_quant_conv.weight"""]
     __a       =   vae_state_dict["""post_quant_conv.bias"""]

     # Retrieves the keys for the encoder down blocks only
     __a       =   len({""".""".join(layer.split("""."""		)[:3]		) for layer in vae_state_dict if """encoder.down""" in layer}		)
     __a       =   {
         layer_id: [key for key in vae_state_dict if F'''down.{layer_id}''' in key] for layer_id in range(__SCREAMING_SNAKE_CASE		)
     }

     # Retrieves the keys for the decoder up blocks only
     __a       =   len({""".""".join(layer.split("""."""		)[:3]		) for layer in vae_state_dict if """decoder.up""" in layer}		)
     __a       =   {
         layer_id: [key for key in vae_state_dict if F'''up.{layer_id}''' in key] for layer_id in range(__SCREAMING_SNAKE_CASE		)
     }

     for i in range(__SCREAMING_SNAKE_CASE		):
          __a       =   [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:
               __a       =   vae_state_dict.pop(
                   F'''encoder.down.{i}.downsample.conv.weight'''		)
               __a       =   vae_state_dict.pop(
                   F'''encoder.down.{i}.downsample.conv.bias'''		)

          __a       =   renew_vae_resnet_paths(__SCREAMING_SNAKE_CASE		)
          __a       =   {"""old""": F'''down.{i}.block''', """new""": F'''down_blocks.{i}.resnets'''}
          assign_to_checkpoint(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  additional_replacements=[meta_path]				,  config=__SCREAMING_SNAKE_CASE		)

     __a       =   [key for key in vae_state_dict if """encoder.mid.block""" in key]
     __a       =   2
     for i in range(1				,  num_mid_res_blocks + 1		):
          __a       =   [key for key in mid_resnets if F'''encoder.mid.block_{i}''' in key]

          __a       =   renew_vae_resnet_paths(__SCREAMING_SNAKE_CASE		)
          __a       =   {"""old""": F'''mid.block_{i}''', """new""": F'''mid_block.resnets.{i - 1}'''}
          assign_to_checkpoint(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  additional_replacements=[meta_path]				,  config=__SCREAMING_SNAKE_CASE		)

     __a       =   [key for key in vae_state_dict if """encoder.mid.attn""" in key]
     __a       =   renew_vae_attention_paths(__SCREAMING_SNAKE_CASE		)
     __a       =   {"""old""": """mid.attn_1""", """new""": """mid_block.attentions.0"""}
     assign_to_checkpoint(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  additional_replacements=[meta_path]				,  config=__SCREAMING_SNAKE_CASE		)
     conv_attn_to_linear(__SCREAMING_SNAKE_CASE		)

     for i in range(__SCREAMING_SNAKE_CASE		):
          __a       =   num_up_blocks - 1 - i
          __a       =   [
              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:
               __a       =   vae_state_dict[
                   F'''decoder.up.{block_id}.upsample.conv.weight'''
               ]
               __a       =   vae_state_dict[
                   F'''decoder.up.{block_id}.upsample.conv.bias'''
               ]

          __a       =   renew_vae_resnet_paths(__SCREAMING_SNAKE_CASE		)
          __a       =   {"""old""": F'''up.{block_id}.block''', """new""": F'''up_blocks.{i}.resnets'''}
          assign_to_checkpoint(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  additional_replacements=[meta_path]				,  config=__SCREAMING_SNAKE_CASE		)

     __a       =   [key for key in vae_state_dict if """decoder.mid.block""" in key]
     __a       =   2
     for i in range(1				,  num_mid_res_blocks + 1		):
          __a       =   [key for key in mid_resnets if F'''decoder.mid.block_{i}''' in key]

          __a       =   renew_vae_resnet_paths(__SCREAMING_SNAKE_CASE		)
          __a       =   {"""old""": F'''mid.block_{i}''', """new""": F'''mid_block.resnets.{i - 1}'''}
          assign_to_checkpoint(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  additional_replacements=[meta_path]				,  config=__SCREAMING_SNAKE_CASE		)

     __a       =   [key for key in vae_state_dict if """decoder.mid.attn""" in key]
     __a       =   renew_vae_attention_paths(__SCREAMING_SNAKE_CASE		)
     __a       =   {"""old""": """mid.attn_1""", """new""": """mid_block.attentions.0"""}
     assign_to_checkpoint(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  additional_replacements=[meta_path]				,  config=__SCREAMING_SNAKE_CASE		)
     conv_attn_to_linear(__SCREAMING_SNAKE_CASE		)
     return new_checkpoint




def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  )				->					Union[str, Any]:
     """simple docstring"""




     __a       =   requests.get(
         """ https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"""		)
     __a       =   io.BytesIO(r.content		)

     __a       =   OmegaConf.load(__SCREAMING_SNAKE_CASE		)
     __a       =   512
     __a       =   """cuda""" if torch.cuda.is_available() else """cpu"""
     if checkpoint_path.endswith("""safetensors"""		):
          from safetensors import safe_open

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

     # Convert the VAE model.
     __a       =   create_vae_diffusers_config(__SCREAMING_SNAKE_CASE				,  image_size=__SCREAMING_SNAKE_CASE		)
     __a       =   custom_convert_ldm_vae_checkpoint(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)

     __a       =   AutoencoderKL(**__SCREAMING_SNAKE_CASE		)
     vae.load_state_dict(__SCREAMING_SNAKE_CASE		)
     vae.save_pretrained(__SCREAMING_SNAKE_CASE		)


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE_      =			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.')

 SCREAMING_SNAKE_CASE_      =			parser.parse_args()

 vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)










'''simple docstring'''







import unittest

from transformers import (
    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
    TextClassificationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, nested_simplify, require_tf, require_torch, slow

from .test_pipelines_common import ANY


# These 2 model types require different inputs than those of the usual text models.
SCREAMING_SNAKE_CASE_      =			{'LayoutLMv2Config', 'LayoutLMv3Config'}


@is_pipeline_test
class  lowerCAmelCase_						(       unittest.TestCase      ):





  """simple docstring"""



  a_       :List[Any]												=MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
  a_       :List[Any]												=TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING

  if model_mapping is not None:
    a_       :List[Any]												={config: model for config, model in model_mapping.items() if config.__name__ not in _TO_SKIP}
  if tf_model_mapping is not None:
    a_       :Tuple												={
        config: model for config, model in tf_model_mapping.items() if config.__name__ not in _TO_SKIP
    }

  @require_torch
  def __a						(					self     :							int						):





       '''simple docstring'''

       __a       =   pipeline(
           task="""text-classification"""				,							model="""hf-internal-testing/tiny-random-distilbert"""				,							framework="""pt"""						)

       __a       =   text_classifier("""This is great !"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """LABEL_0""", """score""": 0.5_0_4}]						)

       __a       =   text_classifier("""This is great !"""				,							top_k=2						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """LABEL_0""", """score""": 0.5_0_4}, {"""label""": """LABEL_1""", """score""": 0.4_9_6}]						)

       __a       =   text_classifier(["""This is great !""", """This is bad"""]				,							top_k=2						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[
               [{"""label""": """LABEL_0""", """score""": 0.5_0_4}, {"""label""": """LABEL_1""", """score""": 0.4_9_6}],
               [{"""label""": """LABEL_0""", """score""": 0.5_0_4}, {"""label""": """LABEL_1""", """score""": 0.4_9_6}],
           ]				,							)

       __a       =   text_classifier("""This is great !"""				,							top_k=1						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """LABEL_0""", """score""": 0.5_0_4}]						)

       # Legacy behavior
       __a       =   text_classifier("""This is great !"""				,							return_all_scores=SCREAMING_SNAKE_CASE__						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """LABEL_0""", """score""": 0.5_0_4}]						)

       __a       =   text_classifier("""This is great !"""				,							return_all_scores=SCREAMING_SNAKE_CASE__						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[[{"""label""": """LABEL_0""", """score""": 0.5_0_4}, {"""label""": """LABEL_1""", """score""": 0.4_9_6}]]						)

       __a       =   text_classifier(["""This is great !""", """Something else"""]				,							return_all_scores=SCREAMING_SNAKE_CASE__						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[
               [{"""label""": """LABEL_0""", """score""": 0.5_0_4}, {"""label""": """LABEL_1""", """score""": 0.4_9_6}],
               [{"""label""": """LABEL_0""", """score""": 0.5_0_4}, {"""label""": """LABEL_1""", """score""": 0.4_9_6}],
           ]				,							)

       __a       =   text_classifier(["""This is great !""", """Something else"""]				,							return_all_scores=SCREAMING_SNAKE_CASE__						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[
               {"""label""": """LABEL_0""", """score""": 0.5_0_4},
               {"""label""": """LABEL_0""", """score""": 0.5_0_4},
           ]				,							)

  @require_torch
  def __a						(					self     :							Optional[int]						):





       '''simple docstring'''

       import torch

       __a       =   pipeline(
           task="""text-classification"""				,							model="""hf-internal-testing/tiny-random-distilbert"""				,							framework="""pt"""				,							device=torch.device("""cpu"""						)				,							)

       __a       =   text_classifier("""This is great !"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """LABEL_0""", """score""": 0.5_0_4}]						)

  @require_tf
  def __a						(					self     :							List[str]						):





       '''simple docstring'''

       __a       =   pipeline(
           task="""text-classification"""				,							model="""hf-internal-testing/tiny-random-distilbert"""				,							framework="""tf"""						)

       __a       =   text_classifier("""This is great !"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """LABEL_0""", """score""": 0.5_0_4}]						)

  @slow
  @require_torch
  def __a						(					self     :							int						):





       '''simple docstring'''

       __a       =   pipeline("""text-classification"""						)

       __a       =   text_classifier("""This is great !"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """POSITIVE""", """score""": 1.0}]						)
       __a       =   text_classifier("""This is bad !"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """NEGATIVE""", """score""": 1.0}]						)
       __a       =   text_classifier("""Birds are a type of animal"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """POSITIVE""", """score""": 0.9_8_8}]						)

  @slow
  @require_tf
  def __a						(					self     :							Union[str, Any]						):





       '''simple docstring'''

       __a       =   pipeline("""text-classification"""				,							framework="""tf"""						)

       __a       =   text_classifier("""This is great !"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """POSITIVE""", """score""": 1.0}]						)
       __a       =   text_classifier("""This is bad !"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """NEGATIVE""", """score""": 1.0}]						)
       __a       =   text_classifier("""Birds are a type of animal"""						)
       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": """POSITIVE""", """score""": 0.9_8_8}]						)

  def __a						(					self     :							Tuple				,							SCREAMING_SNAKE_CASE__     :							int				,							SCREAMING_SNAKE_CASE__     :							str				,							SCREAMING_SNAKE_CASE__     :							Tuple						):





       '''simple docstring'''

       __a       =   TextClassificationPipeline(model=SCREAMING_SNAKE_CASE__				,							tokenizer=SCREAMING_SNAKE_CASE__						)
       return text_classifier, ["HuggingFace is in", "This is another test"]


  def __a						(					self     :							Any				,							SCREAMING_SNAKE_CASE__     :							Any				,							SCREAMING_SNAKE_CASE__     :							Tuple						):





       '''simple docstring'''

       __a       =   text_classifier.model
       # Small inputs because BartTokenizer tiny has maximum position embeddings = 22
       __a       =   """HuggingFace is in"""
       __a       =   text_classifier(SCREAMING_SNAKE_CASE__						)

       self.assertEqual(nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": ANY(SCREAMING_SNAKE_CASE__						), """score""": ANY(SCREAMING_SNAKE_CASE__						)}]						)
       self.assertTrue(outputs[0]["""label"""] in model.config.idalabel.values()						)

       __a       =   ["""HuggingFace is in """, """Paris is in France"""]
       __a       =   text_classifier(SCREAMING_SNAKE_CASE__						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": ANY(SCREAMING_SNAKE_CASE__						), """score""": ANY(SCREAMING_SNAKE_CASE__						)}, {"""label""": ANY(SCREAMING_SNAKE_CASE__						), """score""": ANY(SCREAMING_SNAKE_CASE__						)}]				,							)
       self.assertTrue(outputs[0]["""label"""] in model.config.idalabel.values()						)
       self.assertTrue(outputs[1]["""label"""] in model.config.idalabel.values()						)

       # Forcing to get all results with `top_k=None`
       # This is NOT the legacy format
       __a       =   text_classifier(SCREAMING_SNAKE_CASE__				,							top_k=SCREAMING_SNAKE_CASE__						)
       __a       =   len(model.config.idalabel.values()						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[[{"""label""": ANY(SCREAMING_SNAKE_CASE__						), """score""": ANY(SCREAMING_SNAKE_CASE__						)}] * N, [{"""label""": ANY(SCREAMING_SNAKE_CASE__						), """score""": ANY(SCREAMING_SNAKE_CASE__						)}] * N]				,							)

       __a       =   {"""text""": """HuggingFace is in """, """text_pair""": """Paris is in France"""}
       __a       =   text_classifier(SCREAMING_SNAKE_CASE__						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							{"""label""": ANY(SCREAMING_SNAKE_CASE__						), """score""": ANY(SCREAMING_SNAKE_CASE__						)}				,							)
       self.assertTrue(outputs["""label"""] in model.config.idalabel.values()						)

       # This might be used a text pair, but tokenizer + pipe interaction
       # makes it hard to understand that it's not using the pair properly
       # https://github.com/huggingface/transformers/issues/17305
       # We disabled this usage instead as it was outputting wrong outputs.
       __a       =   [["""HuggingFace is in """, """Paris is in France"""]]
       with self.assertRaises(SCREAMING_SNAKE_CASE__						):
            text_classifier(SCREAMING_SNAKE_CASE__						)

       # This used to be valid for doing text pairs
       # We're keeping it working because of backward compatibility
       __a       =   text_classifier([[["""HuggingFace is in """, """Paris is in France"""]]]						)
       self.assertEqual(
           nested_simplify(SCREAMING_SNAKE_CASE__						)				,							[{"""label""": ANY(SCREAMING_SNAKE_CASE__						), """score""": ANY(SCREAMING_SNAKE_CASE__						)}]				,							)
       self.assertTrue(outputs[0]["""label"""] in model.config.idalabel.values()						)






