



from abc import ABC, abstractmethod
from typing import List, Optional

class       _SCREAMING_SNAKE_CASE	(snake_case__ ):





     def __init__(       self :							Any				)->Tuple:
      self.test()





     def 	__snake_case    (       self :							List[str]				)->str:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	0
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	False
      while not completed:
       if counter == 1:
        self.reset()
       __SCREAMING_SNAKE_CASE       :    Any       =	self.advance()
       if not self.does_advance(UpperCAmelCase_				):
        raise Exception(
            "Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true."				)

       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    str       =	self.update(UpperCAmelCase_				)
       counter += 1

       if counter > 1_0_0_0_0:
        raise Exception("update() does not fulfill the constraint."				)

      if self.remaining() != 0:
       raise Exception("Custom Constraint is not defined correctly."				)





     @abstractmethod
     def 	__snake_case    (       self :							List[Any]				)->Optional[Any]:
      raise NotImplementedError(
          F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""				)





     @abstractmethod
     def 	__snake_case    (       self :							Dict     ,  UpperCamelCase :							int				)->Tuple:
      raise NotImplementedError(
          F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""				)





     @abstractmethod
     def 	__snake_case    (       self :							List[Any]     ,  UpperCamelCase :							int				)->List[str]:
      raise NotImplementedError(
          F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""				)





     @abstractmethod
     def 	__snake_case    (       self :							Optional[int]				)->int:
      raise NotImplementedError(
          F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""				)





     @abstractmethod
     def 	__snake_case    (       self :							Any				)->List[str]:
      raise NotImplementedError(
          F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""				)






     @abstractmethod
     def 	__snake_case    (       self :							List[Any]     ,  UpperCamelCase :							Optional[Any]=False				)->List[Any]:
      raise NotImplementedError(
          F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""				)






class       _SCREAMING_SNAKE_CASE	(snake_case__ ):





     def __init__(       self :							Dict     ,  UpperCamelCase :							List[int]				)->str:
      super(UpperCAmelCase_     ,  self				).__init__()

      if not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				) or len(UpperCAmelCase_				) == 0:
       raise ValueError(F"""`token_ids` has to be a non-empty list, but is {token_ids}."""				)
      if any((not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				) or token_id < 0) for token_id in token_ids				):
       raise ValueError(F"""Each list in `token_ids` has to be a list of positive integers, but is {token_ids}."""				)

      __SCREAMING_SNAKE_CASE       :    str       =	token_ids

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	len(self.token_ids				)
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	-1  # the index of the currently fulfilled step
      __SCREAMING_SNAKE_CASE       :    Tuple       =	False





     def 	__snake_case    (       self :							int				)->Dict:



      if self.completed:
       return None
      return self.token_ids[self.fulfilled_idx + 1]





     def 	__snake_case    (       self :							Tuple     ,  UpperCamelCase :							int				)->Optional[int]:



      if not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
       raise ValueError(F"""`token_id` has to be an `int`, but is {token_id} of type {type(UpperCAmelCase_				)}"""				)

      if self.completed:
       return False

      return token_id == self.token_ids[self.fulfilled_idx + 1]





     def 	__snake_case    (       self :							Tuple     ,  UpperCamelCase :							int				)->List[Any]:



      if not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
       raise ValueError(F"""`token_id` has to be an `int`, but is {token_id} of type {type(UpperCAmelCase_				)}"""				)

      __SCREAMING_SNAKE_CASE       :    Dict       =	False
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	False
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	False

      if self.does_advance(UpperCAmelCase_				):
       self.fulfilled_idx += 1
       __SCREAMING_SNAKE_CASE       :    List[str]       =	True
       if self.fulfilled_idx == (self.seqlen - 1):
        __SCREAMING_SNAKE_CASE       :    Dict       =	True
       __SCREAMING_SNAKE_CASE       :    str       =	completed
      else:
       # failed to make progress.
       __SCREAMING_SNAKE_CASE       :    int       =	True
       self.reset()
      return stepped, completed, reset





     def 	__snake_case    (       self :							Dict				)->Optional[Any]:
      __SCREAMING_SNAKE_CASE       :    Dict       =	False
      __SCREAMING_SNAKE_CASE       :    List[str]       =	0





     def 	__snake_case    (       self :							List[Any]				)->Dict:
      return self.seqlen - (self.fulfilled_idx + 1)






     def 	__snake_case    (       self :							Any     ,  UpperCamelCase :							Dict=False				)->List[str]:
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	PhrasalConstraint(self.token_ids				)

      if stateful:
       __SCREAMING_SNAKE_CASE       :    Dict       =	self.seqlen
       __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.fulfilled_idx
       __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	self.completed

      return new_constraint






