backend/synthesis_service.py

from pathlib import Path
from typing import Dict, Iterable, Optional

from backend.magpie_adapter import MagpieAdapter
from backend.modal_client import ModalSynthesisClient
from backend.omnivoice_adapter import OmniVoiceAdapter
from backend.render_pipeline import RenderPipeline
from backend.synthesis_catalog import LOCAL_BACKEND, MODAL_BACKEND, MAGPIE_MODEL, OMNIVOICE_MODEL
from backend.types import RenderJob, VoiceConfig


class SynthesisService:
    def __init__(
        self,
        *,
        session_root: Path,
        local_synthesizers: Optional[Dict[str, object]] = None,
        modal_client: Optional[ModalSynthesisClient] = None,
    ) -> None:
        self.session_root = Path(session_root)
        self.session_root.mkdir(parents=True, exist_ok=True)
        self.local_synthesizers = local_synthesizers or {
            OMNIVOICE_MODEL: OmniVoiceAdapter(),
            MAGPIE_MODEL: MagpieAdapter(),
        }
        self.modal_client = modal_client or ModalSynthesisClient.from_env()
        self.pipeline = RenderPipeline(session_root=self.session_root)
        self._active_backends: Dict[str, str] = {}
        self._modal_job_ids: Dict[str, str] = {}
        self._modal_status: Dict[str, str] = {}

    def get_job(self, session_id: str) -> RenderJob:
        if self._active_backends.get(session_id) == MODAL_BACKEND:
            return RenderJob(
                session_id=session_id,
                status=self._modal_status.get(session_id, "idle"),
            )
        return self.pipeline.get_job(session_id)

    def generate_preview(
        self,
        *,
        text: str,
        output_path: Path,
        voice_config: VoiceConfig,
        diffusion_steps: int,
        speed: float,
    ) -> Dict[str, object]:
        if voice_config.backend == MODAL_BACKEND:
            return self.modal_client.generate_preview(
                text=text,
                output_path=output_path,
                voice_config=voice_config,
                diffusion_steps=diffusion_steps,
                speed=speed,
            )
        synthesizer = self._local_synthesizer(voice_config.model)
        result = synthesizer.synthesize(
            text=text,
            output_path=output_path,
            voice_config=voice_config,
            diffusion_steps=diffusion_steps,
            speed=speed,
        )
        result.setdefault("backend", LOCAL_BACKEND)
        result.setdefault("model", voice_config.model)
        return result

    def render(
        self,
        *,
        session_id: str,
        book: Dict[str, object],
        chapters: list[Dict[str, object]],
        voice_config: Dict[str, object] | VoiceConfig,
        diffusion_steps: int,
        speed: float,
    ) -> Iterable[Dict[str, object]]:
        voice = voice_config if isinstance(voice_config, VoiceConfig) else VoiceConfig.from_dict(voice_config)
        self._active_backends[session_id] = voice.backend
        if voice.backend == MODAL_BACKEND:
            yield from self._render_modal(
                session_id=session_id,
                book=book,
                chapters=chapters,
                voice_config=voice,
                diffusion_steps=diffusion_steps,
                speed=speed,
            )
            return

        synthesizer = self._local_synthesizer(voice.model)
        yield from self.pipeline.render(
            session_id=session_id,
            book=book,
            chapters=chapters,
            voice_config=voice.to_dict(),
            diffusion_steps=diffusion_steps,
            speed=speed,
            synthesizer=synthesizer,
        )

    def pause(self, session_id: str) -> Dict[str, object]:
        if self._active_backends.get(session_id) == MODAL_BACKEND:
            raise ValueError("Pause is only supported for local renders")
        return self.pipeline.pause(session_id)

    def resume(self, session_id: str) -> Dict[str, object]:
        if self._active_backends.get(session_id) == MODAL_BACKEND:
            raise ValueError("Resume is only supported for local renders")
        return self.pipeline.resume(session_id)

    def cancel(self, session_id: str) -> Dict[str, object]:
        if self._active_backends.get(session_id) == MODAL_BACKEND:
            job_id = self._modal_job_ids.get(session_id)
            if not job_id:
                return {"type": "cancelled", "session_id": session_id, "backend": MODAL_BACKEND}
            self._modal_status[session_id] = "cancelled"
            result = self.modal_client.cancel_render(job_id)
            return {**result, "session_id": session_id}
        self.pipeline.cancel(session_id)
        return {"type": "cancelled", "session_id": session_id, "backend": LOCAL_BACKEND}

    def _render_modal(
        self,
        *,
        session_id: str,
        book: Dict[str, object],
        chapters: list[Dict[str, object]],
        voice_config: VoiceConfig,
        diffusion_steps: int,
        speed: float,
    ) -> Iterable[Dict[str, object]]:
        render_dir = self.session_root / session_id / "renders"
        render_dir.mkdir(parents=True, exist_ok=True)
        self._modal_status[session_id] = "pending"
        job_id = self.modal_client.submit_render(
            session_id=session_id,
            book=book,
            chapters=chapters,
            voice_config=voice_config,
            diffusion_steps=diffusion_steps,
            speed=speed,
        )
        self._modal_job_ids[session_id] = job_id
        self._modal_status[session_id] = "running"
        for event in self.modal_client.render(
            session_id=session_id,
            job_id=job_id,
            render_dir=render_dir,
            voice_config=voice_config,
        ):
            event_type = str(event.get("type", "running"))
            if event_type not in {"log"}:
                self._modal_status[session_id] = event_type
            yield event
        if self._modal_status.get(session_id) not in {"failed", "cancelled"}:
            self._modal_status[session_id] = "completed"

    def _local_synthesizer(self, model: str):
        synthesizer = self.local_synthesizers.get(model)
        if synthesizer is None:
            raise ValueError(f"Unsupported synthesis model: {model}")
        return synthesizer