


from __future__ import annotations

import copy
import inspect
import unittest

import numpy as np

from transformers import is_tf_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow
from transformers.utils import cached_property

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


if is_tf_available():
 import tensorflow as tf

 from transformers import (
     TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST,
     TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
     TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
     TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
     TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
     LayoutLMvaConfig,
     TFLayoutLMvaForQuestionAnswering,
     TFLayoutLMvaForSequenceClassification,
     TFLayoutLMvaForTokenClassification,
     TFLayoutLMvaModel,
 )

if is_vision_available():
 from PIL import Image

 from transformers import LayoutLMvaImageProcessor


class 		_SCREAMING_SNAKE_CASE     :


     def __init__(  self							,					A_							,					A_=2							,					A_=3							,					A_=4							,					A_=2							,					A_=7							,					A_=True							,					A_=True							,					A_=True							,					A_=True							,					A_=99							,					A_=36							,					A_=2							,					A_=4							,					A_=37							,					A_="gelu"							,					A_=0.1							,					A_=0.1							,					A_=5_12							,					A_=16							,					A_=2							,					A_=0.0_2							,					A_=6							,					A_=6							,					A_=3							,					A_=4							,					A_=None							,					A_=10_00							,					):
       _UpperCAmelCase		:							Optional[int]							=     parent
       _UpperCAmelCase		:							Any							=     batch_size
       _UpperCAmelCase		:							Dict							=     num_channels
       _UpperCAmelCase		:							Tuple							=     image_size
       _UpperCAmelCase		:							str							=     patch_size
       _UpperCAmelCase		:							Tuple							=     is_training
       _UpperCAmelCase		:							Tuple							=     use_input_mask
       _UpperCAmelCase		:							Union[str, Any]							=     use_token_type_ids
       _UpperCAmelCase		:							Any							=     use_labels
       _UpperCAmelCase		:							Any							=     vocab_size
       _UpperCAmelCase		:							str							=     hidden_size
       _UpperCAmelCase		:							Union[str, Any]							=     num_hidden_layers
       _UpperCAmelCase		:							Union[str, Any]							=     num_attention_heads
       _UpperCAmelCase		:							List[str]							=     intermediate_size
       _UpperCAmelCase		:							int							=     hidden_act
       _UpperCAmelCase		:							Optional[Any]							=     hidden_dropout_prob
       _UpperCAmelCase		:							Any							=     attention_probs_dropout_prob
       _UpperCAmelCase		:							int							=     max_position_embeddings
       _UpperCAmelCase		:							Optional[Any]							=     type_vocab_size
       _UpperCAmelCase		:							List[Any]							=     type_sequence_label_size
       _UpperCAmelCase		:							Optional[int]							=     initializer_range
       _UpperCAmelCase		:							Optional[Any]							=     coordinate_size
       _UpperCAmelCase		:							str							=     shape_size
       _UpperCAmelCase		:							List[Any]							=     num_labels
       _UpperCAmelCase		:							List[Any]							=     num_choices
       _UpperCAmelCase		:							List[str]							=     scope
       _UpperCAmelCase		:							Tuple							=     range_bbox

       # LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
       _UpperCAmelCase		:							Optional[Any]							=     text_seq_length
       _UpperCAmelCase		:							Any							=     (image_size // patch_size) ** 2 + 1
       _UpperCAmelCase		:							str							=     self.text_seq_length + self.image_seq_length


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							Optional[int]							=     ids_tensor([self.batch_size, self.text_seq_length]							,					self.vocab_size    )

       _UpperCAmelCase		:							List[Any]							=     ids_tensor([self.batch_size, self.text_seq_length, 4]							,					self.range_bbox    )
       _UpperCAmelCase		:							List[Any]							=     bbox.numpy()
       # Ensure that bbox is legal
       for i in range(bbox.shape[0]    ):
         for j in range(bbox.shape[1]    ):
           if bbox[i, j, 3] < bbox[i, j, 1]:
             _UpperCAmelCase		:							List[str]							=     bbox[i, j, 3]
             _UpperCAmelCase		:							Any							=     bbox[i, j, 1]
             _UpperCAmelCase		:							Optional[int]							=     tmp_coordinate
           if bbox[i, j, 2] < bbox[i, j, 0]:
             _UpperCAmelCase		:							Union[str, Any]							=     bbox[i, j, 2]
             _UpperCAmelCase		:							List[str]							=     bbox[i, j, 0]
             _UpperCAmelCase		:							int							=     tmp_coordinate
       _UpperCAmelCase		:							Any							=     tf.constant(A_    )

       _UpperCAmelCase		:							Union[str, Any]							=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

       _UpperCAmelCase		:							str							=     None
       if self.use_input_mask:
         _UpperCAmelCase		:							List[str]							=     random_attention_mask([self.batch_size, self.text_seq_length]    )

       _UpperCAmelCase		:							Tuple							=     None
       if self.use_token_type_ids:
         _UpperCAmelCase		:							str							=     ids_tensor([self.batch_size, self.text_seq_length]							,					self.type_vocab_size    )

       _UpperCAmelCase		:							Dict							=     None
       _UpperCAmelCase		:							Tuple							=     None
       if self.use_labels:
         _UpperCAmelCase		:							str							=     ids_tensor([self.batch_size]							,					self.type_sequence_label_size    )
         _UpperCAmelCase		:							Tuple							=     ids_tensor([self.batch_size, self.text_seq_length]							,					self.num_labels    )

       _UpperCAmelCase		:							Any							=     LayoutLMvaConfig(
           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							,					initializer_range=self.initializer_range							,					coordinate_size=self.coordinate_size							,					shape_size=self.shape_size							,					input_size=self.image_size							,					patch_size=self.patch_size							,					)

       return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels


     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_    ):
       _UpperCAmelCase		:							Optional[Any]							=     TFLayoutLMvaModel(config=A_    )

       # text + image
       _UpperCAmelCase		:							Any							=     model(A_							,					pixel_values=A_							,					training=A_    )
       _UpperCAmelCase		:							int							=     model(
           A_							,					bbox=A_							,					pixel_values=A_							,					attention_mask=A_							,					token_type_ids=A_							,					training=A_							,					)
       _UpperCAmelCase		:							List[Any]							=     model(A_							,					bbox=A_							,					pixel_values=A_							,					training=A_    )

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

       # text only
       _UpperCAmelCase		:							List[Any]							=     model(A_							,					training=A_    )
       self.parent.assertEqual(
           result.last_hidden_state.shape							,					(self.batch_size, self.text_seq_length, self.hidden_size)    )

       # image only
       _UpperCAmelCase		:							Dict							=     model({"""pixel_values""": pixel_values}							,					training=A_    )
       self.parent.assertEqual(
           result.last_hidden_state.shape							,					(self.batch_size, self.image_seq_length, self.hidden_size)    )


     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_    ):
       _UpperCAmelCase		:							Optional[int]							=     self.num_labels
       _UpperCAmelCase		:							Tuple							=     TFLayoutLMvaForSequenceClassification(config=A_    )
       _UpperCAmelCase		:							Optional[int]							=     model(
           A_							,					bbox=A_							,					pixel_values=A_							,					attention_mask=A_							,					token_type_ids=A_							,					labels=A_							,					training=A_							,					)
       self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_labels)    )


     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_    ):
       _UpperCAmelCase		:							Any							=     self.num_labels
       _UpperCAmelCase		:							List[Any]							=     TFLayoutLMvaForTokenClassification(config=A_    )
       _UpperCAmelCase		:							Optional[int]							=     model(
           A_							,					bbox=A_							,					pixel_values=A_							,					attention_mask=A_							,					token_type_ids=A_							,					labels=A_							,					training=A_							,					)
       self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.text_seq_length, self.num_labels)    )


     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_    ):
       _UpperCAmelCase		:							Union[str, Any]							=     2
       _UpperCAmelCase		:							Optional[Any]							=     TFLayoutLMvaForQuestionAnswering(config=A_    )
       _UpperCAmelCase		:							Tuple							=     model(
           A_							,					bbox=A_							,					pixel_values=A_							,					attention_mask=A_							,					token_type_ids=A_							,					start_positions=A_							,					end_positions=A_							,					training=A_							,					)
       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        __snake_case(  self    ):
       _UpperCAmelCase		:							Any							=     self.prepare_config_and_inputs()
       ((_UpperCAmelCase),(_UpperCAmelCase),(_UpperCAmelCase),(_UpperCAmelCase),(_UpperCAmelCase),(_UpperCAmelCase),(_UpperCAmelCase),(_UpperCAmelCase))		:							int							=     config_and_inputs
       _UpperCAmelCase		:							Tuple							=     {
           """input_ids""": input_ids,
           """bbox""": bbox,
           """pixel_values""": pixel_values,
           """token_type_ids""": token_type_ids,
           """attention_mask""": input_mask,
       }
       return config, inputs_dict


