

'''simple docstring'''






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






class      lowercase_      (UpperCAmelCase_ ):




  """simple docstring"""




  def __init__( self  :	Tuple     ):
   self.test()






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

    __lowercase    ,      __lowercase    ,      __lowercase				    =					self.update(__a     )
    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    SCREAMING_SNAKE_CASE ( self  :	Dict     ):
   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__  :	Any     ):
   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__  :	Optional[int]     ):
   raise NotImplementedError(
       F"{self.__class__} is an abstract class. Only classes inheriting this class can be called."     )






  @abstractmethod
  def    SCREAMING_SNAKE_CASE ( self  :	str     ):
   raise NotImplementedError(
       F"{self.__class__} is an abstract class. Only classes inheriting this class can be called."     )






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






  @abstractmethod
  def    SCREAMING_SNAKE_CASE ( self  :	int   ,lowercase__  :	Union[str, Any]=False     ):
   raise NotImplementedError(
       F"{self.__class__} is an abstract class. Only classes inheriting this class can be called."     )






class      lowercase_      (UpperCAmelCase_ ):




  """simple docstring"""




  def __init__( self  :	Dict   ,lowercase__  :	Dict     ):
   super(__a   ,self     ).__init__()

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

   __lowercase				    =					token_ids

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






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



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






  def    SCREAMING_SNAKE_CASE ( self  :	Tuple   ,lowercase__  :	Dict     ):



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

   if self.completed:
    return False

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






  def    SCREAMING_SNAKE_CASE ( self  :	List[str]   ,lowercase__  :	Tuple     ):



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

   __lowercase				    =					False
   __lowercase				    =					False
   __lowercase				    =					False

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






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






  def    SCREAMING_SNAKE_CASE ( self  :	Optional[int]     ):
   return self.seqlen - (self.fulfilled_idx + 1)






  def    SCREAMING_SNAKE_CASE ( self  :	Optional[Any]   ,lowercase__  :	Tuple=False     ):
   __lowercase				    =					PhrasalConstraint(self.token_ids     )

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

   return new_constraint






class      lowercase_      :




  """simple docstring"""




  def __init__( self  :	str   ,lowercase__  :	Union[str, Any]   ,lowercase__  :	Any=True     ):
   __lowercase				    =					max([len(__a     ) for one in nested_token_ids]     )

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

     __lowercase				    =					level[token_id]

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

   __lowercase				    =					root






  def    SCREAMING_SNAKE_CASE ( self  :	Dict   ,lowercase__  :	Optional[int]     ):
   __lowercase				    =					self.trie

   for current_token in current_seq:
    __lowercase				    =					start[current_token]

   __lowercase				    =					list(start.keys()     )

   return next_tokens






  def    SCREAMING_SNAKE_CASE ( self  :	Dict   ,lowercase__  :	int     ):
   __lowercase				    =					self.next_tokens(__a     )

   return len(__a     ) == 0






  def    SCREAMING_SNAKE_CASE ( self  :	List[str]   ,lowercase__  :	str     ):
   __lowercase				    =					list(root.values()     )
   if len(__a     ) == 0:
    return 1
   else:
    return sum([self.count_leaves(__a     ) for nn in next_nodes]     )






  def    SCREAMING_SNAKE_CASE ( self  :	Optional[Any]   ,lowercase__  :	Union[str, Any]   ,lowercase__  :	Optional[int]     ):
   __lowercase				    =					self.count_leaves(__a     )
   return len(__a     ) != leaf_count






