
import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


_UpperCAmelCase   :					Optional[Any]    =			[nn.LayerNorm]

_UpperCAmelCase   :					Optional[Any]    =			logging.get_logger(__name__)

_UpperCAmelCase   :					Any    =			version.parse(version.parse(torch.__version__).base_version)

_UpperCAmelCase   :					List[Any]    =			parsed_torch_version_base >= version.parse("2.0")
_UpperCAmelCase   :					Optional[int]    =			parsed_torch_version_base >= version.parse("1.12")
_UpperCAmelCase   :					int    =			parsed_torch_version_base >= version.parse("1.11")
_UpperCAmelCase   :					Dict    =			parsed_torch_version_base < version.parse("1.11")


def  A			(   lowercase       ,			lowercase       ,			lowercase       ,			lowercase       ,			self					)  ->    Optional[Any]:
		'''simple docstring'''

		from torch import _softmax_backward_data

		if is_torch_less_than_1_11:
				return _softmax_backward_data(_lowercase       ,			_lowercase       ,			parent.dim       ,			self					)
		else:
				return _softmax_backward_data(_lowercase       ,			_lowercase       ,			parent.dim       ,			self.dtype					)


def  A			(   lowercase       ,			lowercase       ,			lowercase = 0					)  ->    Optional[int]:
		'''simple docstring'''

		UpperCamelCase            =	index.to(layer.weight.device					)
		UpperCamelCase            =	layer.weight.index_select(_lowercase       ,			_lowercase					).clone().detach()
		if layer.bias is not None:
				if dim == 1:
						UpperCamelCase            =	layer.bias.clone().detach()
				else:
						UpperCamelCase            =	layer.bias[index].clone().detach()
		UpperCamelCase            =	list(layer.weight.size()					)
		UpperCamelCase            =	len(_lowercase					)
		UpperCamelCase            =	nn.Linear(new_size[1]       ,			new_size[0]       ,			bias=layer.bias is not None					).to(layer.weight.device					)
		UpperCamelCase            =	False
		new_layer.weight.copy_(W.contiguous()					)
		UpperCamelCase            =	True
		if layer.bias is not None:
				UpperCamelCase            =	False
				new_layer.bias.copy_(b.contiguous()					)
				UpperCamelCase            =	True
		return new_layer


class 							lowercase   (						nn.Module  ):



		def __init__(					self							, A_							, A_			)    ->   List[Any]:
				"""simple docstring"""
				super().__init__()
				UpperCamelCase            =	nf
				UpperCamelCase            =	nn.Parameter(torch.empty(__UpperCAmelCase							, __UpperCAmelCase			)			)
				UpperCamelCase            =	nn.Parameter(torch.zeros(__UpperCAmelCase			)			)
				nn.init.normal_(self.weight							, std=0.02			)



		def __UpperCamelCase					(					self							, A_			)    ->   Optional[int]:
				"""simple docstring"""
				UpperCamelCase            =	x.size()[:-1] + (self.nf,)
				UpperCamelCase            =	torch.addmm(self.bias							, x.view(-1							, x.size(-1			)			)							, self.weight			)
				UpperCamelCase            =	x.view(__UpperCAmelCase			)
				return x






def  A			(   lowercase       ,			lowercase       ,			lowercase = 1					)  ->    Union[str, Any]:
		'''simple docstring'''

		UpperCamelCase            =	index.to(layer.weight.device					)
		UpperCamelCase            =	layer.weight.index_select(_lowercase       ,			_lowercase					).clone().detach()
		if dim == 0:
				UpperCamelCase            =	layer.bias.clone().detach()
		else:
				UpperCamelCase            =	layer.bias[index].clone().detach()
		UpperCamelCase            =	list(layer.weight.size()					)
		UpperCamelCase            =	len(_lowercase					)
		UpperCamelCase            =	ConvaD(new_size[1]       ,			new_size[0]					).to(layer.weight.device					)
		UpperCamelCase            =	False
		new_layer.weight.copy_(W.contiguous()					)
		UpperCamelCase            =	True
		UpperCamelCase            =	False
		new_layer.bias.copy_(b.contiguous()					)
		UpperCamelCase            =	True
		return new_layer


