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