core/workflow.py

"""Full media workflow orchestration outside the HTTP layer."""

from __future__ import annotations

from pathlib import Path

from config import Settings
from core.executor import InferenceCommand
from core.manager import TaskManager
from core.schemas import WorkflowRequest
from utils.files import OutputManager


class WorkflowService:
    """Run image, voice, music, video, and subtitle generation in order."""

    def __init__(
        self, settings: Settings, tasks: TaskManager, outputs: OutputManager
    ) -> None:
        self.settings = settings
        self.tasks = tasks
        self.outputs = outputs

    async def run(self, payload: WorkflowRequest, request_id: str) -> dict[str, Path]:
        """Queue the entire pipeline as one non-interleavable operation."""
        width = payload.width if payload.width is not None else self.settings.image_width
        height = payload.height if payload.height is not None else self.settings.image_height
        voice = payload.voice or self.settings.kokoro_default_voice
        targets = {
            "image": self.outputs.allocate("images"),
            "voice": self.outputs.allocate("audio"),
            "music": self.outputs.allocate("music"),
            "video": self.outputs.allocate("videos"),
            "subtitle": self.outputs.allocate("subtitles"),
        }

        def execute() -> dict[str, Path]:
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="flux",
                    method_name="generate",
                    arguments={
                        "prompt": payload.image_prompt,
                        "output_path": targets["image"],
                        "width": width,
                        "height": height,
                        "steps": self.settings.image_steps,
                        "seed": payload.seed,
                        "guidance_scale": self.settings.flux_guidance_scale,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_flux_duration,
                ),
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="kokoro",
                    method_name="synthesize",
                    arguments={
                        "text": payload.script,
                        "voice": voice,
                        "speed": 1.0,
                        "output_path": targets["voice"],
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_kokoro_duration,
                ),
            )
            words = len(payload.script.split())
            music_duration = min(30.0, max(5.0, words / 2.5))
            music_prompt = payload.music_prompt or (
                f"cinematic instrumental background score for {payload.title}, no vocals"
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="musicgen",
                    method_name="generate",
                    arguments={
                        "prompt": music_prompt,
                        "duration": music_duration,
                        "guidance_scale": self.settings.music_guidance_scale,
                        "seed": payload.seed,
                        "output_path": targets["music"],
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_musicgen_duration,
                ),
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="wan",
                    method_name="generate",
                    arguments={
                        "image_path": targets["image"],
                        "prompt": payload.video_prompt,
                        "negative_prompt": "low quality, distorted, static",
                        "output_path": targets["video"],
                        "steps": self.settings.wan_steps,
                        "frames": self.settings.video_frames,
                        "fps": self.settings.video_fps,
                        "seed": payload.seed,
                        "guidance_scale": self.settings.wan_guidance_scale,
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_wan_duration,
                    gpu_size="xlarge",
                ),
            )
            self.tasks.invoke_direct(
                InferenceCommand(
                    model_name="whisper",
                    method_name="transcribe",
                    arguments={
                        "source": targets["voice"],
                        "subtitle_path": targets["subtitle"],
                        "language": None,
                        "task": "transcribe",
                    },
                    request_id=request_id,
                    duration_seconds=self.settings.zerogpu_whisper_duration,
                ),
            )
            return targets

        try:
            return await self.tasks.run_exclusive(request_id, "workflow", execute)
        except Exception:
            for target in targets.values():
                self.outputs.remove(target)
            raise