






"""simple docstring"""
import argparse
import os

import torch
from torchvision.datasets.utils import download_url

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, TransformeraDModel


lowerCAmelCase__							= {512: "DiT-XL-2-512x512.pt", 256: "DiT-XL-2-256x256.pt"}





def 							lowercase__		(				lowerCamelCase					):
				_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						f"""pretrained_models/{model_name}"""
				if not os.path.isfile(_UpperCamelCase					):
								os.makedirs('pretrained_models',	exist_ok=_UpperCamelCase					)
								_SCREAMING_SNAKE_CASE      :		List[Any]							=						f"""https://dl.fbaipublicfiles.com/DiT/models/{model_name}"""
								download_url(_UpperCamelCase,	'pretrained_models'					)
				_SCREAMING_SNAKE_CASE      :		List[Any]							=						torch.load(_UpperCamelCase,	map_location=lambda lowerCamelCase,	lowerCamelCase					: storage					)
				return model





def 							lowercase__		(				lowerCamelCase					):
				_SCREAMING_SNAKE_CASE      :		Dict							=						download_model(pretrained_models[args.image_size]					)

				_SCREAMING_SNAKE_CASE      :		str							=						state_dict['x_embedder.proj.weight']
				_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						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					):
								_SCREAMING_SNAKE_CASE      :		str							=						state_dict[
								    't_embedder.mlp.0.weight'
								]
								_SCREAMING_SNAKE_CASE      :		Tuple							=						state_dict[
								    't_embedder.mlp.0.bias'
								]
								_SCREAMING_SNAKE_CASE      :		Optional[int]							=						state_dict[
								    't_embedder.mlp.2.weight'
								]
								_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						state_dict[
								    't_embedder.mlp.2.bias'
								]
								_SCREAMING_SNAKE_CASE      :		Any							=						state_dict[
								    'y_embedder.embedding_table.weight'
								]

								_SCREAMING_SNAKE_CASE      :		Tuple							=						state_dict[
								    f"""blocks.{depth}.adaLN_modulation.1.weight"""
								]
								_SCREAMING_SNAKE_CASE      :		Optional[int]							=						state_dict[
								    f"""blocks.{depth}.adaLN_modulation.1.bias"""
								]

								_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						torch.chunk(state_dict[f"""blocks.{depth}.attn.qkv.weight"""],	3,	dim=0					)
								_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						torch.chunk(state_dict[f"""blocks.{depth}.attn.qkv.bias"""],	3,	dim=0					)

								_SCREAMING_SNAKE_CASE      :		str							=						q
								_SCREAMING_SNAKE_CASE      :		Dict							=						q_bias
								_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						k
								_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						k_bias
								_SCREAMING_SNAKE_CASE      :		Tuple							=						v
								_SCREAMING_SNAKE_CASE      :		Tuple							=						v_bias

								_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						state_dict[
								    f"""blocks.{depth}.attn.proj.weight"""
								]
								_SCREAMING_SNAKE_CASE      :		List[Any]							=						state_dict[f"""blocks.{depth}.attn.proj.bias"""]

								_SCREAMING_SNAKE_CASE      :		Dict							=						state_dict[f"""blocks.{depth}.mlp.fc1.weight"""]
								_SCREAMING_SNAKE_CASE      :		str							=						state_dict[f"""blocks.{depth}.mlp.fc1.bias"""]
								_SCREAMING_SNAKE_CASE      :		int							=						state_dict[f"""blocks.{depth}.mlp.fc2.weight"""]
								_SCREAMING_SNAKE_CASE      :		Tuple							=						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'					)

				_SCREAMING_SNAKE_CASE      :		Tuple							=						state_dict['final_layer.adaLN_modulation.1.weight']
				_SCREAMING_SNAKE_CASE      :		Any							=						state_dict['final_layer.adaLN_modulation.1.bias']
				_SCREAMING_SNAKE_CASE      :		Tuple							=						state_dict['final_layer.linear.weight']
				_SCREAMING_SNAKE_CASE      :		List[Any]							=						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
				_SCREAMING_SNAKE_CASE      :		Optional[int]							=						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=_UpperCamelCase,	num_attention_heads=16,	activation_fn='gelu-approximate',	num_embeds_ada_norm=1_000,	norm_type='ada_norm_zero',	norm_elementwise_affine=_UpperCamelCase,	)
				transformer.load_state_dict(_UpperCamelCase,	strict=_UpperCamelCase					)

				_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						DDIMScheduler(
				    num_train_timesteps=1_000,	beta_schedule='linear',	prediction_type='epsilon',	clip_sample=_UpperCamelCase,	)

				_SCREAMING_SNAKE_CASE      :		Any							=						AutoencoderKL.from_pretrained(args.vae_model					)

				_SCREAMING_SNAKE_CASE      :		Optional[int]							=						DiTPipeline(transformer=_UpperCamelCase,	vae=_UpperCamelCase,	scheduler=_UpperCamelCase					)

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


