


import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


UpperCAmelCase_      : str  					=       logging.get_logger(__name__)

enable_full_determinism()
class  SCREAMING_SNAKE_CASE__  (				lowercase__						,    lowercase__						,    unittest.TestCase			):
      snake_case__				:     Tuple        =						UNetaDModel
      snake_case__				:     int        =						'''sample'''
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[int]     )		->			int:
             a_					:		Optional[Any]							=  4
             a_					:		Tuple							=  3
             a_					:		Dict							=  (3_2, 3_2)

             a_					:		List[Any]							=  floats_tensor((batch_size, num_channels) + sizes     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		str							=  torch.tensor([1_0]     ).to(SCREAMING_SNAKE_CASE__     )

             return {"sample": noise, "timestep": time_step}
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			List[Any]:
             return (3, 3_2, 3_2)
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			Optional[Any]:
             return (3, 3_2, 3_2)

      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Optional[Any]:
             a_					:		List[str]							=  {
                 'block_out_channels': (3_2, 6_4),
                 'down_block_types': ('DownBlock2D', 'AttnDownBlock2D'),
                 'up_block_types': ('AttnUpBlock2D', 'UpBlock2D'),
                 'attention_head_dim': 3,
                 'out_channels': 3,
                 'in_channels': 3,
                 'layers_per_block': 2,
                 'sample_size': 3_2,
             }
             a_					:		List[Any]							=  self.dummy_input
             return init_dict, inputs_dict
class  SCREAMING_SNAKE_CASE__  (				lowercase__						,    lowercase__						,    unittest.TestCase			):
      snake_case__				:     Tuple        =						UNetaDModel
      snake_case__				:     Any        =						'''sample'''
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			List[Any]:
             a_					:		Tuple							=  4
             a_					:		Any							=  4
             a_					:		Dict							=  (3_2, 3_2)

             a_					:		str							=  floats_tensor((batch_size, num_channels) + sizes     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		List[str]							=  torch.tensor([1_0]     ).to(SCREAMING_SNAKE_CASE__     )

             return {"sample": noise, "timestep": time_step}
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			str:
             return (4, 3_2, 3_2)
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			List[Any]:
             return (4, 3_2, 3_2)
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			List[Any]:
             a_					:		Any							=  {
                 'sample_size': 3_2,
                 'in_channels': 4,
                 'out_channels': 4,
                 'layers_per_block': 2,
                 'block_out_channels': (3_2, 6_4),
                 'attention_head_dim': 3_2,
                 'down_block_types': ('DownBlock2D', 'DownBlock2D'),
                 'up_block_types': ('UpBlock2D', 'UpBlock2D'),
             }
             a_					:		int							=  self.dummy_input
             return init_dict, inputs_dict
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Union[str, Any]:
             a_							,		a_					:		Dict							=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'     ,							output_loading_info=SCREAMING_SNAKE_CASE__     )

             self.assertIsNotNone(SCREAMING_SNAKE_CASE__     )
             self.assertEqual(len(loading_info['missing_keys']     )     ,							0     )

             model.to(SCREAMING_SNAKE_CASE__     )
             a_					:		List[Any]							=  model(**self.dummy_input     ).sample

             assert image is not None, "Make sure output is not None"
      @unittest.skipIf(torch_device != 'cuda'     ,							'This test is supposed to run on GPU'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			str:
             a_							,		a_					:		List[Any]							=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'     ,							output_loading_info=SCREAMING_SNAKE_CASE__     )
             model.to(SCREAMING_SNAKE_CASE__     )
             a_					:		str							=  model(**self.dummy_input     ).sample

             assert image is not None, "Make sure output is not None"
      @unittest.skipIf(torch_device != 'cuda'     ,							'This test is supposed to run on GPU'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			Optional[Any]:
             # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
             a_							,		a_					:		Optional[int]							=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'     ,							output_loading_info=SCREAMING_SNAKE_CASE__     )
             model_accelerate.to(SCREAMING_SNAKE_CASE__     )
             model_accelerate.eval()

             a_					:		Optional[Any]							=  torch.randn(
                 1     ,							model_accelerate.config.in_channels     ,							model_accelerate.config.sample_size     ,							model_accelerate.config.sample_size     ,							generator=torch.manual_seed(0     )     ,							)
             a_					:		Union[str, Any]							=  noise.to(SCREAMING_SNAKE_CASE__     )
             a_					:		Any							=  torch.tensor([1_0] * noise.shape[0]     ).to(SCREAMING_SNAKE_CASE__     )

             a_					:		Dict							=  model_accelerate(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )['sample']

             # two models don't need to stay in the device at the same time
             del model_accelerate
             torch.cuda.empty_cache()
             gc.collect()

             a_							,		a_					:		Optional[Any]							=  UNetaDModel.from_pretrained(
                 'fusing/unet-ldm-dummy-update'     ,							output_loading_info=SCREAMING_SNAKE_CASE__     ,							low_cpu_mem_usage=SCREAMING_SNAKE_CASE__     )
             model_normal_load.to(SCREAMING_SNAKE_CASE__     )
             model_normal_load.eval()
             a_					:		List[Any]							=  model_normal_load(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )['sample']

             assert torch_all_close(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							rtol=1E-3     )

      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			Optional[int]:
             a_					:		Dict							=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'     )
             model.eval()
             model.to(SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[Any]							=  torch.randn(
                 1     ,							model.config.in_channels     ,							model.config.sample_size     ,							model.config.sample_size     ,							generator=torch.manual_seed(0     )     ,							)
             a_					:		Dict							=  noise.to(SCREAMING_SNAKE_CASE__     )
             a_					:		Any							=  torch.tensor([1_0] * noise.shape[0]     ).to(SCREAMING_SNAKE_CASE__     )

             with torch.no_grad():
                    a_					:		Union[str, Any]							=  model(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ).sample

             a_					:		Optional[int]							=  output[0, -1, -3:, -3:].flatten().cpu()
             # fmt: off
             a_					:		List[Any]							=  torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]     )
             # fmt: on

             self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							rtol=1E-3     )     )
class  SCREAMING_SNAKE_CASE__  (				lowercase__						,    lowercase__						,    unittest.TestCase			):
      snake_case__				:     List[Any]        =						UNetaDModel
      snake_case__				:     str        =						'''sample'''
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       List[str]=(3_2, 3_2)     )		->			List[str]:
             a_					:		Any							=  4
             a_					:		Any							=  3

             a_					:		int							=  floats_tensor((batch_size, num_channels) + sizes     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		Tuple							=  torch.tensor(batch_size * [1_0]     ).to(dtype=torch.intaa     ,							device=SCREAMING_SNAKE_CASE__     )

             return {"sample": noise, "timestep": time_step}
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Union[str, Any]:
             return (3, 3_2, 3_2)
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Tuple:
             return (3, 3_2, 3_2)
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Dict:
             a_					:		Optional[Any]							=  {
                 'block_out_channels': [3_2, 6_4, 6_4, 6_4],
                 'in_channels': 3,
                 'layers_per_block': 1,
                 'out_channels': 3,
                 'time_embedding_type': 'fourier',
                 'norm_eps': 1E-6,
                 'mid_block_scale_factor': math.sqrt(2.0     ),
                 'norm_num_groups': None,
                 'down_block_types': [
                     'SkipDownBlock2D',
                     'AttnSkipDownBlock2D',
                     'SkipDownBlock2D',
                     'SkipDownBlock2D',
                 ],
                 'up_block_types': [
                     'SkipUpBlock2D',
                     'SkipUpBlock2D',
                     'AttnSkipUpBlock2D',
                     'SkipUpBlock2D',
                 ],
             }
             a_					:		List[str]							=  self.dummy_input
             return init_dict, inputs_dict
      @slow
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Union[str, Any]:
             a_							,		a_					:		List[Any]							=  UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'     ,							output_loading_info=SCREAMING_SNAKE_CASE__     )
             self.assertIsNotNone(SCREAMING_SNAKE_CASE__     )
             self.assertEqual(len(loading_info['missing_keys']     )     ,							0     )

             model.to(SCREAMING_SNAKE_CASE__     )
             a_					:		Union[str, Any]							=  self.dummy_input
             a_					:		Dict							=  floats_tensor((4, 3) + (2_5_6, 2_5_6)     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		Optional[Any]							=  noise
             a_					:		List[str]							=  model(**SCREAMING_SNAKE_CASE__     )

             assert image is not None, "Make sure output is not None"
      @slow
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			str:
             a_					:		List[str]							=  UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'     )
             model.to(SCREAMING_SNAKE_CASE__     )

             a_					:		int							=  4
             a_					:		Union[str, Any]							=  3
             a_					:		List[str]							=  (2_5_6, 2_5_6)

             a_					:		Union[str, Any]							=  torch.ones((batch_size, num_channels) + sizes     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		Union[str, Any]							=  torch.tensor(batch_size * [1E-4]     ).to(SCREAMING_SNAKE_CASE__     )

             with torch.no_grad():
                    a_					:		Dict							=  model(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ).sample

             a_					:		Union[str, Any]							=  output[0, -3:, -3:, -1].flatten().cpu()
             # fmt: off
             a_					:		str							=  torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -10980.7129, -20028.8535, 8148.2822, 2342.2905, 567.7608]     )
             # fmt: on

             self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							rtol=1E-2     )     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Any:
             a_					:		str							=  UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'     )
             model.to(SCREAMING_SNAKE_CASE__     )

             a_					:		List[str]							=  4
             a_					:		List[str]							=  3
             a_					:		Union[str, Any]							=  (3_2, 3_2)

             a_					:		Dict							=  torch.ones((batch_size, num_channels) + sizes     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		List[str]							=  torch.tensor(batch_size * [1E-4]     ).to(SCREAMING_SNAKE_CASE__     )

             with torch.no_grad():
                    a_					:		str							=  model(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ).sample

             a_					:		int							=  output[0, -3:, -3:, -1].flatten().cpu()
             # fmt: off
             a_					:		str							=  torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]     )
             # fmt: on

             self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							rtol=1E-2     )     )

      def        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			Optional[int]:
             # not required for this model
             pass






