code: 
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  import argparse
  import os
  
  import torch
  from torchvision.datasets.utils import download_url
  
  from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, TransformeraDModel
  
  
  _lowerCamelCase							:		Optional[Any]							=						{512: 'DiT-XL-2-512x512.pt', 256: 'DiT-XL-2-256x256.pt'}
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = f'''pretrained_models/{model_name}'''
  		if not os.path.isfile(lowercase__						):
  				os.makedirs("pretrained_models"		,							exist_ok=lowercase__						)
  				_snake_case			 = f'''https://dl.fbaipublicfiles.com/DiT/models/{model_name}'''
  				download_url(lowercase__		,							"pretrained_models"						)
  		_snake_case			 = torch.load(lowercase__		,							map_location=lambda lowercase_		,							lowercase_						: storage						)
  		return model
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = download_model(pretrained_models[args.image_size]						)
  
  		_snake_case			 = state_dict["x_embedder.proj.weight"]
  		_snake_case			 = state_dict["x_embedder.proj.bias"]
  		state_dict.pop("x_embedder.proj.weight"						)
  		state_dict.pop("x_embedder.proj.bias"						)
  
  		for depth in range(28						):
  				_snake_case			 = state_dict[
  				    "t_embedder.mlp.0.weight"
  				]
  				_snake_case			 = state_dict[
  				    "t_embedder.mlp.0.bias"
  				]
  				_snake_case			 = state_dict[
  				    "t_embedder.mlp.2.weight"
  				]
  				_snake_case			 = state_dict[
  				    "t_embedder.mlp.2.bias"
  				]
  				_snake_case			 = state_dict[
  				    "y_embedder.embedding_table.weight"
  				]
  
  				_snake_case			 = state_dict[
  				    f'''blocks.{depth}.adaLN_modulation.1.weight'''
  				]
  				_snake_case			 = state_dict[
  				    f'''blocks.{depth}.adaLN_modulation.1.bias'''
  				]
  
  				_snake_case ,					_snake_case ,					_snake_case			 = torch.chunk(state_dict[f'''blocks.{depth}.attn.qkv.weight''']		,							3		,							dim=0						)
  				_snake_case ,					_snake_case ,					_snake_case			 = torch.chunk(state_dict[f'''blocks.{depth}.attn.qkv.bias''']		,							3		,							dim=0						)
  
  				_snake_case			 = q
  				_snake_case			 = q_bias
  				_snake_case			 = k
  				_snake_case			 = k_bias
  				_snake_case			 = v
  				_snake_case			 = v_bias
  
  				_snake_case			 = state_dict[
  				    f'''blocks.{depth}.attn.proj.weight'''
  				]
  				_snake_case			 = state_dict[f'''blocks.{depth}.attn.proj.bias''']
  
  				_snake_case			 = state_dict[f'''blocks.{depth}.mlp.fc1.weight''']
  				_snake_case			 = state_dict[f'''blocks.{depth}.mlp.fc1.bias''']
  				_snake_case			 = state_dict[f'''blocks.{depth}.mlp.fc2.weight''']
  				_snake_case			 = state_dict[f'''blocks.{depth}.mlp.fc2.bias''']
  
  				state_dict.pop(f'''blocks.{depth}.attn.qkv.weight'''						)
  				state_dict.pop(f'''blocks.{depth}.attn.qkv.bias'''						)
  				state_dict.pop(f'''blocks.{depth}.attn.proj.weight'''						)
  				state_dict.pop(f'''blocks.{depth}.attn.proj.bias'''						)
  				state_dict.pop(f'''blocks.{depth}.mlp.fc1.weight'''						)
  				state_dict.pop(f'''blocks.{depth}.mlp.fc1.bias'''						)
  				state_dict.pop(f'''blocks.{depth}.mlp.fc2.weight'''						)
  				state_dict.pop(f'''blocks.{depth}.mlp.fc2.bias'''						)
  				state_dict.pop(f'''blocks.{depth}.adaLN_modulation.1.weight'''						)
  				state_dict.pop(f'''blocks.{depth}.adaLN_modulation.1.bias'''						)
  
  		state_dict.pop("t_embedder.mlp.0.weight"						)
  		state_dict.pop("t_embedder.mlp.0.bias"						)
  		state_dict.pop("t_embedder.mlp.2.weight"						)
  		state_dict.pop("t_embedder.mlp.2.bias"						)
  		state_dict.pop("y_embedder.embedding_table.weight"						)
  
  		_snake_case			 = state_dict["final_layer.adaLN_modulation.1.weight"]
  		_snake_case			 = state_dict["final_layer.adaLN_modulation.1.bias"]
  		_snake_case			 = state_dict["final_layer.linear.weight"]
  		_snake_case			 = state_dict["final_layer.linear.bias"]
  
  		state_dict.pop("final_layer.linear.weight"						)
  		state_dict.pop("final_layer.linear.bias"						)
  		state_dict.pop("final_layer.adaLN_modulation.1.weight"						)
  		state_dict.pop("final_layer.adaLN_modulation.1.bias"						)
  
  		# DiT XL/2
  		_snake_case			 = TransformeraDModel(
  		    sample_size=args.image_size // 8		,							num_layers=28		,							attention_head_dim=72		,							in_channels=4		,							out_channels=8		,							patch_size=2		,							attention_bias=lowercase__		,							num_attention_heads=16		,							activation_fn="gelu-approximate"		,							num_embeds_ada_norm=1000		,							norm_type="ada_norm_zero"		,							norm_elementwise_affine=lowercase__		,							)
  		transformer.load_state_dict(lowercase__		,							strict=lowercase__						)
  
  		_snake_case			 = DDIMScheduler(
  		    num_train_timesteps=1000		,							beta_schedule="linear"		,							prediction_type="epsilon"		,							clip_sample=lowercase__		,							)
  
  		_snake_case			 = AutoencoderKL.from_pretrained(args.vae_model						)
  
  		_snake_case			 = DiTPipeline(transformer=lowercase__		,							vae=lowercase__		,							scheduler=lowercase__						)
  
  		if args.save:
  				pipeline.save_pretrained(args.checkpoint_path						)
  
  
  if __name__ == "__main__":
  						_lowerCamelCase							:		Tuple							=						argparse.ArgumentParser()
  
  						parser.add_argument(
  						    '--image_size',
  						    default=256,
  						    type=int,
  						    required=False,
  						    help='Image size of pretrained model, either 256 or 512.',
  						)
  						parser.add_argument(
  						    '--vae_model',
  						    default='stabilityai/sd-vae-ft-ema',
  						    type=str,
  						    required=False,
  						    help='Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.',
  						)
  						parser.add_argument(
  						    '--save', default=True, type=bool, required=False, help='Whether to save the converted pipeline or not.'
  						)
  						parser.add_argument(
  						    '--checkpoint_path', default=None, type=str, required=True, help='Path to the output pipeline.'
  						)
  
  						_lowerCamelCase							:		str							=						parser.parse_args()
  						main(args)
  
  
  
  
  
  
style_context: 
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  import importlib
  import importlib.metadata
  import os
  import warnings
  from distutils.util import strtobool
  from functools import lru_cache
  
  import torch
  from packaging import version
  from packaging.version import parse
  
  from .environment import parse_flag_from_env
  from .versions import compare_versions, is_torch_version
  
  
  try:
  						import torch_xla.core.xla_model as xm  # noqa: F401
  
  						_lowerCamelCase							:		int										=						True
  except ImportError:
  						_lowerCamelCase							:		List[str]										=						False
  
  
  # Cache this result has it's a C FFI call which can be pretty time-consuming
  _lowerCamelCase							:		List[Any]										=						torch.distributed.is_available()
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = importlib.util.find_spec(lowercase_						) is not None
  		if package_exists:
  				try:
  						_snake_case			 = importlib.metadata.metadata(lowercase_						)
  						return True
  				except importlib.metadata.PackageNotFoundError:
  						return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _torch_distributed_available
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		try:
  				pass
  		except ImportError:
  				print(
  				    "Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not"
  				    " detected. If you see \"ValueError: Invalid backend: 'ccl'\" error, please install Intel(R) oneCCL"
  				    " Bindings for PyTorch*."						)
  		return (
  		    importlib.util.find_spec("torch_ccl"						) is not None
  		    or importlib.util.find_spec("oneccl_bindings_for_pytorch"						) is not None
  		)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return importlib.metadata.version("oneccl_bind_pt"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("transformer_engine"						)
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=True						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if torch.cuda.is_available():
  				return False
  		if _tpu_available and check_device:
  				try:
  						# Will raise a RuntimeError if no XLA configuration is found
  						_snake_case			 = xm.xla_device()
  						return True
  				except RuntimeError:
  						return False
  		return _tpu_available
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("deepspeed"						)
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if is_tpu_available():
  				return not ignore_tpu
  		if torch.cuda.is_available():
  				return torch.cuda.is_bfaa_supported()
  		if is_npu_available():
  				return False
  		return True
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = _is_package_available("bitsandbytes"						)
  		if package_exists:
  				_snake_case			 = version.parse(importlib.metadata.version("bitsandbytes"						)						)
  				return compare_versions(lowercase_		,							">="		,							"0.39.0"						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = _is_package_available("bitsandbytes"						)
  		if package_exists:
  				_snake_case			 = version.parse(importlib.metadata.version("bitsandbytes"						)						)
  				return compare_versions(lowercase_		,							">="		,							"0.37.2"						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("bitsandbytes"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if strtobool(os.environ.get("ACCELERATE_USE_MEGATRON_LM"		,							"False"						)						) == 1:
  				_snake_case			 = importlib.util.find_spec("megatron"						) is not None
  				if package_exists:
  						try:
  								_snake_case			 = parse(importlib.metadata.version("megatron-lm"						)						)
  								return compare_versions(lowercase_		,							">="		,							"2.2.0"						)
  						except Exception as e:
  								warnings.warn(f'''Parse Megatron version failed. Exception:{e}'''						)
  								return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("safetensors"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("transformers"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("datasets"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("aim"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("tensorboard"						) or _is_package_available("tensorboardX"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("wandb"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("comet_ml"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("boto3"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if _is_package_available("rich"						):
  				if "ACCELERATE_DISABLE_RICH" in os.environ:
  						warnings.warn(
  						    "`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`."						)
  						return not parse_flag_from_env("ACCELERATE_DISABLE_RICH"		,							lowercase_						)
  				return parse_flag_from_env("ACCELERATE_ENABLE_RICH"		,							lowercase_						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("sagemaker"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("tqdm"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("mlflow"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return is_torch_version(">="		,							"1.12"						) and torch.backends.mps.is_available() and torch.backends.mps.is_built()
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def get_major_and_minor_from_version(lowercase_						):
  				return str(version.parse(lowercase_						).major						) + "." + str(version.parse(lowercase_						).minor						)
  
  		_snake_case			 = importlib.metadata.version("torch"						)
  		if importlib.util.find_spec("intel_extension_for_pytorch"						) is None:
  				return False
  		_snake_case			 = "N/A"
  		try:
  				_snake_case			 = importlib.metadata.version("intel_extension_for_pytorch"						)
  		except importlib.metadata.PackageNotFoundError:
  				return False
  		_snake_case			 = get_major_and_minor_from_version(_torch_version						)
  		_snake_case			 = get_major_and_minor_from_version(_ipex_version						)
  		if torch_major_and_minor != ipex_major_and_minor:
  				warnings.warn(
  				    f'''Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,'''
  				    f''' but PyTorch {_torch_version} is found. Please switch to the matching version and run again.'''						)
  				return False
  		return True
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if importlib.util.find_spec("torch"						) is None or importlib.util.find_spec("torch_npu"						) is None:
  				return False
  
  		import torch
  		import torch_npu  # noqa: F401
  
  		if check_device:
  				try:
  						# Will raise a RuntimeError if no NPU is found
  						_snake_case			 = torch.npu.device_count()
  						return torch.npu.is_available()
  				except RuntimeError:
  						return False
  		return hasattr(lowercase_		,							"npu"						) and torch.npu.is_available()
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if not parse_flag_from_env("ACCELERATE_USE_XPU"		,							default=lowercase_						):
  				return False
  		if is_ipex_available():
  				import torch
  
  				if is_torch_version("<="		,							"1.12"						):
  						return False
  		else:
  				return False
  
  		import intel_extension_for_pytorch  # noqa: F401
  
  		if check_device:
  				try:
  						# Will raise a RuntimeError if no XPU  is found
  						_snake_case			 = torch.xpu.device_count()
  						return torch.xpu.is_available()
  				except RuntimeError:
  						return False
  		return hasattr(lowercase_		,							"xpu"						) and torch.xpu.is_available()
  
  
  
  
  
  