

"""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


_lowerCAmelCase						=		logging.get_logger(__name__)


# General docstring
_lowerCAmelCase						=		"""MobileNetV1Config"""

# Base docstring
_lowerCAmelCase						=		"""google/mobilenet_v1_1.0_224"""
_lowerCAmelCase						=		[1, 1_0_2_4, 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 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase=None   ):

    '''simple docstring'''



    _lowerCAmelCase       :   Optional[int]  =    {}

    if isinstance(_lowerCamelCase ,   _lowerCamelCase   ):
        _lowerCAmelCase       :   Tuple  =    model.mobilenet_va
    else:
        _lowerCAmelCase       :   Dict  =    model

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

    for i in range(13   ):
        _lowerCAmelCase       :   int  =    i + 1
        _lowerCAmelCase       :   Optional[Any]  =    i * 2

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

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

    if isinstance(_lowerCamelCase ,   _lowerCamelCase   ):
        _lowerCAmelCase       :   Optional[Any]  =    'MobilenetV1/Logits/Conv2d_1c_1x1/'
        _lowerCAmelCase       :   Tuple  =    model.classifier.weight
        _lowerCAmelCase       :   Dict  =    model.classifier.bias

    return tf_to_pt_map

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase   ):

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

    # Build TF to PyTorch weights loading map
    _lowerCAmelCase       :   str  =    _build_tf_to_pytorch_map(_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase   )

    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

        _lowerCAmelCase       :   Dict  =    tf_weights[name]

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

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

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

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase   ):

    '''simple docstring'''



    _lowerCAmelCase       :   Tuple  =    features.shape[-2:]
    _lowerCAmelCase       :   List[Any]  =    conv_layer.stride
    _lowerCAmelCase       :   List[Any]  =    conv_layer.kernel_size

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

    if in_width % stride_width == 0:
        _lowerCAmelCase       :   List[str]  =    max(kernel_width - stride_width ,   0   )
    else:
        _lowerCAmelCase       :   Optional[Any]  =    max(kernel_width - (in_width % stride_width) ,   0   )

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

    _lowerCAmelCase       :   int  =    (pad_left, pad_right, pad_top, pad_bottom)
    return nn.functional.pad(_lowerCamelCase ,   _lowerCamelCase ,   '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__()
      _lowerCAmelCase       :   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."""      )

      _lowerCAmelCase       :   Union[str, Any]  =    0 if config.tf_padding else int((kernel_size - 1) / 2      )

      _lowerCAmelCase       :   Optional[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:
          _lowerCAmelCase       :   Union[str, Any]  =    nn.BatchNormad(
              num_features=_A    ,eps=config.layer_norm_eps    ,momentum=0.9_9_9_7    ,affine=_A    ,track_running_stats=_A    ,)
      else:
          _lowerCAmelCase       :   Dict  =    None

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







  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

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







class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		MobileNetVaConfig
  _UpperCAmelCase          =		load_tf_weights_in_mobilenet_va
  _UpperCAmelCase          =		"mobilenet_v1"
  _UpperCAmelCase          =		"pixel_values"
  _UpperCAmelCase          =		False

  def __lowerCamelCase		(		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      )





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

  def __init__(		self    ,_A    ,_A = True      ):
      '''simple docstring'''

      super().__init__(_A      )
      _lowerCAmelCase       :   Optional[Any]  =    config

      _lowerCAmelCase       :   Any  =    32
      _lowerCAmelCase       :   List[Any]  =    max(int(depth * config.depth_multiplier      )    ,config.min_depth      )

      _lowerCAmelCase       :   Any  =    MobileNetVaConvLayer(
          _A    ,in_channels=config.num_channels    ,out_channels=_A    ,kernel_size=3    ,stride=2    ,)

      _lowerCAmelCase       :   List[str]  =    [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

      _lowerCAmelCase       :   Tuple  =    nn.ModuleList()
      for i in range(13      ):
          _lowerCAmelCase       :   Any  =    out_channels

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

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

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

  def __lowerCamelCase		(		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 __lowerCamelCase		(		self    ,_A = None    ,_A = None    ,_A = None    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    (
          output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
      )
      _lowerCAmelCase       :   Union[str, 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'      )

      _lowerCAmelCase       :   Any  =    self.conv_stem(_A      )

      _lowerCAmelCase       :   Optional[int]  =    () if output_hidden_states else None

      for i, layer_module in enumerate(self.layer      ):
          _lowerCAmelCase       :   Tuple  =    layer_module(_A      )

          if output_hidden_states:
              _lowerCAmelCase       :   int  =    all_hidden_states + (hidden_states,)

      _lowerCAmelCase       :   Dict  =    hidden_states

      if self.pooler is not None:
          _lowerCAmelCase       :   Union[str, Any]  =    torch.flatten(self.pooler(_A      )    ,start_dim=1      )
      else:
          _lowerCAmelCase       :   Tuple  =    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    "							,							a__							,							)
class 	__UpperCamelCase    (    a__						):

  def __init__(		self    ,_A      ):
      '''simple docstring'''

      super().__init__(_A      )

      _lowerCAmelCase       :   Dict  =    config.num_labels
      _lowerCAmelCase       :   List[str]  =    MobileNetVaModel(_A      )

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

      # Classifier head
      _lowerCAmelCase       :   Dict  =    nn.Dropout(config.classifier_dropout_prob    ,inplace=_A      )
      _lowerCAmelCase       :   List[str]  =    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 __lowerCamelCase		(		self    ,_A = None    ,_A = None    ,_A = None    ,_A = None    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    return_dict if return_dict is not None else self.config.use_return_dict

      _lowerCAmelCase       :   Any  =    self.mobilenet_va(_A    ,output_hidden_states=_A    ,return_dict=_A      )

      _lowerCAmelCase       :   int  =    outputs.pooler_output if return_dict else outputs[1]

      _lowerCAmelCase       :   Optional[int]  =    self.classifier(self.dropout(_A      )      )

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

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

      if not return_dict:
          _lowerCAmelCase       :   Union[str, 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    ,)







"""simple docstring"""






from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel







class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		["vqvae"]

  def __init__(		self    ,_A    ,_A    ,_A    ,_A    ,):
      '''simple docstring'''

      super().__init__()
      self.register_modules(unet=_A    ,scheduler=_A    ,mel=_A    ,vqvae=_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return 50 if isinstance(self.scheduler    ,_A      ) else 1000

  @torch.no_grad()
  def __call__(		self    ,_A = 1    ,_A = None    ,_A = None    ,_A = 0    ,_A = 0    ,_A = None    ,_A = None    ,_A = 0    ,_A = 0    ,_A = None    ,_A = 0    ,_A = None    ,_A = None    ,_A=True    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   List[str]  =    steps or self.get_default_steps()
      self.scheduler.set_timesteps(_A      )
      _lowerCAmelCase       :   Optional[Any]  =    step_generator or generator
      # For backwards compatibility
      if type(self.unet.config.sample_size      ) == int:
          _lowerCAmelCase       :   Tuple  =    (self.unet.config.sample_size, self.unet.config.sample_size)
      if noise is None:
          _lowerCAmelCase       :   Optional[Any]  =    randn_tensor(
              (
                  batch_size,
                  self.unet.config.in_channels,
                  self.unet.config.sample_size[0],
                  self.unet.config.sample_size[1],
              )    ,generator=_A    ,device=self.device    ,)
      _lowerCAmelCase       :   Dict  =    noise
      _lowerCAmelCase       :   Optional[Any]  =    None

      if audio_file is not None or raw_audio is not None:
          self.mel.load_audio(_A    ,_A      )
          _lowerCAmelCase       :   Union[str, Any]  =    self.mel.audio_slice_to_image(_A      )
          _lowerCAmelCase       :   int  =    np.frombuffer(input_image.tobytes()    ,dtype='uint8'      ).reshape(
              (input_image.height, input_image.width)      )
          _lowerCAmelCase       :   int  =    (input_image / 255) * 2 - 1
          _lowerCAmelCase       :   str  =    torch.tensor(input_image[np.newaxis, :, :]    ,dtype=torch.float      ).to(self.device      )

          if self.vqvae is not None:
              _lowerCAmelCase       :   List[Any]  =    self.vqvae.encode(torch.unsqueeze(_A    ,0      )      ).latent_dist.sample(
                  generator=_A      )[0]
              _lowerCAmelCase       :   Tuple  =    self.vqvae.config.scaling_factor * input_images

          if start_step > 0:
              _lowerCAmelCase       :   List[Any]  =    self.scheduler.add_noise(_A    ,_A    ,self.scheduler.timesteps[start_step - 1]      )

          _lowerCAmelCase       :   Optional[Any]  =    (
              self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
          )
          _lowerCAmelCase       :   Optional[Any]  =    int(mask_start_secs * pixels_per_second      )
          _lowerCAmelCase       :   Optional[int]  =    int(mask_end_secs * pixels_per_second      )
          _lowerCAmelCase       :   int  =    self.scheduler.add_noise(_A    ,_A    ,torch.tensor(self.scheduler.timesteps[start_step:]      )      )

      for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]      )      ):
          if isinstance(self.unet    ,_A      ):
              _lowerCAmelCase       :   str  =    self.unet(_A    ,_A    ,_A      )['sample']
          else:
              _lowerCAmelCase       :   Any  =    self.unet(_A    ,_A      )['sample']

          if isinstance(self.scheduler    ,_A      ):
              _lowerCAmelCase       :   Union[str, Any]  =    self.scheduler.step(
                  model_output=_A    ,timestep=_A    ,sample=_A    ,eta=_A    ,generator=_A    ,)['prev_sample']
          else:
              _lowerCAmelCase       :   Any  =    self.scheduler.step(
                  model_output=_A    ,timestep=_A    ,sample=_A    ,generator=_A    ,)['prev_sample']

          if mask is not None:
              if mask_start > 0:
                  _lowerCAmelCase       :   Any  =    mask[:, step, :, :mask_start]
              if mask_end > 0:
                  _lowerCAmelCase       :   Optional[Any]  =    mask[:, step, :, -mask_end:]

      if self.vqvae is not None:
          # 0.18215 was scaling factor used in training to ensure unit variance
          _lowerCAmelCase       :   Union[str, Any]  =    1 / self.vqvae.config.scaling_factor * images
          _lowerCAmelCase       :   Any  =    self.vqvae.decode(_A      )['sample']

      _lowerCAmelCase       :   Any  =    (images / 2 + 0.5).clamp(0    ,1      )
      _lowerCAmelCase       :   Tuple  =    images.cpu().permute(0    ,2    ,3    ,1      ).numpy()
      _lowerCAmelCase       :   Any  =    (images * 255).round().astype('uint8'      )
      _lowerCAmelCase       :   Any  =    list(
          (Image.fromarray(_[:, :, 0]      ) for _ in images)
          if images.shape[3] == 1
          else (Image.fromarray(_A    ,mode='RGB'      ).convert('L'      ) for _ in images)      )

      _lowerCAmelCase       :   Dict  =    [self.mel.image_to_audio(_A      ) for _ in images]
      if not return_dict:
          return images, (self.mel.get_sample_rate(), audios)

      return BaseOutput(**AudioPipelineOutput(np.array(_A      )[:, np.newaxis, :]      )    ,**ImagePipelineOutput(_A      )      )

  @torch.no_grad()
  def __lowerCamelCase		(		self    ,_A    ,_A = 50      ):
      '''simple docstring'''

      assert isinstance(self.scheduler    ,_A      )
      self.scheduler.set_timesteps(_A      )
      _lowerCAmelCase       :   Dict  =    np.array(
          [np.frombuffer(image.tobytes()    ,dtype='uint8'      ).reshape((1, image.height, image.width)      ) for image in images]      )
      _lowerCAmelCase       :   Dict  =    (sample / 255) * 2 - 1
      _lowerCAmelCase       :   List[str]  =    torch.Tensor(_A      ).to(self.device      )

      for t in self.progress_bar(torch.flip(self.scheduler.timesteps    ,(0,)      )      ):
          _lowerCAmelCase       :   Tuple  =    t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
          _lowerCAmelCase       :   Optional[int]  =    self.scheduler.alphas_cumprod[t]
          _lowerCAmelCase       :   Dict  =    (
              self.scheduler.alphas_cumprod[prev_timestep]
              if prev_timestep >= 0
              else self.scheduler.final_alpha_cumprod
          )
          _lowerCAmelCase       :   Union[str, Any]  =    1 - alpha_prod_t
          _lowerCAmelCase       :   Union[str, Any]  =    self.unet(_A    ,_A      )['sample']
          _lowerCAmelCase       :   Optional[int]  =    (1 - alpha_prod_t_prev) ** 0.5 * model_output
          _lowerCAmelCase       :   Any  =    (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
          _lowerCAmelCase       :   Dict  =    sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

      return sample







  @staticmethod
  def __lowerCamelCase		(		_A    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    acos(torch.dot(torch.flatten(_A      )    ,torch.flatten(_A      )      ) / torch.norm(_A      ) / torch.norm(_A      )      )
      return sin((1 - alpha) * theta      ) * xa / sin(_A      ) + sin(alpha * theta      ) * xa / sin(_A      )