@require_tf
class 		_SCREAMING_SNAKE_CASE     (						A							,	A							,	unittest.TestCase     ):
     __SCREAMING_SNAKE_CASE     					=    (
         (
             TFLayoutLMvaModel,
             TFLayoutLMvaForQuestionAnswering,
             TFLayoutLMvaForSequenceClassification,
             TFLayoutLMvaForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     __SCREAMING_SNAKE_CASE     					=    (
         {'''document-question-answering''': TFLayoutLMvaForQuestionAnswering, '''feature-extraction''': TFLayoutLMvaModel}
         if is_tf_available()
         else {}
     )

     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False
     __SCREAMING_SNAKE_CASE     					=    False


     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_							,					A_    ):
       return True


     def        __snake_case(  self							,					A_							,					A_							,					A_=False    ):
       _UpperCAmelCase		:							Optional[Any]							=     copy.deepcopy(A_    )

       if model_class in get_values(A_    ):
         _UpperCAmelCase		:							str							=     {
             k: tf.tile(tf.expand_dims(A_							,					1    )							,					(1, self.model_tester.num_choices) + (1,) * (v.ndim - 1)    )
             if isinstance(A_							,					tf.Tensor    ) and v.ndim > 0
             else v
             for k, v in inputs_dict.items()
         }

       if return_labels:
         if model_class in get_values(A_    ):
           _UpperCAmelCase		:							str							=     tf.ones(self.model_tester.batch_size							,					dtype=tf.intaa    )
         elif model_class in get_values(A_    ):
           _UpperCAmelCase		:							Tuple							=     tf.zeros(self.model_tester.batch_size							,					dtype=tf.intaa    )
           _UpperCAmelCase		:							Optional[int]							=     tf.zeros(self.model_tester.batch_size							,					dtype=tf.intaa    )
         elif model_class in get_values(A_    ):
           _UpperCAmelCase		:							List[str]							=     tf.zeros(self.model_tester.batch_size							,					dtype=tf.intaa    )
         elif model_class in get_values(A_    ):
           _UpperCAmelCase		:							str							=     tf.zeros(
               (self.model_tester.batch_size, self.model_tester.text_seq_length)							,					dtype=tf.intaa    )

       return inputs_dict


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							str							=     TFLayoutLMvaModelTester(self    )
       _UpperCAmelCase		:							Tuple							=     ConfigTester(self							,					config_class=A_							,					hidden_size=37    )


     def        __snake_case(  self    ):
       self.config_tester.run_common_tests()


     def        __snake_case(  self    ):
       _UpperCAmelCase,_UpperCAmelCase		:							int							=     self.model_tester.prepare_config_and_inputs_for_common()
       for model_class in self.all_model_classes:
         _UpperCAmelCase		:							List[str]							=     model_class(A_    )
         if getattr(A_							,					"""hf_compute_loss"""							,					A_    ):
           # The number of elements in the loss should be the same as the number of elements in the label
           _UpperCAmelCase		:							str							=     self._prepare_for_class(inputs_dict.copy()							,					A_							,					return_labels=A_    )
           _UpperCAmelCase		:							List[Any]							=     prepared_for_class[
               sorted(prepared_for_class.keys() - inputs_dict.keys()							,					reverse=A_    )[0]
           ]
           _UpperCAmelCase		:							List[str]							=     added_label.shape.as_list()[:1]

           # Test that model correctly compute the loss with kwargs
           _UpperCAmelCase		:							Union[str, Any]							=     self._prepare_for_class(inputs_dict.copy()							,					A_							,					return_labels=A_    )
           _UpperCAmelCase		:							Tuple							=     prepared_for_class.pop("""input_ids"""    )

           _UpperCAmelCase		:							Union[str, Any]							=     model(A_							,					**A_    )[0]
           self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )

           # Test that model correctly compute the loss when we mask some positions
           _UpperCAmelCase		:							int							=     self._prepare_for_class(inputs_dict.copy()							,					A_							,					return_labels=A_    )
           _UpperCAmelCase		:							Optional[int]							=     prepared_for_class.pop("""input_ids"""    )
           if "labels" in prepared_for_class:
             _UpperCAmelCase		:							str							=     prepared_for_class["""labels"""].numpy()
             if len(labels.shape    ) > 1 and labels.shape[1] != 1:
               _UpperCAmelCase		:							Dict							=     -1_00
               _UpperCAmelCase		:							int							=     tf.convert_to_tensor(A_    )
               _UpperCAmelCase		:							Tuple							=     model(A_							,					**A_    )[0]
               self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )
               self.assertTrue(not np.any(np.isnan(loss.numpy()    )    )    )

                # Test that model correctly compute the loss with a dict
           _UpperCAmelCase		:							Optional[int]							=     self._prepare_for_class(inputs_dict.copy()							,					A_							,					return_labels=A_    )
           _UpperCAmelCase		:							Any							=     model(A_    )[0]
           self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )

           # Test that model correctly compute the loss with a tuple
           _UpperCAmelCase		:							int							=     self._prepare_for_class(inputs_dict.copy()							,					A_							,					return_labels=A_    )

           # Get keys that were added with the _prepare_for_class function
           _UpperCAmelCase		:							str							=     prepared_for_class.keys() - inputs_dict.keys()
           _UpperCAmelCase		:							Dict							=     inspect.signature(model.call    ).parameters
           _UpperCAmelCase		:							Tuple							=     list(signature.keys()    )

           # Create a dictionary holding the location of the tensors in the tuple
           _UpperCAmelCase		:							Tuple							=     {0: """input_ids"""}
           for label_key in label_keys:
             _UpperCAmelCase		:							str							=     signature_names.index(A_    )
             _UpperCAmelCase		:							List[Any]							=     label_key
           _UpperCAmelCase		:							int							=     sorted(tuple_index_mapping.items()    )
           # Initialize a list with their default values, update the values and convert to a tuple
           _UpperCAmelCase		:							List[Any]							=     []

           for name in signature_names:
             if name != "kwargs":
               list_input.append(signature[name].default    )

           for index, value in sorted_tuple_index_mapping:
             _UpperCAmelCase		:							str							=     prepared_for_class[value]

           _UpperCAmelCase		:							Any							=     tuple(A_    )

           # Send to model
           _UpperCAmelCase		:							str							=     model(tuple_input[:-1]    )[0]

           self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )


     def        __snake_case(  self    ):
       (
           (
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),
       )		:							Dict							=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(A_							,					A_							,					A_							,					A_							,					A_							,					A_    )


     def        __snake_case(  self    ):
       (
           (
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),
       )		:							Any							=     self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
         _UpperCAmelCase		:							List[str]							=     type
         self.model_tester.create_and_check_model(A_							,					A_							,					A_							,					A_							,					A_							,					A_    )


     def        __snake_case(  self    ):
       (
           (
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),
       )		:							List[Any]							=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(
           A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_    )


     def        __snake_case(  self    ):
       (
           (
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),
       )		:							List[Any]							=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(
           A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_    )


     def        __snake_case(  self    ):
       (
           (
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),(
           _UpperCAmelCase
       ),
       )		:							Optional[Any]							=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(
           A_							,					A_							,					A_							,					A_							,					A_							,					A_							,					A_    )




     @slow
     def        __snake_case(  self    ):
       for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _UpperCAmelCase		:							int							=     TFLayoutLMvaModel.from_pretrained(A_    )
         self.assertIsNotNone(A_    )




