





'''simple docstring'''

import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


lowercase__		:    List[str]  =       logging.get_logger(__name__)







def       a__       (       lowercase      :  Optional[int],      lowercase      :  Optional[int]   )	->	Optional[int]:





     """simple docstring"""

     _UpperCamelCase       					= nn.functional.normalize(lowercase   )
     _UpperCamelCase       					= nn.functional.normalize(lowercase   )
     return torch.mm(lowercase,      normalized_text_embeds.t()   )




class        __lowerCAmelCase  (					A__				):






       """simple docstring"""
       _snake_case			:  Dict		       =					CLIPConfig

       _snake_case			:  Tuple		       =					['CLIPEncoderLayer']





       def __init__(							self     :	List[str] ,				lowerCAmelCase__     :	CLIPConfig  )					->  int:

            '''simple docstring'''




            super().__init__(_lowerCAmelCase  )

            _UpperCamelCase       					= CLIPVisionModel(config.vision_config  )
            _UpperCamelCase       					= nn.Linear(config.vision_config.hidden_size ,				config.projection_dim ,				bias=_lowerCAmelCase  )

            _UpperCamelCase       					= nn.Parameter(torch.ones(17 ,				config.projection_dim  ) ,				requires_grad=_lowerCAmelCase  )
            _UpperCamelCase       					= nn.Parameter(torch.ones(3 ,				config.projection_dim  ) ,				requires_grad=_lowerCAmelCase  )

            _UpperCamelCase       					= nn.Parameter(torch.ones(17  ) ,				requires_grad=_lowerCAmelCase  )
            _UpperCamelCase       					= nn.Parameter(torch.ones(3  ) ,				requires_grad=_lowerCAmelCase  )





       @torch.no_grad()
       def snake_case__		(							self     :	Optional[Any] ,				lowerCAmelCase__     :	int ,				lowerCAmelCase__     :	Dict  )					->  Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase       					= self.vision_model(_lowerCAmelCase  )[1]  # pooled_output
            _UpperCamelCase       					= self.visual_projection(_lowerCAmelCase  )

            # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
            _UpperCamelCase       					= cosine_distance(_lowerCAmelCase ,				self.special_care_embeds  ).cpu().float().numpy()
            _UpperCamelCase       					= cosine_distance(_lowerCAmelCase ,				self.concept_embeds  ).cpu().float().numpy()

            _UpperCamelCase       					= []
            _UpperCamelCase       					= image_embeds.shape[0]
            for i in range(_lowerCAmelCase  ):
                 _UpperCamelCase       					= {'''special_scores''': {}, '''special_care''': [], '''concept_scores''': {}, '''bad_concepts''': []}

                 # increase this value to create a stronger `nfsw` filter
                 # at the cost of increasing the possibility of filtering benign images
                 _UpperCamelCase       					= 0.0

                 for concept_idx in range(len(special_cos_dist[0]  )  ):
                      _UpperCamelCase       					= special_cos_dist[i][concept_idx]
                      _UpperCamelCase       					= self.special_care_embeds_weights[concept_idx].item()
                      _UpperCamelCase       					= round(concept_cos - concept_threshold + adjustment ,				3  )
                      if result_img["special_scores"][concept_idx] > 0:
                           result_img["special_care"].append({concept_idx, result_img['''special_scores'''][concept_idx]}  )
                           _UpperCamelCase       					= 0.01

                 for concept_idx in range(len(cos_dist[0]  )  ):
                      _UpperCamelCase       					= cos_dist[i][concept_idx]
                      _UpperCamelCase       					= self.concept_embeds_weights[concept_idx].item()
                      _UpperCamelCase       					= round(concept_cos - concept_threshold + adjustment ,				3  )
                      if result_img["concept_scores"][concept_idx] > 0:
                           result_img["bad_concepts"].append(_lowerCAmelCase  )

                 result.append(_lowerCAmelCase  )

            _UpperCamelCase       					= [len(res['''bad_concepts''']  ) > 0 for res in result]

            for idx, has_nsfw_concept in enumerate(_lowerCAmelCase  ):
                 if has_nsfw_concept:
                      if torch.is_tensor(_lowerCAmelCase  ) or torch.is_tensor(images[0]  ):
                           _UpperCamelCase       					= torch.zeros_like(images[idx]  )  # black image
                      else:
                           _UpperCamelCase       					= np.zeros(images[idx].shape  )  # black image

            if any(_lowerCAmelCase  ):
                 logger.warning(
                     '''Potential NSFW content was detected in one or more images. A black image will be returned instead.'''
                     ''' Try again with a different prompt and/or seed.'''  )

            return images, has_nsfw_concepts







       @torch.no_grad()
       def snake_case__		(							self     :	Tuple ,				lowerCAmelCase__     :	torch.FloatTensor ,				lowerCAmelCase__     :	torch.FloatTensor  )					->  Union[str, Any]:

            '''simple docstring'''




            _UpperCamelCase       					= self.vision_model(_lowerCAmelCase  )[1]  # pooled_output
            _UpperCamelCase       					= self.visual_projection(_lowerCAmelCase  )

            _UpperCamelCase       					= cosine_distance(_lowerCAmelCase ,				self.special_care_embeds  )
            _UpperCamelCase       					= cosine_distance(_lowerCAmelCase ,				self.concept_embeds  )

            # increase this value to create a stronger `nsfw` filter
            # at the cost of increasing the possibility of filtering benign images
            _UpperCamelCase       					= 0.0

            _UpperCamelCase       					= special_cos_dist - self.special_care_embeds_weights + adjustment
            # special_scores = special_scores.round(decimals=3)
            _UpperCamelCase       					= torch.any(special_scores > 0 ,				dim=1  )
            _UpperCamelCase       					= special_care * 0.01
            _UpperCamelCase       					= special_adjustment.unsqueeze(1  ).expand(-1 ,				cos_dist.shape[1]  )

            _UpperCamelCase       					= (cos_dist - self.concept_embeds_weights) + special_adjustment
            # concept_scores = concept_scores.round(decimals=3)
            _UpperCamelCase       					= torch.any(concept_scores > 0 ,				dim=1  )

            _UpperCamelCase       					= 0.0  # black image

            return images, has_nsfw_concepts










