
"""simple docstring"""





import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


_a		       =							datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
    import pyspark


@dataclass
class A_	(datasets.BuilderConfig      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Optional[datasets.Features]       =		None







def 	__a     (		__lowerCamelCase,		__lowerCamelCase,		):
   import pyspark

   def generate_fn():
      UpperCAmelCase_					:					List[str]			     =    df.select("*",		pyspark.sql.functions.spark_partition_id().alias("part_id"  )  )
      for partition_id in partition_order:
         UpperCAmelCase_					:					List[Any]			     =    df_with_partition_id.select("*"  ).where(f"""part_id = {partition_id}"""  ).drop("part_id"  )
         UpperCAmelCase_					:					Tuple			     =    partition_df.collect()
         UpperCAmelCase_					:					List[Any]			     =    0
         for row in rows:
            yield f"""{partition_id}_{row_id}""", row.asDict()
            row_id += 1

   return generate_fn


class A_	(_BaseExamplesIterable      ):




  '''simple docstring'''







  def __init__(		self				,   lowercase_				,   lowercase_=None				,   ):

     """simple docstring"""
     UpperCAmelCase_					:					Any			     =    df
     UpperCAmelCase_					:					Optional[Any]			     =    partition_order or range(self.df.rdd.getNumPartitions()  )
     UpperCAmelCase_					:					str			     =    _generate_iterable_examples(self.df				,   self.partition_order  )







  def __iter__(		self  ):

     """simple docstring"""
     yield from self.generate_examples_fn()







  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    list(range(self.df.rdd.getNumPartitions()  )  )
     generator.shuffle(lowercase_  )
     return SparkExamplesIterable(self.df				,   partition_order=lowercase_  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    self.split_shard_indices_by_worker(lowercase_				,   lowercase_  )
     return SparkExamplesIterable(self.df				,   partition_order=lowercase_  )







  @property
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return len(self.partition_order  )




class A_	(datasets.DatasetBuilder      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						List[str]       =		SparkConfig







  def __init__(		self				,   lowercase_				,   lowercase_ = None				,   lowercase_ = None				,   **lowercase_				,   ):

     """simple docstring"""
     import pyspark

     UpperCAmelCase_					:					Optional[Any]			     =    pyspark.sql.SparkSession.builder.getOrCreate()
     UpperCAmelCase_					:					Optional[Any]			     =    df
     UpperCAmelCase_					:					Tuple			     =    working_dir

     super().__init__(
         cache_dir=lowercase_				,   config_name=str(self.df.semanticHash()  )				,   **lowercase_				,   )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     # Returns the path of the created file.
     def create_cache_and_write_probe(lowercase_  ):
        # makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
        # already exist.
        os.makedirs(self._cache_dir				,   exist_ok=lowercase_  )
        UpperCAmelCase_					:					Union[str, Any]			     =    os.path.join(self._cache_dir				,   "fs_test" + uuid.uuida().hex  )
        # Opening the file in append mode will create a new file unless it already exists, in which case it will not
        # change the file contents.
        open(lowercase_				,   "a"  )
        return [probe_file]

     if self._spark.conf.get("spark.master"				,   ""  ).startswith("local"  ):
        return

     # If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
     # accessible to the driver.
     # TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
     if self._cache_dir:
        UpperCAmelCase_					:					List[Any]			     =    (
            self._spark.sparkContext.parallelize(range(1  )				,   1  ).mapPartitions(lowercase_  ).collect()
        )
        if os.path.isfile(probe[0]  ):
           return

     raise ValueError(
         "When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir"  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return datasets.DatasetInfo(features=self.config.features  )







  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     return [datasets.SplitGenerator(name=datasets.Split.TRAIN  )]







  def        UpperCamelCase__					(		self				,   lowercase_  ):

     """simple docstring"""
     import pyspark

     def get_arrow_batch_size(lowercase_  ):
        for batch in it:
           yield pa.RecordBatch.from_pydict({"batch_bytes": [batch.nbytes]}  )

     UpperCAmelCase_					:					int			     =    self.df.count()
     UpperCAmelCase_					:					Optional[Any]			     =    df_num_rows if df_num_rows <= 100 else 100
     # Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
     UpperCAmelCase_					:					Dict			     =    (
         self.df.limit(lowercase_  )
         .repartition(1  )
         .mapInArrow(lowercase_				,   "batch_bytes: long"  )
         .agg(pyspark.sql.functions.sum("batch_bytes"  ).alias("sample_bytes"  )  )
         .collect()[0]
         .sample_bytes
         / sample_num_rows
     )
     UpperCAmelCase_					:					Tuple			     =    approx_bytes_per_row * df_num_rows
     if approx_total_size > max_shard_size:
        # Make sure there is at least one row per partition.
        UpperCAmelCase_					:					str			     =    min(lowercase_				,   int(approx_total_size / max_shard_size  )  )
        UpperCAmelCase_					:					Dict			     =    self.df.repartition(lowercase_  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   ):

     """simple docstring"""
     import pyspark

     UpperCAmelCase_					:					str			     =    ParquetWriter if file_format == "parquet" else ArrowWriter
     UpperCAmelCase_					:					Optional[Any]			     =    os.path.join(self._working_dir				,   os.path.basename(lowercase_  )  ) if self._working_dir else fpath
     UpperCAmelCase_					:					Union[str, Any]			     =    file_format == "parquet"

     # Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
     # pickling the SparkContext.
     UpperCAmelCase_					:					Union[str, Any]			     =    self.config.features
     UpperCAmelCase_					:					str			     =    self._writer_batch_size
     UpperCAmelCase_					:					str			     =    self._fs.storage_options

     def write_arrow(lowercase_  ):
        # Within the same SparkContext, no two task attempts will share the same attempt ID.
        UpperCAmelCase_					:					Dict			     =    pyspark.TaskContext().taskAttemptId()
        UpperCAmelCase_					:					Union[str, Any]			     =    next(lowercase_				,   lowercase_  )
        if first_batch is None:
           # Some partitions might not receive any data.
           return pa.RecordBatch.from_arrays(
               [[task_id], [0], [0]]				,   names=["task_id", "num_examples", "num_bytes"]				,   )
        UpperCAmelCase_					:					List[Any]			     =    0
        UpperCAmelCase_					:					int			     =    writer_class(
            features=lowercase_				,   path=working_fpath.replace("SSSSS"				,   F"""{shard_id:05d}"""  ).replace("TTTTT"				,   F"""{task_id:05d}"""  )				,   writer_batch_size=lowercase_				,   storage_options=lowercase_				,   embed_local_files=lowercase_				,   )
        UpperCAmelCase_					:					str			     =    pa.Table.from_batches([first_batch]  )
        writer.write_table(lowercase_  )
        for batch in it:
           if max_shard_size is not None and writer._num_bytes >= max_shard_size:
              UpperCAmelCase_     ,   UpperCAmelCase_					:					Optional[Any]			     =    writer.finalize()
              writer.close()
              yield pa.RecordBatch.from_arrays(
                  [[task_id], [num_examples], [num_bytes]]				,   names=["task_id", "num_examples", "num_bytes"]				,   )
              shard_id += 1
              UpperCAmelCase_					:					int			     =    writer_class(
                  features=writer._features				,   path=working_fpath.replace("SSSSS"				,   F"""{shard_id:05d}"""  ).replace("TTTTT"				,   F"""{task_id:05d}"""  )				,   writer_batch_size=lowercase_				,   storage_options=lowercase_				,   embed_local_files=lowercase_				,   )
           UpperCAmelCase_					:					List[str]			     =    pa.Table.from_batches([batch]  )
           writer.write_table(lowercase_  )

        if writer._num_bytes > 0:
           UpperCAmelCase_     ,   UpperCAmelCase_					:					List[Any]			     =    writer.finalize()
           writer.close()
           yield pa.RecordBatch.from_arrays(
               [[task_id], [num_examples], [num_bytes]]				,   names=["task_id", "num_examples", "num_bytes"]				,   )

        if working_fpath != fpath:
           for file in os.listdir(os.path.dirname(lowercase_  )  ):
              UpperCAmelCase_					:					Union[str, Any]			     =    os.path.join(os.path.dirname(lowercase_  )				,   os.path.basename(lowercase_  )  )
              shutil.move(lowercase_				,   lowercase_  )

     UpperCAmelCase_					:					str			     =    (
         self.df.mapInArrow(lowercase_				,   "task_id: long, num_examples: long, num_bytes: long"  )
         .groupBy("task_id"  )
         .agg(
             pyspark.sql.functions.sum("num_examples"  ).alias("total_num_examples"  )				,   pyspark.sql.functions.sum("num_bytes"  ).alias("total_num_bytes"  )				,   pyspark.sql.functions.count("num_bytes"  ).alias("num_shards"  )				,   pyspark.sql.functions.collect_list("num_examples"  ).alias("shard_lengths"  )				,   )
         .collect()
     )
     for row in stats:
        yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_ = "arrow"				,   lowercase_ = None				,   lowercase_ = None				,   **lowercase_				,   ):

     """simple docstring"""
     self._validate_cache_dir()

     UpperCAmelCase_					:					int			     =    convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE  )
     self._repartition_df_if_needed(lowercase_  )
     UpperCAmelCase_					:					int			     =    not is_remote_filesystem(self._fs  )
     UpperCAmelCase_					:					Dict			     =    os.path.join if is_local else posixpath.join

     UpperCAmelCase_					:					Dict			     =    "-TTTTT-SSSSS-of-NNNNN"
     UpperCAmelCase_					:					Union[str, Any]			     =    F"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
     UpperCAmelCase_					:					List[Any]			     =    path_join(self._output_dir				,   lowercase_  )

     UpperCAmelCase_					:					Union[str, Any]			     =    0
     UpperCAmelCase_					:					Dict			     =    0
     UpperCAmelCase_					:					int			     =    0
     UpperCAmelCase_					:					Tuple			     =    []
     UpperCAmelCase_					:					List[Any]			     =    []

     for task_id, content in self._prepare_split_single(lowercase_				,   lowercase_				,   lowercase_  ):
        (
            (
            UpperCAmelCase_
        )     ,   (
            UpperCAmelCase_
        )     ,   (
            UpperCAmelCase_
        )     ,   (
            UpperCAmelCase_
        )     ,   
        )					:					Tuple			     =    content
        if num_bytes > 0:
           total_num_examples += num_examples
           total_num_bytes += num_bytes
           total_shards += num_shards
           task_id_and_num_shards.append((task_id, num_shards)  )
           all_shard_lengths.extend(lowercase_  )

     UpperCAmelCase_					:					Tuple			     =    total_num_examples
     UpperCAmelCase_					:					Union[str, Any]			     =    total_num_bytes

     # should rename everything at the end
     logger.debug(F"""Renaming {total_shards} shards."""  )
     if total_shards > 1:
        UpperCAmelCase_					:					Union[str, Any]			     =    all_shard_lengths

        # Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
        # pickling error due to pickling the SparkContext.
        UpperCAmelCase_					:					Tuple			     =    self._fs

        # use the -SSSSS-of-NNNNN pattern
        def _rename_shard(
            lowercase_				,   lowercase_				,   lowercase_				,   ):
           rename(
               lowercase_				,   fpath.replace("SSSSS"				,   F"""{shard_id:05d}"""  ).replace("TTTTT"				,   F"""{task_id:05d}"""  )				,   fpath.replace("TTTTT-SSSSS"				,   F"""{global_shard_id:05d}"""  ).replace("NNNNN"				,   F"""{total_shards:05d}"""  )				,   )

        UpperCAmelCase_					:					str			     =    []
        UpperCAmelCase_					:					Optional[Any]			     =    0
        for i in range(len(lowercase_  )  ):
           UpperCAmelCase_     ,   UpperCAmelCase_					:					Union[str, Any]			     =    task_id_and_num_shards[i]
           for shard_id in range(lowercase_  ):
              args.append([task_id, shard_id, global_shard_id]  )
              global_shard_id += 1
        self._spark.sparkContext.parallelize(lowercase_				,   len(lowercase_  )  ).map(lambda lowercase_  : _rename_shard(*lowercase_  )  ).collect()
     else:
        # don't use any pattern
        UpperCAmelCase_					:					Union[str, Any]			     =    0
        UpperCAmelCase_					:					Optional[Any]			     =    task_id_and_num_shards[0][0]
        self._rename(
            fpath.replace("SSSSS"				,   F"""{shard_id:05d}"""  ).replace("TTTTT"				,   F"""{task_id:05d}"""  )				,   fpath.replace(lowercase_				,   ""  )				,   )







  def        UpperCamelCase__					(		self				,   lowercase_				,   ):

     """simple docstring"""
     return SparkExamplesIterable(self.df  )




"""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)


