



'''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      						=	logging.get_logger(__name__)

enable_full_determinism()






class        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_						, unittest.TestCase     ):
  '''simple docstring'''




  __lowercase							:      int				=      UNetaDModel
  __lowercase							:      Optional[int]				=      '''sample'''




  @property
  def 							UpperCAmelCase_	(    self      )    ->		str:
      lowerCAmelCase__ :  List[str]							      =					4
      lowerCAmelCase__ :  str							      =					3
      lowerCAmelCase__ :  List[str]							      =					(32, 32)

      lowerCAmelCase__ :  Optional[Any]							      =					floats_tensor((batch_size, num_channels) + sizes      ).to(__UpperCAmelCase      )
      lowerCAmelCase__ :  Tuple							      =					torch.tensor([10]      ).to(__UpperCAmelCase      )

      return {"sample": noise, "timestep": time_step}




  @property
  def 							UpperCAmelCase_	(    self      )    ->		List[Any]:
      return (3, 32, 32)




  @property
  def 							UpperCAmelCase_	(    self      )    ->		int:
      return (3, 32, 32)






  def 							UpperCAmelCase_	(    self      )    ->		Dict:
      lowerCAmelCase__ :  List[str]							      =					{
          """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,
      }
      lowerCAmelCase__ :  str							      =					self.dummy_input
      return init_dict, inputs_dict







class        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_						, unittest.TestCase     ):
  '''simple docstring'''




  __lowercase							:      Any				=      UNetaDModel
  __lowercase							:      Tuple				=      '''sample'''




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Any:
      lowerCAmelCase__ :  Any							      =					4
      lowerCAmelCase__ :  Dict							      =					4
      lowerCAmelCase__ :  str							      =					(32, 32)

      lowerCAmelCase__ :  Any							      =					floats_tensor((batch_size, num_channels) + sizes      ).to(__UpperCAmelCase      )
      lowerCAmelCase__ :  Union[str, Any]							      =					torch.tensor([10]      ).to(__UpperCAmelCase      )

      return {"sample": noise, "timestep": time_step}




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Any:
      return (4, 32, 32)




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Union[str, Any]:
      return (4, 32, 32)




  def 							UpperCAmelCase_	(    self      )    ->		Optional[Any]:
      lowerCAmelCase__ :  str							      =					{
          """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"""),
      }
      lowerCAmelCase__ :  List[str]							      =					self.dummy_input
      return init_dict, inputs_dict




  def 							UpperCAmelCase_	(    self      )    ->		List[str]:
      lowerCAmelCase__      ,   lowerCAmelCase__ :  Tuple							      =					UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""  ,output_loading_info=__UpperCAmelCase      )

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

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

      lowerCAmelCase__ :  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      )  ,)
      lowerCAmelCase__ :  List[str]							      =					noise.to(__UpperCAmelCase      )
      lowerCAmelCase__ :  int							      =					torch.tensor([10] * noise.shape[0]      ).to(__UpperCAmelCase      )

      lowerCAmelCase__ :  List[Any]							      =					model_accelerate(__UpperCAmelCase  ,__UpperCAmelCase      )["""sample"""]

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

      lowerCAmelCase__      ,   lowerCAmelCase__ :  Optional[Any]							      =					UNetaDModel.from_pretrained(
          """fusing/unet-ldm-dummy-update"""  ,output_loading_info=__UpperCAmelCase  ,low_cpu_mem_usage=__UpperCAmelCase      )
      model_normal_load.to(__UpperCAmelCase      )
      model_normal_load.eval()
      lowerCAmelCase__ :  Dict							      =					model_normal_load(__UpperCAmelCase  ,__UpperCAmelCase      )["""sample"""]

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






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

      lowerCAmelCase__ :  int							      =					torch.randn(
          1  ,model.config.in_channels  ,model.config.sample_size  ,model.config.sample_size  ,generator=torch.manual_seed(0      )  ,)
      lowerCAmelCase__ :  Optional[int]							      =					noise.to(__UpperCAmelCase      )
      lowerCAmelCase__ :  str							      =					torch.tensor([10] * noise.shape[0]      ).to(__UpperCAmelCase      )

      with torch.no_grad():
          lowerCAmelCase__ :  List[str]							      =					model(__UpperCAmelCase  ,__UpperCAmelCase      ).sample

      lowerCAmelCase__ :  Optional[Any]							      =					output[0, -1, -3:, -3:].flatten().cpu()
      # fmt: off
      lowerCAmelCase__ :  int							      =					torch.tensor([-1_3.3_2_5_8, -2_0.1_1_0_0, -1_5.9_8_7_3, -1_7.6_6_1_7, -2_3.0_5_9_6, -1_7.9_4_1_9, -1_3.3_6_7_5, -1_6.1_8_8_9, -1_2.3_8_0_0]      )
      # fmt: on

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







class        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_						, unittest.TestCase     ):
  '''simple docstring'''




  __lowercase							:      List[str]				=      UNetaDModel
  __lowercase							:      Union[str, Any]				=      '''sample'''




  @property
  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase=(32, 32)      )    ->		Dict:
      lowerCAmelCase__ :  Any							      =					4
      lowerCAmelCase__ :  Optional[Any]							      =					3

      lowerCAmelCase__ :  Union[str, Any]							      =					floats_tensor((batch_size, num_channels) + sizes      ).to(__UpperCAmelCase      )
      lowerCAmelCase__ :  str							      =					torch.tensor(batch_size * [10]      ).to(dtype=torch.intaa  ,device=__UpperCAmelCase      )

      return {"sample": noise, "timestep": time_step}




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Dict:
      return (3, 32, 32)




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Any:
      return (3, 32, 32)




  def 							UpperCAmelCase_	(    self      )    ->		Optional[int]:
      lowerCAmelCase__ :  Union[str, Any]							      =					{
          """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""",
          ],
      }
      lowerCAmelCase__ :  Tuple							      =					self.dummy_input
      return init_dict, inputs_dict




  @slow
  def 							UpperCAmelCase_	(    self      )    ->		str:
      lowerCAmelCase__      ,   lowerCAmelCase__ :  str							      =					UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""  ,output_loading_info=__UpperCAmelCase      )
      self.assertIsNotNone(__UpperCAmelCase      )
      self.assertEqual(len(loading_info["""missing_keys"""]      )  ,0      )

      model.to(__UpperCAmelCase      )
      lowerCAmelCase__ :  Union[str, Any]							      =					self.dummy_input
      lowerCAmelCase__ :  List[Any]							      =					floats_tensor((4, 3) + (256, 256)      ).to(__UpperCAmelCase      )
      lowerCAmelCase__ :  Tuple							      =					noise
      lowerCAmelCase__ :  Union[str, Any]							      =					model(**__UpperCAmelCase      )

      assert image is not None, "Make sure output is not None"




  @slow
  def 							UpperCAmelCase_	(    self      )    ->		str:
      lowerCAmelCase__ :  Optional[Any]							      =					UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""      )
      model.to(__UpperCAmelCase      )

      lowerCAmelCase__ :  Optional[int]							      =					4
      lowerCAmelCase__ :  Any							      =					3
      lowerCAmelCase__ :  Optional[int]							      =					(256, 256)

      lowerCAmelCase__ :  int							      =					torch.ones((batch_size, num_channels) + sizes      ).to(__UpperCAmelCase      )
      lowerCAmelCase__ :  List[Any]							      =					torch.tensor(batch_size * [1E-4]      ).to(__UpperCAmelCase      )

      with torch.no_grad():
          lowerCAmelCase__ :  Dict							      =					model(__UpperCAmelCase  ,__UpperCAmelCase      ).sample

      lowerCAmelCase__ :  int							      =					output[0, -3:, -3:, -1].flatten().cpu()
      # fmt: off
      lowerCAmelCase__ :  Any							      =					torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8]      )
      # fmt: on

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




  def 							UpperCAmelCase_	(    self      )    ->		Tuple:
      lowerCAmelCase__ :  Dict							      =					UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""      )
      model.to(__UpperCAmelCase      )

      lowerCAmelCase__ :  int							      =					4
      lowerCAmelCase__ :  List[str]							      =					3
      lowerCAmelCase__ :  Union[str, Any]							      =					(32, 32)

      lowerCAmelCase__ :  int							      =					torch.ones((batch_size, num_channels) + sizes      ).to(__UpperCAmelCase      )
      lowerCAmelCase__ :  Optional[Any]							      =					torch.tensor(batch_size * [1E-4]      ).to(__UpperCAmelCase      )

      with torch.no_grad():
          lowerCAmelCase__ :  Any							      =					model(__UpperCAmelCase  ,__UpperCAmelCase      ).sample

      lowerCAmelCase__ :  str							      =					output[0, -3:, -3:, -1].flatten().cpu()
      # fmt: off
      lowerCAmelCase__ :  str							      =					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(__UpperCAmelCase  ,__UpperCAmelCase  ,rtol=1E-2      )      )






  def 							UpperCAmelCase_	(    self      )    ->		str:
      # not required for this model
      pass







'''simple docstring'''





import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
 from sagemaker import Session, TrainingJobAnalytics
 from sagemaker.huggingface import HuggingFace






@pytest.mark.skipif(
    literal_eval(os.getenv('''TEST_SAGEMAKER'''						, '''False'''     )     ) is not True						, reason='''Skipping test because should only be run when releasing minor transformers version'''						, )
@pytest.mark.usefixtures('''sm_env'''     )
@parameterized_class(
    [
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_glue.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_5_0, '''eval_accuracy''': 0.7, '''eval_loss''': 0.6},
        },
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_ddp.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.7, '''eval_loss''': 0.6},
        },
        {
            '''framework''': '''tensorflow''',
            '''script''': '''run_tf_dist.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.6, '''eval_loss''': 0.7},
        },
    ]     )
