



'''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  snake_case_     (					_lowerCAmelCase  :						str   ,      _lowerCAmelCase  :						Union[str, Any]  )      ->   Any:
							UpperCAmelCase :       int 							=		checkpoint

							UpperCAmelCase :       Any 							=		{}

							UpperCAmelCase :       Optional[Any] 							=		vae_state_dict['''encoder.conv_in.weight''']
							UpperCAmelCase :       Dict 							=		vae_state_dict['''encoder.conv_in.bias''']
							UpperCAmelCase :       Tuple 							=		vae_state_dict['''encoder.conv_out.weight''']
							UpperCAmelCase :       Optional[int] 							=		vae_state_dict['''encoder.conv_out.bias''']
							UpperCAmelCase :       int 							=		vae_state_dict['''encoder.norm_out.weight''']
							UpperCAmelCase :       Union[str, Any] 							=		vae_state_dict['''encoder.norm_out.bias''']

							UpperCAmelCase :       List[Any] 							=		vae_state_dict['''decoder.conv_in.weight''']
							UpperCAmelCase :       List[Any] 							=		vae_state_dict['''decoder.conv_in.bias''']
							UpperCAmelCase :       str 							=		vae_state_dict['''decoder.conv_out.weight''']
							UpperCAmelCase :       Any 							=		vae_state_dict['''decoder.conv_out.bias''']
							UpperCAmelCase :       List[Any] 							=		vae_state_dict['''decoder.norm_out.weight''']
							UpperCAmelCase :       Any 							=		vae_state_dict['''decoder.norm_out.bias''']

							UpperCAmelCase :       int 							=		vae_state_dict['''quant_conv.weight''']
							UpperCAmelCase :       Union[str, Any] 							=		vae_state_dict['''quant_conv.bias''']
							UpperCAmelCase :       Any 							=		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 :       Any 							=		len({'''.'''.join(layer.split('''.'''  )[:3]  ) for layer in vae_state_dict if '''decoder.up''' in layer}  )
							UpperCAmelCase :       List[str] 							=		{
							    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 :       Dict 							=		[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 :       Optional[int] 							=		vae_state_dict.pop(
																					    f"""encoder.down.{i}.downsample.conv.weight"""  )
																					UpperCAmelCase :       List[Any] 							=		vae_state_dict.pop(
																					    f"""encoder.down.{i}.downsample.conv.bias"""  )

														UpperCAmelCase :       Tuple 							=		renew_vae_resnet_paths(_lowerCAmelCase  )
														UpperCAmelCase :       str 							=		{'''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 :       Union[str, Any] 							=		2
							for i in range(1   ,      num_mid_res_blocks + 1  ):
														UpperCAmelCase :       Dict 							=		[key for key in mid_resnets if f"""encoder.mid.block_{i}""" in key]

														UpperCAmelCase :       Tuple 							=		renew_vae_resnet_paths(_lowerCAmelCase  )
														UpperCAmelCase :       Dict 							=		{'''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 :       Any 							=		[key for key in vae_state_dict if '''encoder.mid.attn''' in key]
							UpperCAmelCase :       List[Any] 							=		renew_vae_attention_paths(_lowerCAmelCase  )
							UpperCAmelCase :       Dict 							=		{'''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 :       List[Any] 							=		num_up_blocks - 1 - i
														UpperCAmelCase :       Dict 							=		[
														    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 :       List[Any] 							=		vae_state_dict[
																					    f"""decoder.up.{block_id}.upsample.conv.weight"""
																					]
																					UpperCAmelCase :       List[Any] 							=		vae_state_dict[
																					    f"""decoder.up.{block_id}.upsample.conv.bias"""
																					]

														UpperCAmelCase :       List[Any] 							=		renew_vae_resnet_paths(_lowerCAmelCase  )
														UpperCAmelCase :       int 							=		{'''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 :       Dict 							=		[key for key in vae_state_dict if '''decoder.mid.block''' in key]
							UpperCAmelCase :       Tuple 							=		2
							for i in range(1   ,      num_mid_res_blocks + 1  ):
														UpperCAmelCase :       List[str] 							=		[key for key in mid_resnets if f"""decoder.mid.block_{i}""" in key]

														UpperCAmelCase :       List[str] 							=		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 :       int 							=		[key for key in vae_state_dict if '''decoder.mid.attn''' in key]
							UpperCAmelCase :       Dict 							=		renew_vae_attention_paths(_lowerCAmelCase  )
							UpperCAmelCase :       Tuple 							=		{'''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  snake_case_     (					_lowerCAmelCase  :						str   ,      _lowerCAmelCase  :						str   ,      )      ->   Tuple:
							# Only support V1
							UpperCAmelCase :       str 							=		requests.get(
							    ''' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'''  )
							UpperCAmelCase :       Optional[Any] 							=		io.BytesIO(r.content  )

							UpperCAmelCase :       int 							=		OmegaConf.load(_lowerCAmelCase  )
							UpperCAmelCase :       Optional[int] 							=		512
							UpperCAmelCase :       List[str] 							=		'''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 :       str 							=		f.get_tensor(_lowerCAmelCase  )
							else:
														UpperCAmelCase :       Dict 							=		torch.load(_lowerCAmelCase   ,      map_location=_lowerCAmelCase  )['''state_dict''']

							# Convert the VAE model.
							UpperCAmelCase :       List[Any] 							=		create_vae_diffusers_config(_lowerCAmelCase   ,      image_size=_lowerCAmelCase  )
							UpperCAmelCase :       Tuple 							=		custom_convert_ldm_vae_checkpoint(_lowerCAmelCase   ,      _lowerCAmelCase  )

							UpperCAmelCase :       Tuple 							=		AutoencoderKL(**_lowerCAmelCase  )
							vae.load_state_dict(_lowerCAmelCase  )
							vae.save_pretrained(_lowerCAmelCase  )


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

						UpperCamelCase__:   str								=			parser.parse_args()

						vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)









import os
import re
import shutil
from argparse import ArgumentParser, Namespace

from datasets.commands import BaseDatasetsCLICommand
from datasets.utils.logging import get_logger


lowerCAmelCase       =							'<<<<<<< This should probably be modified because it mentions: '

lowerCAmelCase       =							'=======\n>>>>>>>\n'

lowerCAmelCase       =							[
    'TextEncoderConfig',
    'ByteTextEncoder',
    'SubwordTextEncoder',
    'encoder_config',
    'maybe_build_from_corpus',
    'manual_dir',
]

lowerCAmelCase       =							[
    # (pattern, replacement)
    # Order is important here for some replacements
    (R'tfds\.core', R'datasets'),
    (R'tf\.io\.gfile\.GFile', R'open'),
    (R'tf\.([\w\d]+)', R'datasets.Value(\'\1\')'),
    (R'tfds\.features\.Text\(\)', R'datasets.Value(\'string\')'),
    (R'tfds\.features\.Text\(', R'datasets.Value(\'string\'),'),
    (R'features\s*=\s*tfds.features.FeaturesDict\(', R'features=datasets.Features('),
    (R'tfds\.features\.FeaturesDict\(', R'dict('),
    (R'The TensorFlow Datasets Authors', R'The TensorFlow Datasets Authors and the HuggingFace Datasets Authors'),
    (R'tfds\.', R'datasets.'),
    (R'dl_manager\.manual_dir', R'self.config.data_dir'),
    (R'self\.builder_config', R'self.config'),
]
def 			_a (      SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		return ConvertCommand(args.tfds_path						,				args.datasets_directory     )


class        _a							(			UpperCamelCase__			):

							@staticmethod
							def       lowerCamelCase_ (							UpperCamelCase_:	ArgumentParser							)			->   int:





									"""simple docstring"""







									lowercase__            =      parser.add_parser(
									    '''convert'''    ,						help='''Convert a TensorFlow Datasets dataset to a HuggingFace Datasets dataset.'''    ,						)
									train_parser.add_argument(
									    '''--tfds_path'''    ,						type=UpperCamelCase_    ,						required=UpperCamelCase_    ,						help='''Path to a TensorFlow Datasets folder to convert or a single tfds file to convert.'''    ,						)
									train_parser.add_argument(
									    '''--datasets_directory'''    ,						type=UpperCamelCase_    ,						required=UpperCamelCase_    ,						help='''Path to the HuggingFace Datasets folder.'''							)
									train_parser.set_defaults(func=UpperCamelCase_							)

							def __init__(							self:	int    ,						UpperCamelCase_:	str    ,						UpperCamelCase_:	str    ,						*UpperCamelCase_:	Optional[int]							)			->   Union[str, Any]:





									"""simple docstring"""







									lowercase__            =      get_logger('''datasets-cli/converting'''							)

									lowercase__            =      tfds_path
									lowercase__            =      datasets_directory


							def       lowerCamelCase_ (							self:	int							)			->   Optional[int]:





									"""simple docstring"""







									if os.path.isdir(self._tfds_path							):
											lowercase__            =      os.path.abspath(self._tfds_path							)
									elif os.path.isfile(self._tfds_path							):
											lowercase__            =      os.path.dirname(self._tfds_path							)
									else:
											raise ValueError('''--tfds_path is neither a directory nor a file. Please check path.'''							)

									lowercase__            =      os.path.abspath(self._datasets_directory							)

									self._logger.info(f'Converting datasets from {abs_tfds_path} to {abs_datasets_path}'							)

									lowercase__            =      []
									lowercase__            =      []
									lowercase__            =      {}

									if os.path.isdir(self._tfds_path							):
											lowercase__            =      os.listdir(UpperCamelCase_							)
									else:
											lowercase__            =      [os.path.basename(self._tfds_path							)]

									for f_name in file_names:
											self._logger.info(f'Looking at file {f_name}'							)
											lowercase__            =      os.path.join(UpperCamelCase_    ,						UpperCamelCase_							)
											lowercase__            =      os.path.join(UpperCamelCase_    ,						UpperCamelCase_							)

											if not os.path.isfile(UpperCamelCase_							) or "__init__" in f_name or "_test" in f_name or ".py" not in f_name:
													self._logger.info('''Skipping file'''							)
													continue

											with open(UpperCamelCase_    ,						encoding='''utf-8'''							) as f:
													lowercase__            =      f.readlines()

											lowercase__            =      []
											lowercase__            =      False
											lowercase__            =      False
											lowercase__            =      []
											for line in lines:
													lowercase__            =      line

													# Convert imports
													if "import tensorflow.compat.v2 as tf" in out_line:
															continue
													elif "@tfds.core" in out_line:
															continue
													elif "builder=self" in out_line:
															continue
													elif "import tensorflow_datasets.public_api as tfds" in out_line:
															lowercase__            =      '''import datasets\n'''
													elif "import tensorflow" in out_line:
															# order is important here
															lowercase__            =      ''''''
															continue
													elif "from absl import logging" in out_line:
															lowercase__            =      '''from datasets import logging\n'''
													elif "getLogger" in out_line:
															lowercase__            =      out_line.replace('''getLogger'''    ,						'''get_logger'''							)
													elif any(expression in out_line for expression in TO_HIGHLIGHT							):
															lowercase__            =      True
															lowercase__            =      list(filter(lambda UpperCamelCase_							: e in out_line    ,						UpperCamelCase_							)							)
															out_lines.append(HIGHLIGHT_MESSAGE_PRE + str(UpperCamelCase_							) + '''\n'''							)
															out_lines.append(UpperCamelCase_							)
															out_lines.append(UpperCamelCase_							)
															continue
													else:
															for pattern, replacement in TO_CONVERT:
																	lowercase__            =      re.sub(UpperCamelCase_    ,						UpperCamelCase_    ,						UpperCamelCase_							)

                # Take care of saving utilities (to later move them together with main script)
													if "tensorflow_datasets" in out_line:
															lowercase__            =      re.match(r'''from\stensorflow_datasets.*import\s([^\.\r\n]+)'''    ,						UpperCamelCase_							)
															tfds_imports.extend(imp.strip() for imp in match.group(1							).split(''','''							)							)
															lowercase__            =      '''from . import ''' + match.group(1							)

													# Check we have not forget anything
													if "tf." in out_line or "tfds." in out_line or "tensorflow_datasets" in out_line:
															raise ValueError(f'Error converting {out_line.strip()}'							)

													if "GeneratorBasedBuilder" in out_line or "BeamBasedBuilder" in out_line:
															lowercase__            =      True
													out_lines.append(UpperCamelCase_							)

											if is_builder or "wmt" in f_name:
													# We create a new directory for each dataset
													lowercase__            =      f_name.replace('''.py'''    ,						''''''							)
													lowercase__            =      os.path.join(UpperCamelCase_    ,						UpperCamelCase_							)
													lowercase__            =      os.path.join(UpperCamelCase_    ,						UpperCamelCase_							)
													os.makedirs(UpperCamelCase_    ,						exist_ok=UpperCamelCase_							)
													self._logger.info(f'Adding directory {output_dir}'							)
													imports_to_builder_map.update({imp: output_dir for imp in tfds_imports}							)
											else:
													# Utilities will be moved at the end
													utils_files.append(UpperCamelCase_							)

											if needs_manual_update:
													with_manual_update.append(UpperCamelCase_							)

											with open(UpperCamelCase_    ,						'''w'''    ,						encoding='''utf-8'''							) as f:
													f.writelines(UpperCamelCase_							)
											self._logger.info(f'Converted in {output_file}'							)

									for utils_file in utils_files:
											try:
													lowercase__            =      os.path.basename(UpperCamelCase_							)
													lowercase__            =      imports_to_builder_map[f_name.replace('''.py'''    ,						''''''							)]
													self._logger.info(f'Moving {dest_folder} to {utils_file}'							)
													shutil.copy(UpperCamelCase_    ,						UpperCamelCase_							)
											except KeyError:
													self._logger.error(f'Cannot find destination folder for {utils_file}. Please copy manually.'							)

									if with_manual_update:
											for file_path in with_manual_update:
													self._logger.warning(
													    f'You need to manually update file {file_path} to remove configurations using \'TextEncoderConfig\'.'							)


