From bad78c67505c4b77b17db74d5eb0458ccd748511 Mon Sep 17 00:00:00 2001
From: firestar5683 <168790843+firestar5683@users.noreply.github.com>
Date: Sat, 19 Sep 2026 13:54:31 -0700
Subject: [PATCH] Release duplicate Torch diffusion weights before loading host
 VAE

---
 .../ace_chestnut_20260916/host_hook_adapter.py  | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/roadscore/experiments/ace_chestnut_20260916/host_hook_adapter.py b/roadscore/experiments/ace_chestnut_20260916/host_hook_adapter.py
index 18e50c1770..4c0b67f979 100644
--- a/roadscore/experiments/ace_chestnut_20260916/host_hook_adapter.py
+++ b/roadscore/experiments/ace_chestnut_20260916/host_hook_adapter.py
@@ -79,8 +79,23 @@ class HostHookAdapter:
         status, ok = lm.initialize(str(self.base / 'models/ace/checkpoints'), 'acestep-5Hz-lm-1.7B', backend='mlx', device='mps')
         if not ok:
             raise RuntimeError(status)
+        original_init = handler._init_mlx_dit
+        def initialize_without_duplicate(*args, **kwargs):
+            ok = original_init(*args, **kwargs)
+            if not ok or handler.mlx_decoder is None:
+                raise RuntimeError('MLX conversion failed; refusing Torch diffusion fallback')
+            import gc
+            import torch
+            handler.model.decoder = None
+            def no_torch_diffusion(*args, **kwargs):
+                raise RuntimeError('Torch diffusion disabled after verified MLX conversion')
+            handler.model.generate_audio = no_torch_diffusion
+            gc.collect()
+            torch.mps.empty_cache()
+            return ok
+        handler._init_mlx_dit = initialize_without_duplicate
         status, ok = handler.initialize_service(str(self.base / 'models/ace'), config_path='acestep-v15-turbo',
-            device='mps', use_mlx_dit=True, offload_to_cpu=True, offload_dit_to_cpu=False)
+            device='mps', use_mlx_dit=True, offload_to_cpu=True, offload_dit_to_cpu=True)
         if not ok or not handler.use_mlx_dit or handler.mlx_decoder is None:
             raise RuntimeError(f'No safe MLX preparation boundary: {status}')
         self.handler, self.lm = handler, lm