diodiogod
51e97f70f6
Fix RVC HuBERT subdirectory double-nesting issue Technical fix: - Add check to detect if .pt file is already in correct subdirectory - Prevent creating nested subdirectories (e.g., chinese_hubert_base/chinese_hubert_base/) - Use existing subdirectory path when model is already correctly placed
281 lines
14 KiB
Python
281 lines
14 KiB
Python
|
|
import hashlib
|
|
import torch.nn.functional as F
|
|
import librosa
|
|
import torch
|
|
import os
|
|
|
|
from .infer_pack.loaders import HubertModelWithFinalProj
|
|
|
|
def get_hash(model_path):
|
|
try:
|
|
with open(model_path, 'rb') as f:
|
|
f.seek(- 10000 * 1024, 2)
|
|
model_hash = hashlib.md5(f.read()).hexdigest()
|
|
except:
|
|
model_hash = hashlib.md5(open(model_path, 'rb').read()).hexdigest()
|
|
|
|
return model_hash
|
|
|
|
def load_hubert(model_path: str, config):
|
|
try:
|
|
print(f"🔧 Loading HuBERT model: {model_path}")
|
|
if model_path.endswith(".safetensors"):
|
|
# Check if this safetensors file needs migration to subdirectory
|
|
from engines.rvc.hubert_models import HUBERT_MODELS
|
|
|
|
model_filename = os.path.basename(model_path)
|
|
model_dir_path = os.path.dirname(model_path)
|
|
|
|
# Check if this is a transformers model that should be in a subdirectory
|
|
for key, info in HUBERT_MODELS.items():
|
|
if info.get('is_transformers') and info.get('filename') == model_filename and info.get('model_dir'):
|
|
# Check if currently in flat directory (not already in subdirectory)
|
|
if os.path.basename(model_dir_path) != info['model_dir']:
|
|
target_subdir = os.path.join(model_dir_path, info['model_dir'])
|
|
target_path = os.path.join(target_subdir, model_filename)
|
|
|
|
# Migrate to subdirectory
|
|
print(f"🔄 Migrating {model_filename} to subdirectory: {info['model_dir']}")
|
|
os.makedirs(target_subdir, exist_ok=True)
|
|
|
|
try:
|
|
import shutil
|
|
shutil.move(model_path, target_path)
|
|
model_path = target_path # Update path to new location
|
|
print(f"✅ Migrated to: {target_path}")
|
|
|
|
# Download config if missing
|
|
from engines.rvc.hubert_downloader import _download_transformers_config
|
|
config_path = os.path.join(target_subdir, 'config.json')
|
|
if not os.path.exists(config_path) and info.get('repo_id'):
|
|
print(f"📥 Downloading config.json for {key}...")
|
|
_download_transformers_config(info['repo_id'], target_subdir)
|
|
except Exception as migrate_error:
|
|
print(f"⚠️ Migration failed: {migrate_error}")
|
|
break
|
|
|
|
try:
|
|
model = HubertModelWithFinalProj.from_safetensors(model_path, device=config.device)
|
|
print(f"✅ HuBERT safetensors model loaded successfully")
|
|
return model
|
|
except Exception as e:
|
|
print(f"❌ HuBERT safetensors loading error: {e}")
|
|
print(f"🔧 Error type: {type(e).__name__}")
|
|
raise
|
|
else:
|
|
# Try loading .pt file directly first (safer approach)
|
|
print(f"🔧 Attempting direct .pt loading: {model_path}")
|
|
try:
|
|
import torch
|
|
checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
|
|
|
|
# Try to create HuBERT model directly from .pt checkpoint
|
|
# This will work if the checkpoint contains the full model with config
|
|
if hasattr(checkpoint, 'eval'):
|
|
# The checkpoint itself is a model
|
|
model = checkpoint.to(config.device)
|
|
model.eval()
|
|
print(f"✅ Direct .pt model loading successful")
|
|
return model
|
|
|
|
except Exception as direct_error:
|
|
print(f"⚠️ Direct .pt loading failed: {direct_error}")
|
|
print(f"🔄 Falling back to safetensors conversion")
|
|
|
|
# Fallback: Convert .pt file to .safetensors format for compatibility
|
|
print(f"🔄 Converting .pt Hubert model to safetensors format: {model_path}")
|
|
|
|
# Determine if this is a transformers model that needs a subdirectory
|
|
# Import here to avoid circular dependency
|
|
from engines.rvc.hubert_models import HUBERT_MODELS
|
|
|
|
# Try to identify which model this is by filename
|
|
model_filename = os.path.basename(model_path)
|
|
target_subdir = None
|
|
model_key = None
|
|
|
|
for key, info in HUBERT_MODELS.items():
|
|
# Check if this filename matches a transformers model
|
|
if info.get('is_transformers') and info.get('filename'):
|
|
# Match both .pt and .safetensors versions
|
|
expected_pt = info['filename'].replace('.safetensors', '.pt')
|
|
expected_safetensors = info['filename']
|
|
if model_filename in [expected_pt, expected_safetensors, info['filename']]:
|
|
target_subdir = info.get('model_dir')
|
|
model_key = key
|
|
break
|
|
|
|
# Generate safetensors path - use subdirectory if this is a transformers model
|
|
if target_subdir:
|
|
base_dir = os.path.dirname(model_path)
|
|
|
|
# Check if we're ALREADY in the correct subdirectory (avoid double-nesting)
|
|
if os.path.basename(base_dir) == target_subdir:
|
|
# Already in subdirectory, use current location
|
|
subdir_path = base_dir
|
|
safetensors_filename = model_filename.replace(".pt", ".safetensors")
|
|
safetensors_path = os.path.join(subdir_path, safetensors_filename)
|
|
print(f"🔧 Transformers model ({model_key}) already in subdirectory: {target_subdir}")
|
|
else:
|
|
# Need to create/use subdirectory
|
|
subdir_path = os.path.join(base_dir, target_subdir)
|
|
os.makedirs(subdir_path, exist_ok=True)
|
|
safetensors_filename = model_filename.replace(".pt", ".safetensors")
|
|
safetensors_path = os.path.join(subdir_path, safetensors_filename)
|
|
print(f"🔧 Transformers model detected ({model_key}), saving to subdirectory: {target_subdir}")
|
|
|
|
# Check if safetensors already exists in flat directory and needs migration
|
|
flat_safetensors = model_path.replace(".pt", ".safetensors")
|
|
if os.path.exists(flat_safetensors) and not os.path.exists(safetensors_path):
|
|
print(f"🔄 Migrating existing safetensors to subdirectory...")
|
|
try:
|
|
import shutil
|
|
shutil.move(flat_safetensors, safetensors_path)
|
|
print(f"✅ Migrated: {os.path.basename(flat_safetensors)} -> {target_subdir}/")
|
|
|
|
# Also download config.json if missing
|
|
try:
|
|
from engines.rvc.hubert_downloader import _download_transformers_config
|
|
info = HUBERT_MODELS.get(model_key)
|
|
if info and info.get('repo_id'):
|
|
config_dir = os.path.dirname(safetensors_path)
|
|
config_path = os.path.join(config_dir, 'config.json')
|
|
if not os.path.exists(config_path):
|
|
print(f"📥 Downloading config.json for {model_key}...")
|
|
_download_transformers_config(info['repo_id'], config_dir)
|
|
except Exception as config_error:
|
|
print(f"⚠️ Could not download config after migration: {config_error}")
|
|
|
|
except Exception as migrate_error:
|
|
print(f"⚠️ Migration failed: {migrate_error}")
|
|
else:
|
|
# Use same directory for non-transformers models
|
|
safetensors_path = model_path.replace(".pt", ".safetensors")
|
|
if not safetensors_path.endswith(".safetensors"):
|
|
safetensors_path = model_path + ".safetensors"
|
|
|
|
# For transformers models in subdirectories, ensure config exists even if model already exists
|
|
if target_subdir and model_key and os.path.exists(safetensors_path):
|
|
try:
|
|
from engines.rvc.hubert_downloader import _download_transformers_config
|
|
info = HUBERT_MODELS.get(model_key)
|
|
if info and info.get('repo_id'):
|
|
config_dir = os.path.dirname(safetensors_path)
|
|
config_path = os.path.join(config_dir, 'config.json')
|
|
if not os.path.exists(config_path):
|
|
print(f"📥 Safetensors exists but config.json missing, downloading for {model_key}...")
|
|
_download_transformers_config(info['repo_id'], config_dir)
|
|
except Exception as config_error:
|
|
print(f"⚠️ Could not download missing config: {config_error}")
|
|
|
|
# If safetensors version doesn't exist, create it
|
|
if not os.path.exists(safetensors_path):
|
|
try:
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
print(f"🔧 Converting {model_path} to {safetensors_path}")
|
|
|
|
# Try to extract just the model weights, ignoring fairseq objects
|
|
try:
|
|
# First attempt: load with torch pickle_module to handle fairseq objects
|
|
import pickle
|
|
import io
|
|
|
|
with open(model_path, 'rb') as f:
|
|
# Load raw data and try to extract model state_dict only
|
|
checkpoint = torch.load(f, map_location='cpu', weights_only=False)
|
|
|
|
# Extract state_dict from various possible formats
|
|
if hasattr(checkpoint, 'state_dict'):
|
|
state_dict = checkpoint.state_dict()
|
|
elif isinstance(checkpoint, dict):
|
|
if 'model' in checkpoint:
|
|
if hasattr(checkpoint['model'], 'state_dict'):
|
|
state_dict = checkpoint['model'].state_dict()
|
|
else:
|
|
state_dict = checkpoint['model']
|
|
elif 'state_dict' in checkpoint:
|
|
state_dict = checkpoint['state_dict']
|
|
else:
|
|
# Assume the dict itself is the state_dict
|
|
state_dict = {k: v for k, v in checkpoint.items()
|
|
if isinstance(v, torch.Tensor)}
|
|
else:
|
|
raise ValueError("Cannot extract state_dict from checkpoint")
|
|
|
|
except Exception as load_error:
|
|
print(f"⚠️ Standard loading failed: {load_error}")
|
|
# Fallback: try to manually extract tensors
|
|
raise NotImplementedError(f"Cannot convert complex .pt file without fairseq: {model_path}")
|
|
|
|
# Save as safetensors (without config metadata - this is the limitation)
|
|
save_file(state_dict, safetensors_path)
|
|
print(f"✅ Converted to safetensors: {safetensors_path}")
|
|
|
|
# For transformers models, download config.json to the subdirectory
|
|
if target_subdir and model_key:
|
|
try:
|
|
from engines.rvc.hubert_downloader import _download_transformers_config
|
|
info = HUBERT_MODELS.get(model_key)
|
|
if info and info.get('repo_id'):
|
|
config_dir = os.path.dirname(safetensors_path)
|
|
config_path = os.path.join(config_dir, 'config.json')
|
|
if not os.path.exists(config_path):
|
|
print(f"📥 Downloading config.json for {model_key}...")
|
|
_download_transformers_config(info['repo_id'], config_dir)
|
|
else:
|
|
print(f"✅ Config already exists: {config_path}")
|
|
except Exception as config_error:
|
|
print(f"⚠️ Could not download config: {config_error}")
|
|
else:
|
|
print(f"⚠️ Note: Config metadata not preserved - may cause loading issues")
|
|
|
|
except Exception as conv_error:
|
|
print(f"❌ Failed to convert model: {conv_error}")
|
|
raise NotImplementedError(f"Cannot load .pt file without fairseq: {model_path}")
|
|
|
|
# Load the safetensors version
|
|
return HubertModelWithFinalProj.from_safetensors(safetensors_path, device=config.device)
|
|
except Exception as e:
|
|
print(e)
|
|
return None
|
|
|
|
def change_rms(data1, sr1, data2, sr2, rate): # 1是输入音频，2是输出音频,rate是2的占比
|
|
# print(data1.max(),data2.max())
|
|
|
|
# Manual RMS calculation to avoid librosa/numba issues in Python 3.13
|
|
def calculate_rms(y, frame_length, hop_length):
|
|
"""Calculate RMS without librosa to avoid numba issues"""
|
|
import numpy as np
|
|
|
|
# Pad the signal
|
|
y = np.asarray(y)
|
|
n_frames = 1 + (len(y) - frame_length) // hop_length
|
|
|
|
# Create frames
|
|
frames = np.lib.stride_tricks.sliding_window_view(
|
|
y, window_shape=frame_length
|
|
)[::hop_length]
|
|
|
|
# Calculate RMS for each frame
|
|
rms = np.sqrt(np.mean(frames**2, axis=1))
|
|
return rms.reshape(1, -1) # Shape: (1, n_frames)
|
|
|
|
rms1 = calculate_rms(data1, sr1 // 2 * 2, sr1 // 2) # 每半秒一个点
|
|
rms2 = calculate_rms(data2, sr2 // 2 * 2, sr2 // 2)
|
|
rms1 = torch.from_numpy(rms1)
|
|
rms1 = F.interpolate(
|
|
rms1.unsqueeze(0), size=data2.shape[0], mode="linear"
|
|
).squeeze()
|
|
rms2 = torch.from_numpy(rms2)
|
|
rms2 = F.interpolate(
|
|
rms2.unsqueeze(0), size=data2.shape[0], mode="linear"
|
|
).squeeze()
|
|
rms2 = torch.max(rms2, torch.zeros_like(rms2) + 1e-6)
|
|
data2 *= (
|
|
torch.pow(rms1, torch.tensor(1 - rate))
|
|
* torch.pow(rms2, torch.tensor(rate - 1))
|
|
).numpy()
|
|
return data2 |