


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

						return shapes


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

						return tuple(reversed(SCREAMING_SNAKE_CASE_  )  )


@torch.jit.ignore
def     lowercase			(SCREAMING_SNAKE_CASE_		:				Sequence[int] ,			SCREAMING_SNAKE_CASE_		:				Sequence[int] ,			SCREAMING_SNAKE_CASE_		:				Sequence[int] ,			SCREAMING_SNAKE_CASE_		:				Optional[Sequence[bool]] = None ,			SCREAMING_SNAKE_CASE_		:				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(SCREAMING_SNAKE_CASE_		:				List[bool]  ) -> None:
												SCREAMING_SNAKE_CASE													=					True
												for i in range(len(SCREAMING_SNAKE_CASE_  )  ):
																		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(SCREAMING_SNAKE_CASE_  )
						if end_edges is None:
												SCREAMING_SNAKE_CASE													=					[e == (d - 1) for e, d in zip(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )]
												reduce_edge_list(SCREAMING_SNAKE_CASE_  )

						# Base cases. Either start/end are empty and we're done, or the final,
						# one-dimensional tensor can be simply sliced
						if len(SCREAMING_SNAKE_CASE_  ) == 0:
												return [()]
						elif len(SCREAMING_SNAKE_CASE_  ) == 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(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  ):
												if s == e:
																		path_list.append(slice(SCREAMING_SNAKE_CASE_ ,			s + 1  )  )
												else:
																		break

						SCREAMING_SNAKE_CASE													=					tuple(SCREAMING_SNAKE_CASE_  )
						SCREAMING_SNAKE_CASE													=					len(SCREAMING_SNAKE_CASE_  )

						# start == end, and we're done
						if divergence_idx == len(SCREAMING_SNAKE_CASE_  ):
												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(SCREAMING_SNAKE_CASE_ ,			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(SCREAMING_SNAKE_CASE_ ,			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     lowercase			(SCREAMING_SNAKE_CASE_		:				torch.Tensor ,			SCREAMING_SNAKE_CASE_		:				int ,			SCREAMING_SNAKE_CASE_		:				int ,			SCREAMING_SNAKE_CASE_		:				int  )   ->						torch.Tensor:
						SCREAMING_SNAKE_CASE													=					t.shape[:no_batch_dims]
						SCREAMING_SNAKE_CASE													=					list(_flat_idx_to_idx(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )  )
						# _get_minimal_slice_set is inclusive
						SCREAMING_SNAKE_CASE													=					list(_flat_idx_to_idx(flat_end - 1 ,			SCREAMING_SNAKE_CASE_  )  )

						# Get an ordered list of slices to perform
						SCREAMING_SNAKE_CASE													=					_get_minimal_slice_set(
						    SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			)

						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     lowercase			(SCREAMING_SNAKE_CASE_		:				Callable ,			SCREAMING_SNAKE_CASE_		:				Dict[str, Any] ,			SCREAMING_SNAKE_CASE_		:				int ,			SCREAMING_SNAKE_CASE_		:				int ,			SCREAMING_SNAKE_CASE_		:				bool = False ,			SCREAMING_SNAKE_CASE_		:				Any = None ,			SCREAMING_SNAKE_CASE_		:				bool = False ,			)   ->						Any:
						if not (len(SCREAMING_SNAKE_CASE_  ) > 0):
												raise ValueError('Must provide at least one input'  )

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

						def _prep_inputs(SCREAMING_SNAKE_CASE_		:				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 ,			SCREAMING_SNAKE_CASE_  )
						SCREAMING_SNAKE_CASE													=					None
						if _out is not None:
												SCREAMING_SNAKE_CASE													=					tensor_tree_map(lambda SCREAMING_SNAKE_CASE_  : 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(SCREAMING_SNAKE_CASE_		:				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(SCREAMING_SNAKE_CASE_  ):
												# Chunk the input
												if not low_mem:
																		SCREAMING_SNAKE_CASE													=					_select_chunk
												else:
																		SCREAMING_SNAKE_CASE													=					partial(
																		    _chunk_slice ,			flat_start=SCREAMING_SNAKE_CASE_ ,			flat_end=min(SCREAMING_SNAKE_CASE_ ,			i + chunk_size  ) ,			no_batch_dims=len(SCREAMING_SNAKE_CASE_  ) ,			)

												SCREAMING_SNAKE_CASE													=					tensor_tree_map(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )

												# Run the layer on the chunk
												SCREAMING_SNAKE_CASE													=					layer(**SCREAMING_SNAKE_CASE_  )

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

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

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

																		assign(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )
												elif isinstance(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  ):
																		for xa, xa in zip(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  ):
																								if _add_into_out:
																														xa[i : i + chunk_size] += xa
																								else:
																														SCREAMING_SNAKE_CASE													=					xa
												elif isinstance(SCREAMING_SNAKE_CASE_ ,			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 SCREAMING_SNAKE_CASE_  : t.view(orig_batch_dims + t.shape[1:]  ) ,			SCREAMING_SNAKE_CASE_  )

						return out


class 					lowerCAmelCase	:




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__ = 512						,       )     ->							Union[str, Any]:
								SCREAMING_SNAKE_CASE													=					max_chunk_size
								SCREAMING_SNAKE_CASE													=					None
								SCREAMING_SNAKE_CASE													=					None







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							int:
								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(lowerCAmelCase__				) -> bool:
														try:
																				with torch.no_grad():
																										fn(*lowerCAmelCase__						,       chunk_size=lowerCAmelCase__				)
																				return True
														except RuntimeError:
																				return False

								SCREAMING_SNAKE_CASE													=					0
								SCREAMING_SNAKE_CASE													=					len(lowerCAmelCase__				) - 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(lowerCAmelCase__				) - 1) // 2

								return candidates[min_viable_chunk_size_index]







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							bool:
								SCREAMING_SNAKE_CASE													=					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__				):
																				SCREAMING_SNAKE_CASE													=					[v for _, v in sorted(aa.items()						,       key=lambda lowerCAmelCase__				: x[0]				)]
																				SCREAMING_SNAKE_CASE													=					[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 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       )     ->							int:
								SCREAMING_SNAKE_CASE													=					True
								SCREAMING_SNAKE_CASE													=					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__				)
														SCREAMING_SNAKE_CASE													=					self._compare_arg_caches(self.cached_arg_data						,       lowerCAmelCase__				)
								else:
														# Otherwise, we can reuse the precomputed value
														SCREAMING_SNAKE_CASE													=					False

								if not consistent:
														SCREAMING_SNAKE_CASE													=					self._determine_favorable_chunk_size(
														    lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       )
														SCREAMING_SNAKE_CASE													=					arg_data

								assert self.cached_chunk_size is not None

								return self.cached_chunk_size






"""simple docstring"""




import unittest

import numpy as np

from diffusers import OnnxStableDiffusionInpaintPipelineLegacy
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    load_numpy,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)


