

"""simple docstring"""





import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel




class 					_lowerCAmelCase  :







 """simple docstring"""







 def __init__(			self							:			Any,					UpperCAmelCase__							:			str,					UpperCAmelCase__							:			Union[str, Any]=1_3,					UpperCAmelCase__							:			str=7,					UpperCAmelCase__							:			int=True,					UpperCAmelCase__							:			Optional[int]=True,					UpperCAmelCase__							:			Any=False,					UpperCAmelCase__							:			str=True,					UpperCAmelCase__							:			Tuple=9_9,					UpperCAmelCase__							:			Union[str, Any]=3_2,					UpperCAmelCase__							:			List[str]=5,					UpperCAmelCase__							:			int=4,					UpperCAmelCase__							:			str=3_7,					UpperCAmelCase__							:			Dict="gelu",					UpperCAmelCase__							:			List[str]=0.1,					UpperCAmelCase__							:			Tuple=0.1,					UpperCAmelCase__							:			Dict=5_1_2,					UpperCAmelCase__							:			List[str]=1_6,					UpperCAmelCase__							:			Optional[Any]=2,					UpperCAmelCase__							:			List[Any]=0.02,					UpperCAmelCase__							:			str=3,					UpperCAmelCase__							:			Dict=4,					UpperCAmelCase__							:			Union[str, Any]=None,					):
    __lowercase								=						parent
    __lowercase								=						batch_size
    __lowercase								=						seq_length
    __lowercase								=						is_training
    __lowercase								=						use_input_mask
    __lowercase								=						use_token_type_ids
    __lowercase								=						use_labels
    __lowercase								=						vocab_size
    __lowercase								=						hidden_size
    __lowercase								=						num_hidden_layers
    __lowercase								=						num_attention_heads
    __lowercase								=						intermediate_size
    __lowercase								=						hidden_act
    __lowercase								=						hidden_dropout_prob
    __lowercase								=						attention_probs_dropout_prob
    __lowercase								=						max_position_embeddings
    __lowercase								=						type_vocab_size
    __lowercase								=						type_sequence_label_size
    __lowercase								=						initializer_range
    __lowercase								=						num_labels
    __lowercase								=						num_choices
    __lowercase								=						scope


 def 					_lowercase (			self							:			Optional[Any]	):
    __lowercase								=						ids_tensor([self.batch_size, self.seq_length],					self.vocab_size	)

    __lowercase								=						None
    if self.use_input_mask:
       __lowercase								=						random_attention_mask([self.batch_size, self.seq_length]	)

    __lowercase								=						None
    if self.use_token_type_ids:
       __lowercase								=						ids_tensor([self.batch_size, self.seq_length],					self.type_vocab_size	)

    __lowercase								=						None
    __lowercase								=						None
    __lowercase								=						None
    if self.use_labels:
       __lowercase								=						ids_tensor([self.batch_size],					self.type_sequence_label_size	)
       __lowercase								=						ids_tensor([self.batch_size, self.seq_length],					self.num_labels	)
       __lowercase								=						ids_tensor([self.batch_size],					self.num_choices	)

    __lowercase								=						self.get_config()

    return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


 def 					_lowercase (			self							:			List[str]	):
    return OpenLlamaConfig(
        vocab_size=self.vocab_size,					hidden_size=self.hidden_size,					num_hidden_layers=self.num_hidden_layers,					num_attention_heads=self.num_attention_heads,					intermediate_size=self.intermediate_size,					hidden_act=self.hidden_act,					hidden_dropout_prob=self.hidden_dropout_prob,					attention_probs_dropout_prob=self.attention_probs_dropout_prob,					max_position_embeddings=self.max_position_embeddings,					type_vocab_size=self.type_vocab_size,					is_decoder=UpperCAmelCase__,					initializer_range=self.initializer_range,					use_stable_embedding=UpperCAmelCase__,					)


 def 					_lowercase (			self							:			Optional[Any],					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Union[str, Any],					UpperCAmelCase__							:			Optional[Any],					UpperCAmelCase__							:			Optional[Any],					UpperCAmelCase__							:			int	):
    __lowercase								=						OpenLlamaModel(config=UpperCAmelCase__	)
    model.to(UpperCAmelCase__	)
    model.eval()
    __lowercase								=						model(UpperCAmelCase__,					attention_mask=UpperCAmelCase__	)
    __lowercase								=						model(UpperCAmelCase__	)
    self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)	)


 def 					_lowercase (			self							:			Union[str, Any],					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			str,					UpperCAmelCase__							:			Optional[Any],					UpperCAmelCase__							:			Optional[Any],					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Tuple,					UpperCAmelCase__							:			Optional[Any],					):
    __lowercase								=						True
    __lowercase								=						OpenLlamaModel(UpperCAmelCase__	)
    model.to(UpperCAmelCase__	)
    model.eval()
    __lowercase								=						model(
        UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					encoder_hidden_states=UpperCAmelCase__,					encoder_attention_mask=UpperCAmelCase__,					)
    __lowercase								=						model(
        UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					encoder_hidden_states=UpperCAmelCase__,					)
    __lowercase								=						model(UpperCAmelCase__,					attention_mask=UpperCAmelCase__	)
    self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)	)


 def 					_lowercase (			self							:			List[Any],					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Optional[Any],					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			List[Any],					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			Any,					):
    __lowercase								=						OpenLlamaForCausalLM(config=UpperCAmelCase__	)
    model.to(UpperCAmelCase__	)
    model.eval()
    __lowercase								=						model(UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					labels=UpperCAmelCase__	)
    self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.seq_length, self.vocab_size)	)


 def 					_lowercase (			self							:			int,					UpperCAmelCase__							:			str,					UpperCAmelCase__							:			List[str],					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			List[str],					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			List[str],					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			Tuple,					):
    __lowercase								=						True
    __lowercase								=						True
    __lowercase								=						OpenLlamaForCausalLM(config=UpperCAmelCase__	)
    model.to(UpperCAmelCase__	)
    model.eval()

    # first forward pass
    __lowercase								=						model(
        UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					encoder_hidden_states=UpperCAmelCase__,					encoder_attention_mask=UpperCAmelCase__,					use_cache=UpperCAmelCase__,					)
    __lowercase								=						outputs.past_key_values

    # create hypothetical multiple next token and extent to next_input_ids
    __lowercase								=						ids_tensor((self.batch_size, 3),					config.vocab_size	)
    __lowercase								=						ids_tensor((self.batch_size, 3),					vocab_size=2	)

    # append to next input_ids and
    __lowercase								=						torch.cat([input_ids, next_tokens],					dim=-1	)
    __lowercase								=						torch.cat([input_mask, next_mask],					dim=-1	)

    __lowercase								=						model(
        UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					encoder_hidden_states=UpperCAmelCase__,					encoder_attention_mask=UpperCAmelCase__,					output_hidden_states=UpperCAmelCase__,					)["hidden_states"][0]
    __lowercase								=						model(
        UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					encoder_hidden_states=UpperCAmelCase__,					encoder_attention_mask=UpperCAmelCase__,					past_key_values=UpperCAmelCase__,					output_hidden_states=UpperCAmelCase__,					)["hidden_states"][0]

    # select random slice
    __lowercase								=						ids_tensor((1,),					output_from_past.shape[-1]	).item()
    __lowercase								=						output_from_no_past[:, -3:, random_slice_idx].detach()
    __lowercase								=						output_from_past[:, :, random_slice_idx].detach()

    self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]	)

    # test that outputs are equal for slice
    self.parent.assertTrue(torch.allclose(UpperCAmelCase__,					UpperCAmelCase__,					atol=1E-3	)	)



 def 					_lowercase (			self							:			Any	):
    __lowercase								=						self.prepare_config_and_inputs()
    (
        (
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,
    )								=						config_and_inputs
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    return config, inputs_dict







@require_torch
class 					_lowerCAmelCase  (	lowercase     ,lowercase     ,lowercase     ,unittest.TestCase   ):







 """simple docstring"""







 __UpperCAmelCase	:					int   		=	(
     (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
 )
 __UpperCAmelCase	:					Union[str, Any]   		=	(OpenLlamaForCausalLM,) if is_torch_available() else ()
 __UpperCAmelCase	:					Any   		=	(
     {
         "feature-extraction": OpenLlamaModel,
         "text-classification": OpenLlamaForSequenceClassification,
         "text-generation": OpenLlamaForCausalLM,
         "zero-shot": OpenLlamaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 __UpperCAmelCase	:					List[str]   		=	False
 __UpperCAmelCase	:					Dict   		=	False


 def 					_lowercase (			self							:			Optional[int]	):
    __lowercase								=						OpenLlamaModelTester(self	)
    __lowercase								=						ConfigTester(self,					config_class=UpperCAmelCase__,					hidden_size=3_7	)


 def 					_lowercase (			self							:			int	):
    self.config_tester.run_common_tests()


 def 					_lowercase (			self							:			Optional[int]	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*UpperCAmelCase__	)


 def 					_lowercase (			self							:			Union[str, Any]	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    for type in ["absolute", "relative_key", "relative_key_query"]:
       __lowercase								=						type
       self.model_tester.create_and_check_model(*UpperCAmelCase__	)


 def 					_lowercase (			self							:			Union[str, Any]	):
    __lowercase		,__lowercase								=						self.model_tester.prepare_config_and_inputs_for_common()
    __lowercase								=						3
    __lowercase								=						input_dict["input_ids"]
    __lowercase								=						input_ids.ne(1	).to(UpperCAmelCase__	)
    __lowercase								=						ids_tensor([self.model_tester.batch_size],					self.model_tester.type_sequence_label_size	)
    __lowercase								=						OpenLlamaForSequenceClassification(UpperCAmelCase__	)
    model.to(UpperCAmelCase__	)
    model.eval()
    __lowercase								=						model(UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					labels=UpperCAmelCase__	)
    self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)	)


 def 					_lowercase (			self							:			List[Any]	):
    __lowercase		,__lowercase								=						self.model_tester.prepare_config_and_inputs_for_common()
    __lowercase								=						3
    __lowercase								=						"single_label_classification"
    __lowercase								=						input_dict["input_ids"]
    __lowercase								=						input_ids.ne(1	).to(UpperCAmelCase__	)
    __lowercase								=						ids_tensor([self.model_tester.batch_size],					self.model_tester.type_sequence_label_size	)
    __lowercase								=						OpenLlamaForSequenceClassification(UpperCAmelCase__	)
    model.to(UpperCAmelCase__	)
    model.eval()
    __lowercase								=						model(UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					labels=UpperCAmelCase__	)
    self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)	)


 def 					_lowercase (			self							:			Union[str, Any]	):
    __lowercase		,__lowercase								=						self.model_tester.prepare_config_and_inputs_for_common()
    __lowercase								=						3
    __lowercase								=						"multi_label_classification"
    __lowercase								=						input_dict["input_ids"]
    __lowercase								=						input_ids.ne(1	).to(UpperCAmelCase__	)
    __lowercase								=						ids_tensor(
        [self.model_tester.batch_size, config.num_labels],					self.model_tester.type_sequence_label_size	).to(torch.float	)
    __lowercase								=						OpenLlamaForSequenceClassification(UpperCAmelCase__	)
    model.to(UpperCAmelCase__	)
    model.eval()
    __lowercase								=						model(UpperCAmelCase__,					attention_mask=UpperCAmelCase__,					labels=UpperCAmelCase__	)
    self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)	)


 @unittest.skip("Open-Llama buffers include complex numbers, which breaks this test"	)
 def 					_lowercase (			self							:			Dict	):
    pass



 @parameterized.expand([("linear",), ("dynamic",)]	)
 def 					_lowercase (			self							:			List[Any],					UpperCAmelCase__							:			List[Any]	):
    __lowercase		,__lowercase								=						self.model_tester.prepare_config_and_inputs_for_common()
    __lowercase								=						ids_tensor([1, 1_0],					config.vocab_size	)
    __lowercase								=						ids_tensor([1, int(config.max_position_embeddings * 1.5	)],					config.vocab_size	)

    set_seed(4_2	)  # Fixed seed at init time so the two models get the same random weights
    __lowercase								=						OpenLlamaModel(UpperCAmelCase__	)
    original_model.to(UpperCAmelCase__	)
    original_model.eval()
    __lowercase								=						original_model(UpperCAmelCase__	).last_hidden_state
    __lowercase								=						original_model(UpperCAmelCase__	).last_hidden_state

    set_seed(4_2	)  # Fixed seed at init time so the two models get the same random weights
    __lowercase								=						{"type": scaling_type, "factor": 10.0}
    __lowercase								=						OpenLlamaModel(UpperCAmelCase__	)
    scaled_model.to(UpperCAmelCase__	)
    scaled_model.eval()
    __lowercase								=						scaled_model(UpperCAmelCase__	).last_hidden_state
    __lowercase								=						scaled_model(UpperCAmelCase__	).last_hidden_state

    # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
    # maximum sequence length, so the outputs for the short input should match.
    if scaling_type == "dynamic":
       self.assertTrue(torch.allclose(UpperCAmelCase__,					UpperCAmelCase__,					atol=1E-5	)	)
    else:
       self.assertFalse(torch.allclose(UpperCAmelCase__,					UpperCAmelCase__,					atol=1E-5	)	)

    # The output should be different for long inputs
    self.assertFalse(torch.allclose(UpperCAmelCase__,					UpperCAmelCase__,					atol=1E-5	)	)








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

