yoinked commited on
Commit
dc036e7
1 Parent(s): 8e25602

memoize and also based

Browse files
Files changed (1) hide show
  1. app.py +19 -3
app.py CHANGED
@@ -7,7 +7,20 @@ from PIL import Image
7
  from huggingface_hub import hf_hub_download
8
  from onnxruntime import InferenceSession
9
 
10
-
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  # noinspection PyUnresolvedReferences
12
  def make_square(img, target_size):
13
  old_size = img.shape[:2]
@@ -119,7 +132,7 @@ WAIFU_MODELS: Mapping[str, WaifuDiffusionInterrogator] = {
119
  }
120
  RE_SPECIAL = re.compile(r'([\\()])')
121
 
122
-
123
  def image_to_wd14_tags(image: Image.Image, model_name: str, threshold: float,
124
  use_spaces: bool, use_escape: bool, include_ranks=False, score_descend=True) \
125
  -> Tuple[Mapping[str, float], str, Mapping[str, float]]:
@@ -139,12 +152,15 @@ def image_to_wd14_tags(image: Image.Image, model_name: str, threshold: float,
139
  tag_outformat = tag
140
  if use_spaces:
141
  tag_outformat = tag_outformat.replace('_', ' ')
 
 
 
142
  if use_escape:
143
  tag_outformat = re.sub(RE_SPECIAL, r'\\\1', tag_outformat)
144
  if include_ranks:
145
  tag_outformat = f"({tag_outformat}:{score:.3f})"
146
  text_items.append(tag_outformat)
147
- output_text = ', '.join(text_items)
148
 
149
  return ratings, output_text, filtered_tags
150
 
 
7
  from huggingface_hub import hf_hub_download
8
  from onnxruntime import InferenceSession
9
 
10
+ from functools import wraps
11
+
12
+ def memoize(func):
13
+ cache = {}
14
+ cache_maxsize = 25
15
+ @wraps(func)
16
+ def memoized_func(*args):
17
+ if args not in cache:
18
+ cache[args] = func(*args)
19
+ if len(cache) > cache_maxsize:
20
+ cache.popitem(last=False)
21
+ return cache[args]
22
+
23
+ return memoized_func
24
  # noinspection PyUnresolvedReferences
25
  def make_square(img, target_size):
26
  old_size = img.shape[:2]
 
132
  }
133
  RE_SPECIAL = re.compile(r'([\\()])')
134
 
135
+ @memoize
136
  def image_to_wd14_tags(image: Image.Image, model_name: str, threshold: float,
137
  use_spaces: bool, use_escape: bool, include_ranks=False, score_descend=True) \
138
  -> Tuple[Mapping[str, float], str, Mapping[str, float]]:
 
152
  tag_outformat = tag
153
  if use_spaces:
154
  tag_outformat = tag_outformat.replace('_', ' ')
155
+ tag_outformat.replace(' ', '-')
156
+ else:
157
+ tag_outformat.replace('_', '-')
158
  if use_escape:
159
  tag_outformat = re.sub(RE_SPECIAL, r'\\\1', tag_outformat)
160
  if include_ranks:
161
  tag_outformat = f"({tag_outformat}:{score:.3f})"
162
  text_items.append(tag_outformat)
163
+ output_text = ' '.join(text_items)
164
 
165
  return ratings, output_text, filtered_tags
166