if __name__ == "__main__":
			lowerCAmelCase__							= 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__							= parser.parse_args()
			main(args)









"""simple docstring"""
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class _lowerCAmelCase	(						__UpperCAmelCase      ,     __UpperCAmelCase      ,     unittest.TestCase					):
						SCREAMING_SNAKE_CASE_:				int   				=							StableDiffusionParadigmsPipeline
						SCREAMING_SNAKE_CASE_:				str   				=							TEXT_TO_IMAGE_PARAMS
						SCREAMING_SNAKE_CASE_:				Any   				=							TEXT_TO_IMAGE_BATCH_PARAMS
						SCREAMING_SNAKE_CASE_:				str   				=							TEXT_TO_IMAGE_IMAGE_PARAMS
						SCREAMING_SNAKE_CASE_:				List[Any]   				=							TEXT_TO_IMAGE_IMAGE_PARAMS







						def  A			(						self       )		->       int:
										torch.manual_seed(0       )
										_SCREAMING_SNAKE_CASE      :		Dict							=						UNetaDConditionModel(
										    block_out_channels=(3_2, 6_4)      ,							layers_per_block=2      ,							sample_size=3_2      ,							in_channels=4      ,							out_channels=4      ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')      ,							up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')      ,							cross_attention_dim=3_2      ,							attention_head_dim=(2, 4)      ,							use_linear_projection=lowerCAmelCase_      ,							)
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						DDIMParallelScheduler(
										    beta_start=0.00_085      ,							beta_end=0.012      ,							beta_schedule='scaled_linear'      ,							clip_sample=lowerCAmelCase_      ,							set_alpha_to_one=lowerCAmelCase_      ,							)
										torch.manual_seed(0       )
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						AutoencoderKL(
										    block_out_channels=[3_2, 6_4]      ,							in_channels=3      ,							out_channels=3      ,							down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']      ,							up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']      ,							latent_channels=4      ,							sample_size=1_2_8      ,							)
										torch.manual_seed(0       )
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						CLIPTextConfig(
										    bos_token_id=0      ,							eos_token_id=2      ,							hidden_size=3_2      ,							intermediate_size=3_7      ,							layer_norm_eps=1e-05      ,							num_attention_heads=4      ,							num_hidden_layers=5      ,							pad_token_id=1      ,							vocab_size=1_0_0_0      ,							hidden_act='gelu'      ,							projection_dim=5_1_2      ,							)
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						CLIPTextModel(lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

										_SCREAMING_SNAKE_CASE      :		List[Any]							=						{
										    'unet': unet,
										    'scheduler': scheduler,
										    'vae': vae,
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'safety_checker': None,
										    'feature_extractor': None,
										}
										return components







						def  A			(						self      ,							lowerCAmelCase_      ,							lowerCAmelCase_=0       )		->       Tuple:
										if str(lowerCAmelCase_       ).startswith('mps'       ):
														_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						torch.manual_seed(lowerCAmelCase_       )
										else:
														_SCREAMING_SNAKE_CASE      :		int							=						torch.Generator(device=lowerCAmelCase_       ).manual_seed(lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		str							=						{
										    'prompt': 'a photograph of an astronaut riding a horse',
										    'generator': generator,
										    'num_inference_steps': 1_0,
										    'guidance_scale': 6.0,
										    'output_type': 'numpy',
										    'parallel': 3,
										    'debug': True,
										}
										return inputs







						def  A			(						self       )		->       int:
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						'cpu'  # ensure determinism for the device-dependent torch.Generator
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						self.get_dummy_components()
										_SCREAMING_SNAKE_CASE      :		Any							=						StableDiffusionParadigmsPipeline(**lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						sd_pipe.to(lowerCAmelCase_       )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						self.get_dummy_inputs(lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		Dict							=						sd_pipe(**lowerCAmelCase_       ).images
										_SCREAMING_SNAKE_CASE      :		Dict							=						image[0, -3:, -3:, -1]
										assert image.shape == (1, 6_4, 6_4, 3)

										_SCREAMING_SNAKE_CASE      :		int							=						np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023]       )

										assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







						def  A			(						self       )		->       Tuple:
										_SCREAMING_SNAKE_CASE      :		Tuple							=						'cpu'  # ensure determinism for the device-dependent torch.Generator
										_SCREAMING_SNAKE_CASE      :		str							=						self.get_dummy_components()
										torch.manual_seed(0       )
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						DDPMParallelScheduler()
										torch.manual_seed(0       )
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						StableDiffusionParadigmsPipeline(**lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		int							=						sd_pipe.to(lowerCAmelCase_       )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		Tuple							=						self.get_dummy_inputs(lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						sd_pipe(**lowerCAmelCase_       ).images
										_SCREAMING_SNAKE_CASE      :		Dict							=						image[0, -3:, -3:, -1]
										assert image.shape == (1, 6_4, 6_4, 3)

										_SCREAMING_SNAKE_CASE      :		Dict							=						np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468]       )

										assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







						def  A			(						self       )		->       Optional[Any]:
										super().test_inference_batch_consistent(batch_sizes=[1, 2]       )







						def  A			(						self       )		->       Tuple:
										super().test_inference_batch_single_identical(batch_size=2      ,							expected_max_diff=3e-3       )







						def  A			(						self       )		->       List[str]:
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						'cpu'  # ensure determinism for the device-dependent torch.Generator
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						self.get_dummy_components()
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						StableDiffusionParadigmsPipeline(**lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		List[str]							=						sd_pipe.to(lowerCAmelCase_       )
										sd_pipe.set_progress_bar_config(disable=lowerCAmelCase_       )

										_SCREAMING_SNAKE_CASE      :		int							=						self.get_dummy_inputs(lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		Tuple							=						'french fries'
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						sd_pipe(**lowerCAmelCase_      ,							negative_prompt=lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						output.images
										_SCREAMING_SNAKE_CASE      :		Dict							=						image[0, -3:, -3:, -1]

										assert image.shape == (1, 6_4, 6_4, 3)

										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015]       )

										assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2







@slow
@require_torch_gpu
class _lowerCAmelCase	(						unittest.TestCase					):







						def  A			(						self       )		->       Tuple:
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







						def  A			(						self      ,							lowerCAmelCase_=0       )		->       int:
										_SCREAMING_SNAKE_CASE      :		str							=						torch.Generator(device=lowerCAmelCase_       ).manual_seed(lowerCAmelCase_       )
										_SCREAMING_SNAKE_CASE      :		List[str]							=						{
										    'prompt': 'a photograph of an astronaut riding a horse',
										    'generator': generator,
										    'num_inference_steps': 1_0,
										    'guidance_scale': 7.5,
										    'output_type': 'numpy',
										    'parallel': 3,
										    'debug': True,
										}
										return inputs







						def  A			(						self       )		->       List[str]:
										_SCREAMING_SNAKE_CASE      :		Any							=						'stabilityai/stable-diffusion-2-base'
										_SCREAMING_SNAKE_CASE      :		Dict							=						DDIMParallelScheduler.from_pretrained(lowerCAmelCase_      ,							subfolder='scheduler'       )
										_SCREAMING_SNAKE_CASE      :		List[str]							=						StableDiffusionParadigmsPipeline.from_pretrained(lowerCAmelCase_      ,							scheduler=lowerCAmelCase_      ,							safety_checker=lowerCAmelCase_       )
										pipe.to(lowerCAmelCase_       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase_       )
										pipe.enable_attention_slicing()

										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						self.get_inputs()
										_SCREAMING_SNAKE_CASE      :		str							=						pipe(**lowerCAmelCase_       ).images
										_SCREAMING_SNAKE_CASE      :		Dict							=						image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 5_1_2, 5_1_2, 3)

										_SCREAMING_SNAKE_CASE      :		Any							=						np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741]       )

										assert np.abs(expected_slice - image_slice       ).max() < 1e-2

