From 7f98198c1499908c29cebb566b51732d709255f7 Mon Sep 17 00:00:00 2001 From: Fro Zen Date: Thu, 4 Dec 2025 15:56:32 -0600 Subject: [PATCH 1/6] Backport ARM build --- .github/workflows/build.yml | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 6312e01..d340c4f 100644 --- a/.github/workflows/build.yml +++ b/.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 } + - { result: Windows-x64, runs-on: windows-2022 } + - { result: Linux-arm, runs-on: ubuntu-22.04-arm } runs-on: ${{ matrix.runs-on }} @@ -59,20 +60,33 @@ 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: 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 }} \ No newline at end of file From 70f752749649a797ede0bf869d174d29b2545a77 Mon Sep 17 00:00:00 2001 From: Fro Zen Date: Thu, 4 Dec 2025 15:58:44 -0600 Subject: [PATCH 2/6] oops --- .github/workflows/build.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d340c4f..e85529b 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -35,6 +35,11 @@ jobs: uses: actions/checkout@v5 with: submodules: 'recursive' + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: 3.12.8 - name: Create Executable uses: sayyid5416/pyinstaller@v1 From ed73128e0529a816ee7d710ef5e5eebf79af5d47 Mon Sep 17 00:00:00 2001 From: Fro Zen Date: Thu, 4 Dec 2025 15:59:50 -0600 Subject: [PATCH 3/6] oops --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e85529b..31c3b04 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -37,7 +37,7 @@ jobs: submodules: 'recursive' - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: 3.12.8 From 2e6737f4c70147c8a9b5b5c8654935b0f46b3288 Mon Sep 17 00:00:00 2001 From: dfgHiatus Date: Fri, 19 Dec 2025 12:17:16 -1000 Subject: [PATCH 4/6] Merge GPU accel to Linux, add temp directories, update gitignore --- .github/workflows/build.yml | 12 ++---- .gitignore | 6 ++- requirements.txt | 6 +-- trainermin.py | 82 ++++++++++++++++++++++++++++++++----- 4 files changed, 82 insertions(+), 24 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 31c3b04..5a6b31a 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,9 +24,9 @@ jobs: strategy: matrix: include: - - { result: Linux-x64, runs-on: ubuntu-22.04 } - - { result: Windows-x64, runs-on: windows-2022 } - - { result: Linux-arm, runs-on: ubuntu-22.04-arm } + - { 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 }} @@ -35,17 +35,13 @@ jobs: uses: actions/checkout@v5 with: submodules: 'recursive' - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: 3.12.8 - name: Create Executable uses: sayyid5416/pyinstaller@v1 with: spec: 'BabbleTrainer.spec' python_ver: '3.12.8' + python_arch: ${{ matrix.arch }} requirements: 'requirements.txt' upload_exe_with_name: ${{ matrix.result }} diff --git a/.gitignore b/.gitignore index 62c8935..617fc53 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,5 @@ -.idea/ \ No newline at end of file +/babble_data +/__pycache__ +/build +/dist +/venv diff --git a/requirements.txt b/requirements.txt index b6924ef..9d8d56a 100644 --- a/requirements.txt +++ b/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" diff --git a/trainermin.py b/trainermin.py index 11b5cf5..49399ad 100644 --- a/trainermin.py +++ b/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,6 +99,13 @@ 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) + # 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) @@ -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() From 8b9aec925ca69b71a9c957e39908887bfccec144 Mon Sep 17 00:00:00 2001 From: Summer <71572746+SummerSigh@users.noreply.github.com> Date: Thu, 1 Jan 2026 10:09:22 -0800 Subject: [PATCH 5/6] Replace Adaptive with amax --- trainermin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trainermin.py b/trainermin.py index 49399ad..3475162 100644 --- a/trainermin.py +++ b/trainermin.py @@ -106,7 +106,7 @@ def __init__(self): # 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.adaptive = nn.AdaptiveMaxPool2d(output_size=1) self.act = nn.ReLU(inplace=True) self.sigmoid = nn.Sigmoid() @@ -135,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 From d03ddb09af7d8538ff63fa23803e3189423b5bc5 Mon Sep 17 00:00:00 2001 From: dfgHiatus <51272212+dfgHiatus@users.noreply.github.com> Date: Fri, 13 Mar 2026 19:04:02 -0500 Subject: [PATCH 6/6] Make trainers executable on Linux --- .github/workflows/build.yml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 5a6b31a..024adc2 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -82,6 +82,11 @@ jobs: 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: @@ -90,4 +95,4 @@ jobs: artifacts/linux/BabbleTrainer-x64 artifacts/linux-arm/BabbleTrainer-arm64 tag_name: ${{ github.event.inputs.version-number }} - prerelease: ${{ github.event.inputs.publish-pre-release }} \ No newline at end of file + prerelease: ${{ github.event.inputs.publish-pre-release }}