



import gc
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

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


enable_full_determinism()



class   UpperCAmelCase_    (       a   ,					unittest.TestCase):
							lowerCamelCase__   			=		LDMTextToImagePipeline
							lowerCamelCase__   			=		TEXT_TO_IMAGE_PARAMS - {
							    'negative_prompt',
							    'negative_prompt_embeds',
							    'cross_attention_kwargs',
							    'prompt_embeds',
							}
							lowerCamelCase__   			=		PipelineTesterMixin.required_optional_params - {
							    'num_images_per_prompt',
							    'callback',
							    'callback_steps',
							}
							lowerCamelCase__   			=		TEXT_TO_IMAGE_BATCH_PARAMS
							lowerCamelCase__   			=		False







							def    snake_case__      (   self):

														'''simple docstring'''

														torch.manual_seed(0)
														_lowerCAmelCase     :							Union[str, Any]    =    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,			)
														_lowerCAmelCase     :							Optional[int]    =    DDIMScheduler(
														    beta_start=0.00_085,			beta_end=0.012,			beta_schedule="scaled_linear",			clip_sample=__a,			set_alpha_to_one=__a,			)
														torch.manual_seed(0)
														_lowerCAmelCase     :							int    =    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,			)
														torch.manual_seed(0)
														_lowerCAmelCase     :							str    =    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=1000,			)
														_lowerCAmelCase     :							Any    =    CLIPTextModel(__a)
														_lowerCAmelCase     :							Optional[Any]    =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

														_lowerCAmelCase     :							List[Any]    =    {
														    "unet": unet,
														    "scheduler": scheduler,
														    "vqvae": vae,
														    "bert": text_encoder,
														    "tokenizer": tokenizer,
														}
														return components







							def    snake_case__      (   self,			__a,			__a=0):

														'''simple docstring'''

														if str(__a).startswith("mps"):
																					_lowerCAmelCase     :							Any    =    torch.manual_seed(__a)
														else:
																					_lowerCAmelCase     :							Optional[Any]    =    torch.Generator(device=__a).manual_seed(__a)
														_lowerCAmelCase     :							Dict    =    {
														    "prompt": "A painting of a squirrel eating a burger",
														    "generator": generator,
														    "num_inference_steps": 2,
														    "guidance_scale": 6.0,
														    "output_type": "numpy",
														}
														return inputs







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							List[str]    =    "cpu"  # ensure determinism for the device-dependent torch.Generator

														_lowerCAmelCase     :							Dict    =    self.get_dummy_components()
														_lowerCAmelCase     :							Optional[Any]    =    LDMTextToImagePipeline(**__a)
														pipe.to(__a)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							str    =    self.get_dummy_inputs(__a)
														_lowerCAmelCase     :							Optional[int]    =    pipe(**__a).images
														_lowerCAmelCase     :							Dict    =    image[0, -3:, -3:, -1]

														assert image.shape == (1, 16, 16, 3)
														_lowerCAmelCase     :							Any    =    np.array([0.6_101, 0.6_156, 0.5_622, 0.4_895, 0.6_661, 0.3_804, 0.5_748, 0.6_136, 0.5_014])

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



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







							def    snake_case__      (   self):

														'''simple docstring'''

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







							def    snake_case__      (   self,			__a,			__a=torch.floataa,			__a=0):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[int]    =    torch.manual_seed(__a)
														_lowerCAmelCase     :							Any    =    np.random.RandomState(__a).standard_normal((1, 4, 32, 32))
														_lowerCAmelCase     :							Optional[int]    =    torch.from_numpy(__a).to(device=__a,			dtype=__a)
														_lowerCAmelCase     :							Tuple    =    {
														    "prompt": "A painting of a squirrel eating a burger",
														    "latents": latents,
														    "generator": generator,
														    "num_inference_steps": 3,
														    "guidance_scale": 6.0,
														    "output_type": "numpy",
														}
														return inputs







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256").to(__a)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Union[str, Any]    =    self.get_inputs(__a)
														_lowerCAmelCase     :							List[Any]    =    pipe(**__a).images
														_lowerCAmelCase     :							str    =    image[0, -3:, -3:, -1].flatten()

														assert image.shape == (1, 256, 256, 3)
														_lowerCAmelCase     :							List[str]    =    np.array([0.51_825, 0.52_850, 0.52_543, 0.54_258, 0.52_304, 0.52_569, 0.54_363, 0.55_276, 0.56_878])
														_lowerCAmelCase     :							str    =    np.abs(expected_slice - image_slice).max()
														assert max_diff < 1E-3



