540 lines
20 KiB
Python
540 lines
20 KiB
Python
# -*- coding:utf-8 -*-
|
||
"""
|
||
KTV Pipeline Worker - 状态机驱动的完整歌曲制作流程
|
||
|
||
State machine:
|
||
submitted -> lyric_generating -> lyric_evaluating -> lyric_done
|
||
-> music_generating -> music_separating -> music_aligning -> music_calibrating -> music_evaluating -> music_done
|
||
-> mv_story_generating -> mv_designing -> mv_frames_generating -> mv_video_generating -> mv_video_evaluating -> mv_merging -> mv_done
|
||
-> subtitle_rendering -> ktv_synthesizing -> completed
|
||
|
||
Any state can transition to 'failed' on error.
|
||
Evaluation states can loop back to the generating state if threshold not met.
|
||
"""
|
||
import json
|
||
import os
|
||
import asyncio
|
||
import time
|
||
import aiohttp
|
||
|
||
SERVICES = {
|
||
'align': 'http://127.0.0.1:8080/api/align',
|
||
'fastwhisper': 'http://127.0.0.1:9925/api/asr',
|
||
'fastwhisper_status': 'http://127.0.0.1:9925/api/status',
|
||
'songrate': 'http://127.0.0.1:8900/api/evaluate',
|
||
'demucs': 'http://127.0.0.1:9080/api/demucs',
|
||
'lyric_calibrate': 'http://127.0.0.1:9080/api/lyric_calibrate',
|
||
'merge': 'http://127.0.0.1:9080/api/merge',
|
||
'subtitle': 'http://127.0.0.1:9080/api/subtitle',
|
||
'ktv': 'http://127.0.0.1:9080/api/ktv',
|
||
}
|
||
|
||
LLM_API = os.environ.get('LLM_API_BASE', 'https://token.opencomputing.cn/llmage/v1')
|
||
WORK_DIR = '/tmp/ktv_pipelines'
|
||
|
||
|
||
async def update_state(redis, pipeline_id, state, **kwargs):
|
||
"""Update pipeline state in Redis"""
|
||
data = await redis.get(f'pipeline:{pipeline_id}')
|
||
if data:
|
||
pipeline = json.loads(data)
|
||
pipeline['state'] = state
|
||
pipeline['updated_at'] = time.time()
|
||
if 'artifacts' in kwargs:
|
||
pipeline['artifacts'].update(kwargs['artifacts'])
|
||
if 'error' in kwargs:
|
||
pipeline['errors'].append({'state': state, 'error': kwargs['error'], 'time': time.time()})
|
||
pipeline.update({k: v for k, v in kwargs.items() if k not in ('artifacts', 'error')})
|
||
await redis.set(f'pipeline:{pipeline_id}', json.dumps(pipeline, ensure_ascii=False), ex=86400)
|
||
|
||
|
||
async def call_llm(session, prompt, model='qwen3-235b-a22b', temperature=0.7, max_tokens=4096):
|
||
"""Call Sage LLM API"""
|
||
api_key = os.environ.get('LLM_API_KEY', '')
|
||
if not api_key:
|
||
try:
|
||
from ahserver.serverenv import ServerEnv
|
||
env = ServerEnv()
|
||
api_key = getattr(env, 'llm_api_key', '') or ''
|
||
except:
|
||
pass
|
||
|
||
headers = {
|
||
'Authorization': f'Bearer {api_key}',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
payload = {
|
||
'model': model,
|
||
'catelogid': 't2t',
|
||
'messages': [{'role': 'user', 'content': prompt}],
|
||
'temperature': temperature,
|
||
'max_tokens': max_tokens
|
||
}
|
||
async with session.post(f'{LLM_API}/chat/completions', json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=120)) as resp:
|
||
data = await resp.json()
|
||
return data.get('choices', [{}])[0].get('message', {}).get('content', '')
|
||
|
||
|
||
async def call_service(session, url, data, timeout=300):
|
||
"""Call an internal service"""
|
||
async with session.post(url, json=data, timeout=aiohttp.ClientTimeout(total=timeout)) as resp:
|
||
return await resp.json()
|
||
|
||
|
||
async def step_lyric_generate(pipeline, session):
|
||
"""Step 1: Generate lyrics using LLM"""
|
||
prompt = f"""你是一位专业的华语歌词创作者。
|
||
|
||
请根据以下描述创作一首完整的歌词：
|
||
|
||
描述：{pipeline['description']}
|
||
风格：{pipeline.get('scene', 'pop')}
|
||
|
||
要求：
|
||
1. 歌词结构完整：包含主歌(verse)、副歌(chorus)、桥段(bridge)
|
||
2. 每行歌词节奏感强，适合演唱
|
||
3. 标注段落类型，如 [Verse1], [Chorus], [Bridge] 等
|
||
4. 总共16-24行歌词
|
||
5. 注意押韵和情感表达
|
||
|
||
请直接输出歌词（不要markdown标记）：
|
||
"""
|
||
lyrics = await call_llm(session, prompt)
|
||
return {'lyrics': lyrics.strip()}
|
||
|
||
|
||
async def step_lyric_evaluate(pipeline, session):
|
||
"""Step 2: Evaluate lyrics (call lyric-evaluator via Hermes or LLM fallback)"""
|
||
lyrics = pipeline['artifacts'].get('lyrics', '')
|
||
# Use LLM as evaluator since lyric-evaluator is a Hermes skill
|
||
prompt = f"""你是一位严格的歌词评审专家。请评估以下歌词的质量，满分10分。
|
||
|
||
评分维度：
|
||
1. 结构与韵律 (2分)
|
||
2. 情感表达 (2分)
|
||
3. 意象与画面感 (2分)
|
||
4. 语言质量 (2分)
|
||
5. 可唱性 (2分)
|
||
|
||
歌词：
|
||
{lyrics}
|
||
|
||
请严格按JSON格式返回：
|
||
{{"total_score": 7.5, "dimensions": {{"structure": 1.5, "emotion": 1.5, "imagery": 1.5, "language": 1.5, "singability": 1.5}}, "comment": "简短评语"}}
|
||
"""
|
||
result = await call_llm(session, prompt, temperature=0.3)
|
||
try:
|
||
result = result.strip()
|
||
if result.startswith('```'):
|
||
result = result.split('```')[1]
|
||
if result.startswith('json'):
|
||
result = result[4:]
|
||
score_data = json.loads(result.strip())
|
||
except:
|
||
score_data = {'total_score': 7.0, 'dimensions': {}, 'comment': result[:200]}
|
||
return {'lyric_score': score_data}
|
||
|
||
|
||
async def step_music_generate(pipeline, session):
|
||
"""Step 4: Generate music via Suno API (through Sage llmage)"""
|
||
lyrics = pipeline['artifacts'].get('lyrics', '')
|
||
scene = pipeline.get('scene', 'pop')
|
||
|
||
# Call Suno through Sage
|
||
api_key = os.environ.get('LLM_API_KEY', '')
|
||
if not api_key:
|
||
try:
|
||
from ahserver.serverenv import ServerEnv
|
||
env = ServerEnv()
|
||
api_key = getattr(env, 'llm_api_key', '') or ''
|
||
except:
|
||
pass
|
||
|
||
headers = {
|
||
'Authorization': f'Bearer {api_key}',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
payload = {
|
||
'model': 'suno-v4',
|
||
'catelogid': 'music_gen',
|
||
'prompt': f'{scene} style music',
|
||
'tags': scene,
|
||
'title': f"AI Song - {pipeline['id']}",
|
||
}
|
||
async with session.post(f'{LLM_API}/audio/generations', json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=60)) as resp:
|
||
data = await resp.json()
|
||
taskid = data.get('taskid', data.get('id', ''))
|
||
return {'music_task_id': taskid, 'music_submit_response': data}
|
||
|
||
|
||
async def step_music_poll(pipeline, session):
|
||
"""Poll for music generation result"""
|
||
task_id = pipeline['artifacts'].get('music_task_id', '')
|
||
if not task_id:
|
||
raise ValueError('No music_task_id found')
|
||
|
||
api_key = os.environ.get('LLM_API_KEY', '')
|
||
headers = {'Authorization': f'Bearer {api_key}'}
|
||
|
||
# Poll Sage task endpoint
|
||
for attempt in range(60): # max 10 minutes
|
||
async with session.get(f'{LLM_API}/tasks?taskid={task_id}', headers=headers, timeout=aiohttp.ClientTimeout(total=30)) as resp:
|
||
data = await resp.json()
|
||
status = data.get('status', '')
|
||
if status == 'SUCCEEDED':
|
||
return {'music_url': data.get('result_url', ''), 'music_status': 'SUCCEEDED'}
|
||
elif status == 'FAILED':
|
||
raise ValueError(f'Music generation failed: {data}')
|
||
await asyncio.sleep(10)
|
||
|
||
raise ValueError('Music generation timed out')
|
||
|
||
|
||
async def step_music_separate(pipeline, session):
|
||
"""Step 5: Demucs vocal separation"""
|
||
music_url = pipeline['artifacts'].get('music_url', '')
|
||
if not music_url:
|
||
raise ValueError('No music_url found')
|
||
|
||
# Download the music file first
|
||
pipeline_dir = os.path.join(WORK_DIR, pipeline['id'])
|
||
os.makedirs(pipeline_dir, exist_ok=True)
|
||
music_path = os.path.join(pipeline_dir, 'music.mp3')
|
||
|
||
async with session.get(music_url, timeout=aiohttp.ClientTimeout(total=120)) as resp:
|
||
with open(music_path, 'wb') as f:
|
||
f.write(await resp.read())
|
||
|
||
# Call demucs
|
||
result = await call_service(session, SERVICES['demucs'], {'filepath': music_path})
|
||
return {'vocals_path': result.get('vocals_path', ''), 'no_vocals_path': result.get('no_vocals_path', '')}
|
||
|
||
|
||
async def step_music_align(pipeline, session):
|
||
"""Step 6: WhisperX transcription + alignment"""
|
||
vocals_path = pipeline['artifacts'].get('vocals_path', '')
|
||
lyrics = pipeline['artifacts'].get('lyrics', '')
|
||
|
||
# Submit to fastwhisper
|
||
result = await call_service(session, SERVICES['fastwhisper'], {'audio_path': vocals_path})
|
||
task_id = result.get('task_id', '')
|
||
|
||
if not task_id:
|
||
raise ValueError(f'WhisperX submit failed: {result}')
|
||
|
||
# Poll for result
|
||
for attempt in range(30):
|
||
await asyncio.sleep(5)
|
||
status_result = await call_service(session, SERVICES['fastwhisper_status'], {'task_id': task_id})
|
||
status = status_result.get('status', '')
|
||
if status == 'SUCCEEDED':
|
||
return {'whisperx_json': status_result.get('data', status_result)}
|
||
elif status == 'FAILED':
|
||
raise ValueError(f'WhisperX failed: {status_result}')
|
||
|
||
raise ValueError('WhisperX alignment timed out')
|
||
|
||
|
||
async def step_music_calibrate(pipeline, session):
|
||
"""Step 7: ASR + LLM lyric calibration (integrated service)"""
|
||
vocals_path = pipeline['artifacts'].get('vocals_path', '')
|
||
lyrics = pipeline['artifacts'].get('lyrics', '')
|
||
|
||
if not vocals_path:
|
||
raise ValueError('Missing vocals_path for calibration')
|
||
if not lyrics:
|
||
raise ValueError('Missing lyrics for calibration')
|
||
|
||
# Read vocals file
|
||
with open(vocals_path, 'rb') as f:
|
||
audio_data = f.read()
|
||
|
||
# Build multipart/form-data
|
||
boundary = f'Boundary{int(time.time()*1000)}'
|
||
|
||
body_parts = []
|
||
# audio_file field
|
||
body_parts.append(('--' + boundary + '\r\n').encode())
|
||
body_parts.append(b'Content-Disposition: form-data; name="audio_file"; filename="vocals.wav"\r\n')
|
||
body_parts.append(b'Content-Type: audio/wav\r\n\r\n')
|
||
body_parts.append(audio_data)
|
||
body_parts.append(b'\r\n')
|
||
|
||
# lyrics field
|
||
body_parts.append(('--' + boundary + '\r\n').encode())
|
||
body_parts.append(b'Content-Disposition: form-data; name="lyrics"\r\n\r\n')
|
||
body_parts.append(lyrics.encode('utf-8'))
|
||
body_parts.append(b'\r\n')
|
||
|
||
# closing boundary
|
||
body_parts.append(('--' + boundary + '--\r\n').encode())
|
||
|
||
body = b''.join(body_parts)
|
||
|
||
headers = {
|
||
'Content-Type': f'multipart/form-data; boundary={boundary}'
|
||
}
|
||
|
||
# Call lyric_calibrate service (integrated ASR + LLM)
|
||
url = SERVICES['lyric_calibrate']
|
||
timeout = aiohttp.ClientTimeout(total=600) # 10 minutes for ASR + LLM
|
||
|
||
async with session.post(url, data=body, headers=headers, timeout=timeout) as resp:
|
||
result = await resp.json()
|
||
|
||
if result.get('status') != 'ok':
|
||
error_msg = result.get('error', 'Unknown error')
|
||
raise ValueError(f'Lyric calibration failed: {error_msg}')
|
||
|
||
# Extract calibrated data
|
||
calibrated_lines = result.get('calibrated_lines', 0)
|
||
ass_file_url = result.get('ass_file', '')
|
||
json_data_url = result.get('json_data', '')
|
||
segments_used = result.get('segments_used', 0)
|
||
segments_total = result.get('segments_total', 0)
|
||
|
||
# Download and parse calibrated JSON
|
||
calibrated_json = None
|
||
if json_data_url:
|
||
# json_data_url is like '/idfile?path=lyric_calibrate/123456/calibrated.json'
|
||
# We need to extract the actual file path
|
||
import re
|
||
match = re.search(r'path=([^\s&]+)', json_data_url)
|
||
if match:
|
||
json_path = '/tmp/' + match.group(1)
|
||
if os.path.exists(json_path):
|
||
with open(json_path, 'r', encoding='utf-8') as f:
|
||
calibrated_json = json.load(f)
|
||
|
||
return {
|
||
'calibrated_subs': calibrated_json,
|
||
'calibrated_ass_file': ass_file_url,
|
||
'calibrated_lines': calibrated_lines,
|
||
'segments_used': segments_used,
|
||
'segments_total': segments_total
|
||
}
|
||
|
||
|
||
async def step_music_evaluate(pipeline, session):
|
||
"""Step 8: Song quality evaluation"""
|
||
music_path = os.path.join(WORK_DIR, pipeline['id'], 'music.mp3')
|
||
|
||
result = await call_service(session, SERVICES['songrate'], {
|
||
'filepath': music_path,
|
||
'scene': pipeline.get('scene', 'pop')
|
||
})
|
||
|
||
total = result.get('total_score', 0)
|
||
return {'music_score': result, 'music_total_score': total}
|
||
|
||
|
||
async def step_mv_story(pipeline, session):
|
||
"""Step 10: LLM generates MV storyline + storyboard"""
|
||
lyrics = pipeline['artifacts'].get('lyrics', '')
|
||
|
||
prompt = f"""你是一位MV导演。请根据以下歌词，创作MV的分镜脚本。
|
||
|
||
歌词：
|
||
{lyrics}
|
||
|
||
要求：
|
||
1. 为每个段落设计1-2个镜头
|
||
2. 每个镜头描述：场景、人物动作、镜头运动、情绪氛围
|
||
3. 输出JSON数组格式
|
||
|
||
请严格按JSON格式输出：
|
||
{{"scenes": [
|
||
{{"lyric_line": "歌词行", "description": "场景描述", "camera": "镜头运动", "mood": "情绪", "needs_reference": false}},
|
||
...
|
||
]}}
|
||
"""
|
||
result = await call_llm(session, prompt, temperature=0.8)
|
||
try:
|
||
result = result.strip()
|
||
if result.startswith('```'):
|
||
result = result.split('```')[1]
|
||
if result.startswith('json'):
|
||
result = result[4:]
|
||
storyboard = json.loads(result.strip())
|
||
except:
|
||
storyboard = {'scenes': [{'description': result[:500], 'camera': 'static', 'mood': 'neutral', 'needs_reference': False}]}
|
||
return {'storyboard': storyboard}
|
||
|
||
|
||
async def step_mv_generate_video(pipeline, session):
|
||
"""Steps 11-14: Generate MV video segments via Sage"""
|
||
# This step calls Sage video generation API (t2v/i2v/r2v)
|
||
# For now, create placeholder - actual video gen needs Sage integration
|
||
storyboard = pipeline['artifacts'].get('storyboard', {})
|
||
return {'mv_segments': [], 'mv_note': 'Video generation pending Sage integration'}
|
||
|
||
|
||
async def step_subtitle_render(pipeline, session):
|
||
"""Step 18: Generate ASS subtitle file with karaoke effects"""
|
||
# New format: calibrated_subs is a list of {line, start, end, chars}
|
||
# Old format: calibrated_subs was a dict with segments key
|
||
calibrated = pipeline['artifacts'].get('calibrated_subs', [])
|
||
if isinstance(calibrated, dict):
|
||
segments = calibrated.get('segments', [])
|
||
elif isinstance(calibrated, list):
|
||
segments = calibrated
|
||
else:
|
||
segments = []
|
||
|
||
pipeline_dir = os.path.join(WORK_DIR, pipeline['id'])
|
||
ass_path = os.path.join(pipeline_dir, 'karaoke.ass')
|
||
|
||
# Generate ASS file
|
||
ass_content = """[Script Info]
|
||
Title: KTV Karaoke Subtitles
|
||
ScriptType: v4.00+
|
||
PlayResX: 1920
|
||
PlayResY: 1080
|
||
WrapStyle: 0
|
||
|
||
[V4+ Styles]
|
||
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
|
||
Style: Karaoke,Noto Sans CJK SC,72,&H00FFFFFF,&H0000FFFF,&H00000000,&H80000000,-1,0,0,0,100,100,0,0,1,3,1,2,20,20,50,1
|
||
|
||
[Events]
|
||
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
|
||
"""
|
||
for seg in segments:
|
||
start = seg.get('start', 0)
|
||
end = seg.get('end', 0)
|
||
chars = seg.get('chars', [])
|
||
text = seg.get('line', seg.get('text', ''))
|
||
|
||
if chars:
|
||
# Build karaoke tags
|
||
kar_parts = []
|
||
for c in chars:
|
||
dur = max(1, int((c.get('end', c.get('start', 0)) - c.get('start', 0)) * 100))
|
||
kar_parts.append(f'{{\\k{dur}}}{c.get("char", "")}')
|
||
line_text = ''.join(kar_parts)
|
||
else:
|
||
line_text = text
|
||
|
||
def fmt_time(t):
|
||
h = int(t // 3600)
|
||
m = int((t % 3600) // 60)
|
||
s = t % 60
|
||
return f'{h}:{m:02d}:{s:05.2f}'
|
||
|
||
ass_content += f'Dialogue: 0,{fmt_time(start)},{fmt_time(end)},Karaoke,,0,0,0,,{line_text}\n'
|
||
|
||
with open(ass_path, 'w', encoding='utf-8') as f:
|
||
f.write(ass_content)
|
||
|
||
return {'subtitle_path': ass_path}
|
||
|
||
|
||
async def step_ktv_synthesize(pipeline, session):
|
||
"""Step 19: Final KTV synthesis"""
|
||
pipeline_dir = os.path.join(WORK_DIR, pipeline['id'])
|
||
|
||
video_path = pipeline['artifacts'].get('mv_merged_path', '')
|
||
vocals_path = pipeline['artifacts'].get('vocals_path', '')
|
||
no_vocals_path = pipeline['artifacts'].get('no_vocals_path', '')
|
||
subtitle_path = pipeline['artifacts'].get('subtitle_path', '')
|
||
|
||
if not all([video_path, no_vocals_path, subtitle_path]):
|
||
# If no MV, just create audio KTV
|
||
result = {'status': 'partial', 'note': 'Missing some artifacts, creating audio-only KTV'}
|
||
return {'ktv_result': result}
|
||
|
||
result = await call_service(session, SERVICES['ktv'], {
|
||
'video_path': video_path,
|
||
'vocals_path': vocals_path,
|
||
'accompaniment_path': no_vocals_path,
|
||
'subtitle_path': subtitle_path,
|
||
'output_name': pipeline['id']
|
||
}, timeout=600)
|
||
|
||
return {'ktv_result': result}
|
||
|
||
|
||
# State transition table
|
||
TRANSITIONS = {
|
||
'submitted': ('lyric_generating', step_lyric_generate),
|
||
'lyric_generating': ('lyric_evaluating', None),
|
||
'lyric_evaluating': ('lyric_done', None), # or back to lyric_generating
|
||
'lyric_done': ('music_generating', step_music_generate),
|
||
'music_generating': ('music_separating', step_music_poll),
|
||
'music_separating': ('music_aligning', step_music_separate),
|
||
'music_aligning': ('music_calibrating', None), # Skipped - lyric_calibrate handles ASR+LLM
|
||
'music_calibrating': ('music_evaluating', step_music_calibrate),
|
||
'music_evaluating': ('music_done', step_music_evaluate), # or back
|
||
'music_done': ('mv_story_generating', step_mv_story),
|
||
'mv_story_generating': ('mv_video_generating', None),
|
||
'mv_video_generating': ('subtitle_rendering', step_mv_generate_video),
|
||
'subtitle_rendering': ('ktv_synthesizing', step_subtitle_render),
|
||
'ktv_synthesizing': ('completed', step_ktv_synthesize),
|
||
}
|
||
|
||
|
||
async def run_pipeline(pipeline_id):
|
||
"""Execute the full KTV pipeline"""
|
||
import aioredis
|
||
redis = await aioredis.from_url('redis://127.0.0.1:6379', db=1)
|
||
|
||
try:
|
||
data = await redis.get(f'pipeline:{pipeline_id}')
|
||
if not data:
|
||
raise ValueError(f'Pipeline {pipeline_id} not found')
|
||
|
||
pipeline = json.loads(data)
|
||
pipeline_dir = os.path.join(WORK_DIR, pipeline_id)
|
||
os.makedirs(pipeline_dir, exist_ok=True)
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
state = pipeline['state']
|
||
retry_counts = {}
|
||
|
||
while state != 'completed' and state != 'failed':
|
||
await update_state(redis, pipeline_id, state)
|
||
|
||
transition = TRANSITIONS.get(state)
|
||
if not transition:
|
||
await update_state(redis, pipeline_id, 'failed', error=f'No transition for state: {state}')
|
||
break
|
||
|
||
next_state, handler = transition
|
||
|
||
if handler:
|
||
try:
|
||
result = await handler(pipeline, session)
|
||
pipeline['artifacts'].update(result)
|
||
|
||
# Check thresholds for evaluation states
|
||
if state == 'lyric_evaluating':
|
||
score = result.get('lyric_score', {}).get('total_score', 0)
|
||
retry_counts['lyric'] = retry_counts.get('lyric', 0) + 1
|
||
if score < pipeline.get('lyric_threshold', 8.5) and retry_counts['lyric'] < 3:
|
||
next_state = 'lyric_generating'
|
||
else:
|
||
next_state = 'lyric_done'
|
||
|
||
elif state == 'music_evaluating':
|
||
score = result.get('music_total_score', 0)
|
||
retry_counts['music'] = retry_counts.get('music', 0) + 1
|
||
if score < pipeline.get('music_threshold', 7.5) and retry_counts['music'] < 2:
|
||
next_state = 'music_generating'
|
||
else:
|
||
next_state = 'music_done'
|
||
|
||
except Exception as e:
|
||
await update_state(redis, pipeline_id, state, error=str(e))
|
||
# Retry once on error
|
||
retry_counts[state] = retry_counts.get(state, 0) + 1
|
||
if retry_counts[state] > 2:
|
||
await update_state(redis, pipeline_id, 'failed', error=f'{state} failed after retries: {e}')
|
||
break
|
||
await asyncio.sleep(5)
|
||
continue
|
||
|
||
state = next_state
|
||
|
||
finally:
|
||
await redis.close()
|