




'''simple docstring'''





import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()




class  _snake_case		(	unittest.TestCase							):




			@property
			def    lowerCAmelCase__				(					self							)				->				Any:

										'''simple docstring'''




										torch.manual_seed(0							)
										snake_case_						=      UNetaDModel(
										    block_out_channels=(32, 64)				,							layers_per_block=2				,							sample_size=32				,							in_channels=3				,							out_channels=3				,							down_block_types=("DownBlock2D", "AttnDownBlock2D")				,							up_block_types=("AttnUpBlock2D", "UpBlock2D")				,							)
										return model




			def    lowerCAmelCase__				(					self							)				->				Optional[Any]:

										'''simple docstring'''




										snake_case_						=      self.dummy_uncond_unet
										snake_case_						=      KarrasVeScheduler()

										snake_case_						=      KarrasVePipeline(unet=a__				,							scheduler=a__							)
										pipe.to(a__							)
										pipe.set_progress_bar_config(disable=a__							)

										snake_case_						=      torch.manual_seed(0							)
										snake_case_						=      pipe(num_inference_steps=2				,							generator=a__				,							output_type="numpy"							).images

										snake_case_						=      torch.manual_seed(0							)
										snake_case_						=      pipe(num_inference_steps=2				,							generator=a__				,							output_type="numpy"				,							return_dict=a__							)[0]

										snake_case_						=      image[0, -3:, -3:, -1]
										snake_case_						=      image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 32, 32, 3)
										snake_case_						=      np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]							)

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




@slow
@require_torch
class  _snake_case		(	unittest.TestCase							):




			def    lowerCAmelCase__				(					self							)				->				int:

										'''simple docstring'''




										snake_case_						=      "google/ncsnpp-celebahq-256"
										snake_case_						=      UNetaDModel.from_pretrained(a__							)
										snake_case_						=      KarrasVeScheduler()

										snake_case_						=      KarrasVePipeline(unet=a__				,							scheduler=a__							)
										pipe.to(a__							)
										pipe.set_progress_bar_config(disable=a__							)

										snake_case_						=      torch.manual_seed(0							)
										snake_case_						=      pipe(num_inference_steps=20				,							generator=a__				,							output_type="numpy"							).images

										snake_case_						=      image[0, -3:, -3:, -1]
										assert image.shape == (1, 256, 256, 3)
										snake_case_						=      np.array([0.5_7_8, 0.5_8_1_1, 0.5_9_2_4, 0.5_8_0_9, 0.5_8_7, 0.5_8_8_6, 0.5_8_6_1, 0.5_8_0_2, 0.5_8_6]							)

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












"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__A		:       List[Any]  				=		{
    '''configuration_mega''': ['''MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MegaConfig''', '''MegaOnnxConfig'''],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__A		:       List[str]  				=		[
						    '''MEGA_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''MegaForCausalLM''',
						    '''MegaForMaskedLM''',
						    '''MegaForMultipleChoice''',
						    '''MegaForQuestionAnswering''',
						    '''MegaForSequenceClassification''',
						    '''MegaForTokenClassification''',
						    '''MegaModel''',
						    '''MegaPreTrainedModel''',
						]

if TYPE_CHECKING:
						from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_mega import (
												    MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
												    MegaForCausalLM,
												    MegaForMaskedLM,
												    MegaForMultipleChoice,
												    MegaForQuestionAnswering,
												    MegaForSequenceClassification,
												    MegaForTokenClassification,
												    MegaModel,
												    MegaPreTrainedModel,
												)

else:
						import sys

						__A		:       Any  				=		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



