






"""simple docstring"""

import unittest

import torch

from diffusers import DDIMScheduler, DDPMScheduler, UNetaDModel
from diffusers.training_utils import set_seed
from diffusers.utils.testing_utils import slow


lowercase__					:						List[Any]												=			False




class  _UpperCAmelCase (    unittest.TestCase):

 def 		_snake_case      (  self	: Optional[int]						,  lowercase_	: str=32       ):
     set_seed(0       )
     snake_case_				:  Any       =			UNetaDModel(sample_size=lowercase_						,  in_channels=3						,  out_channels=3       )
     snake_case_				:  Optional[int]       =			torch.optim.SGD(model.parameters()						,  lr=0.00_01       )
     return model, optimizer






 @slow
 def 		_snake_case      (  self	: str       ):
     snake_case_				:  Dict       =			'''cpu'''  # ensure full determinism without setting the CUBLAS_WORKSPACE_CONFIG env variable
     snake_case_				:  List[str]       =			DDPMScheduler(
         num_train_timesteps=1000						,  beta_start=0.00_01						,  beta_end=0.02						,  beta_schedule='''linear'''						,  clip_sample=lowercase_						,  )
     snake_case_				:  Union[str, Any]       =			DDIMScheduler(
         num_train_timesteps=1000						,  beta_start=0.00_01						,  beta_end=0.02						,  beta_schedule='''linear'''						,  clip_sample=lowercase_						,  )

     assert ddpm_scheduler.config.num_train_timesteps == ddim_scheduler.config.num_train_timesteps

     # shared batches for DDPM and DDIM
     set_seed(0       )
     snake_case_				:  List[str]       =			[torch.randn((4, 3, 32, 32)       ).clip(-1						,  1       ).to(lowercase_       ) for _ in range(4       )]
     snake_case_				:  str       =			[torch.randn((4, 3, 32, 32)       ).to(lowercase_       ) for _ in range(4       )]
     snake_case_				:  Optional[int]       =			[torch.randint(0						,  1000						,  (4,)       ).long().to(lowercase_       ) for _ in range(4       )]

     # train with a DDPM scheduler
     snake_case_				:  Optional[int]       =			self.get_model_optimizer(resolution=32       )
     model.train().to(lowercase_       )
     for i in range(4       ):
         optimizer.zero_grad()
         snake_case_				:  Union[str, Any]       =			ddpm_scheduler.add_noise(clean_images[i]						,  noise[i]						,  timesteps[i]       )
         snake_case_				:  Any       =			model(lowercase_						,  timesteps[i]       ).sample
         snake_case_				:  int       =			torch.nn.functional.mse_loss(lowercase_						,  noise[i]       )
         loss.backward()
         optimizer.step()
     del model, optimizer

     # recreate the model and optimizer, and retry with DDIM
     snake_case_				:  List[str]       =			self.get_model_optimizer(resolution=32       )
     model.train().to(lowercase_       )
     for i in range(4       ):
         optimizer.zero_grad()
         snake_case_				:  Tuple       =			ddim_scheduler.add_noise(clean_images[i]						,  noise[i]						,  timesteps[i]       )
         snake_case_				:  Union[str, Any]       =			model(lowercase_						,  timesteps[i]       ).sample
         snake_case_				:  int       =			torch.nn.functional.mse_loss(lowercase_						,  noise[i]       )
         loss.backward()
         optimizer.step()
     del model, optimizer

     self.assertTrue(torch.allclose(lowercase_						,  lowercase_						,  atol=1E-5       )       )
     self.assertTrue(torch.allclose(lowercase_						,  lowercase_						,  atol=1E-5       )       )










"""simple docstring"""

import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / '''utils'''))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


lowercase__					:						Any												=			get_tests_dir('''fixtures''')




