




'''simple docstring'''


import unittest

from transformers import XLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, 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 (
				    XLMForMultipleChoice,
				    XLMForQuestionAnswering,
				    XLMForQuestionAnsweringSimple,
				    XLMForSequenceClassification,
				    XLMForTokenClassification,
				    XLMModel,
				    XLMWithLMHeadModel,
				)
				from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST

class      A       :







							def __init__(    self			:    List[str]				,     lowerCAmelCase_			:    Tuple				,     lowerCAmelCase_			:    Dict=13				,     lowerCAmelCase_			:    int=7				,     lowerCAmelCase_			:    Any=True				,     lowerCAmelCase_			:    Any=True				,     lowerCAmelCase_			:    Any=True				,     lowerCAmelCase_			:    List[str]=True				,     lowerCAmelCase_			:    List[str]=True				,     lowerCAmelCase_			:    Union[str, Any]=False				,     lowerCAmelCase_			:    Optional[Any]=False				,     lowerCAmelCase_			:    int=False				,     lowerCAmelCase_			:    Tuple=2				,     lowerCAmelCase_			:    Optional[Any]=99				,     lowerCAmelCase_			:    Tuple=0				,     lowerCAmelCase_			:    List[Any]=32				,     lowerCAmelCase_			:    List[Any]=5				,     lowerCAmelCase_			:    str=4				,     lowerCAmelCase_			:    List[Any]=0.1				,     lowerCAmelCase_			:    Optional[int]=0.1				,     lowerCAmelCase_			:    Tuple=5_12				,     lowerCAmelCase_			:    Union[str, Any]=2				,     lowerCAmelCase_			:    List[Any]=0.0_2				,     lowerCAmelCase_			:    Optional[Any]=2				,     lowerCAmelCase_			:    Tuple=4				,     lowerCAmelCase_			:    str="last"				,     lowerCAmelCase_			:    List[str]=True				,     lowerCAmelCase_			:    int=None				,     lowerCAmelCase_			:    Tuple=0				,     )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							parent
														_a      	=							batch_size
														_a      	=							seq_length
														_a      	=							is_training
														_a      	=							use_input_lengths
														_a      	=							use_token_type_ids
														_a      	=							use_labels
														_a      	=							gelu_activation
														_a      	=							sinusoidal_embeddings
														_a      	=							causal
														_a      	=							asm
														_a      	=							n_langs
														_a      	=							vocab_size
														_a      	=							n_special
														_a      	=							hidden_size
														_a      	=							num_hidden_layers
														_a      	=							num_attention_heads
														_a      	=							hidden_dropout_prob
														_a      	=							attention_probs_dropout_prob
														_a      	=							max_position_embeddings
														_a      	=							type_sequence_label_size
														_a      	=							initializer_range
														_a      	=							num_labels
														_a      	=							num_choices
														_a      	=							summary_type
														_a      	=							use_proj
														_a      	=							scope
														_a      	=							bos_token_id







							def   __lowerCAmelCase     (    self			:    Tuple   )     ->				Union[str, Any]:


														"""simple docstring"""


														_a      	=							ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size   )
														_a      	=							random_attention_mask([self.batch_size, self.seq_length]   )

														_a      	=							None
														if self.use_input_lengths:
																					_a      	=							(
																					    ids_tensor([self.batch_size]				,     vocab_size=2   ) + self.seq_length - 2
																					)  # small variation of seq_length

														_a      	=							None
														if self.use_token_type_ids:
																					_a      	=							ids_tensor([self.batch_size, self.seq_length]				,     self.n_langs   )

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

														_a      	=							self.get_config()

														return (
														    config,
														    input_ids,
														    token_type_ids,
														    input_lengths,
														    sequence_labels,
														    token_labels,
														    is_impossible_labels,
														    choice_labels,
														    input_mask,
														)







							def   __lowerCAmelCase     (    self			:    List[str]   )     ->				Dict:


														"""simple docstring"""


														return XLMConfig(
														    vocab_size=self.vocab_size				,     n_special=self.n_special				,     emb_dim=self.hidden_size				,     n_layers=self.num_hidden_layers				,     n_heads=self.num_attention_heads				,     dropout=self.hidden_dropout_prob				,     attention_dropout=self.attention_probs_dropout_prob				,     gelu_activation=self.gelu_activation				,     sinusoidal_embeddings=self.sinusoidal_embeddings				,     asm=self.asm				,     causal=self.causal				,     n_langs=self.n_langs				,     max_position_embeddings=self.max_position_embeddings				,     initializer_range=self.initializer_range				,     summary_type=self.summary_type				,     use_proj=self.use_proj				,     num_labels=self.num_labels				,     bos_token_id=self.bos_token_id				,     )







							def   __lowerCAmelCase     (    self			:    Union[str, Any]				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    int				,     )     ->				str:


														"""simple docstring"""


														_a      	=							XLMModel(config=lowerCAmelCase_   )
														model.to(lowerCAmelCase_   )
														model.eval()
														_a      	=							model(lowerCAmelCase_				,     lengths=lowerCAmelCase_				,     langs=lowerCAmelCase_   )
														_a      	=							model(lowerCAmelCase_				,     langs=lowerCAmelCase_   )
														_a      	=							model(lowerCAmelCase_   )
														self.parent.assertEqual(result.last_hidden_state.shape				,     (self.batch_size, self.seq_length, self.hidden_size)   )







							def   __lowerCAmelCase     (    self			:    int				,     lowerCAmelCase_			:    str				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    str				,     )     ->				int:


														"""simple docstring"""


														_a      	=							XLMWithLMHeadModel(lowerCAmelCase_   )
														model.to(lowerCAmelCase_   )
														model.eval()

														_a      	=							model(lowerCAmelCase_				,     token_type_ids=lowerCAmelCase_				,     labels=lowerCAmelCase_   )
														self.parent.assertEqual(result.loss.shape				,     ()   )
														self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size)   )







							def   __lowerCAmelCase     (    self			:    Union[str, Any]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Tuple				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    Dict				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    str				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Optional[int]				,     )     ->				int:


														"""simple docstring"""


														_a      	=							XLMForQuestionAnsweringSimple(lowerCAmelCase_   )
														model.to(lowerCAmelCase_   )
														model.eval()

														_a      	=							model(lowerCAmelCase_   )

														_a      	=							model(lowerCAmelCase_				,     start_positions=lowerCAmelCase_				,     end_positions=lowerCAmelCase_   )
														_a      	=							outputs
														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   __lowerCAmelCase     (    self			:    str				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    Any				,     )     ->				Tuple:


														"""simple docstring"""


														_a      	=							XLMForQuestionAnswering(lowerCAmelCase_   )
														model.to(lowerCAmelCase_   )
														model.eval()

														_a      	=							model(lowerCAmelCase_   )

														_a      	=							model(
														    lowerCAmelCase_				,     start_positions=lowerCAmelCase_				,     end_positions=lowerCAmelCase_				,     cls_index=lowerCAmelCase_				,     is_impossible=lowerCAmelCase_				,     p_mask=lowerCAmelCase_				,     )

														_a      	=							model(
														    lowerCAmelCase_				,     start_positions=lowerCAmelCase_				,     end_positions=lowerCAmelCase_				,     cls_index=lowerCAmelCase_				,     is_impossible=lowerCAmelCase_				,     )

														((_a)	,				)      	=							result_with_labels.to_tuple()

														_a      	=							model(lowerCAmelCase_				,     start_positions=lowerCAmelCase_				,     end_positions=lowerCAmelCase_   )

														((_a)	,				)      	=							result_with_labels.to_tuple()

														self.parent.assertEqual(result_with_labels.loss.shape				,     ()   )
														self.parent.assertEqual(result.start_top_log_probs.shape				,     (self.batch_size, model.config.start_n_top)   )
														self.parent.assertEqual(result.start_top_index.shape				,     (self.batch_size, model.config.start_n_top)   )
														self.parent.assertEqual(
														    result.end_top_log_probs.shape				,     (self.batch_size, model.config.start_n_top * model.config.end_n_top)   )
														self.parent.assertEqual(
														    result.end_top_index.shape				,     (self.batch_size, model.config.start_n_top * model.config.end_n_top)   )
														self.parent.assertEqual(result.cls_logits.shape				,     (self.batch_size,)   )







							def   __lowerCAmelCase     (    self			:    str				,     lowerCAmelCase_			:    Dict				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    str				,     lowerCAmelCase_			:    Tuple				,     )     ->				Dict:


														"""simple docstring"""


														_a      	=							XLMForSequenceClassification(lowerCAmelCase_   )
														model.to(lowerCAmelCase_   )
														model.eval()

														_a      	=							model(lowerCAmelCase_   )
														_a      	=							model(lowerCAmelCase_				,     labels=lowerCAmelCase_   )
														self.parent.assertEqual(result.loss.shape				,     ()   )
														self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.type_sequence_label_size)   )







							def   __lowerCAmelCase     (    self			:    List[Any]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    str				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    str				,     )     ->				Optional[Any]:


														"""simple docstring"""


														_a      	=							self.num_labels
														_a      	=							XLMForTokenClassification(lowerCAmelCase_   )
														model.to(lowerCAmelCase_   )
														model.eval()

														_a      	=							model(lowerCAmelCase_				,     attention_mask=lowerCAmelCase_				,     labels=lowerCAmelCase_   )
														self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.num_labels)   )







							def   __lowerCAmelCase     (    self			:    Union[str, Any]				,     lowerCAmelCase_			:    str				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Tuple				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    Tuple				,     lowerCAmelCase_			:    Union[str, Any]				,     )     ->				Any:


														"""simple docstring"""


														_a      	=							self.num_choices
														_a      	=							XLMForMultipleChoice(config=lowerCAmelCase_   )
														model.to(lowerCAmelCase_   )
														model.eval()
														_a      	=							input_ids.unsqueeze(1   ).expand(-1				,     self.num_choices				,     -1   ).contiguous()
														_a      	=							token_type_ids.unsqueeze(1   ).expand(-1				,     self.num_choices				,     -1   ).contiguous()
														_a      	=							input_mask.unsqueeze(1   ).expand(-1				,     self.num_choices				,     -1   ).contiguous()
														_a      	=							model(
														    lowerCAmelCase_				,     attention_mask=lowerCAmelCase_				,     token_type_ids=lowerCAmelCase_				,     labels=lowerCAmelCase_				,     )
														self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.num_choices)   )







							def   __lowerCAmelCase     (    self			:    int   )     ->				Optional[int]:


														"""simple docstring"""


														_a      	=							self.prepare_config_and_inputs()
														(
														    (
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				
														)      	=							config_and_inputs
														_a      	=							{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''lengths''': input_lengths}
														return config, inputs_dict



@require_torch
class      A       (						_a	,_a	,_a	,unittest.TestCase       ):
							lowercase_							     =							(
							    (
							        XLMModel,
							        XLMWithLMHeadModel,
							        XLMForQuestionAnswering,
							        XLMForSequenceClassification,
							        XLMForQuestionAnsweringSimple,
							        XLMForTokenClassification,
							        XLMForMultipleChoice,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase_							     =							(
							    (XLMWithLMHeadModel,) if is_torch_available() else ()
							)  # TODO (PVP): Check other models whether language generation is also applicable
							lowercase_							     =							(
							    {
							        'feature-extraction': XLMModel,
							        'fill-mask': XLMWithLMHeadModel,
							        'question-answering': XLMForQuestionAnsweringSimple,
							        'text-classification': XLMForSequenceClassification,
							        'text-generation': XLMWithLMHeadModel,
							        'token-classification': XLMForTokenClassification,
							        'zero-shot': XLMForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)







							def   __lowerCAmelCase     (    self			:    Optional[int]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    Tuple   )     ->				Optional[int]:


														"""simple docstring"""


														if (
														    pipeline_test_casse_name == "QAPipelineTests"
														    and tokenizer_name is not None
														    and not tokenizer_name.endswith('''Fast'''   )
														):
																					# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
																					# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
																					# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
																					return True

														return False







							def   __lowerCAmelCase     (    self			:    List[str]				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    int=False   )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							super()._prepare_for_class(lowerCAmelCase_				,     lowerCAmelCase_				,     return_labels=lowerCAmelCase_   )

														if return_labels:
																					if model_class.__name__ == "XLMForQuestionAnswering":
																												_a      	=							torch.zeros(
																												    self.model_tester.batch_size				,     dtype=torch.long				,     device=lowerCAmelCase_   )
																												_a      	=							torch.zeros(
																												    self.model_tester.batch_size				,     dtype=torch.long				,     device=lowerCAmelCase_   )

														return inputs_dict







							def   __lowerCAmelCase     (    self			:    Any   )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							XLMModelTester(self   )
														_a      	=							ConfigTester(self				,     config_class=lowerCAmelCase_				,     emb_dim=37   )







							def   __lowerCAmelCase     (    self			:    Optional[int]   )     ->				str:


														"""simple docstring"""


														self.config_tester.run_common_tests()







							def   __lowerCAmelCase     (    self			:    List[Any]   )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_xlm_model(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Optional[int]   )     ->				List[str]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_xlm_lm_head(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Optional[Any]   )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_xlm_simple_qa(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Dict   )     ->				Optional[int]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_xlm_qa(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Union[str, Any]   )     ->				int:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_xlm_sequence_classif(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    List[str]   )     ->				List[str]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_xlm_token_classif(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    List[Any]   )     ->				int:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_xlm_for_multiple_choice(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Any				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    str				,     lowerCAmelCase_			:    str=False				,     lowerCAmelCase_			:    str=1   )     ->				int:


														"""simple docstring"""


														self.assertIsInstance(lowerCAmelCase_				,     lowerCAmelCase_   )
														self.assertListEqual(
														    [isinstance(lowerCAmelCase_				,     lowerCAmelCase_   ) for iter_attentions in attentions]				,     [True] * len(lowerCAmelCase_   )   )
														self.assertEqual(len(lowerCAmelCase_   )				,     (max_length - min_length) * num_beam_groups   )

														for idx, iter_attentions in enumerate(lowerCAmelCase_   ):
																					# adds PAD dummy token
																					_a      	=							min_length + idx + 1
																					_a      	=							min_length + idx + 1

																					_a      	=							(
																					    batch_size * num_beam_groups,
																					    config.num_attention_heads,
																					    tgt_len,
																					    src_len,
																					)
																					# check attn size
																					self.assertListEqual(
																					    [layer_attention.shape for layer_attention in iter_attentions]				,     [expected_shape] * len(lowerCAmelCase_   )   )







							def   __lowerCAmelCase     (    self			:    str				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    Tuple				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    int=False				,     lowerCAmelCase_			:    List[str]=1   )     ->				Optional[Any]:


														"""simple docstring"""


														self.assertIsInstance(lowerCAmelCase_				,     lowerCAmelCase_   )
														self.assertListEqual(
														    [isinstance(lowerCAmelCase_				,     lowerCAmelCase_   ) for iter_hidden_states in hidden_states]				,     [True] * len(lowerCAmelCase_   )				,     )
														self.assertEqual(len(lowerCAmelCase_   )				,     (max_length - min_length) * num_beam_groups   )

														for idx, iter_hidden_states in enumerate(lowerCAmelCase_   ):
																					# adds PAD dummy token
																					_a      	=							min_length + idx + 1
																					_a      	=							(batch_size * num_beam_groups, seq_len, config.hidden_size)
																					# check hidden size
																					self.assertListEqual(
																					    [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states]				,     [expected_shape] * len(lowerCAmelCase_   )				,     )
														pass







							@slow
							def   __lowerCAmelCase     (    self			:    Optional[Any]   )     ->				List[Any]:


														"""simple docstring"""


														for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					_a      	=							XLMModel.from_pretrained(lowerCAmelCase_   )
																					self.assertIsNotNone(lowerCAmelCase_   )



@require_torch
class      A       (						unittest.TestCase       ):







							@slow
							def   __lowerCAmelCase     (    self			:    int   )     ->				Optional[int]:


														"""simple docstring"""


														_a      	=							XLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048'''   )
														model.to(lowerCAmelCase_   )
														_a      	=							torch.tensor([[14, 4_47]]				,     dtype=torch.long				,     device=lowerCAmelCase_   )  # the president
														_a      	=							[
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														    14,
														    4_47,
														]  # the president the president the president the president the president the president the president the president the president the president
														# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
														_a      	=							model.generate(lowerCAmelCase_				,     do_sample=lowerCAmelCase_   )
														self.assertListEqual(output_ids[0].cpu().numpy().tolist()				,     lowerCAmelCase_   )







'''simple docstring'''







import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()

class 	_snake_case(			UpperCAmelCase       , UpperCAmelCase       , UpperCAmelCase       , unittest.TestCase      ):
 __snake_case:							Union[str, Any]						=					StableUnCLIPPipeline
 __snake_case:							List[str]						=					TEXT_TO_IMAGE_PARAMS
 __snake_case:							str						=					TEXT_TO_IMAGE_BATCH_PARAMS
 __snake_case:							Dict						=					TEXT_TO_IMAGE_IMAGE_PARAMS
 __snake_case:							List[str]						=					TEXT_TO_IMAGE_IMAGE_PARAMS

 # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
 __snake_case:							str						=					False






 def 						_UpperCamelCase (self       :  Optional[int]    )       ->   int:






    """simple docstring"""


    A__       =     32
    A__       =     embedder_hidden_size

    # prior components

    torch.manual_seed(0    )
    A__       =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

    torch.manual_seed(0    )
    A__       =     CLIPTextModelWithProjection(
        CLIPTextConfig(
            bos_token_id=0    ,		eos_token_id=2    ,		hidden_size=a    ,		projection_dim=a    ,		intermediate_size=37    ,		layer_norm_eps=1e-05    ,		num_attention_heads=4    ,		num_hidden_layers=5    ,		pad_token_id=1    ,		vocab_size=10_00    ,		)    )

    torch.manual_seed(0    )
    A__       =     PriorTransformer(
        num_attention_heads=2    ,		attention_head_dim=12    ,		embedding_dim=a    ,		num_layers=1    ,		)

    torch.manual_seed(0    )
    A__       =     DDPMScheduler(
        variance_type='fixed_small_log'    ,		prediction_type='sample'    ,		num_train_timesteps=10_00    ,		clip_sample=a    ,		clip_sample_range=5.0    ,		beta_schedule='squaredcos_cap_v2'    ,		)

    # regular denoising components

    torch.manual_seed(0    )
    A__       =     StableUnCLIPImageNormalizer(embedding_dim=a    )
    A__       =     DDPMScheduler(beta_schedule='squaredcos_cap_v2'    )

    torch.manual_seed(0    )
    A__       =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

    torch.manual_seed(0    )
    A__       =     CLIPTextModel(
        CLIPTextConfig(
            bos_token_id=0    ,		eos_token_id=2    ,		hidden_size=a    ,		projection_dim=32    ,		intermediate_size=37    ,		layer_norm_eps=1e-05    ,		num_attention_heads=4    ,		num_hidden_layers=5    ,		pad_token_id=1    ,		vocab_size=10_00    ,		)    )

    torch.manual_seed(0    )
    A__       =     UNetaDConditionModel(
        sample_size=32    ,		in_channels=4    ,		out_channels=4    ,		down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')    ,		up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')    ,		block_out_channels=(32, 64)    ,		attention_head_dim=(2, 4)    ,		class_embed_type='projection'    ,		projection_class_embeddings_input_dim=embedder_projection_dim * 2    ,		cross_attention_dim=a    ,		layers_per_block=1    ,		upcast_attention=a    ,		use_linear_projection=a    ,		)

    torch.manual_seed(0    )
    A__       =     DDIMScheduler(
        beta_schedule='scaled_linear'    ,		beta_start=0.0_0085    ,		beta_end=0.012    ,		prediction_type='v_prediction'    ,		set_alpha_to_one=a    ,		steps_offset=1    ,		)

    torch.manual_seed(0    )
    A__       =     AutoencoderKL()

    A__       =     {
        # prior components
        'prior_tokenizer': prior_tokenizer,
        'prior_text_encoder': prior_text_encoder,
        'prior': prior,
        'prior_scheduler': prior_scheduler,
        # image noising components
        'image_normalizer': image_normalizer,
        'image_noising_scheduler': image_noising_scheduler,
        # regular denoising components
        'tokenizer': tokenizer,
        'text_encoder': text_encoder,
        'unet': unet,
        'scheduler': scheduler,
        'vae': vae,
    }

    return components






 def 						_UpperCamelCase (self       :  List[Any]    ,		a       :  Optional[Any]    ,		a       :  str=0    )       ->   Dict:






    """simple docstring"""


    if str(a    ).startswith('mps'    ):
       A__       =     torch.manual_seed(a    )
    else:
       A__       =     torch.Generator(device=a    ).manual_seed(a    )
    A__       =     {
        'prompt': 'A painting of a squirrel eating a burger',
        'generator': generator,
        'num_inference_steps': 2,
        'prior_num_inference_steps': 2,
        'output_type': 'numpy',
    }
    return inputs






 def 						_UpperCamelCase (self       :  List[str]    )       ->   Optional[int]:






    """simple docstring"""


    A__       =     torch_device == 'cpu'

    self._test_attention_slicing_forward_pass(test_max_difference=a    )






 def 						_UpperCamelCase (self       :  Dict    )       ->   str:






    """simple docstring"""


    A__       =     torch_device in ['cpu', 'mps']

    self._test_inference_batch_single_identical(test_max_difference=a    )




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






 def 						_UpperCamelCase (self       :  Any    )       ->   Optional[Any]:






    """simple docstring"""


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






 def 						_UpperCamelCase (self       :  Dict    )       ->   Union[str, Any]:






    """simple docstring"""


    A__       =     load_numpy(
        'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'    )

    A__       =     StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'    ,		torch_dtype=torch.floataa    )
    pipe.to(a    )
    pipe.set_progress_bar_config(disable=a    )
    # stable unclip will oom when integration tests are run on a V100,
    # so turn on memory savings
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()

    A__       =     torch.Generator(device='cpu'    ).manual_seed(0    )
    A__       =     pipe('anime turle'    ,		generator=a    ,		output_type='np'    )

    A__       =     output.images[0]

    assert image.shape == (7_68, 7_68, 3)

    assert_mean_pixel_difference(a    ,		a    )






 def 						_UpperCamelCase (self       :  Optional[Any]    )       ->   str:






    """simple docstring"""


    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

    A__       =     StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'    ,		torch_dtype=torch.floataa    )
    A__       =     pipe.to(a    )
    pipe.set_progress_bar_config(disable=a    )
    pipe.enable_attention_slicing()
    pipe.enable_sequential_cpu_offload()

    A__       =     pipe(
        'anime turtle'    ,		prior_num_inference_steps=2    ,		num_inference_steps=2    ,		output_type='np'    ,		)

    A__       =     torch.cuda.max_memory_allocated()
    # make sure that less than 7 GB is allocated
    assert mem_bytes < 7 * 10**9


