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


__UpperCamelCase	: List[str]    		=  logging.get_logger(__name__)


# General docstring
__UpperCamelCase	: Dict    		=  'MobileNetV1Config'

# Base docstring
__UpperCamelCase	: Tuple    		=  'google/mobilenet_v1_1.0_224'
__UpperCamelCase	: Dict    		=  [1, 1024, 7, 7]

# Image classification docstring
__UpperCamelCase	: int    		=  'google/mobilenet_v1_1.0_224'
__UpperCamelCase	: List[str]    		=  'tabby, tabby cat'


__UpperCamelCase	: Tuple    		=  [
    '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    A			(					_lowercase	,						_lowercase	,						_lowercase=None		):
						SCREAMING_SNAKE_CASE							:					Optional[int]           =							{}

						if isinstance(snake_case_	,						snake_case_		):
												SCREAMING_SNAKE_CASE							:					str           =							model.mobilenet_va
						else:
												SCREAMING_SNAKE_CASE							:					Any           =							model

						SCREAMING_SNAKE_CASE							:					List[str]           =							"""MobilenetV1/Conv2d_0/"""
						SCREAMING_SNAKE_CASE							:					Dict           =							backbone.conv_stem.convolution.weight
						SCREAMING_SNAKE_CASE							:					Tuple           =							backbone.conv_stem.normalization.bias
						SCREAMING_SNAKE_CASE							:					int           =							backbone.conv_stem.normalization.weight
						SCREAMING_SNAKE_CASE							:					int           =							backbone.conv_stem.normalization.running_mean
						SCREAMING_SNAKE_CASE							:					Tuple           =							backbone.conv_stem.normalization.running_var

						for i in range(13		):
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							i + 1
												SCREAMING_SNAKE_CASE							:					Tuple           =							i * 2

												SCREAMING_SNAKE_CASE							:					int           =							backbone.layer[pt_index]
												SCREAMING_SNAKE_CASE							:					List[str]           =							f"""MobilenetV1/Conv2d_{tf_index}_depthwise/"""
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							pointer.convolution.weight
												SCREAMING_SNAKE_CASE							:					int           =							pointer.normalization.bias
												SCREAMING_SNAKE_CASE							:					int           =							pointer.normalization.weight
												SCREAMING_SNAKE_CASE							:					Tuple           =							pointer.normalization.running_mean
												SCREAMING_SNAKE_CASE							:					Any           =							pointer.normalization.running_var

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

						if isinstance(snake_case_	,						snake_case_		):
												SCREAMING_SNAKE_CASE							:					str           =							"""MobilenetV1/Logits/Conv2d_1c_1x1/"""
												SCREAMING_SNAKE_CASE							:					Dict           =							model.classifier.weight
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							model.classifier.bias

						return tf_to_pt_map
def    A			(					_lowercase	,						_lowercase	,						_lowercase		):
						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
						SCREAMING_SNAKE_CASE							:					Optional[Any]           =							tf.train.list_variables(snake_case_		)
						SCREAMING_SNAKE_CASE							:					List[str]           =							{}
						for name, shape in init_vars:
												logger.info(f"""Loading TF weight {name} with shape {shape}"""		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							tf.train.load_variable(snake_case_	,						snake_case_		)
												SCREAMING_SNAKE_CASE							:					Tuple           =							array

						# Build TF to PyTorch weights loading map
						SCREAMING_SNAKE_CASE							:					Any           =							_build_tf_to_pytorch_map(snake_case_	,						snake_case_	,						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

												SCREAMING_SNAKE_CASE							:					Tuple           =							tf_weights[name]

												if "depthwise_weights" in name:
																		logger.info('''Transposing depthwise'''		)
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							np.transpose(snake_case_	,						(2, 3, 0, 1)		)
												elif "weights" in name:
																		logger.info('''Transposing'''		)
																		if len(pointer.shape		) == 2:  # copying into linear layer
																								SCREAMING_SNAKE_CASE							:					Any           =							array.squeeze().transpose()
																		else:
																								SCREAMING_SNAKE_CASE							:					Optional[int]           =							np.transpose(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}"""		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							torch.from_numpy(snake_case_		)

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

						logger.info(f"""Weights not copied to PyTorch model: {', '.join(tf_weights.keys()		)}"""		)
						return model
def    A			(					_lowercase	,						_lowercase		):
						SCREAMING_SNAKE_CASE							:					Tuple           =							features.shape[-2:]
						SCREAMING_SNAKE_CASE							:					Any           =							conv_layer.stride
						SCREAMING_SNAKE_CASE							:					Dict           =							conv_layer.kernel_size

						if in_height % stride_height == 0:
												SCREAMING_SNAKE_CASE							:					Tuple           =							max(kernel_height - stride_height	,						0		)
						else:
												SCREAMING_SNAKE_CASE							:					Dict           =							max(kernel_height - (in_height % stride_height)	,						0		)

						if in_width % stride_width == 0:
												SCREAMING_SNAKE_CASE							:					Any           =							max(kernel_width - stride_width	,						0		)
						else:
												SCREAMING_SNAKE_CASE							:					Tuple           =							max(kernel_width - (in_width % stride_width)	,						0		)

						SCREAMING_SNAKE_CASE							:					Dict           =							pad_along_width // 2
						SCREAMING_SNAKE_CASE							:					Any           =							pad_along_width - pad_left
						SCREAMING_SNAKE_CASE							:					Optional[int]           =							pad_along_height // 2
						SCREAMING_SNAKE_CASE							:					Tuple           =							pad_along_height - pad_top

						SCREAMING_SNAKE_CASE							:					List[str]           =							(pad_left, pad_right, pad_top, pad_bottom)
						return nn.functional.pad(snake_case_	,						snake_case_	,						'''constant'''	,						0.0		)





class       lowercase__		(     nn.Module):


						def __init__( self			:       Any					,	UpperCamelCase__			:       List[str]					,	UpperCamelCase__			:       str					,	UpperCamelCase__			:       List[Any]					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       Tuple = 1					,	UpperCamelCase__			:       List[str] = 1					,	UpperCamelCase__			:       List[Any] = False					,	UpperCamelCase__			:       str = True					,	UpperCamelCase__			:       Optional[Any] = True					,	):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					List[Any]           =							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."""		)

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

												SCREAMING_SNAKE_CASE							:					List[str]           =							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:
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							nn.BatchNormad(
																		    num_features=_a					,	eps=config.layer_norm_eps					,	momentum=0.9997					,	affine=_a					,	track_running_stats=_a					,	)
												else:
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							None

												if use_activation:
																		if isinstance(_a					,	_a		):
																								SCREAMING_SNAKE_CASE							:					List[str]           =							ACTaFN[use_activation]
																		elif isinstance(config.hidden_act					,	_a		):
																								SCREAMING_SNAKE_CASE							:					List[str]           =							ACTaFN[config.hidden_act]
																		else:
																								SCREAMING_SNAKE_CASE							:					Optional[int]           =							config.hidden_act
												else:
																		SCREAMING_SNAKE_CASE							:					Tuple           =							None



						def  __A ( self			:       int					,	UpperCamelCase__			:       Dict		):





												'''simple docstring'''





												if self.config.tf_padding:
																		SCREAMING_SNAKE_CASE							:					Tuple           =							apply_tf_padding(_a					,	self.convolution		)
												SCREAMING_SNAKE_CASE							:					Tuple           =							self.convolution(_a		)
												if self.normalization is not None:
																		SCREAMING_SNAKE_CASE							:					Dict           =							self.normalization(_a		)
												if self.activation is not None:
																		SCREAMING_SNAKE_CASE							:					str           =							self.activation(_a		)
												return features






class       lowercase__		(     UpperCamelCase__):
						UpperCamelCase_  =	MobileNetVaConfig
						UpperCamelCase_  =	load_tf_weights_in_mobilenet_va
						UpperCamelCase_  =	"""mobilenet_v1"""
						UpperCamelCase_  =	"""pixel_values"""
						UpperCamelCase_  =	False


						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       List[str]		):





												'''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		)






__UpperCamelCase	: Optional[int]    		=  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'

__UpperCamelCase	: 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."""			,							UpperCamelCase__			,							)
class       lowercase__		(     UpperCamelCase__):


						def __init__( self			:       Dict					,	UpperCamelCase__			:       Optional[Any]					,	UpperCamelCase__			:       Optional[int] = True		):





												'''simple docstring'''





												super().__init__(_a		)
												SCREAMING_SNAKE_CASE							:					Any           =							config

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							32
												SCREAMING_SNAKE_CASE							:					str           =							max(int(depth * config.depth_multiplier		)					,	config.min_depth		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							MobileNetVaConvLayer(
												    _a					,	in_channels=config.num_channels					,	out_channels=_a					,	kernel_size=3					,	stride=2					,	)

												SCREAMING_SNAKE_CASE							:					Dict           =							[1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

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

																		if strides[i] == 2 or i == 0:
																								depth *= 2
																								SCREAMING_SNAKE_CASE							:					int           =							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					,	)		)

												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							nn.AdaptiveAvgPoolad((1, 1)		) if add_pooling_layer else None

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


						def  __A ( self			:       str					,	UpperCamelCase__			:       Dict		):





												'''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  __A ( self			:       Any					,	UpperCamelCase__			:       str = None					,	UpperCamelCase__			:       Union[str, Any] = None					,	UpperCamelCase__			:       str = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							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'''		)

												SCREAMING_SNAKE_CASE							:					str           =							self.conv_stem(_a		)

												SCREAMING_SNAKE_CASE							:					int           =							() if output_hidden_states else None

												for i, layer_module in enumerate(self.layer		):
																		SCREAMING_SNAKE_CASE							:					str           =							layer_module(_a		)

																		if output_hidden_states:
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							all_hidden_states + (hidden_states,)

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							hidden_states

												if self.pooler is not None:
																		SCREAMING_SNAKE_CASE							:					int           =							torch.flatten(self.pooler(_a		)					,	start_dim=1		)
												else:
																		SCREAMING_SNAKE_CASE							:					Dict           =							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    """			,							UpperCamelCase__			,							)
class       lowercase__		(     UpperCamelCase__):


						def __init__( self			:       Dict					,	UpperCamelCase__			:       List[Any]		):





												'''simple docstring'''





												super().__init__(_a		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							config.num_labels
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							MobileNetVaModel(_a		)

												SCREAMING_SNAKE_CASE							:					List[str]           =							self.mobilenet_va.layer[-1].convolution.out_channels

												# Classifier head
												SCREAMING_SNAKE_CASE							:					List[Any]           =							nn.Dropout(config.classifier_dropout_prob					,	inplace=_a		)
												SCREAMING_SNAKE_CASE							:					Optional[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  __A ( self			:       str					,	UpperCamelCase__			:       List[str] = None					,	UpperCamelCase__			:       str = None					,	UpperCamelCase__			:       str = None					,	UpperCamelCase__			:       Dict = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Tuple           =							return_dict if return_dict is not None else self.config.use_return_dict

												SCREAMING_SNAKE_CASE							:					Tuple           =							self.mobilenet_va(_a					,	output_hidden_states=_a					,	return_dict=_a		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							outputs.pooler_output if return_dict else outputs[1]

												SCREAMING_SNAKE_CASE							:					Any           =							self.classifier(self.dropout(_a		)		)

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

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

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

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






import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin





@skip_mps
class       lowercase__		(     UpperCamelCase_			,							UpperCamelCase_			,							unittest.TestCase):
						UpperCamelCase_  =	IFPipeline
						UpperCamelCase_  =	TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
						UpperCamelCase_  =	TEXT_TO_IMAGE_BATCH_PARAMS
						UpperCamelCase_  =	PipelineTesterMixin.required_optional_params - {"""latents"""}


						def  __A ( self			:       Tuple		):





												'''simple docstring'''





												return self._get_dummy_components()


						def  __A ( self			:       int					,	UpperCamelCase__			:       Dict					,	UpperCamelCase__			:       int=0		):





												'''simple docstring'''





												if str(UpperCamelCase__		).startswith('''mps'''		):
																		SCREAMING_SNAKE_CASE							:					Tuple           =							torch.manual_seed(UpperCamelCase__		)
												else:
																		SCREAMING_SNAKE_CASE							:					Dict           =							torch.Generator(device=UpperCamelCase__		).manual_seed(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					int           =							{
												    '''prompt''': '''A painting of a squirrel eating a burger''',
												    '''generator''': generator,
												    '''num_inference_steps''': 2,
												    '''output_type''': '''numpy''',
												}

												return inputs


						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												self._test_save_load_optional_components()


						@unittest.skipIf(torch_device != '''cuda'''					,	reason='''float16 requires CUDA'''		)
						def  __A ( self			:       Any		):





												'''simple docstring'''





												super().test_save_load_floataa(expected_max_diff=1E-1		)


						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												self._test_attention_slicing_forward_pass(expected_max_diff=1E-2		)


						def  __A ( self			:       List[Any]		):





												'''simple docstring'''





												self._test_save_load_local()


						def  __A ( self			:       List[str]		):





												'''simple docstring'''





												self._test_inference_batch_single_identical(
												    expected_max_diff=1E-2					,	)



						@unittest.skipIf(
						    torch_device != '''cuda''' or not is_xformers_available()					,	reason='''XFormers attention is only available with CUDA and `xformers` installed'''					,	)
						def  __A ( self			:       Tuple		):





												'''simple docstring'''





												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3		)






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


						def  __A ( self			:       Optional[Any]		):





												'''simple docstring'''





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


						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''					,	variant='''fp16'''					,	torch_dtype=torch.floataa		)

												SCREAMING_SNAKE_CASE							:					str           =							IFSuperResolutionPipeline.from_pretrained(
												    '''DeepFloyd/IF-II-L-v1.0'''					,	variant='''fp16'''					,	torch_dtype=torch.floataa					,	text_encoder=UpperCamelCase__					,	tokenizer=UpperCamelCase__		)

												# pre compute text embeddings and remove T5 to save memory

												pipe_a.text_encoder.to('''cuda'''		)

												SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE							:					Dict           =							pipe_a.encode_prompt('''anime turtle'''					,	device='''cuda'''		)

												del pipe_a.tokenizer
												del pipe_a.text_encoder
												gc.collect()

												SCREAMING_SNAKE_CASE							:					Tuple           =							None
												SCREAMING_SNAKE_CASE							:					str           =							None

												pipe_a.enable_model_cpu_offload()
												pipe_a.enable_model_cpu_offload()

												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()		)
												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()		)

												self._test_if(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)

												pipe_a.remove_all_hooks()
												pipe_a.remove_all_hooks()

												# img2img

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							IFImgaImgPipeline(**pipe_a.components		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							IFImgaImgSuperResolutionPipeline(**pipe_a.components		)

												pipe_a.enable_model_cpu_offload()
												pipe_a.enable_model_cpu_offload()

												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()		)
												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()		)

												self._test_if_imgaimg(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)

												pipe_a.remove_all_hooks()
												pipe_a.remove_all_hooks()

												# inpainting

												SCREAMING_SNAKE_CASE							:					Tuple           =							IFInpaintingPipeline(**pipe_a.components		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							IFInpaintingSuperResolutionPipeline(**pipe_a.components		)

												pipe_a.enable_model_cpu_offload()
												pipe_a.enable_model_cpu_offload()

												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()		)
												pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()		)

												self._test_if_inpainting(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)


						def  __A ( self			:       Tuple					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       Dict		):





												'''simple docstring'''





												_start_torch_memory_measurement()

												SCREAMING_SNAKE_CASE							:					List[str]           =							torch.Generator(device='''cpu'''		).manual_seed(0		)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	num_inference_steps=2					,	generator=UpperCamelCase__					,	output_type='''np'''					,	)

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					Tuple           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 13 * 10**9

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)

												# pipeline 2

												_start_torch_memory_measurement()

												SCREAMING_SNAKE_CASE							:					Tuple           =							torch.Generator(device='''cpu'''		).manual_seed(0		)

												SCREAMING_SNAKE_CASE							:					int           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Any           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	generator=UpperCamelCase__					,	num_inference_steps=2					,	output_type='''np'''					,	)

												SCREAMING_SNAKE_CASE							:					Tuple           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 4 * 10**9

												SCREAMING_SNAKE_CASE							:					int           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)


						def  __A ( self			:       Optional[int]					,	UpperCamelCase__			:       str					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       Optional[Any]					,	UpperCamelCase__			:       str		):





												'''simple docstring'''





												_start_torch_memory_measurement()

												SCREAMING_SNAKE_CASE							:					Tuple           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Any           =							torch.Generator(device='''cpu'''		).manual_seed(0		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	num_inference_steps=2					,	generator=UpperCamelCase__					,	output_type='''np'''					,	)

												SCREAMING_SNAKE_CASE							:					int           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 10 * 10**9

												SCREAMING_SNAKE_CASE							:					Tuple           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)

												# pipeline 2

												_start_torch_memory_measurement()

												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							torch.Generator(device='''cpu'''		).manual_seed(0		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							floats_tensor((1, 3, 256, 256)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					str           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	original_image=UpperCamelCase__					,	generator=UpperCamelCase__					,	num_inference_steps=2					,	output_type='''np'''					,	)

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					List[Any]           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 4 * 10**9

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)



						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       Optional[Any]		):





												'''simple docstring'''





												_start_torch_memory_measurement()

												SCREAMING_SNAKE_CASE							:					str           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(1		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							torch.Generator(device='''cpu'''		).manual_seed(0		)
												SCREAMING_SNAKE_CASE							:					int           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	mask_image=UpperCamelCase__					,	num_inference_steps=2					,	generator=UpperCamelCase__					,	output_type='''np'''					,	)

												SCREAMING_SNAKE_CASE							:					Any           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					Any           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 10 * 10**9

												SCREAMING_SNAKE_CASE							:					List[Any]           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)

												# pipeline 2

												_start_torch_memory_measurement()

												SCREAMING_SNAKE_CASE							:					str           =							torch.Generator(device='''cpu'''		).manual_seed(0		)

												SCREAMING_SNAKE_CASE							:					List[str]           =							floats_tensor((1, 3, 64, 64)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					int           =							floats_tensor((1, 3, 256, 256)					,	rng=random.Random(0		)		).to(UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							floats_tensor((1, 3, 256, 256)					,	rng=random.Random(1		)		).to(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Dict           =							pipe_a(
												    prompt_embeds=UpperCamelCase__					,	negative_prompt_embeds=UpperCamelCase__					,	image=UpperCamelCase__					,	mask_image=UpperCamelCase__					,	original_image=UpperCamelCase__					,	generator=UpperCamelCase__					,	num_inference_steps=2					,	output_type='''np'''					,	)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							output.images[0]

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

												SCREAMING_SNAKE_CASE							:					Any           =							torch.cuda.max_memory_allocated()
												assert mem_bytes < 4 * 10**9

												SCREAMING_SNAKE_CASE							:					Tuple           =							load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'''		)
												assert_mean_pixel_difference(UpperCamelCase__					,	UpperCamelCase__		)






def    A			(					):
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()





