





import json
import os
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from requests.exceptions import HTTPError

from transformers.utils import (
    CONFIG_NAME,
    FLAX_WEIGHTS_NAME,
    TF2_WEIGHTS_NAME,
    TRANSFORMERS_CACHE,
    WEIGHTS_NAME,
    cached_file,
    get_file_from_repo,
    has_file,
)


__A									=				"hf-internal-testing/tiny-random-bert"
__A									=				os.path.join(TRANSFORMERS_CACHE, "models--hf-internal-testing--tiny-random-bert")
__A									=				"9b8c223d42b2188cb49d29af482996f9d0f3e5a6"



class    _SCREAMING_SNAKE_CASE						(	unittest.TestCase    ):







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =cached_file(UpperCAmelCase_       , UpperCAmelCase_)
   # Should have downloaded the file in here
   self.assertTrue(os.path.isdir(UpperCAmelCase_))
   # Cache should contain at least those three subfolders:
   for subfolder in ["blobs", "refs", "snapshots"]:
     self.assertTrue(os.path.isdir(os.path.join(UpperCAmelCase_       , UpperCAmelCase_)))
   with open(os.path.join(UpperCAmelCase_       , "refs"       , "main")) as f:
     lowerCamelCase__:      List[Any]			  =f.read()
   self.assertEqual(UpperCAmelCase_       , os.path.join(UpperCAmelCase_       , "snapshots"       , UpperCAmelCase_       , UpperCAmelCase_))
   self.assertTrue(os.path.isfile(UpperCAmelCase_))

   # File is cached at the same place the second time.
   lowerCamelCase__:      Optional[Any]			  =cached_file(UpperCAmelCase_       , UpperCAmelCase_)
   self.assertEqual(UpperCAmelCase_       , UpperCAmelCase_)

   # Using a specific revision to test the full commit hash.
   lowerCamelCase__:      str			  =cached_file(UpperCAmelCase_       , UpperCAmelCase_       , revision="9b8c223")
   self.assertEqual(UpperCAmelCase_       , os.path.join(UpperCAmelCase_       , "snapshots"       , UpperCAmelCase_       , UpperCAmelCase_))

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->str:



   '''simple docstring'''



   with self.assertRaisesRegex(UpperCAmelCase_       , "is not a valid model identifier"):
     lowerCamelCase__:      List[Any]			  =cached_file("tiny-random-bert"       , UpperCAmelCase_)

   with self.assertRaisesRegex(UpperCAmelCase_       , "is not a valid git identifier"):
     lowerCamelCase__:      List[str]			  =cached_file(UpperCAmelCase_       , UpperCAmelCase_       , revision="aaaa")

   with self.assertRaisesRegex(UpperCAmelCase_       , "does not appear to have a file named"):
     lowerCamelCase__:      Optional[Any]			  =cached_file(UpperCAmelCase_       , "conf")

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->List[str]:



   '''simple docstring'''



   with self.assertRaisesRegex(UpperCAmelCase_       , "does not appear to have a file named"):
     lowerCamelCase__:      Dict			  =cached_file(UpperCAmelCase_       , "conf")

   with open(os.path.join(UpperCAmelCase_       , "refs"       , "main")) as f:
     lowerCamelCase__:      List[str]			  =f.read()
   self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase_       , ".no_exist"       , UpperCAmelCase_       , "conf")))

   lowerCamelCase__:      Union[str, Any]			  =cached_file(UpperCAmelCase_       , "conf"       , _raise_exceptions_for_missing_entries=UpperCAmelCase_)
   self.assertIsNone(UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =cached_file(UpperCAmelCase_       , "conf"       , local_files_only=UpperCAmelCase_       , _raise_exceptions_for_missing_entries=UpperCAmelCase_)
   self.assertIsNone(UpperCAmelCase_)

   lowerCamelCase__:      List[Any]			  =mock.Mock()
   lowerCamelCase__:      int			  =500
   lowerCamelCase__:      Union[str, Any]			  ={}
   lowerCamelCase__:      Any			  =HTTPError
   lowerCamelCase__:      List[Any]			  ={}

   # Under the mock environment we get a 500 error when trying to reach the tokenizer.
   with mock.patch("requests.Session.request"       , return_value=UpperCAmelCase_) as mock_head:
     lowerCamelCase__:      int			  =cached_file(UpperCAmelCase_       , "conf"       , _raise_exceptions_for_connection_errors=UpperCAmelCase_)
     self.assertIsNone(UpperCAmelCase_)
     # This check we did call the fake head request
     mock_head.assert_called()

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Tuple:



   '''simple docstring'''



   self.assertTrue(has_file("hf-internal-testing/tiny-bert-pt-only"       , UpperCAmelCase_))
   self.assertFalse(has_file("hf-internal-testing/tiny-bert-pt-only"       , UpperCAmelCase_))
   self.assertFalse(has_file("hf-internal-testing/tiny-bert-pt-only"       , UpperCAmelCase_))

 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->List[str]:



   '''simple docstring'''



   self.assertIsNone(get_file_from_repo("bert-base-cased"       , "ahah.txt"))

   # The function raises if the repository does not exist.
   with self.assertRaisesRegex(UpperCAmelCase_       , "is not a valid model identifier"):
     get_file_from_repo("bert-base-case"       , UpperCAmelCase_)

   # The function raises if the revision does not exist.
   with self.assertRaisesRegex(UpperCAmelCase_       , "is not a valid git identifier"):
     get_file_from_repo("bert-base-cased"       , UpperCAmelCase_       , revision="ahaha")

   lowerCamelCase__:      List[str]			  =get_file_from_repo("bert-base-cased"       , UpperCAmelCase_)
   # The name is the cached name which is not very easy to test, so instead we load the content.
   lowerCamelCase__:      List[str]			  =json.loads(open(UpperCAmelCase_       , "r").read())
   self.assertEqual(config["hidden_size"]       , 768)







 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->int:



   '''simple docstring'''



   with tempfile.TemporaryDirectory() as tmp_dir:
     lowerCamelCase__:      List[Any]			  =Path(UpperCAmelCase_) / "a.txt"
     filename.touch()
     self.assertEqual(get_file_from_repo(UpperCAmelCase_       , "a.txt")       , str(UpperCAmelCase_))

     self.assertIsNone(get_file_from_repo(UpperCAmelCase_       , "b.txt"))








import torch

from diffusers import DDPMParallelScheduler

from .test_schedulers import SchedulerCommonTest



class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''
 lowercase_		=			(DDPMParallelScheduler,)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any       , **UpperCAmelCase_				:   Any)       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Any			  ={
       "num_train_timesteps": 1_000,
       "beta_start": 0.0001,
       "beta_end": 0.02,
       "beta_schedule": "linear",
       "variance_type": "fixed_small",
       "clip_sample": True,
   }

   config.update(**UpperCAmelCase_)
   return config

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Dict:



   '''simple docstring'''



   for timesteps in [1, 5, 100, 1_000]:
     self.check_over_configs(num_train_timesteps=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Optional[int]:



   '''simple docstring'''



   for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1]       , [0.002, 0.02, 0.2, 2]):
     self.check_over_configs(beta_start=UpperCAmelCase_       , beta_end=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Any:



   '''simple docstring'''



   for schedule in ["linear", "squaredcos_cap_v2"]:
     self.check_over_configs(beta_schedule=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Optional[int]:



   '''simple docstring'''



   for variance in ["fixed_small", "fixed_large", "other"]:
     self.check_over_configs(variance_type=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->Optional[Any]:



   '''simple docstring'''



   for clip_sample in [True, False]:
     self.check_over_configs(clip_sample=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Tuple:



   '''simple docstring'''



   self.check_over_configs(thresholding=UpperCAmelCase_)
   for threshold in [0.5, 1.0, 2.0]:
     for prediction_type in ["epsilon", "sample", "v_prediction"]:
       self.check_over_configs(
           thresholding=UpperCAmelCase_       , prediction_type=UpperCAmelCase_       , sample_max_value=UpperCAmelCase_       , )

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Optional[int]:



   '''simple docstring'''



   for prediction_type in ["epsilon", "sample", "v_prediction"]:
     self.check_over_configs(prediction_type=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->int:



   '''simple docstring'''



   for t in [0, 500, 999]:
     self.check_over_forward(time_step=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->str:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =self.scheduler_classes[0]
   lowerCamelCase__:      Tuple			  =self.get_scheduler_config()
   lowerCamelCase__:      Any			  =scheduler_class(**UpperCAmelCase_)

   assert torch.sum(torch.abs(scheduler._get_variance(0) - 0.0)) < 1E-5
   assert torch.sum(torch.abs(scheduler._get_variance(487) - 0.0_0979)) < 1E-5
   assert torch.sum(torch.abs(scheduler._get_variance(999) - 0.02)) < 1E-5

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->str:



   '''simple docstring'''



   lowerCamelCase__:      int			  =self.scheduler_classes[0]
   lowerCamelCase__:      Tuple			  =self.get_scheduler_config()
   lowerCamelCase__:      Tuple			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      str			  =len(UpperCAmelCase_)

   lowerCamelCase__:      Optional[int]			  =self.dummy_model()
   lowerCamelCase__:      int			  =self.dummy_sample_deter
   lowerCamelCase__:      Union[str, Any]			  =self.dummy_sample_deter + 0.1
   lowerCamelCase__:      Optional[Any]			  =self.dummy_sample_deter - 0.1

   lowerCamelCase__:      Optional[Any]			  =samplea.shape[0]
   lowerCamelCase__:      List[Any]			  =torch.stack([samplea, samplea, samplea]       , dim=0)
   lowerCamelCase__:      Union[str, Any]			  =torch.arange(UpperCAmelCase_)[0:3, None].repeat(1       , UpperCAmelCase_)

   lowerCamelCase__:      Optional[int]			  =model(samples.flatten(0       , 1)       , timesteps.flatten(0       , 1))
   lowerCamelCase__:      Tuple			  =scheduler.batch_step_no_noise(UpperCAmelCase_       , timesteps.flatten(0       , 1)       , samples.flatten(0       , 1))

   lowerCamelCase__:      List[str]			  =torch.sum(torch.abs(UpperCAmelCase_))
   lowerCamelCase__:      Any			  =torch.mean(torch.abs(UpperCAmelCase_))

   assert abs(result_sum.item() - 1153.1833) < 1E-2
   assert abs(result_mean.item() - 0.5005) < 1E-3

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =self.scheduler_classes[0]
   lowerCamelCase__:      Optional[Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Optional[int]			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =len(UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =self.dummy_model()
   lowerCamelCase__:      List[Any]			  =self.dummy_sample_deter
   lowerCamelCase__:      int			  =torch.manual_seed(0)

   for t in reversed(range(UpperCAmelCase_)):
     # 1. predict noise residual
     lowerCamelCase__:      Tuple			  =model(UpperCAmelCase_       , UpperCAmelCase_)

     # 2. predict previous mean of sample x_t-1
     lowerCamelCase__:      Optional[Any]			  =scheduler.step(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , generator=UpperCAmelCase_).prev_sample

     lowerCamelCase__:      Any			  =pred_prev_sample

   lowerCamelCase__:      Any			  =torch.sum(torch.abs(UpperCAmelCase_))
   lowerCamelCase__:      List[str]			  =torch.mean(torch.abs(UpperCAmelCase_))

   assert abs(result_sum.item() - 258.9606) < 1E-2
   assert abs(result_mean.item() - 0.3372) < 1E-3

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =self.scheduler_classes[0]
   lowerCamelCase__:      Any			  =self.get_scheduler_config(prediction_type="v_prediction")
   lowerCamelCase__:      Any			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      str			  =len(UpperCAmelCase_)

   lowerCamelCase__:      str			  =self.dummy_model()
   lowerCamelCase__:      str			  =self.dummy_sample_deter
   lowerCamelCase__:      Dict			  =torch.manual_seed(0)

   for t in reversed(range(UpperCAmelCase_)):
     # 1. predict noise residual
     lowerCamelCase__:      Union[str, Any]			  =model(UpperCAmelCase_       , UpperCAmelCase_)

     # 2. predict previous mean of sample x_t-1
     lowerCamelCase__:      Dict			  =scheduler.step(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , generator=UpperCAmelCase_).prev_sample

     lowerCamelCase__:      List[str]			  =pred_prev_sample

   lowerCamelCase__:      List[Any]			  =torch.sum(torch.abs(UpperCAmelCase_))
   lowerCamelCase__:      Tuple			  =torch.mean(torch.abs(UpperCAmelCase_))

   assert abs(result_sum.item() - 202.0296) < 1E-2
   assert abs(result_mean.item() - 0.2631) < 1E-3

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple)       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      str			  =self.scheduler_classes[0]
   lowerCamelCase__:      Union[str, Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Any			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      List[Any]			  =[100, 87, 50, 1, 0]

   scheduler.set_timesteps(timesteps=UpperCAmelCase_)

   lowerCamelCase__:      Union[str, Any]			  =scheduler.timesteps

   for i, timestep in enumerate(UpperCAmelCase_):
     if i == len(UpperCAmelCase_) - 1:
       lowerCamelCase__:      Dict			  =-1
     else:
       lowerCamelCase__:      Union[str, Any]			  =timesteps[i + 1]

     lowerCamelCase__:      Tuple			  =scheduler.previous_timestep(UpperCAmelCase_)
     lowerCamelCase__:      str			  =prev_t.item()

     self.assertEqual(UpperCAmelCase_       , UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =self.scheduler_classes[0]
   lowerCamelCase__:      List[Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Dict			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Optional[Any]			  =[100, 87, 50, 51, 0]

   with self.assertRaises(UpperCAmelCase_       , msg="`custom_timesteps` must be in descending order."):
     scheduler.set_timesteps(timesteps=UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =self.scheduler_classes[0]
   lowerCamelCase__:      Any			  =self.get_scheduler_config()
   lowerCamelCase__:      int			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Optional[int]			  =[100, 87, 50, 1, 0]
   lowerCamelCase__:      int			  =len(UpperCAmelCase_)

   with self.assertRaises(UpperCAmelCase_       , msg="Can only pass one of `num_inference_steps` or `custom_timesteps`."):
     scheduler.set_timesteps(num_inference_steps=UpperCAmelCase_       , timesteps=UpperCAmelCase_)







 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =self.scheduler_classes[0]
   lowerCamelCase__:      Optional[Any]			  =self.get_scheduler_config()
   lowerCamelCase__:      Optional[Any]			  =scheduler_class(**UpperCAmelCase_)

   lowerCamelCase__:      Dict			  =[scheduler.config.num_train_timesteps]

   with self.assertRaises(
       UpperCAmelCase_       , msg="`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}"       , ):
     scheduler.set_timesteps(timesteps=UpperCAmelCase_)

