Add GUI support for ONNX in WD 14

library/wd14_caption_gui.py

@@ -25,6 +25,9 @@ def caption_images(
     frequency_tags,
     prefix,
     postfix,
+    onnx,
+    append_tags,
+    force_download
 ):
     # Check for images_dir_input
     if train_data_dir == '':
@@ -54,6 +57,12 @@ def caption_images(
         run_cmd += f' --remove_underscore'
     if frequency_tags:
         run_cmd += f' --frequency_tags'
+    if onnx:
+        run_cmd += f' --onnx'
+    if append_tags:
+        run_cmd += f' --append_tags'
+    if force_download:
+        run_cmd += f' --force_download'
 
     if not undesired_tags == '':
         run_cmd += f' --undesired_tags="{undesired_tags}"'
@@ -132,6 +141,20 @@ def gradio_wd14_caption_gui_tab(headless=False):
                 interactive=True,
             )
 
+        with gr.Row():
+            onnx = gr.Checkbox(
+                label='Use onnx',
+                value=False,
+                interactive=True,
+                info="https://github.com/onnx/onnx"
+            )
+            append_tags = gr.Checkbox(
+                label='Append TAGs',
+                value=False,
+                interactive=True,
+                info="This option appends the tags to the existing tags, instead of replacing them."
+            )
+
         with gr.Row():
             replace_underscores = gr.Checkbox(
                 label='Replace underscores in filenames with spaces',
@@ -168,6 +191,12 @@ def gradio_wd14_caption_gui_tab(headless=False):
                 ],
                 value='SmilingWolf/wd-v1-4-convnextv2-tagger-v2',
             )
+            
+            force_download = gr.Checkbox(
+                label='Force model re-download',
+                value=False,
+                info='Usefull to force model re download when switching to onnx',
+            )
 
             general_threshold = gr.Slider(
                 value=0.35,
@@ -215,6 +244,9 @@ def gradio_wd14_caption_gui_tab(headless=False):
                 frequency_tags,
                 prefix,
                 postfix,
+                onnx,
+                append_tags,
+                force_download
             ],
             show_progress=False,
         )

requirements.txt

@@ -19,13 +19,14 @@ lycoris_lora==1.9.0
 # timm==0.6.12
 # fairscale==0.4.13
 # for WD14 captioning (tensorflow)
-# tensorflow==2.10.1
+# tensorflow==2.14.0
 # for WD14 captioning (onnx)
-# onnx==1.14.1
-# onnxruntime-gpu==1.16.0
+onnx==1.14.1
+onnxruntime-gpu==1.16.0
 # onnxruntime==1.16.0
-# this is for onnx: 
-# protobuf==3.20.3
+# this is for onnx:
+# tensorboard==2.14.1
+protobuf==3.20.3
 # open clip for SDXL
 open-clip-torch==2.20.0
 opencv-python==4.7.0.68

requirements_linux.txt

@@ -1,4 +1,4 @@
 torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118 # no_verify leave this to specify not checking this a verification stage
 xformers==0.0.21 bitsandbytes==0.41.1
-tensorboard==2.12.3 tensorflow==2.12.0
+tensorboard==2.14.1 tensorflow==2.14.0
 -r requirements.txt

requirements_linux_docker.txt

@@ -1,5 +1,5 @@
 xformers==0.0.20
 bitsandbytes==0.41.1
 accelerate==0.19.0
-tensorboard==2.12.1
-tensorflow==2.12.0
+tensorboard==2.14.1
+tensorflow==2.14.0

requirements_linux_ipex.txt

@@ -1,3 +1,3 @@
 torch==2.0.1a0+cxx11.abi torchvision==0.15.2a0+cxx11.abi intel_extension_for_pytorch==2.0.110+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/
-tensorboard==2.12.3 tensorflow==2.12.0 intel-extension-for-tensorflow[gpu]
+tensorboard==2.14.1 tensorflow==2.14.0 intel-extension-for-tensorflow[gpu]
 -r requirements.txt

requirements_macos_amd64.txt

@@ -1,4 +1,4 @@
 torch==2.0.0 torchvision==0.15.1 -f https://download.pytorch.org/whl/cpu/torch_stable.html
 xformers bitsandbytes==0.41.1
-tensorflow-macos tensorboard==2.12.1
+tensorflow-macos tensorboard==2.14.1
 -r requirements.txt

requirements_macos_arm64.txt

@@ -1,4 +1,4 @@
 torch==2.0.0 torchvision==0.15.1 -f https://download.pytorch.org/whl/cpu/torch_stable.html
 xformers bitsandbytes==0.41.1
-tensorflow-macos tensorflow-metal tensorboard==2.12.1
+tensorflow-macos tensorflow-metal tensorboard==2.14.1
 -r requirements.txt

requirements_runpod.txt

@@ -1,5 +1,5 @@
 torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118 # no_verify leave this to specify not checking this a verification stage
 xformers==0.0.21 bitsandbytes==0.41.1
-tensorboard==2.12.3 tensorflow==2.12.0 wheel
+tensorboard==2.14.1 tensorflow==2.14.0 wheel
 tensorrt
 -r requirements.txt

requirements_windows_torch2.txt

@@ -2,5 +2,5 @@ torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorc
 xformers==0.0.21
 bitsandbytes==0.35.0 # no_verify
 # https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.41.1-py3-none-win_amd64.whl # no_verify
-tensorboard==2.12.3 tensorflow==2.12.0
+tensorboard==2.14.1 tensorflow==2.14.0
 -r requirements.txt