





'''simple docstring'''

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


__lowerCAmelCase							:     Union[str, Any]		     =logging.get_logger(__name__)

enable_full_determinism()






class 					UpperCAmelCase     (	_a      ,	_a      ,	unittest.TestCase    ):
 __lowercase       		=    UNetaDModel
 __lowercase       		=    """sample"""
 @property
 def 					UpperCAmelCase_							(    self				:int							)->							Union[str, Any]:
     A__    =      4
     A__    =      3
     A__    =      (32, 32)

     A__    =      floats_tensor((batch_size, num_channels) + sizes							).to(_a							)
     A__    =      torch.tensor([10]							).to(_a							)

     return {"sample": noise, "timestep": time_step}
 @property
 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							Tuple:
     return (3, 32, 32)
 @property
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							Dict:
     return (3, 32, 32)



 def 					UpperCAmelCase_							(    self				:Tuple							)->							Dict:
     A__    =      {
         "block_out_channels": (32, 64),
         "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": 32,
     }
     A__    =      self.dummy_input
     return init_dict, inputs_dict






class 					UpperCAmelCase     (	_a      ,	_a      ,	unittest.TestCase    ):
 __lowercase       		=    UNetaDModel
 __lowercase       		=    """sample"""
 @property
 def 					UpperCAmelCase_							(    self				:int							)->							str:
     A__    =      4
     A__    =      4
     A__    =      (32, 32)

     A__    =      floats_tensor((batch_size, num_channels) + sizes							).to(_a							)
     A__    =      torch.tensor([10]							).to(_a							)

     return {"sample": noise, "timestep": time_step}
 @property
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							int:
     return (4, 32, 32)
 @property
 def 					UpperCAmelCase_							(    self				:Tuple							)->							Optional[int]:
     return (4, 32, 32)
 def 					UpperCAmelCase_							(    self				:Union[str, Any]							)->							int:
     A__    =      {
         "sample_size": 32,
         "in_channels": 4,
         "out_channels": 4,
         "layers_per_block": 2,
         "block_out_channels": (32, 64),
         "attention_head_dim": 32,
         "down_block_types": ("DownBlock2D", "DownBlock2D"),
         "up_block_types": ("UpBlock2D", "UpBlock2D"),
     }
     A__    =      self.dummy_input
     return init_dict, inputs_dict
 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							int:
     A__    =      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"						,	output_loading_info=_a							)

     self.assertIsNotNone(_a							)
     self.assertEqual(len(loading_info["missing_keys"]							)						,	0							)

     model.to(_a							)
     A__    =      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 					UpperCAmelCase_							(    self				:int							)->							Dict:
     A__    =      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"						,	output_loading_info=_a							)
     model.to(_a							)
     A__    =      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 					UpperCAmelCase_							(    self				:Union[str, Any]							)->							Optional[int]:
     # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
     A__    =      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"						,	output_loading_info=_a							)
     model_accelerate.to(_a							)
     model_accelerate.eval()

     A__    =      torch.randn(
         1						,	model_accelerate.config.in_channels						,	model_accelerate.config.sample_size						,	model_accelerate.config.sample_size						,	generator=torch.manual_seed(0							)						,	)
     A__    =      noise.to(_a							)
     A__    =      torch.tensor([10] * noise.shape[0]							).to(_a							)

     A__    =      model_accelerate(_a						,	_a							)["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__    =      UNetaDModel.from_pretrained(
         "fusing/unet-ldm-dummy-update"						,	output_loading_info=_a						,	low_cpu_mem_usage=_a							)
     model_normal_load.to(_a							)
     model_normal_load.eval()
     A__    =      model_normal_load(_a						,	_a							)["sample"]

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



 def 					UpperCAmelCase_							(    self				:Optional[int]							)->							Optional[Any]:
     A__    =      UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"							)
     model.eval()
     model.to(_a							)

     A__    =      torch.randn(
         1						,	model.config.in_channels						,	model.config.sample_size						,	model.config.sample_size						,	generator=torch.manual_seed(0							)						,	)
     A__    =      noise.to(_a							)
     A__    =      torch.tensor([10] * noise.shape[0]							).to(_a							)

     with torch.no_grad():
         A__    =      model(_a						,	_a							).sample

     A__    =      output[0, -1, -3:, -3:].flatten().cpu()
     # fmt: off
     A__    =      torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]							)
     # fmt: on

     self.assertTrue(torch_all_close(_a						,	_a						,	rtol=1E-3							)							)






