
'''simple docstring'''


import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest

class 		__lowercase			(					_lowercase      ):
			lowerCamelCase     :						Tuple		 =     (VQDiffusionScheduler,)

			def     UpperCAmelCase__	(self     ,				**A ):
								lowerCamelCase_   : str       =						{
								    '''num_vec_classes''': 4_0_9_7,
								    '''num_train_timesteps''': 1_0_0,
								}

								config.update(**A )
								return config

			def     UpperCAmelCase__	(self     ,				A ):
								lowerCamelCase_   : Dict       =						4
								lowerCamelCase_   : Optional[Any]       =						8
								lowerCamelCase_   : List[Any]       =						8

								lowerCamelCase_   : str       =						torch.randint(0     ,				A     ,				(batch_size, height * width) )

								return sample

			@property
			def     UpperCAmelCase__	(self ):
								assert False

			def     UpperCAmelCase__	(self     ,				A ):
								def model(A     ,				A     ,				*A ):
													lowerCamelCase_, lowerCamelCase_   : List[Any]       =						sample.shape
													lowerCamelCase_   : Optional[Any]       =						torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels) )
													lowerCamelCase_   : Tuple       =						F.log_softmax(logits.double()     ,				dim=1 ).float()
													return return_value

								return model

			def     UpperCAmelCase__	(self ):
								for timesteps in [2, 5, 1_0_0, 1_0_0_0]:
													self.check_over_configs(num_train_timesteps=A )

			def     UpperCAmelCase__	(self ):
								for num_vec_classes in [5, 1_0_0, 1_0_0_0, 4_0_0_0]:
													self.check_over_configs(num_vec_classes=A )

			def     UpperCAmelCase__	(self ):
								for t in [0, 5_0, 9_9]:
													self.check_over_forward(time_step=A )






			def     UpperCAmelCase__	(self ):
								pass




'''simple docstring'''


from typing import TYPE_CHECKING

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


__lowercase							:				List[str]					=       {
    '''configuration_sam''': [
        '''SAM_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''SamConfig''',
        '''SamMaskDecoderConfig''',
        '''SamPromptEncoderConfig''',
        '''SamVisionConfig''',
    ],
    '''processing_sam''': ['''SamProcessor'''],
}


try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__lowercase							:				List[Any]					=       [
					    '''SAM_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''SamModel''',
					    '''SamPreTrainedModel''',
					]
try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__lowercase							:				int					=       [
					    '''TF_SAM_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''TFSamModel''',
					    '''TFSamPreTrainedModel''',
					]
try:
					if not is_vision_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__lowercase							:				str					=       ['''SamImageProcessor''']


if TYPE_CHECKING:
					from .configuration_sam import (
					    SAM_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    SamConfig,
					    SamMaskDecoderConfig,
					    SamPromptEncoderConfig,
					    SamVisionConfig,
					)
					from .processing_sam import SamProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_sam import SAM_PRETRAINED_MODEL_ARCHIVE_LIST, SamModel, SamPreTrainedModel

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_sam import TF_SAM_PRETRAINED_MODEL_ARCHIVE_LIST, TFSamModel, TFSamPreTrainedModel

					try:
										if not is_vision_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .image_processing_sam import SamImageProcessor

else:
					import sys

					__lowercase							:				List[str]					=       _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


