merge upstream changes

.github/workflows/build.yml

@@ -24,8 +24,9 @@ jobs:
     strategy:
       matrix:
         include:
-          - { result: Linux, runs-on: ubuntu-22.04 }
-          - { result: Windows, runs-on: windows-2022 }
+          - { result: Linux-x64, runs-on: ubuntu-22.04, arch: x64 }
+          - { result: Windows-x64, runs-on: windows-2022, arch: x64 }
+          - { result: Linux-arm, runs-on: ubuntu-22.04-arm, arch: arm64 }
 
     runs-on: ${{ matrix.runs-on }}
 
@@ -40,6 +41,7 @@ jobs:
         with:
           spec: 'BabbleTrainer.spec'
           python_ver: '3.12.8'
+          python_arch: ${{ matrix.arch }}
           requirements: 'requirements.txt'
           upload_exe_with_name: ${{ matrix.result }}
 
@@ -59,20 +61,38 @@ jobs:
       - name: Download Windows artifact
         uses: actions/download-artifact@v4
         with:
-          name: Windows
+          name: Windows-x64
           path: ./artifacts/windows
 
       - name: Download Linux artifact
         uses: actions/download-artifact@v4
         with:
-          name: Linux
+          name: Linux-x64
           path: ./artifacts/linux
 
+      - name: Download Linux ARM artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: Linux-arm
+          path: ./artifacts/linux-arm
+
+      - name: Rename Artifacts
+        run: |
+          sudo mv artifacts/windows/BabbleTrainer.exe artifacts/windows/BabbleTrainer-x64.exe
+          sudo mv artifacts/linux/BabbleTrainer artifacts/linux/BabbleTrainer-x64
+          sudo mv artifacts/linux-arm/BabbleTrainer artifacts/linux-arm/BabbleTrainer-arm64
+
+      - name: Change Permissions
+        run: |
+          chmod +x artifacts/linux/BabbleTrainer-x64
+          chmod +x artifacts/linux-arm/BabbleTrainer-arm64
+
       - name: Publish release
         uses: softprops/action-gh-release@v2
         with:
           files: |
-            artifacts/windows/BabbleTrainer.exe
-            artifacts/linux/BabbleTrainer
+            artifacts/windows/BabbleTrainer-x64.exe
+            artifacts/linux/BabbleTrainer-x64
+            artifacts/linux-arm/BabbleTrainer-arm64
           tag_name: ${{ github.event.inputs.version-number }}
-          prerelease: ${{ github.event.inputs.publish-pre-release }}
+          prerelease: ${{ github.event.inputs.publish-pre-release }}

.gitignore

@@ -1 +1,5 @@
-.idea/
+/babble_data
+/__pycache__
+/build
+/dist
+/venv

requirements.txt

@@ -1,7 +1,5 @@
---extra-index-url https://download.pytorch.org/whl/cpu
-
-# PyTorch: CPU wheels for non-Windows
-torch==2.4.1; platform_system != "Windows"
+# PyTorch: GPU/CPU wheels for non-Windows
+torch==2.9.1; platform_system != "Windows"
 
 # PyTorch for Windows via DirectML backend
 torch-directml; platform_system == "Windows"

trainermin.py

@@ -1,3 +1,6 @@
+import os
+from pathlib import Path
+import tempfile
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -15,6 +18,13 @@
 from collections import deque
 from PIL import Image, ImageFile
 
+# We want to store the temporary '.pth'-files somewhere before we can
+# merge them later. Usually it's fine to put it next to the binary but
+# some setups may want to overwrite this. E.g. the binary is stored 
+# in a immutable location such as system installs.
+tmp_dir = tempfile.TemporaryDirectory(prefix="babble-trainer")
+baseline_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
+
 # Constants
 FLOAT_TO_INT_CONSTANT = 1
 
@@ -26,14 +36,56 @@
 # Optimized alignment parameters
 WIN_SIZE_MUL = 10  # Window size multiplier for perfect accuracy
 
-DEVICE = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = None
+if sys.platform == 'win32':
+    try:
+        import torch_directml
+        import time
+        import torch
+
+        best_time = float("inf")
+        best_idx = None
 
-DEVICE = "cpu"
+        for i in range(torch_directml.device_count()):
+            d = torch_directml.device(i)
+            torch.randn(1, device=d)
 
