Update handler.py

handler.py

@@ -1,19 +1,80 @@
-from typing import  Dict, List, Any
+
 import timm
 import torch
+from PIL import Image
 from timm.utils import ParseKwargs
 from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, DEFAULT_CROP_PCT
 
+###
+
+import os
+import time
+from contextlib import suppress
+from functools import partial
+
+import numpy as np
+import pandas as pd
+import torch
+
+from timm.data import create_dataset, create_loader, resolve_data_config, ImageNetInfo, infer_imagenet_subset
+from timm.layers import apply_test_time_pool
+from timm.models import create_model
+from timm.utils import AverageMeter, setup_default_logging, set_jit_fuser, ParseKwargs
+
+try:
+    from apex import amp
+    has_apex = True
+except ImportError:
+    has_apex = False
+
+has_native_amp = False
+try:
+    if getattr(torch.cuda.amp, 'autocast') is not None:
+        has_native_amp = True
+except AttributeError:
+    pass
+
+# try:
+#     from functorch.compile import memory_efficient_fusion
+#     has_functorch = True
+# except ImportError as e:
+#     has_functorch = False
+
+has_compile = hasattr(torch, 'compile')
+
+import PIL
+import requests
+import io
+import base64
+
+
+
+# ImageFile.LOAD_TRUNCATED_IMAGES = True
+###
+
 class EndpointHandler():
     def __init__(self, path=""):
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        # self.aiGeneratorModel = timm.create_model('eva02_base_patch14_448.mim_in22k_ft_in22k_in1k', num_classes=9, in_chans=3, checkpoint_path=path + '/AIModelDetector.pth-6ff3631e.pth')
-        aiArtModel = timm.create_model('eva02_base_patch14_448.mim_in22k_ft_in22k_in1k', num_classes=3, in_chans=3, checkpoint_path=path + '/AIArtDetector.pth-af59f7fa.pth')
-        # aiGeneratorModel = aiGeneratorModel.to(self.device)
-        aiArtModel = aiArtModel.to(self.device)
-        # aiGeneratorModel.eval()
-        aiArtModel.eval()
+        
+        if torch.cuda.is_available():
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.benchmark = True
+
+        # May sacrifice a bit of accuracy, depending on our needs
+        assert has_native_amp, 'Please update PyTorch to a version with native AMP (or use APEX).'
+        amp_dtype = torch.float16
+        amp_autocast = partial(torch.autocast, device_type=self.device.type, dtype=amp_dtype)
+
+        # data_config = resolve_data_config(vars(args), model=model)
+        
+        # self.aiGeneratorModel = timm.create_model('eva02_base_patch14_448.mim_in22k_ft_in22k_in1k', num_classes=9, in_chans=3, checkpoint_path=path + 'AIModelDetector.pth-6ff3631e.pth')
+        self.aiArtModel = timm.create_model('eva02_base_patch14_448.mim_in22k_ft_in22k_in1k', num_classes=3, in_chans=3, checkpoint_path=path + 'AIArtDetector.pth-af59f7fa.pth')
+        # self.aiGeneratorModel = self.aiGeneratorModel.to(self.device)
+        self.aiArtModel = self.aiArtModel.to(self.device)
+        # self.aiGeneratorModel.eval()
+        self.aiArtModel.eval()
 
+        
         self.transform = timm.data.create_transform(input_size=(3, 448, 448),
         is_training=False,
         use_prefetcher=False,
@@ -36,22 +97,95 @@ class EndpointHandler():
         crop_mode='squash',
         tf_preprocessing=False,
         separate=False)
+        
+        # assert has_compile, 'A version of torch w/ torch.compile() is required for --compile, possibly a nightly.'
+        # torch._dynamo.reset()
+        # model = torch.compile(model, backend=args.torchcompile)
 