def  A			(   lowercase       ,			lowercase       ,			lowercase = None					)  ->    str:
		'''simple docstring'''

		if isinstance(_lowercase       ,			nn.Linear					):
				return prune_linear_layer(_lowercase       ,			_lowercase       ,			dim=0 if dim is None else dim					)
		elif isinstance(_lowercase       ,			_lowercase					):
				return prune_convad_layer(_lowercase       ,			_lowercase       ,			dim=1 if dim is None else dim					)
		else:
				raise ValueError(f'''Can\'t prune layer of class {layer.__class__}'''					)


def  A			(   lowercase       ,			lowercase       ,			lowercase       ,			*lowercase					)  ->    Any:
		'''simple docstring'''

		assert len(_lowercase					) > 0, f'''{input_tensors} has to be a tuple/list of tensors'''

		# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
		UpperCamelCase            =	len(inspect.signature(_lowercase					).parameters					)
		if num_args_in_forward_chunk_fn != len(_lowercase					):
				raise ValueError(
				    f'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(_lowercase					)} input '''
				    'tensors are given'					)

		if chunk_size > 0:
				UpperCamelCase            =	input_tensors[0].shape[chunk_dim]
				for input_tensor in input_tensors:
						if input_tensor.shape[chunk_dim] != tensor_shape:
								raise ValueError(
								    f'''All input tenors have to be of the same shape: {tensor_shape}, '''
								    f'''found shape {input_tensor.shape[chunk_dim]}'''					)

				if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
						raise ValueError(
						    f'''The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk '''
						    f'''size {chunk_size}'''					)

				UpperCamelCase            =	input_tensors[0].shape[chunk_dim] // chunk_size

				# chunk input tensor into tuples
				UpperCamelCase            =	tuple(input_tensor.chunk(_lowercase       ,			dim=_lowercase					) for input_tensor in input_tensors					)
				# apply forward fn to every tuple
				UpperCamelCase            =	tuple(forward_fn(*_lowercase					) for input_tensors_chunk in zip(*_lowercase					)					)
				# concatenate output at same dimension
				return torch.cat(_lowercase       ,			dim=_lowercase					)

		return forward_fn(*_lowercase					)


def  A			(   lowercase       ,			lowercase       ,			lowercase       ,			lowercase					)  ->    Tuple:
		'''simple docstring'''

		UpperCamelCase            =	torch.ones(_lowercase       ,			_lowercase					)
		UpperCamelCase            =	set(_lowercase					) - already_pruned_heads  # Convert to set and remove already pruned heads
		for head in heads:
				# Compute how many pruned heads are before the head and move the index accordingly
				UpperCamelCase            =	head - sum(1 if h < head else 0 for h in already_pruned_heads					)
				UpperCamelCase            =	0
		UpperCamelCase            =	mask.view(-1					).contiguous().eq(1					)
		UpperCamelCase            =	torch.arange(len(_lowercase					)					)[mask].long()
		return heads, index


def  A			(   *lowercase       ,			lowercase = None					)  ->    Any:
		'''simple docstring'''

		return torch.meshgrid(*_lowercase       ,			indexing=_lowercase					)


def  A			(   lowercase					)  ->    Optional[int]:
		'''simple docstring'''

		return tensor.device, storage_ptr(_lowercase					), storage_size(_lowercase					)

