From 1386580cb6da269488ae1312354c2ce1f0ba9010 Mon Sep 17 00:00:00 2001 From: mm65x Date: Fri, 20 Mar 2026 10:14:29 +0000 Subject: [PATCH 1/2] optimize ace_step: remove torch runtime dependency + unified conversion pipeline --- mlx_audio/tts/models/ace_step/ace_step.py | 12 +-- mlx_audio/tts/models/ace_step/convert.py | 109 ++++++++++++++++++++++ 2 files changed, 115 insertions(+), 6 deletions(-) create mode 100644 mlx_audio/tts/models/ace_step/convert.py diff --git a/mlx_audio/tts/models/ace_step/ace_step.py b/mlx_audio/tts/models/ace_step/ace_step.py index bfa644f86..2792400b1 100644 --- a/mlx_audio/tts/models/ace_step/ace_step.py +++ b/mlx_audio/tts/models/ace_step/ace_step.py @@ -163,15 +163,15 @@ def post_load_hook(cls, model: "Model", model_path: Path) -> "Model": # Load silence latent - try turbo subdirectory first, then root silence_path = model_path / "acestep-v15-turbo" / "silence_latent.pt" - if not silence_path.exists(): + if not Path(str(silence_path).replace(".pt", ".npy")).exists(): silence_path = model_path / "silence_latent.pt" - if silence_path.exists(): - import torch + if Path(str(silence_path).replace(".pt", ".npy")).exists(): + import numpy as np - silence_pt = torch.load(silence_path, map_location="cpu", weights_only=True) - silence_pt = silence_pt.transpose(1, 2) # [1, 64, T] -> [1, T, 64] - model.silence_latent = mx.array(silence_pt.numpy()) + silence_pt = np.load(str(silence_path).replace(".pt", ".npy")) + silence_pt = silence_pt.transpose(0, 2, 1) # [1, 64, T] -> [1, T, 64] + model.silence_latent = mx.array(silence_pt) else: model.silence_latent = mx.zeros((1, 3000, 64)) diff --git a/mlx_audio/tts/models/ace_step/convert.py b/mlx_audio/tts/models/ace_step/convert.py new file mode 100644 index 000000000..f7e36822f --- /dev/null +++ b/mlx_audio/tts/models/ace_step/convert.py @@ -0,0 +1,109 @@ +import sys +import os +import torch +import transformers +import dataclasses + +def convert_ace_step(model_repo, output_dir): + import mlx.core as mx + import safetensors.torch + import json + from mlx_audio.tts.models.ace_step.config import ModelConfig + from diffusers.models import AutoencoderOobleck + + print(f"Downloading {model_repo}...") + from huggingface_hub import snapshot_download + local_dir = snapshot_download(model_repo) + # The models are nested in a subfolder for ACE-Step + turbo_dir = os.path.join(local_dir, "acestep-v15-turbo") + + print("Loading raw state dict to bypass PyTorch bugs...") + state_dict = safetensors.torch.load_file(os.path.join(turbo_dir, "model.safetensors")) + + with open(os.path.join(turbo_dir, "config.json")) as f: + config_dict = json.load(f) + config = ModelConfig.from_dict(config_dict) + + weights = [] + # Extract Decoder manually + for key, value in state_dict.items(): + if not key.startswith("decoder.") and not key.startswith("encoder."): continue + np_val = value.cpu().float().numpy() + + if key.startswith("decoder."): + new_key = key.replace("decoder.", "dit.") + if "proj_in.1." in new_key: + new_key = new_key.replace("proj_in.1.", "proj_in.") + if new_key.endswith(".weight"): + np_val = np_val.swapaxes(1, 2) + elif "proj_out.1." in new_key: + new_key = new_key.replace("proj_out.1.", "proj_out.") + if new_key.endswith(".weight"): + np_val = np_val.transpose(1, 2, 0) + elif "rotary_emb" in new_key: + continue + weights.append((new_key, mx.array(np_val))) + + elif key.startswith("encoder."): + new_key = key + if "rotary_emb" in new_key: + continue + weights.append((new_key, mx.array(np_val))) + + out_path = output_dir + mx.save_safetensors(os.path.join(out_path, "model.safetensors"), dict(weights)) + + print("Loading VAE...") + pt_vae = AutoencoderOobleck.from_pretrained(f"{local_dir}/vae") + vae_weights = [] + + import numpy as np + def _fuse_weight_norm(weight_g, weight_v, eps=1e-9): + v_flat = weight_v.reshape(weight_v.shape[0], -1) + norm = np.linalg.norm(v_flat, axis=1).reshape(weight_g.shape) + return weight_v * (weight_g / np.maximum(norm, eps)) + + vae_state = pt_vae.state_dict() + processed = set() + all_keys = sorted(vae_state.keys()) + for key in all_keys: + if key in processed: continue + if key.endswith(".weight_g"): + base = key[: -len(".weight_g")] + v_key = base + ".weight_v" + g = vae_state[key].detach().cpu().float().numpy() + v = vae_state[v_key].detach().cpu().float().numpy() + w = _fuse_weight_norm(g, v) + if "conv_t1" in base: w = w.transpose(1, 2, 0) + else: w = w.swapaxes(1, 2) + vae_weights.append((base + ".weight", mx.array(w))) + processed.add(key) + processed.add(v_key) + continue + if key.endswith(".weight_v"): continue + val = vae_state[key].detach().cpu().float().numpy() + if key.endswith(".alpha") or key.endswith(".beta"): + val = val.squeeze() + if "conv" in key and key.endswith(".weight"): + if "conv_t1" in key: + val = val.transpose(1, 2, 0) + else: + val = val.swapaxes(1, 2) + vae_weights.append((key, mx.array(val))) + processed.add(key) + + mx.save_safetensors(os.path.join(out_path, "vae.safetensors"), dict(vae_weights)) + + with open(os.path.join(out_path, "config.json"), "w") as f: + # Save raw dict since we are bypassing to_dict + json.dump(config_dict, f, indent=4) + + silence_path = os.path.join(turbo_dir, "silence_latent.pt") + if os.path.exists(silence_path): + pt_silence = torch.load(silence_path, map_location="cpu", weights_only=True) + np.save(os.path.join(out_path, "silence_latent.npy"), pt_silence.numpy()) + + print("Success!") + +print("Starting conversion...") +convert_ace_step("ACE-Step/Ace-Step1.5", "/tmp/ace_step-mlx-converted") From f861698ee9c38a2057544f1e8759303e0d82f513 Mon Sep 17 00:00:00 2001 From: mm65x Date: Sat, 21 Mar 2026 18:05:04 +0000 Subject: [PATCH 2/2] fix ace-step model conversion layout --- mlx_audio/tts/models/ace_step/convert.py | 169 ++++++++++++----------- 1 file changed, 89 insertions(+), 80 deletions(-) diff --git a/mlx_audio/tts/models/ace_step/convert.py b/mlx_audio/tts/models/ace_step/convert.py index f7e36822f..e338067f0 100644 --- a/mlx_audio/tts/models/ace_step/convert.py +++ b/mlx_audio/tts/models/ace_step/convert.py @@ -1,35 +1,40 @@ -import sys +import argparse +import json import os +import shutil +from pathlib import Path + +import numpy as np import torch -import transformers -import dataclasses -def convert_ace_step(model_repo, output_dir): + +def convert_ace_step(model_repo: str, output_dir: str, local_files_only: bool = False): import mlx.core as mx import safetensors.torch - import json - from mlx_audio.tts.models.ace_step.config import ModelConfig from diffusers.models import AutoencoderOobleck - - print(f"Downloading {model_repo}...") from huggingface_hub import snapshot_download - local_dir = snapshot_download(model_repo) - # The models are nested in a subfolder for ACE-Step - turbo_dir = os.path.join(local_dir, "acestep-v15-turbo") - + + print(f"Downloading {model_repo}...") + local_dir = snapshot_download(model_repo, local_files_only=local_files_only) + turbo_dir = Path(local_dir) / "acestep-v15-turbo" + vae_dir = Path(local_dir) / "vae" + text_dir = Path(local_dir) / "Qwen3-Embedding-0.6B" + out_dir = Path(output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + print("Loading raw state dict to bypass PyTorch bugs...") - state_dict = safetensors.torch.load_file(os.path.join(turbo_dir, "model.safetensors")) - - with open(os.path.join(turbo_dir, "config.json")) as f: + state_dict = safetensors.torch.load_file(str(turbo_dir / "model.safetensors")) + + with open(turbo_dir / "config.json") as f: config_dict = json.load(f) - config = ModelConfig.from_dict(config_dict) - - weights = [] - # Extract Decoder manually + + weights = {} for key, value in state_dict.items(): - if not key.startswith("decoder.") and not key.startswith("encoder."): continue + if not key.startswith("decoder.") and not key.startswith("encoder."): + continue + np_val = value.cpu().float().numpy() - + if key.startswith("decoder."): new_key = key.replace("decoder.", "dit.") if "proj_in.1." in new_key: @@ -42,68 +47,72 @@ def convert_ace_step(model_repo, output_dir): np_val = np_val.transpose(1, 2, 0) elif "rotary_emb" in new_key: continue - weights.append((new_key, mx.array(np_val))) - - elif key.startswith("encoder."): - new_key = key - if "rotary_emb" in new_key: + weights[new_key] = mx.array(np_val) + else: + if "rotary_emb" in key: continue - weights.append((new_key, mx.array(np_val))) - - out_path = output_dir - mx.save_safetensors(os.path.join(out_path, "model.safetensors"), dict(weights)) - + weights[key] = mx.array(np_val) + + mx.save_safetensors(str(out_dir / "model.safetensors"), weights) + print("Loading VAE...") - pt_vae = AutoencoderOobleck.from_pretrained(f"{local_dir}/vae") - vae_weights = [] - - import numpy as np - def _fuse_weight_norm(weight_g, weight_v, eps=1e-9): - v_flat = weight_v.reshape(weight_v.shape[0], -1) - norm = np.linalg.norm(v_flat, axis=1).reshape(weight_g.shape) - return weight_v * (weight_g / np.maximum(norm, eps)) - - vae_state = pt_vae.state_dict() - processed = set() - all_keys = sorted(vae_state.keys()) - for key in all_keys: - if key in processed: continue - if key.endswith(".weight_g"): - base = key[: -len(".weight_g")] - v_key = base + ".weight_v" - g = vae_state[key].detach().cpu().float().numpy() - v = vae_state[v_key].detach().cpu().float().numpy() - w = _fuse_weight_norm(g, v) - if "conv_t1" in base: w = w.transpose(1, 2, 0) - else: w = w.swapaxes(1, 2) - vae_weights.append((base + ".weight", mx.array(w))) - processed.add(key) - processed.add(v_key) - continue - if key.endswith(".weight_v"): continue - val = vae_state[key].detach().cpu().float().numpy() - if key.endswith(".alpha") or key.endswith(".beta"): - val = val.squeeze() - if "conv" in key and key.endswith(".weight"): - if "conv_t1" in key: - val = val.transpose(1, 2, 0) - else: - val = val.swapaxes(1, 2) - vae_weights.append((key, mx.array(val))) - processed.add(key) - - mx.save_safetensors(os.path.join(out_path, "vae.safetensors"), dict(vae_weights)) - - with open(os.path.join(out_path, "config.json"), "w") as f: - # Save raw dict since we are bypassing to_dict + pt_vae = AutoencoderOobleck.from_pretrained(str(vae_dir)) + vae_weights = {} + + for key, value in pt_vae.state_dict().items(): + np_val = value.detach().cpu().float().numpy() + vae_weights[key] = mx.array(np_val) + + vae_out_dir = out_dir / "vae" + vae_out_dir.mkdir(exist_ok=True) + mx.save_safetensors( + str(vae_out_dir / "diffusion_pytorch_model.safetensors"), vae_weights + ) + shutil.copy2(vae_dir / "config.json", vae_out_dir / "config.json") + + with open(out_dir / "config.json", "w") as f: json.dump(config_dict, f, indent=4) - - silence_path = os.path.join(turbo_dir, "silence_latent.pt") - if os.path.exists(silence_path): + + silence_path = turbo_dir / "silence_latent.pt" + if silence_path.exists(): pt_silence = torch.load(silence_path, map_location="cpu", weights_only=True) - np.save(os.path.join(out_path, "silence_latent.npy"), pt_silence.numpy()) + np.save(out_dir / "silence_latent.npy", pt_silence.numpy()) + + print("Copying Qwen3 text encoder...") + if text_dir.exists(): + text_out_dir = out_dir / "Qwen3-Embedding-0.6B" + text_out_dir.mkdir(exist_ok=True) + + text_state = safetensors.torch.load_file(str(text_dir / "model.safetensors")) + text_weights = {} + for key, value in text_state.items(): + clean_key = key[6:] if key.startswith("model.") else key + text_weights[clean_key] = mx.array(value.cpu().float().numpy()) + + mx.save_safetensors(str(text_out_dir / "model.safetensors"), text_weights) + + for fname in [ + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + ]: + src = text_dir / fname + if src.exists(): + shutil.copy2(src, text_out_dir / fname) + + print(f"Success! Wrote converted files to {out_dir}") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("model_repo", nargs="?", default="ACE-Step/Ace-Step1.5") + parser.add_argument("output_dir", nargs="?", default="/tmp/ace_step-mlx-converted") + parser.add_argument("--local-files-only", action="store_true") + args = parser.parse_args() + convert_ace_step(args.model_repo, args.output_dir, args.local_files_only) - print("Success!") -print("Starting conversion...") -convert_ace_step("ACE-Step/Ace-Step1.5", "/tmp/ace_step-mlx-converted") +if __name__ == "__main__": + main()