
from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
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 import TF_MODEL_FOR_PRETRAINING_MAPPING
     from transformers.models.albert.modeling_tf_albert import (
         TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFAlbertForMaskedLM,
         TFAlbertForMultipleChoice,
         TFAlbertForPreTraining,
         TFAlbertForQuestionAnswering,
         TFAlbertForSequenceClassification,
         TFAlbertForTokenClassification,
         TFAlbertModel,
     )





class   __UpperCamelCase				:






     """simple docstring"""







     def __init__(							self    :    Tuple      ,		_A    :    Union[str, Any]      ,		_A    :    List[str]=13      ,		_A    :    List[str]=7      ,		_A    :    Any=True      ,		_A    :    Any=True      ,		_A    :    Any=True      ,		_A    :    int=True      ,		_A    :    Tuple=99      ,		_A    :    List[Any]=16      ,		_A    :    Union[str, Any]=32      ,		_A    :    Union[str, Any]=2      ,		_A    :    int=4      ,		_A    :    Any=37      ,		_A    :    Dict="gelu"      ,		_A    :    Any=0.1      ,		_A    :    Dict=0.1      ,		_A    :    List[str]=512      ,		_A    :    Any=16      ,		_A    :    List[Any]=2      ,		_A    :    Dict=0.02      ,		_A    :    Optional[Any]=3      ,		_A    :    str=4      ,		_A    :    List[str]=None      ,		):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						parent
       __SCREAMING_SNAKE_CASE     :					List[str]			=						13
       __SCREAMING_SNAKE_CASE     :					str			=						7
       __SCREAMING_SNAKE_CASE     :					Any			=						True
       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						True
       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						True
       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						True
       __SCREAMING_SNAKE_CASE     :					str			=						99
       __SCREAMING_SNAKE_CASE     :					int			=						16
       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						32
       __SCREAMING_SNAKE_CASE     :					Any			=						2
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						4
       __SCREAMING_SNAKE_CASE     :					int			=						37
       __SCREAMING_SNAKE_CASE     :					Any			=						'''gelu'''
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						0.1
       __SCREAMING_SNAKE_CASE     :					Any			=						0.1
       __SCREAMING_SNAKE_CASE     :					str			=						512
       __SCREAMING_SNAKE_CASE     :					Tuple			=						16
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						2
       __SCREAMING_SNAKE_CASE     :					Tuple			=						0.02
       __SCREAMING_SNAKE_CASE     :					str			=						3
       __SCREAMING_SNAKE_CASE     :					Any			=						4
       __SCREAMING_SNAKE_CASE     :					Any			=						None



     def 				UpperCAmelCase__    (							self    :    Tuple				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						ids_tensor([self.batch_size, self.seq_length]      ,		self.vocab_size				)

       __SCREAMING_SNAKE_CASE     :					Dict			=						None
       if self.use_input_mask:
         __SCREAMING_SNAKE_CASE     :					List[str]			=						random_attention_mask([self.batch_size, self.seq_length]				)

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

       __SCREAMING_SNAKE_CASE     :					int			=						None
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						None
       __SCREAMING_SNAKE_CASE     :					str			=						None
       if self.use_labels:
         __SCREAMING_SNAKE_CASE     :					str			=						ids_tensor([self.batch_size]      ,		self.type_sequence_label_size				)
         __SCREAMING_SNAKE_CASE     :					Optional[int]			=						ids_tensor([self.batch_size, self.seq_length]      ,		self.num_labels				)
         __SCREAMING_SNAKE_CASE     :					Tuple			=						ids_tensor([self.batch_size]      ,		self.num_choices				)

       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						AlbertConfig(
           vocab_size=self.vocab_size      ,		hidden_size=self.hidden_size      ,		embedding_size=self.embedding_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      ,		initializer_range=self.initializer_range      ,		)

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



     def 				UpperCAmelCase__    (							self    :    int      ,		_A    :    Any      ,		_A    :    Optional[int]      ,		_A    :    Union[str, Any]      ,		_A    :    Dict      ,		_A    :    int      ,		_A    :    Optional[int]      ,		_A    :    Union[str, Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Any			=						TFAlbertModel(config=lowerCAmelCase__				)
       # inputs = {'input_ids': input_ids,
       #           'attention_mask': input_mask,
       #           'token_type_ids': token_type_ids}
       # sequence_output, pooled_output = model(**inputs)
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __SCREAMING_SNAKE_CASE     :					Tuple			=						model(lowerCAmelCase__				)

       __SCREAMING_SNAKE_CASE     :					int			=						[input_ids, input_mask]
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						model(lowerCAmelCase__				)

       __SCREAMING_SNAKE_CASE     :					Any			=						model(lowerCAmelCase__				)

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



     def 				UpperCAmelCase__    (							self    :    Tuple      ,		_A    :    Union[str, Any]      ,		_A    :    List[Any]      ,		_A    :    Any      ,		_A    :    List[Any]      ,		_A    :    Optional[Any]      ,		_A    :    Union[str, Any]      ,		_A    :    Optional[int]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Dict			=						self.num_labels
       __SCREAMING_SNAKE_CASE     :					Tuple			=						TFAlbertForPreTraining(config=lowerCAmelCase__				)
       __SCREAMING_SNAKE_CASE     :					List[str]			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						model(lowerCAmelCase__				)
       self.parent.assertEqual(result.prediction_logits.shape      ,		(self.batch_size, self.seq_length, self.vocab_size)				)
       self.parent.assertEqual(result.sop_logits.shape      ,		(self.batch_size, self.num_labels)				)



     def 				UpperCAmelCase__    (							self    :    Optional[int]      ,		_A    :    str      ,		_A    :    Union[str, Any]      ,		_A    :    List[Any]      ,		_A    :    Any      ,		_A    :    Optional[int]      ,		_A    :    Dict      ,		_A    :    str				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						TFAlbertForMaskedLM(config=lowerCAmelCase__				)
       __SCREAMING_SNAKE_CASE     :					List[str]			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __SCREAMING_SNAKE_CASE     :					str			=						model(lowerCAmelCase__				)
       self.parent.assertEqual(result.logits.shape      ,		(self.batch_size, self.seq_length, self.vocab_size)				)



     def 				UpperCAmelCase__    (							self    :    Optional[int]      ,		_A    :    int      ,		_A    :    List[str]      ,		_A    :    str      ,		_A    :    str      ,		_A    :    Optional[Any]      ,		_A    :    Optional[int]      ,		_A    :    List[str]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Tuple			=						self.num_labels
       __SCREAMING_SNAKE_CASE     :					str			=						TFAlbertForSequenceClassification(config=lowerCAmelCase__				)
       __SCREAMING_SNAKE_CASE     :					str			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __SCREAMING_SNAKE_CASE     :					Tuple			=						model(lowerCAmelCase__				)
       self.parent.assertEqual(result.logits.shape      ,		(self.batch_size, self.num_labels)				)



     def 				UpperCAmelCase__    (							self    :    str      ,		_A    :    Union[str, Any]      ,		_A    :    Union[str, Any]      ,		_A    :    List[str]      ,		_A    :    int      ,		_A    :    Dict      ,		_A    :    List[str]      ,		_A    :    Dict				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[Any]			=						TFAlbertForQuestionAnswering(config=lowerCAmelCase__				)
       __SCREAMING_SNAKE_CASE     :					Any			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
       __SCREAMING_SNAKE_CASE     :					int			=						model(lowerCAmelCase__				)
       self.parent.assertEqual(result.start_logits.shape      ,		(self.batch_size, self.seq_length)				)
       self.parent.assertEqual(result.end_logits.shape      ,		(self.batch_size, self.seq_length)				)



     def 				UpperCAmelCase__    (							self    :    Optional[Any]      ,		_A    :    Dict      ,		_A    :    List[Any]      ,		_A    :    List[str]      ,		_A    :    Any      ,		_A    :    Any      ,		_A    :    Union[str, Any]      ,		_A    :    int				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						self.num_choices
       __SCREAMING_SNAKE_CASE     :					Any			=						TFAlbertForMultipleChoice(config=lowerCAmelCase__				)
       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						tf.tile(tf.expand_dims(lowerCAmelCase__      ,		1				)      ,		(1, self.num_choices, 1)				)
       __SCREAMING_SNAKE_CASE     :					str			=						tf.tile(tf.expand_dims(lowerCAmelCase__      ,		1				)      ,		(1, self.num_choices, 1)				)
       __SCREAMING_SNAKE_CASE     :					int			=						tf.tile(tf.expand_dims(lowerCAmelCase__      ,		1				)      ,		(1, self.num_choices, 1)				)
       __SCREAMING_SNAKE_CASE     :					str			=						{
           '''input_ids''': multiple_choice_inputs_ids,
           '''attention_mask''': multiple_choice_input_mask,
           '''token_type_ids''': multiple_choice_token_type_ids,
       }
       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						model(lowerCAmelCase__				)
       self.parent.assertListEqual(list(result['''logits'''].shape				)      ,		[self.batch_size, self.num_choices]				)



     def 				UpperCAmelCase__    (							self    :    Any      ,		_A    :    Optional[Any]      ,		_A    :    Dict      ,		_A    :    str      ,		_A    :    Optional[Any]      ,		_A    :    Tuple      ,		_A    :    Union[str, Any]      ,		_A    :    Optional[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						self.num_labels
       __SCREAMING_SNAKE_CASE     :					Tuple			=						TFAlbertForTokenClassification(config=lowerCAmelCase__				)
       __SCREAMING_SNAKE_CASE     :					List[str]			=						{
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }
       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						model(lowerCAmelCase__				)
       self.parent.assertListEqual(list(result['''logits'''].shape				)      ,		[self.batch_size, self.seq_length, self.num_labels]				)




     def 				UpperCAmelCase__    (							self    :    Tuple				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[Any]			=						self.prepare_config_and_inputs()
       (
           (
           __SCREAMING_SNAKE_CASE 
       ),						(
           __SCREAMING_SNAKE_CASE 
       ),						(
           __SCREAMING_SNAKE_CASE 
       ),						(
           __SCREAMING_SNAKE_CASE 
       ),						(
           __SCREAMING_SNAKE_CASE 
       ),						(
           __SCREAMING_SNAKE_CASE 
       ),						(
           __SCREAMING_SNAKE_CASE 
       ),						 
       )     :					List[str]			=						config_and_inputs
       __SCREAMING_SNAKE_CASE     :					int			=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_tf
class   __UpperCamelCase				(    lowercase__				,		lowercase__				,		unittest.TestCase   ):






     """simple docstring"""







     lowerCAmelCase_     =	(
         (
             TFAlbertModel,
             TFAlbertForPreTraining,
             TFAlbertForMaskedLM,
             TFAlbertForSequenceClassification,
             TFAlbertForQuestionAnswering,
             TFAlbertForTokenClassification,
             TFAlbertForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     lowerCAmelCase_     =	(
         {
             '''feature-extraction''': TFAlbertModel,
             '''fill-mask''': TFAlbertForMaskedLM,
             '''question-answering''': TFAlbertForQuestionAnswering,
             '''text-classification''': TFAlbertForSequenceClassification,
             '''token-classification''': TFAlbertForTokenClassification,
             '''zero-shot''': TFAlbertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     lowerCAmelCase_     =	False
     lowerCAmelCase_     =	False



     def 				UpperCAmelCase__    (							self    :    Union[str, Any]      ,		_A    :    Union[str, Any]      ,		_A    :    List[Any]      ,		_A    :    Tuple=False				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					List[Any]			=						super()._prepare_for_class(lowerCAmelCase__      ,		lowerCAmelCase__      ,		return_labels=lowerCAmelCase__				)

       if return_labels:
         if model_class in get_values(lowerCAmelCase__				):
           __SCREAMING_SNAKE_CASE     :					int			=						tf.zeros(self.model_tester.batch_size      ,		dtype=tf.intaa				)

       return inputs_dict



     def 				UpperCAmelCase__    (							self    :    Dict				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					str			=						TFAlbertModelTester(self				)
       __SCREAMING_SNAKE_CASE     :					Dict			=						ConfigTester(self      ,		config_class=lowerCAmelCase__      ,		hidden_size=37				)



     def 				UpperCAmelCase__    (							self    :    List[Any]				):







       """simple docstring"""




       self.config_tester.run_common_tests()



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







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_model(*lowerCAmelCase__				)



     def 				UpperCAmelCase__    (							self    :    int				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Any			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_pretraining(*lowerCAmelCase__				)



     def 				UpperCAmelCase__    (							self    :    Any				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_masked_lm(*lowerCAmelCase__				)



     def 				UpperCAmelCase__    (							self    :    int				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					int			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_multiple_choice(*lowerCAmelCase__				)



     def 				UpperCAmelCase__    (							self    :    Any				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					str			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_sequence_classification(*lowerCAmelCase__				)



     def 				UpperCAmelCase__    (							self    :    Optional[int]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_question_answering(*lowerCAmelCase__				)




     @slow
     def 				UpperCAmelCase__    (							self    :    Tuple				):







       """simple docstring"""




       for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __SCREAMING_SNAKE_CASE     :					Any			=						TFAlbertModel.from_pretrained(lowerCAmelCase__				)
         self.assertIsNotNone(lowerCAmelCase__				)






@require_tf
class   __UpperCamelCase				(    unittest.TestCase   ):






     """simple docstring"""







     @slow
     def 				UpperCAmelCase__    (							self    :    int				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					str			=						TFAlbertForPreTraining.from_pretrained('''albert-base-v2'''				)
       __SCREAMING_SNAKE_CASE     :					int			=						tf.constant([[0, 1, 2, 3, 4, 5]]				)
       __SCREAMING_SNAKE_CASE     :					List[str]			=						model(lowerCAmelCase__				)[0]

       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						[1, 6, 3_0000]
       self.assertEqual(output.shape      ,		lowerCAmelCase__				)

       __SCREAMING_SNAKE_CASE     :					Optional[int]			=						tf.constant(
           [
               [
                   [4.59_56_68, 0.74_46_27_54, -1.81_81_47],
                   [4.5_95_43_47, 0.7_45_41_84, -1.8_18_82_58],
                   [4.5_95_49_05, 0.7_44_82_35, -1.8_18_23_16],
               ]
           ]				)
       tf.debugging.assert_near(output[:, :3, :3]      ,		lowerCAmelCase__      ,		atol=1e-4				)


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   __UpperCamelCase				(    lowerCAmelCase__   ):






     """simple docstring"""







     def __init__(							self    :    Optional[int]      ,		_A    :    AutoencoderKL      ,		_A    :    CLIPTextModel      ,		_A    :    CLIPTokenizer      ,		_A    :    UNetaDConditionModel      ,		_A    :    Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler]      ,		):







       """simple docstring"""




       super().__init__()

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



     def 				UpperCAmelCase__    (							self    :    Optional[int]      ,		_A    :    Dict				):







       """simple docstring"""




       with torch.no_grad():
         __SCREAMING_SNAKE_CASE     :					Dict			=						self.vae.encode(tfms.ToTensor()(_A				).unsqueeze(0				).to(self.device				) * 2 - 1				)
         __SCREAMING_SNAKE_CASE     :					str			=						0.1_82_15 * latent.latent_dist.sample()
       return latent



     def 				UpperCAmelCase__    (							self    :    Any      ,		_A    :    List[Any]				):







       """simple docstring"""




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



     def 				UpperCAmelCase__    (							self    :    Dict      ,		_A    :    List[Any]				):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Any			=						self.tokenizer(
           _A      ,		padding='''max_length'''      ,		max_length=self.tokenizer.model_max_length      ,		truncation=_A      ,		return_tensors='''pt'''      ,		)

       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						self.text_encoder(text_input.input_ids.to(self.device				)				)[0]

       __SCREAMING_SNAKE_CASE     :					Tuple			=						self.tokenizer(
           ''''''      ,		padding='''max_length'''      ,		max_length=self.tokenizer.model_max_length      ,		truncation=_A      ,		return_tensors='''pt'''      ,		)

       __SCREAMING_SNAKE_CASE     :					str			=						self.text_encoder(uncond_input.input_ids.to(self.device				)				)[0]

       return torch.cat([uncond_embedding, text_embedding]				)




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







       """simple docstring"""




       __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						steps - int(kmin * steps				)
       __SCREAMING_SNAKE_CASE     :					List[Any]			=						steps - int(kmax * steps				)

       __SCREAMING_SNAKE_CASE     :					List[str]			=						self.prep_text(_A				)

       self.scheduler.set_timesteps(_A				)

       __SCREAMING_SNAKE_CASE,						__SCREAMING_SNAKE_CASE     :					Dict			=						img.size
       __SCREAMING_SNAKE_CASE     :					Dict			=						self.encode(_A				)

       torch.manual_seed(_A				)
       __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						torch.randn(
           (1, self.unet.config.in_channels, height // 8, width // 8)      ,		).to(self.device				)

       __SCREAMING_SNAKE_CASE     :					Dict			=						self.scheduler.add_noise(
           _A      ,		_A      ,		timesteps=self.scheduler.timesteps[tmax]      ,		)

       __SCREAMING_SNAKE_CASE     :					List[str]			=						torch.cat([latents] * 2				)

       __SCREAMING_SNAKE_CASE     :					List[Any]			=						self.scheduler.scale_model_input(_A      ,		self.scheduler.timesteps[tmax]				)

       with torch.no_grad():
         __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						self.unet(
             _A      ,		self.scheduler.timesteps[tmax]      ,		encoder_hidden_states=_A      ,		).sample

       __SCREAMING_SNAKE_CASE,						__SCREAMING_SNAKE_CASE     :					List[str]			=						pred.chunk(2				)
       __SCREAMING_SNAKE_CASE     :					str			=						pred_uncond + guidance_scale * (pred_text - pred_uncond)

       __SCREAMING_SNAKE_CASE     :					Any			=						self.scheduler.step(_A      ,		self.scheduler.timesteps[tmax]      ,		_A				).prev_sample

       for i, t in enumerate(tqdm(self.scheduler.timesteps				)				):
         if i > tmax:
           if i < tmin:  # layout generation phase
             __SCREAMING_SNAKE_CASE     :					List[str]			=						self.scheduler.add_noise(
                 _A      ,		_A      ,		timesteps=_A      ,		)

             __SCREAMING_SNAKE_CASE     :					Union[str, Any]			=						(mix_factor * latents) + (
                 1 - mix_factor
             ) * orig_latents  # interpolating between layout noise and conditionally generated noise to preserve layout sematics
             __SCREAMING_SNAKE_CASE     :					Optional[Any]			=						torch.cat([input] * 2				)

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

           __SCREAMING_SNAKE_CASE     :					str			=						self.scheduler.scale_model_input(_A      ,		_A				)

           with torch.no_grad():
             __SCREAMING_SNAKE_CASE     :					int			=						self.unet(
                 _A      ,		_A      ,		encoder_hidden_states=_A      ,		).sample

           __SCREAMING_SNAKE_CASE,						__SCREAMING_SNAKE_CASE     :					Optional[int]			=						pred.chunk(2				)
           __SCREAMING_SNAKE_CASE     :					Optional[int]			=						pred_uncond + guidance_scale * (pred_text - pred_uncond)

           __SCREAMING_SNAKE_CASE     :					Any			=						self.scheduler.step(_A      ,		_A      ,		_A				).prev_sample

       return self.decode(_A				)
