Initialize database and add storage directories; include audiobook routes

app/api/audiobook_routes.py

@@ -0,0 +1,212 @@
+"""
+Audiobook creation routes for the CSM-1B TTS API.
+"""
+import os
+import uuid
+import logging
+from datetime import datetime
+from typing import Optional, List
+from fastapi import APIRouter, Request, HTTPException, BackgroundTasks, UploadFile, File, Form, Depends
+from fastapi.responses import FileResponse, JSONResponse
+from sqlalchemy.orm import Session
+from app.models.database import Audiobook, AudiobookStatus, AudiobookChunk
+from app.services.storage import storage
+from app.db import get_db
+import torchaudio
+
+# Set up logging
+logger = logging.getLogger(__name__)
+router = APIRouter(prefix="/audiobook", tags=["Audiobook"])
+
+async def process_audiobook(
+    request: Request,
+    book_id: str,
+    text_content: str,
+    voice_id: int,
+    db: Session
+):
+    """Process audiobook in the background."""
+    try:
+        # Get the book from database
+        book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
+        if not book:
+            logger.error(f"Book {book_id} not found")
+            return False
+
+        # Update status to processing
+        book.status = AudiobookStatus.PROCESSING
+        db.commit()
+
+        logger.info(f"Starting processing for audiobook {book_id}")
+
+        # Get the generator from app state
+        generator = request.app.state.generator
+        if generator is None:
+            raise Exception("TTS model not available")
+
+        # Get voice info
+        voice_info = request.app.state.get_voice_info(voice_id)
+        if not voice_info:
+            raise Exception(f"Voice ID {voice_id} not found")
+
+        # Generate audio for the entire text
+        logger.info(f"Generating audio for entire text of book {book_id}")
+        audio = generator.generate(
+            text=text_content,
+            speaker=voice_info["speaker_id"],
+            max_audio_length_ms=min(300000, len(text_content) * 80)  # Big text = big audio
+        )
+
+        if audio is None:
+            raise Exception("Failed to generate audio")
+
+        # Save the audio using storage service
+        audio_to_save = audio.unsqueeze(0).cpu() if len(audio.shape) == 1 else audio.cpu()
+        audio_bytes = audio_to_save.numpy().tobytes()
+        audio_path = await storage.save_audio_file(book_id, audio_bytes)
+
+        # Update book status in database
+        book.status = AudiobookStatus.COMPLETED
+        book.audio_file_path = audio_path
+        db.commit()
+
+        logger.info(f"Successfully created audiobook {book_id}")
+        return True
+
+    except Exception as e:
+        logger.error(f"Error processing audiobook {book_id}: {e}")
+        
+        # Update status to failed in database
+        book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
+        if book:
+            book.status = AudiobookStatus.FAILED
+            book.error_message = str(e)
+            db.commit()
+        
+        return False
+
+@router.post("/")
+async def create_audiobook(
+    request: Request,
+    background_tasks: BackgroundTasks,
+    title: str = Form(...),
+    author: str = Form(...),
+    voice_id: int = Form(0),
+    text_file: Optional[UploadFile] = File(None),
+    text_content: Optional[str] = Form(None),
+    db: Session = Depends(get_db)
+):
+    """Create a new audiobook from text."""
+    try:
+        # Validate input
+        if not text_file and not text_content:
+            raise HTTPException(status_code=400, detail="Either text_file or text_content is required")
+
+        # Generate unique ID
+        book_id = str(uuid.uuid4())
+
+        # Handle text content
+        if text_file:
+            text_file_path = await storage.save_text_file(book_id, text_file)
+            with open(text_file_path, "r", encoding="utf-8") as f:
+                text_content = f.read()
+        else:
+            text_file_path = await storage.save_text_content(book_id, text_content)
+
+        # Create book in database
+        book = Audiobook(
+            id=book_id,
+            title=title,
+            author=author,
+            voice_id=voice_id,
+            status=AudiobookStatus.PENDING,
+            text_file_path=text_file_path,
+            text_content=text_content if len(text_content) <= 10000 else None  # Store small texts directly
+        )
+        db.add(book)
+        db.commit()
+
+        # Process in background
+        background_tasks.add_task(process_audiobook, request, book_id, text_content, voice_id, db)
+
+        return JSONResponse(content={"message": "Audiobook creation started", "book_id": book_id})
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error creating audiobook: {str(e)}")
+
+@router.get("/{book_id}")
+async def get_audiobook(book_id: str, db: Session = Depends(get_db)):
+    """Get audiobook information."""
+    book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
+    if not book:
+        raise HTTPException(status_code=404, detail="Audiobook not found")
+
+    return {
+        "id": book.id,
+        "title": book.title,
+        "author": book.author,
+        "voice_id": book.voice_id,
+        "status": book.status.value,
+        "created_at": book.created_at.isoformat(),
+        "updated_at": book.updated_at.isoformat(),
+        "error_message": book.error_message
+    }
+
+@router.get("/{book_id}/audio")
+async def get_audiobook_audio(book_id: str, db: Session = Depends(get_db)):
+    """Get the audiobook audio file."""
+    book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
+    if not book:
+        raise HTTPException(status_code=404, detail="Audiobook not found")
+
+    if book.status != AudiobookStatus.COMPLETED or not book.audio_file_path:
+        raise HTTPException(status_code=400, detail="Audiobook is not yet completed")
+
+    audio_path = await storage.get_audio_file(book_id)
+    if not audio_path:
+        raise HTTPException(status_code=404, detail="Audio file not found")
+
+    return FileResponse(
+        str(audio_path),
+        media_type="audio/wav",
+        filename=f"{book.title}.wav"
+    )
+
+@router.get("/")
+async def get_audiobooks(db: Session = Depends(get_db)):
+    """Get all audiobooks."""
+    books = db.query(Audiobook).order_by(Audiobook.created_at.desc()).all()
+    return {
+        "audiobooks": [
+            {
+                "id": book.id,
+                "title": book.title,
+                "author": book.author,
+                "voice_id": book.voice_id,
+                "status": book.status.value,
+                "created_at": book.created_at.isoformat(),
+                "updated_at": book.updated_at.isoformat(),
+                "error_message": book.error_message
+            }
+            for book in books
+        ]
+    }
+
+@router.delete("/{book_id}")
+async def delete_audiobook(book_id: str, db: Session = Depends(get_db)):
+    """Delete an audiobook."""
+    book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
+    if not book:
+        raise HTTPException(status_code=404, detail="Audiobook not found")
+
+    try:
+        # Delete associated files
+        await storage.delete_book_files(book_id)
+
+        # Delete from database
+        db.delete(book)
+        db.commit()
+
+        return {"message": "Audiobook deleted successfully"}
+    except Exception as e:
+        db.rollback()
+        raise HTTPException(status_code=500, detail=f"Error deleting audiobook: {str(e)}") 

