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()