class  _UpperCAmelCase (    unittest.TestCase):

 def 		_snake_case      (  self	: Tuple       ):
     # A mock response for an HTTP head request to emulate server down
     snake_case_				:  str       =			mock.Mock()
     snake_case_				:  Optional[Any]       =			500
     snake_case_				:  str       =			{}
     snake_case_				:  Optional[int]       =			HTTPError
     snake_case_				:  Tuple       =			{}

     # Download this model to make sure it's in the cache.
     snake_case_				:  List[str]       =			WavaVecaFeatureExtractor.from_pretrained('''hf-internal-testing/tiny-random-wav2vec2'''       )
     # Under the mock environment we get a 500 error when trying to reach the model.
     with mock.patch('''requests.Session.request'''						,  return_value=lowercase_       ) as mock_head:
         snake_case_				:  int       =			WavaVecaFeatureExtractor.from_pretrained('''hf-internal-testing/tiny-random-wav2vec2'''       )
         # This check we did call the fake head request
         mock_head.assert_called()






 def 		_snake_case      (  self	: Union[str, Any]       ):
     # This test is for deprecated behavior and can be removed in v5
     snake_case_				:  str       =			WavaVecaFeatureExtractor.from_pretrained(
         '''https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json'''       )





@is_staging_test
class  _UpperCAmelCase (    unittest.TestCase):

 @classmethod
 def 		_snake_case      (  cls	: int       ):
     snake_case_				:  Dict       =			TOKEN
     HfFolder.save_token(lowercase_       )

 @classmethod
 def 		_snake_case      (  cls	: str       ):
     try:
         delete_repo(token=cls._token						,  repo_id='''test-feature-extractor'''       )
     except HTTPError:
         pass

     try:
         delete_repo(token=cls._token						,  repo_id='''valid_org/test-feature-extractor-org'''       )
     except HTTPError:
         pass

     try:
         delete_repo(token=cls._token						,  repo_id='''test-dynamic-feature-extractor'''       )
     except HTTPError:
         pass

 def 		_snake_case      (  self	: Any       ):
     snake_case_				:  int       =			WavaVecaFeatureExtractor.from_pretrained(lowercase_       )
     feature_extractor.push_to_hub('''test-feature-extractor'''						,  use_auth_token=self._token       )

     snake_case_				:  Optional[Any]       =			WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )

     # Reset repo
     delete_repo(token=self._token						,  repo_id='''test-feature-extractor'''       )

     # Push to hub via save_pretrained
     with tempfile.TemporaryDirectory() as tmp_dir:
         feature_extractor.save_pretrained(
             lowercase_						,  repo_id='''test-feature-extractor'''						,  push_to_hub=lowercase_						,  use_auth_token=self._token       )

     snake_case_				:  Optional[int]       =			WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )

 def 		_snake_case      (  self	: Dict       ):
     snake_case_				:  Union[str, Any]       =			WavaVecaFeatureExtractor.from_pretrained(lowercase_       )
     feature_extractor.push_to_hub('''valid_org/test-feature-extractor'''						,  use_auth_token=self._token       )

     snake_case_				:  Optional[int]       =			WavaVecaFeatureExtractor.from_pretrained('''valid_org/test-feature-extractor'''       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )

     # Reset repo
     delete_repo(token=self._token						,  repo_id='''valid_org/test-feature-extractor'''       )

     # Push to hub via save_pretrained
     with tempfile.TemporaryDirectory() as tmp_dir:
         feature_extractor.save_pretrained(
             lowercase_						,  repo_id='''valid_org/test-feature-extractor-org'''						,  push_to_hub=lowercase_						,  use_auth_token=self._token       )

     snake_case_				:  int       =			WavaVecaFeatureExtractor.from_pretrained('''valid_org/test-feature-extractor-org'''       )
     for k, v in feature_extractor.__dict__.items():
         self.assertEqual(lowercase_						,  getattr(lowercase_						,  lowercase_       )       )






 def 		_snake_case      (  self	: str       ):
     CustomFeatureExtractor.register_for_auto_class()
     snake_case_				:  Optional[Any]       =			CustomFeatureExtractor.from_pretrained(lowercase_       )

     feature_extractor.push_to_hub('''test-dynamic-feature-extractor'''						,  use_auth_token=self._token       )

     # This has added the proper auto_map field to the config
     self.assertDictEqual(
         feature_extractor.auto_map						,  {'''AutoFeatureExtractor''': '''custom_feature_extraction.CustomFeatureExtractor'''}						,  )

     snake_case_				:  Any       =			AutoFeatureExtractor.from_pretrained(
         f"{USER}/test-dynamic-feature-extractor"						,  trust_remote_code=lowercase_       )
     # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
     self.assertEqual(new_feature_extractor.__class__.__name__						,  '''CustomFeatureExtractor'''       )





