






from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
							from .deepspeed import (
							    DeepSpeedEngineWrapper,
							    DeepSpeedOptimizerWrapper,
							    DeepSpeedSchedulerWrapper,
							    DummyOptim,
							    DummyScheduler,
							    HfDeepSpeedConfig,
							)

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers









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

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()



class      SCREAMING_SNAKE_CASE  (				_SCREAMING_SNAKE_CASE       ,       unittest.TestCase						):





	"""simple docstring"""




	_A       :				int    =				PriorTransformer
	_A       :				Dict    =				"""hidden_states"""





	@property
	def    lowerCamelCase(self							):
		A_		:						Tuple	    =   4
		A_		:						Dict	    =   8
		A_		:						str	    =   7

		A_		:						Dict	    =   floats_tensor((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						str	    =   floats_tensor((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						List[str]	    =   floats_tensor((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	def    lowerCamelCase(self     ,     lowerCAmelCase_=0							):
		torch.manual_seed(lowerCAmelCase_							)
		A_		:						Optional[Any]	    =   4
		A_		:						Tuple	    =   8
		A_		:						Tuple	    =   7

		A_		:						int	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						List[str]	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						Tuple	    =   torch.randn((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	@property
	def    lowerCamelCase(self							):
		return (4, 8)





	@property
	def    lowerCamelCase(self							):
		return (4, 8)





	def    lowerCamelCase(self							):
		A_		:						int	    =   {
		    """num_attention_heads""": 2,
		    """attention_head_dim""": 4,
		    """num_layers""": 2,
		    """embedding_dim""": 8,
		    """num_embeddings""": 7,
		    """additional_embeddings""": 4,
		}
		A_		:						Dict	    =   self.dummy_input
		return init_dict, inputs_dict





	def    lowerCamelCase(self							):
		A_    ,     A_		:						Optional[int]	    =   PriorTransformer.from_pretrained(
		    """hf-internal-testing/prior-dummy"""     ,     output_loading_info=lowerCAmelCase_							)
		self.assertIsNotNone(lowerCAmelCase_							)
		self.assertEqual(len(loading_info["""missing_keys"""]							)     ,     0							)

		model.to(lowerCAmelCase_							)
		A_		:						List[Any]	    =   model(**self.dummy_input							)[0]

		assert hidden_states is not None, "Make sure output is not None"





	def    lowerCamelCase(self							):
		A_    ,     A_		:						Union[str, Any]	    =   self.prepare_init_args_and_inputs_for_common()

		A_		:						Union[str, Any]	    =   self.model_class(**lowerCAmelCase_							)
		A_		:						Tuple	    =   inspect.signature(model.forward							)
		# signature.parameters is an OrderedDict => so arg_names order is deterministic
		A_		:						Tuple	    =   [*signature.parameters.keys()]

		A_		:						List[str]	    =   ["""hidden_states""", """timestep"""]
		self.assertListEqual(arg_names[:2]     ,     lowerCAmelCase_							)





	def    lowerCamelCase(self							):
		A_		:						int	    =   PriorTransformer.from_pretrained("""hf-internal-testing/prior-dummy"""							)
		A_		:						Optional[Any]	    =   model.to(lowerCAmelCase_							)

		if hasattr(lowerCAmelCase_     ,     """set_default_attn_processor"""							):
			model.set_default_attn_processor()

		A_		:						int	    =   self.get_dummy_seed_input()

		with torch.no_grad():
			A_		:						int	    =   model(**lowerCAmelCase_							)[0]

		A_		:						Any	    =   output[0, :5].flatten().cpu()
		print(lowerCAmelCase_							)

		# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
		# the expected output slices are not the same for CPU and GPU.
		A_		:						Optional[int]	    =   torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239]							)
		self.assertTrue(torch_all_close(lowerCAmelCase_     ,     lowerCAmelCase_     ,     rtol=1e-2							)							)






@slow
class      SCREAMING_SNAKE_CASE  (				unittest.TestCase						):





	"""simple docstring"""




	def    lowerCamelCase(self     ,     lowerCAmelCase_=1     ,     lowerCAmelCase_=768     ,     lowerCAmelCase_=77     ,     lowerCAmelCase_=0							):
		torch.manual_seed(lowerCAmelCase_							)
		A_		:						Tuple	    =   batch_size
		A_		:						Any	    =   embedding_dim
		A_		:						Tuple	    =   num_embeddings

		A_		:						Any	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						Union[str, Any]	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						Tuple	    =   torch.randn((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	def    lowerCamelCase(self							):
		# clean up the VRAM after each test
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	@parameterized.expand(
	    [
	        # fmt: off
	        [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
	        [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
	        # fmt: on
	    ]							)
	def    lowerCamelCase(self     ,     lowerCAmelCase_     ,     lowerCAmelCase_							):
		A_		:						List[Any]	    =   PriorTransformer.from_pretrained("""kandinsky-community/kandinsky-2-1-prior"""     ,     subfolder="""prior"""							)
		model.to(lowerCAmelCase_							)
		A_		:						int	    =   self.get_dummy_seed_input(seed=lowerCAmelCase_							)

		with torch.no_grad():
			A_		:						Tuple	    =   model(**lowerCAmelCase_							)[0]

		assert list(sample.shape							) == [1, 768]

		A_		:						Tuple	    =   sample[0, :8].flatten().cpu()
		print(lowerCAmelCase_							)
		A_		:						Optional[int]	    =   torch.tensor(lowerCAmelCase_							)

		assert torch_all_close(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1e-3							)

