



'''simple docstring'''



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class 							__UpperCamelCase    (		lowerCAmelCase_		,				unittest.TestCase ):
	A_           =   KandinskyInpaintPipeline
	A_           =   ["prompt", "image_embeds", "negative_image_embeds", "image", "mask_image"]
	A_           =   [
	    "prompt",
	    "negative_prompt",
	    "image_embeds",
	    "negative_image_embeds",
	    "image",
	    "mask_image",
	]
	A_           =   [
	    "generator",
	    "height",
	    "width",
	    "latents",
	    "guidance_scale",
	    "negative_prompt",
	    "num_inference_steps",
	    "return_dict",
	    "guidance_scale",
	    "num_images_per_prompt",
	    "output_type",
	    "return_dict",
	]
	A_           =   False




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						return 32




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						return 32




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						return self.time_input_dim




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						return self.time_input_dim * 4




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						return 100




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Dict    =  XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'							)
						return tokenizer




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						torch.manual_seed(0							)
						__a      :	Union[str, Any]    =  MCLIPConfig(
						    numDims=self.cross_attention_dim						,    transformerDimensions=self.text_embedder_hidden_size						,    hidden_size=self.text_embedder_hidden_size						,    intermediate_size=37						,    num_attention_heads=4						,    num_hidden_layers=5						,    vocab_size=1005						,    )

						__a      :	List[Any]    =  MultilingualCLIP(__a							)
						__a      :	int    =  text_encoder.eval()

						return text_encoder




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						torch.manual_seed(0							)

						__a      :	Optional[Any]    =  {
						    'in_channels': 9,
						    # Out channels is double in channels because predicts mean and variance
						    'out_channels': 8,
						    'addition_embed_type': 'text_image',
						    'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
						    'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
						    'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
						    'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
						    'layers_per_block': 1,
						    'encoder_hid_dim': self.text_embedder_hidden_size,
						    'encoder_hid_dim_type': 'text_image_proj',
						    'cross_attention_dim': self.cross_attention_dim,
						    'attention_head_dim': 4,
						    'resnet_time_scale_shift': 'scale_shift',
						    'class_embed_type': None,
						}

						__a      :	str    =  UNetaDConditionModel(**__a							)
						return model




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						return {
						    "block_out_channels": [32, 64],
						    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
						    "in_channels": 3,
						    "latent_channels": 4,
						    "layers_per_block": 1,
						    "norm_num_groups": 8,
						    "norm_type": "spatial",
						    "num_vq_embeddings": 12,
						    "out_channels": 3,
						    "up_block_types": [
						        "AttnUpDecoderBlock2D",
						        "UpDecoderBlock2D",
						    ],
						    "vq_embed_dim": 4,
						}




	@property
	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						torch.manual_seed(0							)
						__a      :	List[Any]    =  VQModel(**self.dummy_movq_kwargs							)
						return model




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	Tuple    =  self.dummy_text_encoder
						__a      :	List[str]    =  self.dummy_tokenizer
						__a      :	Optional[Any]    =  self.dummy_unet
						__a      :	str    =  self.dummy_movq

						__a      :	str    =  DDIMScheduler(
						    num_train_timesteps=1000						,    beta_schedule='linear'						,    beta_start=0.00085						,    beta_end=0.012						,    clip_sample=__a						,    set_alpha_to_one=__a						,    steps_offset=1						,    prediction_type='epsilon'						,    thresholding=__a						,    )

						__a      :	Optional[int]    =  {
						    'text_encoder': text_encoder,
						    'tokenizer': tokenizer,
						    'unet': unet,
						    'scheduler': scheduler,
						    'movq': movq,
						}

						return components




	def 				__UpperCAmelCase							(       self						,    __a						,    __a=0							):

						'''simple docstring'''





						__a      :	List[Any]    =  floats_tensor((1, self.cross_attention_dim)						,    rng=random.Random(__a							)							).to(__a							)
						__a      :	Optional[int]    =  floats_tensor((1, self.cross_attention_dim)						,    rng=random.Random(seed + 1							)							).to(__a							)
						# create init_image
						__a      :	List[str]    =  floats_tensor((1, 3, 64, 64)						,    rng=random.Random(__a							)							).to(__a							)
						__a      :	Optional[Any]    =  image.cpu().permute(0						,    2						,    3						,    1							)[0]
						__a      :	str    =  Image.fromarray(np.uinta(__a							)							).convert('RGB'							).resize((256, 256)							)
						# create mask
						__a      :	Tuple    =  np.ones((64, 64)						,    dtype=np.floataa							)
						__a      :	Any    =  0

						if str(__a							).startswith('mps'							):
											__a      :	str    =  torch.manual_seed(__a							)
						else:
											__a      :	str    =  torch.Generator(device=__a							).manual_seed(__a							)
						__a      :	List[Any]    =  {
						    'prompt': 'horse',
						    'image': init_image,
						    'mask_image': mask,
						    'image_embeds': image_embeds,
						    'negative_image_embeds': negative_image_embeds,
						    'generator': generator,
						    'height': 64,
						    'width': 64,
						    'num_inference_steps': 2,
						    'guidance_scale': 4.0,
						    'output_type': 'np',
						}
						return inputs




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	List[Any]    =  'cpu'

						__a      :	List[Any]    =  self.get_dummy_components()

						__a      :	str    =  self.pipeline_class(**__a							)
						__a      :	List[str]    =  pipe.to(__a							)

						pipe.set_progress_bar_config(disable=__a							)

						__a      :	Tuple    =  pipe(**self.get_dummy_inputs(__a							)							)
						__a      :	Any    =  output.images

						__a      :	Optional[int]    =  pipe(
						    **self.get_dummy_inputs(__a							)						,    return_dict=__a						,    )[0]

						__a      :	int    =  image[0, -3:, -3:, -1]
						__a      :	Tuple    =  image_from_tuple[0, -3:, -3:, -1]

						print(f"""image.shape {image.shape}"""							)

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

						__a      :	Union[str, Any]    =  np.array(
						    [0.8326919, 0.73790467, 0.20918581, 0.9309612, 0.5511791, 0.43713328, 0.5513321, 0.49922934, 0.59497786]							)

						assert (
						    np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2
						), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
						assert (
						    np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1E-2
						), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""





	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						super().test_inference_batch_single_identical(expected_max_diff=3E-3							)





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




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





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





	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	List[str]    =  load_numpy(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
						    '/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy'							)

						__a      :	Optional[Any]    =  load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'							)
						__a      :	Union[str, Any]    =  np.ones((768, 768)						,    dtype=np.floataa							)
						__a      :	List[str]    =  0

						__a      :	Dict    =  'a hat'

						__a      :	Any    =  KandinskyPriorPipeline.from_pretrained(
						    'kandinsky-community/kandinsky-2-1-prior'						,    torch_dtype=torch.floataa							)
						pipe_prior.to(__a							)

						__a      :	List[str]    =  KandinskyInpaintPipeline.from_pretrained(
						    'kandinsky-community/kandinsky-2-1-inpaint'						,    torch_dtype=torch.floataa							)
						__a      :	Optional[Any]    =  pipeline.to(__a							)
						pipeline.set_progress_bar_config(disable=__a							)

						__a      :	Union[str, Any]    =  torch.Generator(device='cpu'							).manual_seed(0							)
						__a      :	Optional[Any]    =  pipe_prior(
						    __a						,    generator=__a						,    num_inference_steps=5						,    negative_prompt=''						,    ).to_tuple()

						__a      :	Tuple    =  pipeline(
						    __a						,    image=__a						,    mask_image=__a						,    image_embeds=__a						,    negative_image_embeds=__a						,    generator=__a						,    num_inference_steps=100						,    height=768						,    width=768						,    output_type='np'						,    )

						__a      :	Dict    =  output.images[0]

						assert image.shape == (768, 768, 3)

						assert_mean_pixel_difference(__a						,    __a							)







