Upload folder using huggingface_hub

415215e verified 11 months ago

7.83 kB

import sys
import os

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)

from transformers import PreTrainedModel, PretrainedConfig, AutoConfig
import torch
import numpy as np
from f5_tts.infer.utils_infer import (
    infer_process,
    load_model,
    load_vocoder,
    preprocess_ref_audio_text,
)
from f5_tts.model import DiT
import soundfile as sf
import io
from pydub import AudioSegment, silence
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

class INF5Config(PretrainedConfig):
    model_type = "inf5"

    def __init__(self, ckpt_repo_id: str = None, vocab_repo_id: str = None,
                 ckpt_filename: str = None, vocab_filename: str = "vocab.txt",
                 speed: float = 1.0, remove_sil: bool = True, **kwargs):
        super().__init__(**kwargs)
        # If not specified, use the model's own repo for both
        self.ckpt_repo_id = ckpt_repo_id
        self.vocab_repo_id = vocab_repo_id
        self.ckpt_filename = ckpt_filename
        self.vocab_filename = vocab_filename
        self.speed = speed
        self.remove_sil = remove_sil

class INF5Model(PreTrainedModel):
    config_class = INF5Config

    def __init__(self, config):
        super().__init__(config)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Load vocoder
        self.vocoder = torch.compile(
            load_vocoder(vocoder_name="vocos", is_local=False, device=device)
        )

        # Determine which repo to use for vocab
        # Default to the model's own repo if not specified
        vocab_repo = config.vocab_repo_id or config.name_or_path
        
        # Download vocab.txt from HF Hub
        vocab_path = hf_hub_download(repo_id=vocab_repo, filename=config.vocab_filename)

        # Determine which repo to use for checkpoint
        ckpt_repo = config.ckpt_repo_id or config.name_or_path

        ckpt_candidates = [
            "model_last.pt",  # Try this first since it's in your repo
            "checkpoints/model.safetensors",
            "model.safetensors",
            "checkpoints/pytorch_model.bin",
            "pytorch_model.bin",
            "checkpoints/model.pt",
            "model.pt",
            "checkpoints/checkpoint.pt",
            "checkpoint.pt"
        ]

        # If a specific checkpoint filename is provided, use only that
        if config.ckpt_filename:
            ckpt_candidates = [config.ckpt_filename]

        ckpt_path = None

        for fname in ckpt_candidates:
            try:
                ckpt_path = hf_hub_download(repo_id=ckpt_repo, filename=fname)
                print(f"Found checkpoint on hub: {fname} -> {ckpt_path}")
                break
            except Exception as e:
                # ignore and try next candidate; but log for debugging
                # common failures: file not found, LFS not enabled, permission issues
                # print(f"Attempt to download {fname} failed: {e}")
                continue

        if ckpt_path is None:
            raise RuntimeError(
                "Could not find a checkpoint file on the Hub. "
                "Tried: " + ", ".join(ckpt_candidates) + ".\n"
                "If your checkpoint is stored under a different path or name, "
                "update ckpt_candidates or pass the path via config (e.g. config.ckpt_filename). "
                "If the file is >5GB, ensure Git LFS is enabled for the repo (hf lfs-enable-largefiles)."
            )

        # Pass ckpt_path to load_model. Use keyword to avoid mismatch in positional args.
        self.ema_model = torch.compile(
            load_model(
                DiT,
                dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),
                mel_spec_type="vocos",
                vocab_file=vocab_path,
                device=device,
                ckpt_path=ckpt_path
            )
        )

        # Optionally: if load_model returns an uninitialized model and you want to load a state dict:
        # state_dict = load_file(ckpt_path, device=str(device))
        # self.ema_model.load_state_dict(state_dict, strict=False)

    
    def forward(self, text: str, ref_audio_path: str, ref_text: str):
        """
        Generate speech given a reference audio & text input.
        
        Args:
            text (str): The text to be synthesized.
            ref_audio_path (str): Path to the reference audio file.
            ref_text (str): The reference text.
        Returns:
            np.array: Generated waveform.
        """

        if not os.path.exists(ref_audio_path):
            raise FileNotFoundError(f"Reference audio file {ref_audio_path} not found.")

        # Load reference audio & text
        ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_path, ref_text)

        
        self.ema_model.to(self.device)
        self.vocoder.to(self.device)
        
        # Perform inference
        audio, final_sample_rate, _ = infer_process(
            ref_audio,
            ref_text,
            text,
            self.ema_model,
            self.vocoder,
            mel_spec_type="vocos",
            speed=self.config.speed,
            device=self.device,
        )

        # Convert to pydub format and remove silence if needed
        buffer = io.BytesIO()
        sf.write(buffer, audio, samplerate=24000, format="WAV")
        buffer.seek(0)
        audio_segment = AudioSegment.from_file(buffer, format="wav")

        if self.config.remove_sil:
            non_silent_segs = silence.split_on_silence(
                audio_segment,
                min_silence_len=1000,
                silence_thresh=-50,
                keep_silence=500,
                seek_step=10,
            )
            non_silent_wave = sum(non_silent_segs, AudioSegment.silent(duration=0))
            audio_segment = non_silent_wave

        # Normalize loudness
        target_dBFS = -20.0
        change_in_dBFS = target_dBFS - audio_segment.dBFS
        audio_segment = audio_segment.apply_gain(change_in_dBFS)

        return np.array(audio_segment.get_array_of_samples())



