Upload folder using huggingface_hub

app.py

@@ -0,0 +1,272 @@
+import argparse
+parser = argparse.ArgumentParser(description="WhisperVQ Application")
+parser.add_argument('--log-path', type=str,
+                    default='whisper.log', help='The log file path')
+parser.add_argument('--log-level', type=str, default='INFO',
+                    choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'TRACE'], help='The log level')
+parser.add_argument('--port', type=int, default=3348,
+                    help='The port to run the WhisperVQ app on')
+parser.add_argument('--package-dir', type=str, default="",
+                    help='The package-dir to be extended to sys.path')
+args = parser.parse_args()
+import sys
+sys.path.insert(0, args.environment)
+import tempfile
+from typing import Tuple
+from enum import Enum
+import io
+import logging
+from custom_component import CustomRQBottleneckTransformer
+from whisperspeech.vq_stoks import RQBottleneckTransformer
+from huggingface_hub import hf_hub_download
+import uvicorn
+from transformers import WhisperModel, WhisperProcessor
+from fastapi.responses import JSONResponse
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from contextlib import asynccontextmanager
+import torchaudio
+import torch
+import os
+import time
+import psutil
+import threading
+
+
+logging.basicConfig(level=args.log_level, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+                    handlers=[
+                        logging.FileHandler(args.log_path),
+                        # logging.StreamHandler()
+                    ])
+logger = logging.getLogger(__name__)
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Use the first GPU
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if not os.path.exists(os.path.dirname(os.path.realpath(__file__))+"/whisper-vq-stoks-v3-7lang-fixed.model"):
+    hf_hub_download(
+        repo_id="jan-hq/WhisperVQ",
+        filename="whisper-vq-stoks-v3-7lang-fixed.model",
+        local_dir=".",
+    )
+vq_model = CustomRQBottleneckTransformer.load_vq_only(
+    os.path.dirname(os.path.realpath(__file__)) +
+    "/whisper-vq-stoks-v3-7lang-fixed.model"
+).to(device)
+vq_model.load_encoder(device)
+vq_model.eval()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+
+    yield
+    # on shutdown
+
+
+# vq_model = torch.compile(vq_model)
+
+
+class AudioFormat(str, Enum):
+    WAV = "wav"    # Supported by both backends
+    MP3 = "mp3"    # Supported by ffmpeg
+    FLAC = "flac"  # Supported by both
+    AAC = "aac"    # Supported by ffmpeg
+    OGG = "ogg"    # Supported by ffmpeg
+    OPUS = "opus"  # Supported by ffmpeg
+    PCM = "pcm"    # Raw PCM data
+
+
+# Format to backend mapping
+FORMAT_BACKENDS = {
+    AudioFormat.WAV: ["soundfile", "ffmpeg"],
+    AudioFormat.MP3: ["ffmpeg"],
+    AudioFormat.FLAC: ["soundfile", "ffmpeg"],
+    AudioFormat.AAC: ["ffmpeg"],
+    AudioFormat.OGG: ["ffmpeg"],
+    AudioFormat.OPUS: ["ffmpeg"],
+    AudioFormat.PCM: ["soundfile"]
+}
+
+
+class AudioProcessor:
+    def __init__(self):
+        self.available_backends = torchaudio.list_audio_backends()
+        logger.info(f"Available backends: {self.available_backends}")
+
+        # Verify ffmpeg support
+        self.has_ffmpeg = "ffmpeg" in self.available_backends
+        if not self.has_ffmpeg:
+            logger.warning(
+                "FFMPEG backend not available. Some formats may not be supported")
+
+    def _get_best_backend(self, format: AudioFormat) -> str:
+        """Determine the best backend for the given format"""
+        supported_backends = FORMAT_BACKENDS[format]
+        for backend in supported_backends:
+            if backend in self.available_backends:
+                return backend
+        raise ValueError(f"No available backend supports format {format}")
+
+    async def load_audio(
+        self,
+        file_obj: bytes,
+        format: AudioFormat,
+        target_sr: int = 16000
+    ) -> Tuple[torch.Tensor, int]:
+        """
+        Load audio from bytes object with format handling
+
+        Args:
+            file_obj: Audio file bytes
+            format: Audio format enum
+            target_sr: Target sample rate (default: 16000)
+
+        Returns:
+            Tuple[torch.Tensor, int]: Audio tensor and sample rate
+        """
+        try:
+            # Get appropriate backend
+            backend = self._get_best_backend(format)
+            torchaudio.set_audio_backend(backend)
+            logger.info(f"Using {backend} backend for {format} format")
+
+            if format == AudioFormat.PCM:
+                # Handle raw PCM
+                wav = torch.frombuffer(file_obj, dtype=torch.int16)
+                wav = wav.float() / 32768.0  # Normalize to [-1, 1]
+                wav = wav.unsqueeze(0)  # Add channel dimension
+                sr = target_sr
+            else:
+                # For formats that might need ffmpeg processing
+                if os.name == "nt":  # for windows
+                    wav, sr = torchaudio.load(io.BytesIO(file_obj))
+                else:
+                    with tempfile.NamedTemporaryFile(suffix=f".{format}") as temp_file:
+                        # Write bytes to temporary file
+                        temp_file.write(file_obj)
+                        temp_file.flush()
+
+                        # Load audio
+                        wav, sr = torchaudio.load(temp_file.name)
+
+            # Convert to mono if stereo
+            if wav.shape[0] > 1:
+                wav = torch.mean(wav, dim=0, keepdim=True)
+
+            # Resample if needed
+            if sr != target_sr:
+                wav = torchaudio.functional.resample(wav, sr, target_sr)
+                sr = target_sr
+
+            return wav, sr
+
+        except Exception as e:
+            logger.error(f"Error loading audio: {e}")
+            raise HTTPException(
+                status_code=400,
+                detail=f"Error processing {format} audio: {str(e)}"
+            )
+
+    def get_format_info(self) -> dict:
+        """Get information about supported formats"""
+        supported_formats = {}
+        for format in AudioFormat:
+            try:
+                backend = self._get_best_backend(format)
+                supported_formats[format] = {
+                    "supported": True,
+                    "backend": backend
+                }
+            except ValueError:
+                supported_formats[format] = {
+                    "supported": False,
+                    "backend": None
+                }
+        return supported_formats
+
+
+audio_processor = AudioProcessor()
+
+app = FastAPI(lifespan=lifespan)
+
+
+@app.get("/supported_formats")
+async def get_supported_formats():
+    """Endpoint to check supported formats"""
+    return audio_processor.get_format_info()
+
+
+@app.post("/tokenize/{format}")
+async def tokenize_audio(format: AudioFormat = "wav", file: UploadFile = File(...)):
+    try:
+        # Read file
+        file_obj = await file.read()
+
+        # Load and process audio
+        wav, sr = await audio_processor.load_audio(file_obj, format)
+
+        # Ensure we're using CUDA if available
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        wav = wav.to(device)
+
+        # Generate tokens
+        with torch.no_grad():
+            codes = vq_model.encode_audio(wav)
+            codes = codes[0].cpu().tolist()
+
+        # Format result
+        result = ''.join(f'<|sound_{num:04d}|>' for num in codes)
+
+        return JSONResponse(content={
+            "model_name": "whisper-vq-stoks-v3-7lang-fixed.model",
+            "tokens": f'<|sound_start|>{result}<|sound_end|>',
+            "format": format,
+            "sample_rate": sr,
+            "backend_used": audio_processor._get_best_backend(format)
+        })
+
+    except Exception as e:
+        logger.error(f"Error processing request: {e}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error processing request: {str(e)}"
+        )
+
+
+def self_terminate():
+    time.sleep(1)
+    parent = psutil.Process(psutil.Process(os.getpid()).ppid())
+    parent.kill()
+
+
+@app.post("/kill")
+async def kill():
+    threading.Thread(target=self_terminate, daemon=True).start()
+    return {"success": True}
+
+if __name__ == "__main__":
+    import uvicorn
+    from uvicorn.config import LOGGING_CONFIG
+
+    LOGGING_CONFIG["handlers"]["default"] = {
+        "class": "logging.FileHandler",
+        "filename": args.log_path,
+        "formatter": "default"
+    }
+    LOGGING_CONFIG["handlers"]["access"] = {
+        "class": "logging.FileHandler",
+        "filename": args.log_path,
+        "formatter": "access"
+    }
+    LOGGING_CONFIG["loggers"]["uvicorn.error"]["level"] = args.log_level
+    LOGGING_CONFIG["loggers"]["uvicorn.access"]["level"] = args.log_level
+
+# Print supported formats at startup
+    processor = AudioProcessor()
+    format_info = processor.get_format_info()
+    logger.info("Supported formats:")
+    for format, info in format_info.items():
+        logger.info(f"{format}: {info}")
+
+    uvicorn.run(app, host="0.0.0.0", port=args.port)