'''simple docstring'''

import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

lowercase__		:    Union[str, Any] =       logging.get_logger(__name__)







def       a__       (       lowercase      :  Dict,      lowercase      :  bool = True   )	->	int:





     """simple docstring"""

     _UpperCamelCase       					= {}
     for name in model.attn_processors.keys():
          _UpperCamelCase       					= name.endswith('''attn2.processor'''   ) and not (
              name.startswith('''transformer_in'''   ) or '''temp_attentions''' in name.split('''.'''   )
          )
          _UpperCamelCase       					= model.config.cross_attention_dim if has_cross_attention else None
          if name.startswith('''mid_block'''   ):
               _UpperCamelCase       					= model.config.block_out_channels[-1]
          elif name.startswith('''up_blocks'''   ):
               _UpperCamelCase       					= int(name[len('''up_blocks.'''   )]   )
               _UpperCamelCase       					= list(reversed(model.config.block_out_channels   )   )[block_id]
          elif name.startswith('''down_blocks'''   ):
               _UpperCamelCase       					= int(name[len('''down_blocks.'''   )]   )
               _UpperCamelCase       					= model.config.block_out_channels[block_id]
          elif name.startswith('''transformer_in'''   ):
               # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
               _UpperCamelCase       					= 8 * model.config.attention_head_dim

          _UpperCamelCase       					= LoRAAttnProcessor(hidden_size=lowercase,      cross_attention_dim=lowercase   )
          _UpperCamelCase       					= lora_attn_procs[name].to(model.device   )

          if mock_weights:
               # add 1 to weights to mock trained weights
               with torch.no_grad():
                    lora_attn_procs[name].to_q_lora.up.weight += 1
                    lora_attn_procs[name].to_k_lora.up.weight += 1
                    lora_attn_procs[name].to_v_lora.up.weight += 1
                    lora_attn_procs[name].to_out_lora.up.weight += 1

     return lora_attn_procs