@nightly
@require_torch_gpu
class   UpperCAmelCase_    (       unittest.TestCase):







							def    snake_case__      (   self):

														'''simple docstring'''

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







							def    snake_case__      (   self,			__a,			__a=torch.floataa,			__a=0):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[Any]    =    torch.manual_seed(__a)
														_lowerCAmelCase     :							Dict    =    np.random.RandomState(__a).standard_normal((1, 4, 32, 32))
														_lowerCAmelCase     :							str    =    torch.from_numpy(__a).to(device=__a,			dtype=__a)
														_lowerCAmelCase     :							Dict    =    {
														    "prompt": "A painting of a squirrel eating a burger",
														    "latents": latents,
														    "generator": generator,
														    "num_inference_steps": 50,
														    "guidance_scale": 6.0,
														    "output_type": "numpy",
														}
														return inputs







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256").to(__a)
														pipe.set_progress_bar_config(disable=__a)

														_lowerCAmelCase     :							Optional[Any]    =    self.get_inputs(__a)
														_lowerCAmelCase     :							List[str]    =    pipe(**__a).images[0]

														_lowerCAmelCase     :							Optional[int]    =    load_numpy(
														    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy")
														_lowerCAmelCase     :							Dict    =    np.abs(expected_image - image).max()
														assert max_diff < 1E-3





from __future__ import absolute_import, division, print_function, unicode_literals

from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy



@add_start_docstrings(
    'The RoBERTa Model transformer with early exiting (DeeRoBERTa). '   ,					a   ,					)
class   UpperCAmelCase_    (       a):
							lowerCamelCase__   			=		RobertaConfig
							lowerCamelCase__   			=		'roberta'







							def __init__(   self,			__a):

														'''simple docstring'''

														super().__init__(__a)

														_lowerCAmelCase     :							Optional[Any]    =    RobertaEmbeddings(__a)
														self.init_weights()



@add_start_docstrings(
    'RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,\n    also takes care of multi-layer training. '   ,					a   ,					)
class   UpperCAmelCase_    (       a):
							lowerCamelCase__   			=		RobertaConfig
							lowerCamelCase__   			=		'roberta'







							def __init__(   self,			__a):

														'''simple docstring'''

														super().__init__(__a)
														_lowerCAmelCase     :							Optional[int]    =    config.num_labels
														_lowerCAmelCase     :							Optional[int]    =    config.num_hidden_layers

														_lowerCAmelCase     :							Optional[int]    =    DeeRobertaModel(__a)
														_lowerCAmelCase     :							Union[str, Any]    =    nn.Dropout(config.hidden_dropout_prob)
														_lowerCAmelCase     :							List[str]    =    nn.Linear(config.hidden_size,			self.config.num_labels)







							@add_start_docstrings_to_model_forward(__a)
							def    snake_case__      (   self,			__a=None,			__a=None,			__a=None,			__a=None,			__a=None,			__a=None,			__a=None,			__a=-1,			__a=False,			):

														'''simple docstring'''

														_lowerCAmelCase     :							Union[str, Any]    =    self.num_layers
														try:
																					_lowerCAmelCase     :							List[Any]    =    self.roberta(
																					    __a,			attention_mask=__a,			token_type_ids=__a,			position_ids=__a,			head_mask=__a,			inputs_embeds=__a,			)

																					_lowerCAmelCase     :							List[Any]    =    outputs[1]

																					_lowerCAmelCase     :							Dict    =    self.dropout(__a)
																					_lowerCAmelCase     :							Dict    =    self.classifier(__a)
																					_lowerCAmelCase     :							Optional[Any]    =    (logits,) + outputs[2:]  # add hidden states and attention if they are here
														except HighwayException as e:
																					_lowerCAmelCase     :							Tuple    =    e.message
																					_lowerCAmelCase     :							Union[str, Any]    =    e.exit_layer
																					_lowerCAmelCase     :							List[Any]    =    outputs[0]

														if not self.training:
																					_lowerCAmelCase     :							int    =    entropy(__a)
																					_lowerCAmelCase     :							List[Any]    =    []
																					_lowerCAmelCase     :							str    =    []
														if labels is not None:
																					if self.num_labels == 1:
																												#  We are doing regression
																												_lowerCAmelCase     :							Optional[Any]    =    MSELoss()
																												_lowerCAmelCase     :							int    =    loss_fct(logits.view(-1),			labels.view(-1))
																					else:
																												_lowerCAmelCase     :							Optional[Any]    =    CrossEntropyLoss()
																												_lowerCAmelCase     :							Optional[Any]    =    loss_fct(logits.view(-1,			self.num_labels),			labels.view(-1))

																					# work with highway exits
																					_lowerCAmelCase     :							Optional[int]    =    []
																					for highway_exit in outputs[-1]:
																												_lowerCAmelCase     :							Any    =    highway_exit[0]
																												if not self.training:
																																			highway_logits_all.append(__a)
																																			highway_entropy.append(highway_exit[2])
																												if self.num_labels == 1:
																																			#  We are doing regression
																																			_lowerCAmelCase     :							List[str]    =    MSELoss()
																																			_lowerCAmelCase     :							List[Any]    =    loss_fct(highway_logits.view(-1),			labels.view(-1))
																												else:
																																			_lowerCAmelCase     :							Dict    =    CrossEntropyLoss()
																																			_lowerCAmelCase     :							Optional[Any]    =    loss_fct(highway_logits.view(-1,			self.num_labels),			labels.view(-1))
																												highway_losses.append(__a)

																					if train_highway:
																												_lowerCAmelCase     :							int    =    (sum(highway_losses[:-1]),) + outputs
																												# exclude the final highway, of course
																					else:
																												_lowerCAmelCase     :							Any    =    (loss,) + outputs
														if not self.training:
																					_lowerCAmelCase     :							Optional[Any]    =    outputs + ((original_entropy, highway_entropy), exit_layer)
																					if output_layer >= 0:
																												_lowerCAmelCase     :							Optional[Any]    =    (
																												    (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
																												)  # use the highway of the last layer

														return outputs  # (loss), logits, (hidden_states), (attentions), entropy
