ace-step-mlx-quant.patch

diff --git a/acestep/api/model_download.py b/acestep/api/model_download.py
index 92257c9..00f3176 100644
--- a/acestep/api/model_download.py
+++ b/acestep/api/model_download.py
@@ -101,6 +101,9 @@ def ensure_model_downloaded(model_name: str, checkpoint_dir: str) -> str:
         print(f"[Model Download] Model {model_name} already exists at {model_path}")
         return model_path
 
+    if os.environ.get("ACESTEP_NO_DOWNLOAD") == "1":  # Audante patch: models come only with Audante's checksums
+        raise RuntimeError(f"{model_name} is not installed ({model_path}); ACESTEP_NO_DOWNLOAD=1 forbids downloading it")
+
     repo_id = MODEL_REPO_MAPPING.get(model_name, DEFAULT_REPO_ID)
     print(f"[Model Download] Model {model_name} not found, checking network...")
 
diff --git a/acestep/api/startup_llm_init.py b/acestep/api/startup_llm_init.py
index dd5199d..8759630 100644
--- a/acestep/api/startup_llm_init.py
+++ b/acestep/api/startup_llm_init.py
@@ -81,6 +81,10 @@ def initialize_llm_at_startup(
             ensure_model_downloaded(lm_model_name, checkpoint_dir)
         except Exception as exc:
             print(f"[API Server] Warning: Failed to download LLM model: {exc}")
+            if os.environ.get("ACESTEP_NO_DOWNLOAD") == "1":  # Audante patch: say why there is no language model
+                app.state._llm_initialized = False
+                app.state._llm_init_error = f"{exc}" + ("" if is_supported else f" ({warning_msg})")
+                return
 
         llm_status, llm_ok = llm_handler.initialize(
             checkpoint_dir=checkpoint_dir,
diff --git a/acestep/core/generation/handler/init_service_loader.py b/acestep/core/generation/handler/init_service_loader.py
index 040c5ce..68eca1d 100644
--- a/acestep/core/generation/handler/init_service_loader.py
+++ b/acestep/core/generation/handler/init_service_loader.py
@@ -11,6 +11,55 @@ from acestep import gpu_config
 from .init_service_loader_components import InitServiceLoaderComponentsMixin
 
 
+# Audante patch: a checkpoint may carry its DiT decoder as a pre-quantized MLX file (mlx_dit.safetensors with
+# mlx_dit.json, next to a model.safetensors without decoder.* weights). Then the PyTorch decoder stays empty and
+# mlx_dit_init.py fills the MLX decoder from that file instead of converting it from PyTorch.
+def audante_mlx_dit_file(checkpoint_path, device):
+    """The checkpoint's pre-quantized MLX DiT, or None when it has none."""
+    path = os.path.join(checkpoint_path, "mlx_dit.safetensors")
+    if not os.path.exists(path):
+        return None
+    from acestep.models.mlx import mlx_available
+
+    if device not in ("mps", "cpu") or not mlx_available():
+        raise RuntimeError(f"{path}: this checkpoint's DiT runs only in MLX (Apple Silicon), device here: {device}")
+    return path
+
+
+class _AudanteSkipDecoderInit:
+    """While a checkpoint without decoder weights loads, transformers would initialise them at random (time and
+    memory for nothing): skip that, the decoder is emptied right after loading."""
+
+    def __init__(self, active):
+        self.active = bool(active)
+
+    def __enter__(self):
+        if self.active:
+            from transformers import modeling_utils as mu
+
+            self.orig = orig = mu.PreTrainedModel._initialize_missing_keys
+            mu.PreTrainedModel._initialize_missing_keys = lambda model, keys, quantized: orig(
+                model, [k for k in keys if not k.startswith("decoder.")], quantized)
+        return self
+
+    def __exit__(self, *exc):
+        if self.active:
+            from transformers import modeling_utils as mu
+
+            mu.PreTrainedModel._initialize_missing_keys = self.orig
+        return False
+
+
+def audante_empty_decoder(model):
+    """Replace every tensor of the PyTorch DiT decoder by an empty one (frees its memory; MLX runs the DiT)."""
+    for param in model.decoder.parameters():
+        param.data = torch.empty(0, dtype=param.dtype, device=param.device)
+    for name, buf in list(model.decoder.named_buffers()):
+        owner, _, attr = name.rpartition(".")
+        module = model.decoder.get_submodule(owner) if owner else model.decoder
+        setattr(module, attr, torch.empty(0, dtype=buf.dtype, device=buf.device))
+
+
 class InitServiceLoaderMixin(InitServiceLoaderComponentsMixin):
     """Helpers for heavy model component loading."""
 
@@ -169,15 +218,17 @@ class InitServiceLoaderMixin(InitServiceLoaderComponentsMixin):
 
         last_attn_error = None
         self.model = None
+        self._audante_mlx_dit = audante_mlx_dit_file(model_checkpoint_path, device)  # Audante patch
         for candidate in attn_candidates:
             try:
                 logger.info(f"[initialize_service] Attempting to load model with attention implementation: {candidate}")
-                self.model = AutoModel.from_pretrained(
-                    model_checkpoint_path,
-                    trust_remote_code=True,
-                    attn_implementation=candidate,
-                    dtype=self.dtype,
-                )
+                with _AudanteSkipDecoderInit(self._audante_mlx_dit):
+                    self.model = AutoModel.from_pretrained(
+                        model_checkpoint_path,
+                        trust_remote_code=True,
+                        attn_implementation=candidate,
+                        dtype=self.dtype,
+                    )
                 attn_implementation = candidate
                 break
             except Exception as exc:
@@ -193,6 +244,9 @@ class InitServiceLoaderMixin(InitServiceLoaderComponentsMixin):
         self.config = self.model.config
         self._sync_alignment_config()
         self._apply_cuda_bool_argsort_workaround()
+        if self._audante_mlx_dit:  # Audante patch: the DiT comes from the MLX file
+            audante_empty_decoder(self.model)
+            logger.info("[initialize_service] Audante patch: PyTorch decoder left empty (pre-quantized MLX DiT)")
 
         if not self.offload_to_cpu:
             self.model = self.model.to(device).to(self.dtype)
diff --git a/acestep/core/generation/handler/mlx_dit_init.py b/acestep/core/generation/handler/mlx_dit_init.py
index 63cedde..ff7e77a 100644
--- a/acestep/core/generation/handler/mlx_dit_init.py
+++ b/acestep/core/generation/handler/mlx_dit_init.py
@@ -1,8 +1,46 @@
 """MLX DiT initialization helpers for Apple Silicon acceleration."""
 
+import json
+import os
+
 from loguru import logger
 
 
+def audante_load_mlx_dit(handler, mlx_decoder, convert_and_load):
+    """Audante patch: fill the MLX DiT from the checkpoint's pre-quantized file (every Linear of the decoder
+    group-wise affine at `bits`, the layers in `keep8` at 8 bits), else convert it from PyTorch as upstream does."""
+    path = getattr(handler, "_audante_mlx_dit", None)
+    if not path:
+        convert_and_load(handler.model, mlx_decoder)
+        return
+    import mlx.core as mx
+    import mlx.nn as nn
+
+    with open(os.path.join(os.path.dirname(path), "mlx_dit.json")) as f:
+        spec = json.load(f)
+    weights = mx.load(path)
+    group, bits, keep8 = spec["group_size"], spec["bits"], tuple(spec.get("keep8", ()))
+    nn.quantize(mlx_decoder, group, bits, class_predicate=lambda p, m: f"{p}.scales" in weights and {
+        "group_size": group, "bits": 8 if p.startswith(keep8) else bits})
+    mlx_decoder.load_weights(list(weights.items()), strict=True)
+    mx.eval(mlx_decoder.parameters())
+    logger.info(f"[MLX-DiT] Audante patch: pre-quantized {bits}-bit DiT from {path}")
+
+
+def audante_free_torch_decoder(handler):
+    """Audante patch: with the original files and ACESTEP_FREE_TORCH_DIT=1, empty the PyTorch decoder once the MLX
+    DiT runs (6.3 GB in float32; only the Gradio UI's lyric score and LRC tools would call it). Same audio."""
+    if getattr(handler, "_audante_mlx_dit", None) or os.environ.get("ACESTEP_FREE_TORCH_DIT") != "1":
+        return
+    import torch
+    from acestep.core.generation.handler.init_service_loader import audante_empty_decoder
+
+    audante_empty_decoder(handler.model)
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+    logger.info("[MLX-DiT] Audante patch: PyTorch decoder freed")
+
+
 class MlxDitInitMixin:
     """Initialize native MLX DiT decoder state used by generation runtime."""
 
@@ -26,17 +64,20 @@ class MlxDitInitMixin:
             from acestep.models.mlx.dit_convert import convert_and_load
 
             mlx_decoder = MLXDiTDecoder.from_config(self.config)
-            convert_and_load(self.model, mlx_decoder)
+            audante_load_mlx_dit(self, mlx_decoder, convert_and_load)  # Audante patch
             mlx_decoder.materialize_static_buffers()
             self.mlx_decoder = mlx_decoder
             self.use_mlx_dit = True
             self.mlx_dit_compiled = compile_model
+            audante_free_torch_decoder(self)  # Audante patch
             logger.info(
                 "[MLX-DiT] Native MLX DiT decoder initialized successfully "
                 f"(mx.compile={compile_model})."
             )
             return True
         except Exception as exc:  # noqa: BLE001
+            if getattr(self, "_audante_mlx_dit", None):  # Audante patch: no PyTorch decoder to fall back to
+                raise
             logger.warning(f"[MLX-DiT] Failed to initialize MLX decoder (non-fatal): {exc}")
             self.mlx_decoder = None
             self.use_mlx_dit = False
diff --git a/acestep/llm_inference.py b/acestep/llm_inference.py
index 3690af5..9e3ae07 100644
--- a/acestep/llm_inference.py
+++ b/acestep/llm_inference.py
@@ -1352,6 +1352,11 @@ class LLMHandler:
         phase1_time = 0.0
         phase2_time = 0.0
 
+        # Audante patch: upstream seeds the MLX LM only in batch mode, so on a Mac a fixed seed gives other audio codes
+        # each time (MLX seeds from the clock). With ACESTEP_MLX_LM_SEED=1 a single song is seeded too.
+        self._audante_seed = (int(seeds[0]) if not is_batch and seeds and os.environ.get("ACESTEP_MLX_LM_SEED") == "1"
+                              else None)
+
         # Handle seeds for batch mode
         if is_batch:
             if seeds is None:
@@ -2942,14 +2947,17 @@ class LLMHandler:
         """
         try:
             import mlx.core as mx
-            from mlx_lm.utils import load as mlx_load
+            from pathlib import Path as _Path
+            from mlx_lm.utils import load_model as mlx_load_model
 
             logger.info(f"Loading MLX model from {model_path}")
             start_time = time.time()
 
             # Try standard mlx-lm load first
+            # Audante patch: load_model, not load: the model only (initialize() has loaded the tokenizer already,
+            # load would read it a second time) and never a download
             try:
-                self._mlx_model, _ = mlx_load(model_path)
+                self._mlx_model, _ = mlx_load_model(_Path(model_path))
             except Exception as first_err:
                 # The ACE-Step 5Hz LM checkpoints store safetensors keys without
                 # the "model." prefix (e.g. "layers.0.xxx" instead of
@@ -4069,6 +4077,9 @@ class LLMHandler:
 
         # Single mode
         formatted_prompt = formatted_prompt_list[0]
+        if getattr(self, "_audante_seed", None) is not None:  # Audante patch: metadata with the seed, codes with seed + 1
+            mx.random.seed(self._audante_seed + (1 if generation_phase == "codes" else 0))
+            logger.info(f"Audante patch: MLX LM seeded ({generation_phase})")
         return self._run_mlx_single(
             formatted_prompt=formatted_prompt,
             temperature=temperature,
diff --git a/acestep/model_downloader.py b/acestep/model_downloader.py
index 30d9f9a..51778d5 100644
--- a/acestep/model_downloader.py
+++ b/acestep/model_downloader.py
@@ -232,6 +232,11 @@ def _smart_download(
     Returns:
         Tuple of (success, message)
     """
+    if os.environ.get("ACESTEP_NO_DOWNLOAD") == "1":  # Audante patch: models come only with Audante's checksums
+        error_msg = f"{repo_id} is not installed ({local_dir}), and ACESTEP_NO_DOWNLOAD=1 forbids downloading it"
+        logger.error(f"[Model Download] {error_msg}")
+        return False, error_msg
+
     # Ensure directory exists
     local_dir.mkdir(parents=True, exist_ok=True)
 
diff --git a/acestep/models/mlx/__init__.py b/acestep/models/mlx/__init__.py
index 8135ab1..dae4006 100644
--- a/acestep/models/mlx/__init__.py
+++ b/acestep/models/mlx/__init__.py
@@ -17,11 +17,23 @@ def is_mlx_available() -> bool:
         # Verify we can actually create arrays (Metal backend works)
         _ = mx.array([1.0])
         mx.eval(_)
+        _audante_cache_limit(mx)
         return True
     except Exception:
         return False
 
 
+def _audante_cache_limit(mx) -> None:
+    """Audante patch: MLX keeps freed buffers for reuse and lets that cache grow towards the whole memory; cap it
+    at ACESTEP_MLX_CACHE_GB (same audio, a few GB less at the peak)."""
+    import os
+
+    gb = os.environ.get("ACESTEP_MLX_CACHE_GB", "").strip()
+    if gb:
+        mx.set_cache_limit(int(float(gb) * 2**30))
+        print(f"Audante patch: MLX cache limit {gb} GB", flush=True)
+
+
 _MLX_AVAILABLE = None