speed up

3、tagger.ps1

@@ -33,6 +33,7 @@ $TagConfig = @{
 #region Environment Setup
 # Activate python venv
 Set-Location $PSScriptRoot
+$env:PYTHONPATH = "$PSScriptRoot;$env:PYTHONPATH"
 $VenvPaths = @(
     "./venv/Scripts/activate",
     "./.venv/Scripts/activate",
@@ -87,7 +88,7 @@ if ($TagConfig.tag_replacement) { [void]$ExtArgs.Add("--tag_replacement=$($TagCo
 Write-Output "Starting tagger..."
 
 # Run tagger
-python -m utils.wdtagger `
+accelerate launch --num_cpu_threads_per_process=8 "./utils/wdtagger.py" `
     $Config.train_data_dir `
     --thresh=$($Config.thresh) `
     --caption_extension .txt `

module/lanceImport.py

@@ -323,7 +323,11 @@ def load_metadata(
                     try:
                         with iio.imopen(file_path, "r") as file:
                             first_frame = file.read(index=0)
-                            channels = first_frame.shape[2] if len(first_frame.shape) > 2 else 1
+                            channels = (
+                                first_frame.shape[2]
+                                if len(first_frame.shape) > 2
+                                else 1
+                            )
                             depth = first_frame.dtype.itemsize * 8
                     except Exception as e:
                         console.print(
@@ -644,13 +648,14 @@ def transform2lance(
             schema, process(data, save_binary, import_mode)
         )
 
-        if not_save_disk:
-            table = reader.read_all()
-            return lance.dataset(table)
-
         dataset_path = Path(dataset_dir) / f"{output_name}.lance"
+        mode = "append" if dataset_path.exists() else "create"
+
         lancedataset = lance.write_dataset(
-            reader, str(dataset_path), schema, mode="overwrite"
+            reader,
+            str(dataset_path),
+            schema,
+            mode=mode if not_save_disk else "overwrite",
         )
 
         try:

requirements-uv.txt

@@ -1,5 +1,7 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements.txt -o requirements-uv.txt --index-strategy unsafe-best-match --no-build-isolation -p 3.10
+accelerate==1.4.0
+    # via -r requirements.txt
 aiohappyeyeballs==2.4.6
     # via aiohttp
 aiohttp==3.11.12
@@ -111,7 +113,7 @@ httpx==0.27.2
     #   mistralai
     #   openai
 huggingface-hub==0.29.1
-    # via -r requirements.txt
+    # via accelerate
 humanfriendly==10.0
     # via coloredlogs
 idna==3.10
@@ -136,7 +138,7 @@ lxml==4.9.4
     #   pyquery
 markdown-it-py==3.0.0
     # via rich
-markupsafe==3.0.2
+markupsafe==3.0.1
     # via jinja2
 mdurl==0.1.2
     # via markdown-it-py
@@ -156,16 +158,21 @@ networkx==3.4.2
     # via torch
 numpy==1.26.4
     # via
+    #   accelerate
     #   imageio
     #   onnxruntime-gpu
+    #   opencv-python
     #   pandas
     #   pylance
 onnxruntime-gpu==1.20.1
     # via -r requirements.txt
 openai==1.59.7
     # via -r requirements.txt
+opencv-python==4.11.0.86
+    # via -r requirements.txt
 packaging==24.2
     # via
+    #   accelerate
     #   deprecation
     #   hbutils
     #   huggingface-hub
@@ -202,6 +209,8 @@ protobuf==5.29.2
     #   grpcio-status
     #   onnxruntime-gpu
     #   proto-plus
+psutil==7.0.0
+    # via accelerate
 pyanimeinfo==0.0.4
     # via -r requirements.txt
 pyarrow==18.1.0
@@ -249,7 +258,9 @@ pytimeparse==1.1.8
 pytz==2024.2
     # via pandas
 pyyaml==6.0.2
-    # via huggingface-hub
+    # via
+    #   accelerate
+    #   huggingface-hub
 requests==2.32.3
     # via
     #   dashscope
@@ -265,6 +276,8 @@ rich-pixels==3.0.1
     # via -r requirements.txt
 rsa==4.9
     # via google-auth
+safetensors==0.5.2
+    # via accelerate
 setuptools==75.6.0
     # via
     #   hbutils
@@ -283,7 +296,9 @@ sympy==1.13.1
 toml==0.10.2
     # via -r requirements.txt
 torch==2.5.1+cu124
-    # via -r requirements.txt
+    # via
+    #   -r requirements.txt
+    #   accelerate
 tqdm==4.67.1
     # via
     #   google-generativeai

requirements.txt

@@ -21,6 +21,7 @@ dashscope
 pyanimeinfo
 
 #WDtagger
-huggingface_hub
+opencv-python
+accelerate
 torch==2.5.1+cu124
 onnxruntime-gpu>=1.20