Upload folder using huggingface_hub

infer.py

@@ -1,13 +1,21 @@
+import os
 import time
+import glob
 import numpy as np
 from PIL import Image, ImageDraw
 
 import mlx.core as mx
 from mlxDeepDanBooru.mlx_deep_danbooru_model import mlxDeepDanBooruModel
 
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed, wait, FIRST_COMPLETED
+from copy import deepcopy
 
-model_path = "models/model-resnet_custom_v3_mlx.npz"
-tags_path = 'models/tags-resnet_custom_v3_mlx.npy'
+ROOTDIR = os.path.dirname(os.path.abspath(__file__))
+IMAGEDIR = f'{ROOTDIR}/example'
+
+
+model_path = f"{ROOTDIR}/models/model-resnet_custom_v3_mlx.npz"
+tags_path = f'{ROOTDIR}/models/tags-resnet_custom_v3_mlx.npy'
 
 mlx_dan = mlxDeepDanBooruModel()
 mlx_dan.load_weights(model_path)
@@ -15,50 +23,54 @@ mx.eval(mlx_dan.parameters())
 
 
 model_tags = np.load(tags_path)
-print(f'total tags: {len(model_tags)}')
+#print(f'total tags: {len(model_tags)}')
 
 def danbooru_tags(fpath):
+    results = {}
     tags = []
+
     pic = Image.open(fpath).convert("RGB").resize((512, 512))
     a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
 
     x = mx.array(a)
     y = mlx_dan(x)[0]
 
-    for n in range(10):
-        mlx_dan(x)
-    for i, p in enumerate(y):
-        if p >= 0.5:
-            #print(model_tags[i].item(), p)
-            tags.append(model_tags[i].item())
+    try:
+        for n in range(10):
+            mlx_dan(x)
+        for i, p in enumerate(y):
+           if p >= 0.55:             
+                #print(model_tags[i].item(), p)
+                tags.append(model_tags[i].item())
+    except Exception as err:
+        print(err)
+
+    results[fpath] = tags
+    return results
 
-    return tags
 
-image_count = 0
 def image_infer(fpath):
-    global image_count
     tags = danbooru_tags(fpath)
-    image_count += 1
     return tags
 
-
 t1 = time.time()
-tags_1 = image_infer("example/1.png")
-tags_2 = image_infer("example/2.png")
+
+tags_1 = image_infer(f'{IMAGEDIR}/1.png')
+tags_2 = image_infer(f'{IMAGEDIR}/2.png')
 
 t2 = time.time()
 
 print(tags_1)
 print(tags_2)
-# print(tags_3)
-# print(tags_4)
-# print(tags_5)
 
-print("-----------")
-print(f'infer speed(with mlx): {(t2 - t1)/image_count} seconds per image')
+print(f'2 images: infer speed(with mlx): {(t2 - t1)/2} seconds per image')
+
+
+
 
 
 
+    
 
 
 

infer_multiprocessing.py

@@ -0,0 +1,93 @@
+import os
+import time
+import glob
+import numpy as np
+from PIL import Image, ImageDraw
+
+import mlx.core as mx
+from mlxDeepDanBooru.mlx_deep_danbooru_model import mlxDeepDanBooruModel
+
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed, wait, FIRST_COMPLETED
+from copy import deepcopy
+
+ROOTDIR = os.path.dirname(os.path.abspath(__file__))
+IMAGEDIR = f'{ROOTDIR}/example'
+
+worker_count = os.cpu_count()
+# worker_count depends on your unified-memory size
+# if oom, decrease the number
+
+model_path = f"{ROOTDIR}/models/model-resnet_custom_v3_mlx.npz"
+tags_path = f'{ROOTDIR}/models/tags-resnet_custom_v3_mlx.npy'
+
+mlx_dan = mlxDeepDanBooruModel()
+mlx_dan.load_weights(model_path)
+mx.eval(mlx_dan.parameters())
+
+
+model_tags = np.load(tags_path)
+#print(f'total tags: {len(model_tags)}')
+
+def danbooru_tags(fpath):
+    results = {}
+    tags = []
+
+    pic = Image.open(fpath).convert("RGB").resize((512, 512))
+    a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
+
+    x = mx.array(a)
+    y = mlx_dan(x)[0]
+
+    try:
+        for n in range(10):
+            mlx_dan(x)
+        for i, p in enumerate(y):
+           if p >= 0.55:
+                
+                #print(model_tags[i].item(), p)
+                tags.append(model_tags[i].item())
+    except Exception as err:
+        print(err)
+
+    results[fpath] = tags
+    return results
+
+
+def image_infer(fpath):
+    tags = danbooru_tags(fpath)
+    return tags
+
+
+def batch_infer(image_list):
+    workers = min(len(image_list), worker_count)
+    print(f'workers: {workers}: {os.cpu_count()}')
+    with ProcessPoolExecutor(max_workers=workers) as executor:
+        process_results = list(executor.map(image_infer, image_list))
+        return process_results
+
+
+
+if __name__ == '__main__':
+    image_list = []
+    for root, dirs, files in os.walk(IMAGEDIR, True):
+        for file in files:
+            if not file[-4:].lower() in [".png", ".jpg", "jpeg"]:
+                continue
+            fpath = os.path.join(root, file).replace("\\","/")
+            image_list.append(fpath)
+    
+    #print(image_list)
+
+
+    t1 = time.time()
+    lines = batch_infer(image_list)
+    t2 = time.time()
+
+    for line in lines:
+        print(line)
+        print("-----------")
+
+    print(f'{len(image_list)} images: infer speed(with mlx): {(t2 - t1)/len(image_list)} seconds per image')
+
+
+