




import argparse
import os

import torch
from torchvision.datasets.utils import download_url

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, TransformeraDModel


UpperCamelCase__       							=      {5_1_2: 'DiT-XL-2-512x512.pt', 2_5_6: 'DiT-XL-2-256x256.pt'}




def      lowerCAmelCase_				(  __A							)  ->       Dict:







      '''simple docstring'''


      UpperCAmelCase__			=						f"""pretrained_models/{model_name}"""
      if not os.path.isfile(__A							):
            os.makedirs("pretrained_models",						exist_ok=__A							)
            UpperCAmelCase__			=						f"""https://dl.fbaipublicfiles.com/DiT/models/{model_name}"""
            download_url(__A,						"pretrained_models"							)
      UpperCAmelCase__			=						torch.load(__A,						map_location=lambda __A,						__A							: storage							)
      return model




def      lowerCAmelCase_				(  __A							)  ->       Union[str, Any]:







      '''simple docstring'''


      UpperCAmelCase__			=						download_model(pretrained_models[args.image_size]							)

      UpperCAmelCase__			=						state_dict["x_embedder.proj.weight"]
      UpperCAmelCase__			=						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							):
            UpperCAmelCase__			=						state_dict[
                "t_embedder.mlp.0.weight"
            ]
            UpperCAmelCase__			=						state_dict[
                "t_embedder.mlp.0.bias"
            ]
            UpperCAmelCase__			=						state_dict[
                "t_embedder.mlp.2.weight"
            ]
            UpperCAmelCase__			=						state_dict[
                "t_embedder.mlp.2.bias"
            ]
            UpperCAmelCase__			=						state_dict[
                "y_embedder.embedding_table.weight"
            ]

            UpperCAmelCase__			=						state_dict[
                f"""blocks.{depth}.adaLN_modulation.1.weight"""
            ]
            UpperCAmelCase__			=						state_dict[
                f"""blocks.{depth}.adaLN_modulation.1.bias"""
            ]

            UpperCAmelCase__	,   UpperCAmelCase__	,   UpperCAmelCase__			=						torch.chunk(state_dict[f"""blocks.{depth}.attn.qkv.weight"""],						3,						dim=0							)
            UpperCAmelCase__	,   UpperCAmelCase__	,   UpperCAmelCase__			=						torch.chunk(state_dict[f"""blocks.{depth}.attn.qkv.bias"""],						3,						dim=0							)

            UpperCAmelCase__			=						q
            UpperCAmelCase__			=						q_bias
            UpperCAmelCase__			=						k
            UpperCAmelCase__			=						k_bias
            UpperCAmelCase__			=						v
            UpperCAmelCase__			=						v_bias

            UpperCAmelCase__			=						state_dict[
                f"""blocks.{depth}.attn.proj.weight"""
            ]
            UpperCAmelCase__			=						state_dict[f"""blocks.{depth}.attn.proj.bias"""]

            UpperCAmelCase__			=						state_dict[f"""blocks.{depth}.mlp.fc1.weight"""]
            UpperCAmelCase__			=						state_dict[f"""blocks.{depth}.mlp.fc1.bias"""]
            UpperCAmelCase__			=						state_dict[f"""blocks.{depth}.mlp.fc2.weight"""]
            UpperCAmelCase__			=						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"							)

      UpperCAmelCase__			=						state_dict["final_layer.adaLN_modulation.1.weight"]
      UpperCAmelCase__			=						state_dict["final_layer.adaLN_modulation.1.bias"]
      UpperCAmelCase__			=						state_dict["final_layer.linear.weight"]
      UpperCAmelCase__			=						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
      UpperCAmelCase__			=						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=__A,						num_attention_heads=16,						activation_fn="gelu-approximate",						num_embeds_ada_norm=1_000,						norm_type="ada_norm_zero",						norm_elementwise_affine=__A,						)
      transformer.load_state_dict(__A,						strict=__A							)

      UpperCAmelCase__			=						DDIMScheduler(
          num_train_timesteps=1_000,						beta_schedule="linear",						prediction_type="epsilon",						clip_sample=__A,						)

      UpperCAmelCase__			=						AutoencoderKL.from_pretrained(args.vae_model							)

      UpperCAmelCase__			=						DiTPipeline(transformer=__A,						vae=__A,						scheduler=__A							)

      if args.save:
            pipeline.save_pretrained(args.checkpoint_path							)


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

      parser.add_argument(
          '--image_size',
          default=2_5_6,
          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.'
      )

      UpperCamelCase__       							=      parser.parse_args()
      main(args)