'''simple docstring'''



from typing import Optional, Union

import torch
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_mobilenet_va import MobileNetVaConfig


__lowercase     :						str	=     logging.get_logger(__name__)


# General docstring
__lowercase     :						List[str]	=     'MobileNetV1Config'

# Base docstring
__lowercase     :						Tuple	=     'google/mobilenet_v1_1.0_224'
__lowercase     :						List[Any]	=     [1, 10_24, 7, 7]

# Image classification docstring
__lowercase     :						int	=     'google/mobilenet_v1_1.0_224'
__lowercase     :						Any	=     'tabby, tabby cat'


__lowercase     :						Dict	=     [
    'google/mobilenet_v1_1.0_224',
    'google/mobilenet_v1_0.75_192',
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
]






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       List[Any]						,				_SCREAMING_SNAKE_CASE				:       int						,				_SCREAMING_SNAKE_CASE				:       Optional[Any]=None       ):
					__a      :	Dict    =  {}

					if isinstance(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       ):
										__a      :	Optional[Any]    =  model.mobilenet_va
					else:
										__a      :	List[Any]    =  model

					__a      :	Dict    =  'MobilenetV1/Conv2d_0/'
					__a      :	Dict    =  backbone.conv_stem.convolution.weight
					__a      :	Optional[Any]    =  backbone.conv_stem.normalization.bias
					__a      :	int    =  backbone.conv_stem.normalization.weight
					__a      :	int    =  backbone.conv_stem.normalization.running_mean
					__a      :	Tuple    =  backbone.conv_stem.normalization.running_var

					for i in range(13       ):
										__a      :	int    =  i + 1
										__a      :	Dict    =  i * 2

										__a      :	Dict    =  backbone.layer[pt_index]
										__a      :	Dict    =  F"""MobilenetV1/Conv2d_{tf_index}_depthwise/"""
										__a      :	Union[str, Any]    =  pointer.convolution.weight
										__a      :	Optional[Any]    =  pointer.normalization.bias
										__a      :	Union[str, Any]    =  pointer.normalization.weight
										__a      :	List[Any]    =  pointer.normalization.running_mean
										__a      :	Tuple    =  pointer.normalization.running_var

										__a      :	List[str]    =  backbone.layer[pt_index + 1]
										__a      :	Optional[Any]    =  F"""MobilenetV1/Conv2d_{tf_index}_pointwise/"""
										__a      :	Optional[int]    =  pointer.convolution.weight
										__a      :	List[str]    =  pointer.normalization.bias
										__a      :	Dict    =  pointer.normalization.weight
										__a      :	Dict    =  pointer.normalization.running_mean
										__a      :	Optional[int]    =  pointer.normalization.running_var

					if isinstance(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       ):
										__a      :	Any    =  'MobilenetV1/Logits/Conv2d_1c_1x1/'
										__a      :	Optional[int]    =  model.classifier.weight
										__a      :	List[Any]    =  model.classifier.bias

					return tf_to_pt_map






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       str						,				_SCREAMING_SNAKE_CASE				:       List[str]						,				_SCREAMING_SNAKE_CASE				:       Dict       ):
					try:
										import numpy as np
										import tensorflow as tf
					except ImportError:
										logger.error(
										    'Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see '
										    'https://www.tensorflow.org/install/ for installation instructions.'       )
										raise

					# Load weights from TF model
					__a      :	Union[str, Any]    =  tf.train.list_variables(_SCREAMING_SNAKE_CASE       )
					__a      :	Optional[int]    =  {}
					for name, shape in init_vars:
										logger.info(F"""Loading TF weight {name} with shape {shape}"""       )
										__a      :	List[str]    =  tf.train.load_variable(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
										__a      :	Optional[Any]    =  array

					# Build TF to PyTorch weights loading map
					__a      :	Optional[int]    =  _build_tf_to_pytorch_map(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )

					for name, pointer in tf_to_pt_map.items():
										logger.info(F"""Importing {name}"""       )
										if name not in tf_weights:
															logger.info(F"""{name} not in tf pre-trained weights, skipping"""       )
															continue

										__a      :	Union[str, Any]    =  tf_weights[name]

										if "depthwise_weights" in name:
															logger.info('Transposing depthwise'       )
															__a      :	Optional[Any]    =  np.transpose(_SCREAMING_SNAKE_CASE						,				(2, 3, 0, 1)       )
										elif "weights" in name:
															logger.info('Transposing'       )
															if len(pointer.shape       ) == 2:  # copying into linear layer
																				__a      :	Union[str, Any]    =  array.squeeze().transpose()
															else:
																				__a      :	Dict    =  np.transpose(_SCREAMING_SNAKE_CASE						,				(3, 2, 0, 1)       )

										if pointer.shape != array.shape:
															raise ValueError(F"""Pointer shape {pointer.shape} and array shape {array.shape} mismatched"""       )

										logger.info(F"""Initialize PyTorch weight {name} {array.shape}"""       )
										__a      :	List[str]    =  torch.from_numpy(_SCREAMING_SNAKE_CASE       )

										tf_weights.pop(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE       )
										tf_weights.pop(name + '/RMSProp'						,				_SCREAMING_SNAKE_CASE       )
										tf_weights.pop(name + '/RMSProp_1'						,				_SCREAMING_SNAKE_CASE       )
										tf_weights.pop(name + '/ExponentialMovingAverage'						,				_SCREAMING_SNAKE_CASE       )

					logger.info(F"""Weights not copied to PyTorch model: {", ".join(tf_weights.keys()       )}"""       )
					return model






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE				:       torch.Tensor						,				_SCREAMING_SNAKE_CASE				:       nn.Convad       ):
					__a   ,    __a      :	Any    =  features.shape[-2:]
					__a   ,    __a      :	int    =  conv_layer.stride
					__a   ,    __a      :	Any    =  conv_layer.kernel_size

					if in_height % stride_height == 0:
										__a      :	int    =  max(kernel_height - stride_height						,				0       )
					else:
										__a      :	int    =  max(kernel_height - (in_height % stride_height)						,				0       )

					if in_width % stride_width == 0:
										__a      :	Any    =  max(kernel_width - stride_width						,				0       )
					else:
										__a      :	str    =  max(kernel_width - (in_width % stride_width)						,				0       )

					__a      :	int    =  pad_along_width // 2
					__a      :	Dict    =  pad_along_width - pad_left
					__a      :	List[str]    =  pad_along_height // 2
					__a      :	Union[str, Any]    =  pad_along_height - pad_top

					__a      :	str    =  (pad_left, pad_right, pad_top, pad_bottom)
					return nn.functional.pad(_SCREAMING_SNAKE_CASE						,				_SCREAMING_SNAKE_CASE						,				'constant'						,				0.0       )





