

'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import OpenAIGPTConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers.models.openai.modeling_tf_openai import (
           TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFOpenAIGPTDoubleHeadsModel,
           TFOpenAIGPTForSequenceClassification,
           TFOpenAIGPTLMHeadModel,
           TFOpenAIGPTModel,
       )


class 					lowerCAmelCase       :







       def __init__(							self     :	Dict  , __lowercase     :	str  , ):







            """simple docstring"""







            __lowercase     						=parent
            __lowercase     						=13
            __lowercase     						=7
            __lowercase     						=True
            __lowercase     						=True
            __lowercase     						=True
            __lowercase     						=True
            __lowercase     						=True
            __lowercase     						=99
            __lowercase     						=32
            __lowercase     						=2
            __lowercase     						=4
            __lowercase     						=37
            __lowercase     						='gelu'
            __lowercase     						=0.1
            __lowercase     						=0.1
            __lowercase     						=512
            __lowercase     						=16
            __lowercase     						=2
            __lowercase     						=0.0_2
            __lowercase     						=3
            __lowercase     						=4
            __lowercase     						=None
            __lowercase     						=self.vocab_size - 1







       def      snake_case				(							self     :	Optional[Any] ):







            """simple docstring"""







            __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
            if self.use_mc_token_ids:
                 __lowercase     						=ids_tensor([self.batch_size, self.num_choices]  , self.seq_length )

            __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     						=OpenAIGPTConfig(
                vocab_size=self.vocab_size  , n_embd=self.hidden_size  , n_layer=self.num_hidden_layers  , n_head=self.num_attention_heads  , n_positions=self.max_position_embeddings  , pad_token_id=self.pad_token_id  , )

            __lowercase     						=ids_tensor([self.num_hidden_layers, self.num_attention_heads]  , 2 )

            return (
                config,
                input_ids,
                input_mask,
                head_mask,
                token_type_ids,
                mc_token_ids,
                sequence_labels,
                token_labels,
                choice_labels,
            )







       def      snake_case				(							self     :	List[Any]  , __lowercase     :	Optional[int]  , __lowercase     :	Optional[Any]  , __lowercase     :	Optional[Any]  , __lowercase     :	Optional[int]  , __lowercase     :	Dict  , *__lowercase     :	Any ):







            """simple docstring"""







            __lowercase     						=TFOpenAIGPTModel(config=__lowercase )
            __lowercase     						={'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
            __lowercase     						=model(__lowercase )

            __lowercase     						=[input_ids, input_mask]
            __lowercase     						=model(__lowercase )

            __lowercase     						=model(__lowercase )

            self.parent.assertEqual(result.last_hidden_state.shape  , (self.batch_size, self.seq_length, self.hidden_size) )







       def      snake_case				(							self     :	Optional[Any]  , __lowercase     :	Tuple  , __lowercase     :	str  , __lowercase     :	Dict  , __lowercase     :	Optional[Any]  , __lowercase     :	Union[str, Any]  , *__lowercase     :	Dict ):







            """simple docstring"""







            __lowercase     						=TFOpenAIGPTLMHeadModel(config=__lowercase )
            __lowercase     						={'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
            __lowercase     						=model(__lowercase )
            self.parent.assertEqual(result.logits.shape  , (self.batch_size, self.seq_length, self.vocab_size) )







       def      snake_case				(							self     :	int  , __lowercase     :	Dict  , __lowercase     :	List[Any]  , __lowercase     :	List[Any]  , __lowercase     :	List[str]  , __lowercase     :	List[Any]  , __lowercase     :	List[Any]  , *__lowercase     :	Union[str, Any] ):







            """simple docstring"""







            __lowercase     						=TFOpenAIGPTDoubleHeadsModel(config=__lowercase )

            __lowercase     						=tf.tile(tf.expand_dims(__lowercase  , 1 )  , (1, self.num_choices, 1) )
            __lowercase     						=tf.tile(tf.expand_dims(__lowercase  , 1 )  , (1, self.num_choices, 1) )
            __lowercase     						=tf.tile(tf.expand_dims(__lowercase  , 1 )  , (1, self.num_choices, 1) )

            __lowercase     						={
                'input_ids': multiple_choice_inputs_ids,
                'mc_token_ids': mc_token_ids,
                'attention_mask': multiple_choice_input_mask,
                'token_type_ids': multiple_choice_token_type_ids,
            }
            __lowercase     						=model(__lowercase )
            self.parent.assertEqual(
                result.logits.shape  , (self.batch_size, self.num_choices, self.seq_length, self.vocab_size) )
            self.parent.assertEqual(result.mc_logits.shape  , (self.batch_size, self.num_choices) )







       def      snake_case				(							self     :	Optional[Any]  , __lowercase     :	List[Any]  , __lowercase     :	int  , __lowercase     :	List[str]  , __lowercase     :	Tuple  , __lowercase     :	int  , *__lowercase     :	Any ):







            """simple docstring"""







            __lowercase     						=self.num_labels
            __lowercase     						=ids_tensor([self.batch_size]  , self.type_sequence_label_size )
            __lowercase     						={
                'input_ids': input_ids,
                'attention_mask': input_mask,
                'token_type_ids': token_type_ids,
                'labels': sequence_labels,
            }
            __lowercase     						=TFOpenAIGPTForSequenceClassification(__lowercase )
            __lowercase     						=model(__lowercase )
            self.parent.assertEqual(result.logits.shape  , (self.batch_size, self.num_labels) )







       def      snake_case				(							self     :	List[Any] ):







            """simple docstring"""







            __lowercase     						=self.prepare_config_and_inputs()

            (
                (
                __lowercase
            )						,       (
                __lowercase
            )						,       (
                __lowercase
            )						,       (
                __lowercase
            )						,       (
                __lowercase
            )						,       (
                __lowercase
            )						,       (
                __lowercase
            )						,       (
                __lowercase
            )						,       (
                __lowercase
            )						,       
            )     						=config_and_inputs

            __lowercase     						={'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
            return config, inputs_dict






@require_tf
class 					lowerCAmelCase       ( A				,     A				,     unittest.TestCase				):
       lowerCAmelCase_				=	(
           (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
           if is_tf_available()
           else ()
       )
       lowerCAmelCase_				=	(
           (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
       )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
       lowerCAmelCase_				=	(
           {
               "feature-extraction": TFOpenAIGPTModel,
               "text-classification": TFOpenAIGPTForSequenceClassification,
               "text-generation": TFOpenAIGPTLMHeadModel,
               "zero-shot": TFOpenAIGPTForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       lowerCAmelCase_				=	False
       lowerCAmelCase_				=	False







       def      snake_case				(							self     :	str  , __lowercase     :	List[str]  , __lowercase     :	Union[str, Any]  , __lowercase     :	List[Any]  , __lowercase     :	List[str]  , __lowercase     :	Union[str, Any] ):







            """simple docstring"""







            if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
                 # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
                 # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
                 # tiny config could not be created.
                 return True

            return False







       def      snake_case				(							self     :	Union[str, Any] ):







            """simple docstring"""







            __lowercase     						=TFOpenAIGPTModelTester(self )
            __lowercase     						=ConfigTester(self  , config_class=__lowercase  , n_embd=37 )







       def      snake_case				(							self     :	int ):







            """simple docstring"""







            self.config_tester.run_common_tests()







       def      snake_case				(							self     :	Any ):







            """simple docstring"""







            __lowercase     						=self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_model(*__lowercase )







       def      snake_case				(							self     :	str ):







            """simple docstring"""







            __lowercase     						=self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_lm_head(*__lowercase )







       def      snake_case				(							self     :	Dict ):







            """simple docstring"""







            __lowercase     						=self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_double_head(*__lowercase )







       def      snake_case				(							self     :	Union[str, Any] ):







            """simple docstring"""







            __lowercase     						=self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*__lowercase )







       @slow
       def      snake_case				(							self     :	List[Any] ):







            """simple docstring"""







            for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 __lowercase     						=TFOpenAIGPTModel.from_pretrained(__lowercase )
                 self.assertIsNotNone(__lowercase )






@require_tf
class 					lowerCAmelCase       ( unittest.TestCase				):







       @slow
       def      snake_case				(							self     :	Dict ):







            """simple docstring"""







            __lowercase     						=TFOpenAIGPTLMHeadModel.from_pretrained('openai-gpt' )
            __lowercase     						=tf.convert_to_tensor([[481, 4735, 544]]  , dtype=tf.intaa )  # the president is
            __lowercase     						=[
                481,
                4735,
                544,
                246,
                963,
                870,
                762,
                239,
                244,
                40477,
                244,
                249,
                719,
                881,
                487,
                544,
                240,
                244,
                603,
                481,
            ]  # the president is a very good man. " \n " i\'m sure he is, " said the

            __lowercase     						=model.generate(__lowercase  , do_sample=__lowercase )
            self.assertListEqual(output_ids[0].numpy().tolist()  , __lowercase )





'''simple docstring'''

from typing import Union

import torch
from PIL import Image
from torchvision import transforms as tfms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)


class 					lowerCAmelCase       ( A				):







       def __init__(							self     :	Union[str, Any]  , __lowercase     :	AutoencoderKL  , __lowercase     :	CLIPTextModel  , __lowercase     :	CLIPTokenizer  , __lowercase     :	UNetaDConditionModel  , __lowercase     :	Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler]  , ):







            """simple docstring"""







            super().__init__()

            self.register_modules(vae=__lowercase  , text_encoder=__lowercase  , tokenizer=__lowercase  , unet=__lowercase  , scheduler=__lowercase )







       def      snake_case				(							self     :	str  , __lowercase     :	Optional[int] ):







            """simple docstring"""







            with torch.no_grad():
                 __lowercase     						=self.vae.encode(tfms.ToTensor()(__lowercase ).unsqueeze(0 ).to(self.device ) * 2 - 1 )
                 __lowercase     						=0.1_8_2_1_5 * latent.latent_dist.sample()
            return latent







       def      snake_case				(							self     :	List[str]  , __lowercase     :	Tuple ):







            """simple docstring"""







            __lowercase     						=(1 / 0.1_8_2_1_5) * latent
            with torch.no_grad():
                 __lowercase     						=self.vae.decode(__lowercase ).sample
            __lowercase     						=(img / 2 + 0.5).clamp(0  , 1 )
            __lowercase     						=img.detach().cpu().permute(0  , 2  , 3  , 1 ).numpy()
            __lowercase     						=(img * 255).round().astype('uint8' )
            return Image.fromarray(img[0] )







       def      snake_case				(							self     :	int  , __lowercase     :	Optional[int] ):







            """simple docstring"""







            __lowercase     						=self.tokenizer(
                __lowercase  , padding='max_length'  , max_length=self.tokenizer.model_max_length  , truncation=__lowercase  , return_tensors='pt'  , )

            __lowercase     						=self.text_encoder(text_input.input_ids.to(self.device ) )[0]

            __lowercase     						=self.tokenizer(
                ''  , padding='max_length'  , max_length=self.tokenizer.model_max_length  , truncation=__lowercase  , return_tensors='pt'  , )

            __lowercase     						=self.text_encoder(uncond_input.input_ids.to(self.device ) )[0]

            return torch.cat([uncond_embedding, text_embedding] )







       def __call__(							self     :	Any  , __lowercase     :	Image.Image  , __lowercase     :	str  , __lowercase     :	float = 0.3  , __lowercase     :	float = 0.6  , __lowercase     :	float = 0.5  , __lowercase     :	int = 42  , __lowercase     :	int = 50  , __lowercase     :	float = 7.5  , ):







            """simple docstring"""







            __lowercase     						=steps - int(kmin * steps )
            __lowercase     						=steps - int(kmax * steps )

            __lowercase     						=self.prep_text(__lowercase )

            self.scheduler.set_timesteps(__lowercase )

            __lowercase						,       __lowercase     						=img.size
            __lowercase     						=self.encode(__lowercase )

            torch.manual_seed(__lowercase )
            __lowercase     						=torch.randn(
                (1, self.unet.config.in_channels, height // 8, width // 8)  , ).to(self.device )

            __lowercase     						=self.scheduler.add_noise(
                __lowercase  , __lowercase  , timesteps=self.scheduler.timesteps[tmax]  , )

            __lowercase     						=torch.cat([latents] * 2 )

            __lowercase     						=self.scheduler.scale_model_input(__lowercase  , self.scheduler.timesteps[tmax] )

            with torch.no_grad():
                 __lowercase     						=self.unet(
                     __lowercase  , self.scheduler.timesteps[tmax]  , encoder_hidden_states=__lowercase  , ).sample

            __lowercase						,       __lowercase     						=pred.chunk(2 )
            __lowercase     						=pred_uncond + guidance_scale * (pred_text - pred_uncond)

            __lowercase     						=self.scheduler.step(__lowercase  , self.scheduler.timesteps[tmax]  , __lowercase ).prev_sample

            for i, t in enumerate(tqdm(self.scheduler.timesteps ) ):
                 if i > tmax:
                      if i < tmin:  # layout generation phase
                           __lowercase     						=self.scheduler.add_noise(
                               __lowercase  , __lowercase  , timesteps=__lowercase  , )

                           __lowercase     						=(mix_factor * latents) + (
                               1 - mix_factor
                           ) * orig_latents  # interpolating between layout noise and conditionally generated noise to preserve layout sematics
                           __lowercase     						=torch.cat([input] * 2 )

                      else:  # content generation phase
                           __lowercase     						=torch.cat([latents] * 2 )

                      __lowercase     						=self.scheduler.scale_model_input(__lowercase  , __lowercase )

                      with torch.no_grad():
                           __lowercase     						=self.unet(
                               __lowercase  , __lowercase  , encoder_hidden_states=__lowercase  , ).sample

                      __lowercase						,       __lowercase     						=pred.chunk(2 )
                      __lowercase     						=pred_uncond + guidance_scale * (pred_text - pred_uncond)

                      __lowercase     						=self.scheduler.step(__lowercase  , __lowercase  , __lowercase ).prev_sample

            return self.decode(__lowercase )


