260 lines
8.2 KiB
Python
260 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from typing import Awaitable, Callable
|
|
from urllib.parse import urlparse
|
|
|
|
from app.config import settings
|
|
from app.services.generation.pipeline.owner_service import GenerationOwner
|
|
from app.services.image_gen import download_image
|
|
from app.services.provider_limit import provider_limit
|
|
from app.services.resource_accounting_service import safe_file_size
|
|
from app.services.video_cover_service import create_video_cover_for_local_video
|
|
from app.services.video_gen import download_video
|
|
from app.services.video_upscale.media_service import build_part_mp4_path, probe_video
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class DownloadedGenerationResult:
|
|
url: str
|
|
storage_path: str | None
|
|
file_size_bytes: int
|
|
resource_type: str
|
|
storage_type: str = "local"
|
|
cover_url: str | None = None
|
|
cover_storage_path: str | None = None
|
|
|
|
|
|
def _to_aware_utc(value: datetime | None) -> datetime | None:
|
|
if value is None:
|
|
return None
|
|
if value.tzinfo is None:
|
|
return value.replace(tzinfo=timezone.utc)
|
|
return value.astimezone(timezone.utc)
|
|
|
|
|
|
def _build_storage_date_dir(record: GenerationOwner) -> str:
|
|
fixed = (getattr(record, "download_storage_date_dir", None) or "").strip().strip("/")
|
|
if fixed:
|
|
return fixed
|
|
created_at = _to_aware_utc(getattr(record, "created_at", None)) or datetime.now(timezone.utc)
|
|
return created_at.strftime("%Y/%m/%d")
|
|
|
|
|
|
|
|
|
|
def _normalize_image_extension(record: GenerationOwner) -> str:
|
|
output_format = ""
|
|
try:
|
|
snapshot = json.loads(getattr(record, "engine_snapshot_json", None) or "{}")
|
|
if isinstance(snapshot, dict):
|
|
output_format = str(snapshot.get("output_format") or "").strip().lower()
|
|
except Exception:
|
|
output_format = ""
|
|
if output_format in {"jpg", "jpeg"}:
|
|
return "jpg"
|
|
if output_format in {"png", "webp"}:
|
|
return output_format
|
|
|
|
remote_url = str(getattr(record, "remote_result_url", None) or "")
|
|
try:
|
|
suffix = os.path.splitext(urlparse(remote_url).path or "")[1].lower().lstrip(".")
|
|
except Exception:
|
|
suffix = ""
|
|
if suffix in {"jpg", "jpeg"}:
|
|
return "jpg"
|
|
if suffix in {"png", "webp"}:
|
|
return suffix
|
|
return "jpg"
|
|
|
|
def _make_part_path(final_path: str) -> str:
|
|
return f"{final_path}.{uuid.uuid4().hex}.part"
|
|
|
|
|
|
def _is_valid_file(path: str | None) -> bool:
|
|
if not path:
|
|
return False
|
|
try:
|
|
return os.path.isfile(path) and os.path.getsize(path) > 0
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _safe_remove(path: str | None) -> None:
|
|
if not path:
|
|
return
|
|
try:
|
|
if os.path.exists(path):
|
|
os.remove(path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
async def _download_image_atomically(
|
|
remote_url: str,
|
|
final_path: str,
|
|
*,
|
|
execution_guard: Callable[[], Awaitable[None]] | None = None,
|
|
) -> str:
|
|
if _is_valid_file(final_path):
|
|
return final_path
|
|
|
|
os.makedirs(os.path.dirname(final_path), exist_ok=True)
|
|
part_path = _make_part_path(final_path)
|
|
try:
|
|
await download_image(
|
|
remote_url,
|
|
part_path,
|
|
execution_guard=execution_guard,
|
|
)
|
|
if not _is_valid_file(part_path):
|
|
raise RuntimeError("图片下载完成但临时文件为空")
|
|
if execution_guard is not None:
|
|
await execution_guard()
|
|
os.replace(part_path, final_path)
|
|
return final_path
|
|
except Exception:
|
|
_safe_remove(part_path)
|
|
raise
|
|
|
|
|
|
async def _download_video_atomically(
|
|
remote_url: str,
|
|
final_path: str,
|
|
*,
|
|
execution_guard: Callable[[], Awaitable[None]] | None = None,
|
|
) -> str:
|
|
if _is_valid_file(final_path):
|
|
return final_path
|
|
|
|
os.makedirs(os.path.dirname(final_path), exist_ok=True)
|
|
part_path = _make_part_path(final_path)
|
|
try:
|
|
await download_video(
|
|
remote_url,
|
|
part_path,
|
|
execution_guard=execution_guard,
|
|
)
|
|
if not _is_valid_file(part_path):
|
|
raise RuntimeError("视频下载完成但临时文件为空")
|
|
if execution_guard is not None:
|
|
await execution_guard()
|
|
os.replace(part_path, final_path)
|
|
return final_path
|
|
except Exception:
|
|
_safe_remove(part_path)
|
|
raise
|
|
|
|
|
|
|
|
|
|
async def download_video_upscale_source(
|
|
record: GenerationOwner,
|
|
*,
|
|
execution_guard: Callable[[], Awaitable[None]] | None = None,
|
|
) -> DownloadedGenerationResult:
|
|
"""下载超分源视频。
|
|
|
|
源视频只供后处理使用，不生成封面，也不作为用户 GeneratedResource。
|
|
临时文件保持 .part.mp4 后缀，校验成功后原子重命名为 .source.mp4。
|
|
"""
|
|
if not record.remote_result_url:
|
|
raise ValueError("缺少远程结果URL")
|
|
|
|
date_dir = _build_storage_date_dir(record)
|
|
dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, "_upscale_source", date_dir)
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
dest = os.path.join(dest_dir, f"{record.id}.source.mp4")
|
|
|
|
if not _is_valid_file(dest):
|
|
part_path = build_part_mp4_path(dest)
|
|
try:
|
|
async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY):
|
|
await download_video(
|
|
record.remote_result_url,
|
|
part_path,
|
|
execution_guard=execution_guard,
|
|
)
|
|
if not _is_valid_file(part_path):
|
|
raise RuntimeError("超分源视频下载完成但临时文件为空")
|
|
await probe_video(part_path)
|
|
if execution_guard is not None:
|
|
await execution_guard()
|
|
os.replace(part_path, dest)
|
|
except Exception:
|
|
_safe_remove(part_path)
|
|
raise
|
|
else:
|
|
await probe_video(dest)
|
|
|
|
return DownloadedGenerationResult(
|
|
url=f"/generate/videos/_upscale_source/{date_dir}/{record.id}.source.mp4",
|
|
storage_path=dest,
|
|
file_size_bytes=safe_file_size(dest),
|
|
resource_type="video",
|
|
)
|
|
|
|
|
|
async def download_generation_result(
|
|
record: GenerationOwner,
|
|
*,
|
|
execution_guard: Callable[[], Awaitable[None]] | None = None,
|
|
) -> DownloadedGenerationResult:
|
|
if not record.remote_result_url:
|
|
raise ValueError("缺少远程结果URL")
|
|
|
|
date_dir = _build_storage_date_dir(record)
|
|
|
|
if record.gen_type == "image":
|
|
dest_dir = os.path.join(settings.STORAGE_IMAGE_LOCAL_PATH, date_dir)
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
extension = _normalize_image_extension(record)
|
|
dest = os.path.join(dest_dir, f"{record.id}.{extension}")
|
|
|
|
async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY):
|
|
await _download_image_atomically(
|
|
record.remote_result_url if record.remote_result_url else "",
|
|
dest,
|
|
execution_guard=execution_guard,
|
|
)
|
|
|
|
return DownloadedGenerationResult(
|
|
url=f"/generate/images/{date_dir}/{record.id}.{extension}",
|
|
storage_path=dest,
|
|
file_size_bytes=safe_file_size(dest),
|
|
resource_type="image",
|
|
)
|
|
|
|
dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, date_dir)
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
dest = os.path.join(dest_dir, f"{record.id}.mp4")
|
|
|
|
async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY):
|
|
await _download_video_atomically(
|
|
record.remote_result_url if record.remote_result_url else "",
|
|
dest,
|
|
execution_guard=execution_guard,
|
|
)
|
|
|
|
cover_url, cover_storage_path = create_video_cover_for_local_video(
|
|
record_id=record.id,
|
|
video_path=dest,
|
|
date_dir=date_dir,
|
|
log_prefix=f"生成资源视频封面 task_id={record.id}",
|
|
)
|
|
if execution_guard is not None:
|
|
await execution_guard()
|
|
|
|
return DownloadedGenerationResult(
|
|
url=f"/generate/videos/{date_dir}/{record.id}.mp4",
|
|
storage_path=dest,
|
|
file_size_bytes=safe_file_size(dest),
|
|
resource_type="video",
|
|
cover_url=cover_url,
|
|
cover_storage_path=cover_storage_path,
|
|
)
|