@skip_mps
class        __lowerCAmelCase  (					__magic_name__		,		__magic_name__		,		unittest.TestCase				):






       """simple docstring"""
       _snake_case			:  int		       =					UNetaDConditionModel
       _snake_case			:  List[str]		       =					'sample'





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

            '''simple docstring'''




            _UpperCamelCase       					= 4
            _UpperCamelCase       					= 4
            _UpperCamelCase       					= 4
            _UpperCamelCase       					= (32, 32)

            _UpperCamelCase       					= floats_tensor((batch_size, num_channels, num_frames) + sizes  ).to(lowerCAmelCase__  )
            _UpperCamelCase       					= torch.tensor([10]  ).to(lowerCAmelCase__  )
            _UpperCamelCase       					= floats_tensor((batch_size, 4, 32)  ).to(lowerCAmelCase__  )

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





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

            '''simple docstring'''




            return (4, 4, 32, 32)





       @property
       def snake_case__		(							self     :	Optional[Any]  )					->  Union[str, Any]:

            '''simple docstring'''




            return (4, 4, 32, 32)





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

            '''simple docstring'''




            _UpperCamelCase       					= {
                '''block_out_channels''': (32, 64),
                '''down_block_types''': (
                    '''CrossAttnDownBlock3D''',
                    '''DownBlock3D''',
                ),
                '''up_block_types''': ('''UpBlock3D''', '''CrossAttnUpBlock3D'''),
                '''cross_attention_dim''': 32,
                '''attention_head_dim''': 8,
                '''out_channels''': 4,
                '''in_channels''': 4,
                '''layers_per_block''': 1,
                '''sample_size''': 32,
            }
            _UpperCamelCase       					= self.dummy_input
            return init_dict, inputs_dict





       @unittest.skipIf(
           torch_device != '''cuda''' or not is_xformers_available() ,				reason='''XFormers attention is only available with CUDA and `xformers` installed''' ,				)
       def snake_case__		(							self     :	Union[str, Any]  )					->  int:

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )

            model.enable_xformers_memory_efficient_attention()

            assert (
                model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
                == "XFormersAttnProcessor"
            ), "xformers is not enabled"





       def snake_case__		(							self     :	Tuple  )					->  Tuple:

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 32

            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )
            model.eval()

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  )

                 if isinstance(lowerCAmelCase__ ,				lowerCAmelCase__  ):
                      _UpperCamelCase       					= output.sample

            self.assertIsNotNone(lowerCAmelCase__  )
            _UpperCamelCase       					= inputs_dict['''sample'''].shape
            self.assertEqual(output.shape ,				lowerCAmelCase__ ,				'''Input and output shapes do not match'''  )





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

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )
            model.eval()

            with torch.no_grad():
                 # Warmup pass when using mps (see #372)
                 if torch_device == "mps" and isinstance(lowerCAmelCase__ ,				lowerCAmelCase__  ):
                      model(**self.dummy_input  )

                 _UpperCamelCase       					= model(**lowerCAmelCase__  )
                 if isinstance(lowerCAmelCase__ ,				lowerCAmelCase__  ):
                      _UpperCamelCase       					= first.sample

                 _UpperCamelCase       					= model(**lowerCAmelCase__  )
                 if isinstance(lowerCAmelCase__ ,				lowerCAmelCase__  ):
                      _UpperCamelCase       					= second.sample

            _UpperCamelCase       					= first.cpu().numpy()
            _UpperCamelCase       					= second.cpu().numpy()
            _UpperCamelCase       					= out_a[~np.isnan(lowerCAmelCase__  )]
            _UpperCamelCase       					= out_a[~np.isnan(lowerCAmelCase__  )]
            _UpperCamelCase       					= np.amax(np.abs(out_a - out_a  )  )
            self.assertLessEqual(lowerCAmelCase__ ,				1e-5  )





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

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 8

            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )
            model.eval()

            model.set_attention_slice('''auto'''  )
            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  )
            assert output is not None

            model.set_attention_slice('''max'''  )
            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  )
            assert output is not None

            model.set_attention_slice(2  )
            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  )
            assert output is not None





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

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 8

            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

            _UpperCamelCase       					= create_lora_layers(lowerCAmelCase__  )

            # make sure we can set a list of attention processors
            model.set_attn_processor(lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )

            # test that attn processors can be set to itself
            model.set_attn_processor(model.attn_processors  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.0}  ).sample
                 _UpperCamelCase       					= model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.5}  ).sample
                 _UpperCamelCase       					= model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.5}  ).sample

            assert (samplea - samplea).abs().max() < 3e-3
            assert (samplea - samplea).abs().max() < 3e-3

            # sample 2 and sample 3 should be different
            assert (samplea - samplea).abs().max() > 3e-3





       def snake_case__		(							self     :	str  )					->  Tuple:

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 8

            torch.manual_seed(0  )
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

            _UpperCamelCase       					= create_lora_layers(lowerCAmelCase__  )
            model.set_attn_processor(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.5}  ).sample

            with tempfile.TemporaryDirectory() as tmpdirname:
                 model.save_attn_procs(lowerCAmelCase__  )
                 self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__ ,				'''pytorch_lora_weights.bin'''  )  )  )
                 torch.manual_seed(0  )
                 _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
                 new_model.to(lowerCAmelCase__  )
                 new_model.load_attn_procs(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= new_model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.5}  ).sample

            assert (sample - new_sample).abs().max() < 1e-3

            # LoRA and no LoRA should NOT be the same
            assert (sample - old_sample).abs().max() > 1e-4





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

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 8

            torch.manual_seed(0  )
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

            _UpperCamelCase       					= create_lora_layers(lowerCAmelCase__  )
            model.set_attn_processor(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.5}  ).sample

            with tempfile.TemporaryDirectory() as tmpdirname:
                 model.save_attn_procs(lowerCAmelCase__ ,				safe_serialization=lowerCAmelCase__  )
                 self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__ ,				'''pytorch_lora_weights.safetensors'''  )  )  )
                 torch.manual_seed(0  )
                 _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
                 new_model.to(lowerCAmelCase__  )
                 new_model.load_attn_procs(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= new_model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.5}  ).sample

            assert (sample - new_sample).abs().max() < 3e-3

            # LoRA and no LoRA should NOT be the same
            assert (sample - old_sample).abs().max() > 1e-4





       def snake_case__		(							self     :	Tuple  )					->  List[Any]:

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 8

            torch.manual_seed(0  )
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )

            _UpperCamelCase       					= create_lora_layers(lowerCAmelCase__ ,				mock_weights=lowerCAmelCase__  )
            model.set_attn_processor(lowerCAmelCase__  )
            # Saving as torch, properly reloads with directly filename
            with tempfile.TemporaryDirectory() as tmpdirname:
                 model.save_attn_procs(lowerCAmelCase__  )
                 self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__ ,				'''pytorch_lora_weights.bin'''  )  )  )
                 torch.manual_seed(0  )
                 _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
                 new_model.to(lowerCAmelCase__  )
                 new_model.load_attn_procs(lowerCAmelCase__ ,				weight_name='''pytorch_lora_weights.bin'''  )





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

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 8

            torch.manual_seed(0  )
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )

            _UpperCamelCase       					= create_lora_layers(lowerCAmelCase__ ,				mock_weights=lowerCAmelCase__  )
            model.set_attn_processor(lowerCAmelCase__  )
            # Saving as torch, properly reloads with directly filename
            with tempfile.TemporaryDirectory() as tmpdirname:
                 model.save_attn_procs(lowerCAmelCase__  )
                 self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase__ ,				'''pytorch_lora_weights.bin'''  )  )  )
                 torch.manual_seed(0  )
                 _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
                 new_model.to(lowerCAmelCase__  )
                 with self.assertRaises(lowerCAmelCase__  ) as e:
                      new_model.load_attn_procs(lowerCAmelCase__ ,				use_safetensors=lowerCAmelCase__  )
                 self.assertIn('''Error no file named pytorch_lora_weights.safetensors''' ,				str(e.exception  )  )





       def snake_case__		(							self     :	Tuple  )					->  List[Any]:

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 8

            torch.manual_seed(0  )
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

            _UpperCamelCase       					= create_lora_layers(lowerCAmelCase__  )
            model.set_attn_processor(lowerCAmelCase__  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__ ,				cross_attention_kwargs={'''scale''': 0.0}  ).sample

            model.set_attn_processor(AttnProcessor()  )

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

            assert (sample - new_sample).abs().max() < 1e-4
            assert (sample - old_sample).abs().max() < 3e-3





       @unittest.skipIf(
           torch_device != '''cuda''' or not is_xformers_available() ,				reason='''XFormers attention is only available with CUDA and `xformers` installed''' ,				)
       def snake_case__		(							self     :	List[str]  )					->  Tuple:

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()

            _UpperCamelCase       					= 4

            torch.manual_seed(0  )
            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )
            _UpperCamelCase       					= create_lora_layers(lowerCAmelCase__  )
            model.set_attn_processor(lowerCAmelCase__  )

            # default
            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

                 model.enable_xformers_memory_efficient_attention()
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

                 model.disable_xformers_memory_efficient_attention()
                 _UpperCamelCase       					= model(**lowerCAmelCase__  ).sample

            assert (sample - on_sample).abs().max() < 1e-4
            assert (sample - off_sample).abs().max() < 1e-4







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

            '''simple docstring'''




            _UpperCamelCase					,       _UpperCamelCase       					= self.prepare_init_args_and_inputs_for_common()
            _UpperCamelCase       					= 32

            _UpperCamelCase       					= self.model_class(**lowerCAmelCase__  )
            model.to(lowerCAmelCase__  )
            model.eval()

            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  )[0]

            model.enable_forward_chunking()
            with torch.no_grad():
                 _UpperCamelCase       					= model(**lowerCAmelCase__  )[0]

            self.assertEqual(output.shape ,				output_a.shape ,				'''Shape doesn\'t match'''  )
            assert np.abs(output.cpu() - output_a.cpu()  ).max() < 1e-2


# (todo: sayakpaul) implement SLOW tests.