import numpy as np
import PIL
import torch
import torch.utils.checkpoint

from ...models import UNetaDModel, VQModel
from ...schedulers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
)
from ...utils import PIL_INTERPOLATION, randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
def 							_UpperCAmelCase				(UpperCamelCase__     :  Union[str, Any]						):
       _A    ,			_A      :       Any				    = image.size
       _A    ,			_A      :       str				    = (x - x % 32 for x in (w, h))  # resize to integer multiple of 32
       _A      :       List[str]				    = image.resize((w, h)				,							resample=PIL_INTERPOLATION["lanczos"]						)
       _A      :       Any				    = np.array(UpperCamelCase__						).astype(np.floataa						) / 2_55.0
       _A      :       Optional[Any]				    = image[None].transpose(0				,							3				,							1				,							2						)
       _A      :       Union[str, Any]				    = torch.from_numpy(UpperCamelCase__						)
       return 2.0 * image - 1.0







class      lowerCAmelCase__       (				a):




     '''simple docstring'''


     def __init__(						self ,						__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						)			->				Optional[int]:
            super().__init__()
            self.register_modules(vqvae=__lowerCamelCase ,						unet=__lowerCamelCase ,						scheduler=__lowerCamelCase)


     @torch.no_grad()
     def __call__(						self ,						__lowerCamelCase = None ,						__lowerCamelCase = 1 ,						__lowerCamelCase = 1_0_0 ,						__lowerCamelCase = 0.0 ,						__lowerCamelCase = None ,						__lowerCamelCase = "pil" ,						__lowerCamelCase = True ,						)			->				Union[Tuple, ImagePipelineOutput]:
            if isinstance(__lowerCamelCase ,						PIL.Image.Image):
                   _A      :       Tuple				    = 1
            elif isinstance(__lowerCamelCase ,						torch.Tensor):
                   _A      :       Union[str, Any]				    = image.shape[0]
            else:
                   raise ValueError(F"`image` has to be of type `PIL.Image.Image` or `torch.Tensor` but is {type(__lowerCamelCase)}")

            if isinstance(__lowerCamelCase ,						PIL.Image.Image):
                   _A      :       Union[str, Any]				    = preprocess(__lowerCamelCase)

            _A    ,			_A      :       Union[str, Any]				    = image.shape[-2:]

            # in_channels should be 6: 3 for latents, 3 for low resolution image
            _A      :       Optional[Any]				    = (batch_size, self.unet.config.in_channels // 2, height, width)
            _A      :       str				    = next(self.unet.parameters()).dtype

            _A      :       Union[str, Any]				    = randn_tensor(__lowerCamelCase ,						generator=__lowerCamelCase ,						device=self.device ,						dtype=__lowerCamelCase)

            _A      :       List[Any]				    = image.to(device=self.device ,						dtype=__lowerCamelCase)

            # set timesteps and move to the correct device
            self.scheduler.set_timesteps(__lowerCamelCase ,						device=self.device)
            _A      :       Any				    = self.scheduler.timesteps

            # scale the initial noise by the standard deviation required by the scheduler
            _A      :       List[str]				    = latents * self.scheduler.init_noise_sigma

            # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature.
            # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
            # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
            # and should be between [0, 1]
            _A      :       str				    = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
            _A      :       Optional[int]				    = {}
            if accepts_eta:
                   _A      :       List[Any]				    = eta

            for t in self.progress_bar(__lowerCamelCase):
                   # concat latents and low resolution image in the channel dimension.
                   _A      :       List[Any]				    = torch.cat([latents, image] ,						dim=1)
                   _A      :       str				    = self.scheduler.scale_model_input(__lowerCamelCase ,						__lowerCamelCase)
                   # predict the noise residual
                   _A      :       Any				    = self.unet(__lowerCamelCase ,						__lowerCamelCase).sample
                   # compute the previous noisy sample x_t -> x_t-1
                   _A      :       Optional[int]				    = self.scheduler.step(__lowerCamelCase ,						__lowerCamelCase ,						__lowerCamelCase ,						**__lowerCamelCase).prev_sample

            # decode the image latents with the VQVAE
            _A      :       Union[str, Any]				    = self.vqvae.decode(__lowerCamelCase).sample
            _A      :       Dict				    = torch.clamp(__lowerCamelCase ,						-1.0 ,						1.0)
            _A      :       Tuple				    = image / 2 + 0.5
            _A      :       int				    = image.cpu().permute(0 ,						2 ,						3 ,						1).numpy()

            if output_type == "pil":
                   _A      :       Optional[int]				    = self.numpy_to_pil(__lowerCamelCase)

            if not return_dict:
                   return (image,)

            return ImagePipelineOutput(images=__lowerCamelCase)
