'''simple docstring'''
from abc import ABC, abstractmethod
from typing import List, Optional







class lowerCamelCase							(	_a					):




  '''simple docstring'''



  def __init__(      self			:       Tuple  )				->Tuple:
        # test for the above condition
        self.test()






  def  lowerCAmelCase__ (      self			:       Any  )				->Union[str, Any]:
        UpperCAmelCase_            =  0
        UpperCAmelCase_            =  False
        while not completed:
              if counter == 1:
                    self.reset()
              UpperCAmelCase_            =  self.advance()
              if not self.does_advance(snake_case_  ):
                    raise Exception(
                        '''Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.'''  )

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

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

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






  @abstractmethod
  def  lowerCAmelCase__ (      self			:       Tuple  )				->Any:
        raise NotImplementedError(
            f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""  )






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






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






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






  @abstractmethod
  def  lowerCAmelCase__ (      self			:       int  )				->Any:
        raise NotImplementedError(
            f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""  )






  @abstractmethod
  def  lowerCAmelCase__ (      self			:       Dict  ,							UpperCAmelCase__			:       Optional[int]=False  )				->Union[str, Any]:
        raise NotImplementedError(
            f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""  )







class lowerCamelCase							(	_a					):




  '''simple docstring'''



  def __init__(      self			:       Optional[int]  ,							UpperCAmelCase__			:       List[str]  )				->List[Any]:
        super(snake_case_  ,							self  ).__init__()

        if not isinstance(snake_case_  ,							snake_case_  ) or len(snake_case_  ) == 0:
              raise ValueError(f"""`token_ids` has to be a non-empty list, but is {token_ids}."""  )
        if any((not isinstance(snake_case_  ,							snake_case_  ) 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  lowerCAmelCase__ (      self			:       Tuple  )				->str:
        if self.completed:
              return None
        return self.token_ids[self.fulfilled_idx + 1]






  def  lowerCAmelCase__ (      self			:       Optional[Any]  ,							UpperCAmelCase__			:       Dict  )				->str:
        if not isinstance(snake_case_  ,							snake_case_  ):
              raise ValueError(f"""`token_id` has to be an `int`, but is {token_id} of type {type(snake_case_  )}"""  )

        if self.completed:
              return False

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






  def  lowerCAmelCase__ (      self			:       str  ,							UpperCAmelCase__			:       List[str]  )				->Optional[Any]:
        if not isinstance(snake_case_  ,							snake_case_  ):
              raise ValueError(f"""`token_id` has to be an `int`, but is {token_id} of type {type(snake_case_  )}"""  )

        UpperCAmelCase_            =  False
        UpperCAmelCase_            =  False
        UpperCAmelCase_            =  False

        if self.does_advance(snake_case_  ):
              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  lowerCAmelCase__ (      self			:       Dict  )				->List[Any]:
        UpperCAmelCase_            =  False
        UpperCAmelCase_            =  0






  def  lowerCAmelCase__ (      self			:       Tuple  )				->Optional[Any]:
        return self.seqlen - (self.fulfilled_idx + 1)






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       List[str]=False  )				->int:
        UpperCAmelCase_            =  PhrasalConstraint(self.token_ids  )

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

        return new_constraint







class lowerCamelCase							:




  '''simple docstring'''



  def __init__(      self			:       List[str]  ,							UpperCAmelCase__			:       Tuple  ,							UpperCAmelCase__			:       Dict=True  )				->List[Any]:
        UpperCAmelCase_            =  max([len(snake_case_  ) for one in nested_token_ids]  )

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

                    UpperCAmelCase_            =  level[token_id]

        if no_subsets and self.has_subsets(snake_case_  ,							snake_case_  ):
              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  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       Dict  )				->Union[str, Any]:
        UpperCAmelCase_            =  self.trie

        for current_token in current_seq:
              UpperCAmelCase_            =  start[current_token]

        UpperCAmelCase_            =  list(start.keys()  )

        return next_tokens






  def  lowerCAmelCase__ (      self			:       Tuple  ,							UpperCAmelCase__			:       List[Any]  )				->Optional[int]:
        UpperCAmelCase_            =  self.next_tokens(snake_case_  )

        return len(snake_case_  ) == 0






  def  lowerCAmelCase__ (      self			:       Dict  ,							UpperCAmelCase__			:       Optional[int]  )				->Optional[int]:
        UpperCAmelCase_            =  list(root.values()  )
        if len(snake_case_  ) == 0:
              return 1
        else:
              return sum([self.count_leaves(snake_case_  ) for nn in next_nodes]  )






  def  lowerCAmelCase__ (      self			:       Optional[Any]  ,							UpperCAmelCase__			:       int  ,							UpperCAmelCase__			:       Optional[Any]  )				->Union[str, Any]:
        UpperCAmelCase_            =  self.count_leaves(snake_case_  )
        return len(snake_case_  ) != leaf_count







