import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Union[dict, list, tuple, torch.Tensor]     )							-> List[Tuple[int, ...]]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				[]
  if isinstance(__UpperCamelCase			,					__UpperCamelCase     ):
    for v in tree.values():
      shapes.extend(_fetch_dims(__UpperCamelCase     )     )
  elif isinstance(__UpperCamelCase			,					(list, tuple)     ):
    for t in tree:
      shapes.extend(_fetch_dims(__UpperCamelCase     )     )
  elif isinstance(__UpperCamelCase			,					torch.Tensor     ):
    shapes.append(tree.shape     )
  else:
    raise ValueError("""Not supported"""     )

  return shapes



@torch.jit.ignore
def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	int			,					__UpperCamelCase   :	Tuple[int, ...]     )							-> Tuple[int, ...]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				[]
  for d in reversed(__UpperCamelCase     ):
    idx.append(flat_idx % d     )
    SCREAMING_SNAKE_CASE__      							=				flat_idx // d

  return tuple(reversed(__UpperCamelCase     )     )



@torch.jit.ignore
def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Sequence[int]			,					__UpperCamelCase   :	Sequence[int]			,					__UpperCamelCase   :	Sequence[int]			,					__UpperCamelCase   :	Optional[Sequence[bool]] = None			,					__UpperCamelCase   :	Optional[Sequence[bool]] = None			,					)							-> List[Tuple[slice, ...]]:







  """simple docstring"""





  def reduce_edge_list(__UpperCamelCase   :	List[bool]     ) -> None:
    SCREAMING_SNAKE_CASE__      							=				True
    for i in range(len(__UpperCamelCase     )     ):
      SCREAMING_SNAKE_CASE__      							=				-1 * (i + 1)
      l[reversed_idx] &= tally
      SCREAMING_SNAKE_CASE__      							=				l[reversed_idx]

  if start_edges is None:
    SCREAMING_SNAKE_CASE__      							=				[s == 0 for s in start]
    reduce_edge_list(__UpperCamelCase     )
  if end_edges is None:
    SCREAMING_SNAKE_CASE__      							=				[e == (d - 1) for e, d in zip(__UpperCamelCase			,					__UpperCamelCase     )]
    reduce_edge_list(__UpperCamelCase     )

  # Base cases. Either start/end are empty and we're done, or the final,
  # one-dimensional tensor can be simply sliced
  if len(__UpperCamelCase     ) == 0:
    return [()]
  elif len(__UpperCamelCase     ) == 1:
    return [(slice(start[0]			,					end[0] + 1     ),)]

  SCREAMING_SNAKE_CASE__      							=				[]
  SCREAMING_SNAKE_CASE__      							=				[]

  # Dimensions common to start and end can be selected directly
  for s, e in zip(__UpperCamelCase			,					__UpperCamelCase     ):
    if s == e:
      path_list.append(slice(__UpperCamelCase			,					s + 1     )     )
    else:
      break

  SCREAMING_SNAKE_CASE__      							=				tuple(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				len(__UpperCamelCase     )

  # start == end, and we're done
  if divergence_idx == len(__UpperCamelCase     ):
    return [path]

  def upper() -> Tuple[Tuple[slice, ...], ...]:
    assert start_edges is not None
    assert end_edges is not None

    SCREAMING_SNAKE_CASE__      							=				start[divergence_idx]
    return tuple(
        path + (slice(__UpperCamelCase			,					sdi + 1     ),) + s
        for s in _get_minimal_slice_set(
            start[divergence_idx + 1 :]			,					[d - 1 for d in dims[divergence_idx + 1 :]]			,					dims[divergence_idx + 1 :]			,					start_edges=start_edges[divergence_idx + 1 :]			,					end_edges=[True for _ in end_edges[divergence_idx + 1 :]]			,					)     )

  def lower() -> Tuple[Tuple[slice, ...], ...]:
    assert start_edges is not None
    assert end_edges is not None

    SCREAMING_SNAKE_CASE__      							=				end[divergence_idx]
    return tuple(
        path + (slice(__UpperCamelCase			,					edi + 1     ),) + s
        for s in _get_minimal_slice_set(
            [0 for _ in start[divergence_idx + 1 :]]			,					end[divergence_idx + 1 :]			,					dims[divergence_idx + 1 :]			,					start_edges=[True for _ in start_edges[divergence_idx + 1 :]]			,					end_edges=end_edges[divergence_idx + 1 :]			,					)     )

  # If both start and end are at the edges of the subtree rooted at
  # divergence_idx, we can just select the whole subtree at once
  if start_edges[divergence_idx] and end_edges[divergence_idx]:
    slices.append(path + (slice(start[divergence_idx]			,					end[divergence_idx] + 1     ),)     )
  # If just start is at the edge, we can grab almost all of the subtree,
  # treating only the ragged bottom edge as an edge case
  elif start_edges[divergence_idx]:
    slices.append(path + (slice(start[divergence_idx]			,					end[divergence_idx]     ),)     )
    slices.extend(lower()     )
  # Analogous to the previous case, but the top is ragged this time
  elif end_edges[divergence_idx]:
    slices.extend(upper()     )
    slices.append(path + (slice(start[divergence_idx] + 1			,					end[divergence_idx] + 1     ),)     )
  # If both sides of the range are ragged, we need to handle both sides
  # separately. If there's contiguous meat in between them, we can index it
  # in one big chunk
  else:
    slices.extend(upper()     )
    SCREAMING_SNAKE_CASE__      							=				end[divergence_idx] - start[divergence_idx]
    if middle_ground > 1:
      slices.append(path + (slice(start[divergence_idx] + 1			,					end[divergence_idx]     ),)     )
    slices.extend(lower()     )

  return slices



@torch.jit.ignore
def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	torch.Tensor			,					__UpperCamelCase   :	int			,					__UpperCamelCase   :	int			,					__UpperCamelCase   :	int     )							-> torch.Tensor:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				t.shape[:no_batch_dims]
  SCREAMING_SNAKE_CASE__      							=				list(_flat_idx_to_idx(__UpperCamelCase			,					__UpperCamelCase     )     )
  # _get_minimal_slice_set is inclusive
  SCREAMING_SNAKE_CASE__      							=				list(_flat_idx_to_idx(flat_end - 1			,					__UpperCamelCase     )     )

  # Get an ordered list of slices to perform
  SCREAMING_SNAKE_CASE__      							=				_get_minimal_slice_set(
      __UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase			,					)

  SCREAMING_SNAKE_CASE__      							=				[t[s] for s in slices]

  return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]     ) for s in sliced_tensors]     )



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Callable			,					__UpperCamelCase   :	Dict[str, Any]			,					__UpperCamelCase   :	int			,					__UpperCamelCase   :	int			,					__UpperCamelCase   :	bool = False			,					__UpperCamelCase   :	Any = None			,					__UpperCamelCase   :	bool = False			,					)							-> Any:







  """simple docstring"""





  if not (len(__UpperCamelCase     ) > 0):
    raise ValueError("""Must provide at least one input"""     )

  SCREAMING_SNAKE_CASE__      							=				[shape[:no_batch_dims] for shape in _fetch_dims(__UpperCamelCase     )]
  SCREAMING_SNAKE_CASE__      							=				tuple([max(__UpperCamelCase     ) for s in zip(*__UpperCamelCase     )]     )

  def _prep_inputs(__UpperCamelCase   :	torch.Tensor     ) -> torch.Tensor:
    if not low_mem:
      if not sum(t.shape[:no_batch_dims]     ) == no_batch_dims:
        SCREAMING_SNAKE_CASE__      							=				t.expand(orig_batch_dims + t.shape[no_batch_dims:]     )
      SCREAMING_SNAKE_CASE__      							=				t.reshape(-1			,					*t.shape[no_batch_dims:]     )
    else:
      SCREAMING_SNAKE_CASE__      							=				t.expand(orig_batch_dims + t.shape[no_batch_dims:]     )
    return t

  SCREAMING_SNAKE_CASE__      							=				tensor_tree_map(_prep_inputs			,					__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				None
  if _out is not None:
    SCREAMING_SNAKE_CASE__      							=				tensor_tree_map(lambda __UpperCamelCase     : t.view([-1] + list(t.shape[no_batch_dims:]     )     )			,					_out     )

  SCREAMING_SNAKE_CASE__      							=				1
  for d in orig_batch_dims:
    flat_batch_dim *= d

  SCREAMING_SNAKE_CASE__      							=				flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

  def _select_chunk(__UpperCamelCase   :	torch.Tensor     ) -> torch.Tensor:
    return t[i : i + chunk_size] if t.shape[0] != 1 else t

  SCREAMING_SNAKE_CASE__      							=				0
  SCREAMING_SNAKE_CASE__      							=				prepped_outputs
  for _ in range(__UpperCamelCase     ):
    # Chunk the input
    if not low_mem:
      SCREAMING_SNAKE_CASE__      							=				_select_chunk
    else:
      SCREAMING_SNAKE_CASE__      							=				partial(
          _chunk_slice			,					flat_start=__UpperCamelCase			,					flat_end=min(__UpperCamelCase			,					i + chunk_size     )			,					no_batch_dims=len(__UpperCamelCase     )			,					)

    SCREAMING_SNAKE_CASE__      							=				tensor_tree_map(__UpperCamelCase			,					__UpperCamelCase     )

    # Run the layer on the chunk
    SCREAMING_SNAKE_CASE__      							=				layer(**__UpperCamelCase     )

    # Allocate space for the output
    if out is None:
      SCREAMING_SNAKE_CASE__      							=				tensor_tree_map(lambda __UpperCamelCase     : t.new_zeros((flat_batch_dim,) + t.shape[1:]     )			,					__UpperCamelCase     )

    # Put the chunk in its pre-allocated space
    if isinstance(__UpperCamelCase			,					__UpperCamelCase     ):

      def assign(__UpperCamelCase   :	dict			,					__UpperCamelCase   :	dict     ) -> None:
        for k, v in da.items():
          if isinstance(__UpperCamelCase			,					__UpperCamelCase     ):
            assign(__UpperCamelCase			,					da[k]     )
          else:
            if _add_into_out:
              v[i : i + chunk_size] += da[k]
            else:
              SCREAMING_SNAKE_CASE__      							=				da[k]

      assign(__UpperCamelCase			,					__UpperCamelCase     )
    elif isinstance(__UpperCamelCase			,					__UpperCamelCase     ):
      for xa, xa in zip(__UpperCamelCase			,					__UpperCamelCase     ):
        if _add_into_out:
          xa[i : i + chunk_size] += xa
        else:
          SCREAMING_SNAKE_CASE__      							=				xa
    elif isinstance(__UpperCamelCase			,					torch.Tensor     ):
      if _add_into_out:
        out[i : i + chunk_size] += output_chunk
      else:
        SCREAMING_SNAKE_CASE__      							=				output_chunk
    else:
      raise ValueError("""Not supported"""     )

    i += chunk_size

  SCREAMING_SNAKE_CASE__      							=				tensor_tree_map(lambda __UpperCamelCase     : t.view(orig_batch_dims + t.shape[1:]     )			,					__UpperCamelCase     )

  return out







class __snake_case    :




   def __init__(   self      :    List[str]	,     _lowercase      :    int = 5_12	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				max_chunk_size
     SCREAMING_SNAKE_CASE__      							=				None
     SCREAMING_SNAKE_CASE__      							=				None




   def      __a						(   self      :    Optional[Any]	,     _lowercase      :    Callable	,     _lowercase      :    tuple	,     _lowercase      :    int				):



     """simple docstring"""






     logging.info("""Tuning chunk size..."""				)

     if min_chunk_size >= self.max_chunk_size:
       return min_chunk_size

     SCREAMING_SNAKE_CASE__      							=				[2**l for l in range(int(math.log(self.max_chunk_size	,     2				)				) + 1				)]
     SCREAMING_SNAKE_CASE__      							=				[c for c in candidates if c > min_chunk_size]
     SCREAMING_SNAKE_CASE__      							=				[min_chunk_size] + candidates
     candidates[-1] += 4

     def test_chunk_size(_lowercase      :    int				) -> bool:
       try:
         with torch.no_grad():
           fn(*_lowercase	,     chunk_size=_lowercase				)
         return True
       except RuntimeError:
         return False

     SCREAMING_SNAKE_CASE__      							=				0
     SCREAMING_SNAKE_CASE__      							=				len(_lowercase				) - 1
     while i > min_viable_chunk_size_index:
       SCREAMING_SNAKE_CASE__      							=				test_chunk_size(candidates[i]				)
       if not viable:
         SCREAMING_SNAKE_CASE__      							=				(min_viable_chunk_size_index + i) // 2
       else:
         SCREAMING_SNAKE_CASE__      							=				i
         SCREAMING_SNAKE_CASE__      							=				(i + len(_lowercase				) - 1) // 2

     return candidates[min_viable_chunk_size_index]




   def      __a						(   self      :    List[Any]	,     _lowercase      :    Iterable	,     _lowercase      :    Iterable				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				True
     for aa, aa in zip(_lowercase	,     _lowercase				):
       assert type(_lowercase				) == type(_lowercase				)
       if isinstance(_lowercase	,     (list, tuple)				):
         consistent &= self._compare_arg_caches(_lowercase	,     _lowercase				)
       elif isinstance(_lowercase	,     _lowercase				):
         SCREAMING_SNAKE_CASE__      							=				[v for _, v in sorted(aa.items()	,     key=lambda _lowercase				: x[0]				)]
         SCREAMING_SNAKE_CASE__      							=				[v for _, v in sorted(aa.items()	,     key=lambda _lowercase				: x[0]				)]
         consistent &= self._compare_arg_caches(_lowercase	,     _lowercase				)
       else:
         consistent &= aa == aa

     return consistent




   def      __a						(   self      :    Tuple	,     _lowercase      :    Callable	,     _lowercase      :    tuple	,     _lowercase      :    int	,     ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				True
     SCREAMING_SNAKE_CASE__      							=				tree_map(lambda _lowercase				: a.shape if isinstance(_lowercase	,     torch.Tensor				) else a	,     _lowercase	,     _lowercase				)
     if self.cached_arg_data is not None:
       # If args have changed shape/value, we need to re-tune
       assert len(self.cached_arg_data				) == len(_lowercase				)
       SCREAMING_SNAKE_CASE__      							=				self._compare_arg_caches(self.cached_arg_data	,     _lowercase				)
     else:
       # Otherwise, we can reuse the precomputed value
       SCREAMING_SNAKE_CASE__      							=				False

     if not consistent:
       SCREAMING_SNAKE_CASE__      							=				self._determine_favorable_chunk_size(
           _lowercase	,     _lowercase	,     _lowercase	,     )
       SCREAMING_SNAKE_CASE__      							=				arg_data

     assert self.cached_chunk_size is not None

     return self.cached_chunk_size

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 __snake_case    (						unittest.TestCase		):




   def      __a						(   self      :    str				):



     """simple docstring"""






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




   @property
   def      __a						(   self      :    str				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				1
     SCREAMING_SNAKE_CASE__      							=				3
     SCREAMING_SNAKE_CASE__      							=				(32, 32)

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




   @property
   def      __a						(   self      :    Dict				):



     """simple docstring"""






     torch.manual_seed(0				)
     SCREAMING_SNAKE_CASE__      							=				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=_lowercase	,     only_cross_attention=(True, True, False)	,     num_class_embeds=1_00	,     )
     return model




   @property
   def      __a						(   self      :    Optional[int]				):



     """simple docstring"""






     torch.manual_seed(0				)
     SCREAMING_SNAKE_CASE__      							=				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      :    Dict				):



     """simple docstring"""






     torch.manual_seed(0				)
     SCREAMING_SNAKE_CASE__      							=				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(_lowercase				)




   def      __a						(   self      :    Dict				):



     """simple docstring"""






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

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

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

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

     SCREAMING_SNAKE_CASE__      							=				output.images

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

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

     SCREAMING_SNAKE_CASE__      							=				low_res_image.size[0] * 4
     assert image.shape == (1, expected_height_width, expected_height_width, 3)
     SCREAMING_SNAKE_CASE__      							=				np.array([0.31_13, 0.39_10, 0.42_72, 0.48_59, 0.50_61, 0.46_52, 0.53_62, 0.57_15, 0.56_61]				)

     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      :    Dict				):



     """simple docstring"""






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

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

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

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

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




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



     """simple docstring"""






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

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

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

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

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

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







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




   def      __a						(   self      :    Optional[int]				):



     """simple docstring"""






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




   def      __a						(   self      :    str				):



     """simple docstring"""






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

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

     SCREAMING_SNAKE_CASE__      							=				"""a cat sitting on a park bench"""

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

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




   def      __a						(   self      :    List[Any]				):



     """simple docstring"""






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

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

     SCREAMING_SNAKE_CASE__      							=				"""a cat sitting on a park bench"""

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

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




   def      __a						(   self      :    Any				):



     """simple docstring"""






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

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

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

     SCREAMING_SNAKE_CASE__      							=				"""a cat sitting on a park bench"""

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

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