'''simple docstring'''






import json
import os
from collections import Counter

import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch import nn
from torch.utils.data import Dataset


_UpperCAmelCase    :     Tuple         =       {1: (1, 1), 2: (2, 1), 3: (3, 1), 4: (2, 2), 5: (5, 1), 6: (3, 2), 7: (7, 1), 8: (4, 2), 9: (3, 3)}







class lowercase_       (		nn.Module			):







   """simple docstring"""

   def __init__(      self      :			List[str], UpperCamelCase__      :			Optional[int]     )  ->			Dict:
        super().__init__()
        _A    		=  torchvision.models.resnetaaa(pretrained=UpperCamelCase__     )
        _A    		=  list(model.children()     )[:-2]
        _A    		=  nn.Sequential(*UpperCamelCase__     )
        _A    		=  nn.AdaptiveAvgPoolad(POOLING_BREAKDOWN[args.num_image_embeds]     )

   def 	__UpperCAmelCase   (      self      :			List[Any], UpperCamelCase__      :			Any     )  ->			List[Any]:
        # Bx3x224x224 -> Bx2048x7x7 -> Bx2048xN -> BxNx2048
        _A    		=  self.pool(self.model(UpperCamelCase__     )     )
        _A    		=  torch.flatten(UpperCamelCase__, start_dim=2     )
        _A    		=  out.transpose(1, 2     ).contiguous()
        return out  # BxNx2048







class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   def __init__(      self      :			List[Any], UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			List[Any], UpperCamelCase__      :			Any, UpperCamelCase__      :			Dict, UpperCamelCase__      :			Any     )  ->			Any:
        _A    		=  [json.loads(UpperCamelCase__     ) for l in open(UpperCamelCase__     )]
        _A    		=  os.path.dirname(UpperCamelCase__     )
        _A    		=  tokenizer
        _A    		=  labels
        _A    		=  len(UpperCamelCase__     )
        _A    		=  max_seq_length

        _A    		=  transforms
   def __len__(      self      :			Optional[int]     )  ->			int:
        return len(self.data     )
   def __getitem__(      self      :			Optional[int], UpperCamelCase__      :			Dict     )  ->			Union[str, Any]:
        _A    		=  torch.LongTensor(self.tokenizer.encode(self.data[index]['text'], add_special_tokens=UpperCamelCase__     )     )
        _A		,       _A		,       _A    		=  sentence[0], sentence[1:-1], sentence[-1]
        _A    		=  sentence[: self.max_seq_length]

        _A    		=  torch.zeros(self.n_classes     )
        _A    		=  1

        _A    		=  Image.open(os.path.join(self.data_dir, self.data[index]['img']     )     ).convert('RGB'     )
        _A    		=  self.transforms(UpperCamelCase__     )

        return {
            "image_start_token": start_token,
            "image_end_token": end_token,
            "sentence": sentence,
            "image": image,
            "label": label,
        }

   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			Tuple:
        _A    		=  Counter()
        for row in self.data:
             label_freqs.update(row['label']     )
        return label_freqs






def  _SCREAMING_SNAKE_CASE		(      __snake_case							:			Union[str, Any]				):
     _A    		=  [len(row['sentence']				) for row in batch]
     _A		,       _A    		=  len(__snake_case				), max(__snake_case				)

     _A    		=  torch.zeros(__snake_case   ,		__snake_case   ,		dtype=torch.long				)
     _A    		=  torch.zeros(__snake_case   ,		__snake_case   ,		dtype=torch.long				)

     for i_batch, (input_row, length) in enumerate(zip(__snake_case   ,		__snake_case				)				):
          _A    		=  input_row['sentence']
          _A    		=  1

     _A    		=  torch.stack([row['image'] for row in batch]				)
     _A    		=  torch.stack([row['label'] for row in batch]				)
     _A    		=  torch.stack([row['image_start_token'] for row in batch]				)
     _A    		=  torch.stack([row['image_end_token'] for row in batch]				)

     return text_tensor, mask_tensor, img_tensor, img_start_token, img_end_token, tgt_tensor