class      lowercase_      (UpperCAmelCase_ ):




  """simple docstring"""




  def __init__( self  :	Optional[int]   ,lowercase__  :	Dict     ):
   super(__a   ,self     ).__init__()

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

   __lowercase				    =					DisjunctiveTrie(__a     )
   __lowercase				    =					nested_token_ids

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






  def    SCREAMING_SNAKE_CASE ( self  :	Union[str, Any]     ):
   __lowercase				    =					self.trie.next_tokens(self.current_seq     )

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






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



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

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

   return token_id in next_tokens






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



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

   __lowercase				    =					False
   __lowercase				    =					False
   __lowercase				    =					False

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

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

   return stepped, completed, reset






  def    SCREAMING_SNAKE_CASE ( self  :	List[Any]     ):
   __lowercase				    =					False
   __lowercase				    =					[]






  def    SCREAMING_SNAKE_CASE ( self  :	Dict     ):



   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__  :	str=False     ):
   __lowercase				    =					DisjunctiveConstraint(self.token_ids     )

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

   return new_constraint






class      lowercase_      :




  """simple docstring"""




  def __init__( self  :	Any   ,lowercase__  :	Optional[Any]     ):
   __lowercase				    =					constraints

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

   self.init_state()






  def    SCREAMING_SNAKE_CASE ( self  :	Dict     ):
   __lowercase				    =					[]
   __lowercase				    =					None
   __lowercase				    =					[constraint.copy(stateful=__a     ) for constraint in self.constraints]






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

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






  def    SCREAMING_SNAKE_CASE ( self  :	Union[str, Any]   ,lowercase__  :	Optional[int]     ):
   self.init_state()

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

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






  def    SCREAMING_SNAKE_CASE ( self  :	Tuple   ,lowercase__  :	Optional[int]     ):



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

   __lowercase    ,      __lowercase				    =					False, False

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

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

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

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

      if not complete and stepped:
       __lowercase				    =					pending_constraint

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

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

        __lowercase				    =					True

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

   return complete, stepped






  def    SCREAMING_SNAKE_CASE ( self  :	Any   ,lowercase__  :	int=True     ):
   __lowercase				    =					ConstraintListState(self.constraints     )  # we actually never though self.constraints objects
   # throughout this process. So it's at initialization state.

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

   return new_state






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class  lowerCamelCase__	(    unittest.TestCase):
     """simple docstring"""



     def     _a   (self     ):






           '''simple docstring'''



           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()







     @property
     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		1
           lowerCamelCase             =		3
           lowerCamelCase             =		(32, 32)

           lowerCamelCase             =		floats_tensor((batch_size, num_channels) + sizes					,      rng=random.Random(0     )     ).to(__a     )
           return image







     @property
     def     _a   (self     ):






           '''simple docstring'''



           torch.manual_seed(0     )
           lowerCamelCase             =		UNetaDConditionModel(
               block_out_channels=(32, 32, 64)					,      layers_per_block=2					,      sample_size=32					,      in_channels=7					,      out_channels=4					,      down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")					,      up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")					,      cross_attention_dim=32					,      attention_head_dim=8					,      use_linear_projection=__a					,      only_cross_attention=(True, True, False)					,      num_class_embeds=1_00					,      )
           return model







     @property
     def     _a   (self     ):






           '''simple docstring'''



           torch.manual_seed(0     )
           lowerCamelCase             =		AutoencoderKL(
               block_out_channels=[32, 32, 64]					,      in_channels=3					,      out_channels=3					,      down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]					,      up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]					,      latent_channels=4					,      )
           return model







     @property
     def     _a   (self     ):






           '''simple docstring'''



           torch.manual_seed(0     )
           lowerCamelCase             =		CLIPTextConfig(
               bos_token_id=0					,      eos_token_id=2					,      hidden_size=32					,      intermediate_size=37					,      layer_norm_eps=1E-05					,      num_attention_heads=4					,      num_hidden_layers=5					,      pad_token_id=1					,      vocab_size=10_00					,      hidden_act="gelu"					,      projection_dim=5_12					,      )
           return CLIPTextModel(__a     )







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowerCamelCase             =		self.dummy_cond_unet_upscale
           lowerCamelCase             =		DDPMScheduler()
           lowerCamelCase             =		DDIMScheduler(prediction_type="v_prediction"     )
           lowerCamelCase             =		self.dummy_vae
           lowerCamelCase             =		self.dummy_text_encoder
           lowerCamelCase             =		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

           lowerCamelCase             =		self.dummy_image.cpu().permute(0					,      2					,      3					,      1     )[0]
           lowerCamelCase             =		Image.fromarray(np.uinta(__a     )     ).convert("RGB"     ).resize((64, 64)     )

           # make sure here that pndm scheduler skips prk
           lowerCamelCase             =		StableDiffusionUpscalePipeline(
               unet=__a					,      low_res_scheduler=__a					,      scheduler=__a					,      vae=__a					,      text_encoder=__a					,      tokenizer=__a					,      max_noise_level=3_50					,      )
           lowerCamelCase             =		sd_pipe.to(__a     )
           sd_pipe.set_progress_bar_config(disable=__a     )

           lowerCamelCase             =		"A painting of a squirrel eating a burger"
           lowerCamelCase             =		torch.Generator(device=__a     ).manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      guidance_scale=6.0					,      noise_level=20					,      num_inference_steps=2					,      output_type="np"					,      )

           lowerCamelCase             =		output.images

           lowerCamelCase             =		torch.Generator(device=__a     ).manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      guidance_scale=6.0					,      noise_level=20					,      num_inference_steps=2					,      output_type="np"					,      return_dict=__a					,      )[0]

           lowerCamelCase             =		image[0, -3:, -3:, -1]
           lowerCamelCase             =		image_from_tuple[0, -3:, -3:, -1]

           lowerCamelCase             =		low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)
           lowerCamelCase             =		np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]     )

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







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		"cpu"  # ensure determinism for the device-dependent torch.Generator
           lowerCamelCase             =		self.dummy_cond_unet_upscale
           lowerCamelCase             =		DDPMScheduler()
           lowerCamelCase             =		DDIMScheduler(prediction_type="v_prediction"     )
           lowerCamelCase             =		self.dummy_vae
           lowerCamelCase             =		self.dummy_text_encoder
           lowerCamelCase             =		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

           lowerCamelCase             =		self.dummy_image.cpu().permute(0					,      2					,      3					,      1     )[0]
           lowerCamelCase             =		Image.fromarray(np.uinta(__a     )     ).convert("RGB"     ).resize((64, 64)     )

           # make sure here that pndm scheduler skips prk
           lowerCamelCase             =		StableDiffusionUpscalePipeline(
               unet=__a					,      low_res_scheduler=__a					,      scheduler=__a					,      vae=__a					,      text_encoder=__a					,      tokenizer=__a					,      max_noise_level=3_50					,      )
           lowerCamelCase             =		sd_pipe.to(__a     )
           sd_pipe.set_progress_bar_config(disable=__a     )

           lowerCamelCase             =		"A painting of a squirrel eating a burger"
           lowerCamelCase             =		sd_pipe(
               2 * [prompt]					,      image=2 * [low_res_image]					,      guidance_scale=6.0					,      noise_level=20					,      num_inference_steps=2					,      output_type="np"					,      )
           lowerCamelCase             =		output.images
           assert image.shape[0] == 2

           lowerCamelCase             =		torch.Generator(device=__a     ).manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      num_images_per_prompt=2					,      guidance_scale=6.0					,      noise_level=20					,      num_inference_steps=2					,      output_type="np"					,      )
           lowerCamelCase             =		output.images
           assert image.shape[0] == 2







     @unittest.skipIf(torch_device != "cuda"					,      "This test requires a GPU"     )
     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		self.dummy_cond_unet_upscale
           lowerCamelCase             =		DDPMScheduler()
           lowerCamelCase             =		DDIMScheduler(prediction_type="v_prediction"     )
           lowerCamelCase             =		self.dummy_vae
           lowerCamelCase             =		self.dummy_text_encoder
           lowerCamelCase             =		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

           lowerCamelCase             =		self.dummy_image.cpu().permute(0					,      2					,      3					,      1     )[0]
           lowerCamelCase             =		Image.fromarray(np.uinta(__a     )     ).convert("RGB"     ).resize((64, 64)     )

           # put models in fp16, except vae as it overflows in fp16
           lowerCamelCase             =		unet.half()
           lowerCamelCase             =		text_encoder.half()

           # make sure here that pndm scheduler skips prk
           lowerCamelCase             =		StableDiffusionUpscalePipeline(
               unet=__a					,      low_res_scheduler=__a					,      scheduler=__a					,      vae=__a					,      text_encoder=__a					,      tokenizer=__a					,      max_noise_level=3_50					,      )
           lowerCamelCase             =		sd_pipe.to(__a     )
           sd_pipe.set_progress_bar_config(disable=__a     )

           lowerCamelCase             =		"A painting of a squirrel eating a burger"
           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		sd_pipe(
               [prompt]					,      image=__a					,      generator=__a					,      num_inference_steps=2					,      output_type="np"					,      ).images

           lowerCamelCase             =		low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)