class 							__UpperCamelCase    (		nn.Module ):




	def __init__(       self						,    __a						,    __a						,    __a						,    __a						,    __a = 1						,    __a = 1						,    __a = False						,    __a = True						,    __a = True						,    ):

						'''simple docstring'''





						super().__init__()
						__a      :	Optional[int]    =  config

						if in_channels % groups != 0:
											raise ValueError(f"""Input channels ({in_channels}) are not divisible by {groups} groups."""							)
						if out_channels % groups != 0:
											raise ValueError(f"""Output channels ({out_channels}) are not divisible by {groups} groups."""							)

						__a      :	Dict    =  0 if config.tf_padding else int((kernel_size - 1) / 2							)

						__a      :	Union[str, Any]    =  nn.Convad(
						    in_channels=__a						,    out_channels=__a						,    kernel_size=__a						,    stride=__a						,    padding=__a						,    groups=__a						,    bias=__a						,    padding_mode='zeros'						,    )

						if use_normalization:
											__a      :	List[str]    =  nn.BatchNormad(
											    num_features=__a						,    eps=config.layer_norm_eps						,    momentum=0.9997						,    affine=__a						,    track_running_stats=__a						,    )
						else:
											__a      :	Tuple    =  None

						if use_activation:
											if isinstance(__a						,    __a							):
																__a      :	Tuple    =  ACTaFN[use_activation]
											elif isinstance(config.hidden_act						,    __a							):
																__a      :	Union[str, Any]    =  ACTaFN[config.hidden_act]
											else:
																__a      :	Dict    =  config.hidden_act
						else:
											__a      :	List[Any]    =  None





	def 				__UpperCAmelCase							(       self						,    __a							):

						'''simple docstring'''





						if self.config.tf_padding:
											__a      :	Union[str, Any]    =  apply_tf_padding(__a						,    self.convolution							)
						__a      :	Union[str, Any]    =  self.convolution(__a							)
						if self.normalization is not None:
											__a      :	str    =  self.normalization(__a							)
						if self.activation is not None:
											__a      :	Optional[int]    =  self.activation(__a							)
						return features





