Delete modeling_nsfw_image_detection.py

modeling_nsfw_image_detection.py

@@ -1,38 +0,0 @@
-import torch
-from PIL import Image
-import timm
-from timm import create_model
-from transformers import PreTrainedModel
-from typing import List
-from .configuration_nsfw_image_detection import NSFWImageDetectionConfig
-
-
-class NSFWImageDetector(PreTrainedModel):
-    config_class = NSFWImageDetectionConfig
-
-    def __init__(self, config: NSFWImageDetectionConfig):
-        super().__init__(config)
-        self.model = create_model(
-            'hf-hub:Marqo/nsfw-image-detection-384',
-            pretrained=True,
-        ).eval()
-        self.data_config = timm.data.resolve_model_data_config(self.model)
-        self.transforms = timm.data.create_transform(**self.data_config, is_training=False)
-        self.model = self.model.to(self.device)
-
-    def forward(
-        self,
-        images: List[Image.Image]
-    ) -> torch.Tensor:
-        image_inputs = [self.transforms(image) for image in images]
-        image_tensor = torch.stack(image_inputs, dim=0).to(self.device)
-        outputs = self.model(image_tensor).softmax(dim=-1).cpu()
-        return outputs
-            
-    def __call__(self, images: List[Image.Image]):
-        return self.forward(images)
-
-    @classmethod
-    def from_pretrained(self, *args, **kwargs):
-        config = NSFWImageDetectionConfig()
-        return NSFWImageDetector(config)