





'''simple docstring'''







from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


lowerCAmelCase_					:    str      	=			logging.get_logger(__name__)

lowerCAmelCase_					:    Optional[int]      	=			{
    'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json',
    # See all Marian models at https://huggingface.co/models?filter=marian
}
class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_      ):






				"""simple docstring"""







				__a            ='marian'
				__a            =['past_key_values']
				__a            ={'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}




				def __init__(  self     :					List[str]					,				__a     :					List[str]=5_81_01					,				__a     :					str=None					,				__a     :					int=10_24					,				__a     :					Optional[Any]=12					,				__a     :					int=40_96					,				__a     :					List[str]=16					,				__a     :					Optional[int]=12					,				__a     :					str=40_96					,				__a     :					Union[str, Any]=16					,				__a     :					List[str]=0.0					,				__a     :					str=0.0					,				__a     :					Optional[int]=True					,				__a     :					Optional[int]=True					,				__a     :					List[str]="gelu"					,				__a     :					Dict=10_24					,				__a     :					Optional[Any]=0.1					,				__a     :					Union[str, Any]=0.0					,				__a     :					Tuple=0.0					,				__a     :					Any=0.02					,				__a     :					str=5_81_00					,				__a     :					int=False					,				__a     :					int=5_81_00					,				__a     :					int=0					,				__a     :					Dict=0					,				__a     :					List[str]=True					,				**__a     :					Union[str, Any]					,				):
								_a								= vocab_size
								_a								= decoder_vocab_size or vocab_size
								_a								= max_position_embeddings
								_a								= d_model
								_a								= encoder_ffn_dim
								_a								= encoder_layers
								_a								= encoder_attention_heads
								_a								= decoder_ffn_dim
								_a								= decoder_layers
								_a								= decoder_attention_heads
								_a								= dropout
								_a								= attention_dropout
								_a								= activation_dropout
								_a								= activation_function
								_a								= init_std
								_a								= encoder_layerdrop
								_a								= decoder_layerdrop
								_a								= use_cache
								_a								= encoder_layers
								_a								= scale_embedding  # scale factor will be sqrt(d_model) if True
								_a								= share_encoder_decoder_embeddings
								super().__init__(
								    pad_token_id=__a					,				eos_token_id=__a					,				is_encoder_decoder=__a					,				decoder_start_token_id=__a					,				forced_eos_token_id=__a					,				**__a					,				)


class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_      ):






				"""simple docstring"""







				@property
				# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
				def   UpperCamelCase__   (  self     :					str					):
								if self.task in ["default", "seq2seq-lm"]:
												_a								= OrderedDict(
												    [
												        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
												        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
												    ]					)

												if self.use_past:
																_a								= {0: "batch"}
																_a								= {0: "batch", 1: "past_decoder_sequence + sequence"}
												else:
																_a								= {0: "batch", 1: "decoder_sequence"}
																_a								= {0: "batch", 1: "decoder_sequence"}

												if self.use_past:
																self.fill_with_past_key_values_(__a					,				direction="inputs"					)
								elif self.task == "causal-lm":
												# TODO: figure this case out.
												_a								= OrderedDict(
												    [
												        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
												        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
												    ]					)
												if self.use_past:
																_a	,					_a								= self.num_layers
																for i in range(__a					):
																				_a								= {0: "batch", 2: "past_sequence + sequence"}
																				_a								= {0: "batch", 2: "past_sequence + sequence"}
								else:
												_a								= OrderedDict(
												    [
												        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
												        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
												        ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
												        ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
												    ]					)

								return common_inputs




				@property
				# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
				def   UpperCamelCase__   (  self     :					List[str]					):
								if self.task in ["default", "seq2seq-lm"]:
												_a								= super().outputs
								else:
												_a								= super(__a					,				self					).outputs
												if self.use_past:
																_a	,					_a								= self.num_layers
																for i in range(__a					):
																				_a								= {0: "batch", 2: "past_sequence + sequence"}
																				_a								= {0: "batch", 2: "past_sequence + sequence"}
								return common_outputs




				def   UpperCamelCase__   (  self     :					Tuple					,				__a     :					PreTrainedTokenizer					,				__a     :					int = -1					,				__a     :					int = -1					,				__a     :					bool = False					,				__a     :					Optional[TensorType] = None					,				):
								_a								= self._generate_dummy_inputs_for_encoder_and_decoder(
								    __a					,				__a					,				__a					,				__a					,				__a					)

								# Generate decoder inputs
								_a								= seq_length if not self.use_past else 1
								_a								= self._generate_dummy_inputs_for_encoder_and_decoder(
								    __a					,				__a					,				__a					,				__a					,				__a					)
								_a								= {f'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
								_a								= dict(**__a					,				**__a					)

								if self.use_past:
												if not is_torch_available():
																raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."					)
												else:
																import torch
												_a	,					_a								= common_inputs["input_ids"].shape
												_a								= common_inputs["decoder_input_ids"].shape[1]
												_a	,					_a								= self.num_attention_heads
												_a								= (
												    batch,
												    num_encoder_attention_heads,
												    encoder_seq_length,
												    self._config.hidden_size // num_encoder_attention_heads,
												)
												_a								= decoder_seq_length + 3
												_a								= (
												    batch,
												    num_decoder_attention_heads,
												    decoder_past_length,
												    self._config.hidden_size // num_decoder_attention_heads,
												)

												_a								= torch.cat(
												    [common_inputs["decoder_attention_mask"], torch.ones(__a					,				__a					)]					,				dim=1					)

												_a								= []
												# If the number of encoder and decoder layers are present in the model configuration, both are considered
												_a	,					_a								= self.num_layers
												_a								= min(__a					,				__a					)
												_a								= max(__a					,				__a					) - min_num_layers
												_a								= "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

												for _ in range(__a					):
																common_inputs["past_key_values"].append(
																    (
																        torch.zeros(__a					),
																        torch.zeros(__a					),
																        torch.zeros(__a					),
																        torch.zeros(__a					),
																    )					)
												# TODO: test this.
												_a								= encoder_shape if remaining_side_name == "encoder" else decoder_shape
												for _ in range(__a					,				__a					):
																common_inputs["past_key_values"].append((torch.zeros(__a					), torch.zeros(__a					))					)
								return common_inputs




				def   UpperCamelCase__   (  self     :					Union[str, Any]					,				__a     :					PreTrainedTokenizer					,				__a     :					int = -1					,				__a     :					int = -1					,				__a     :					bool = False					,				__a     :					Optional[TensorType] = None					,				):
								_a								= self._generate_dummy_inputs_for_encoder_and_decoder(
								    __a					,				__a					,				__a					,				__a					,				__a					)

								if self.use_past:
												if not is_torch_available():
																raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."					)
												else:
																import torch
												_a	,					_a								= common_inputs["input_ids"].shape
												# Not using the same length for past_key_values
												_a								= seqlen + 2
												_a	,					_a								= self.num_layers
												_a	,					_a								= self.num_attention_heads
												_a								= (
												    batch,
												    num_encoder_attention_heads,
												    past_key_values_length,
												    self._config.hidden_size // num_encoder_attention_heads,
												)

												_a								= common_inputs["attention_mask"].dtype
												_a								= torch.cat(
												    [common_inputs["attention_mask"], torch.ones(__a					,				__a					,				dtype=__a					)]					,				dim=1					)
												_a								= [
												    (torch.zeros(__a					), torch.zeros(__a					)) for _ in range(__a					)
												]
								return common_inputs




				def   UpperCamelCase__   (  self     :					str					,				__a     :					PreTrainedTokenizer					,				__a     :					int = -1					,				__a     :					int = -1					,				__a     :					bool = False					,				__a     :					Optional[TensorType] = None					,				):
								# Copied from OnnxConfig.generate_dummy_inputs
								# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
								# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
								_a								= compute_effective_axis_dimension(
								    __a					,				fixed_dimension=OnnxConfig.default_fixed_batch					,				num_token_to_add=0					)

								# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
								_a								= tokenizer.num_special_tokens_to_add(__a					)
								_a								= compute_effective_axis_dimension(
								    __a					,				fixed_dimension=OnnxConfig.default_fixed_sequence					,				num_token_to_add=__a					)

								# Generate dummy inputs according to compute batch and sequence
								_a								= [" ".join([tokenizer.unk_token]					) * seq_length] * batch_size
								_a								= dict(tokenizer(__a					,				return_tensors=__a					)					)
								return common_inputs




				def   UpperCamelCase__   (  self     :					Optional[int]					,				__a     :					PreTrainedTokenizer					,				__a     :					int = -1					,				__a     :					int = -1					,				__a     :					bool = False					,				__a     :					Optional[TensorType] = None					,				):
								if self.task in ["default", "seq2seq-lm"]:
												_a								= self._generate_dummy_inputs_for_default_and_seqaseq_lm(
												    __a					,				batch_size=__a					,				seq_length=__a					,				is_pair=__a					,				framework=__a					)

								else:
												_a								= self._generate_dummy_inputs_for_causal_lm(
												    __a					,				batch_size=__a					,				seq_length=__a					,				is_pair=__a					,				framework=__a					)

								return common_inputs




				def   UpperCamelCase__   (  self     :					Dict					,				__a     :					Dict					,				__a     :					Optional[int]					,				__a     :					str					,				__a     :					Tuple					):
								if self.task in ["default", "seq2seq-lm"]:
												_a								= super()._flatten_past_key_values_(__a					,				__a					,				__a					,				__a					)
								else:
												_a								= super(__a					,				self					)._flatten_past_key_values_(
												    __a					,				__a					,				__a					,				__a					)




				@property
				def   UpperCamelCase__   (  self     :					Optional[Any]					):
								return 1e-4

import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class 	lowerCamelCase_     (   UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_						,	unittest.TestCase ):






    '''simple docstring'''






    a__	:						int						=      StableUnCLIPImgaImgPipeline
    a__	:						Optional[int]						=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS
    a__	:						Union[str, Any]						=      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    a__	:						Optional[Any]						=      frozenset(
        [] )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    a__	:						int						=      frozenset([] )







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :Tuple			  =		32
        __UpperCamelCase     :Optional[int]			  =		embedder_hidden_size

        # image encoding components

        __UpperCamelCase     :Union[str, Any]			  =		CLIPImageProcessor(crop_size=32 , size=32)

        torch.manual_seed(0)
        __UpperCamelCase     :Union[str, Any]			  =		CLIPVisionModelWithProjection(
            CLIPVisionConfig(
                hidden_size=__lowercase , projection_dim=__lowercase , num_hidden_layers=5 , num_attention_heads=4 , image_size=32 , intermediate_size=37 , patch_size=1 , ))

        # regular denoising components

        torch.manual_seed(0)
        __UpperCamelCase     :str			  =		StableUnCLIPImageNormalizer(embedding_dim=__lowercase)
        __UpperCamelCase     :Optional[int]			  =		DDPMScheduler(beta_schedule='''squaredcos_cap_v2''')

        torch.manual_seed(0)
        __UpperCamelCase     :Union[str, Any]			  =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

        torch.manual_seed(0)
        __UpperCamelCase     :Dict			  =		CLIPTextModel(
            CLIPTextConfig(
                bos_token_id=0 , eos_token_id=2 , hidden_size=__lowercase , projection_dim=32 , intermediate_size=37 , layer_norm_eps=1E-0_5 , num_attention_heads=4 , num_hidden_layers=5 , pad_token_id=1 , vocab_size=1_000 , ))

        torch.manual_seed(0)
        __UpperCamelCase     :List[Any]			  =		UNetaDConditionModel(
            sample_size=32 , in_channels=4 , out_channels=4 , down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''') , up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''') , block_out_channels=(32, 64) , attention_head_dim=(2, 4) , class_embed_type='''projection''' , projection_class_embeddings_input_dim=embedder_projection_dim * 2 , cross_attention_dim=__lowercase , layers_per_block=1 , upcast_attention=__lowercase , use_linear_projection=__lowercase , )

        torch.manual_seed(0)
        __UpperCamelCase     :Tuple			  =		DDIMScheduler(
            beta_schedule='''scaled_linear''' , beta_start=0.0_00_85 , beta_end=0.0_12 , prediction_type='''v_prediction''' , set_alpha_to_one=__lowercase , steps_offset=1 , )

        torch.manual_seed(0)
        __UpperCamelCase     :List[str]			  =		AutoencoderKL()

        __UpperCamelCase     :Tuple			  =		{
            # image encoding components
            '''feature_extractor''': feature_extractor,
            '''image_encoder''': image_encoder.eval(),
            # image noising components
            '''image_normalizer''': image_normalizer.eval(),
            '''image_noising_scheduler''': image_noising_scheduler,
            # regular denoising components
            '''tokenizer''': tokenizer,
            '''text_encoder''': text_encoder.eval(),
            '''unet''': unet.eval(),
            '''scheduler''': scheduler,
            '''vae''': vae.eval(),
        }

        return components







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase=0 , __lowercase=True)       ->	str:
        if str(__lowercase).startswith('''mps'''):
            __UpperCamelCase     :Union[str, Any]			  =		torch.manual_seed(__lowercase)
        else:
            __UpperCamelCase     :int			  =		torch.Generator(device=__lowercase).manual_seed(__lowercase)

        __UpperCamelCase     :int			  =		floats_tensor((1, 3, 32, 32) , rng=random.Random(__lowercase)).to(__lowercase)

        if pil_image:
            __UpperCamelCase     :List[Any]			  =		input_image * 0.5 + 0.5
            __UpperCamelCase     :Optional[Any]			  =		input_image.clamp(0 , 1)
            __UpperCamelCase     :int			  =		input_image.cpu().permute(0 , 2 , 3 , 1).float().numpy()
            __UpperCamelCase     :Optional[Any]			  =		DiffusionPipeline.numpy_to_pil(__lowercase)[0]

        return {
            "prompt": "An anime racoon running a marathon",
            "image": input_image,
            "generator": generator,
            "num_inference_steps": 2,
            "output_type": "np",
        }







    @skip_mps
    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :Dict			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Tuple			  =		self.get_dummy_components()
        __UpperCamelCase     :Any			  =		StableUnCLIPImgaImgPipeline(**__lowercase)
        __UpperCamelCase     :Optional[Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :List[Any]			  =		self.get_dummy_inputs(__lowercase)
        inputs.update({'''image_embeds''': None})
        __UpperCamelCase     :Any			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :List[str]			  =		image[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :List[Any]			  =		np.array([0.38_72, 0.72_24, 0.56_01, 0.47_41, 0.68_72, 0.58_14, 0.46_36, 0.38_67, 0.50_78])

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







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Optional[Any]			  =		torch_device in ['''cpu''', '''mps''']

        self._test_attention_slicing_forward_pass(test_max_difference=__lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        __UpperCamelCase     :Optional[Any]			  =		torch_device in ['''cpu''', '''mps''']

        self._test_inference_batch_single_identical(test_max_difference=__lowercase)







    @unittest.skipIf(
        torch_device != '''cuda''' or not is_xformers_available() , reason='''XFormers attention is only available with CUDA and `xformers` installed''' , )
    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        self._test_xformers_attention_forwardGenerator_pass(test_max_difference=__lowercase)


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






    '''simple docstring'''






    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :int			  =		load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png''')

        __UpperCamelCase     :Any			  =		load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy''')

        __UpperCamelCase     :List[Any]			  =		StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-l-img2img''' , torch_dtype=torch.floataa)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :int			  =		torch.Generator(device='''cpu''').manual_seed(0)
        __UpperCamelCase     :Dict			  =		pipe(__lowercase , '''anime turle''' , generator=__lowercase , output_type='''np''')

        __UpperCamelCase     :Dict			  =		output.images[0]

        assert image.shape == (768, 768, 3)

        assert_mean_pixel_difference(__lowercase , __lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :Optional[Any]			  =		load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png''')

        __UpperCamelCase     :Any			  =		load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy''')

        __UpperCamelCase     :Any			  =		StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-h-img2img''' , torch_dtype=torch.floataa)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :int			  =		torch.Generator(device='''cpu''').manual_seed(0)
        __UpperCamelCase     :Optional[int]			  =		pipe(__lowercase , '''anime turle''' , generator=__lowercase , output_type='''np''')

        __UpperCamelCase     :List[Any]			  =		output.images[0]

        assert image.shape == (768, 768, 3)

        assert_mean_pixel_difference(__lowercase , __lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :Dict			  =		load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png''')

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

        __UpperCamelCase     :List[Any]			  =		StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-h-img2img''' , torch_dtype=torch.floataa)
        __UpperCamelCase     :Union[str, Any]			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :Optional[Any]			  =		pipe(
            __lowercase , '''anime turtle''' , num_inference_steps=2 , output_type='''np''' , )

        __UpperCamelCase     :int			  =		torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 10**9





