Views
No views yet
mlx-DeepDanBooru is available for: MacBook Pro / Air, Mac mini, iMac.no pytorch neededmlx-DeepDanBooru Model implementation is inspired by a PyTorch implementation of AUTOMATIC1111/TorchDeepDanbooru1conda create -n mlx026 python=3.12
2conda activate mlx026
3#
4pip install numpy
5pip install pillowpip install mlxmlx-DeepDanbooru is base on mlx version: 0.26.11#
2python ./infer.py /Volumes/HDD4/hazhu/example/1.png
3#
4# inference folder's all images one-by-one
5# python ./infer.py /Volumes/HDD4/hazhu/example
6
7#
8# inference folder's all images parallelly
9# python ./infer_multiprocessing.py /Volumes/HDD4/hazhu/example1python ./segment_infer.py /Volumes/HDD4/hazhu/example
2#segment_infer.py will crop the image into 3 segments: top, middle, bottom; then inference each segment one-by-one, and merge tags at last.
segment_infer.py can enable cache, as the same input will always get the same tags, you can cache the results for speedup.1import numpy as np
2from PIL import Image
3
4# using apple silicon's MLX
5# not Pytorch
6import mlx.core as mx
7from mlxDeepDanBooru.mlx_deep_danbooru_model import mlxDeepDanBooruModel
8
9model_path = "models/model-resnet_custom_v3_fp32_mlx.safetensors"
10model_tags = np.load('models/tags-resnet_custom_v3_mlx.npy')
11
12mlx_dan = mlxDeepDanBooruModel()
13mlx_dan.load_weights(model_path)
14mx.eval(mlx_dan.parameters())
15
16def danbooru_tags(fpath):
17 tags = []
18
19 try:
20 pic = Image.open(fpath).convert("RGB").resize((512, 512))
21 a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
22 #
23 x = mx.array(a)
24 y = mlx_dan(x)[0]
25 ylen = len(y)
26 ylst = y.tolist()
27 for i in range(ylen):
28 if ylst[i] >= 0.55:
29 # 0.55 can be changed for demand: 0.0 ~ 1.0
30 #print(model_tags[i].item(), p)
31 tags.append(model_tags[i].item())
32
33 except Exception as err:
34 print(err)
35 tags = []
36
37 return tags
38
39def image_infer(fpath):
40 tags = danbooru_tags(fpath)
41 return tags
42
43tags_1 = image_infer("example/1.png")
44tags_2 = image_infer("example/2.png")
45
46print(tags_1)
47# will show tags: ['1girl', 'beach', 'black_hair', 'blurry', 'blurry_background', 'blurry_foreground', 'building', 'bush', 'christmas_tree', 'day', 'depth_of_field', 'field', 'grass', 'lake', 'looking_at_viewer', 'mountain', 'nature', 'outdoors', 'palm_leaf', 'palm_tree', 'park', 'park_bench', 'path', 'photo_background', 'plant', 'river', 'road', 'skirt', 'sky', 'smile', 'striped', 'striped_dress', 'striped_shirt', 'tree', 'vertical-striped_shirt', 'vertical_stripes', 'rating:safe']
48
49print(tags_2)
50# will show tags: ['1girl', '3d', 'blurry', 'blurry_background', 'blurry_foreground', 'brown_eyes', 'brown_hair', 'bush', 'christmas_tree', 'cosplay_photo', 'day', 'depth_of_field', 'field', 'floral_print', 'foliage', 'forest', 'garden', 'grass', 'jungle', 'lips', 'long_hair', 'long_sleeves', 'looking_at_viewer', 'nature', 'on_grass', 'outdoors', 'palm_tree', 'park', 'path', 'plant', 'potted_plant', 'realistic', 'smile', 'solo', 'tree', 'upper_body', 'white_dress', 'rating:safe']
511024x1024 pixel imageexample folder, on Mac Mini M4, mlx-DeepDanBooru inference Speed (image one-by-one):SPEED: 0.12 seconds per image3840x5760 pixel imagemlx-DeepDanBooru with multiprocessing(4 parallel tasks), run infer_multiprocessing.py:SPEED: 0.18 seconds per image 3840x5760), speed (seconds per image, shorter is better):
Pytorch + CPU : 0.54
Pytorch + MPS : 0.34
mlx-DeepDanBooru : 0.18
