Update app.py

app.py

@@ -2,25 +2,12 @@ import os, re, cv2
 from typing import Mapping, Tuple, Dict
 import gradio as gr
 import numpy as np
+import io
 import pandas as pd
 from PIL import Image
 from huggingface_hub import hf_hub_download
 from onnxruntime import InferenceSession
 
-from functools import wraps
-
-def memoize(func):
-    cache = {}
-    cache_maxsize = 25
-    @wraps(func)
-    def memoized_func(*args):
-        if args not in cache:
-            cache[args] = func(*args)
-            if len(cache) > cache_maxsize:
-                cache.popitem(last=False)
-        return cache[args]
-
-    return memoized_func
 # noinspection PyUnresolvedReferences
 def make_square(img, target_size):
     old_size = img.shape[:2]
@@ -64,7 +51,7 @@ class WaifuDiffusionInterrogator:
 
         self.__initialized = False
         self._model, self._tags = None, None
-
+        self.cache = {}
     def _init(self) -> None:
         if self.__initialized:
             return
@@ -111,8 +98,14 @@ class WaifuDiffusionInterrogator:
         full_tags['confidence'] = confidence[0]
 
         return full_tags
-
     def interrogate(self, image: Image) -> Tuple[Dict[str, float], Dict[str, float]]:
+        
+        imgByteArr = io.BytesIO()
+        image.save(imgByteArr, format=image.format)
+        imgByteArr = imgByteArr.getvalue()
+        if imgByteArr in cache:
+            return cache[imgByteArr]
+
         full_tags = self._calculation(image)
 
         # first 4 items are for rating (general, sensitive, questionable, explicit)
@@ -120,6 +113,9 @@ class WaifuDiffusionInterrogator:
 
         # rest are regular tags
         tags = dict(full_tags[full_tags['category'] != 9][['name', 'confidence']].values)
+        self.cache[imgByteArr] = (ratings, tags)
+        if len(cache) > 25:
+            cache.popitem(last=False)
 
         return ratings, tags
 
@@ -132,7 +128,7 @@ WAIFU_MODELS: Mapping[str, WaifuDiffusionInterrogator] = {
 }
 RE_SPECIAL = re.compile(r'([\\()])')
 
-@memoize
+
 def image_to_wd14_tags(image: Image.Image, model_name: str, threshold: float,
                        use_spaces: bool, use_escape: bool, include_ranks=False, score_descend=True) \
         -> Tuple[Mapping[str, float], str, Mapping[str, float]]: