mirror of
https://github.com/Nighthawk42/rvc_pth2onnx.git
synced 2026-08-30 04:30:20 +00:00
457 lines
20 KiB
Python
457 lines
20 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
Exports RVC (Retrieval-based Voice Conversion) models from PyTorch (.pth)
|
|
to ONNX format, simplifying the model optionally.
|
|
Prompts the user for the input model path and optionally verifies the output.
|
|
Logs output to console and 'conversion.log'. Keeps window open on exit.
|
|
Includes workaround for onnxsim dynamic shape issue and weight_norm FutureWarning.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
import argparse
|
|
import time
|
|
import warnings # Added for warning suppression
|
|
|
|
# Third-party libraries
|
|
import onnx
|
|
import onnxsim
|
|
import torch
|
|
|
|
# Optional verification library
|
|
try:
|
|
import onnxruntime as ort
|
|
ORT_AVAILABLE = True
|
|
except ImportError:
|
|
ORT_AVAILABLE = False
|
|
|
|
|
|
# --- Configure Logging (File and Console) ---
|
|
log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
|
root_logger = logging.getLogger()
|
|
root_logger.setLevel(logging.INFO)
|
|
|
|
# Console Handler
|
|
console_handler = logging.StreamHandler(sys.stdout)
|
|
console_handler.setFormatter(log_formatter)
|
|
root_logger.addHandler(console_handler)
|
|
|
|
# File Handler (log to conversion.log in script's directory)
|
|
try:
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
log_file_path = os.path.join(script_dir, "conversion.log")
|
|
# 'a' appends to the log, 'w' overwrites
|
|
file_handler = logging.FileHandler(log_file_path, mode='a', encoding='utf-8')
|
|
file_handler.setFormatter(log_formatter)
|
|
root_logger.addHandler(file_handler)
|
|
# Use print only for the initial notification about the log file location
|
|
print(f"Logging detailed output to: {log_file_path}")
|
|
except Exception as e:
|
|
print(f"Warning: Could not configure file logging. Error: {e}", file=sys.stderr)
|
|
|
|
|
|
# --- RVC Model Import ---
|
|
# Suppress the specific FutureWarning during import and model loading if needed
|
|
# Note: Ideally, the RVC library itself should be updated.
|
|
try:
|
|
# This context manager is generally better applied specifically around the
|
|
# code that triggers the warning (model instantiation/loading).
|
|
# Placing it here might suppress warnings during import of the module itself.
|
|
# with warnings.catch_warnings():
|
|
# warnings.filterwarnings("ignore", category=FutureWarning, message=".*`torch.nn.utils.weight_norm` is deprecated.*")
|
|
from infer.lib.infer_pack.models_onnx import SynthesizerTrnMsNSFsidM
|
|
except ImportError as e:
|
|
logging.error(f"Error importing SynthesizerTrnMsNSFsidM: {e}", exc_info=True)
|
|
logging.error("Please ensure the script is run from within the RVC "
|
|
"project structure or that the RVC library's root "
|
|
"directory is in your PYTHONPATH.")
|
|
input("Import error occurred. Press Enter to exit...")
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
logging.error(f"An unexpected error occurred during import: {e}", exc_info=True)
|
|
input("Import error occurred. Press Enter to exit...")
|
|
sys.exit(1)
|
|
|
|
|
|
class RvcOnnxExporter:
|
|
"""Handles the export of an RVC model to ONNX format."""
|
|
|
|
INPUT_NAMES = ["phone", "phone_lengths", "pitch", "pitchf", "ds", "rnd"]
|
|
OUTPUT_NAMES = ["audio"]
|
|
|
|
def __init__(self, opset_version: int = 18, simplify: bool = True):
|
|
"""
|
|
Initializes the exporter.
|
|
|
|
Args:
|
|
opset_version (int): The ONNX opset version for export.
|
|
simplify (bool): Whether to simplify the exported model using onnxsim.
|
|
"""
|
|
self.opset_version = opset_version
|
|
self.simplify = simplify
|
|
self.config = None
|
|
self.model_version = "v1"
|
|
self.vec_channels = 256
|
|
self.n_spk = 1
|
|
self.posterior_channels = 192 # Default value used in error case before
|
|
self.dummy_input_shapes = {} # Store shapes for simplify step
|
|
|
|
def _load_checkpoint(self, model_path: str) -> dict | None:
|
|
"""Loads the PyTorch checkpoint file."""
|
|
if not os.path.exists(model_path):
|
|
logging.error(f"Model file not found: {model_path}")
|
|
return None
|
|
try:
|
|
logging.info(f"Loading checkpoint: {model_path}")
|
|
checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
|
|
|
|
required_keys = ['config', 'weight']
|
|
if not all(key in checkpoint for key in required_keys):
|
|
missing = [k for k in required_keys if k not in checkpoint]
|
|
logging.error(f"Checkpoint missing required keys: {missing}. "
|
|
f"Found keys: {list(checkpoint.keys())}")
|
|
return None
|
|
if "emb_g.weight" not in checkpoint["weight"]:
|
|
logging.error("Speaker embedding 'emb_g.weight' not found in "
|
|
"checkpoint weights.")
|
|
return None
|
|
logging.info("Checkpoint loaded successfully.")
|
|
return checkpoint
|
|
except Exception as e:
|
|
logging.error(f"Failed to load checkpoint '{model_path}': {e}",
|
|
exc_info=True)
|
|
return None
|
|
|
|
def _prepare_config(self, checkpoint: dict) -> bool:
|
|
"""Determines model version, channels, and speaker count from config."""
|
|
try:
|
|
logging.info("Preparing model configuration...")
|
|
self.config = list(checkpoint["config"])
|
|
self.model_version = checkpoint.get("version", "v1")
|
|
self.vec_channels = 768 if self.model_version == "v2" else 256
|
|
logging.info(f"Detected model version: {self.model_version}, "
|
|
f"Feature channels: {self.vec_channels}")
|
|
|
|
# Use config[2] for posterior_channels if valid, otherwise log error
|
|
if len(self.config) > 2 and isinstance(self.config[2], int) and self.config[2] > 0:
|
|
self.posterior_channels = self.config[2]
|
|
logging.info(f"Using posterior encoder channels (config[2]): "
|
|
f"{self.posterior_channels} for 'rnd' input.")
|
|
else:
|
|
# Log the error but might proceed if the old default was somehow correct (unlikely)
|
|
# Consider returning False here for stricter validation
|
|
config2_val = self.config[2] if len(self.config) > 2 else 'N/A'
|
|
logging.error("Could not determine valid posterior encoder channels "
|
|
f"from config[2] (value: {config2_val}). "
|
|
f"Using fallback: {self.posterior_channels}")
|
|
# return False # Uncomment for stricter validation
|
|
|
|
n_spk_inferred = checkpoint["weight"]["emb_g.weight"].shape[0]
|
|
if len(self.config) < 3:
|
|
logging.error(f"Config list has fewer than 3 elements "
|
|
f"({len(self.config)}), cannot access speaker "
|
|
f"count at index -3.")
|
|
return False
|
|
|
|
n_spk_config = self.config[-3]
|
|
if n_spk_config != n_spk_inferred:
|
|
logging.warning(f"Mismatch: Config speaker count ({n_spk_config})"
|
|
f" != Inferred count ({n_spk_inferred}). "
|
|
f"Updating config.")
|
|
self.config[-3] = n_spk_inferred
|
|
|
|
self.n_spk = n_spk_inferred
|
|
logging.info(f"Using speaker count (n_spk): {self.n_spk}")
|
|
self.config = tuple(self.config)
|
|
logging.info("Model configuration prepared successfully.")
|
|
return True
|
|
|
|
except (IndexError, KeyError, TypeError, ValueError) as e:
|
|
logging.error(f"Failed to prepare config from checkpoint: {e}",
|
|
exc_info=True)
|
|
return False
|
|
|
|
def _create_dummy_inputs(self, seq_len: int = 200) -> tuple | None:
|
|
"""Creates dummy input tensors for ONNX tracing and stores their shapes."""
|
|
try:
|
|
logging.info(f"Creating dummy inputs (sequence length: {seq_len})...")
|
|
# Create tensors
|
|
dummy_phone = torch.rand(1, seq_len, self.vec_channels, dtype=torch.float32)
|
|
dummy_phone_lengths = torch.tensor([seq_len], dtype=torch.long)
|
|
dummy_pitch = torch.randint(low=5, high=255, size=(1, seq_len), dtype=torch.long)
|
|
dummy_pitchf = torch.rand(1, seq_len, dtype=torch.float32)
|
|
dummy_ds = torch.tensor([0], dtype=torch.long) # Speaker ID
|
|
dummy_rnd = torch.rand(1, self.posterior_channels, seq_len, dtype=torch.float32)
|
|
|
|
inputs_tuple = (dummy_phone, dummy_phone_lengths, dummy_pitch,
|
|
dummy_pitchf, dummy_ds, dummy_rnd)
|
|
|
|
# Store shapes mapped to names for use in simplification step
|
|
self.dummy_input_shapes = {
|
|
name: list(tensor.shape)
|
|
for name, tensor in zip(self.INPUT_NAMES, inputs_tuple)
|
|
}
|
|
logging.info(f"Dummy input shapes: {self.dummy_input_shapes}")
|
|
logging.info("Dummy inputs created successfully.")
|
|
return inputs_tuple
|
|
except Exception as e:
|
|
logging.error(f"Failed to create dummy inputs: {e}", exc_info=True)
|
|
return None
|
|
|
|
def _instantiate_and_load_model(self, checkpoint: dict) -> SynthesizerTrnMsNSFsidM | None:
|
|
"""Instantiates the RVC model and loads weights, suppressing specific warnings."""
|
|
try:
|
|
logging.info("Instantiating model SynthesizerTrnMsNSFsidM...")
|
|
# Suppress the specific weight_norm FutureWarning during model init and loading
|
|
# Note: The ideal solution is to update the RVC library code itself.
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", category=FutureWarning, message=".*`torch.nn.utils.weight_norm` is deprecated.*")
|
|
net_g = SynthesizerTrnMsNSFsidM(
|
|
*self.config,
|
|
is_half=False,
|
|
version=self.model_version
|
|
)
|
|
logging.info("Loading state dictionary into the model...")
|
|
net_g.load_state_dict(checkpoint["weight"], strict=False)
|
|
|
|
net_g.eval()
|
|
logging.info("Model instantiated and weights loaded successfully.")
|
|
return net_g
|
|
except Exception as e:
|
|
logging.error(f"Failed to instantiate or load model: {e}",
|
|
exc_info=True)
|
|
return None
|
|
|
|
def _simplify_model(self, onnx_path: str) -> bool:
|
|
"""Simplifies the exported ONNX model using onnxsim, providing input shapes."""
|
|
try:
|
|
logging.info("Simplifying ONNX model using onnxsim...")
|
|
if not os.path.exists(onnx_path):
|
|
logging.error(f"ONNX file not found at {onnx_path} before simplification.")
|
|
return False
|
|
if not self.dummy_input_shapes:
|
|
logging.error("Cannot simplify model: Dummy input shapes were not generated.")
|
|
return False # Treat as failure if shapes are missing
|
|
|
|
logging.info(f"Providing input shapes to onnxsim: {self.dummy_input_shapes}")
|
|
onnx_model = onnx.load(onnx_path)
|
|
# Provide the input shapes derived from dummy data
|
|
model_opt, check_ok = onnxsim.simplify(
|
|
onnx_model,
|
|
input_shapes=self.dummy_input_shapes, # Pass the shapes here
|
|
check_n=3,
|
|
perform_optimization=True
|
|
)
|
|
|
|
if check_ok:
|
|
onnx.save(model_opt, onnx_path)
|
|
logging.info("Simplified ONNX model saved successfully.")
|
|
return True
|
|
else:
|
|
logging.error("ONNX simplification check failed. Keeping the "
|
|
"original (unsimplified) model.")
|
|
return True # Treat simplification check failure as non-critical
|
|
except Exception as e:
|
|
# Log the specific error from onnxsim
|
|
logging.error(f"Error during ONNX simplification: {e}", exc_info=True)
|
|
logging.warning("Proceeding with the unsimplified ONNX model.")
|
|
return True # Simplification process failed, but export is usable
|
|
|
|
def export(self, model_path: str, exported_path: str) -> bool:
|
|
"""
|
|
Executes the full export process.
|
|
|
|
Args:
|
|
model_path (str): Path to the input PyTorch model checkpoint (.pth).
|
|
exported_path (str): Path to save the exported ONNX model.
|
|
|
|
Returns:
|
|
bool: True if base export was successful, False otherwise.
|
|
"""
|
|
checkpoint = self._load_checkpoint(model_path)
|
|
if not checkpoint: return False
|
|
|
|
if not self._prepare_config(checkpoint): return False
|
|
|
|
dummy_inputs = self._create_dummy_inputs() # Generates shapes needed later
|
|
if not dummy_inputs: return False
|
|
|
|
net_g = self._instantiate_and_load_model(checkpoint)
|
|
if not net_g: return False
|
|
|
|
dynamic_axes = {
|
|
name: {1: "sequence_length"} for name in ["phone", "pitch", "pitchf"]
|
|
}
|
|
dynamic_axes["rnd"] = {2: "sequence_length"} # Dim 2 for rnd noise
|
|
dynamic_axes[self.OUTPUT_NAMES[0]] = {1: "audio_length"} # Output audio
|
|
|
|
try:
|
|
logging.info(f"Exporting model to ONNX format at: {exported_path} "
|
|
f"(Opset: {self.opset_version})")
|
|
torch.onnx.export(
|
|
net_g, dummy_inputs, exported_path,
|
|
input_names=self.INPUT_NAMES, output_names=self.OUTPUT_NAMES,
|
|
dynamic_axes=dynamic_axes, do_constant_folding=False, # Be cautious with folding + dynamic axes
|
|
opset_version=self.opset_version, verbose=False,
|
|
)
|
|
logging.info(f"Initial ONNX model exported successfully to: {exported_path}")
|
|
|
|
except Exception as e:
|
|
logging.error(f"ONNX export failed: {e}", exc_info=True)
|
|
if os.path.exists(exported_path):
|
|
try: os.remove(exported_path)
|
|
except OSError as oe: logging.error(f"Could not remove intermediate ONNX file {exported_path}: {oe}")
|
|
return False
|
|
|
|
if self.simplify:
|
|
# _simplify_model now uses self.dummy_input_shapes generated earlier
|
|
self._simplify_model(exported_path)
|
|
|
|
logging.info("ONNX export process finished.")
|
|
return True
|
|
|
|
|
|
def get_model_path_from_user() -> str | None:
|
|
"""Prompts the user to provide a valid .pth model file path."""
|
|
while True:
|
|
# Use print for direct user interaction
|
|
print("\nPlease provide the path to the RVC model checkpoint file (.pth).")
|
|
print("You can type the full path or drag and drop the file onto "
|
|
"this window and press Enter.")
|
|
print("Type 'quit' or 'exit' to cancel.")
|
|
raw_path = input("Model Path: ").strip()
|
|
|
|
if raw_path.lower() in ['quit', 'exit']:
|
|
logging.info("Operation cancelled by user.")
|
|
return None
|
|
|
|
cleaned_path = raw_path.strip("'\"")
|
|
if not cleaned_path:
|
|
print("ERROR: No path provided. Please try again or type 'quit'.")
|
|
continue
|
|
|
|
try: absolute_path = os.path.abspath(cleaned_path)
|
|
except Exception as e: print(f"ERROR: Could not resolve path '{cleaned_path}'. Error: {e}"); continue
|
|
|
|
if not os.path.exists(absolute_path): print(f"ERROR: File not found at '{absolute_path}'. Check the path."); continue
|
|
if not os.path.isfile(absolute_path): print(f"ERROR: Path '{absolute_path}' is a directory. Need a file."); continue
|
|
if not absolute_path.lower().endswith(".pth"): print(f"ERROR: File '{os.path.basename(absolute_path)}' needs .pth extension."); continue
|
|
|
|
logging.info(f"Input model accepted: {absolute_path}")
|
|
return absolute_path
|
|
|
|
|
|
def verify_onnx_model(onnx_path: str):
|
|
"""Attempts to load the ONNX model using onnxruntime for verification."""
|
|
# Use print for separators in direct user feedback section
|
|
print("-" * 30)
|
|
logging.info(f"Attempting to verify ONNX model: {onnx_path}")
|
|
|
|
if not ORT_AVAILABLE:
|
|
logging.warning("ONNX Runtime (onnxruntime) library not found.")
|
|
logging.warning("Verification skipped. Install with: pip install onnxruntime")
|
|
print("-" * 30)
|
|
return
|
|
|
|
if not os.path.exists(onnx_path):
|
|
logging.error(f"Verification failed: ONNX file not found at {onnx_path}")
|
|
print("-" * 30)
|
|
return
|
|
|
|
try:
|
|
providers = ['CPUExecutionProvider']
|
|
ort_session = ort.InferenceSession(onnx_path, providers=providers)
|
|
inputs = ort_session.get_inputs()
|
|
outputs = ort_session.get_outputs()
|
|
logging.info(f"ONNX Runtime loaded model successfully.")
|
|
logging.info(f"Detected {len(inputs)} inputs: {[inp.name for inp in inputs]}")
|
|
logging.info(f"Detected {len(outputs)} outputs: {[out.name for out in outputs]}")
|
|
logging.info("Verification successful.")
|
|
|
|
except ort.OrtLoadError as load_error:
|
|
logging.error(f"ONNX Runtime failed to load the model: {load_error}", exc_info=True)
|
|
logging.error("Verification failed: Model loading error.")
|
|
except Exception as e:
|
|
logging.error(f"An unexpected error occurred during ONNX verification: {e}", exc_info=True)
|
|
logging.error("Verification failed.")
|
|
finally:
|
|
# Ensure separator is printed even on error during verification
|
|
print("-" * 30)
|
|
|
|
|
|
def main():
|
|
"""Main execution function."""
|
|
# Use print for initial user-facing messages
|
|
print("-" * 30)
|
|
print("RVC ONNX Exporter")
|
|
print("-" * 30)
|
|
logging.info("Exporter script started.") # Log that script has begun
|
|
|
|
checkpoint_path = get_model_path_from_user()
|
|
if not checkpoint_path:
|
|
return 0 # User quit
|
|
|
|
input_dir = os.path.dirname(checkpoint_path)
|
|
base_name = os.path.splitext(os.path.basename(checkpoint_path))[0]
|
|
onnx_output_path = os.path.join(input_dir, f"{base_name}.onnx")
|
|
|
|
# Log essential paths
|
|
print("-" * 30)
|
|
logging.info(f"Input PyTorch Model: {checkpoint_path}")
|
|
logging.info(f"Output ONNX Model: {onnx_output_path}")
|
|
print("-" * 30)
|
|
|
|
# --- Initialize and Run Exporter ---
|
|
exporter = RvcOnnxExporter(opset_version=18, simplify=True)
|
|
export_successful = exporter.export(checkpoint_path, onnx_output_path)
|
|
|
|
# --- Final Status & Optional Verification ---
|
|
print("-" * 30) # Use print for visual separation for user
|
|
if export_successful:
|
|
logging.info(f"Successfully exported model to: {onnx_output_path}")
|
|
|
|
while True:
|
|
# Use print for direct user interaction prompt
|
|
verify_choice = input("Verify the exported ONNX model using ONNX Runtime? (y/n): ").lower().strip()
|
|
if verify_choice in ['y', 'yes']:
|
|
verify_onnx_model(onnx_output_path)
|
|
break
|
|
elif verify_choice in ['n', 'no']:
|
|
logging.info("Skipping ONNX model verification.")
|
|
print("-" * 30) # Print separator after skipping
|
|
break
|
|
else:
|
|
print("Invalid input. Please enter 'y' or 'n'.") # Use print for error feedback
|
|
else:
|
|
logging.error("Model export failed. See log messages above for details.")
|
|
|
|
logging.info("Exporter script finished.")
|
|
return 0 if export_successful else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit_code = 1 # Default to error
|
|
try:
|
|
# Start logging session info
|
|
logging.info(f"--- Starting new conversion run ---")
|
|
logging.info(f"Python version: {sys.version}")
|
|
logging.info(f"Torch version: {torch.__version__}")
|
|
logging.info(f"ONNX version: {onnx.__version__}")
|
|
logging.info(f"ONNXSim version: {onnxsim.__version__}")
|
|
logging.info(f"ONNX Runtime Available: {ORT_AVAILABLE}" + (f" (Version: {ort.__version__})" if ORT_AVAILABLE else ""))
|
|
|
|
exit_code = main()
|
|
|
|
except Exception as e:
|
|
logging.error(f"An unexpected critical error occurred in main execution: {e}", exc_info=True)
|
|
# Ensure the error is logged before the exit prompt
|
|
|
|
finally:
|
|
# This block ensures the prompt appears even if main() raises an exception
|
|
logging.info(f"--- Ending conversion run (Exit Code: {exit_code}) ---")
|
|
print("\nScript execution complete.") # Use print for final user message
|
|
input("Press Enter to exit...") # Keep window open
|
|
sys.exit(exit_code) |