Fine-tune a Whisper TTS model with Unsloth and Deploy it to Koyeb

# Install Dependencies
!pip install requests pydub openai-whisper
!echo "unsloth" | sudo -S apt update && sudo apt install -y ffmpeg

import os
import requests
from pydub import AudioSegment
from pydub.silence import detect_nonsilent
import whisper

# LibriVox URLs (Only the first 3 Chapters of the book)
LIBRIVOX_URLS=[
"<https://www.archive.org/download/adventuressherlockholmes_v4_1501_librivox/adventuresofsherlockholmes_01_doyle_128kb.mp3>"
    ]

# Download the LibriVox URLs
for url in LIBRIVOX_URLS:
    # Create a directory to store the downloaded files
    filename = url.split("/")[-1]
    os.makedirs("audio_files", exist_ok=True)
    print(f"Downloading {filename}")
    response = requests.get(url)
    with open(f"audio_files/{filename}", "wb") as f:
        f.write(response.content)
        print(f"Downloaded {filename}")

# Convert the audio files to WAV for
for file in os.listdir("audio_files"):
    if file.endswith(".mp3"):
        print(f"Converting {file} to WAV format")
        audio = AudioSegment.from_mp3(f"audio_files/{file}")
        audio.export(f"audio_files/{file.replace('.mp3', '.wav')}", format="wav")
        print(f"Converted {file} to WAV format")

# Split the audio files into chunks at natural pauses (silence)
target_chunk_length_ms = 7 * 1000  # Target 7 seconds per chunk
min_silence_len = 500  # Minimum silence length in ms to be considered a pause
silence_thresh = -40  # Silence threshold in dBFS (adjust if needed)

for file in os.listdir("audio_files"):
    if file.endswith(".wav") and not "chunk" in file:
        audio = AudioSegment.from_wav(f"audio_files/{file}")
        print(f"Processing {file}, duration: {len(audio)/1000:.2f} seconds")
        # Detect non-silent chunks
        print("Detecting non-silent chunks")
        nonsilent_ranges = detect_nonsilent(audio, min_silence_len=min_silence_len, silence_thresh=silence_thresh)
        if not nonsilent_ranges:
            print(f"Warning: No speech detected in {file}")
            continue
        # Group non-silent ranges into chunks of approximately target_chunk_length_ms
        print("Grouping non-silent ranges into chunks")
        chunks = []
        current_chunk_start = nonsilent_ranges[0][0]
        current_chunk_end = nonsilent_ranges[0][1]
        for start, end in nonsilent_ranges[1:]:
            # Check if adding this segment would exceed target length
            if end - current_chunk_start > target_chunk_length_ms:
                # Save current chunk and start a new one
                print("Saving current chunk")
                chunks.append((current_chunk_start, current_chunk_end))
                current_chunk_start = start
                current_chunk_end = end
            else:
                # Extend current chunk
                current_chunk_end = end
        # Add the last chunk
        print("Adding last chunk")
        chunks.append((current_chunk_start, current_chunk_end))
        # Export chunks
        for i, (start, end) in enumerate(chunks):
            chunk = audio[start:end]
            chunk_name = f"audio_files/{file.replace('.wav', '')}_chunk_{i:03d}.wav"
            chunk.export(chunk_name, format="wav")
            print(f"Exported chunk {i} from {file} ({len(chunk)/1000:.2f} seconds, {start/1000:.2f}s - {end/1000:.2f}s)")
        print(f"Finished splitting {file} into {len(chunks)} chunks")

# Transcribe the audio files
whisper_model = whisper.load_model("medium.en")
for file in os.listdir("audio_files"):
    if file.endswith(".wav") and "chunk" in file:
        print(f"Transcribing {file}")
        filename = file.replace(".wav", ".txt")
        result = whisper_model.transcribe(f"audio_files/{file}")
        with open(f"audio_files/{filename}", "w") as f:
            f.write(result["text"])
        print(f"Transcribed {filename}")

# Generate the dataset into a TSV file with the following columns: filename, text
tsv_file = "dataset.tsv"
with open(tsv_file, "w") as f:
    f.write("audio\\ttext\\n")
    for file in os.listdir("audio_files"):
        if file.endswith(".txt"):
            wav_file = file.replace(".txt", ".wav")
            wav_file = "audio_files/" + wav_file
            with open(f"audio_files/{file}", "r") as fr:
                text = fr.read()
            f.write(f"{wav_file}\\t{text}\\n")
            print(f"Added {wav_file} to {tsv_file}")
