






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 							snake_case_       :

						"""simple docstring"""



						def __init__(self:					List[str] ,					__UpperCAmelCase:					Any ,					__UpperCAmelCase:					int=13 ,					__UpperCAmelCase:					int=7 ,					__UpperCAmelCase:					List[str]=6 ,					__UpperCAmelCase:					Union[str, Any]=17 ,					__UpperCAmelCase:					str=23 ,					__UpperCAmelCase:					Optional[Any]=11 ,					__UpperCAmelCase:					Union[str, Any]=True ,					)    ->     Union[str, Any]:




											'''simple docstring'''



											__a					:						List[str]      =					parent
											__a					:						List[Any]      =					batch_size
											__a					:						Tuple      =					seq_length
											__a					:						Dict      =					act_dim
											__a					:						Optional[int]      =					state_dim
											__a					:						Tuple      =					hidden_size
											__a					:						int      =					max_length
											__a					:						Optional[Any]      =					is_training



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




											'''simple docstring'''



											__a					:						Tuple      =					floats_tensor((self.batch_size, self.seq_length, self.state_dim)   )
											__a					:						Union[str, Any]      =					floats_tensor((self.batch_size, self.seq_length, self.act_dim)   )
											__a					:						Union[str, Any]      =					floats_tensor((self.batch_size, self.seq_length, 1)   )
											__a					:						Any      =					floats_tensor((self.batch_size, self.seq_length, 1)   )
											__a					:						Tuple      =					ids_tensor((self.batch_size, self.seq_length) ,					vocab_size=1000   )
											__a					:						Union[str, Any]      =					random_attention_mask((self.batch_size, self.seq_length)   )

											__a					:						int      =					self.get_config()

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



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




											'''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        UpperCAmelCase__					(self:					Optional[Any] ,					__UpperCAmelCase:					Optional[int] ,					__UpperCAmelCase:					Optional[int] ,					__UpperCAmelCase:					Tuple ,					__UpperCAmelCase:					Dict ,					__UpperCAmelCase:					Dict ,					__UpperCAmelCase:					Any ,					__UpperCAmelCase:					Union[str, Any] ,					)    ->     Any:




											'''simple docstring'''



											__a					:						Dict      =					DecisionTransformerModel(config=__UpperCAmelCase   )
											model.to(__UpperCAmelCase   )
											model.eval()
											__a					:						Optional[Any]      =					model(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase   )

											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        UpperCAmelCase__					(self:					int   )    ->     Any:




											'''simple docstring'''



											__a					:						str      =					self.prepare_config_and_inputs()
											(
											    (
											    __a
											)      ,			(
											    __a
											)      ,			(
											    __a
											)      ,			(
											    __a
											)      ,			(
											    __a
											)      ,			(
											    __a
											)      ,			(
											    __a
											)      ,			
											)					:						Tuple      =					config_and_inputs
											__a					:						List[str]      =					{
											    "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 							snake_case_       (    __UpperCamelCase		,							__UpperCamelCase		,							__UpperCamelCase		,							unittest.TestCase      ):

						"""simple docstring"""



						snake_case__        =      (DecisionTransformerModel,) if is_torch_available() else ()
						snake_case__        =      ()
						snake_case__        =      {"""feature-extraction""": DecisionTransformerModel} if is_torch_available() else {}

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

						# Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False



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




											'''simple docstring'''



											__a					:						List[Any]      =					DecisionTransformerModelTester(self   )
											__a					:						Any      =					ConfigTester(self ,					config_class=__UpperCAmelCase ,					hidden_size=37   )



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




											'''simple docstring'''



											self.config_tester.run_common_tests()



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




											'''simple docstring'''



											__a					:						Any      =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__UpperCAmelCase   )



						@slow
						def        UpperCAmelCase__					(self:					Tuple   )    ->     List[str]:




											'''simple docstring'''



											for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__a					:						List[str]      =					DecisionTransformerModel.from_pretrained(__UpperCAmelCase   )
																self.assertIsNotNone(__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					int   )    ->     int:




											'''simple docstring'''



											__a      ,			__a					:						List[str]      =					self.model_tester.prepare_config_and_inputs_for_common()

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

																__a					:						Any      =					[
																    "states",
																    "actions",
																    "rewards",
																    "returns_to_go",
																    "timesteps",
																    "attention_mask",
																]

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




@require_torch
class 							snake_case_       (    unittest.TestCase      ):

						"""simple docstring"""



						@slow
						def        UpperCAmelCase__					(self:					str   )    ->     List[str]:




											'''simple docstring'''



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

											__a					:						Tuple      =					torch.tensor(
											    [[0.24_27_93, -0.28_69_30_74, 0.8_74_26_13], [0.67_81_52_74, -0.08_10_10_85, -0.12_95_21_47]] ,					device=__UpperCAmelCase   )

											__a					:						Any      =					torch.tensor(__UpperCAmelCase ,					device=__UpperCAmelCase ,					dtype=torch.floataa   ).reshape(1 ,					1 ,					1   )
											__a					:						Tuple      =					state
											__a					:						Union[str, Any]      =					torch.zeros(1 ,					0 ,					config.act_dim ,					device=__UpperCAmelCase ,					dtype=torch.floataa   )
											__a					:						Union[str, Any]      =					torch.zeros(1 ,					0 ,					device=__UpperCAmelCase ,					dtype=torch.floataa   )
											__a					:						Union[str, Any]      =					torch.tensor(0 ,					device=__UpperCAmelCase ,					dtype=torch.long   ).reshape(1 ,					1   )

											for step in range(__UpperCAmelCase   ):
																__a					:						Optional[Any]      =					torch.cat([actions, torch.zeros(1 ,					1 ,					config.act_dim ,					device=__UpperCAmelCase   )] ,					dim=1   )
																__a					:						Optional[int]      =					torch.cat([rewards, torch.zeros(1 ,					1 ,					device=__UpperCAmelCase   )] ,					dim=1   )

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

																with torch.no_grad():
																					__a      ,			__a      ,			__a					:						Tuple      =					model(
																					    states=__UpperCAmelCase ,					actions=__UpperCAmelCase ,					rewards=__UpperCAmelCase ,					returns_to_go=__UpperCAmelCase ,					timesteps=__UpperCAmelCase ,					attention_mask=__UpperCAmelCase ,					return_dict=__UpperCAmelCase ,					)

																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					:						Optional[int]      =					(  # env.step(action)
																    torch.randn(1 ,					1 ,					config.state_dim   ).to(device=__UpperCAmelCase ,					dtype=torch.floataa   ),
																    1.0,
																    False,
																    {},
																)

																__a					:						str      =					action_pred[0, -1]
																__a					:						Tuple      =					torch.cat([states, state] ,					dim=1   )
																__a					:						List[str]      =					returns_to_go[0, -1] - reward
																__a					:						str      =					torch.cat([returns_to_go, pred_return.reshape(1 ,					1 ,					1   )] ,					dim=1   )
																__a					:						Dict      =					torch.cat(
																    [timesteps, torch.ones((1, 1) ,					device=__UpperCAmelCase ,					dtype=torch.long   ) * (step + 1)] ,					dim=1   )













from collections import defaultdict
from typing import Optional

from ..image_utils import load_image
from ..utils import (
    add_end_docstrings,
    is_torch_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_torch_available():
			import torch

			from ..models.auto.modeling_auto import MODEL_FOR_MASK_GENERATION_MAPPING

UpperCAmelCase__      			=						logging.get_logger(__name__)




@add_end_docstrings(__UpperCamelCase      )
class 							snake_case_       (    __UpperCamelCase      ):

						"""simple docstring"""



						def __init__(self:					int ,					**__UpperCAmelCase:					Optional[Any]   )    ->     str:




											'''simple docstring'''



											super().__init__(**__UpperCAmelCase   )
											requires_backends(self ,					"vision"   )
											requires_backends(self ,					"torch"   )

											if self.framework != "pt":
																raise ValueError(f'The {self.__class__} is only available in PyTorch.'   )

											self.check_model_type(__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					Optional[int] ,					**__UpperCAmelCase:					Optional[int]   )    ->     List[Any]:




											'''simple docstring'''



											__a					:						Any      =					{}
											__a					:						Union[str, Any]      =					{}
											__a					:						List[str]      =					{}
											# preprocess args
											if "points_per_batch" in kwargs:
																__a					:						Dict      =					kwargs["points_per_batch"]
											if "points_per_crop" in kwargs:
																__a					:						Tuple      =					kwargs["points_per_crop"]
											if "crops_n_layers" in kwargs:
																__a					:						Dict      =					kwargs["crops_n_layers"]
											if "crop_overlap_ratio" in kwargs:
																__a					:						int      =					kwargs["crop_overlap_ratio"]
											if "crop_n_points_downscale_factor" in kwargs:
																__a					:						Optional[int]      =					kwargs["crop_n_points_downscale_factor"]
											# postprocess args
											if "pred_iou_thresh" in kwargs:
																__a					:						int      =					kwargs["pred_iou_thresh"]
											if "stability_score_offset" in kwargs:
																__a					:						Any      =					kwargs["stability_score_offset"]
											if "mask_threshold" in kwargs:
																__a					:						int      =					kwargs["mask_threshold"]
											if "stability_score_thresh" in kwargs:
																__a					:						Any      =					kwargs["stability_score_thresh"]
											if "crops_nms_thresh" in kwargs:
																__a					:						int      =					kwargs["crops_nms_thresh"]
											if "output_rle_mask" in kwargs:
																__a					:						Optional[Any]      =					kwargs["output_rle_mask"]
											if "output_bboxes_mask" in kwargs:
																__a					:						Optional[Any]      =					kwargs["output_bboxes_mask"]
											return preprocess_kwargs, forward_params, postprocess_kwargs



						def __call__(self:					Union[str, Any] ,					__UpperCAmelCase:					str ,					*__UpperCAmelCase:					Optional[Any] ,					__UpperCAmelCase:					str=None ,					__UpperCAmelCase:					Union[str, Any]=None ,					**__UpperCAmelCase:					Optional[int]   )    ->     Tuple:




											'''simple docstring'''



											return super().__call__(__UpperCAmelCase ,					*__UpperCAmelCase ,					num_workers=__UpperCAmelCase ,					batch_size=__UpperCAmelCase ,					**__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					Optional[Any] ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					int=64 ,					__UpperCAmelCase:					int = 0 ,					__UpperCAmelCase:					float = 512 / 1500 ,					__UpperCAmelCase:					Optional[int] = 32 ,					__UpperCAmelCase:					Optional[int] = 1 ,					)    ->     str:




											'''simple docstring'''



											__a					:						Union[str, Any]      =					load_image(__UpperCAmelCase   )
											__a					:						List[str]      =					self.image_processor.size["longest_edge"]
											__a      ,			__a      ,			__a      ,			__a					:						int      =					self.image_processor.generate_crop_boxes(
											    __UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase   )
											__a					:						Union[str, Any]      =					self.image_processor(images=__UpperCAmelCase ,					return_tensors="pt"   )

											with self.device_placement():
																if self.framework == "pt":
																					__a					:						str      =					self.get_inference_context()
																					with inference_context():
																										__a					:						int      =					self._ensure_tensor_on_device(__UpperCAmelCase ,					device=self.device   )
																										__a					:						Union[str, Any]      =					self.model.get_image_embeddings(model_inputs.pop("pixel_values"   )   )
																										__a					:						Optional[Any]      =					image_embeddings

											__a					:						Dict      =					grid_points.shape[1]
											__a					:						Any      =					points_per_batch if points_per_batch is not None else n_points

											if points_per_batch <= 0:
																raise ValueError(
																    "Cannot have points_per_batch<=0. Must be >=1 to returned batched outputs. "
																    "To return all points at once, set points_per_batch to None"   )

											for i in range(0 ,					__UpperCAmelCase ,					__UpperCAmelCase   ):
																__a					:						List[str]      =					grid_points[:, i : i + points_per_batch, :, :]
																__a					:						str      =					input_labels[:, i : i + points_per_batch]
																__a					:						str      =					i == n_points - points_per_batch
																yield {
																    "input_points": batched_points,
																    "input_labels": labels,
																    "input_boxes": crop_boxes,
																    "is_last": is_last,
																    **model_inputs,
																}



						def        UpperCAmelCase__					(self:					Any ,					__UpperCAmelCase:					Optional[Any] ,					__UpperCAmelCase:					Any=0.88 ,					__UpperCAmelCase:					Optional[int]=0.95 ,					__UpperCAmelCase:					Optional[int]=0 ,					__UpperCAmelCase:					int=1 ,					)    ->     List[str]:




											'''simple docstring'''



											__a					:						List[Any]      =					model_inputs.pop("input_boxes"   )
											__a					:						Optional[Any]      =					model_inputs.pop("is_last"   )
											__a					:						Dict      =					model_inputs.pop("original_sizes"   ).tolist()
											__a					:						int      =					model_inputs.pop("reshaped_input_sizes"   ).tolist()

											__a					:						Optional[Any]      =					self.model(**__UpperCAmelCase   )

											# post processing happens here in order to avoid CPU GPU copies of ALL the masks
											__a					:						Union[str, Any]      =					model_outputs["pred_masks"]
											__a					:						Optional[Any]      =					self.image_processor.post_process_masks(
											    __UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					binarize=__UpperCAmelCase   )
											__a					:						Optional[Any]      =					model_outputs["iou_scores"]
											__a      ,			__a      ,			__a					:						Dict      =					self.image_processor.filter_masks(
											    masks[0] ,					iou_scores[0] ,					original_sizes[0] ,					input_boxes[0] ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					)
											return {
											    "masks": masks,
											    "is_last": is_last,
											    "boxes": boxes,
											    "iou_scores": iou_scores,
											}



						def        UpperCAmelCase__					(self:					Any ,					__UpperCAmelCase:					str ,					__UpperCAmelCase:					List[str]=False ,					__UpperCAmelCase:					int=False ,					__UpperCAmelCase:					List[str]=0.7 ,					)    ->     Dict:




											'''simple docstring'''



											__a					:						Dict      =					[]
											__a					:						Optional[Any]      =					[]
											__a					:						Optional[Any]      =					[]
											for model_output in model_outputs:
																all_scores.append(model_output.pop("iou_scores"   )   )
																all_masks.extend(model_output.pop("masks"   )   )
																all_boxes.append(model_output.pop("boxes"   )   )

											__a					:						Any      =					torch.cat(__UpperCAmelCase   )
											__a					:						Union[str, Any]      =					torch.cat(__UpperCAmelCase   )
											__a      ,			__a      ,			__a      ,			__a					:						Tuple      =					self.image_processor.post_process_for_mask_generation(
											    __UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase   )

											__a					:						Optional[int]      =					defaultdict(__UpperCAmelCase   )
											for output in model_outputs:
																for k, v in output.items():
																					extra[k].append(__UpperCAmelCase   )

											__a					:						int      =					{}
											if output_rle_mask:
																__a					:						Optional[int]      =					rle_mask

											if output_bboxes_mask:
																__a					:						List[str]      =					bounding_boxes

											return {"masks": output_masks, "scores": iou_scores, **optional, **extra}





