
'''simple docstring'''






from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class     __snake_case		(   a__    ,   a__    ,   a__):
					_lowerCAmelCase			     =	[R'''h\.\d+\.attn\.bias''', R'''h\.\d+\.attn\.masked_bias''']
					@register_to_config
					def __init__(     self,    A,    A,    A = None,    A = 5_0257,    A = 1024,    A = 768,    A = 12,    A = 12,    A = None,    A = "gelu_new",    A = 0.1,    A = 0.1,    A = 0.1,    A = 1e-5,    A = 0.02,    A = True,    A = True,    A = False,    A = False,    ):

									"""simple docstring"""


									super().__init__()

									lowerCamelCase					:			str									= prefix_length

									if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
													raise ValueError(
													    F'''`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'''
													    F''' `n_embd`: {n_embd} are not equal.'''      )

									lowerCamelCase					:			str									= prefix_inner_dim
									lowerCamelCase					:			int									= prefix_hidden_dim

									lowerCamelCase					:			Optional[int]									= (
									    nn.Linear(self.prefix_inner_dim,    self.prefix_hidden_dim      )
									    if self.prefix_hidden_dim is not None
									    else nn.Identity()
									)
									lowerCamelCase					:			int									= (
									    nn.Linear(self.prefix_hidden_dim,    A      ) if self.prefix_hidden_dim is not None else nn.Identity()
									)

									lowerCamelCase					:			int									= GPTaConfig(
									    vocab_size=A,    n_positions=A,    n_embd=A,    n_layer=A,    n_head=A,    n_inner=A,    activation_function=A,    resid_pdrop=A,    embd_pdrop=A,    attn_pdrop=A,    layer_norm_epsilon=A,    initializer_range=A,    scale_attn_weights=A,    use_cache=A,    scale_attn_by_inverse_layer_idx=A,    reorder_and_upcast_attn=A,    )
									lowerCamelCase					:			Optional[int]									= GPTaLMHeadModel(A      )
					def     UpperCAmelCase_							(     self,    A,    A,    A = None,    A = None,    ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= self.transformer.transformer.wte(A      )
									lowerCamelCase					:			Union[str, Any]									= self.encode_prefix(A      )
									lowerCamelCase					:			List[Any]									= self.decode_prefix(A      )
									lowerCamelCase					:			int									= torch.cat((prefix_embeds, embedding_text),    dim=1      )

									if labels is not None:
													lowerCamelCase					:			Optional[Any]									= self.get_dummy_token(input_ids.shape[0],    input_ids.device      )
													lowerCamelCase					:			Optional[Any]									= torch.cat((dummy_token, input_ids),    dim=1      )
									lowerCamelCase					:			List[str]									= self.transformer(inputs_embeds=A,    labels=A,    attention_mask=A      )
									if self.prefix_hidden_dim is not None:
													return out, hidden
									else:
													return out
					def     UpperCAmelCase_							(     self,    A,    A      ):

									"""simple docstring"""


									return torch.zeros(A,    self.prefix_length,    dtype=torch.intaa,    device=A      )
					def     UpperCAmelCase_							(     self,    A      ):

									"""simple docstring"""


									return self.encode_prefix(A      )
					@torch.no_grad()
					def     UpperCAmelCase_							(     self,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= torch.split(A,    1,    dim=0      )
									lowerCamelCase					:			List[str]									= []
									lowerCamelCase					:			Any									= []
									for feature in features:
													lowerCamelCase					:			Tuple									= self.decode_prefix(feature.to(A      )      )  # back to the clip feature
													# Only support beam search for now
													lowerCamelCase					:			Union[str, Any]									= self.generate_beam(
													    input_embeds=A,    device=A,    eos_token_id=A      )
													generated_tokens.append(output_tokens[0]      )
													generated_seq_lengths.append(seq_lengths[0]      )
									lowerCamelCase					:			str									= torch.stack(A      )
									lowerCamelCase					:			Any									= torch.stack(A      )
									return generated_tokens, generated_seq_lengths




					@torch.no_grad()
					def     UpperCAmelCase_							(     self,    A=None,    A=None,    A=None,    A = 5,    A = 67,    A = 1.0,    A = None,    ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= eos_token_id
									lowerCamelCase					:			int									= None
									lowerCamelCase					:			List[Any]									= None
									lowerCamelCase					:			Union[str, Any]									= torch.ones(A,    device=A,    dtype=torch.int      )
									lowerCamelCase					:			Dict									= torch.zeros(A,    device=A,    dtype=torch.bool      )

									if input_embeds is not None:
													lowerCamelCase					:			Union[str, Any]									= input_embeds
									else:
													lowerCamelCase					:			Union[str, Any]									= self.transformer.transformer.wte(A      )

									for i in range(A      ):
													lowerCamelCase					:			List[str]									= self.transformer(inputs_embeds=A      )
													lowerCamelCase					:			Optional[Any]									= outputs.logits
													lowerCamelCase					:			Tuple									= logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
													lowerCamelCase					:			Any									= logits.softmax(-1      ).log()

													if scores is None:
																	lowerCamelCase					:			Tuple									= logits.topk(A,    -1      )
																	lowerCamelCase					:			Optional[Any]									= generated.expand(A,    *generated.shape[1:]      )
																	lowerCamelCase					:			Any									= next_tokens.permute(1,    0      ), scores.squeeze(0      )
																	if tokens is None:
																					lowerCamelCase					:			List[Any]									= next_tokens
																	else:
																					lowerCamelCase					:			str									= tokens.expand(A,    *tokens.shape[1:]      )
																					lowerCamelCase					:			Optional[int]									= torch.cat((tokens, next_tokens),    dim=1      )
													else:
																	lowerCamelCase					:			int									= -float(np.inf      )
																	lowerCamelCase					:			int									= 0
																	lowerCamelCase					:			List[str]									= scores[:, None] + logits
																	seq_lengths[~is_stopped] += 1
																	lowerCamelCase					:			str									= scores_sum / seq_lengths[:, None]
																	lowerCamelCase					:			Optional[Any]									= scores_sum_average.view(-1      ).topk(A,    -1      )
																	lowerCamelCase					:			str									= next_tokens // scores_sum.shape[1]
																	lowerCamelCase					:			int									= seq_lengths[next_tokens_source]
																	lowerCamelCase					:			List[Any]									= next_tokens % scores_sum.shape[1]
																	lowerCamelCase					:			List[Any]									= next_tokens.unsqueeze(1      )
																	lowerCamelCase					:			str									= tokens[next_tokens_source]
																	lowerCamelCase					:			int									= torch.cat((tokens, next_tokens),    dim=1      )
																	lowerCamelCase					:			int									= generated[next_tokens_source]
																	lowerCamelCase					:			int									= scores_sum_average * seq_lengths
																	lowerCamelCase					:			int									= is_stopped[next_tokens_source]

													lowerCamelCase					:			Optional[int]									= self.transformer.transformer.wte(next_tokens.squeeze()      ).view(generated.shape[0],    1,    -1      )
													lowerCamelCase					:			Any									= torch.cat((generated, next_token_embed),    dim=1      )
													lowerCamelCase					:			List[str]									= is_stopped + next_tokens.eq(A      ).squeeze()
													if is_stopped.all():
																	break

									lowerCamelCase					:			Tuple									= scores / seq_lengths
									lowerCamelCase					:			Optional[Any]									= scores.argsort(descending=A      )
									# tokens tensors are already padded to max_seq_length
									lowerCamelCase					:			List[str]									= [tokens[i] for i in order]
									lowerCamelCase					:			int									= torch.stack(A,    dim=0      )
									lowerCamelCase					:			List[str]									= torch.tensor([seq_lengths[i] for i in order],    dtype=seq_lengths.dtype      )
									return output_texts, seq_lengths


'''simple docstring'''






import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
					import onnxruntime as ort




class     __snake_case		(   a__    ,   unittest.TestCase):
					_lowerCAmelCase			     =	'''hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'''
					def     UpperCAmelCase_							(     self,    A=0      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= floats_tensor((1, 3, 128, 128),    rng=random.Random(A      )      )
									lowerCamelCase					:			Tuple									= np.random.RandomState(A      )
									lowerCamelCase					:			str									= {
									    'prompt': 'A painting of a squirrel eating a burger',
									    'image': image,
									    'generator': generator,
									    'num_inference_steps': 3,
									    'strength': 0.75,
									    'guidance_scale': 7.5,
									    'output_type': 'numpy',
									}
									return inputs
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Union[str, Any]									= self.get_dummy_inputs()
									lowerCamelCase					:			List[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Union[str, Any]									= image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Union[str, Any]									= np.array([0.6_9643, 0.5_8484, 0.5_0314, 0.5_8760, 0.5_5368, 0.5_9643, 0.5_1529, 0.4_1217, 0.4_9087]      )
									assert np.abs(image_slice - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Any									= PNDMScheduler.from_config(pipe.scheduler.config,    skip_prk_steps=A      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Dict									= self.get_dummy_inputs()
									lowerCamelCase					:			Dict									= pipe(**A      ).images
									lowerCamelCase					:			Optional[Any]									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Union[str, Any]									= np.array([0.6_1737, 0.5_4642, 0.5_3183, 0.5_4465, 0.5_2742, 0.6_0525, 0.4_9969, 0.4_0655, 0.4_8154]      )

									assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Optional[int]									= LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									# warmup pass to apply optimizations
									lowerCamelCase					:			Union[str, Any]									= pipe(**self.get_dummy_inputs()      )

									lowerCamelCase					:			Any									= self.get_dummy_inputs()
									lowerCamelCase					:			Optional[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Dict									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Dict									= np.array([0.5_2761, 0.5_9977, 0.4_9033, 0.4_9619, 0.5_4282, 0.5_0311, 0.4_7600, 0.4_0918, 0.4_5203]      )

									assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			int									= EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			str									= self.get_dummy_inputs()
									lowerCamelCase					:			List[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Tuple									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			List[str]									= np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]      )

									assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-1
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Union[str, Any]									= EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Union[str, Any]									= self.get_dummy_inputs()
									lowerCamelCase					:			List[Any]									= pipe(**A      ).images
									lowerCamelCase					:			Optional[Any]									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			Optional[int]									= np.array([0.5_2911, 0.6_0004, 0.4_9229, 0.4_9805, 0.5_4502, 0.5_0680, 0.4_7777, 0.4_1028, 0.4_5304]      )

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




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint,    provider='CPUExecutionProvider'      )
									lowerCamelCase					:			Any									= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			str									= self.get_dummy_inputs()
									lowerCamelCase					:			Tuple									= pipe(**A      ).images
									lowerCamelCase					:			Union[str, Any]									= image[0, -3:, -3:, -1]

									assert image.shape == (1, 128, 128, 3)
									lowerCamelCase					:			List[Any]									= np.array([0.6_5331, 0.5_8277, 0.4_8204, 0.5_6059, 0.5_3665, 0.5_6235, 0.5_0969, 0.4_0009, 0.4_6552]      )

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





@nightly
@require_onnxruntime
@require_torch_gpu
class     __snake_case		(   unittest.TestCase):
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return (
									    "CUDAExecutionProvider",
									    {
									        "gpu_mem_limit": "15000000000",  # 15GB
									        "arena_extend_strategy": "kSameAsRequested",
									    },
									)
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= ort.SessionOptions()
									lowerCamelCase					:			Union[str, Any]									= False
									return options
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    '/img2img/sketch-mountains-input.jpg'      )
									lowerCamelCase					:			int									= init_image.resize((768, 512)      )
									# using the PNDM scheduler by default
									lowerCamelCase					:			Union[str, Any]									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(
									    'CompVis/stable-diffusion-v1-4',    revision='onnx',    safety_checker=A,    feature_extractor=A,    provider=self.gpu_provider,    sess_options=self.gpu_options,    )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Optional[Any]									= 'A fantasy landscape, trending on artstation'

									lowerCamelCase					:			Dict									= np.random.RandomState(0      )
									lowerCamelCase					:			Any									= pipe(
									    prompt=A,    image=A,    strength=0.75,    guidance_scale=7.5,    num_inference_steps=10,    generator=A,    output_type='np',    )
									lowerCamelCase					:			Union[str, Any]									= output.images
									lowerCamelCase					:			Union[str, Any]									= images[0, 255:258, 383:386, -1]

									assert images.shape == (1, 512, 768, 3)
									lowerCamelCase					:			List[Any]									= np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019]      )
									# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
									    '/img2img/sketch-mountains-input.jpg'      )
									lowerCamelCase					:			Any									= init_image.resize((768, 512)      )
									lowerCamelCase					:			List[Any]									= LMSDiscreteScheduler.from_pretrained(
									    'runwayml/stable-diffusion-v1-5',    subfolder='scheduler',    revision='onnx'      )
									lowerCamelCase					:			Dict									= OnnxStableDiffusionImgaImgPipeline.from_pretrained(
									    'runwayml/stable-diffusion-v1-5',    revision='onnx',    scheduler=A,    safety_checker=A,    feature_extractor=A,    provider=self.gpu_provider,    sess_options=self.gpu_options,    )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			Union[str, Any]									= 'A fantasy landscape, trending on artstation'

									lowerCamelCase					:			List[Any]									= np.random.RandomState(0      )
									lowerCamelCase					:			Tuple									= pipe(
									    prompt=A,    image=A,    strength=0.75,    guidance_scale=7.5,    num_inference_steps=20,    generator=A,    output_type='np',    )
									lowerCamelCase					:			List[Any]									= output.images
									lowerCamelCase					:			Union[str, Any]									= images[0, 255:258, 383:386, -1]

									assert images.shape == (1, 512, 768, 3)
									lowerCamelCase					:			int									= np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431]      )
									# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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