def    a__    (		):
  _UpperCAmelCase		:							Dict							=     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""     )
  return image





@require_tf
class 		_SCREAMING_SNAKE_CASE     (						unittest.TestCase     ):


     @cached_property
     def        __snake_case(  self    ):
       return LayoutLMvaImageProcessor(apply_ocr=A_    ) if is_vision_available() else None




     @slow
     def        __snake_case(  self    ):
       _UpperCAmelCase		:							Dict							=     TFLayoutLMvaModel.from_pretrained("""microsoft/layoutlmv3-base"""    )

       _UpperCAmelCase		:							List[Any]							=     self.default_image_processor
       _UpperCAmelCase		:							Union[str, Any]							=     prepare_img()
       _UpperCAmelCase		:							int							=     image_processor(images=A_							,					return_tensors="""tf"""    ).pixel_values

       _UpperCAmelCase		:							Union[str, Any]							=     tf.constant([[1, 2]]    )
       _UpperCAmelCase		:							List[str]							=     tf.expand_dims(tf.constant([[1, 2, 3, 4], [5, 6, 7, 8]]    )							,					axis=0    )

       # forward pass
       _UpperCAmelCase		:							str							=     model(input_ids=A_							,					bbox=A_							,					pixel_values=A_							,					training=A_    )

       # verify the logits
       _UpperCAmelCase		:							Tuple							=     (1, 1_99, 7_68)
       self.assertEqual(outputs.last_hidden_state.shape							,					A_    )

       _UpperCAmelCase		:							int							=     tf.constant(
           [[-0.0_5_2_9, 0.3_6_1_8, 0.1_6_3_2], [-0.1_5_8_7, -0.1_6_6_7, -0.0_4_0_0], [-0.1_5_5_7, -0.1_6_7_1, -0.0_5_0_5]]    )

       self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]							,					A_							,					atol=1e-4    )    )