"""simple docstring"""




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class        __SCREAMING_SNAKE_CASE     (   _lowerCAmelCase						,       _lowerCAmelCase						,       unittest.TestCase							):
				UpperCAmelCase	=      StableDiffusionParadigmsPipeline
				UpperCAmelCase	=      TEXT_TO_IMAGE_PARAMS
				UpperCAmelCase	=      TEXT_TO_IMAGE_BATCH_PARAMS
				UpperCAmelCase	=      TEXT_TO_IMAGE_IMAGE_PARAMS
				UpperCAmelCase	=      TEXT_TO_IMAGE_IMAGE_PARAMS







				def   lowercase_		(    self    :List[Any]  )      ->					Tuple:




									"""simple docstring"""



									torch.manual_seed(0  )
									lowerCamelCase__		:   int  				=					UNetaDConditionModel(
									    block_out_channels=(32, 64)      ,layers_per_block=2      ,sample_size=32      ,in_channels=4      ,out_channels=4      ,down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')      ,up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')      ,cross_attention_dim=32      ,attention_head_dim=(2, 4)      ,use_linear_projection=__UpperCAmelCase      ,)
									lowerCamelCase__		:   List[Any]  				=					DDIMParallelScheduler(
									    beta_start=0.00_085      ,beta_end=0.012      ,beta_schedule='''scaled_linear'''      ,clip_sample=__UpperCAmelCase      ,set_alpha_to_one=__UpperCAmelCase      ,)
									torch.manual_seed(0  )
									lowerCamelCase__		:   Dict  				=					AutoencoderKL(
									    block_out_channels=[32, 64]      ,in_channels=3      ,out_channels=3      ,down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']      ,up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']      ,latent_channels=4      ,sample_size=1_28      ,)
									torch.manual_seed(0  )
									lowerCamelCase__		:   Any  				=					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      ,)
									lowerCamelCase__		:   Tuple  				=					CLIPTextModel(__UpperCAmelCase  )
									lowerCamelCase__		:   int  				=					CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

									lowerCamelCase__		:   List[Any]  				=					{
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									    '''text_encoder''': text_encoder,
									    '''tokenizer''': tokenizer,
									    '''safety_checker''': None,
									    '''feature_extractor''': None,
									}
									return components







				def   lowercase_		(    self    :Tuple      ,__UpperCAmelCase    :int      ,__UpperCAmelCase    :Dict=0  )      ->					Optional[int]:




									"""simple docstring"""



									if str(__UpperCAmelCase  ).startswith('''mps'''  ):
														lowerCamelCase__		:   Dict  				=					torch.manual_seed(__UpperCAmelCase  )
									else:
														lowerCamelCase__		:   List[str]  				=					torch.Generator(device=__UpperCAmelCase  ).manual_seed(__UpperCAmelCase  )
									lowerCamelCase__		:   Optional[int]  				=					{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''generator''': generator,
									    '''num_inference_steps''': 10,
									    '''guidance_scale''': 6.0,
									    '''output_type''': '''numpy''',
									    '''parallel''': 3,
									    '''debug''': True,
									}
									return inputs







				def   lowercase_		(    self    :Optional[int]  )      ->					str:




									"""simple docstring"""



									lowerCamelCase__		:   Any  				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCamelCase__		:   int  				=					self.get_dummy_components()
									lowerCamelCase__		:   int  				=					StableDiffusionParadigmsPipeline(**__UpperCAmelCase  )
									lowerCamelCase__		:   Union[str, Any]  				=					sd_pipe.to(__UpperCAmelCase  )
									sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

									lowerCamelCase__		:   Dict  				=					self.get_dummy_inputs(__UpperCAmelCase  )
									lowerCamelCase__		:   Tuple  				=					sd_pipe(**__UpperCAmelCase  ).images
									lowerCamelCase__		:   List[Any]  				=					image[0, -3:, -3:, -1]
									assert image.shape == (1, 64, 64, 3)

									lowerCamelCase__		:   Tuple  				=					np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023]  )

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







				def   lowercase_		(    self    :int  )      ->					str:




									"""simple docstring"""



									lowerCamelCase__		:   Tuple  				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCamelCase__		:   Optional[int]  				=					self.get_dummy_components()
									torch.manual_seed(0  )
									lowerCamelCase__		:   Any  				=					DDPMParallelScheduler()
									torch.manual_seed(0  )
									lowerCamelCase__		:   str  				=					StableDiffusionParadigmsPipeline(**__UpperCAmelCase  )
									lowerCamelCase__		:   Tuple  				=					sd_pipe.to(__UpperCAmelCase  )
									sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

									lowerCamelCase__		:   int  				=					self.get_dummy_inputs(__UpperCAmelCase  )
									lowerCamelCase__		:   Optional[int]  				=					sd_pipe(**__UpperCAmelCase  ).images
									lowerCamelCase__		:   int  				=					image[0, -3:, -3:, -1]
									assert image.shape == (1, 64, 64, 3)

									lowerCamelCase__		:   Optional[int]  				=					np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468]  )

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







				def   lowercase_		(    self    :List[str]  )      ->					Dict:




									"""simple docstring"""



									super().test_inference_batch_consistent(batch_sizes=[1, 2]  )







				def   lowercase_		(    self    :Any  )      ->					str:




									"""simple docstring"""



									super().test_inference_batch_single_identical(batch_size=2      ,expected_max_diff=3E-3  )







				def   lowercase_		(    self    :int  )      ->					List[str]:




									"""simple docstring"""



									lowerCamelCase__		:   Optional[Any]  				=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowerCamelCase__		:   str  				=					self.get_dummy_components()
									lowerCamelCase__		:   int  				=					StableDiffusionParadigmsPipeline(**__UpperCAmelCase  )
									lowerCamelCase__		:   List[Any]  				=					sd_pipe.to(__UpperCAmelCase  )
									sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

									lowerCamelCase__		:   int  				=					self.get_dummy_inputs(__UpperCAmelCase  )
									lowerCamelCase__		:   int  				=					'''french fries'''
									lowerCamelCase__		:   Any  				=					sd_pipe(**__UpperCAmelCase      ,negative_prompt=__UpperCAmelCase  )
									lowerCamelCase__		:   Optional[Any]  				=					output.images
									lowerCamelCase__		:   str  				=					image[0, -3:, -3:, -1]

									assert image.shape == (1, 64, 64, 3)

									lowerCamelCase__		:   Tuple  				=					np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015]  )

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





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







				def   lowercase_		(    self    :Union[str, Any]  )      ->					Optional[int]:




									"""simple docstring"""



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







				def   lowercase_		(    self    :Optional[int]      ,__UpperCAmelCase    :List[str]=0  )      ->					str:




									"""simple docstring"""



									lowerCamelCase__		:   str  				=					torch.Generator(device=__UpperCAmelCase  ).manual_seed(__UpperCAmelCase  )
									lowerCamelCase__		:   List[Any]  				=					{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''generator''': generator,
									    '''num_inference_steps''': 10,
									    '''guidance_scale''': 7.5,
									    '''output_type''': '''numpy''',
									    '''parallel''': 3,
									    '''debug''': True,
									}
									return inputs







				def   lowercase_		(    self    :Optional[Any]  )      ->					Optional[int]:




									"""simple docstring"""



									lowerCamelCase__		:   Optional[Any]  				=					'''stabilityai/stable-diffusion-2-base'''
									lowerCamelCase__		:   Tuple  				=					DDIMParallelScheduler.from_pretrained(__UpperCAmelCase      ,subfolder='''scheduler'''  )
									lowerCamelCase__		:   Optional[Any]  				=					StableDiffusionParadigmsPipeline.from_pretrained(__UpperCAmelCase      ,scheduler=__UpperCAmelCase      ,safety_checker=__UpperCAmelCase  )
									pipe.to(__UpperCAmelCase  )
									pipe.set_progress_bar_config(disable=__UpperCAmelCase  )
									pipe.enable_attention_slicing()

									lowerCamelCase__		:   Optional[Any]  				=					self.get_inputs()
									lowerCamelCase__		:   Tuple  				=					pipe(**__UpperCAmelCase  ).images
									lowerCamelCase__		:   List[Any]  				=					image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 5_12, 5_12, 3)

									lowerCamelCase__		:   List[str]  				=					np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741]  )

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