print(f"Finished generating {tsv_file}")

!pip install torchcodec

from unsloth import FastModel
from transformers import CsmForConditionalGeneration
import torch

model, processor = FastModel.from_pretrained(
    model_name = "unsloth/csm-1b",
    max_seq_length= 2048, # Choose any for long context!
    dtype = None, # Leave as None for auto-detection
    auto_model = CsmForConditionalGeneration,
    load_in_4bit = False, # Select True for 4bit - reduces memory usage
)

model = FastModel.get_peft_model(
    model,
    r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 32,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)

# Dataset Prep functions
from datasets import load_dataset, Audio, Dataset
import os
from transformers import AutoProcessor
from pathlib import Path

# Load model
processor = AutoProcessor.from_pretrained("unsloth/csm-1b")

# Load the TSV dataset
raw_ds = load_dataset("csv", data_files="dataset.tsv", delimiter="\\t", split="train")
# Resolve the path of the audio filename
raw_ds = raw_ds.map(lambda x: {"audio": str(Path(x["audio"]).resolve())})

# Getting the speaker id is important for multi-speaker models and speaker consistency
speaker_key = "source"
if "source" not in raw_ds.column_names and "speaker_id" not in raw_ds.column_names:
    print("Unsloth: No speaker found, adding default \\"source\\" of 0 for all examples")
    new_column = ["0"] * len(raw_ds)
    raw_ds = raw_ds.add_column("source", new_column)
elif "source" not in raw_ds.column_names and "speaker_id" in raw_ds.column_names:
    speaker_key = "speaker_id"

# Convert the audio path to an Audio file
target_sampling_rate = 24000
raw_ds = raw_ds.cast_column("audio", Audio(sampling_rate=target_sampling_rate))

def preprocess_example(example):
    conversation = [
        {
            "role": str(example[speaker_key]),
            "content": [
                {"type": "text", "text": example["text"]},
                {"type": "audio", "path": example["audio"]["array"]},
            ],
        }
    ]

    try:
        model_inputs = processor.apply_chat_template(
            conversation,
            tokenize=True,
            return_dict=True,
            output_labels=True,
            text_kwargs = {
                "padding": "max_length", # pad to the max_length
                "max_length": 256, # this should be the max length of audio
                "pad_to_multiple_of": 8,
                "padding_side": "right",
            },
            audio_kwargs = {
                "sampling_rate": 24_000,
                "max_length": 265633, # max input_values length of the whole dataset
                "padding": "max_length",
            },
            common_kwargs = {"return_tensors": "pt"},
        )
    except Exception as e:
        print(f"Error processing example with text '{example['text'][:50]}...': {e}")
        return None

    required_keys = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"]
    processed_example = {}
    # print(model_inputs.keys())
    for key in required_keys:
        if key not in model_inputs:
            print(f"Warning: Required key '{key}' not found in processor output for example.")
            return None

        value = model_inputs[key][0]
        processed_example[key] = value
    # Final check (optional but good)
    if not all(isinstance(processed_example[key], torch.Tensor) for key in processed_example):
         print(f"Error: Not all required keys are tensors in final processed example. Keys: {list(processed_example.keys())}")
         return None

    return processed_example

processed_ds = raw_ds.map(
    preprocess_example,
    remove_columns=raw_ds.column_names,
    desc="Preprocessing dataset",
)

ValueError: expected sequence of length X at dim 1 (got Y)

from transformers import TrainingArguments, Trainer
from unsloth import is_bfloat16_supported

# Define a Trainer
trainer = Trainer(
    model = model,
    train_dataset = processed_ds,
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        num_train_epochs=1,
        learning_rate = 2e-4,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01, # Turn this on if overfitting
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
        report_to = "none", # Use this for WandB etc
    ),
)

# Train the model
trainer_stats = trainer.train()

# Save to HF
model.save_pretrained_merged("SesameSherlockHolmes", processor, save_method = "merged_16bit",)
model.push_to_hub_merged("username/SesameSherlockHolmes", processor, save_method = "merged_16bit", token = "")

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel, ConfigDict
import torch
from transformers import CsmForConditionalGeneration, AutoProcessor
import os
import tempfile
from contextlib import asynccontextmanager

# Global variables for model and processor
model = None
processor = None
device = None

