

"""simple docstring"""






import collections.abc
from typing import Optional, Tuple, Union

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

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, 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_poolformer import PoolFormerConfig


_lowerCAmelCase						=		logging.get_logger(__name__)

# General docstring
_lowerCAmelCase						=		"""PoolFormerConfig"""

# Base docstring
_lowerCAmelCase						=		"""sail/poolformer_s12"""
_lowerCAmelCase						=		[1, 5_1_2, 7, 7]

# Image classification docstring
_lowerCAmelCase						=		"""sail/poolformer_s12"""
_lowerCAmelCase						=		"""tabby, tabby cat"""

_lowerCAmelCase						=		[
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase = 0.0 ,   _lowerCamelCase = False   ):

    '''simple docstring'''



    if drop_prob == 0.0 or not training:
        return input
    _lowerCAmelCase       :   List[str]  =    1 - drop_prob
    _lowerCAmelCase       :   List[str]  =    (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    _lowerCAmelCase       :   str  =    keep_prob + torch.rand(_lowerCamelCase ,   dtype=input.dtype ,   device=input.device   )
    random_tensor.floor_()  # binarize
    _lowerCAmelCase       :   Any  =    input.div(_lowerCamelCase   ) * random_tensor
    return output







class 	__UpperCamelCase    (    nn.Module						):

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

      super().__init__()
      _lowerCAmelCase       :   int  =    drop_prob

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

      return drop_path(_A    ,self.drop_prob    ,self.training      )







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

      return "p={}".format(self.drop_prob      )







class 	__UpperCamelCase    (    nn.Module						):

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

      super().__init__()
      _lowerCAmelCase       :   Optional[Any]  =    patch_size if isinstance(_A    ,collections.abc.Iterable      ) else (patch_size, patch_size)
      _lowerCAmelCase       :   Union[str, Any]  =    stride if isinstance(_A    ,collections.abc.Iterable      ) else (stride, stride)
      _lowerCAmelCase       :   Optional[Any]  =    padding if isinstance(_A    ,collections.abc.Iterable      ) else (padding, padding)

      _lowerCAmelCase       :   List[Any]  =    nn.Convad(_A    ,_A    ,kernel_size=_A    ,stride=_A    ,padding=_A      )
      _lowerCAmelCase       :   Any  =    norm_layer(_A      ) if norm_layer else nn.Identity()







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

      _lowerCAmelCase       :   str  =    self.projection(_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    self.norm(_A      )
      return embeddings







class 	__UpperCamelCase    (    nn.GroupNorm						):

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

      super().__init__(1    ,_A    ,**_A      )







class 	__UpperCamelCase    (    nn.Module						):

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

      super().__init__()
      _lowerCAmelCase       :   Union[str, Any]  =    nn.AvgPoolad(_A    ,stride=1    ,padding=pool_size // 2    ,count_include_pad=_A      )







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

      return self.pool(_A      ) - hidden_states







class 	__UpperCamelCase    (    nn.Module						):

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

      super().__init__()
      _lowerCAmelCase       :   str  =    nn.Convad(_A    ,_A    ,1      )
      _lowerCAmelCase       :   Optional[Any]  =    nn.Convad(_A    ,_A    ,1      )
      _lowerCAmelCase       :   Union[str, Any]  =    PoolFormerDropPath(_A      )
      if isinstance(config.hidden_act    ,_A      ):
          _lowerCAmelCase       :   Optional[int]  =    ACTaFN[config.hidden_act]
      else:
          _lowerCAmelCase       :   str  =    config.hidden_act







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

      _lowerCAmelCase       :   Tuple  =    self.conva(_A      )
      _lowerCAmelCase       :   Optional[Any]  =    self.act_fn(_A      )
      _lowerCAmelCase       :   List[str]  =    self.drop(_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    self.conva(_A      )
      _lowerCAmelCase       :   Any  =    self.drop(_A      )

      return hidden_states







class 	__UpperCamelCase    (    nn.Module						):

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

      super().__init__()
      _lowerCAmelCase       :   int  =    PoolFormerPooling(_A      )
      _lowerCAmelCase       :   int  =    PoolFormerOutput(_A    ,_A    ,_A    ,_A      )
      _lowerCAmelCase       :   List[Any]  =    PoolFormerGroupNorm(_A      )
      _lowerCAmelCase       :   Dict  =    PoolFormerGroupNorm(_A      )

      # Useful for training neural nets
      _lowerCAmelCase       :   Optional[Any]  =    PoolFormerDropPath(_A      ) if drop_path > 0.0 else nn.Identity()
      _lowerCAmelCase       :   Any  =    config.use_layer_scale
      if config.use_layer_scale:
          _lowerCAmelCase       :   List[str]  =    nn.Parameter(
              config.layer_scale_init_value * torch.ones((_A)      )    ,requires_grad=_A      )
          _lowerCAmelCase       :   Optional[Any]  =    nn.Parameter(
              config.layer_scale_init_value * torch.ones((_A)      )    ,requires_grad=_A      )







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

      if self.use_layer_scale:
          _lowerCAmelCase       :   Optional[int]  =    self.pooling(self.before_norm(_A      )      )
          _lowerCAmelCase       :   List[str]  =    self.layer_scale_a.unsqueeze(-1      ).unsqueeze(-1      ) * pooling_output
          # First residual connection
          _lowerCAmelCase       :   Union[str, Any]  =    hidden_states + self.drop_path(_A      )
          _lowerCAmelCase       :   Union[str, Any]  =    ()

          _lowerCAmelCase       :   Optional[int]  =    self.output(self.after_norm(_A      )      )
          _lowerCAmelCase       :   Optional[Any]  =    self.layer_scale_a.unsqueeze(-1      ).unsqueeze(-1      ) * layer_output
          # Second residual connection
          _lowerCAmelCase       :   int  =    hidden_states + self.drop_path(_A      )

          _lowerCAmelCase       :   int  =    (output,) + outputs
          return outputs

      else:
          _lowerCAmelCase       :   List[Any]  =    self.drop_path(self.pooling(self.before_norm(_A      )      )      )
          # First residual connection
          _lowerCAmelCase       :   int  =    pooling_output + hidden_states
          _lowerCAmelCase       :   List[str]  =    ()

          # Second residual connection inside the PoolFormerOutput block
          _lowerCAmelCase       :   Tuple  =    self.drop_path(self.output(self.after_norm(_A      )      )      )
          _lowerCAmelCase       :   str  =    hidden_states + layer_output

          _lowerCAmelCase       :   Union[str, Any]  =    (output,) + outputs
          return outputs







class 	__UpperCamelCase    (    nn.Module						):

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

      super().__init__()
      _lowerCAmelCase       :   Optional[int]  =    config
      # stochastic depth decay rule
      _lowerCAmelCase       :   str  =    [x.item() for x in torch.linspace(0    ,config.drop_path_rate    ,sum(config.depths      )      )]

      # patch embeddings
      _lowerCAmelCase       :   Optional[Any]  =    []
      for i in range(config.num_encoder_blocks      ):
          embeddings.append(
              PoolFormerEmbeddings(
                  patch_size=config.patch_sizes[i]    ,stride=config.strides[i]    ,padding=config.padding[i]    ,num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]    ,hidden_size=config.hidden_sizes[i]    ,)      )
      _lowerCAmelCase       :   Dict  =    nn.ModuleList(_A      )

      # Transformer blocks
      _lowerCAmelCase       :   List[Any]  =    []
      _lowerCAmelCase       :   Tuple  =    0
      for i in range(config.num_encoder_blocks      ):
          # each block consists of layers
          _lowerCAmelCase       :   int  =    []
          if i != 0:
              cur += config.depths[i - 1]
          for j in range(config.depths[i]      ):
              layers.append(
                  PoolFormerLayer(
                      _A    ,num_channels=config.hidden_sizes[i]    ,pool_size=config.pool_size    ,hidden_size=config.hidden_sizes[i]    ,intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio      )    ,drop_path=dpr[cur + j]    ,)      )
          blocks.append(nn.ModuleList(_A      )      )

      _lowerCAmelCase       :   Tuple  =    nn.ModuleList(_A      )







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

      _lowerCAmelCase       :   Dict  =    () if output_hidden_states else None

      _lowerCAmelCase       :   str  =    pixel_values
      for idx, layers in enumerate(zip(self.patch_embeddings    ,self.block      )      ):
          _lowerCAmelCase,					_lowerCAmelCase       :   Optional[int]  =    layers
          # Get patch embeddings from hidden_states
          _lowerCAmelCase       :   Dict  =    embedding_layer(_A      )
          # Send the embeddings through the blocks
          for _, blk in enumerate(_A      ):
              _lowerCAmelCase       :   Optional[int]  =    blk(_A      )
              _lowerCAmelCase       :   int  =    layer_outputs[0]

          if output_hidden_states:
              _lowerCAmelCase       :   List[str]  =    all_hidden_states + (hidden_states,)

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

      return BaseModelOutputWithNoAttention(last_hidden_state=_A    ,hidden_states=_A      )







class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		PoolFormerConfig
  _UpperCAmelCase          =		"poolformer"
  _UpperCAmelCase          =		"pixel_values"
  _UpperCAmelCase          =		True

  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.LayerNorm      ):
          module.bias.data.zero_()
          module.weight.data.fill_(1.0      )







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

      if isinstance(_A    ,_A      ):
          _lowerCAmelCase       :   Any  =    value





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

    Parameters:
        config ([`PoolFormerConfig`]): 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
            [`PoolFormerImageProcessor.__call__`] for details.
"""







@add_start_docstrings(
    "The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."							,							a__							,							)
class 	__UpperCamelCase    (    a__						):

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

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

      _lowerCAmelCase       :   int  =    PoolFormerEncoder(_A      )

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

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

      return self.embeddings.patch_embeddings







  @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       :   List[str]  =    (
          output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
      )
      _lowerCAmelCase       :   List[str]  =    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       :   List[Any]  =    self.encoder(
          _A    ,output_hidden_states=_A    ,return_dict=_A    ,)
      _lowerCAmelCase       :   Optional[int]  =    encoder_outputs[0]

      if not return_dict:
          return (sequence_output, None) + encoder_outputs[1:]

      return BaseModelOutputWithNoAttention(
          last_hidden_state=_A    ,hidden_states=encoder_outputs.hidden_states    ,)







class 	__UpperCamelCase    (    nn.Module						):

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

      super().__init__()
      _lowerCAmelCase       :   Dict  =    nn.Linear(config.hidden_size    ,config.hidden_size      )







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

      _lowerCAmelCase       :   int  =    self.dense(_A      )
      return output







@add_start_docstrings(
    "\n    PoolFormer Model transformer with an image classification head on top\n    "							,							a__							,							)
class 	__UpperCamelCase    (    a__						):

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

      super().__init__(_A      )
      _lowerCAmelCase       :   Optional[int]  =    config.num_labels
      _lowerCAmelCase       :   Optional[int]  =    PoolFormerModel(_A      )

      # Final norm
      _lowerCAmelCase       :   Tuple  =    PoolFormerGroupNorm(config.hidden_sizes[-1]      )
      # Classifier head
      _lowerCAmelCase       :   Tuple  =    (
          nn.Linear(config.hidden_sizes[-1]    ,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       :   Dict  =    self.poolformer(
          _A    ,output_hidden_states=_A    ,return_dict=_A    ,)

      _lowerCAmelCase       :   Tuple  =    outputs[0]

      _lowerCAmelCase       :   Any  =    self.classifier(self.norm(_A      ).mean([-2, -1]      )      )

      _lowerCAmelCase       :   List[Any]  =    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       :   str  =    'single_label_classification'
              else:
                  _lowerCAmelCase       :   Optional[int]  =    'multi_label_classification'

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

      if not return_dict:
          _lowerCAmelCase       :   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      )





