146 lines
7.1 KiB
Python
146 lines
7.1 KiB
Python
"""Pipeline state machine - step definitions, dependency graph, and ordering."""
|
|
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
# Step definitions per mode: (step_name, dependencies, display_name)
|
|
MODE_STEPS: Dict[str, List[Tuple[str, List[str], str]]] = {
|
|
"audio_lyrics": [ # Mode A
|
|
("audio_preparing", [], "音频准备"),
|
|
("demucs_separating", ["audio_preparing"], "人声分离"),
|
|
("lyric_generating", ["demucs_separating"], "歌词生成"),
|
|
("lyric_evaluating", ["lyric_generating"], "歌词评估"),
|
|
("music_generating", ["lyric_evaluating"], "音乐生成"),
|
|
("music_polling", ["music_generating"], "音乐轮询"),
|
|
("lyric_calibrating", ["music_polling", "demucs_separating"], "歌词校准"),
|
|
("subtitle_rendering", ["lyric_calibrating"], "字幕渲染"),
|
|
("subtitle_exporting", ["subtitle_rendering"], "字幕导出"),
|
|
("character_designing", ["lyric_calibrating"], "角色设计"),
|
|
("character_image_generating", ["character_designing"], "角色图生成"),
|
|
("storyboard_generating", ["character_designing", "music_polling"], "分镜剧本"),
|
|
("scene_video_generating", ["storyboard_generating", "character_image_generating"], "分镜视频生成"),
|
|
("scene_video_evaluating", ["scene_video_generating"], "分镜视频评估"),
|
|
("scene_video_concatenating", ["scene_video_evaluating"], "分镜视频拼接"),
|
|
("ktv_synthesizing", ["scene_video_concatenating", "subtitle_rendering", "music_polling"], "KTV合成"),
|
|
],
|
|
"video_lyrics": [ # Mode B
|
|
("video_preparing", [], "视频准备"),
|
|
("demucs_separating", ["video_preparing"], "人声分离"),
|
|
("lyric_generating", ["demucs_separating"], "歌词生成"),
|
|
("lyric_evaluating", ["lyric_generating"], "歌词评估"),
|
|
("music_generating", ["lyric_evaluating"], "音乐生成"),
|
|
("music_polling", ["music_generating"], "音乐轮询"),
|
|
("lyric_calibrating", ["music_polling", "demucs_separating"], "歌词校准"),
|
|
("subtitle_rendering", ["lyric_calibrating"], "字幕渲染"),
|
|
("subtitle_exporting", ["subtitle_rendering"], "字幕导出"),
|
|
("character_designing", ["lyric_calibrating"], "角色设计"),
|
|
("character_image_generating", ["character_designing"], "角色图生成"),
|
|
("storyboard_generating", ["character_designing", "music_polling"], "分镜剧本"),
|
|
("scene_video_generating", ["storyboard_generating", "character_image_generating"], "分镜视频生成"),
|
|
("scene_video_evaluating", ["scene_video_generating"], "分镜视频评估"),
|
|
("scene_video_concatenating", ["scene_video_evaluating"], "分镜视频拼接"),
|
|
("ktv_synthesizing", ["scene_video_concatenating", "subtitle_rendering", "music_polling"], "KTV合成"),
|
|
],
|
|
"lyrics_only": [ # Mode C
|
|
("lyric_generating", [], "歌词生成"),
|
|
("lyric_evaluating", ["lyric_generating"], "歌词评估"),
|
|
("music_generating", ["lyric_evaluating"], "音乐生成"),
|
|
("music_polling", ["music_generating"], "音乐轮询"),
|
|
("lyric_calibrating", ["music_polling"], "歌词校准"),
|
|
("subtitle_rendering", ["lyric_calibrating"], "字幕渲染"),
|
|
("subtitle_exporting", ["subtitle_rendering"], "字幕导出"),
|
|
("character_designing", ["lyric_calibrating"], "角色设计"),
|
|
("character_image_generating", ["character_designing"], "角色图生成"),
|
|
("storyboard_generating", ["character_designing", "music_polling"], "分镜剧本"),
|
|
("scene_video_generating", ["storyboard_generating", "character_image_generating"], "分镜视频生成"),
|
|
("scene_video_evaluating", ["scene_video_generating"], "分镜视频评估"),
|
|
("scene_video_concatenating", ["scene_video_evaluating"], "分镜视频拼接"),
|
|
("ktv_synthesizing", ["scene_video_concatenating", "subtitle_rendering", "music_polling"], "KTV合成"),
|
|
],
|
|
}
|
|
|
|
# Step states
|
|
STATE_PENDING = "pending"
|
|
STATE_RUNNING = "running"
|
|
STATE_COMPLETED = "completed"
|
|
STATE_FAILED = "failed"
|
|
STATE_SKIPPED = "skipped"
|
|
|
|
# Pipeline states
|
|
PIPELINE_SUBMITTED = "submitted"
|
|
PIPELINE_RUNNING = "running"
|
|
PIPELINE_COMPLETED = "completed"
|
|
PIPELINE_FAILED = "failed"
|
|
PIPELINE_PAUSED = "paused" # waiting for user modification
|
|
|
|
|
|
def get_step_graph(mode: str) -> List[Tuple[str, List[str], str]]:
|
|
"""Get step definitions for a mode. Returns [(name, deps, display_name), ...]"""
|
|
if mode not in MODE_STEPS:
|
|
raise ValueError(f"Unknown mode: {mode}. Available: {list(MODE_STEPS.keys())}")
|
|
return MODE_STEPS[mode]
|
|
|
|
|
|
def build_dependency_map(mode: str) -> Dict[str, dict]:
|
|
"""Build a dependency map for a mode.
|
|
Returns: {step_name: {"deps": [...], "dependents": [...], "display_name": "...", "order": int}}
|
|
"""
|
|
steps = get_step_graph(mode)
|
|
dep_map = {}
|
|
for i, (name, deps, display) in enumerate(steps):
|
|
dep_map[name] = {
|
|
"deps": list(deps),
|
|
"dependents": [],
|
|
"display_name": display,
|
|
"order": i + 1,
|
|
}
|
|
# Build reverse mapping
|
|
for name, info in dep_map.items():
|
|
for dep in info["deps"]:
|
|
if dep in dep_map:
|
|
dep_map[dep]["dependents"].append(name)
|
|
return dep_map
|
|
|
|
|
|
def get_cascade_rerun_steps(mode: str, from_step: str) -> List[str]:
|
|
"""Get all steps that need to be rerun when a step is modified.
|
|
BFS from the modified step through dependents.
|
|
Returns ordered list of step names.
|
|
"""
|
|
dep_map = build_dependency_map(mode)
|
|
if from_step not in dep_map:
|
|
return []
|
|
|
|
visited = set()
|
|
queue = [from_step]
|
|
result = []
|
|
|
|
while queue:
|
|
current = queue.pop(0)
|
|
if current in visited:
|
|
continue
|
|
visited.add(current)
|
|
result.append(current)
|
|
for dep in dep_map.get(current, {}).get("dependents", []):
|
|
if dep not in visited:
|
|
queue.append(dep)
|
|
|
|
# Sort by order
|
|
result.sort(key=lambda s: dep_map.get(s, {}).get("order", 999))
|
|
return result
|
|
|
|
|
|
def get_rerun_from_next(mode: str, from_step: str) -> List[str]:
|
|
"""When output is modified, rerun from the NEXT steps (dependents only, not the step itself)."""
|
|
dep_map = build_dependency_map(mode)
|
|
if from_step not in dep_map:
|
|
return []
|
|
|
|
direct_dependents = dep_map[from_step]["dependents"]
|
|
all_steps = set()
|
|
for d in direct_dependents:
|
|
all_steps.update(get_cascade_rerun_steps(mode, d))
|
|
|
|
result = list(all_steps)
|
|
result.sort(key=lambda s: dep_map.get(s, {}).get("order", 999))
|
|
return result
|