


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


lowerCAmelCase							=						logging.get_logger(__name__)


# General docstring
lowerCAmelCase							=						"""MobileNetV1Config"""

# Base docstring
lowerCAmelCase							=						"""google/mobilenet_v1_1.0_224"""
lowerCAmelCase							=						[1, 1_024, 7, 7]

# Image classification docstring
lowerCAmelCase							=						"""google/mobilenet_v1_1.0_224"""
lowerCAmelCase							=						"""tabby, tabby cat"""


lowerCAmelCase							=						[
    """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 						__SCREAMING_SNAKE_CASE       (					lowercase_			,					lowercase_			,					lowercase_=None			)		->    str:




						'''simple docstring'''







						__UpperCAmelCase :       Optional[int]  						=  {}

						if isinstance(lowercase_			,					lowercase_			):
												__UpperCAmelCase :       Dict  						=  model.mobilenet_va
						else:
												__UpperCAmelCase :       List[str]  						=  model

						__UpperCAmelCase :       Union[str, Any]  						=  '''MobilenetV1/Conv2d_0/'''
						__UpperCAmelCase :       Union[str, Any]  						=  backbone.conv_stem.convolution.weight
						__UpperCAmelCase :       Optional[Any]  						=  backbone.conv_stem.normalization.bias
						__UpperCAmelCase :       Tuple  						=  backbone.conv_stem.normalization.weight
						__UpperCAmelCase :       Optional[int]  						=  backbone.conv_stem.normalization.running_mean
						__UpperCAmelCase :       int  						=  backbone.conv_stem.normalization.running_var

						for i in range(13			):
												__UpperCAmelCase :       List[str]  						=  i + 1
												__UpperCAmelCase :       Optional[int]  						=  i * 2

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

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

						if isinstance(lowercase_			,					lowercase_			):
												__UpperCAmelCase :       Any  						=  '''MobilenetV1/Logits/Conv2d_1c_1x1/'''
												__UpperCAmelCase :       Any  						=  model.classifier.weight
												__UpperCAmelCase :       Optional[Any]  						=  model.classifier.bias

						return tf_to_pt_map




def 						__SCREAMING_SNAKE_CASE       (					lowercase_			,					lowercase_			,					lowercase_			)		->    Any:




						'''simple docstring'''







						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
						__UpperCAmelCase :       Any  						=  tf.train.list_variables(lowercase_			)
						__UpperCAmelCase :       Union[str, Any]  						=  {}
						for name, shape in init_vars:
												logger.info(f"Loading TF weight {name} with shape {shape}"			)
												__UpperCAmelCase :       Optional[int]  						=  tf.train.load_variable(lowercase_			,					lowercase_			)
												__UpperCAmelCase :       Any  						=  array

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

						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

												__UpperCAmelCase :       Dict  						=  tf_weights[name]

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

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

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




def 						__SCREAMING_SNAKE_CASE       (					lowercase_			,					lowercase_			)		->    torch.Tensor:




						'''simple docstring'''







						__UpperCAmelCase						,			__UpperCAmelCase :       Optional[int]  						=  features.shape[-2:]
						__UpperCAmelCase						,			__UpperCAmelCase :       Union[str, Any]  						=  conv_layer.stride
						__UpperCAmelCase						,			__UpperCAmelCase :       Any  						=  conv_layer.kernel_size

						if in_height % stride_height == 0:
												__UpperCAmelCase :       str  						=  max(kernel_height - stride_height			,					0			)
						else:
												__UpperCAmelCase :       List[str]  						=  max(kernel_height - (in_height % stride_height)			,					0			)

						if in_width % stride_width == 0:
												__UpperCAmelCase :       List[Any]  						=  max(kernel_width - stride_width			,					0			)
						else:
												__UpperCAmelCase :       Tuple  						=  max(kernel_width - (in_width % stride_width)			,					0			)

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

						__UpperCAmelCase :       Union[str, Any]  						=  (pad_left, pad_right, pad_top, pad_bottom)
						return nn.functional.pad(lowercase_			,					lowercase_			,					'''constant'''			,					0.0			)




class        lowerCamelCase      (     nn.Module   ):


	def __init__(						self							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__ = 1							,    lowercase__ = 1							,    lowercase__ = False							,    lowercase__ = True							,    lowercase__ = True							,    ):
							super().__init__()
							__UpperCAmelCase :       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.")

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

							__UpperCAmelCase :       List[str]  						=  nn.Convad(
							    in_channels=lowercase__							,    out_channels=lowercase__							,    kernel_size=lowercase__							,    stride=lowercase__							,    padding=lowercase__							,    groups=lowercase__							,    bias=lowercase__							,    padding_mode='''zeros'''							,    )

							if use_normalization:
													__UpperCAmelCase :       Any  						=  nn.BatchNormad(
													    num_features=lowercase__							,    eps=config.layer_norm_eps							,    momentum=0.9_9_9_7							,    affine=lowercase__							,    track_running_stats=lowercase__							,    )
							else:
													__UpperCAmelCase :       Optional[Any]  						=  None

							if use_activation:
													if isinstance(lowercase__							,    lowercase__):
																			__UpperCAmelCase :       Dict  						=  ACTaFN[use_activation]
													elif isinstance(config.hidden_act							,    lowercase__):
																			__UpperCAmelCase :       Optional[int]  						=  ACTaFN[config.hidden_act]
													else:
																			__UpperCAmelCase :       Optional[int]  						=  config.hidden_act
							else:
													__UpperCAmelCase :       List[Any]  						=  None


	def    A(						self							,    lowercase__):
							if self.config.tf_padding:
													__UpperCAmelCase :       Optional[int]  						=  apply_tf_padding(lowercase__							,    self.convolution)
							__UpperCAmelCase :       Dict  						=  self.convolution(lowercase__)
							if self.normalization is not None:
													__UpperCAmelCase :       int  						=  self.normalization(lowercase__)
							if self.activation is not None:
													__UpperCAmelCase :       str  						=  self.activation(lowercase__)
							return features






class        lowerCamelCase      (     _UpperCamelCase   ):
	_lowerCAmelCase    :		Dict        =	MobileNetVaConfig
	_lowerCAmelCase    :		int        =	load_tf_weights_in_mobilenet_va
	_lowerCAmelCase    :		List[str]        =	'''mobilenet_v1'''
	_lowerCAmelCase    :		Optional[Any]        =	'''pixel_values'''
	_lowerCAmelCase    :		Dict        =	False


	def    A(						self							,    lowercase__):
							if isinstance(lowercase__							,    (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(lowercase__							,    nn.BatchNormad):
													module.bias.data.zero_()
													module.weight.data.fill_(1.0)






lowerCAmelCase							=						R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

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




@add_start_docstrings(
    '''The bare MobileNetV1 model outputting raw hidden-states without any specific head on top.'''       ,	_UpperCamelCase       ,	)
class        lowerCamelCase      (     _UpperCamelCase   ):


	def __init__(						self							,    lowercase__							,    lowercase__ = True):
							super().__init__(lowercase__)
							__UpperCAmelCase :       Tuple  						=  config

							__UpperCAmelCase :       Tuple  						=  3_2
							__UpperCAmelCase :       int  						=  max(int(depth * config.depth_multiplier)							,    config.min_depth)

							__UpperCAmelCase :       List[str]  						=  MobileNetVaConvLayer(
							    lowercase__							,    in_channels=config.num_channels							,    out_channels=lowercase__							,    kernel_size=3							,    stride=2							,    )

							__UpperCAmelCase :       int  						=  [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

							__UpperCAmelCase :       Union[str, Any]  						=  nn.ModuleList()
							for i in range(1_3):
													__UpperCAmelCase :       List[Any]  						=  out_channels

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

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

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

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

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


	def    A(						self							,    lowercase__):
							raise NotImplementedError


	@add_start_docstrings_to_model_forward(lowercase__)
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC							,    output_type=lowercase__							,    config_class=_CONFIG_FOR_DOC							,    modality='''vision'''							,    expected_output=_EXPECTED_OUTPUT_SHAPE							,    )
	def    A(						self							,    lowercase__ = None							,    lowercase__ = None							,    lowercase__ = None							,    ):
							__UpperCAmelCase :       Tuple  						=  (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							__UpperCAmelCase :       Tuple  						=  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''')

							__UpperCAmelCase :       Optional[int]  						=  self.conv_stem(lowercase__)

							__UpperCAmelCase :       Dict  						=  () if output_hidden_states else None

							for i, layer_module in enumerate(self.layer):
													__UpperCAmelCase :       str  						=  layer_module(lowercase__)

													if output_hidden_states:
																			__UpperCAmelCase :       Tuple  						=  all_hidden_states + (hidden_states,)

							__UpperCAmelCase :       Union[str, Any]  						=  hidden_states

							if self.pooler is not None:
													__UpperCAmelCase :       Optional[int]  						=  torch.flatten(self.pooler(lowercase__)							,    start_dim=1)
							else:
													__UpperCAmelCase :       List[Any]  						=  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=lowercase__							,    pooler_output=lowercase__							,    hidden_states=lowercase__							,    )






@add_start_docstrings(
    '''
    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''       ,	_UpperCamelCase       ,	)
class        lowerCamelCase      (     _UpperCamelCase   ):


	def __init__(						self							,    lowercase__):
							super().__init__(lowercase__)

							__UpperCAmelCase :       str  						=  config.num_labels
							__UpperCAmelCase :       Optional[int]  						=  MobileNetVaModel(lowercase__)

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

							# Classifier head
							__UpperCAmelCase :       Optional[Any]  						=  nn.Dropout(config.classifier_dropout_prob							,    inplace=lowercase__)
							__UpperCAmelCase :       Tuple  						=  nn.Linear(lowercase__							,    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(lowercase__)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT							,    output_type=lowercase__							,    config_class=_CONFIG_FOR_DOC							,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,    )
	def    A(						self							,    lowercase__ = None							,    lowercase__ = None							,    lowercase__ = None							,    lowercase__ = None							,    ):
							__UpperCAmelCase :       Tuple  						=  return_dict if return_dict is not None else self.config.use_return_dict

							__UpperCAmelCase :       Any  						=  self.mobilenet_va(lowercase__							,    output_hidden_states=lowercase__							,    return_dict=lowercase__)

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

							__UpperCAmelCase :       Optional[int]  						=  self.classifier(self.dropout(lowercase__))

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

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

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

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





from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


lowerCAmelCase							=						logging.get_logger(__name__)  # pylint: disable=invalid-name

lowerCAmelCase							=						"""
    Examples:
        ```py
        >>> import torch
        >>> import numpy as np

        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline
        >>> from transformers import pipeline
        >>> from diffusers.utils import load_image


        >>> def make_hint(image, depth_estimator):
        ...     image = depth_estimator(image)[\"depth\"]
        ...     image = np.array(image)
        ...     image = image[:, :, None]
        ...     image = np.concatenate([image, image, image], axis=2)
        ...     detected_map = torch.from_numpy(image).float() / 255.0
        ...     hint = detected_map.permute(2, 0, 1)
        ...     return hint


        >>> depth_estimator = pipeline(\"depth-estimation\")

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16
        ... )
        >>> pipe_prior = pipe_prior.to(\"cuda\")

        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(
        ...     \"kandinsky-community/kandinsky-2-2-controlnet-depth\", torch_dtype=torch.float16
        ... )
        >>> pipe = pipe.to(\"cuda\")


        >>> img = load_image(
        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"
        ...     \"/kandinsky/cat.png\"
        ... ).resize((768, 768))

        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to(\"cuda\")

        >>> prompt = \"A robot, 4k photo\"
        >>> negative_prior_prompt = \"lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature\"

        >>> generator = torch.Generator(device=\"cuda\").manual_seed(43)

        >>> image_emb, zero_image_emb = pipe_prior(
        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator
        ... ).to_tuple()

        >>> images = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     hint=hint,
        ...     num_inference_steps=50,
        ...     generator=generator,
        ...     height=768,
        ...     width=768,
        ... ).images

        >>> images[0].save(\"robot_cat.png\")
        ```
"""




def 						__SCREAMING_SNAKE_CASE       (					lowercase_			,					lowercase_			,					lowercase_=8			)		->    List[Any]:




						'''simple docstring'''







						__UpperCAmelCase :       int  						=  height // scale_factor**2
						if height % scale_factor**2 != 0:
												new_height += 1
						__UpperCAmelCase :       Union[str, Any]  						=  width // scale_factor**2
						if width % scale_factor**2 != 0:
												new_width += 1
						return new_height * scale_factor, new_width * scale_factor





class        lowerCamelCase      (     _UpperCamelCase   ):


	def __init__(						self							,    lowercase__							,    lowercase__							,    lowercase__							,    ):
							super().__init__()

							self.register_modules(
							    unet=lowercase__							,    scheduler=lowercase__							,    movq=lowercase__							,    )
							__UpperCAmelCase :       Any  						=  2 ** (len(self.movq.config.block_out_channels) - 1)


	def    A(						self							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__):
							if latents is None:
													__UpperCAmelCase :       Any  						=  randn_tensor(lowercase__							,    generator=lowercase__							,    device=lowercase__							,    dtype=lowercase__)
							else:
													if latents.shape != shape:
																			raise ValueError(F"Unexpected latents shape, got {latents.shape}, expected {shape}")
													__UpperCAmelCase :       Union[str, Any]  						=  latents.to(lowercase__)

							__UpperCAmelCase :       Union[str, Any]  						=  latents * scheduler.init_noise_sigma
							return latents


	def    A(						self							,    lowercase__=0):
							if is_accelerate_available():
													from accelerate import cpu_offload
							else:
													raise ImportError('''Please install accelerate via `pip install accelerate`''')

							__UpperCAmelCase :       List[str]  						=  torch.device(F"cuda:{gpu_id}")

							__UpperCAmelCase :       List[Any]  						=  [
							    self.unet,
							    self.movq,
							]
							for cpu_offloaded_model in models:
													if cpu_offloaded_model is not None:
																			cpu_offload(lowercase__							,    lowercase__)


	def    A(						self							,    lowercase__=0):
							if is_accelerate_available() and is_accelerate_version('''>='''							,    '''0.17.0.dev0'''):
													from accelerate import cpu_offload_with_hook
							else:
													raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.''')

							__UpperCAmelCase :       Optional[Any]  						=  torch.device(F"cuda:{gpu_id}")

							if self.device.type != "cpu":
													self.to('''cpu'''							,    silence_dtype_warnings=lowercase__)
													torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

							__UpperCAmelCase :       List[Any]  						=  None
							for cpu_offloaded_model in [self.unet, self.movq]:
													__UpperCAmelCase						,			__UpperCAmelCase :       List[str]  						=  cpu_offload_with_hook(lowercase__							,    lowercase__							,    prev_module_hook=lowercase__)

							# We'll offload the last model manually.
							__UpperCAmelCase :       Any  						=  hook


	@property
	# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
	def    A(						self):
							if not hasattr(self.unet							,    '''_hf_hook'''):
													return self.device
							for module in self.unet.modules():
													if (
													    hasattr(lowercase__							,    '''_hf_hook''')
													    and hasattr(module._hf_hook							,    '''execution_device''')
													    and module._hf_hook.execution_device is not None
													):
																			return torch.device(module._hf_hook.execution_device)
							return self.device


	@torch.no_grad()
	@replace_example_docstring(lowercase__)
	def __call__(						self							,    lowercase__							,    lowercase__							,    lowercase__							,    lowercase__ = 5_1_2							,    lowercase__ = 5_1_2							,    lowercase__ = 1_0_0							,    lowercase__ = 4.0							,    lowercase__ = 1							,    lowercase__ = None							,    lowercase__ = None							,    lowercase__ = "pil"							,    lowercase__ = True							,    ):
							__UpperCAmelCase :       str  						=  self._execution_device

							__UpperCAmelCase :       List[str]  						=  guidance_scale > 1.0

							if isinstance(lowercase__							,    lowercase__):
													__UpperCAmelCase :       Dict  						=  torch.cat(lowercase__							,    dim=0)
							if isinstance(lowercase__							,    lowercase__):
													__UpperCAmelCase :       Tuple  						=  torch.cat(lowercase__							,    dim=0)
							if isinstance(lowercase__							,    lowercase__):
													__UpperCAmelCase :       Any  						=  torch.cat(lowercase__							,    dim=0)

							__UpperCAmelCase :       Union[str, Any]  						=  image_embeds.shape[0] * num_images_per_prompt

							if do_classifier_free_guidance:
													__UpperCAmelCase :       Optional[int]  						=  image_embeds.repeat_interleave(lowercase__							,    dim=0)
													__UpperCAmelCase :       Dict  						=  negative_image_embeds.repeat_interleave(lowercase__							,    dim=0)
													__UpperCAmelCase :       List[Any]  						=  hint.repeat_interleave(lowercase__							,    dim=0)

							__UpperCAmelCase :       Tuple  						=  torch.cat([negative_image_embeds, image_embeds]							,    dim=0).to(dtype=self.unet.dtype							,    device=lowercase__)
							__UpperCAmelCase :       List[Any]  						=  torch.cat([hint, hint]							,    dim=0).to(dtype=self.unet.dtype							,    device=lowercase__)
							self.scheduler.set_timesteps(lowercase__							,    device=lowercase__)
							__UpperCAmelCase :       List[Any]  						=  self.scheduler.timesteps

							__UpperCAmelCase :       Any  						=  self.movq.config.latent_channels

							__UpperCAmelCase						,			__UpperCAmelCase :       List[str]  						=  downscale_height_and_width(lowercase__							,    lowercase__							,    self.movq_scale_factor)

							# create initial latent
							__UpperCAmelCase :       Union[str, Any]  						=  self.prepare_latents(
							    (batch_size, num_channels_latents, height, width)							,    image_embeds.dtype							,    lowercase__							,    lowercase__							,    lowercase__							,    self.scheduler							,    )

							for i, t in enumerate(self.progress_bar(lowercase__)):
													# expand the latents if we are doing classifier free guidance
													__UpperCAmelCase :       List[Any]  						=  torch.cat([latents] * 2) if do_classifier_free_guidance else latents

													__UpperCAmelCase :       Union[str, Any]  						=  {'''image_embeds''': image_embeds, '''hint''': hint}
													__UpperCAmelCase :       Any  						=  self.unet(
													    sample=lowercase__							,    timestep=lowercase__							,    encoder_hidden_states=lowercase__							,    added_cond_kwargs=lowercase__							,    return_dict=lowercase__							,    )[0]

													if do_classifier_free_guidance:
																			__UpperCAmelCase						,			__UpperCAmelCase :       Optional[Any]  						=  noise_pred.split(latents.shape[1]							,    dim=1)
																			__UpperCAmelCase						,			__UpperCAmelCase :       List[str]  						=  noise_pred.chunk(2)
																			__UpperCAmelCase						,			__UpperCAmelCase :       Optional[Any]  						=  variance_pred.chunk(2)
																			__UpperCAmelCase :       Union[str, Any]  						=  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																			__UpperCAmelCase :       int  						=  torch.cat([noise_pred, variance_pred_text]							,    dim=1)

													if not (
													    hasattr(self.scheduler.config							,    '''variance_type''')
													    and self.scheduler.config.variance_type in ["learned", "learned_range"]
													):
																			__UpperCAmelCase						,			__UpperCAmelCase :       Optional[Any]  						=  noise_pred.split(latents.shape[1]							,    dim=1)

													# compute the previous noisy sample x_t -> x_t-1
													__UpperCAmelCase :       Tuple  						=  self.scheduler.step(
													    lowercase__							,    lowercase__							,    lowercase__							,    generator=lowercase__							,    )[0]
							# post-processing
							__UpperCAmelCase :       str  						=  self.movq.decode(lowercase__							,    force_not_quantize=lowercase__)['''sample''']

							if output_type not in ["pt", "np", "pil"]:
													raise ValueError(F"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}")

							if output_type in ["np", "pil"]:
													__UpperCAmelCase :       Dict  						=  image * 0.5 + 0.5
													__UpperCAmelCase :       Union[str, Any]  						=  image.clamp(0							,    1)
													__UpperCAmelCase :       List[str]  						=  image.cpu().permute(0							,    2							,    3							,    1).float().numpy()

							if output_type == "pil":
													__UpperCAmelCase :       List[str]  						=  self.numpy_to_pil(lowercase__)

							if not return_dict:
													return (image,)

							return ImagePipelineOutput(images=lowercase__)

