



import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






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

    """simple docstring"""







    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()

   @property
   def      UpperCamelCase_  (					self							:			str)				->				List[str]:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					1
    _snake_case    :  int					=					3
    _snake_case    :  int					=					(32, 32)

    _snake_case    :  str					=					floats_tensor((batch_size, num_channels) + sizes				,		rng=random.Random(0)).to(lowercase_)
    return image

   @property
   def      UpperCamelCase_  (					self							:			Tuple)				->				str:

    """simple docstring"""







    torch.manual_seed(0)
    _snake_case    :  int					=					UNetaDConditionModel(
        block_out_channels=(32, 64)				,		layers_per_block=2				,		sample_size=32				,		in_channels=4				,		out_channels=4				,		down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")				,		up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")				,		cross_attention_dim=32				,		)
    return model

   @property
   def      UpperCamelCase_  (					self							:			List[str])				->				str:

    """simple docstring"""







    torch.manual_seed(0)
    _snake_case    :  str					=					AutoencoderKL(
        block_out_channels=[32, 64]				,		in_channels=3				,		out_channels=3				,		down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]				,		up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]				,		latent_channels=4				,		)
    return model

   @property
   def      UpperCamelCase_  (					self							:			Tuple)				->				Any:

    """simple docstring"""







    torch.manual_seed(0)
    _snake_case    :  int					=					RobertaSeriesConfig(
        hidden_size=32				,		project_dim=32				,		intermediate_size=37				,		layer_norm_eps=1E-05				,		num_attention_heads=4				,		num_hidden_layers=5				,		pad_token_id=1				,		vocab_size=5006				,		)
    return RobertaSeriesModelWithTransformation(lowercase_)

   @property
   def      UpperCamelCase_  (					self							:			str)				->				str:

    """simple docstring"""







    def extract(*lowerCAmelCase							:			int				,		**lowerCAmelCase							:			str):







     class 		snake_case						:



        '''simple docstring'''






        def __init__(					self							:			int)				->				Any:

         """simple docstring"""







         _snake_case    :  List[Any]					=					torch.ones([0])


        def      UpperCamelCase_  (					self							:			Dict				,		lowerCAmelCase							:			Optional[int])				->				Optional[Any]:

         """simple docstring"""







         self.pixel_values.to(lowercase_)
         return self
     return Out()

    return extract

   def      UpperCamelCase_  (					self							:			str)				->				int:

    """simple docstring"""







    _snake_case    :  Any					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
    _snake_case    :  Union[str, Any]					=					self.dummy_cond_unet
    _snake_case    :  Optional[Any]					=					PNDMScheduler(skip_prk_steps=lowercase_)
    _snake_case    :  Optional[Any]					=					self.dummy_vae
    _snake_case    :  Tuple					=					self.dummy_text_encoder
    _snake_case    :  Tuple					=					XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta""")
    _snake_case    :  Optional[int]					=					77

    _snake_case    :  str					=					self.dummy_image.to(lowercase_)
    _snake_case    :  Optional[int]					=					init_image / 2 + 0.5

    # make sure here that pndm scheduler skips prk
    _snake_case    :  List[Any]					=					AltDiffusionImgaImgPipeline(
        unet=lowercase_				,		scheduler=lowercase_				,		vae=lowercase_				,		text_encoder=lowercase_				,		tokenizer=lowercase_				,		safety_checker=lowercase_				,		feature_extractor=self.dummy_extractor				,		)
    _snake_case    :  Union[str, Any]					=					VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor				,		do_normalize=lowercase_)
    _snake_case    :  Any					=					alt_pipe.to(lowercase_)
    alt_pipe.set_progress_bar_config(disable=lowercase_)

    _snake_case    :  List[Any]					=					"""A painting of a squirrel eating a burger"""
    _snake_case    :  int					=					torch.Generator(device=lowercase_).manual_seed(0)
    _snake_case    :  Any					=					alt_pipe(
        [prompt]				,		generator=lowercase_				,		guidance_scale=6.0				,		num_inference_steps=2				,		output_type="""np"""				,		image=lowercase_				,		)

    _snake_case    :  Any					=					output.images

    _snake_case    :  List[str]					=					torch.Generator(device=lowercase_).manual_seed(0)
    _snake_case    :  Optional[Any]					=					alt_pipe(
        [prompt]				,		generator=lowercase_				,		guidance_scale=6.0				,		num_inference_steps=2				,		output_type="""np"""				,		image=lowercase_				,		return_dict=lowercase_				,		)[0]

    _snake_case    :  List[str]					=					image[0, -3:, -3:, -1]
    _snake_case    :  List[Any]					=					image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 32, 32, 3)
    _snake_case    :  int					=					np.array([0.4_427, 0.3_731, 0.4_249, 0.4_941, 0.4_546, 0.4_148, 0.4_193, 0.4_666, 0.4_499])

    assert np.abs(image_slice.flatten() - expected_slice).max() < 5E-3
    assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 5E-3

   @unittest.skipIf(torch_device != """cuda"""				,		"""This test requires a GPU""")
   def      UpperCamelCase_  (					self							:			Optional[Any])				->				List[str]:

    """simple docstring"""







    _snake_case    :  Any					=					self.dummy_cond_unet
    _snake_case    :  Optional[int]					=					PNDMScheduler(skip_prk_steps=lowercase_)
    _snake_case    :  Optional[int]					=					self.dummy_vae
    _snake_case    :  Tuple					=					self.dummy_text_encoder
    _snake_case    :  Any					=					XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta""")
    _snake_case    :  Union[str, Any]					=					77

    _snake_case    :  Optional[int]					=					self.dummy_image.to(lowercase_)

    # put models in fp16
    _snake_case    :  Tuple					=					unet.half()
    _snake_case    :  Union[str, Any]					=					vae.half()
    _snake_case    :  Optional[Any]					=					bert.half()

    # make sure here that pndm scheduler skips prk
    _snake_case    :  Any					=					AltDiffusionImgaImgPipeline(
        unet=lowercase_				,		scheduler=lowercase_				,		vae=lowercase_				,		text_encoder=lowercase_				,		tokenizer=lowercase_				,		safety_checker=lowercase_				,		feature_extractor=self.dummy_extractor				,		)
    _snake_case    :  Any					=					VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor				,		do_normalize=lowercase_)
    _snake_case    :  List[Any]					=					alt_pipe.to(lowercase_)
    alt_pipe.set_progress_bar_config(disable=lowercase_)

    _snake_case    :  Optional[Any]					=					"""A painting of a squirrel eating a burger"""
    _snake_case    :  List[Any]					=					torch.manual_seed(0)
    _snake_case    :  Any					=					alt_pipe(
        [prompt]				,		generator=lowercase_				,		num_inference_steps=2				,		output_type="""np"""				,		image=lowercase_				,		).images

    assert image.shape == (1, 32, 32, 3)


   @unittest.skipIf(torch_device != """cuda"""				,		"""This test requires a GPU""")
   def      UpperCamelCase_  (					self							:			Optional[Any])				->				str:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
        """/img2img/sketch-mountains-input.jpg""")
    # resize to resolution that is divisible by 8 but not 16 or 32
    _snake_case    :  Tuple					=					init_image.resize((760, 504))

    _snake_case    :  Dict					=					"""BAAI/AltDiffusion"""
    _snake_case    :  Any					=					AltDiffusionImgaImgPipeline.from_pretrained(
        lowercase_				,		safety_checker=lowercase_				,		)
    pipe.to(lowercase_)
    pipe.set_progress_bar_config(disable=lowercase_)
    pipe.enable_attention_slicing()

    _snake_case    :  List[Any]					=					"""A fantasy landscape, trending on artstation"""

    _snake_case    :  Optional[int]					=					torch.manual_seed(0)
    _snake_case    :  str					=					pipe(
        prompt=lowercase_				,		image=lowercase_				,		strength=0.75				,		guidance_scale=7.5				,		generator=lowercase_				,		output_type="""np"""				,		)
    _snake_case    :  List[Any]					=					output.images[0]

    _snake_case    :  Dict					=					image[255:258, 383:386, -1]

    assert image.shape == (504, 760, 3)
    _snake_case    :  Optional[Any]					=					np.array([0.9_358, 0.9_397, 0.9_599, 0.9_901, 1.0_000, 1.0_000, 0.9_882, 1.0_000, 1.0_000])

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







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



   '''simple docstring'''






   def      UpperCamelCase_  (					self							:			int)				->				Optional[Any]:

    """simple docstring"""







    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


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

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
        """/img2img/sketch-mountains-input.jpg""")
    _snake_case    :  Union[str, Any]					=					init_image.resize((768, 512))
    _snake_case    :  str					=					load_numpy(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy""")

    _snake_case    :  Optional[int]					=					"""BAAI/AltDiffusion"""
    _snake_case    :  Any					=					AltDiffusionImgaImgPipeline.from_pretrained(
        lowercase_				,		safety_checker=lowercase_				,		)
    pipe.to(lowercase_)
    pipe.set_progress_bar_config(disable=lowercase_)
    pipe.enable_attention_slicing()

    _snake_case    :  List[Any]					=					"""A fantasy landscape, trending on artstation"""

    _snake_case    :  int					=					torch.manual_seed(0)
    _snake_case    :  Tuple					=					pipe(
        prompt=lowercase_				,		image=lowercase_				,		strength=0.75				,		guidance_scale=7.5				,		generator=lowercase_				,		output_type="""np"""				,		)
    _snake_case    :  str					=					output.images[0]

    assert image.shape == (512, 768, 3)
    # img2img is flaky across GPUs even in fp32, so using MAE here
    assert np.abs(expected_image - image).max() < 1E-2





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
