codemo/fish-speech-1

1from pathlib import Path2 3import click4import hydra5import numpy as np6import soundfile as sf7import torch8import torchaudio9from hydra import compose, initialize10from hydra.utils import instantiate11from loguru import logger12from omegaconf import OmegaConf13 14from tools.file import AUDIO_EXTENSIONS15 16# register eval resolver17OmegaConf.register_new_resolver("eval", eval)18 19 20def load_model(config_name, checkpoint_path, device="cuda"):21    hydra.core.global_hydra.GlobalHydra.instance().clear()22    with initialize(version_base="1.3", config_path="../../fish_speech/configs"):23        cfg = compose(config_name=config_name)24 25    model = instantiate(cfg)26    state_dict = torch.load(27        checkpoint_path,28        map_location=device,29    )30    if "state_dict" in state_dict:31        state_dict = state_dict["state_dict"]32 33    if any("generator" in k for k in state_dict):34        state_dict = {35            k.replace("generator.", ""): v36            for k, v in state_dict.items()37            if "generator." in k38        }39 40    result = model.load_state_dict(state_dict, strict=False)41    model.eval()42    model.to(device)43 44    logger.info(f"Loaded model: {result}")45    return model46 47 48@torch.no_grad()49@click.command()50@click.option(51    "--input-path",52    "-i",53    default="test.wav",54    type=click.Path(exists=True, path_type=Path),55)56@click.option(57    "--output-path", "-o", default="fake.wav", type=click.Path(path_type=Path)58)59@click.option("--config-name", default="firefly_gan_vq")60@click.option(61    "--checkpoint-path",62    default="checkpoints/fish-speech-1.4/firefly-gan-vq-fsq-8x1024-21hz-generator.pth",63)64@click.option(65    "--device",66    "-d",67    default="cuda",68)69def main(input_path, output_path, config_name, checkpoint_path, device):70    model = load_model(config_name, checkpoint_path, device=device)71 72    if input_path.suffix in AUDIO_EXTENSIONS:73        logger.info(f"Processing in-place reconstruction of {input_path}")74 75        # Load audio76        audio, sr = torchaudio.load(str(input_path))77        if audio.shape[0] > 1:78            audio = audio.mean(0, keepdim=True)79        audio = torchaudio.functional.resample(80            audio, sr, model.spec_transform.sample_rate81        )82 83        audios = audio[None].to(device)84        logger.info(85            f"Loaded audio with {audios.shape[2] / model.spec_transform.sample_rate:.2f} seconds"86        )87 88        # VQ Encoder89        audio_lengths = torch.tensor([audios.shape[2]], device=device, dtype=torch.long)90        indices = model.encode(audios, audio_lengths)[0][0]91 92        logger.info(f"Generated indices of shape {indices.shape}")93 94        # Save indices95        np.save(output_path.with_suffix(".npy"), indices.cpu().numpy())96    elif input_path.suffix == ".npy":97        logger.info(f"Processing precomputed indices from {input_path}")98        indices = np.load(input_path)99        indices = torch.from_numpy(indices).to(device).long()100        assert indices.ndim == 2, f"Expected 2D indices, got {indices.ndim}"101    else:102        raise ValueError(f"Unknown input type: {input_path}")103 104    # Restore105    feature_lengths = torch.tensor([indices.shape[1]], device=device)106    fake_audios, _ = model.decode(107        indices=indices[None], feature_lengths=feature_lengths108    )109    audio_time = fake_audios.shape[-1] / model.spec_transform.sample_rate110 111    logger.info(112        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}"113    )114 115    # Save audio116    fake_audio = fake_audios[0, 0].float().cpu().numpy()117    sf.write(output_path, fake_audio, model.spec_transform.sample_rate)118    logger.info(f"Saved audio to {output_path}")119 120 121if __name__ == "__main__":122    main()123