if __name__ == '__main__':
    model = INF5Model(INF5Config())
    model.save_pretrained("INF5")
    model.config.save_pretrained("INF5")
    

    # import numpy as np
    # import soundfile as sf
    # from transformers import AutoConfig, AutoModel
    
    # AutoConfig.register("inf5", INF5Config)
    # AutoModel.register(INF5Config, INF5Model)

    # model = AutoModel.from_pretrained("INF5")
    # audio = model("नमस्ते! संगीत की तरह जीवन भी खूबसूरत होता है, बस इसे सही ताल में जीना आना चाहिए.", 
    #               ref_audio_path="prompts/PAN_F_HAPPY_00001.wav",
    #               ref_text="भਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ  ਹਨ।")
    
    # if audio.dtype == np.int16:
    #     audio = audio.astype(np.float32) / 32768.0  
    # sf.write("samples/namaste.wav", np.array(audio, dtype=np.float32), samplerate=24000)

    # from huggingface_hub import HfApi

    # repo_id = "svp19/INF5"  # Change to your HF repo

    # # Upload model directory to HF
    # api = HfApi()
    # api.upload_folder(
    #     folder_path="INF5",
    #     repo_id=repo_id,
    #     repo_type="model"
    # )
    # print(f"Model pushed to https://huggingface.co/{repo_id} 🚀")

    # print("Verify Upload")
    # from transformers import AutoModel
    # model = AutoModel.from_pretrained(repo_id)
    # print("Success")

	import sys
	import os

	current_dir = os.path.dirname(os.path.abspath(__file__))
	sys.path.append(current_dir)

	from transformers import PreTrainedModel, PretrainedConfig, AutoConfig
	import torch
	import numpy as np
	from f5_tts.infer.utils_infer import (
	infer_process,
	load_model,
	load_vocoder,
	preprocess_ref_audio_text,
	)
	from f5_tts.model import DiT
	import soundfile as sf
	import io
	from pydub import AudioSegment, silence
	from huggingface_hub import hf_hub_download
	from safetensors.torch import load_file

	class INF5Config(PretrainedConfig):
	model_type = "inf5"

	def __init__(self, ckpt_repo_id: str = None, vocab_repo_id: str = None,
	ckpt_filename: str = None, vocab_filename: str = "vocab.txt",
	speed: float = 1.0, remove_sil: bool = True, **kwargs):
	super().__init__(**kwargs)
	# If not specified, use the model's own repo for both
	self.ckpt_repo_id = ckpt_repo_id
	self.vocab_repo_id = vocab_repo_id
	self.ckpt_filename = ckpt_filename
	self.vocab_filename = vocab_filename
	self.speed = speed
	self.remove_sil = remove_sil

	class INF5Model(PreTrainedModel):
	config_class = INF5Config

	def __init__(self, config):
	super().__init__(config)
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	# Load vocoder
	self.vocoder = torch.compile(
	load_vocoder(vocoder_name="vocos", is_local=False, device=device)
	)

	# Determine which repo to use for vocab
	# Default to the model's own repo if not specified
	vocab_repo = config.vocab_repo_id or config.name_or_path

	# Download vocab.txt from HF Hub
	vocab_path = hf_hub_download(repo_id=vocab_repo, filename=config.vocab_filename)

	# Determine which repo to use for checkpoint
	ckpt_repo = config.ckpt_repo_id or config.name_or_path

	ckpt_candidates = [
	"model_last.pt", # Try this first since it's in your repo
	"checkpoints/model.safetensors",
	"model.safetensors",
	"checkpoints/pytorch_model.bin",
	"pytorch_model.bin",
	"checkpoints/model.pt",
	"model.pt",
	"checkpoints/checkpoint.pt",
	"checkpoint.pt"
	]

	# If a specific checkpoint filename is provided, use only that
	if config.ckpt_filename:
	ckpt_candidates = [config.ckpt_filename]

	ckpt_path = None

	for fname in ckpt_candidates:
	try:
	ckpt_path = hf_hub_download(repo_id=ckpt_repo, filename=fname)
	print(f"Found checkpoint on hub: {fname} -> {ckpt_path}")
	break
	except Exception as e:
	# ignore and try next candidate; but log for debugging
	# common failures: file not found, LFS not enabled, permission issues
	# print(f"Attempt to download {fname} failed: {e}")
	continue

	if ckpt_path is None:
	raise RuntimeError(
	"Could not find a checkpoint file on the Hub. "
	"Tried: " + ", ".join(ckpt_candidates) + ".\n"
	"If your checkpoint is stored under a different path or name, "
	"update ckpt_candidates or pass the path via config (e.g. config.ckpt_filename). "
	"If the file is >5GB, ensure Git LFS is enabled for the repo (hf lfs-enable-largefiles)."
	)

	# Pass ckpt_path to load_model. Use keyword to avoid mismatch in positional args.
	self.ema_model = torch.compile(
	load_model(
	DiT,
	dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),
	mel_spec_type="vocos",
	vocab_file=vocab_path,
	device=device,
	ckpt_path=ckpt_path
	)
	)

	# Optionally: if load_model returns an uninitialized model and you want to load a state dict:
	# state_dict = load_file(ckpt_path, device=str(device))
	# self.ema_model.load_state_dict(state_dict, strict=False)


	def forward(self, text: str, ref_audio_path: str, ref_text: str):
	"""
	Generate speech given a reference audio & text input.

	Args:
	text (str): The text to be synthesized.
	ref_audio_path (str): Path to the reference audio file.
	ref_text (str): The reference text.
	Returns:
	np.array: Generated waveform.
	"""

	if not os.path.exists(ref_audio_path):
	raise FileNotFoundError(f"Reference audio file {ref_audio_path} not found.")

	# Load reference audio & text
	ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_path, ref_text)


	self.ema_model.to(self.device)
	self.vocoder.to(self.device)

	# Perform inference
	audio, final_sample_rate, _ = infer_process(
	ref_audio,
	ref_text,
	text,
	self.ema_model,
	self.vocoder,
	mel_spec_type="vocos",
	speed=self.config.speed,
	device=self.device,
	)

	# Convert to pydub format and remove silence if needed
	buffer = io.BytesIO()
	sf.write(buffer, audio, samplerate=24000, format="WAV")
	buffer.seek(0)
	audio_segment = AudioSegment.from_file(buffer, format="wav")

	if self.config.remove_sil:
	non_silent_segs = silence.split_on_silence(
	audio_segment,
	min_silence_len=1000,
	silence_thresh=-50,
	keep_silence=500,
	seek_step=10,
	)
	non_silent_wave = sum(non_silent_segs, AudioSegment.silent(duration=0))
	audio_segment = non_silent_wave

	# Normalize loudness
	target_dBFS = -20.0
	change_in_dBFS = target_dBFS - audio_segment.dBFS
	audio_segment = audio_segment.apply_gain(change_in_dBFS)

	return np.array(audio_segment.get_array_of_samples())



	if __name__ == '__main__':
	model = INF5Model(INF5Config())
	model.save_pretrained("INF5")
	model.config.save_pretrained("INF5")


	# import numpy as np
	# import soundfile as sf
	# from transformers import AutoConfig, AutoModel

	# AutoConfig.register("inf5", INF5Config)
	# AutoModel.register(INF5Config, INF5Model)

	# model = AutoModel.from_pretrained("INF5")
	# audio = model("नमस्ते! संगीत की तरह जीवन भी खूबसूरत होता है, बस इसे सही ताल में जीना आना चाहिए.",
	# ref_audio_path="prompts/PAN_F_HAPPY_00001.wav",
	# ref_text="भਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ ਹਨ।")

	# if audio.dtype == np.int16:
	# audio = audio.astype(np.float32) / 32768.0
	# sf.write("samples/namaste.wav", np.array(audio, dtype=np.float32), samplerate=24000)

	# from huggingface_hub import HfApi

	# repo_id = "svp19/INF5" # Change to your HF repo

	# # Upload model directory to HF
	# api = HfApi()
	# api.upload_folder(
	# folder_path="INF5",
	# repo_id=repo_id,
	# repo_type="model"
	# )
	# print(f"Model pushed to https://huggingface.co/{repo_id} 🚀")

	# print("Verify Upload")
	# from transformers import AutoModel
	# model = AutoModel.from_pretrained(repo_id)
	# print("Success")