# Load the model on startup and clean up on shutdown
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load the model on startup and clean up on shutdown"""
    global model, processor, device

    model_id = "username/SesameSherlockHolmes"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"Loading model on {device}...")
    # load the model and the processor
    processor = AutoProcessor.from_pretrained(model_id)
    model = CsmForConditionalGeneration.from_pretrained(model_id, device_map=device)
    print("Model loaded successfully!")

    yield

    # Cleanup
    print("Shutting down...")

# Create the FastAPI app
app = FastAPI(
    title="Text-to-Speech API",
    description="API for generating audio from text using SesameSherlockHolmes model",
    version="1.0.0",
    lifespan=lifespan
)

# Define the input model
class TextInput(BaseModel):
    text: str
    speaker_id: int = 0

# Define the generate audio endpoint
@app.post("/generate")
async def generate_audio(input_data: TextInput):
    """
    Generate audio from text

    Args:
        input_data: TextInput object containing text and optional speaker_id

    Returns:
        Audio file (WAV format)
    """
    try:
        # Prepare the text with speaker ID
        formatted_text = f"[{input_data.speaker_id}]{input_data.text}"

        # Prepare the inputs
        inputs = processor(formatted_text, add_special_tokens=True).to(device)

        # Generate audio
        audio = model.generate(
            **inputs,
            output_audio=True,
            max_new_tokens=2000,  # 🔧 increase this for longer outputs
            do_sample=True,       # adds a bit of natural variation
            temperature=0.8       # controls expressiveness
        )

        # Create a temporary file to save the audio
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
            temp_filename = temp_file.name
            processor.save_audio(audio, temp_filename)

        # Return the audio file
        return FileResponse(
            temp_filename,
            media_type="audio/wav",
            filename="generated_audio.wav",
            background=None  # We'll handle cleanup manually
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating audio: {str(e)}")

# Run the app
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

import gradio as gr
import requests
import tempfile
import os
from pathlib import Path

# API Configuration
API_URL = "<http://localhost:8000>"

def generate_audio(text: str):
    """
    Generate audio from text using the TTS API

    Args:
        text: The text to convert to speech

    Returns:
        Path to the generated audio file or error message
    """
    if not text.strip():
        return None, "Please enter some text to generate audio."

    try:
        # Prepare the request
        payload = {
            "text": text,
            "speaker_id": 0
        }

        # Send request to API
        response = requests.post(
            f"{API_URL}/generate",
            json=payload,
            timeout=60
        )

        if response.status_code == 200:
            # Save the audio to a temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
                temp_file.write(response.content)
                temp_filename = temp_file.name

            return temp_filename, "Audio generated successfully!"
        else:
            error_msg = response.json().get("detail", "Unknown error")
            return None, f"Error: {error_msg}"

    except requests.exceptions.ConnectionError:
        return None, f"Could not connect to API at {API_URL}. Please make sure the API server is running."
    except requests.exceptions.Timeout:
        return None, "Request timed out. The text might be too long."
    except Exception as e:
        return None, f"Error: {str(e)}"

# Create Gradio Interface
with gr.Blocks(title="Text-to-Speech Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎙️ Text-to-Speech Generator
        Convert text to speech using the SesameSherlockHolmes model.
        """
    )

    # Define the input and output
    with gr.Row():
        # Define the input
        with gr.Column(scale=2):
            # Define the text input
            text_input = gr.Textbox(
                label="Enter Text",
                placeholder="Type the text you want to convert to speech...",
                lines=5,
                max_lines=10
            )

            # Define the buttons
            with gr.Row():
                generate_btn = gr.Button("🎵 Generate Audio", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Clear", size="lg")

        # Define the output
        with gr.Column(scale=2):
            # Define the status output
            status_output = gr.Textbox(
                label="Status",
                interactive=False,
                container=True
            )

            # Define the audio output
            audio_output = gr.Audio(
                label="Generated Audio",
                type="filepath",
                interactive=False
            )

    # Define the generate button
    generate_btn.click(
        fn=generate_audio,
        inputs=[text_input],
        outputs=[audio_output, status_output]
    )

    # Define the clear button
    clear_btn.click(
        fn=lambda: ("", None, ""),
        outputs=[text_input, audio_output, status_output]
    )

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )

FROM python:3.12-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    build-essential \\
    curl \\
    && rm -rf /var/lib/apt/lists/*

# Copy requirements and install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy the API and Gradio application
COPY api.py .
COPY app.py .

# Run the API and Gradio application
CMD python api.py & python app.py

echo "# SesameUnsloth" >> README.md
git init
git add .
git commit -m "First Commit"
git branch -M main
git remote add origin [Your GitHub repository URL]
git push -u origin main