class 					UpperCAmelCase     (	_a      ,	_a      ,	unittest.TestCase    ):
 __lowercase       		=    UNetaDModel
 __lowercase       		=    """sample"""
 @property
 def 					UpperCAmelCase_							(    self				:Optional[int]						,	lowercase_				:List[str]=(32, 32)							)->							int:
     A__    =      4
     A__    =      3

     A__    =      floats_tensor((batch_size, num_channels) + sizes							).to(_a							)
     A__    =      torch.tensor(batch_size * [10]							).to(dtype=torch.intaa						,	device=_a							)

     return {"sample": noise, "timestep": time_step}
 @property
 def 					UpperCAmelCase_							(    self				:Optional[int]							)->							int:
     return (3, 32, 32)
 @property
 def 					UpperCAmelCase_							(    self				:Any							)->							Tuple:
     return (3, 32, 32)
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							Dict:
     A__    =      {
         "block_out_channels": [32, 64, 64, 64],
         "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__    =      self.dummy_input
     return init_dict, inputs_dict
 @slow
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							str:
     A__    =      UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"						,	output_loading_info=_a							)
     self.assertIsNotNone(_a							)
     self.assertEqual(len(loading_info["missing_keys"]							)						,	0							)

     model.to(_a							)
     A__    =      self.dummy_input
     A__    =      floats_tensor((4, 3) + (2_56, 2_56)							).to(_a							)
     A__    =      noise
     A__    =      model(**_a							)

     assert image is not None, "Make sure output is not None"
 @slow
 def 					UpperCAmelCase_							(    self				:Dict							)->							Any:
     A__    =      UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"							)
     model.to(_a							)

     A__    =      4
     A__    =      3
     A__    =      (2_56, 2_56)

     A__    =      torch.ones((batch_size, num_channels) + sizes							).to(_a							)
     A__    =      torch.tensor(batch_size * [1E-4]							).to(_a							)

     with torch.no_grad():
         A__    =      model(_a						,	_a							).sample

     A__    =      output[0, -3:, -3:, -1].flatten().cpu()
     # fmt: off
     A__    =      torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]							)
     # fmt: on

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

     A__    =      4
     A__    =      3
     A__    =      (32, 32)

     A__    =      torch.ones((batch_size, num_channels) + sizes							).to(_a							)
     A__    =      torch.tensor(batch_size * [1E-4]							).to(_a							)

     with torch.no_grad():
         A__    =      model(_a						,	_a							).sample

     A__    =      output[0, -3:, -3:, -1].flatten().cpu()
     # fmt: off
     A__    =      torch.tensor([-0.0_3_2_5, -0.0_9_0_0, -0.0_8_6_9, -0.0_3_3_2, -0.0_7_2_5, -0.0_2_7_0, -0.0_1_0_1, 0.0_2_2_7, 0.0_2_5_6]							)
     # fmt: on

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



 def 					UpperCAmelCase_							(    self				:Any							)->							List[Any]:
     # not required for this model
     pass










'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__lowerCAmelCase							:     Any		     ={
    "configuration_swinv2": ["SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP", "Swinv2Config"],
}


try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __lowerCAmelCase							:     Tuple		     =[
      "SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST",
      "Swinv2ForImageClassification",
      "Swinv2ForMaskedImageModeling",
      "Swinv2Model",
      "Swinv2PreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_swinva import SWINV2_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinvaConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_swinva import (
        SWINV2_PRETRAINED_MODEL_ARCHIVE_LIST,
        SwinvaForImageClassification,
        SwinvaForMaskedImageModeling,
        SwinvaModel,
        SwinvaPreTrainedModel,
    )


else:
  import sys

  __lowerCAmelCase							:     List[Any]		     =_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