class 							__UpperCamelCase    (		lowerCAmelCase_ ):
	A_           =   MobileNetVaConfig
	A_           =   load_tf_weights_in_mobilenet_va
	A_           =   "mobilenet_v1"
	A_           =   "pixel_values"
	A_           =   False




	def 				__UpperCAmelCase							(       self						,    __a							):

						'''simple docstring'''





						if isinstance(__a						,    (nn.Linear, nn.Convad)							):
											module.weight.data.normal_(mean=0.0						,    std=self.config.initializer_range							)
											if module.bias is not None:
																module.bias.data.zero_()
						elif isinstance(__a						,    nn.BatchNormad							):
											module.bias.data.zero_()
											module.weight.data.fill_(1.0							)



__lowercase     :						Any	=     R'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

__lowercase     :						Optional[int]	=     R'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`MobileNetV1ImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'





@add_start_docstrings(
    "The bare MobileNetV1 model outputting raw hidden-states without any specific head on top."		,				lowerCAmelCase_		,				)
class 							__UpperCamelCase    (		lowerCAmelCase_ ):




	def __init__(       self						,    __a						,    __a = True							):

						'''simple docstring'''





						super().__init__(__a							)
						__a      :	Optional[int]    =  config

						__a      :	str    =  32
						__a      :	Dict    =  max(int(depth * config.depth_multiplier							)						,    config.min_depth							)

						__a      :	Union[str, Any]    =  MobileNetVaConvLayer(
						    __a						,    in_channels=config.num_channels						,    out_channels=__a						,    kernel_size=3						,    stride=2						,    )

						__a      :	Tuple    =  [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

						__a      :	Any    =  nn.ModuleList()
						for i in range(13							):
											__a      :	Union[str, Any]    =  out_channels

											if strides[i] == 2 or i == 0:
																depth *= 2
																__a      :	List[Any]    =  max(int(depth * config.depth_multiplier							)						,    config.min_depth							)

											self.layer.append(
											    MobileNetVaConvLayer(
											        __a						,    in_channels=__a						,    out_channels=__a						,    kernel_size=3						,    stride=strides[i]						,    groups=__a						,    )							)

											self.layer.append(
											    MobileNetVaConvLayer(
											        __a						,    in_channels=__a						,    out_channels=__a						,    kernel_size=1						,    )							)

						__a      :	Optional[int]    =  nn.AdaptiveAvgPoolad((1, 1)							) if add_pooling_layer else None

						# Initialize weights and apply final processing
						self.post_init()




	def 				__UpperCAmelCase							(       self						,    __a							):

						'''simple docstring'''





						raise NotImplementedError





	@add_start_docstrings_to_model_forward(__a							)
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC						,    output_type=__a						,    config_class=_CONFIG_FOR_DOC						,    modality='vision'						,    expected_output=_EXPECTED_OUTPUT_SHAPE						,    )
	def 				__UpperCAmelCase							(       self						,    __a = None						,    __a = None						,    __a = None						,    ):

						'''simple docstring'''





						__a      :	Dict    =  (
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						__a      :	int    =  return_dict if return_dict is not None else self.config.use_return_dict

						if pixel_values is None:
											raise ValueError('You have to specify pixel_values'							)

						__a      :	Union[str, Any]    =  self.conv_stem(__a							)

						__a      :	Any    =  () if output_hidden_states else None

						for i, layer_module in enumerate(self.layer							):
											__a      :	List[str]    =  layer_module(__a							)

											if output_hidden_states:
																__a      :	List[Any]    =  all_hidden_states + (hidden_states,)

						__a      :	str    =  hidden_states

						if self.pooler is not None:
											__a      :	Union[str, Any]    =  torch.flatten(self.pooler(__a							)						,    start_dim=1							)
						else:
											__a      :	int    =  None

						if not return_dict:
											return tuple(v for v in [last_hidden_state, pooled_output, all_hidden_states] if v is not None							)

						return BaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=__a						,    pooler_output=__a						,    hidden_states=__a						,    )