class lowerCamelCase							(	_a					):




  '''simple docstring'''



  def __init__(      self			:       Any  ,							UpperCAmelCase__			:       Optional[Any]  )				->str:
        super(snake_case_  ,							self  ).__init__()

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

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






  def  lowerCAmelCase__ (      self			:       int  )				->str:
        UpperCAmelCase_            =  self.trie.next_tokens(self.current_seq  )

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






  def  lowerCAmelCase__ (      self			:       str  ,							UpperCAmelCase__			:       Optional[int]  )				->List[Any]:
        if not isinstance(snake_case_  ,							snake_case_  ):
              raise ValueError(f"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(snake_case_  )}"""  )

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

        return token_id in next_tokens






  def  lowerCAmelCase__ (      self			:       Optional[Any]  ,							UpperCAmelCase__			:       Optional[Any]  )				->List[Any]:
        if not isinstance(snake_case_  ,							snake_case_  ):
              raise ValueError(f"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(snake_case_  )}"""  )

        UpperCAmelCase_            =  False
        UpperCAmelCase_            =  False
        UpperCAmelCase_            =  False

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

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

        return stepped, completed, reset






  def  lowerCAmelCase__ (      self			:       Tuple  )				->List[Any]:
        UpperCAmelCase_            =  False
        UpperCAmelCase_            =  []






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






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  ,							UpperCAmelCase__			:       str=False  )				->Optional[int]:
        UpperCAmelCase_            =  DisjunctiveConstraint(self.token_ids  )

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

        return new_constraint







class lowerCamelCase							:




  '''simple docstring'''



  def __init__(      self			:       Any  ,							UpperCAmelCase__			:       Union[str, Any]  )				->Any:
        UpperCAmelCase_            =  constraints

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

        self.init_state()






  def  lowerCAmelCase__ (      self			:       str  )				->Union[str, Any]:
        UpperCAmelCase_            =  []
        UpperCAmelCase_            =  None
        UpperCAmelCase_            =  [constraint.copy(stateful=snake_case_  ) for constraint in self.constraints]






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->int:
        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  lowerCAmelCase__ (      self			:       str  )				->Tuple:
        UpperCAmelCase_            =  []
        if self.inprogress_constraint is None:
              for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
                    UpperCAmelCase_            =  constraint.advance()
                    if isinstance(snake_case_  ,							snake_case_  ):
                          token_list.append(snake_case_  )
                    elif isinstance(snake_case_  ,							snake_case_  ):
                          token_list.extend(snake_case_  )
        else:
              UpperCAmelCase_            =  self.inprogress_constraint.advance()
              if isinstance(snake_case_  ,							snake_case_  ):
                    token_list.append(snake_case_  )
              elif isinstance(snake_case_  ,							snake_case_  ):
                    token_list.extend(snake_case_  )

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






  def  lowerCAmelCase__ (      self			:       int  ,							UpperCAmelCase__			:       int  )				->Optional[Any]:
        self.init_state()

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

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






  def  lowerCAmelCase__ (      self			:       Dict  ,							UpperCAmelCase__			:       Optional[Any]  )				->Union[str, Any]:
        if not isinstance(snake_case_  ,							snake_case_  ):
              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(snake_case_  )
              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=snake_case_  )  )
                    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(snake_case_  ):
                          UpperCAmelCase_      ,		UpperCAmelCase_      ,		UpperCAmelCase_            =  pending_constraint.update(snake_case_  )

                          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(snake_case_  )
                                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  lowerCAmelCase__ (      self			:       Optional[int]  ,							UpperCAmelCase__			:       Optional[Any]=True  )				->Union[str, Any]:
        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=snake_case_  ) for constraint in self.complete_constraints
              ]
              if self.inprogress_constraint is not None:
                    UpperCAmelCase_            =  self.inprogress_constraint.copy(stateful=snake_case_  )
              UpperCAmelCase_            =  [constraint.copy() for constraint in self.pending_constraints]

        return new_state






'''simple docstring'''
import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







@slow
@require_torch_gpu
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Any  )				->Optional[int]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






  def  lowerCAmelCase__ (      self			:       str  )				->List[str]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2






  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 5e-1






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Optional[int]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_dpmpp_2m'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe(
            [prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=7.5  ,							num_inference_steps=15  ,							output_type='''np'''  ,							use_karras_sigmas=UpperCAmelCase__  ,							)

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array(
            [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2