-    def __call__(self, data):
+        self.supported_formats = ["JPEG", "PNG", "BMP", "TIFF", "WEBP", "RAW"] #GIF requires its own special implementation to get its frames
+        print("initialized handler.py successfully")
+        # self.label_map = {0: 'Dall-E 2', 1: 'DiscoDiff', 2: 'Midjourney', 3: 'NightCafe', 4: 'NovelAI', 5: 'Stable Diffusion', 6: 'StarryAI', 7: 'WomboDream', 8: 'Artbreeder'}
+        
+    def __call__(self, data): 
         """
        data args:
             inputs: Dict[str, Any]
       Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-        # get inputs
-        image = data.pop("inputs", data)
-        
-        image_tensor1 = self.transform(image).to(self.device)
-        with torch.no_grad():
-            output1 = self.aiArtModel(image_tensor1.unsqueeze(0))
+        inputs = data.pop("inputs")
+        if len(inputs) > 50:
+            return {'error': 'Exceeds max limit of images (50)'}
 
-        return output1
+        image_paths = inputs #['https://google_image.png', '']
+        batch_size = 1  # Set your desired batch size
         
+        results = {}
+        for i in range(0, len(image_paths), batch_size): # For each batch
+
+            batch_paths = image_paths[i:i+batch_size]
+            validUrls = []
+            batch_images = []
+            
+            for j, src in enumerate(batch_paths): # Get all valid images open and inputted in batch_images
+                try:
+                    # Image.open(batch_paths[j]).load() # Tests if image is okay to run inference on.
+                    pos = src.find("base64")
+                    if pos != -1:
+                        # Assuming base64_str is the string value without 'data:image/jpeg;base64,'
+                        new = Image.open(io.BytesIO(base64.decodebytes(bytes(src[pos+7:], "utf-8")))).convert("RGB")
+                        # new.load() Necessary? Does this catch any edge cases? Without this, we don't actually load the image pixels.
+                        batch_images.append(new)
+                        validUrls.append(src)
+                    else: 
+                        try:
+                            # r = requests.get(src, stream=True)
+                            # r.raw.decode_content = True
+                            # new = Image.open(r.raw).convert("RGB")
+                            # new = Image.open(urlopen(src))
+                            headers = {
+                                'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36'
+                            }
+            
+                            r = requests.get(src, headers=headers)
+                            new = Image.open(io.BytesIO(r.content)).convert("RGB")
+                            # new.load()
+                            batch_images.append(new)
+                            validUrls.append(src)
+                        except Exception as e: 
+                            results[src] = {'error': 'Failed to process image'}
+                            # invalid_indices.append(j)
+                            continue
+                    # batch_images.append(batch_paths[j])
+                    
+                except Exception as e:
+                    results[src] = {'error': 'Failed to process image w/ base64 in url'}
+                    continue
+
+            #  width, height = new.size
+
+            # if (width < 250 or height < 250) and len(request.data['srcs']) == 1:
+            #     res['error'] = 'Please use a higher quality image'
+            #     return JsonResponse(res, safe=False, status=status.HTTP_400_BAD_REQUEST)
+
+            batch_tensors = torch.stack([self.transform(img).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')) for img in batch_images])
+            # batch_tensors = torch.unsqueeze(batch_tensors, 0)
+                    
+
+            # batch_images = [Image.open(path) for path in batch_paths]
+            # batch_tensors = torch.stack([preprocess(img) for img in batch_images])
 
+            with torch.no_grad():
+                output1 = self.aiGeneratorModel(batch_tensors)
+                for k, tensor in enumerate(output1):
+                    output = tensor.softmax(-1)
+                    output, indice = output.topk(9)
+                    labels = [self.label_map[x] for x in indice.cpu().numpy().tolist()]
+                    probabilities = [round(i * 100, 2) for i in output.cpu().numpy().tolist()]
+                    single_res = {'prob': probabilities, 'indices': labels}
+                    results[validUrls[k]] = single_res
+            
+        return results
     
+# handler = EndpointHandler()
+# handler.__call__({'inputs': ['']})