class       _SCREAMING_SNAKE_CASE	:





     def __init__(       self :							Optional[Any]     ,  UpperCamelCase :							List[List[int]]     ,  UpperCamelCase :							Optional[int]=True				)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE       :    Any       =	max([len(UpperCAmelCase_				) for one in nested_token_ids]				)

      __SCREAMING_SNAKE_CASE       :    Dict       =	{}
      for token_ids in nested_token_ids:
       __SCREAMING_SNAKE_CASE       :    List[str]       =	root
       for tidx, token_id in enumerate(UpperCAmelCase_				):
        if token_id not in level:
         __SCREAMING_SNAKE_CASE       :    Tuple       =	{}

        __SCREAMING_SNAKE_CASE       :    List[Any]       =	level[token_id]

      if no_subsets and self.has_subsets(UpperCAmelCase_     ,  UpperCAmelCase_				):
       raise ValueError(
           "Each list in `nested_token_ids` can't be a complete subset of another list, but is"
           F""" {nested_token_ids}."""				)

      __SCREAMING_SNAKE_CASE       :    List[Any]       =	root





     def 	__snake_case    (       self :							Optional[Any]     ,  UpperCamelCase :							Optional[int]				)->Optional[Any]:
      __SCREAMING_SNAKE_CASE       :    str       =	self.trie

      for current_token in current_seq:
       __SCREAMING_SNAKE_CASE       :    int       =	start[current_token]

      __SCREAMING_SNAKE_CASE       :    Any       =	list(start.keys()				)

      return next_tokens





     def 	__snake_case    (       self :							str     ,  UpperCamelCase :							int				)->Optional[int]:
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.next_tokens(UpperCAmelCase_				)

      return len(UpperCAmelCase_				) == 0





     def 	__snake_case    (       self :							int     ,  UpperCamelCase :							Dict				)->List[str]:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	list(root.values()				)
      if len(UpperCAmelCase_				) == 0:
       return 1
      else:
       return sum([self.count_leaves(UpperCAmelCase_				) for nn in next_nodes]				)






     def 	__snake_case    (       self :							str     ,  UpperCamelCase :							Union[str, Any]     ,  UpperCamelCase :							Union[str, Any]				)->List[Any]:
      __SCREAMING_SNAKE_CASE       :    str       =	self.count_leaves(UpperCAmelCase_				)
      return len(UpperCAmelCase_				) != leaf_count