if is_onnx_available():
	import onnxruntime as ort
@nightly
@require_onnxruntime
@require_torch_gpu
class 					lowerCAmelCase	(					unittest.TestCase    ):




		'''simple docstring'''







		@property
		def 							__A						(	self				)     ->							List[Any]:
								return (
								    "CUDAExecutionProvider",
								    {
								        "gpu_mem_limit": "15000000000",  # 15GB
								        "arena_extend_strategy": "kSameAsRequested",
								    },
								)







		@property
		def 							__A						(	self				)     ->							List[str]:
								SCREAMING_SNAKE_CASE													=					ort.SessionOptions()
								SCREAMING_SNAKE_CASE													=					False
								return options







		def 							__A						(	self				)     ->							Dict:
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo.png'				)
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo_mask.png'				)
								SCREAMING_SNAKE_CASE													=					load_numpy(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/red_cat_sitting_on_a_park_bench_onnx.npy'				)

								# using the PNDM scheduler by default
								SCREAMING_SNAKE_CASE													=					OnnxStableDiffusionInpaintPipelineLegacy.from_pretrained(
								    'CompVis/stable-diffusion-v1-4'						,       revision='onnx'						,       safety_checker=lowerCAmelCase__						,       feature_extractor=lowerCAmelCase__						,       provider=self.gpu_provider						,       sess_options=self.gpu_options						,       )
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					'A red cat sitting on a park bench'

								SCREAMING_SNAKE_CASE													=					np.random.RandomState(0				)
								SCREAMING_SNAKE_CASE													=					pipe(
								    prompt=lowerCAmelCase__						,       image=lowerCAmelCase__						,       mask_image=lowerCAmelCase__						,       strength=0.75						,       guidance_scale=7.5						,       num_inference_steps=15						,       generator=lowerCAmelCase__						,       output_type='np'						,       )

								SCREAMING_SNAKE_CASE													=					output.images[0]

								assert image.shape == (512, 512, 3)
								assert np.abs(expected_image - image				).max() < 1e-2