import copy
from typing import Dict, List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


UpperCAmelCase_      : Optional[int]  					=       {
    'facebook/mask2former-swin-small-coco-instance': (
        'https://huggingface.co/facebook/mask2former-swin-small-coco-instance/blob/main/config.json'
    )
    # See all Mask2Former models at https://huggingface.co/models?filter=mask2former
}

UpperCAmelCase_      : List[str]  					=       logging.get_logger(__name__)
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      snake_case__				:     Any        =						'''mask2former'''
      snake_case__				:     Any        =						['''swin''']
      snake_case__				:     str        =						{'''hidden_size''': '''hidden_dim'''}
      def __init__(       self				:       Any     ,							SCREAMING_SNAKE_CASE__				:       Optional[Dict] = None     ,							SCREAMING_SNAKE_CASE__				:       int = 2_5_6     ,							SCREAMING_SNAKE_CASE__				:       int = 2_5_6     ,							SCREAMING_SNAKE_CASE__				:       int = 2_5_6     ,							SCREAMING_SNAKE_CASE__				:       int = 1_0_2_4     ,							SCREAMING_SNAKE_CASE__				:       str = "relu"     ,							SCREAMING_SNAKE_CASE__				:       int = 6     ,							SCREAMING_SNAKE_CASE__				:       int = 1_0     ,							SCREAMING_SNAKE_CASE__				:       int = 8     ,							SCREAMING_SNAKE_CASE__				:       float = 0.0     ,							SCREAMING_SNAKE_CASE__				:       int = 2_0_4_8     ,							SCREAMING_SNAKE_CASE__				:       bool = False     ,							SCREAMING_SNAKE_CASE__				:       bool = False     ,							SCREAMING_SNAKE_CASE__				:       int = 4     ,							SCREAMING_SNAKE_CASE__				:       int = 2_5_5     ,							SCREAMING_SNAKE_CASE__				:       int = 1_0_0     ,							SCREAMING_SNAKE_CASE__				:       float = 0.1     ,							SCREAMING_SNAKE_CASE__				:       float = 2.0     ,							SCREAMING_SNAKE_CASE__				:       float = 5.0     ,							SCREAMING_SNAKE_CASE__				:       float = 5.0     ,							SCREAMING_SNAKE_CASE__				:       int = 1_2_5_4_4     ,							SCREAMING_SNAKE_CASE__				:       float = 3.0     ,							SCREAMING_SNAKE_CASE__				:       float = 0.75     ,							SCREAMING_SNAKE_CASE__				:       float = 0.02     ,							SCREAMING_SNAKE_CASE__				:       float = 1.0     ,							SCREAMING_SNAKE_CASE__				:       bool = True     ,							SCREAMING_SNAKE_CASE__				:       List[int] = [4, 8, 1_6, 3_2]     ,							SCREAMING_SNAKE_CASE__				:       bool = None     ,							**SCREAMING_SNAKE_CASE__				:       int     ,							)		->			List[Any]:
             if backbone_config is None:
                    logger.info('`backbone_config` is `None`. Initializing the config with the default `Swin` backbone.'     )
                    a_					:		Dict							=  CONFIG_MAPPING['swin'](
                        image_size=2_2_4     ,							in_channels=3     ,							patch_size=4     ,							embed_dim=9_6     ,							depths=[2, 2, 1_8, 2]     ,							num_heads=[3, 6, 1_2, 2_4]     ,							window_size=7     ,							drop_path_rate=0.3     ,							use_absolute_embeddings=SCREAMING_SNAKE_CASE__     ,							out_features=['stage1', 'stage2', 'stage3', 'stage4']     ,							)

             if isinstance(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ):
                    a_					:		Any							=  backbone_config.pop('model_type'     )
                    a_					:		Optional[Any]							=  CONFIG_MAPPING[backbone_model_type]
                    a_					:		List[str]							=  config_class.from_dict(SCREAMING_SNAKE_CASE__     )

             # verify that the backbone is supported
             if backbone_config.model_type not in self.backbones_supported:
                    logger.warning_once(
                        F"""Backbone {backbone_config.model_type} is not a supported model and may not be compatible with Mask2Former. """
                        F"""Supported model types: {",".join(self.backbones_supported     )}"""     )

             a_					:		Dict							=  backbone_config
             a_					:		List[str]							=  feature_size
             a_					:		List[str]							=  mask_feature_size
             a_					:		int							=  hidden_dim
             a_					:		Dict							=  encoder_feedforward_dim
             a_					:		str							=  activation_function
             a_					:		List[str]							=  encoder_layers
             a_					:		List[str]							=  decoder_layers
             a_					:		Dict							=  num_attention_heads
             a_					:		str							=  dropout
             a_					:		Tuple							=  dim_feedforward
             a_					:		List[str]							=  pre_norm
             a_					:		Optional[int]							=  enforce_input_projection
             a_					:		Any							=  common_stride
             a_					:		Optional[int]							=  ignore_value
             a_					:		int							=  num_queries
             a_					:		Tuple							=  no_object_weight
             a_					:		Dict							=  class_weight
             a_					:		Optional[int]							=  mask_weight
             a_					:		Optional[int]							=  dice_weight
             a_					:		str							=  train_num_points
             a_					:		List[str]							=  oversample_ratio
             a_					:		List[Any]							=  importance_sample_ratio
             a_					:		Any							=  init_std
             a_					:		Union[str, Any]							=  init_xavier_std
             a_					:		Union[str, Any]							=  use_auxiliary_loss
             a_					:		Dict							=  feature_strides
             a_					:		List[str]							=  output_auxiliary_logits
             a_					:		Dict							=  decoder_layers

             super().__init__(**SCREAMING_SNAKE_CASE__     )
      @classmethod
      def        SCREAMING_SNAKE_CASE							(       cls				:       str     ,							SCREAMING_SNAKE_CASE__				:       PretrainedConfig     ,							**SCREAMING_SNAKE_CASE__				:       Union[str, Any]     )		->			List[Any]:
             return cls(
                 backbone_config=SCREAMING_SNAKE_CASE__     ,							**SCREAMING_SNAKE_CASE__     ,							)

      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Dict[str, any]:
             a_					:		Optional[int]							=  copy.deepcopy(self.__dict__     )
             a_					:		List[Any]							=  self.backbone_config.to_dict()
             a_					:		Optional[Any]							=  self.__class__.model_type
             return output