app/db.py

@@ -0,0 +1,39 @@
+"""Database connection and session management."""
+import os
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, Session
+from contextlib import contextmanager
+from app.models.database import Base
+
+# Get database URL from environment or use SQLite as default
+DATABASE_URL = os.environ.get(
+    "DATABASE_URL",
+    "sqlite:///app/audiobooks.db"
+)
+
+# Create engine
+engine = create_engine(
+    DATABASE_URL,
+    echo=False,  # Set to True for SQL logging
+    pool_pre_ping=True,  # Enable connection health checks
+)
+
+# Create session factory
+SessionLocal = sessionmaker(
+    bind=engine,
+    autocommit=False,
+    autoflush=False,
+)
+
+@contextmanager
+def get_db() -> Session:
+    """Get database session."""
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+def init_db():
+    """Initialize database."""
+    Base.metadata.create_all(bind=engine) 

app/main.py

@@ -53,6 +53,12 @@ async def lifespan(app: FastAPI):
     app.state.generator = None  # Will be populated later if model loads
     app.state.logger = logger  # Make logger available to routes
     
+    # Initialize database
+    from app.db import init_db
+    logger.info("Initializing database...")
+    init_db()
+    logger.info("Database initialized")
+    
     # Create necessary directories - use persistent locations
     APP_DIR = "/app"
     os.makedirs(os.path.join(APP_DIR, "models"), exist_ok=True)