class        lowerCAmelCase_(				unittest.TestCase     ):
  '''simple docstring'''




  def 							UpperCAmelCase_	(    self      )    ->		Optional[Any]:
      if self.framework == "pytorch":
          subprocess.run(
              F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()  ,encoding="""utf-8"""  ,check=__UpperCAmelCase  ,)
      assert hasattr(self  ,"""env"""      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Optional[Any]:
      lowerCAmelCase__ :  Optional[int]							      =					F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
      # distributed data settings
      lowerCAmelCase__ :  Any							      =					{"""smdistributed""": {"""dataparallel""": {"""enabled""": True}}} if self.script != """run_ddp.py""" else None

      # creates estimator
      return HuggingFace(
          entry_point=self.script  ,source_dir=self.env.test_path  ,role=self.env.role  ,image_uri=self.env.image_uri  ,base_job_name=__UpperCAmelCase  ,instance_count=__UpperCAmelCase  ,instance_type=self.instance_type  ,debugger_hook_config=__UpperCAmelCase  ,hyperparameters={**self.env.distributed_hyperparameters, """model_name_or_path""": self.model_name_or_path}  ,metric_definitions=self.env.metric_definitions  ,distribution=__UpperCAmelCase  ,py_version="""py36"""  ,)




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Optional[Any]:
      TrainingJobAnalytics(__UpperCAmelCase      ).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""      )






  @parameterized.expand([(2,)]      )
  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Any:
      # create estimator
      lowerCAmelCase__ :  List[Any]							      =					self.create_estimator(__UpperCAmelCase      )

      # run training
      estimator.fit()

      # result dataframe
      lowerCAmelCase__ :  Any							      =					TrainingJobAnalytics(estimator.latest_training_job.name      ).dataframe()

      # extract kpis
      lowerCAmelCase__ :  int							      =					list(result_metrics_df[result_metrics_df.metric_name == """eval_accuracy"""]["""value"""]      )
      lowerCAmelCase__ :  List[Any]							      =					list(result_metrics_df[result_metrics_df.metric_name == """eval_loss"""]["""value"""]      )
      # get train time from SageMaker job, this includes starting, preprocessing, stopping
      lowerCAmelCase__ :  List[str]							      =					(
          Session().describe_training_job(estimator.latest_training_job.name      ).get("""TrainingTimeInSeconds"""  ,99_9999      )
      )

      # assert kpis
      assert train_runtime <= self.results["train_runtime"]
      assert all(t >= self.results["""eval_accuracy"""] for t in eval_accuracy      )
      assert all(t <= self.results["""eval_loss"""] for t in eval_loss      )

      # dump tests result into json file to share in PR
      with open(F"""{estimator.latest_training_job.name}.json"""  ,"""w"""      ) as outfile:
          json.dump({"""train_time""": train_runtime, """eval_accuracy""": eval_accuracy, """eval_loss""": eval_loss}  ,__UpperCAmelCase      )


