Xiaoxia AI
260 lines
8.2 KiB
Python
260 lines
8.2 KiB
Python
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
|
|
import oss2
|
|
from worker_app.celery_app import celery_app
|
|
from worker_app.db import SessionLocal
|
|
|
|
from packages.adapters.sqlalchemy_impl import (
|
|
SQLAlchemyAssetRepository,
|
|
SQLAlchemyGeneratedVideoRepository,
|
|
SQLAlchemyGenerationTaskRepository,
|
|
)
|
|
from packages.domain import GeneratedVideo, GenerationTaskStatus
|
|
|
|
OUTPUT_WIDTH = 1280
|
|
OUTPUT_HEIGHT = 720
|
|
OUTPUT_FPS = 25.0
|
|
OUTPUT_DURATION_SECONDS = 5.0
|
|
GENERATED_FILES_DIR = Path(os.getenv("GENERATED_FILES_DIR", "/app/generated"))
|
|
GENERATED_FILES_URL_PREFIX = os.getenv("GENERATED_FILES_URL_PREFIX", "/generated-files")
|
|
PUBLIC_API_BASE_URL = os.getenv("PUBLIC_API_BASE_URL", "https://api.xiaoxiajianji.com").rstrip("/")
|
|
|
|
|
|
def _run_ffmpeg(command: list[str]) -> None:
|
|
subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
|
|
|
|
def _oss_settings() -> tuple[str, str, str, str] | None:
|
|
access_key_id = os.getenv("OSS_ACCESS_KEY_ID")
|
|
access_key_secret = os.getenv("OSS_ACCESS_KEY_SECRET")
|
|
endpoint = os.getenv("OSS_ENDPOINT")
|
|
bucket_name = os.getenv("OSS_BUCKET_NAME")
|
|
if not all([access_key_id, access_key_secret, endpoint, bucket_name]):
|
|
return None
|
|
return access_key_id, access_key_secret, endpoint, bucket_name
|
|
|
|
|
|
def _oss_bucket() -> oss2.Bucket | None:
|
|
settings = _oss_settings()
|
|
if settings is None:
|
|
return None
|
|
access_key_id, access_key_secret, endpoint, bucket_name = settings
|
|
return oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
|
|
|
|
|
|
def _public_oss_url(storage_key: str) -> str:
|
|
settings = _oss_settings()
|
|
if settings is None:
|
|
raise RuntimeError("OSS storage is not configured")
|
|
_, _, endpoint, bucket_name = settings
|
|
normalized_endpoint = endpoint.removeprefix("https://").removeprefix("http://")
|
|
return f"https://{bucket_name}.{normalized_endpoint}/{storage_key}"
|
|
|
|
|
|
def _normalize_storage_key(storage_key_or_url: str) -> str:
|
|
if storage_key_or_url.startswith(("http://", "https://")):
|
|
return urlparse(storage_key_or_url).path.lstrip("/")
|
|
return storage_key_or_url.lstrip("/")
|
|
|
|
|
|
def _download_asset(asset_storage_key: str, local_path: Path) -> bool:
|
|
bucket = _oss_bucket()
|
|
if bucket is None:
|
|
return False
|
|
try:
|
|
bucket.get_object_to_file(_normalize_storage_key(asset_storage_key), str(local_path))
|
|
return local_path.exists() and local_path.stat().st_size > 0
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _local_generated_url(storage_key: str) -> str:
|
|
return f"{PUBLIC_API_BASE_URL}{GENERATED_FILES_URL_PREFIX}/{storage_key}"
|
|
|
|
|
|
def _store_generated_video(local_path: Path, storage_key: str) -> str:
|
|
bucket = _oss_bucket()
|
|
if bucket is not None:
|
|
bucket.put_object_from_file(
|
|
storage_key,
|
|
str(local_path),
|
|
headers={"Content-Type": "video/mp4"},
|
|
)
|
|
return _public_oss_url(storage_key)
|
|
|
|
target_path = GENERATED_FILES_DIR / storage_key
|
|
target_path.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copyfile(local_path, target_path)
|
|
return _local_generated_url(storage_key)
|
|
|
|
|
|
def _probe_duration(local_path: Path) -> float:
|
|
try:
|
|
result = subprocess.run(
|
|
[
|
|
"ffprobe",
|
|
"-v",
|
|
"error",
|
|
"-show_entries",
|
|
"format=duration",
|
|
"-of",
|
|
"default=noprint_wrappers=1:nokey=1",
|
|
str(local_path),
|
|
],
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
return round(float(result.stdout.strip()), 3)
|
|
except Exception:
|
|
return OUTPUT_DURATION_SECONDS
|
|
|
|
|
|
def _create_fallback_clip(output_path: Path, title: str) -> None:
|
|
safe_title = title.replace(":", "\\:").replace("'", "\\'")[:80]
|
|
_run_ffmpeg(
|
|
[
|
|
"ffmpeg",
|
|
"-y",
|
|
"-f",
|
|
"lavfi",
|
|
"-i",
|
|
f"color=c=#111827:s={OUTPUT_WIDTH}x{OUTPUT_HEIGHT}:d={OUTPUT_DURATION_SECONDS}:r={int(OUTPUT_FPS)}",
|
|
"-vf",
|
|
f"drawtext=text='{safe_title}':fontcolor=white:fontsize=48:x=(w-text_w)/2:y=(h-text_h)/2",
|
|
"-c:v",
|
|
"libx264",
|
|
"-pix_fmt",
|
|
"yuv420p",
|
|
"-movflags",
|
|
"+faststart",
|
|
str(output_path),
|
|
]
|
|
)
|
|
|
|
|
|
def _compose_from_asset(input_path: Path, output_path: Path) -> None:
|
|
_run_ffmpeg(
|
|
[
|
|
"ffmpeg",
|
|
"-y",
|
|
"-i",
|
|
str(input_path),
|
|
"-t",
|
|
str(OUTPUT_DURATION_SECONDS),
|
|
"-vf",
|
|
f"scale={OUTPUT_WIDTH}:{OUTPUT_HEIGHT}:force_original_aspect_ratio=decrease,pad={OUTPUT_WIDTH}:{OUTPUT_HEIGHT}:(ow-iw)/2:(oh-ih)/2,setsar=1",
|
|
"-r",
|
|
str(int(OUTPUT_FPS)),
|
|
"-an",
|
|
"-c:v",
|
|
"libx264",
|
|
"-pix_fmt",
|
|
"yuv420p",
|
|
"-movflags",
|
|
"+faststart",
|
|
str(output_path),
|
|
]
|
|
)
|
|
|
|
|
|
@celery_app.task(name="worker.generate_video")
|
|
def generate_video(task_id: str) -> dict:
|
|
"""Generate and persist a real MP4 video for a generation task."""
|
|
db = SessionLocal()
|
|
task_repo = SQLAlchemyGenerationTaskRepository(db)
|
|
asset_repo = SQLAlchemyAssetRepository(db)
|
|
video_repo = SQLAlchemyGeneratedVideoRepository(db)
|
|
|
|
task = task_repo.get(task_id)
|
|
if task is None:
|
|
db.close()
|
|
return {
|
|
"status": "failed",
|
|
"error": "generation task not found",
|
|
"task_id": task_id,
|
|
}
|
|
|
|
try:
|
|
task.status = GenerationTaskStatus.RUNNING
|
|
task.progress = 10.0
|
|
task.started_at = task.started_at or datetime.now(timezone.utc)
|
|
task_repo.update(task)
|
|
|
|
assets = [
|
|
asset for asset in asset_repo.list_by_library(task.asset_library_id) if asset.mime_type.startswith("video")
|
|
]
|
|
|
|
output_name = f"generated-{task.id}.mp4"
|
|
storage_key = (
|
|
f"generated/workspaces/{task.workspace_id}/projects/{task.project_id}/tasks/{task.id}/{output_name}"
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="xiaoxia-generation-") as temp_dir:
|
|
temp_path = Path(temp_dir)
|
|
output_path = temp_path / output_name
|
|
input_path = temp_path / "source.mp4"
|
|
|
|
task.progress = 35.0
|
|
task_repo.update(task)
|
|
|
|
source_downloaded = False
|
|
if assets:
|
|
source_downloaded = _download_asset(assets[0].storage_key, input_path)
|
|
|
|
if source_downloaded:
|
|
_compose_from_asset(input_path, output_path)
|
|
else:
|
|
_create_fallback_clip(output_path, f"Xiaoxia Generated Video {task.id[:8]}")
|
|
|
|
task.progress = 70.0
|
|
task_repo.update(task)
|
|
|
|
file_url = _store_generated_video(output_path, storage_key)
|
|
file_size = output_path.stat().st_size
|
|
duration = _probe_duration(output_path)
|
|
|
|
video = GeneratedVideo.create(
|
|
workspace_id=task.workspace_id,
|
|
project_id=task.project_id,
|
|
generation_task_id=task.id,
|
|
name=output_name,
|
|
file_url=file_url,
|
|
file_size=file_size,
|
|
duration=duration,
|
|
width=OUTPUT_WIDTH,
|
|
height=OUTPUT_HEIGHT,
|
|
fps=OUTPUT_FPS,
|
|
thumbnail_url=None,
|
|
)
|
|
video_repo.create(video)
|
|
|
|
task.status = GenerationTaskStatus.COMPLETED
|
|
task.progress = 100.0
|
|
task.result_count = 1
|
|
task.error_message = ""
|
|
task.completed_at = datetime.now(timezone.utc)
|
|
task_repo.update(task)
|
|
|
|
return {
|
|
"status": "completed",
|
|
"task_id": task.id,
|
|
"video_id": video.id,
|
|
"file_url": file_url,
|
|
}
|
|
except Exception as error:
|
|
task.status = GenerationTaskStatus.FAILED
|
|
task.error_message = str(error)
|
|
task.completed_at = datetime.now(timezone.utc)
|
|
task_repo.update(task)
|
|
return {"status": "failed", "task_id": task.id, "error": str(error)}
|
|
finally:
|
|
db.close()
|