72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
from rich.console import Console
|
|
from dubbing.models import ProjectPaths, StepStatus, Language
|
|
from dubbing.steps import (
|
|
ExtractAudioStep,
|
|
ASRStep,
|
|
TranslateStep,
|
|
TTSStep,
|
|
MixStep,
|
|
FinalizeStep,
|
|
)
|
|
|
|
console = Console()
|
|
|
|
|
|
class Pipeline:
|
|
def __init__(
|
|
self,
|
|
project_name: str,
|
|
model_name: str = "gemini-2.0-flash-lite",
|
|
language: Language = Language.RU,
|
|
):
|
|
from dubbing.config import settings
|
|
|
|
project_dir = settings.projects_dir / project_name
|
|
self.paths = ProjectPaths.from_project(project_dir, language)
|
|
self.model_name = model_name
|
|
self.language = language
|
|
self.steps = self._create_steps()
|
|
|
|
def _create_steps(self) -> list:
|
|
return [
|
|
ExtractAudioStep(self.paths),
|
|
ASRStep(self.paths),
|
|
TranslateStep(self.paths, self.model_name, self.language),
|
|
TTSStep(self.paths, self.language),
|
|
MixStep(self.paths),
|
|
FinalizeStep(self.paths),
|
|
]
|
|
|
|
def get_cache_status(self) -> list[StepStatus]:
|
|
return [
|
|
StepStatus(
|
|
name=step.name,
|
|
cached=step.is_cached(),
|
|
output_exists=step.is_cached(),
|
|
)
|
|
for step in self.steps
|
|
]
|
|
|
|
def clean_from_step(self, step_index: int) -> None:
|
|
for step in self.steps[step_index:]:
|
|
step.clean()
|
|
|
|
async def run(self, rebuild_from: int = -1) -> None:
|
|
if rebuild_from >= 0:
|
|
console.print(
|
|
f"[yellow]Cleaning from step: {self.steps[rebuild_from].name}[/]"
|
|
)
|
|
self.clean_from_step(rebuild_from)
|
|
|
|
for i, step in enumerate(self.steps):
|
|
if step.is_cached():
|
|
console.print(f"[dim]Skipping {step.name} (cached)[/]")
|
|
continue
|
|
|
|
console.print(
|
|
f"\n[bold cyan]Step {i + 1}/{len(self.steps)}: {step.name}[/]"
|
|
)
|
|
await step.run()
|
|
|
|
console.print("\n[bold green]Pipeline completed![/]")
|