-if DEVICE != "mps" and DEVICE != "cuda" and sys.platform == 'win32':
-    try:
-        DEVICE = torch_directml.device(0)
-    except: DEVICE = "cpu"
+            times = []
+            for _ in range(100):
+                x = torch.randn(2048, 2048, device=d)
+                start = time.time()
+                _ = x @ x
+                times.append(time.time() - start)
+
+            avg = sum(times) / len(times)
+            print(i, torch_directml.device_name(i), avg)
+
+
+            if avg < best_time:
+                best_time = avg
+                best_idx = i
+
+        if best_idx is not None:
+            DEVICE = torch_directml.device(best_idx)
+            name = torch_directml.device_name(best_idx)
+            print("Using DirectML device:", name, flush=True)
+        else:
+            DEVICE = "cpu"
+
+    except:
+        DEVICE = "cpu"
+elif sys.platform == "darwin":
+    # Apple. TODO: verify this works.
+    if torch.backends.mps.is_available():
+        DEVICE = torch.device("mps")
+elif sys.platform.startswith("linux"):
+    # Linux. Assume no DirectML, just use whatever is available.
+    if torch.cuda.is_available():
+        # This also may include ROCm, it's just opaque.
+        DEVICE = torch.device("cuda")
+# Fall back to CPU
+if DEVICE is None:
+    DEVICE = torch.device("cpu")
 
 class MicroChad(nn.Module):
     def __init__(self):
@@ -47,7 +99,14 @@ def __init__(self):
         self.fc = nn.Linear(212, 3)
 
         self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
-        self.adaptive = nn.AdaptiveMaxPool2d(output_size=1)
+        # ONNX does not support AdaptiveMaxPool2D:
+        # * https://github.com/pytorch/pytorch/issues/169949
+        # * https://github.com/pytorch/pytorch/issues/5310
+        # AdaptiveMaxPool of output size 1 can be replaced with GlobalMaxPool.
+        # Originally:
+        #   self.adaptive = nn.AdaptiveMaxPool2d(output_size=1)
+        # HOWEVER we can't export so here this lies
+        #self.adaptive = nn.AdaptiveMaxPool2d(output_size=1)
 
         self.act = nn.ReLU(inplace=True)
         self.sigmoid = nn.Sigmoid()
@@ -76,7 +135,7 @@ def forward(self, x, return_blends=True):
         x = self.conv6(x)
         x = self.act(x)
 
-        x = self.adaptive(x)
+        x = torch.amax(x, dim=(2, 3), keepdim=True)
         x = torch.flatten(x, 1)
         if not return_blends:
             return x
@@ -1250,7 +1309,7 @@ def warmup_fn(epoch):
                 optimizerE.step()
                 #progress.set_description("(%d/%d) Loss: %.6f" % (i, max_i, float(loss)))
                 #  optimizerD.step()
-                print("\rBatch %u/%u, Loss: %.6f" % (i, max_i, float(loss)), flush=True)
+                print("Batch %u/%u, Loss: %.6f" % (i, max_i, float(loss)), flush=True)
 
                 # Print statistics
                 running_loss += loss.item()
@@ -1309,9 +1368,9 @@ def main():
     print(model_L, flush=True)
     print(model_R, flush=True)
 
-    model_L.load_state_dict(torch.load("baseline_L.pth", map_location="cpu"))
+    model_L.load_state_dict(torch.load(os.path.join(baseline_dir, "baseline_L.pth"), map_location="cpu", weights_only=False))
     model_L.to(DEVICE)
-    model_R.load_state_dict(torch.load("baseline_R.pth", map_location="cpu"))
+    model_R.load_state_dict(torch.load(os.path.join(baseline_dir, "baseline_R.pth"), map_location="cpu", weights_only=False))
     model_R.to(DEVICE)
     trained_model_L = model_L
     trained_model_R = model_R
@@ -1394,8 +1453,8 @@ def main():
     # Save the final model
     #torch.save(trained_model.state_dict(), "final_model_temporal_que_tuned_2.pth")
 
-    torch.save(trained_model_L.state_dict(), "left_tuned.pth")
-    torch.save(trained_model_R.state_dict(), "right_tuned.pth")
+    torch.save(trained_model_L.state_dict(), os.path.join(tmp_dir.name, "left_tuned.pth"))
+    torch.save(trained_model_R.state_dict(), os.path.join(tmp_dir.name, "right_tuned.pth"))
 
     multi = MultiChad()
     multi.left.load_state_dict(trained_model_L.state_dict())
@@ -1422,6 +1481,7 @@ def main():
         }
     )
     print("Model exported to ONNX: " + sys.argv[2], flush=True)
+    tmp_dir.cleanup()
 
 if __name__ == "__main__":
     main()