def  _SCREAMING_SNAKE_CASE		(      ):
     return [
         "Crime",
         "Drama",
         "Thriller",
         "Action",
         "Comedy",
         "Romance",
         "Documentary",
         "Short",
         "Mystery",
         "History",
         "Family",
         "Adventure",
         "Fantasy",
         "Sci-Fi",
         "Western",
         "Horror",
         "Sport",
         "War",
         "Music",
         "Musical",
         "Animation",
         "Biography",
         "Film-Noir",
     ]





def  _SCREAMING_SNAKE_CASE		(      ):
     return transforms.Compose(
         [
             transforms.Resize(2_5_6				),
             transforms.CenterCrop(2_2_4				),
             transforms.ToTensor(),
             transforms.Normalize(
                 mean=[0.46_77_70_44, 0.44_53_14_29, 0.40_66_10_17]   ,		std=[0.12_22_19_94, 0.12_14_58_35, 0.14_38_04_69]   ,		),
         ]				)




'''simple docstring'''






import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible







class lowercase_       (		unittest.TestCase			):







   """simple docstring"""

   def 	__UpperCAmelCase   (      self      :			Optional[Any]     )  ->			Dict:
        _A    		=  [
            'safety_checker/pytorch_model.bin',
            'safety_checker/model.safetensors',
            'vae/diffusion_pytorch_model.bin',
            'vae/diffusion_pytorch_model.safetensors',
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			Optional[Any]:
        _A    		=  [
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Union[str, Any]:
        _A    		=  [
            'safety_checker/pytorch_model.bin',
            'safety_checker/model.safetensors',
            'vae/diffusion_pytorch_model.bin',
            'vae/diffusion_pytorch_model.safetensors',
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
            'unet/diffusion_pytorch_model.bin',
            # Removed: 'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Dict:
        _A    		=  [
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
        ]
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			str:
        _A    		=  [
            'safety_checker/pytorch_model.bin',
            'safety_checker/model.safetensors',
            'vae/diffusion_pytorch_model.bin',
            'vae/diffusion_pytorch_model.safetensors',
            'text_encoder/pytorch_model.bin',
            # Removed: 'text_encoder/model.safetensors',
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			str     )  ->			Dict:
        _A    		=  [
            'safety_checker/pytorch_model.fp16.bin',
            'safety_checker/model.fp16.safetensors',
            'vae/diffusion_pytorch_model.fp16.bin',
            'vae/diffusion_pytorch_model.fp16.safetensors',
            'text_encoder/pytorch_model.fp16.bin',
            'text_encoder/model.fp16.safetensors',
            'unet/diffusion_pytorch_model.fp16.bin',
            'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			List[str]:
        _A    		=  [
            'unet/diffusion_pytorch_model.fp16.bin',
            'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			List[str]     )  ->			List[str]:
        # pass variant but use the non-variant filenames
        _A    		=  [
            'unet/diffusion_pytorch_model.bin',
            'unet/diffusion_pytorch_model.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			Optional[Any]:
        _A    		=  [
            'safety_checker/pytorch_model.fp16.bin',
            'safety_checker/model.fp16.safetensors',
            'vae/diffusion_pytorch_model.fp16.bin',
            'vae/diffusion_pytorch_model.fp16.safetensors',
            'text_encoder/pytorch_model.fp16.bin',
            'text_encoder/model.fp16.safetensors',
            'unet/diffusion_pytorch_model.fp16.bin',
            # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			Union[str, Any]:
        _A    		=  [
            'text_encoder/pytorch_model.fp16.bin',
            'text_encoder/model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			str:
        # pass variant but use the non-variant filenames
        _A    		=  [
            'text_encoder/pytorch_model.bin',
            'text_encoder/model.safetensors',
        ]
        _A    		=  'fp16'
        self.assertTrue(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )

   def 	__UpperCAmelCase   (      self      :			List[Any]     )  ->			int:
        _A    		=  [
            'safety_checker/pytorch_model.fp16.bin',
            'safety_checker/model.fp16.safetensors',
            'vae/diffusion_pytorch_model.fp16.bin',
            'vae/diffusion_pytorch_model.fp16.safetensors',
            'text_encoder/pytorch_model.fp16.bin',
            # 'text_encoder/model.fp16.safetensors',
            'unet/diffusion_pytorch_model.fp16.bin',
            'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        _A    		=  'fp16'
        self.assertFalse(is_safetensors_compatible(UpperCamelCase__, variant=UpperCamelCase__     )     )