class       _SCREAMING_SNAKE_CASE	(snake_case__ ):





     def __init__(       self :							Optional[int]     ,  UpperCamelCase :							List[List[int]]				)->List[str]:
      super(UpperCAmelCase_     ,  self				).__init__()

      if not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				) or len(UpperCAmelCase_				) == 0:
       raise ValueError(F"""`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}."""				)
      if any(not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				) for token_ids in nested_token_ids				):
       raise ValueError(F"""`nested_token_ids` has to be a list of lists, but is {nested_token_ids}."""				)
      if any(
          any((not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				) or token_id < 0) for token_id in token_ids				)
          for token_ids in nested_token_ids				):
       raise ValueError(
           F"""Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}."""				)

      __SCREAMING_SNAKE_CASE       :    str       =	DisjunctiveTrie(UpperCAmelCase_				)
      __SCREAMING_SNAKE_CASE       :    int       =	nested_token_ids

      __SCREAMING_SNAKE_CASE       :    int       =	self.trie.max_height
      __SCREAMING_SNAKE_CASE       :    Tuple       =	[]
      __SCREAMING_SNAKE_CASE       :    int       =	False





     def 	__snake_case    (       self :							List[str]				)->List[Any]:
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	self.trie.next_tokens(self.current_seq				)

      if len(UpperCAmelCase_				) == 0:
       return None
      else:
       return token_list





     def 	__snake_case    (       self :							Union[str, Any]     ,  UpperCamelCase :							int				)->int:



      if not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
       raise ValueError(F"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(UpperCAmelCase_				)}"""				)

      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	self.trie.next_tokens(self.current_seq				)

      return token_id in next_tokens





     def 	__snake_case    (       self :							Optional[Any]     ,  UpperCamelCase :							int				)->Dict:



      if not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
       raise ValueError(F"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(UpperCAmelCase_				)}"""				)

      __SCREAMING_SNAKE_CASE       :    Dict       =	False
      __SCREAMING_SNAKE_CASE       :    List[str]       =	False
      __SCREAMING_SNAKE_CASE       :    List[str]       =	False

      if self.does_advance(UpperCAmelCase_				):
       self.current_seq.append(UpperCAmelCase_				)
       __SCREAMING_SNAKE_CASE       :    int       =	True
      else:
       __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	True
       self.reset()

      __SCREAMING_SNAKE_CASE       :    Any       =	self.trie.reached_leaf(self.current_seq				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	completed

      return stepped, completed, reset





     def 	__snake_case    (       self :							Optional[Any]				)->Dict:
      __SCREAMING_SNAKE_CASE       :    int       =	False
      __SCREAMING_SNAKE_CASE       :    List[str]       =	[]





     def 	__snake_case    (       self :							Optional[Any]				)->int:



      if self.completed:
       # since this can be completed without reaching max height
       return 0
      else:
       return self.seqlen - len(self.current_seq				)






     def 	__snake_case    (       self :							str     ,  UpperCamelCase :							Optional[Any]=False				)->List[Any]:
      __SCREAMING_SNAKE_CASE       :    Tuple       =	DisjunctiveConstraint(self.token_ids				)

      if stateful:
       __SCREAMING_SNAKE_CASE       :    Dict       =	self.seqlen
       __SCREAMING_SNAKE_CASE       :    Optional[int]       =	self.current_seq
       __SCREAMING_SNAKE_CASE       :    int       =	self.completed

      return new_constraint






class       _SCREAMING_SNAKE_CASE	:





     def __init__(       self :							str     ,  UpperCamelCase :							List[Constraint]				)->Any:
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	constraints

      # max # of steps required to fulfill a given constraint
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	max([c.seqlen for c in constraints]				)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	len(UpperCAmelCase_				)
      __SCREAMING_SNAKE_CASE       :    Dict       =	False

      self.init_state()





     def 	__snake_case    (       self :							int				)->Dict:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	[]
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	None
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	[constraint.copy(stateful=UpperCAmelCase_				) for constraint in self.constraints]





     def 	__snake_case    (       self :							List[Any]				)->List[Any]:
      __SCREAMING_SNAKE_CASE       :    str       =	0
      if self.inprogress_constraint:
       # extra points for having a constraint mid-fulfilled
       add += self.max_seqlen - self.inprogress_constraint.remaining()

      return (len(self.complete_constraints				) * self.max_seqlen) + add





     def 	__snake_case    (       self :							str				)->List[Any]:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	[]
      if self.inprogress_constraint is None:
       for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
        __SCREAMING_SNAKE_CASE       :    Tuple       =	constraint.advance()
        if isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
         token_list.append(UpperCAmelCase_				)
        elif isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
         token_list.extend(UpperCAmelCase_				)
      else:
       __SCREAMING_SNAKE_CASE       :    int       =	self.inprogress_constraint.advance()
       if isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
        token_list.append(UpperCAmelCase_				)
       elif isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
        token_list.extend(UpperCAmelCase_				)

      if len(UpperCAmelCase_				) == 0:
       return None
      else:
       return token_list





     def 	__snake_case    (       self :							Optional[int]     ,  UpperCamelCase :							Optional[List[int]]				)->Dict:
      self.init_state()

      if token_ids is not None:
       for token in token_ids:
        # completes or steps **one** constraint
        __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    List[str]       =	self.add(UpperCAmelCase_				)

        # the entire list of constraints are fulfilled
        if self.completed:
         break





     def 	__snake_case    (       self :							Union[str, Any]     ,  UpperCamelCase :							int				)->List[Any]:



      if not isinstance(UpperCAmelCase_     ,  UpperCAmelCase_				):
       raise ValueError(F"""`token_id` should be an `int`, but is `{token_id}`."""				)

      __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    Any       =	False, False

      if self.completed:
       __SCREAMING_SNAKE_CASE       :    List[Any]       =	True
       __SCREAMING_SNAKE_CASE       :    Dict       =	False
       return complete, stepped

      if self.inprogress_constraint is not None:
       # In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
       # job, simply update the state

       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    Optional[int]       =	self.inprogress_constraint.update(UpperCAmelCase_				)
       if reset:
        # 1. If the next token breaks the progress, then we must restart.
        #     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

        #     But that doesn't mean we self.init_state(), since we only reset the state for this particular
        #     constraint, not the full list of constraints.

        self.pending_constraints.append(self.inprogress_constraint.copy(stateful=UpperCAmelCase_				)				)
        __SCREAMING_SNAKE_CASE       :    Tuple       =	None

       if complete:
        # 2. If the next token completes the constraint, move it to completed list, set
        #     inprogress to None. If there are no pending constraints either, then this full list of constraints
        #     is complete.

        self.complete_constraints.append(self.inprogress_constraint				)
        __SCREAMING_SNAKE_CASE       :    Optional[int]       =	None

        if len(self.pending_constraints				) == 0:
         # we're done!
         __SCREAMING_SNAKE_CASE       :    Tuple       =	True

      else:
       # Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
       # of constraints?

       for cidx, pending_constraint in enumerate(self.pending_constraints				):
        if pending_constraint.does_advance(UpperCAmelCase_				):
         __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    str       =	pending_constraint.update(UpperCAmelCase_				)

         if not stepped:
          raise Exception(
              "`constraint.update(token_id)` is not yielding incremental progress, "
              "even though `constraint.does_advance(token_id)` is true."				)

         if complete:
          self.complete_constraints.append(UpperCAmelCase_				)
          __SCREAMING_SNAKE_CASE       :    List[str]       =	None

         if not complete and stepped:
          __SCREAMING_SNAKE_CASE       :    Optional[int]       =	pending_constraint

         if complete or stepped:
          # If we made any progress at all, then it's at least not a "pending constraint".

          __SCREAMING_SNAKE_CASE       :    List[Any]       =	(
              self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
          )

          if len(self.pending_constraints				) == 0 and self.inprogress_constraint is None:
           # If there's no longer any pending after this and no inprogress either, then we must be
           # complete.

           __SCREAMING_SNAKE_CASE       :    int       =	True

          break  # prevent accidentally stepping through multiple constraints with just one token.

      return complete, stepped






     def 	__snake_case    (       self :							Tuple     ,  UpperCamelCase :							Optional[Any]=True				)->Optional[int]:
      __SCREAMING_SNAKE_CASE       :    str       =	ConstraintListState(self.constraints				)  # we actually never though self.constraints objects
      # throughout this process. So it's at initialization state.

      if stateful:
       __SCREAMING_SNAKE_CASE       :    Any       =	[
           constraint.copy(stateful=UpperCAmelCase_				) for constraint in self.complete_constraints
       ]
       if self.inprogress_constraint is not None:
        __SCREAMING_SNAKE_CASE       :    int       =	self.inprogress_constraint.copy(stateful=UpperCAmelCase_				)
       __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	[constraint.copy() for constraint in self.pending_constraints]

      return new_state











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       _SCREAMING_SNAKE_CASE	(UpperCamelCase ):
     lowerCAmelCase					=			["""vqvae"""]





     def __init__(       self :							Tuple     ,  UpperCamelCase :							AutoencoderKL     ,  UpperCamelCase :							UNetaDConditionModel     ,  UpperCamelCase :							Mel     ,  UpperCamelCase :							Union[DDIMScheduler, DDPMScheduler]     ,  )->Tuple:
      super().__init__()
      self.register_modules(unet=UpperCamelCase     ,  scheduler=UpperCamelCase     ,  mel=UpperCamelCase     ,  vqvae=UpperCamelCase				)





     def 	__snake_case    (       self :							List[Any]				)->int:
      return 5_0 if isinstance(self.scheduler     ,  UpperCamelCase				) else 1_0_0_0





     @torch.no_grad()
     def __call__(       self :							Dict     ,  UpperCamelCase :							int = 1     ,  UpperCamelCase :							str = None     ,  UpperCamelCase :							np.ndarray = None     ,  UpperCamelCase :							int = 0     ,  UpperCamelCase :							int = 0     ,  UpperCamelCase :							int = None     ,  UpperCamelCase :							torch.Generator = None     ,  UpperCamelCase :							float = 0     ,  UpperCamelCase :							float = 0     ,  UpperCamelCase :							torch.Generator = None     ,  UpperCamelCase :							float = 0     ,  UpperCamelCase :							torch.Tensor = None     ,  UpperCamelCase :							torch.Tensor = None     ,  UpperCamelCase :							Any=True     ,  )->Union[
         Union[AudioPipelineOutput, ImagePipelineOutput],
         Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
     ]:
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	steps or self.get_default_steps()
      self.scheduler.set_timesteps(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	step_generator or generator
      # For backwards compatibility
      if type(self.unet.config.sample_size				) == int:
       __SCREAMING_SNAKE_CASE       :    Dict       =	(self.unet.config.sample_size, self.unet.config.sample_size)
      if noise is None:
       __SCREAMING_SNAKE_CASE       :    Any       =	randn_tensor(
           (
               batch_size,
               self.unet.config.in_channels,
               self.unet.config.sample_size[0],
               self.unet.config.sample_size[1],
           )     ,  generator=UpperCamelCase     ,  device=self.device     ,  )
      __SCREAMING_SNAKE_CASE       :    Any       =	noise
      __SCREAMING_SNAKE_CASE       :    Any       =	None

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

       if self.vqvae is not None:
        __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	self.vqvae.encode(torch.unsqueeze(UpperCamelCase     ,  0				)				).latent_dist.sample(
            generator=UpperCamelCase				)[0]
        __SCREAMING_SNAKE_CASE       :    int       =	self.vqvae.config.scaling_factor * input_images

       if start_step > 0:
        __SCREAMING_SNAKE_CASE       :    List[str]       =	self.scheduler.add_noise(UpperCamelCase     ,  UpperCamelCase     ,  self.scheduler.timesteps[start_step - 1]				)

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

      for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]				)				):
       if isinstance(self.unet     ,  UpperCamelCase				):
        __SCREAMING_SNAKE_CASE       :    str       =	self.unet(UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase				)["sample"]
       else:
        __SCREAMING_SNAKE_CASE       :    int       =	self.unet(UpperCamelCase     ,  UpperCamelCase				)["sample"]

       if isinstance(self.scheduler     ,  UpperCamelCase				):
        __SCREAMING_SNAKE_CASE       :    int       =	self.scheduler.step(
            model_output=UpperCamelCase     ,  timestep=UpperCamelCase     ,  sample=UpperCamelCase     ,  eta=UpperCamelCase     ,  generator=UpperCamelCase     ,  )["prev_sample"]
       else:
        __SCREAMING_SNAKE_CASE       :    Tuple       =	self.scheduler.step(
            model_output=UpperCamelCase     ,  timestep=UpperCamelCase     ,  sample=UpperCamelCase     ,  generator=UpperCamelCase     ,  )["prev_sample"]

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

      if self.vqvae is not None:
       # 0.18215 was scaling factor used in training to ensure unit variance
       __SCREAMING_SNAKE_CASE       :    Any       =	1 / self.vqvae.config.scaling_factor * images
       __SCREAMING_SNAKE_CASE       :    Any       =	self.vqvae.decode(UpperCamelCase				)["sample"]

      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	(images / 2 + 0.5).clamp(0     ,  1				)
      __SCREAMING_SNAKE_CASE       :    str       =	images.cpu().permute(0     ,  2     ,  3     ,  1				).numpy()
      __SCREAMING_SNAKE_CASE       :    Tuple       =	(images * 2_5_5).round().astype("uint8"				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	list(
          (Image.fromarray(_[:, :, 0]				) for _ in images)
          if images.shape[3] == 1
          else (Image.fromarray(UpperCamelCase     ,  mode="RGB"				).convert("L"				) for _ in images)				)

      __SCREAMING_SNAKE_CASE       :    List[str]       =	[self.mel.image_to_audio(UpperCamelCase				) for _ in images]
      if not return_dict:
       return images, (self.mel.get_sample_rate(), audios)

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





     @torch.no_grad()
     def 	__snake_case    (       self :							Dict     ,  UpperCamelCase :							List[Image.Image]     ,  UpperCamelCase :							int = 5_0				)->np.ndarray:
      assert isinstance(self.scheduler     ,  UpperCamelCase				)
      self.scheduler.set_timesteps(UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    int       =	np.array(
          [np.frombuffer(image.tobytes()     ,  dtype="uint8"				).reshape((1, image.height, image.width)				) for image in images]				)
      __SCREAMING_SNAKE_CASE       :    Dict       =	(sample / 2_5_5) * 2 - 1
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	torch.Tensor(UpperCamelCase				).to(self.device				)

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

      return sample






     @staticmethod
     def 	__snake_case    (       UpperCamelCase :							torch.Tensor     ,  UpperCamelCase :							torch.Tensor     ,  UpperCamelCase :							float				)->torch.Tensor:
      __SCREAMING_SNAKE_CASE       :    List[str]       =	acos(torch.dot(torch.flatten(UpperCamelCase				)     ,  torch.flatten(UpperCamelCase				)				) / torch.norm(UpperCamelCase				) / torch.norm(UpperCamelCase				)				)
      return sin((1 - alpha) * theta				) * xa / sin(UpperCamelCase				) + sin(alpha * theta				) * xa / sin(UpperCamelCase				)