import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


UpperCamelCase__       							=      namedtuple('ExampleNamedTuple', 'a b c')





class   A							(       unittest.TestCase  ):


   def     lowercase_   (self  :		str	)		->       List[Any]:
         """simple docstring"""



         UpperCAmelCase__			=						torch.randn(5 ,						2	)
         UpperCAmelCase__			=						torch.device("cuda"	) if torch.cuda.is_available() else torch.device("cpu"	)

         UpperCAmelCase__			=						send_to_device(__UpperCAmelCase ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta.cpu() ,						__UpperCAmelCase	)	)

         UpperCAmelCase__			=						send_to_device((tensor, [tensor, tensor], 1) ,						__UpperCAmelCase	)
         self.assertIsInstance(__UpperCAmelCase ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta[0].cpu() ,						__UpperCAmelCase	)	)
         self.assertIsInstance(resulta[1] ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta[1][0].cpu() ,						__UpperCAmelCase	)	)
         self.assertTrue(torch.equal(resulta[1][1].cpu() ,						__UpperCAmelCase	)	)
         self.assertEqual(resulta[2] ,						1	)

         UpperCAmelCase__			=						send_to_device({"a": tensor, "b": [tensor, tensor], "c": 1} ,						__UpperCAmelCase	)
         self.assertIsInstance(__UpperCAmelCase ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta["a"].cpu() ,						__UpperCAmelCase	)	)
         self.assertIsInstance(resulta["b"] ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta["b"][0].cpu() ,						__UpperCAmelCase	)	)
         self.assertTrue(torch.equal(resulta["b"][1].cpu() ,						__UpperCAmelCase	)	)
         self.assertEqual(resulta["c"] ,						1	)

         UpperCAmelCase__			=						send_to_device(ExampleNamedTuple(a=__UpperCAmelCase ,						b=[tensor, tensor] ,						c=1	) ,						__UpperCAmelCase	)
         self.assertIsInstance(__UpperCAmelCase ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta.a.cpu() ,						__UpperCAmelCase	)	)
         self.assertIsInstance(resulta.b ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta.b[0].cpu() ,						__UpperCAmelCase	)	)
         self.assertTrue(torch.equal(resulta.b[1].cpu() ,						__UpperCAmelCase	)	)
         self.assertEqual(resulta.c ,						1	)

         UpperCAmelCase__			=						send_to_device(UserDict({"a": tensor, "b": [tensor, tensor], "c": 1}	) ,						__UpperCAmelCase	)
         self.assertIsInstance(__UpperCAmelCase ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta["a"].cpu() ,						__UpperCAmelCase	)	)
         self.assertIsInstance(resulta["b"] ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta["b"][0].cpu() ,						__UpperCAmelCase	)	)
         self.assertTrue(torch.equal(resulta["b"][1].cpu() ,						__UpperCAmelCase	)	)
         self.assertEqual(resulta["c"] ,						1	)


   def     lowercase_   (self  :		int	)		->       int:
         """simple docstring"""



         with self.assertRaises(__UpperCAmelCase	) as cm:
               UpperCAmelCase__			=						recursively_apply(torch.tensor ,						(torch.tensor(1	), 1) ,						error_on_other_type=__UpperCAmelCase	)
         self.assertEqual(
             str(cm.exception	) ,						"Unsupported types (<class 'int'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed." ,						)


   def     lowercase_   (self  :		List[str]	)		->       Any:
         """simple docstring"""



         UpperCAmelCase__			=						torch.tensor([1, 2, 3, 4, 5]	)
         self.assertEqual(listify(__UpperCAmelCase	) ,						[1, 2, 3, 4, 5]	)

         UpperCAmelCase__			=						torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]	)
         self.assertEqual(listify(__UpperCAmelCase	) ,						[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]	)

         UpperCAmelCase__			=						torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]	)
         self.assertEqual(
             listify(__UpperCAmelCase	) ,						[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]	)


   def     lowercase_   (self  :		Union[str, Any]	)		->       Dict:
         """simple docstring"""



         with patch_environment(aa=1 ,						BB=2	):
               self.assertEqual(os.environ.get("AA"	) ,						"1"	)
               self.assertEqual(os.environ.get("BB"	) ,						"2"	)

         self.assertNotIn("AA" ,						os.environ	)
         self.assertNotIn("BB" ,						os.environ	)


   def     lowercase_   (self  :		int	)		->       str:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						model.forward
         UpperCAmelCase__			=						convert_outputs_to_fpaa(model.forward	)
         UpperCAmelCase__			=						extract_model_from_parallel(__UpperCAmelCase ,						keep_fpaa_wrapper=__UpperCAmelCase	)
         UpperCAmelCase__			=						pickle.dumps(__UpperCAmelCase	)


   @require_cuda
   def     lowercase_   (self  :		List[str]	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						model.forward
         UpperCAmelCase__			=						torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
         UpperCAmelCase__			=						convert_outputs_to_fpaa(model.forward	)
         UpperCAmelCase__			=						extract_model_from_parallel(__UpperCAmelCase ,						keep_fpaa_wrapper=__UpperCAmelCase	)
         UpperCAmelCase__			=						pickle.dumps(__UpperCAmelCase	)


   @require_cuda
   @require_torch_min_version(version="2.0"	)
   def     lowercase_   (self  :		Optional[int]	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						model.forward
         UpperCAmelCase__			=						torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
         UpperCAmelCase__			=						convert_outputs_to_fpaa(model.forward	)
         UpperCAmelCase__			=						torch.compile(model.forward ,						backend="inductor"	)
         UpperCAmelCase__			=						torch.randn(4 ,						1_0	).cuda()
         UpperCAmelCase__			=						model(__UpperCAmelCase	)


   def     lowercase_   (self  :		List[Any]	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
         UpperCAmelCase__			=						torch.nn.parallel.DataParallel(__UpperCAmelCase	)
         UpperCAmelCase__			=						extract_model_from_parallel(__UpperCAmelCase	)

         self.assertEqual(__UpperCAmelCase ,						__UpperCAmelCase	)


   @require_torch_min_version(version="2.0"	)
   def     lowercase_   (self  :		Union[str, Any]	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						torch.compile(__UpperCAmelCase	)

         # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
         UpperCAmelCase__			=						torch.nn.parallel.DataParallel(__UpperCAmelCase	)
         UpperCAmelCase__			=						torch.compile(__UpperCAmelCase	)
         UpperCAmelCase__			=						extract_model_from_parallel(__UpperCAmelCase	)

         self.assertEqual(compiled_model._orig_mod ,						compiled_model_unwrapped._orig_mod	)





   def     lowercase_   (self  :		int	)		->       int:
         """simple docstring"""



         self.assertEqual(find_device([1, "a", torch.tensor([1, 2, 3]	)]	) ,						torch.device("cpu"	)	)
         self.assertEqual(find_device({"a": 1, "b": torch.tensor([1, 2, 3]	)}	) ,						torch.device("cpu"	)	)
         self.assertIsNone(find_device([1, "a"]	)	)



