|
from pathlib import Path
|
|
|
|
import click
|
|
import hydra
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
import torchaudio
|
|
from hydra import compose, initialize
|
|
from hydra.utils import instantiate
|
|
from loguru import logger
|
|
from omegaconf import OmegaConf
|
|
|
|
from tools.file import AUDIO_EXTENSIONS
|
|
|
|
|
|
OmegaConf.register_new_resolver("eval", eval)
|
|
|
|
|
|
def load_model(config_name, checkpoint_path, device="cuda"):
|
|
hydra.core.global_hydra.GlobalHydra.instance().clear()
|
|
with initialize(version_base="1.3", config_path="../../fish_speech/configs"):
|
|
cfg = compose(config_name=config_name)
|
|
|
|
model = instantiate(cfg)
|
|
state_dict = torch.load(
|
|
checkpoint_path,
|
|
map_location=device,
|
|
)
|
|
if "state_dict" in state_dict:
|
|
state_dict = state_dict["state_dict"]
|
|
|
|
if any("generator" in k for k in state_dict):
|
|
state_dict = {
|
|
k.replace("generator.", ""): v
|
|
for k, v in state_dict.items()
|
|
if "generator." in k
|
|
}
|
|
|
|
result = model.load_state_dict(state_dict, strict=False)
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
logger.info(f"Loaded model: {result}")
|
|
return model
|
|
|
|
|
|
@torch.no_grad()
|
|
@click.command()
|
|
@click.option(
|
|
"--input-path",
|
|
"-i",
|
|
default="test.wav",
|
|
type=click.Path(exists=True, path_type=Path),
|
|
)
|
|
@click.option(
|
|
"--output-path", "-o", default="fake.wav", type=click.Path(path_type=Path)
|
|
)
|
|
@click.option("--config-name", default="firefly_gan_vq")
|
|
@click.option(
|
|
"--checkpoint-path",
|
|
default="checkpoints/fish-speech-1.4/firefly-gan-vq-fsq-8x1024-21hz-generator.pth",
|
|
)
|
|
@click.option(
|
|
"--device",
|
|
"-d",
|
|
default="cuda",
|
|
)
|
|
|
|
|
|
def main(input_path, output_path, config_name, checkpoint_path, device):
|
|
model = load_model(config_name, checkpoint_path, device=device)
|
|
|
|
if input_path.suffix in AUDIO_EXTENSIONS:
|
|
logger.info(f"Processing in-place reconstruction of {input_path}")
|
|
|
|
|
|
audio, sr = torchaudio.load(str(input_path))
|
|
if audio.shape[0] > 1:
|
|
audio = audio.mean(0, keepdim=True)
|
|
audio = torchaudio.functional.resample(
|
|
audio, sr, model.spec_transform.sample_rate
|
|
)
|
|
|
|
audios = audio[None].to(device)
|
|
logger.info(
|
|
f"Loaded audio with {audios.shape[2] / model.spec_transform.sample_rate:.2f} seconds"
|
|
)
|
|
|
|
|
|
audio_lengths = torch.tensor([audios.shape[2]], device=device, dtype=torch.long)
|
|
indices = model.encode(audios, audio_lengths)[0][0]
|
|
|
|
logger.info(f"Generated indices of shape {indices.shape}")
|
|
|
|
|
|
np.save(output_path.with_suffix(".npy"), indices.cpu().numpy())
|
|
elif input_path.suffix == ".npy":
|
|
logger.info(f"Processing precomputed indices from {input_path}")
|
|
indices = np.load(input_path)
|
|
indices = torch.from_numpy(indices).to(device).long()
|
|
assert indices.ndim == 2, f"Expected 2D indices, got {indices.ndim}"
|
|
else:
|
|
raise ValueError(f"Unknown input type: {input_path}")
|
|
|
|
|
|
feature_lengths = torch.tensor([indices.shape[1]], device=device)
|
|
fake_audios, _ = model.decode(
|
|
indices=indices[None], feature_lengths=feature_lengths
|
|
)
|
|
audio_time = fake_audios.shape[-1] / model.spec_transform.sample_rate
|
|
|
|
logger.info(
|
|
f"Generated audio of shape {fake_audios.shape}, equivalent to {audio_time:.2f} seconds from {indices.shape[1]} features, features/second: {indices.shape[1] / audio_time:.2f}"
|
|
)
|
|
|
|
|
|
fake_audio = fake_audios[0, 0].float().cpu().numpy()
|
|
sf.write(output_path, fake_audio, model.spec_transform.sample_rate)
|
|
logger.info(f"Saved audio to {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