@@ -63,6 +69,9 @@ async def lifespan(app: FastAPI):
     os.makedirs(os.path.join(APP_DIR, "cloned_voices"), exist_ok=True)
     os.makedirs(os.path.join(APP_DIR, "audio_cache"), exist_ok=True)
     os.makedirs(os.path.join(APP_DIR, "static"), exist_ok=True)
+    os.makedirs(os.path.join(APP_DIR, "storage/audio"), exist_ok=True)  # For audio files
+    os.makedirs(os.path.join(APP_DIR, "storage/text"), exist_ok=True)   # For text files
+    os.makedirs(os.path.join(APP_DIR, "audiobooks"), exist_ok=True)  # Add audiobooks directory
     
     # Set tokenizer cache
     try:
@@ -520,6 +529,11 @@ from app.api.streaming import router as streaming_router
 app.include_router(streaming_router, prefix="/api/v1")
 app.include_router(streaming_router, prefix="/v1")
 
+# Add audiobook routes
+from app.api.audiobook_routes import router as audiobook_router
+app.include_router(audiobook_router, prefix="/api/v1")
+app.include_router(audiobook_router, prefix="/v1")
+
 # Middleware for request timing
 @app.middleware("http")
 async def add_process_time_header(request: Request, call_next):

app/services/storage.py

@@ -0,0 +1,81 @@
+"""Storage service for managing audio files."""
+import os
+import shutil
+from pathlib import Path
+from typing import Optional, BinaryIO
+from fastapi import UploadFile
+import aiofiles
+
+class StorageService:
+    """Service for managing file storage."""
+
+    def __init__(self, base_path: str = "/app/storage"):
+        """Initialize storage service."""
+        self.base_path = Path(base_path)
+        self.audio_path = self.base_path / "audio"
+        self.text_path = self.base_path / "text"
+        
+        # Create directories
+        self.audio_path.mkdir(parents=True, exist_ok=True)
+        self.text_path.mkdir(parents=True, exist_ok=True)
+
+    async def save_audio_file(self, book_id: str, audio_data: bytes) -> str:
+        """Save audio file to storage."""
+        file_path = self.audio_path / f"{book_id}.wav"
+        async with aiofiles.open(file_path, "wb") as f:
+            await f.write(audio_data)
+        return str(file_path)
+
+    async def save_text_file(self, book_id: str, text_file: UploadFile) -> str:
+        """Save text file to storage."""
+        file_path = self.text_path / f"{book_id}.txt"
+        async with aiofiles.open(file_path, "wb") as f:
+            content = await text_file.read()
+            await f.write(content)
+        return str(file_path)
+
+    async def save_text_content(self, book_id: str, text_content: str) -> str:
+        """Save text content to a file."""
+        file_path = self.text_path / f"{book_id}.txt"
+        async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
+            await f.write(text_content)
+        return str(file_path)
+
+    async def get_audio_file(self, book_id: str) -> Optional[Path]:
+        """Get audio file path."""
+        file_path = self.audio_path / f"{book_id}.wav"
+        return file_path if file_path.exists() else None
+
+    async def get_text_file(self, book_id: str) -> Optional[Path]:
+        """Get text file path."""
+        file_path = self.text_path / f"{book_id}.txt"
+        return file_path if file_path.exists() else None
+
+    async def delete_book_files(self, book_id: str):
+        """Delete all files associated with a book."""
+        # Delete audio file
+        audio_file = self.audio_path / f"{book_id}.wav"
+        if audio_file.exists():
+            audio_file.unlink()
+
+        # Delete text file
+        text_file = self.text_path / f"{book_id}.txt"
+        if text_file.exists():
+            text_file.unlink()
+
+    def cleanup_orphaned_files(self, valid_book_ids: set[str]):
+        """Clean up files that don't belong to any book."""
+        # Clean up audio files
+        for file_path in self.audio_path.glob("*.wav"):
+            book_id = file_path.stem
+            if book_id not in valid_book_ids:
+                file_path.unlink()
+
+        # Clean up text files
+        for file_path in self.text_path.glob("*.txt"):
+            book_id = file_path.stem
+            if book_id not in valid_book_ids:
+                file_path.unlink()
+
+# Create global instance
+storage = StorageService()