214 lines
6.7 KiB
Python
214 lines
6.7 KiB
Python
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.api.api_generation_task import ApiGenerationTask
|
|
from app.models.api.api_key_upscale_config import ApiKeyUpscaleConfig
|
|
from app.models.api.api_upscale_link import ApiUpscaleLink
|
|
from app.utils.id_gen import generate_id
|
|
|
|
logger = logging.getLogger("videogen")
|
|
|
|
|
|
async def get_or_create_upscale_config(
|
|
db: AsyncSession,
|
|
api_key_id: str,
|
|
) -> ApiKeyUpscaleConfig:
|
|
"""获取或创建 API Key 的超分配置。"""
|
|
result = await db.execute(
|
|
select(ApiKeyUpscaleConfig).where(
|
|
ApiKeyUpscaleConfig.api_key_id == api_key_id
|
|
).limit(1)
|
|
)
|
|
config = result.scalar_one_or_none()
|
|
|
|
if not config:
|
|
config = ApiKeyUpscaleConfig(
|
|
id=generate_id(),
|
|
api_key_id=api_key_id,
|
|
enabled=False,
|
|
delete_source_after_success=True,
|
|
rules_json="[]",
|
|
)
|
|
db.add(config)
|
|
await db.flush()
|
|
|
|
return config
|
|
|
|
|
|
async def save_upscale_config(
|
|
db: AsyncSession,
|
|
api_key_id: str,
|
|
enabled: bool,
|
|
delete_source_after_success: bool,
|
|
rules: list[dict],
|
|
) -> ApiKeyUpscaleConfig:
|
|
"""保存 API Key 的超分配置。"""
|
|
config = await get_or_create_upscale_config(db, api_key_id)
|
|
config.enabled = enabled
|
|
config.delete_source_after_success = delete_source_after_success
|
|
config.rules_json = json.dumps(rules, ensure_ascii=False)
|
|
await db.flush()
|
|
return config
|
|
|
|
|
|
async def build_api_upscale_snapshot(
|
|
db: AsyncSession,
|
|
api_key_id: str,
|
|
target_resolution: str,
|
|
aspect_ratio: str | None = None,
|
|
) -> tuple[str | None, bool, str | None]:
|
|
"""构建 API 超分快照。
|
|
|
|
读取 api_key_upscale_configs（而非 system_configs），
|
|
匹配目标分辨率对应的超分规则。
|
|
|
|
Returns:
|
|
(provider_generation_resolution, enabled, snapshot_json)
|
|
"""
|
|
config = await get_or_create_upscale_config(db, api_key_id)
|
|
|
|
if not config.enabled:
|
|
return None, False, None
|
|
|
|
try:
|
|
rules = json.loads(config.rules_json) if config.rules_json else []
|
|
except (json.JSONDecodeError, TypeError):
|
|
return None, False, None
|
|
|
|
# 匹配规则
|
|
matched_rule = None
|
|
for rule in rules:
|
|
if rule.get("enabled") and rule.get("target_resolution") == target_resolution:
|
|
matched_rule = rule
|
|
break
|
|
|
|
if not matched_rule:
|
|
return None, False, None
|
|
|
|
snapshot = {
|
|
"enabled": True,
|
|
"delete_source_after_success": config.delete_source_after_success,
|
|
"rule": matched_rule,
|
|
"matched_at": datetime.now(timezone.utc).isoformat(),
|
|
# 兼容现有超分流水线的 processor 字段
|
|
"processor": {
|
|
"max_attempts": 3,
|
|
"processor_key": matched_rule.get("processor_key", "volc_large_model_v1"),
|
|
},
|
|
"target_resolution": target_resolution,
|
|
"provider_generation_resolution": matched_rule.get("provider_generation_resolution", target_resolution),
|
|
"aspect_ratio": aspect_ratio,
|
|
}
|
|
|
|
provider_resolution = matched_rule.get("provider_generation_resolution", target_resolution)
|
|
snapshot_json = json.dumps(snapshot, ensure_ascii=False)
|
|
|
|
return provider_resolution, True, snapshot_json
|
|
|
|
|
|
async def prepare_api_upscale_task(
|
|
db: AsyncSession,
|
|
api_task: ApiGenerationTask,
|
|
source_local_path: str,
|
|
source_width: int = 0,
|
|
source_height: int = 0,
|
|
source_duration: float = 0.0,
|
|
source_file_size_bytes: int = 0
|
|
) -> "VideoUpscaleTask | None":
|
|
"""为 API 任务创建超分子任务。
|
|
|
|
复用现有的 VideoUpscaleTask 表和 upscale 执行流水线。
|
|
如果已存在超分任务则返回 None（避免重复创建）。
|
|
"""
|
|
from app.models.video_upscale_task import VideoUpscaleTask
|
|
from sqlalchemy import select
|
|
|
|
# 检查是否已存在超分任务（避免重复创建）
|
|
existing = await db.execute(
|
|
select(VideoUpscaleTask).where(
|
|
VideoUpscaleTask.api_generation_task_id == api_task.id
|
|
).limit(1)
|
|
)
|
|
if existing.scalar_one_or_none():
|
|
logger.info("Upscale task already exists for API task %s, skipping", api_task.id)
|
|
return None
|
|
|
|
# 解析快照获取处理器配置
|
|
try:
|
|
snapshot = json.loads(api_task.video_upscale_snapshot_json) if api_task.video_upscale_snapshot_json else {}
|
|
except (json.JSONDecodeError, TypeError):
|
|
snapshot = {}
|
|
|
|
rule = snapshot.get("rule", {})
|
|
processor_key = rule.get("processor_key", "volc_large_model_v1")
|
|
target_resolution = rule.get("target_resolution", api_task.resolution or "1080p")
|
|
|
|
# 计算目标尺寸
|
|
target_width, target_height = _resolution_to_dimensions(target_resolution, api_task.aspect_ratio)
|
|
|
|
upscale_task = VideoUpscaleTask(
|
|
id=generate_id(),
|
|
chat_generation_task_id=None,
|
|
generation_record_id=None,
|
|
api_generation_task_id=api_task.id, # 关联 API v3 任务
|
|
processor_key=processor_key,
|
|
target_width=target_width,
|
|
target_height=target_height,
|
|
effective_target_width=target_width,
|
|
effective_target_height=target_height,
|
|
source_local_path=api_task.local_path or source_local_path, # 优先使用已下载的本地文件
|
|
source_remote_url=api_task.remote_result_url, # 火山 MediaKit 需要远程 URL
|
|
input_source_type="provider_remote",
|
|
source_file_size_bytes=source_file_size_bytes,
|
|
source_width=source_width,
|
|
source_height=source_height,
|
|
source_duration_seconds=source_duration,
|
|
status="pending",
|
|
stage="upscale_queued",
|
|
)
|
|
db.add(upscale_task)
|
|
await db.flush()
|
|
|
|
# 创建关联记录
|
|
link = ApiUpscaleLink(
|
|
id=generate_id(),
|
|
api_generation_task_id=api_task.id,
|
|
video_upscale_task_id=upscale_task.id,
|
|
)
|
|
db.add(link)
|
|
await db.flush()
|
|
|
|
logger.info(
|
|
"API upscale task prepared: api_task=%s upscale_task=%s processor=%s",
|
|
api_task.id, upscale_task.id, processor_key,
|
|
)
|
|
return upscale_task
|
|
|
|
|
|
def _resolution_to_dimensions(resolution: str, aspect_ratio: str | None) -> tuple[int, int]:
|
|
"""将分辨率名称转换为像素尺寸。"""
|
|
# 标准分辨率映射
|
|
resolution_map = {
|
|
"480p": (852, 480),
|
|
"720p": (1280, 720),
|
|
"1080p": (1920, 1080),
|
|
"2K": (2560, 1440),
|
|
"4K": (3840, 2160),
|
|
}
|
|
|
|
base = resolution_map.get(resolution, (1920, 1080))
|
|
|
|
# 根据宽高比调整
|
|
if aspect_ratio == "9:16":
|
|
return (base[1], base[0]) # 竖屏
|
|
elif aspect_ratio == "1:1":
|
|
return (base[0], base[0]) # 正方形
|
|
elif aspect_ratio == "4:3":
|
|
return (base[0], int(base[0] * 3 / 4))
|
|
|
|
return base
|