@slow
@require_torch_gpu
class  lowerCamelCase__	(    unittest.TestCase):
     """simple docstring"""



     def     _a   (self     ):






           '''simple docstring'''



           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		load_image(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/sd2-upscale/low_res_cat.png"     )
           lowerCamelCase             =		load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
               "/upsampled_cat.npy"     )

           lowerCamelCase             =		"stabilityai/stable-diffusion-x4-upscaler"
           lowerCamelCase             =		StableDiffusionUpscalePipeline.from_pretrained(__a     )
           pipe.to(__a     )
           pipe.set_progress_bar_config(disable=__a     )
           pipe.enable_attention_slicing()

           lowerCamelCase             =		"a cat sitting on a park bench"

           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		pipe(
               prompt=__a					,      image=__a					,      generator=__a					,      output_type="np"					,      )
           lowerCamelCase             =		output.images[0]

           assert image.shape == (5_12, 5_12, 3)
           assert np.abs(expected_image - image     ).max() < 1E-3







     def     _a   (self     ):






           '''simple docstring'''



           lowerCamelCase             =		load_image(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/sd2-upscale/low_res_cat.png"     )
           lowerCamelCase             =		load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
               "/upsampled_cat_fp16.npy"     )

           lowerCamelCase             =		"stabilityai/stable-diffusion-x4-upscaler"
           lowerCamelCase             =		StableDiffusionUpscalePipeline.from_pretrained(
               __a					,      torch_dtype=torch.floataa					,      )
           pipe.to(__a     )
           pipe.set_progress_bar_config(disable=__a     )
           pipe.enable_attention_slicing()

           lowerCamelCase             =		"a cat sitting on a park bench"

           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		pipe(
               prompt=__a					,      image=__a					,      generator=__a					,      output_type="np"					,      )
           lowerCamelCase             =		output.images[0]

           assert image.shape == (5_12, 5_12, 3)
           assert np.abs(expected_image - image     ).max() < 5E-1







     def     _a   (self     ):






           '''simple docstring'''



           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           lowerCamelCase             =		load_image(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/sd2-upscale/low_res_cat.png"     )

           lowerCamelCase             =		"stabilityai/stable-diffusion-x4-upscaler"
           lowerCamelCase             =		StableDiffusionUpscalePipeline.from_pretrained(
               __a					,      torch_dtype=torch.floataa					,      )
           pipe.to(__a     )
           pipe.set_progress_bar_config(disable=__a     )
           pipe.enable_attention_slicing(1     )
           pipe.enable_sequential_cpu_offload()

           lowerCamelCase             =		"a cat sitting on a park bench"

           lowerCamelCase             =		torch.manual_seed(0     )
           lowerCamelCase             =		pipe(
               prompt=__a					,      image=__a					,      generator=__a					,      num_inference_steps=5					,      output_type="np"					,      )

           lowerCamelCase             =		torch.cuda.max_memory_allocated()
           # make sure that less than 2.9 GB is allocated
           assert mem_bytes < 2.9 * 10**9