





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







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					def __init__(						self    :  Tuple		):






						'''simple docstring'''







						self.test()





					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







						UpperCAmelCase_								=     0
						UpperCAmelCase_								=     False
						while not completed:
							if counter == 1:
								self.reset()
							UpperCAmelCase_								=     self.advance()
							if not self.does_advance(__lowercase		):
								raise Exception(
								    """Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true."""		)

							UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_								=     self.update(__lowercase		)
							counter += 1

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

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





					@abstractmethod
					def        SCREAMING_SNAKE_CASE							(						self    :  Dict		):






						'''simple docstring'''







						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""		)





					@abstractmethod
					def        SCREAMING_SNAKE_CASE							(						self    :  Dict ,					__lowercase    :  int		):






						'''simple docstring'''







						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""		)





					@abstractmethod
					def        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  int		):






						'''simple docstring'''







						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""		)





					@abstractmethod
					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






						'''simple docstring'''







						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""		)





					@abstractmethod
					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






						'''simple docstring'''







						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""		)





					@abstractmethod
					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any] ,					__lowercase    :  int=False		):






						'''simple docstring'''







						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""		)







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					def __init__(						self    :  Any ,					__lowercase    :  List[int]		):






						'''simple docstring'''







						super(__lowercase ,					self		).__init__()

						if not isinstance(__lowercase ,					__lowercase		) or len(__lowercase		) == 0:
							raise ValueError(F"""`token_ids` has to be a non-empty list, but is {token_ids}."""		)
						if any((not isinstance(__lowercase ,					__lowercase		) 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}."""		)

						UpperCAmelCase_								=     token_ids

						UpperCAmelCase_								=     len(self.token_ids		)
						UpperCAmelCase_								=     -1  # the index of the currently fulfilled step
						UpperCAmelCase_								=     False





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple		):






						'''simple docstring'''







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





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  int		):






						'''simple docstring'''







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

						if self.completed:
							return False

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Dict ,					__lowercase    :  int		):






						'''simple docstring'''







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

						UpperCAmelCase_								=     False
						UpperCAmelCase_								=     False
						UpperCAmelCase_								=     False

						if self.does_advance(__lowercase		):
							self.fulfilled_idx += 1
							UpperCAmelCase_								=     True
							if self.fulfilled_idx == (self.seqlen - 1):
								UpperCAmelCase_								=     True
							UpperCAmelCase_								=     completed
						else:
							# failed to make progress.
							UpperCAmelCase_								=     True
							self.reset()
						return stepped, completed, reset





					def        SCREAMING_SNAKE_CASE							(						self    :  str		):






						'''simple docstring'''







						UpperCAmelCase_								=     False
						UpperCAmelCase_								=     0





					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







						return self.seqlen - (self.fulfilled_idx + 1)





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  List[str]=False		):






						'''simple docstring'''







						UpperCAmelCase_								=     PhrasalConstraint(self.token_ids		)

						if stateful:
							UpperCAmelCase_								=     self.seqlen
							UpperCAmelCase_								=     self.fulfilled_idx
							UpperCAmelCase_								=     self.completed

						return new_constraint







class   _UpperCamelCase  :







					'''simple docstring'''

					def __init__(						self    :  Optional[int] ,					__lowercase    :  List[List[int]] ,					__lowercase    :  int=True		):






						'''simple docstring'''







						UpperCAmelCase_								=     max([len(__lowercase		) for one in nested_token_ids]		)

						UpperCAmelCase_								=     {}
						for token_ids in nested_token_ids:
							UpperCAmelCase_								=     root
							for tidx, token_id in enumerate(__lowercase		):
								if token_id not in level:
									UpperCAmelCase_								=     {}

								UpperCAmelCase_								=     level[token_id]

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

						UpperCAmelCase_								=     root





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  Union[str, Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.trie

						for current_token in current_seq:
							UpperCAmelCase_								=     start[current_token]

						UpperCAmelCase_								=     list(start.keys()		)

						return next_tokens





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any] ,					__lowercase    :  int		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.next_tokens(__lowercase		)

						return len(__lowercase		) == 0





					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any] ,					__lowercase    :  List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     list(root.values()		)
						if len(__lowercase		) == 0:
							return 1
						else:
							return sum([self.count_leaves(__lowercase		) for nn in next_nodes]		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  str ,					__lowercase    :  List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.count_leaves(__lowercase		)
						return len(__lowercase		) != leaf_count







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					def __init__(						self    :  Any ,					__lowercase    :  List[List[int]]		):






						'''simple docstring'''







						super(__lowercase ,					self		).__init__()

						if not isinstance(__lowercase ,					__lowercase		) or len(__lowercase		) == 0:
							raise ValueError(F"""`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}."""		)
						if any(not isinstance(__lowercase ,					__lowercase		) 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(__lowercase ,					__lowercase		) 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}."""		)

						UpperCAmelCase_								=     DisjunctiveTrie(__lowercase		)
						UpperCAmelCase_								=     nested_token_ids

						UpperCAmelCase_								=     self.trie.max_height
						UpperCAmelCase_								=     []
						UpperCAmelCase_								=     False





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.trie.next_tokens(self.current_seq		)

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int] ,					__lowercase    :  int		):






						'''simple docstring'''







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

						UpperCAmelCase_								=     self.trie.next_tokens(self.current_seq		)

						return token_id in next_tokens





					def        SCREAMING_SNAKE_CASE							(						self    :  str ,					__lowercase    :  int		):






						'''simple docstring'''







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

						UpperCAmelCase_								=     False
						UpperCAmelCase_								=     False
						UpperCAmelCase_								=     False

						if self.does_advance(__lowercase		):
							self.current_seq.append(__lowercase		)
							UpperCAmelCase_								=     True
						else:
							UpperCAmelCase_								=     True
							self.reset()

						UpperCAmelCase_								=     self.trie.reached_leaf(self.current_seq		)
						UpperCAmelCase_								=     completed

						return stepped, completed, reset





					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







						UpperCAmelCase_								=     False
						UpperCAmelCase_								=     []





					def        SCREAMING_SNAKE_CASE							(						self    :  List[str]		):






						'''simple docstring'''







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





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  List[str]=False		):






						'''simple docstring'''







						UpperCAmelCase_								=     DisjunctiveConstraint(self.token_ids		)

						if stateful:
							UpperCAmelCase_								=     self.seqlen
							UpperCAmelCase_								=     self.current_seq
							UpperCAmelCase_								=     self.completed

						return new_constraint







class   _UpperCamelCase  :







					'''simple docstring'''

					def __init__(						self    :  Dict ,					__lowercase    :  List[Constraint]		):






						'''simple docstring'''







						UpperCAmelCase_								=     constraints

						# max # of steps required to fulfill a given constraint
						UpperCAmelCase_								=     max([c.seqlen for c in constraints]		)
						UpperCAmelCase_								=     len(__lowercase		)
						UpperCAmelCase_								=     False

						self.init_state()





					def        SCREAMING_SNAKE_CASE							(						self    :  List[str]		):






						'''simple docstring'''







						UpperCAmelCase_								=     []
						UpperCAmelCase_								=     None
						UpperCAmelCase_								=     [constraint.copy(stateful=__lowercase		) for constraint in self.constraints]





					def        SCREAMING_SNAKE_CASE							(						self    :  int		):






						'''simple docstring'''







						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  Dict		):






						'''simple docstring'''







						UpperCAmelCase_								=     []
						if self.inprogress_constraint is None:
							for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
								UpperCAmelCase_								=     constraint.advance()
								if isinstance(__lowercase ,					__lowercase		):
									token_list.append(__lowercase		)
								elif isinstance(__lowercase ,					__lowercase		):
									token_list.extend(__lowercase		)
						else:
							UpperCAmelCase_								=     self.inprogress_constraint.advance()
							if isinstance(__lowercase ,					__lowercase		):
								token_list.append(__lowercase		)
							elif isinstance(__lowercase ,					__lowercase		):
								token_list.extend(__lowercase		)

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any] ,					__lowercase    :  Optional[List[int]]		):






						'''simple docstring'''







						self.init_state()

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

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int] ,					__lowercase    :  int		):






						'''simple docstring'''







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

						UpperCAmelCase_						,	UpperCAmelCase_								=     False, False

						if self.completed:
							UpperCAmelCase_								=     True
							UpperCAmelCase_								=     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

							UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_								=     self.inprogress_constraint.update(__lowercase		)
							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=__lowercase		)		)
								UpperCAmelCase_								=     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		)
								UpperCAmelCase_								=     None

								if len(self.pending_constraints		) == 0:
									# we're done!
									UpperCAmelCase_								=     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(__lowercase		):
									UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_								=     pending_constraint.update(__lowercase		)

									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(__lowercase		)
										UpperCAmelCase_								=     None

									if not complete and stepped:
										UpperCAmelCase_								=     pending_constraint

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

										UpperCAmelCase_								=     (
										    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.

											UpperCAmelCase_								=     True

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

						return complete, stepped





					def        SCREAMING_SNAKE_CASE							(						self    :  int ,					__lowercase    :  Tuple=True		):






						'''simple docstring'''







						UpperCAmelCase_								=     ConstraintListState(self.constraints		)  # we actually never though self.constraints objects
						# throughout this process. So it's at initialization state.

						if stateful:
							UpperCAmelCase_								=     [
							    constraint.copy(stateful=__lowercase		) for constraint in self.complete_constraints
							]
							if self.inprogress_constraint is not None:
								UpperCAmelCase_								=     self.inprogress_constraint.copy(stateful=__lowercase		)
							UpperCAmelCase_								=     [constraint.copy() for constraint in self.pending_constraints]

						return new_state









from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import (
    BaseOutput,
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_onnx_available,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)







@dataclass
class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					lowerCamelCase     :     Union[List[PIL.Image.Image], np.ndarray]
					lowerCamelCase     :     Optional[List[bool]]


try:
				if not (is_transformers_available() and is_torch_available()):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
				from .pipeline_cycle_diffusion import CycleDiffusionPipeline
				from .pipeline_stable_diffusion import StableDiffusionPipeline
				from .pipeline_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline
				from .pipeline_stable_diffusion_imgaimg import StableDiffusionImgaImgPipeline
				from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
				from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy
				from .pipeline_stable_diffusion_instruct_pixapix import StableDiffusionInstructPixaPixPipeline
				from .pipeline_stable_diffusion_latent_upscale import StableDiffusionLatentUpscalePipeline
				from .pipeline_stable_diffusion_ldmad import StableDiffusionLDMaDPipeline
				from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline
				from .pipeline_stable_diffusion_panorama import StableDiffusionPanoramaPipeline
				from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline
				from .pipeline_stable_diffusion_sag import StableDiffusionSAGPipeline
				from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
				from .pipeline_stable_unclip import StableUnCLIPPipeline
				from .pipeline_stable_unclip_imgaimg import StableUnCLIPImgaImgPipeline
				from .safety_checker import StableDiffusionSafetyChecker
				from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer

try:
				if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import StableDiffusionImageVariationPipeline
else:
				from .pipeline_stable_diffusion_image_variation import StableDiffusionImageVariationPipeline


try:
				if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.26.0""")):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import (
				    StableDiffusionDepthaImgPipeline,
				    StableDiffusionDiffEditPipeline,
				    StableDiffusionPixaPixZeroPipeline,
				)
else:
				from .pipeline_stable_diffusion_depthaimg import StableDiffusionDepthaImgPipeline
				from .pipeline_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline
				from .pipeline_stable_diffusion_pixapix_zero import StableDiffusionPixaPixZeroPipeline


try:
				if not (
				    is_torch_available()
				    and is_transformers_available()
				    and is_k_diffusion_available()
				    and is_k_diffusion_version(""">=""", """0.0.12""")
				):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
				from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline

try:
				if not (is_transformers_available() and is_onnx_available()):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_onnx_objects import *  # noqa F403
else:
				from .pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline, StableDiffusionOnnxPipeline
				from .pipeline_onnx_stable_diffusion_imgaimg import OnnxStableDiffusionImgaImgPipeline
				from .pipeline_onnx_stable_diffusion_inpaint import OnnxStableDiffusionInpaintPipeline
				from .pipeline_onnx_stable_diffusion_inpaint_legacy import OnnxStableDiffusionInpaintPipelineLegacy
				from .pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

if is_transformers_available() and is_flax_available():
				import flax







				@flax.struct.dataclass
				class   _UpperCamelCase  (       A_  ):







									'''simple docstring'''

									lowerCamelCase     :     np.ndarray
									lowerCamelCase     :     List[bool]


				from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState
				from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline
				from .pipeline_flax_stable_diffusion_imgaimg import FlaxStableDiffusionImgaImgPipeline
				from .pipeline_flax_stable_diffusion_inpaint import FlaxStableDiffusionInpaintPipeline
				from .safety_checker_flax import FlaxStableDiffusionSafetyChecker


