



'''simple docstring'''





import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import DecisionTransformerModel
  from transformers.models.decision_transformer.modeling_decision_transformer import (
      DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
  )






class __A     :



 '''simple docstring'''







 def __init__(self					,  A					,  A=13					,  A=7					,  A=6					,  A=17					,  A=23					,  A=11					,  A=True					,  )       ->	Tuple:


      """simple docstring"""



      _a           =    parent
      _a           =    batch_size
      _a           =    seq_length
      _a           =    act_dim
      _a           =    state_dim
      _a           =    hidden_size
      _a           =    max_length
      _a           =    is_training



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


      """simple docstring"""



      _a           =    floats_tensor((self.batch_size, self.seq_length, self.state_dim)			)
      _a           =    floats_tensor((self.batch_size, self.seq_length, self.act_dim)			)
      _a           =    floats_tensor((self.batch_size, self.seq_length, 1)			)
      _a           =    floats_tensor((self.batch_size, self.seq_length, 1)			)
      _a           =    ids_tensor((self.batch_size, self.seq_length)					,  vocab_size=1_000			)
      _a           =    random_attention_mask((self.batch_size, self.seq_length)			)

      _a           =    self.get_config()

      return (
          config,
          states,
          actions,
          rewards,
          returns_to_go,
          timesteps,
          attention_mask,
      )



 def    a__    (self			)       ->	str:


      """simple docstring"""



      return DecisionTransformerConfig(
          batch_size=self.batch_size					,  seq_length=self.seq_length					,  act_dim=self.act_dim					,  state_dim=self.state_dim					,  hidden_size=self.hidden_size					,  max_length=self.max_length					,  )



 def    a__    (self					,  A					,  A					,  A					,  A					,  A					,  A					,  A					,  )       ->	List[Any]:


      """simple docstring"""



      _a           =    DecisionTransformerModel(config=A			)
      model.to(A			)
      model.eval()
      _a           =    model(A					,  A					,  A					,  A					,  A					,  A			)

      self.parent.assertEqual(result.state_preds.shape					,  states.shape			)
      self.parent.assertEqual(result.action_preds.shape					,  actions.shape			)
      self.parent.assertEqual(result.return_preds.shape					,  returns_to_go.shape			)
      self.parent.assertEqual(
          result.last_hidden_state.shape					,  (self.batch_size, self.seq_length * 3, self.hidden_size)			)  # seq length *3 as there are 3 modelities: states, returns and actions



 def    a__    (self			)       ->	Dict:


      """simple docstring"""



      _a           =    self.prepare_config_and_inputs()
      (
          (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       (
          _a
      )    ,       
      )           =    config_and_inputs
      _a           =    {
          '''states''': states,
          '''actions''': actions,
          '''rewards''': rewards,
          '''returns_to_go''': returns_to_go,
          '''timesteps''': timesteps,
          '''attention_mask''': attention_mask,
      }
      return config, inputs_dict






@require_torch
class __A     (    A  ,  A  ,  A  ,  unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase       :			Optional[Any]		=			(DecisionTransformerModel,) if is_torch_available() else ()
 __lowerCamelCase       :			List[str]		=			()
 __lowerCamelCase       :			Tuple		=			{'feature-extraction': DecisionTransformerModel} if is_torch_available() else {}

 # Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
 __lowerCamelCase       :			str		=			False

 # Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
 __lowerCamelCase       :			List[str]		=			False
 __lowerCamelCase       :			List[str]		=			False
 __lowerCamelCase       :			Tuple		=			False
 __lowerCamelCase       :			str		=			False
 __lowerCamelCase       :			Dict		=			False
 __lowerCamelCase       :			Tuple		=			False
 __lowerCamelCase       :			Tuple		=			False
 __lowerCamelCase       :			Dict		=			False
 __lowerCamelCase       :			List[str]		=			False



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


      """simple docstring"""



      _a           =    DecisionTransformerModelTester(self			)
      _a           =    ConfigTester(self					,  config_class=A					,  hidden_size=37			)



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


      """simple docstring"""



      self.config_tester.run_common_tests()



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


      """simple docstring"""



      _a           =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A			)



 @slow
 def    a__    (self			)       ->	Optional[Any]:


      """simple docstring"""



      for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _a           =    DecisionTransformerModel.from_pretrained(A			)
           self.assertIsNotNone(A			)



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


      """simple docstring"""



      _a    ,       _a           =    self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           _a           =    model_class(A			)
           _a           =    inspect.signature(model.forward			)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           _a           =    [*signature.parameters.keys()]

           _a           =    [
               '''states''',
               '''actions''',
               '''rewards''',
               '''returns_to_go''',
               '''timesteps''',
               '''attention_mask''',
           ]

           self.assertListEqual(arg_names[: len(A			)]					,  A			)






@require_torch
class __A     (    unittest.TestCase		):



 '''simple docstring'''







 @slow
 def    a__    (self			)       ->	Optional[Any]:


      """simple docstring"""



      _a           =    2  # number of steps of autoregressive prediction we will perform
      _a           =    10  # defined by the RL environment, may be normalized
      _a           =    DecisionTransformerModel.from_pretrained('''edbeeching/decision-transformer-gym-hopper-expert'''			)
      _a           =    model.to(A			)
      _a           =    model.config
      torch.manual_seed(0			)
      _a           =    torch.randn(1					,  1					,  config.state_dim			).to(device=A					,  dtype=torch.floataa			)  # env.reset()

      _a           =    torch.tensor(
          [[0.242793, -0.28693074, 0.8742613], [0.67815274, -0.08101085, -0.12952147]]					,  device=A			)

      _a           =    torch.tensor(A					,  device=A					,  dtype=torch.floataa			).reshape(1					,  1					,  1			)
      _a           =    state
      _a           =    torch.zeros(1					,  0					,  config.act_dim					,  device=A					,  dtype=torch.floataa			)
      _a           =    torch.zeros(1					,  0					,  device=A					,  dtype=torch.floataa			)
      _a           =    torch.tensor(0					,  device=A					,  dtype=torch.long			).reshape(1					,  1			)

      for step in range(A			):
           _a           =    torch.cat([actions, torch.zeros(1					,  1					,  config.act_dim					,  device=A			)]					,  dim=1			)
           _a           =    torch.cat([rewards, torch.zeros(1					,  1					,  device=A			)]					,  dim=1			)

           _a           =    torch.ones(1					,  states.shape[1]			).to(dtype=torch.long					,  device=states.device			)

           with torch.no_grad():
                _a    ,       _a    ,       _a           =    model(
                    states=A					,  actions=A					,  rewards=A					,  returns_to_go=A					,  timesteps=A					,  attention_mask=A					,  return_dict=A					,  )

           self.assertEqual(action_pred.shape					,  actions.shape			)
           self.assertTrue(torch.allclose(action_pred[0, -1]					,  expected_outputs[step]					,  atol=1E-4			)			)
           _a    ,       _a    ,       _a    ,       _a           =    (  # env.step(action)
               torch.randn(1					,  1					,  config.state_dim			).to(device=A					,  dtype=torch.floataa			),
               1.0,
               False,
               {},
           )

           _a           =    action_pred[0, -1]
           _a           =    torch.cat([states, state]					,  dim=1			)
           _a           =    returns_to_go[0, -1] - reward
           _a           =    torch.cat([returns_to_go, pred_return.reshape(1					,  1					,  1			)]					,  dim=1			)
           _a           =    torch.cat(
               [timesteps, torch.ones((1, 1)					,  device=A					,  dtype=torch.long			) * (step + 1)]					,  dim=1			)







import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class        __lowercase       (							a_			,							unittest.TestCase  ):



		"""simple docstring"""

		UpperCamelCase     :					List[str]			       =		ConsistencyModelPipeline
		UpperCamelCase     :					int			       =		UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		UpperCamelCase     :					Any			       =		UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

		# Override required_optional_params to remove num_images_per_prompt
		UpperCamelCase     :					Optional[int]			       =		frozenset(
		    [
		        "num_inference_steps",
		        "generator",
		        "latents",
		        "output_type",
		        "return_dict",
		        "callback",
		        "callback_steps",
		    ]  )







		@property
		def 					__A   ( self	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained(
					    """diffusers/consistency-models-test"""		,      subfolder="""test_unet"""		,      )
					return unet







		@property
		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained(
					    """diffusers/consistency-models-test"""		,      subfolder="""test_unet_class_cond"""		,      )
					return unet







		def 					__A   ( self		,      A=False	)  -> Tuple:







					'''simple docstring'''







					if class_cond:
								lowerCamelCase    		=  self.dummy_cond_unet
					else:
								lowerCamelCase    		=  self.dummy_uncond_unet

					# Default to CM multistep sampler
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )

					lowerCamelCase    		=  {
					    """unet""": unet,
					    """scheduler""": scheduler,
					}

					return components







		def 					__A   ( self		,      A		,      A=0	)  -> List[str]:







					'''simple docstring'''







					if str(A	).startswith("""mps"""	):
								lowerCamelCase    		=  torch.manual_seed(A	)
					else:
								lowerCamelCase    		=  torch.Generator(device=A	).manual_seed(A	)

					lowerCamelCase    		=  {
					    """batch_size""": 1,
					    """num_inference_steps""": None,
					    """timesteps""": [22, 0],
					    """generator""": generator,
					    """output_type""": """np""",
					}

					return inputs







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







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components()
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3







		def 					__A   ( self	)  -> Dict:







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components(class_cond=A	)
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  0
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3







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







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components()
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  1
					lowerCamelCase    		=  None
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3







		def 					__A   ( self	)  -> List[str]:







					'''simple docstring'''







					lowerCamelCase    		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase    		=  self.get_dummy_components(class_cond=A	)
					lowerCamelCase    		=  ConsistencyModelPipeline(**A	)
					lowerCamelCase    		=  pipe.to(A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_dummy_inputs(A	)
					lowerCamelCase    		=  1
					lowerCamelCase    		=  None
					lowerCamelCase    		=  0
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 32, 32, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]
					lowerCamelCase    		=  np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3







@slow
@require_torch_gpu
class        __lowercase       (							unittest.TestCase  ):



		"""simple docstring"""

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







					'''simple docstring'''







					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()







		def 					__A   ( self		,      A=0		,      A=False		,      A="cpu"		,      A=torch.floataa		,      A=(1, 3, 64, 64)	)  -> Optional[int]:







					'''simple docstring'''







					lowerCamelCase    		=  torch.manual_seed(A	)

					lowerCamelCase    		=  {
					    """num_inference_steps""": None,
					    """timesteps""": [22, 0],
					    """class_labels""": 0,
					    """generator""": generator,
					    """output_type""": """np""",
					}

					if get_fixed_latents:
								lowerCamelCase    		=  self.get_fixed_latents(seed=A		,      device=A		,      dtype=A		,      shape=A	)
								lowerCamelCase    		=  latents

					return inputs







		def 					__A   ( self		,      A=0		,      A="cpu"		,      A=torch.floataa		,      A=(1, 3, 64, 64)	)  -> Union[str, Any]:







					'''simple docstring'''







					if type(A	) == str:
								lowerCamelCase    		=  torch.device(A	)
					lowerCamelCase    		=  torch.Generator(device=A	).manual_seed(A	)
					lowerCamelCase    		=  randn_tensor(A		,      generator=A		,      device=A		,      dtype=A	)
					return latents







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_inputs()
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 64, 64, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]

					lowerCamelCase    		=  np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 2e-2







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







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_inputs()
					lowerCamelCase    		=  1
					lowerCamelCase    		=  None
					lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 64, 64, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]

					lowerCamelCase    		=  np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 2e-2







		@require_torch_a
		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A		,      torch_dtype=torch.floataa	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_inputs(get_fixed_latents=A		,      device=A	)
					# Ensure usage of flash attention in torch 2.0
					with sdp_kernel(enable_flash=A		,      enable_math=A		,      enable_mem_efficient=A	):
								lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 64, 64, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]

					lowerCamelCase    		=  np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3







		@require_torch_a
		def 					__A   ( self	)  -> List[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  UNetaDModel.from_pretrained("""diffusers/consistency_models"""		,      subfolder="""diffusers_cd_imagenet64_l2"""	)
					lowerCamelCase    		=  CMStochasticIterativeScheduler(
					    num_train_timesteps=40		,      sigma_min=0.002		,      sigma_max=80.0		,      )
					lowerCamelCase    		=  ConsistencyModelPipeline(unet=A		,      scheduler=A	)
					pipe.to(torch_device=A		,      torch_dtype=torch.floataa	)
					pipe.set_progress_bar_config(disable=A	)

					lowerCamelCase    		=  self.get_inputs(get_fixed_latents=A		,      device=A	)
					lowerCamelCase    		=  1
					lowerCamelCase    		=  None
					# Ensure usage of flash attention in torch 2.0
					with sdp_kernel(enable_flash=A		,      enable_math=A		,      enable_mem_efficient=A	):
								lowerCamelCase    		=  pipe(**A	).images
					assert image.shape == (1, 64, 64, 3)

					lowerCamelCase    		=  image[0, -3:, -3:, -1]

					lowerCamelCase    		=  np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]	)

					assert np.abs(image_slice.flatten() - expected_slice	).max() < 1e-3