custom_component.py

@@ -0,0 +1,172 @@
+import torch
+import torch.nn as nn
+import whisper
+from whisper.model import AudioEncoder, ModelDimensions
+from typing import Dict, Optional
+from whisperspeech.vq_stoks import RQBottleneckTransformer, Tunables
+from huggingface_hub import hf_hub_download
+import torch.nn.functional as F
+import os
+from typing import List, Optional, Union
+import io
+import urllib
+from tqdm import tqdm
+import torchaudio
+_HF_MODELS = {  
+    "medium": "https://huggingface.co/jan-hq/WhisperVQ/resolve/main/medium_encoder_only.pt",
+}
+def available_models() -> List[str]:
+    """Returns the names of available models"""
+    return list(_HF_MODELS.keys())
+def _download(url: str, root: str, in_memory: bool) -> Union[bytes, str]:
+    os.makedirs(root, exist_ok=True)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, os.path.basename(url))
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        with open(download_target, "rb") as f:
+            model_bytes = f.read()
+        return model_bytes if in_memory else download_target
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(
+            total=int(source.info().get("Content-Length")),
+            ncols=80,
+            unit="iB",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    model_bytes = open(download_target, "rb").read()
+    return model_bytes if in_memory else download_target
+class CustomWhisperEncoder(nn.Module):
+    """
+    Lightweight wrapper that only loads the AudioEncoder part of Whisper
+    """
+    def __init__(self, name: str, device: str = None, download_root: str = None, in_memory: bool = False,):
+        super().__init__()
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        if download_root is None:
+            default = os.path.join(os.path.expanduser("~"), ".cache")
+            download_root = os.path.dirname(os.path.realpath(__file__)) #os.path.join(os.getenv("XDG_CACHE_HOME", default), "whisper")
+
+        if name in _HF_MODELS:
+            checkpoint_file = _download(_HF_MODELS[name], download_root, in_memory)
+        elif os.path.isfile(name):
+            checkpoint_file = open(name, "rb").read() if in_memory else name
+        else:
+            raise RuntimeError(
+                f"Model {name} not found; available models = {available_models()}"
+            )
+        
+        # Load weights
+        with (
+            io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb")
+        ) as fp:
+            checkpoint = torch.load(fp, map_location=device)
+        del checkpoint_file
+        dims = ModelDimensions(**checkpoint["dims"])
+        self.encoder = AudioEncoder(
+            dims.n_mels,
+            dims.n_audio_ctx,
+            dims.n_audio_state,
+            dims.n_audio_head,
+            dims.n_audio_layer,
+        )
+        
+        self.encoder.load_state_dict(checkpoint["model_state_dict"])
+        
+        if device:
+            self.to(device)
+        
+        self.eval()
+
+    def forward(self, mel: torch.Tensor):
+        return self.encoder(mel)
+    
+class CustomRQBottleneckTransformer(RQBottleneckTransformer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+    @classmethod
+    def load_vq_only(cls, ref="collabora/spear-tts-pytorch:whisper-vq-stoks-medium-en+pl.model",
+                   repo_id=None, filename=None, local_filename=None):
+        if repo_id is None and filename is None and local_filename is None:
+            if ":" in ref:
+                repo_id, filename = ref.split(":", 1)
+            else:
+                local_filename = ref
+        if not local_filename:
+            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
+        
+        # Load the spec
+        spec = torch.load(local_filename)
+        
+        # Create instance with minimal required components
+        instance = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec.get('tunables', {}))))
+        
+        # Load only necessary state dict entries
+        required_components = {
+            'rq', 'mlp', 'mlp_ln'
+        }
+        filtered_state_dict = {
+            k: v for k, v in spec['state_dict'].items()
+            if any(k.startswith(comp) for comp in required_components)
+        }
+        
+        instance.load_state_dict(filtered_state_dict, strict=False)
+        instance.eval()
+        return instance
+
+    def load_encoder(self, device=None):
+        if self.whmodel is not None: return
+        device = device or self.device
+        # Use our custom encoder-only model
+        if self.whmodel is None:
+            encoder = CustomWhisperEncoder(self.whisper_model_name, device=device)
+            self.whmodel = [encoder]
+        multilingual = not self.whisper_model_name.endswith('.en')
+        self.tokenizer = whisper.tokenizer.get_tokenizer(multilingual)
+
+    def optimzed_encode_mel(self, mel):
+        assert len(mel.shape) == 3, "invalid mel spectrogram shape, expect (batch,chn,time)"
+        self.load_encoder()
+        n = mel.shape[-1]
+        if n > whisper.audio.N_FRAMES:
+            padding = 0
+            padded = mel[:,:,:whisper.audio.N_FRAMES]
+        else:
+            padding = -n % whisper.audio.N_FRAMES
+            padded = F.pad(mel, (0, padding), value=-1.5)
+        embs = self.whmodel[0].encoder(padded)#.to(self.whmodel[0].device))#[:,:n//2]
+        stoks = self.quantize(embs)
+        if self.tunables.mask_embs:
+            return stoks[:,:n//2//self.downsample]
+        else:
+            return stoks
+    # overide
+    def encode_audio(self, audio):
+        if isinstance(audio, str):
+            x, sr = torchaudio.load(audio)
+            x = torchaudio.transforms.Resample(sr, 16000)(x)[0]
+            audio = x.unsqueeze(0)
+        return self.optimzed_encode_mel(self.log_mel_spectrogram(audio).to(self.device))
+    
+if __name__ == "__main__":
+    # Load the model
+    vqmodel = CustomRQBottleneckTransformer.load_vq_only(
+        "whisper-vq-stoks-v3-7lang-fixed.model"
+    ).to("cuda")
+    vqmodel.load_encoder('cuda')
+    vqmodel.eval()

download_model.py

@@ -0,0 +1,51 @@
+import urllib
+from tqdm import tqdm
+from huggingface_hub import hf_hub_download
+import os
+
+encoder_url = "https://huggingface.co/jan-hq/WhisperVQ/resolve/main/medium_encoder_only.pt"
+
+
+def _download(url: str, root: str, in_memory: bool):
+    os.makedirs(root, exist_ok=True)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, os.path.basename(url))
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(
+            f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        with open(download_target, "rb") as f:
+            model_bytes = f.read()
+        return model_bytes if in_memory else download_target
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(
+            total=int(source.info().get("Content-Length")),
+            ncols=80,
+            unit="iB",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    model_bytes = open(download_target, "rb").read()
+    return model_bytes if in_memory else download_target
+
+
+if not os.path.exists(os.path.dirname(os.path.realpath(__file__))+"/whisper-vq-stoks-v3-7lang-fixed.model"):
+    hf_hub_download(
+        repo_id="jan-hq/WhisperVQ",
+        filename="whisper-vq-stoks-v3-7lang-fixed.model",
+        local_dir=".",
+    )
+
+_download(encoder_url, os.path.dirname(os.path.realpath(__file__)), False)

medium_encoder_only.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c713c13a07f8980dfd81a11600424755a1ede2f4904803e7af042a7a94ab168d
+size 614547576

requirements.cuda.txt

@@ -0,0 +1,24 @@
+openai-whisper==20231117
+huggingface_hub
+IPython
+pyarrow
+matplotlib
+librosa
+soundfile
+datasets
+encodec
+boto3
+fire
+vector_quantize_pytorch
+webdataset
+whisperspeech
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.2.0
+torchaudio==2.2.0
+numpy==1.26.4
+fastapi
+uvicorn
+
+python-multipart
+transformers
+psutil

requirements.txt

@@ -0,0 +1,23 @@
+openai-whisper==20231117
+huggingface_hub
+IPython
+pyarrow
+matplotlib
+librosa
+soundfile
+datasets
+encodec
+boto3
+fire
+vector_quantize_pytorch
+webdataset
+whisperspeech
+torch==2.2.0
+torchaudio==2.2.0
+numpy==1.26.4
+fastapi
+uvicorn
+
+python-multipart
+transformers
+psutil

whisper-vq-stoks-v3-7lang-fixed.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09e23368136f07ba474dd50fd728f1d216f4542550c456e8065855969b1df730
+size 90921877