import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

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


@skip_mps
class 		_SCREAMING_SNAKE_CASE     (						A							,	A							,	unittest.TestCase     ):
     __SCREAMING_SNAKE_CASE     					=    IFPipeline
     __SCREAMING_SNAKE_CASE     					=    TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
     __SCREAMING_SNAKE_CASE     					=    TEXT_TO_IMAGE_BATCH_PARAMS
     __SCREAMING_SNAKE_CASE     					=    PipelineTesterMixin.required_optional_params - {'''latents'''}


     def        __snake_case(  self    ):
       return self._get_dummy_components()


     def        __snake_case(  self							,					A_							,					A_=0    ):
       if str(A_    ).startswith("""mps"""    ):
         _UpperCAmelCase		:							Tuple							=     torch.manual_seed(A_    )
       else:
         _UpperCAmelCase		:							Optional[Any]							=     torch.Generator(device=A_    ).manual_seed(A_    )

       _UpperCAmelCase		:							Any							=     {
           """prompt""": """A painting of a squirrel eating a burger""",
           """generator""": generator,
           """num_inference_steps""": 2,
           """output_type""": """numpy""",
       }

       return inputs


     def        __snake_case(  self    ):
       self._test_save_load_optional_components()


     @unittest.skipIf(torch_device != """cuda"""							,					reason="""float16 requires CUDA"""    )
     def        __snake_case(  self    ):
       # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
       super().test_save_load_floataa(expected_max_diff=1e-1    )


     def        __snake_case(  self    ):
       self._test_attention_slicing_forward_pass(expected_max_diff=1e-2    )


     def        __snake_case(  self    ):
       self._test_save_load_local()


     def        __snake_case(  self    ):
       self._test_inference_batch_single_identical(
           expected_max_diff=1e-2							,					)




     @unittest.skipIf(
         torch_device != """cuda""" or not is_xformers_available()							,					reason="""XFormers attention is only available with CUDA and `xformers` installed"""							,					)
     def        __snake_case(  self    ):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3    )


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


     def        __snake_case(  self    ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()


     def        __snake_case(  self    ):
       # if

       _UpperCAmelCase		:							List[str]							=     IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""							,					variant="""fp16"""							,					torch_dtype=torch.floataa    )

       _UpperCAmelCase		:							Dict							=     IFSuperResolutionPipeline.from_pretrained(
           """DeepFloyd/IF-II-L-v1.0"""							,					variant="""fp16"""							,					torch_dtype=torch.floataa							,					text_encoder=A_							,					tokenizer=A_    )

       # pre compute text embeddings and remove T5 to save memory

       pipe_a.text_encoder.to("""cuda"""    )

       _UpperCAmelCase,_UpperCAmelCase		:							Dict							=     pipe_a.encode_prompt("""anime turtle"""							,					device="""cuda"""    )

       del pipe_a.tokenizer
       del pipe_a.text_encoder
       gc.collect()

       _UpperCAmelCase		:							List[Any]							=     None
       _UpperCAmelCase		:							Any							=     None

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )

       self._test_if(A_							,					A_							,					A_							,					A_    )

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # img2img

       _UpperCAmelCase		:							Any							=     IFImgaImgPipeline(**pipe_a.components    )
       _UpperCAmelCase		:							Union[str, Any]							=     IFImgaImgSuperResolutionPipeline(**pipe_a.components    )

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )

       self._test_if_imgaimg(A_							,					A_							,					A_							,					A_    )

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # inpainting

       _UpperCAmelCase		:							Optional[int]							=     IFInpaintingPipeline(**pipe_a.components    )
       _UpperCAmelCase		:							Union[str, Any]							=     IFInpaintingSuperResolutionPipeline(**pipe_a.components    )

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()    )

       self._test_if_inpainting(A_							,					A_							,					A_							,					A_    )


     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_    ):
       # pipeline 1

       _start_torch_memory_measurement()

       _UpperCAmelCase		:							Optional[Any]							=     torch.Generator(device="""cpu"""    ).manual_seed(0    )
       _UpperCAmelCase		:							List[str]							=     pipe_a(
           prompt_embeds=A_							,					negative_prompt_embeds=A_							,					num_inference_steps=2							,					generator=A_							,					output_type="""np"""							,					)

       _UpperCAmelCase		:							int							=     output.images[0]

       assert image.shape == (64, 64, 3)

       _UpperCAmelCase		:							Dict							=     torch.cuda.max_memory_allocated()
       assert mem_bytes < 13 * 10**9

       _UpperCAmelCase		:							str							=     load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy"""    )
       assert_mean_pixel_difference(A_							,					A_    )

       # pipeline 2

       _start_torch_memory_measurement()

       _UpperCAmelCase		:							List[Any]							=     torch.Generator(device="""cpu"""    ).manual_seed(0    )

       _UpperCAmelCase		:							str							=     floats_tensor((1, 3, 64, 64)							,					rng=random.Random(0    )    ).to(A_    )

       _UpperCAmelCase		:							str							=     pipe_a(
           prompt_embeds=A_							,					negative_prompt_embeds=A_							,					image=A_							,					generator=A_							,					num_inference_steps=2							,					output_type="""np"""							,					)

       _UpperCAmelCase		:							int							=     output.images[0]

       assert image.shape == (2_56, 2_56, 3)

       _UpperCAmelCase		:							Any							=     torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       _UpperCAmelCase		:							List[Any]							=     load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy"""    )
       assert_mean_pixel_difference(A_							,					A_    )


     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_    ):
       # pipeline 1

       _start_torch_memory_measurement()

       _UpperCAmelCase		:							Optional[int]							=     floats_tensor((1, 3, 64, 64)							,					rng=random.Random(0    )    ).to(A_    )

       _UpperCAmelCase		:							Dict							=     torch.Generator(device="""cpu"""    ).manual_seed(0    )

       _UpperCAmelCase		:							List[str]							=     pipe_a(
           prompt_embeds=A_							,					negative_prompt_embeds=A_							,					image=A_							,					num_inference_steps=2							,					generator=A_							,					output_type="""np"""							,					)

       _UpperCAmelCase		:							int							=     output.images[0]

       assert image.shape == (64, 64, 3)

       _UpperCAmelCase		:							str							=     torch.cuda.max_memory_allocated()
       assert mem_bytes < 10 * 10**9

       _UpperCAmelCase		:							str							=     load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy"""    )
       assert_mean_pixel_difference(A_							,					A_    )

       # pipeline 2

       _start_torch_memory_measurement()

       _UpperCAmelCase		:							List[str]							=     torch.Generator(device="""cpu"""    ).manual_seed(0    )

       _UpperCAmelCase		:							Union[str, Any]							=     floats_tensor((1, 3, 2_56, 2_56)							,					rng=random.Random(0    )    ).to(A_    )
       _UpperCAmelCase		:							Optional[int]							=     floats_tensor((1, 3, 64, 64)							,					rng=random.Random(0    )    ).to(A_    )

       _UpperCAmelCase		:							Optional[int]							=     pipe_a(
           prompt_embeds=A_							,					negative_prompt_embeds=A_							,					image=A_							,					original_image=A_							,					generator=A_							,					num_inference_steps=2							,					output_type="""np"""							,					)

       _UpperCAmelCase		:							Dict							=     output.images[0]

       assert image.shape == (2_56, 2_56, 3)

       _UpperCAmelCase		:							Any							=     torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       _UpperCAmelCase		:							int							=     load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy"""    )
       assert_mean_pixel_difference(A_							,					A_    )




     def        __snake_case(  self							,					A_							,					A_							,					A_							,					A_    ):
       # pipeline 1

       _start_torch_memory_measurement()

       _UpperCAmelCase		:							List[Any]							=     floats_tensor((1, 3, 64, 64)							,					rng=random.Random(0    )    ).to(A_    )
       _UpperCAmelCase		:							Dict							=     floats_tensor((1, 3, 64, 64)							,					rng=random.Random(1    )    ).to(A_    )

       _UpperCAmelCase		:							Optional[Any]							=     torch.Generator(device="""cpu"""    ).manual_seed(0    )
       _UpperCAmelCase		:							str							=     pipe_a(
           prompt_embeds=A_							,					negative_prompt_embeds=A_							,					image=A_							,					mask_image=A_							,					num_inference_steps=2							,					generator=A_							,					output_type="""np"""							,					)

       _UpperCAmelCase		:							Optional[Any]							=     output.images[0]

       assert image.shape == (64, 64, 3)

       _UpperCAmelCase		:							List[str]							=     torch.cuda.max_memory_allocated()
       assert mem_bytes < 10 * 10**9

       _UpperCAmelCase		:							Dict							=     load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy"""    )
       assert_mean_pixel_difference(A_							,					A_    )

       # pipeline 2

       _start_torch_memory_measurement()

       _UpperCAmelCase		:							str							=     torch.Generator(device="""cpu"""    ).manual_seed(0    )

       _UpperCAmelCase		:							Optional[int]							=     floats_tensor((1, 3, 64, 64)							,					rng=random.Random(0    )    ).to(A_    )
       _UpperCAmelCase		:							Tuple							=     floats_tensor((1, 3, 2_56, 2_56)							,					rng=random.Random(0    )    ).to(A_    )
       _UpperCAmelCase		:							str							=     floats_tensor((1, 3, 2_56, 2_56)							,					rng=random.Random(1    )    ).to(A_    )

       _UpperCAmelCase		:							Union[str, Any]							=     pipe_a(
           prompt_embeds=A_							,					negative_prompt_embeds=A_							,					image=A_							,					mask_image=A_							,					original_image=A_							,					generator=A_							,					num_inference_steps=2							,					output_type="""np"""							,					)

       _UpperCAmelCase		:							Any							=     output.images[0]

       assert image.shape == (2_56, 2_56, 3)

       _UpperCAmelCase		:							Optional[int]							=     torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       _UpperCAmelCase		:							Tuple							=     load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy"""    )
       assert_mean_pixel_difference(A_							,					A_    )




def    a__    (		):
  torch.cuda.empty_cache()
  torch.cuda.reset_max_memory_allocated()
  torch.cuda.reset_peak_memory_stats()





