yumoqing
769fc4968e
- Wan2.2-TI2V-5B GPU 视频推理 - ahserver + longtasks 异步任务队列 - OpenAI 兼容 API: POST /api/submit, GET /api/task, GET /api/status - 模型常驻内存，惰性加载 - 全局串行推理锁（GPU 安全） - 支持 t2v/i2v/ti2v/s2v 四种任务类型
132 lines
3.7 KiB
Python
Executable File
132 lines
3.7 KiB
Python
Executable File
# -*- coding:utf-8 -*-
|
|
"""
|
|
Wan2.2-TI2V-5B 视频生成 worker
|
|
调用官方 generate.py 脚本
|
|
"""
|
|
import os
|
|
import json
|
|
import uuid
|
|
import asyncio
|
|
import subprocess
|
|
from datetime import datetime
|
|
from appPublic.log import debug, exception
|
|
|
|
OUTPUT_DIR = '/data/ymq/wan22-outputs'
|
|
REPO_DIR = '/data/ymq/wan22-service/repo'
|
|
MODEL_PATH = '/data/ymq/models/Wan-AI/Wan2.2-TI2V-5B'
|
|
PYTHON = '/share/vllm-0.8.5/bin/python'
|
|
|
|
|
|
async def run_generate(longtasks, payload):
|
|
"""
|
|
Execute video generation via generate.py subprocess.
|
|
|
|
payload: {
|
|
task_id: str,
|
|
prompt: str,
|
|
image: str (optional path for i2v),
|
|
size: str (default "1280*720"),
|
|
frame_num: int (default 81, must be 4n+1),
|
|
sample_steps: int (optional),
|
|
sample_guide_scale: float (optional),
|
|
base_seed: int (optional),
|
|
}
|
|
"""
|
|
task_id = payload.get('task_id', str(uuid.uuid4())[:12])
|
|
prompt = payload.get('prompt', '')
|
|
image = payload.get('image', None)
|
|
size = payload.get('size', '1280*720')
|
|
frame_num = payload.get('frame_num', 81)
|
|
sample_steps = payload.get('sample_steps', None)
|
|
sample_guide_scale = payload.get('sample_guide_scale', None)
|
|
base_seed = payload.get('base_seed', None)
|
|
|
|
# Ensure frame_num is 4n+1
|
|
frame_num = max(17, min(frame_num, 129))
|
|
if (frame_num - 1) % 4 != 0:
|
|
frame_num = ((frame_num - 1) // 4) * 4 + 1
|
|
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
output_file = os.path.join(OUTPUT_DIR, f'{task_id}.mp4')
|
|
|
|
# Build command
|
|
cmd = [
|
|
PYTHON, 'generate.py',
|
|
'--task', 'ti2v-5B',
|
|
'--ckpt_dir', MODEL_PATH,
|
|
'--size', size,
|
|
'--frame_num', str(frame_num),
|
|
'--prompt', prompt,
|
|
'--save_file', output_file,
|
|
'--offload_model', 'True',
|
|
]
|
|
|
|
if image:
|
|
cmd.extend(['--image', image])
|
|
|
|
if sample_steps:
|
|
cmd.extend(['--sample_steps', str(sample_steps)])
|
|
|
|
if sample_guide_scale:
|
|
cmd.extend(['--sample_guide_scale', str(sample_guide_scale)])
|
|
|
|
if base_seed is not None:
|
|
cmd.extend(['--base_seed', str(base_seed)])
|
|
|
|
# Set CUDA_VISIBLE_DEVICES for single GPU
|
|
gpu_id = longtasks.gpu_id if longtasks.gpu_id else 2
|
|
env = os.environ.copy()
|
|
env['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
|
|
|
|
debug(f'Running: {" ".join(cmd)}')
|
|
|
|
try:
|
|
# Run in subprocess
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
cwd=REPO_DIR,
|
|
env=env,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
|
|
stdout, stderr = await proc.communicate()
|
|
|
|
if proc.returncode != 0:
|
|
error_msg = stderr.decode('utf-8', errors='ignore')[-500:]
|
|
exception(f'generate.py failed: {error_msg}')
|
|
return {
|
|
'task_id': task_id,
|
|
'status': 'failed',
|
|
'error': error_msg
|
|
}
|
|
|
|
# Check output file
|
|
if not os.path.exists(output_file):
|
|
return {
|
|
'task_id': task_id,
|
|
'status': 'failed',
|
|
'error': 'Output file not created'
|
|
}
|
|
|
|
file_size = os.path.getsize(output_file)
|
|
|
|
return {
|
|
'task_id': task_id,
|
|
'status': 'completed',
|
|
'video_url': f'/idfile?path={task_id}.mp4',
|
|
'video_path': output_file,
|
|
'size': size,
|
|
'frame_num': frame_num,
|
|
'file_size': file_size,
|
|
'prompt': prompt[:100]
|
|
}
|
|
|
|
except Exception as e:
|
|
exception(f'Generation error: {e}')
|
|
return {
|
|
'task_id': task_id,
|
|
'status': 'failed',
|
|
'error': str(e)
|
|
}
|