'''simple docstring'''





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       __lowerCAmelCase      (						A_   : Union[dict, list, tuple, torch.Tensor]					)    ->  List[Tuple[int, ...]]:
  __UpperCAmelCase						    =       []
  if isinstance(A_						,	A_					):
    for v in tree.values():
      shapes.extend(_fetch_dims(A_					)					)
  elif isinstance(A_						,	(list, tuple)					):
    for t in tree:
      shapes.extend(_fetch_dims(A_					)					)
  elif isinstance(A_						,	torch.Tensor					):
    shapes.append(tree.shape					)
  else:
    raise ValueError("Not supported"					)

  return shapes




@torch.jit.ignore
def       __lowerCAmelCase      (						A_   : int						,	A_   : Tuple[int, ...]					)    ->  Tuple[int, ...]:
  __UpperCAmelCase						    =       []
  for d in reversed(A_					):
    idx.append(flat_idx % d					)
    __UpperCAmelCase						    =       flat_idx // d

  return tuple(reversed(A_					)					)




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

  # start_edges and end_edges both indicate whether, starting from any given
  # dimension, the start/end index is at the top/bottom edge of the
  # corresponding tensor, modeled as a tree
  def reduce_edge_list(A_   : List[bool]					) -> None:
    __UpperCAmelCase						    =       True
    for i in range(len(A_					)					):
      __UpperCAmelCase						    =       -1 * (i + 1)
      l[reversed_idx] &= tally
      __UpperCAmelCase						    =       l[reversed_idx]

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

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

  __UpperCAmelCase						    =       []
  __UpperCAmelCase						    =       []

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

  __UpperCAmelCase						    =       tuple(A_					)
  __UpperCAmelCase						    =       len(A_					)

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

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

    __UpperCAmelCase						    =       start[divergence_idx]
    return tuple(
        path + (slice(A_						,	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

    __UpperCAmelCase						    =       end[divergence_idx]
    return tuple(
        path + (slice(A_						,	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()					)
    __UpperCAmelCase						    =       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       __lowerCAmelCase      (						A_   : torch.Tensor						,	A_   : int						,	A_   : int						,	A_   : int					)    ->  torch.Tensor:
  __UpperCAmelCase						    =       t.shape[:no_batch_dims]
  __UpperCAmelCase						    =       list(_flat_idx_to_idx(A_						,	A_					)					)
  # _get_minimal_slice_set is inclusive
  __UpperCAmelCase						    =       list(_flat_idx_to_idx(flat_end - 1						,	A_					)					)

  # Get an ordered list of slices to perform
  __UpperCAmelCase						    =       _get_minimal_slice_set(
      A_						,	A_						,	A_						,	)

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

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




def       __lowerCAmelCase      (						A_   : Callable						,	A_   : Dict[str, Any]						,	A_   : int						,	A_   : int						,	A_   : bool = False						,	A_   : Any = None						,	A_   : bool = False						,	)    ->  Any:
  if not (len(A_					) > 0):
    raise ValueError("Must provide at least one input"					)

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

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

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

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

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

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

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

    __UpperCAmelCase						    =       tensor_tree_map(A_						,	A_					)

    # Run the layer on the chunk
    __UpperCAmelCase						    =       layer(**A_					)

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

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

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

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

    i += chunk_size

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

  return out





class 			UpperCAmelCase__    :



  """simple docstring"""






  def __init__(				self:			Tuple				,				__lowerCAmelCase:			int = 512				,				) -> Any:




    '''simple docstring'''




    __UpperCAmelCase						    =       max_chunk_size
    __UpperCAmelCase						    =       None
    __UpperCAmelCase						    =       None






  def 	_UpperCAmelCase				(				self:			Optional[int]				,				__lowerCAmelCase:			Callable				,				__lowerCAmelCase:			tuple				,				__lowerCAmelCase:			int   ) -> int:




    '''simple docstring'''




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

    if min_chunk_size >= self.max_chunk_size:
      return min_chunk_size

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

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

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

    return candidates[min_viable_chunk_size_index]






  def 	_UpperCAmelCase				(				self:			Optional[Any]				,				__lowerCAmelCase:			Iterable				,				__lowerCAmelCase:			Iterable   ) -> bool:




    '''simple docstring'''




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

    return consistent






  def 	_UpperCAmelCase				(				self:			List[Any]				,				__lowerCAmelCase:			Callable				,				__lowerCAmelCase:			tuple				,				__lowerCAmelCase:			int				,				) -> int:




    '''simple docstring'''




    __UpperCAmelCase						    =       True
    __UpperCAmelCase						    =       tree_map(lambda __lowerCAmelCase   : a.shape if isinstance(__lowerCAmelCase				,				torch.Tensor   ) else a				,				__lowerCAmelCase				,				__lowerCAmelCase   )
    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(__lowerCAmelCase   )
      __UpperCAmelCase						    =       self._compare_arg_caches(self.cached_arg_data				,				__lowerCAmelCase   )
    else:
      # Otherwise, we can reuse the precomputed value
      __UpperCAmelCase						    =       False

    if not consistent:
      __UpperCAmelCase						    =       self._determine_favorable_chunk_size(
          __lowerCAmelCase				,				__lowerCAmelCase				,				__lowerCAmelCase				,				)
      __UpperCAmelCase						    =       arg_data

    assert self.cached_chunk_size is not None

    return self.cached_chunk_size




import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 			UpperCAmelCase__    (		snake_case		,	unittest.TestCase					):



  """simple docstring"""






  lowerCAmelCase__		:	int    		= DiTPipeline
  lowerCAmelCase__		:	Any    		= CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
  lowerCAmelCase__		:	Optional[Any]    		= PipelineTesterMixin.required_optional_params - {
      'latents',
      'num_images_per_prompt',
      'callback',
      'callback_steps',
  }
  lowerCAmelCase__		:	List[Any]    		= CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
  lowerCAmelCase__		:	Optional[int]    		= False






  def 	_UpperCAmelCase				(				self:			str   ) -> List[Any]:




    '''simple docstring'''




    torch.manual_seed(0   )
    __UpperCAmelCase						    =       TransformeraDModel(
        sample_size=16				,				num_layers=2				,				patch_size=4				,				attention_head_dim=8				,				num_attention_heads=2				,				in_channels=4				,				out_channels=8				,				attention_bias=__lowerCAmelCase				,				activation_fn="gelu-approximate"				,				num_embeds_ada_norm=1_000				,				norm_type="ada_norm_zero"				,				norm_elementwise_affine=__lowerCAmelCase				,				)
    __UpperCAmelCase						    =       AutoencoderKL()
    __UpperCAmelCase						    =       DDIMScheduler()
    __UpperCAmelCase						    =       {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
    return components






  def 	_UpperCAmelCase				(				self:			str				,				__lowerCAmelCase:			str				,				__lowerCAmelCase:			Optional[int]=0   ) -> int:




    '''simple docstring'''




    if str(__lowerCAmelCase   ).startswith("mps"   ):
      __UpperCAmelCase						    =       torch.manual_seed(__lowerCAmelCase   )
    else:
      __UpperCAmelCase						    =       torch.Generator(device=__lowerCAmelCase   ).manual_seed(__lowerCAmelCase   )
    __UpperCAmelCase						    =       {
        "class_labels": [1],
        "generator": generator,
        "num_inference_steps": 2,
        "output_type": "numpy",
    }
    return inputs






  def 	_UpperCAmelCase				(				self:			Tuple   ) -> Dict:




    '''simple docstring'''




    __UpperCAmelCase						    =       "cpu"

    __UpperCAmelCase						    =       self.get_dummy_components()
    __UpperCAmelCase						    =       self.pipeline_class(**__lowerCAmelCase   )
    pipe.to(__lowerCAmelCase   )
    pipe.set_progress_bar_config(disable=__lowerCAmelCase   )

    __UpperCAmelCase						    =       self.get_dummy_inputs(__lowerCAmelCase   )
    __UpperCAmelCase						    =       pipe(**__lowerCAmelCase   ).images
    __UpperCAmelCase						    =       image[0, -3:, -3:, -1]

    self.assertEqual(image.shape				,				(1, 16, 16, 3)   )
    __UpperCAmelCase						    =       np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]   )
    __UpperCAmelCase						    =       np.abs(image_slice.flatten() - expected_slice   ).max()
    self.assertLessEqual(__lowerCAmelCase				,				1E-3   )






  def 	_UpperCAmelCase				(				self:			int   ) -> str:




    '''simple docstring'''




    self._test_inference_batch_single_identical(relax_max_difference=__lowerCAmelCase				,				expected_max_diff=1E-3   )






  @unittest.skipIf(
      torch_device != "cuda" or not is_xformers_available()				,				reason="XFormers attention is only available with CUDA and `xformers` installed"				,				)
  def 	_UpperCAmelCase				(				self:			Optional[int]   ) -> str:




    '''simple docstring'''




    self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3   )






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



  """simple docstring"""






  def 	_UpperCAmelCase				(				self:			Any   ) -> Optional[Any]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Dict   ) -> List[Any]:




    '''simple docstring'''




    __UpperCAmelCase						    =       torch.manual_seed(0   )

    __UpperCAmelCase						    =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"   )
    pipe.to("cuda"   )

    __UpperCAmelCase						    =       ["vase", "umbrella", "white shark", "white wolf"]
    __UpperCAmelCase						    =       pipe.get_label_ids(__lowerCAmelCase   )

    __UpperCAmelCase						    =       pipe(__lowerCAmelCase				,				generator=__lowerCAmelCase				,				num_inference_steps=40				,				output_type="np"   ).images

    for word, image in zip(__lowerCAmelCase				,				__lowerCAmelCase   ):
      __UpperCAmelCase						    =       load_numpy(
          F'''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'''   )
      assert np.abs((expected_image - image).max()   ) < 1E-2






  def 	_UpperCAmelCase				(				self:			Any   ) -> List[str]:




    '''simple docstring'''




    __UpperCAmelCase						    =       DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"   )
    __UpperCAmelCase						    =       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config   )
    pipe.to("cuda"   )

    __UpperCAmelCase						    =       ["vase", "umbrella"]
    __UpperCAmelCase						    =       pipe.get_label_ids(__lowerCAmelCase   )

    __UpperCAmelCase						    =       torch.manual_seed(0   )
    __UpperCAmelCase						    =       pipe(__lowerCAmelCase				,				generator=__lowerCAmelCase				,				num_inference_steps=25				,				output_type="np"   ).images

    for word, image in zip(__lowerCAmelCase				,				__lowerCAmelCase   ):
      __UpperCAmelCase						    =       load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          F'''/dit/{word}_512.npy'''   )

      assert np.abs((expected_image - image).max()   ) < 1E-1