@add_start_docstrings(
    "\n    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "		,				lowerCAmelCase_		,				)
class 							__UpperCamelCase    (		lowerCAmelCase_ ):




	def __init__(       self						,    __a							):

						'''simple docstring'''





						super().__init__(__a							)

						__a      :	Tuple    =  config.num_labels
						__a      :	Tuple    =  MobileNetVaModel(__a							)

						__a      :	Optional[int]    =  self.mobilenet_va.layer[-1].convolution.out_channels

						# Classifier head
						__a      :	Any    =  nn.Dropout(config.classifier_dropout_prob						,    inplace=__a							)
						__a      :	Any    =  nn.Linear(__a						,    config.num_labels							) if config.num_labels > 0 else nn.Identity()

						# Initialize weights and apply final processing
						self.post_init()





	@add_start_docstrings_to_model_forward(__a							)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT						,    output_type=__a						,    config_class=_CONFIG_FOR_DOC						,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT						,    )
	def 				__UpperCAmelCase							(       self						,    __a = None						,    __a = None						,    __a = None						,    __a = None						,    ):

						'''simple docstring'''





						__a      :	Union[str, Any]    =  return_dict if return_dict is not None else self.config.use_return_dict

						__a      :	Dict    =  self.mobilenet_va(__a						,    output_hidden_states=__a						,    return_dict=__a							)

						__a      :	List[str]    =  outputs.pooler_output if return_dict else outputs[1]

						__a      :	int    =  self.classifier(self.dropout(__a							)							)

						__a      :	Tuple    =  None
						if labels is not None:
											if self.config.problem_type is None:
																if self.num_labels == 1:
																					__a      :	str    =  'regression'
																elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																					__a      :	int    =  'single_label_classification'
																else:
																					__a      :	Optional[Any]    =  'multi_label_classification'

											if self.config.problem_type == "regression":
																__a      :	Optional[Any]    =  MSELoss()
																if self.num_labels == 1:
																					__a      :	List[Any]    =  loss_fct(logits.squeeze()						,    labels.squeeze()							)
																else:
																					__a      :	Any    =  loss_fct(__a						,    __a							)
											elif self.config.problem_type == "single_label_classification":
																__a      :	List[str]    =  CrossEntropyLoss()
																__a      :	str    =  loss_fct(logits.view(-1						,    self.num_labels							)						,    labels.view(-1							)							)
											elif self.config.problem_type == "multi_label_classification":
																__a      :	Tuple    =  BCEWithLogitsLoss()
																__a      :	Optional[int]    =  loss_fct(__a						,    __a							)

						if not return_dict:
											__a      :	List[Any]    =  (logits,) + outputs[2:]
											return ((loss,) + output) if loss is not None else output

						return ImageClassifierOutputWithNoAttention(
						    loss=__a						,    logits=__a						,    hidden_states=outputs.hidden_states						,    )



