diff --git a/mytoolkit/api/__init__.py b/mytoolkit/api/__init__.py index 885ed56..ed8e23a 100644 --- a/mytoolkit/api/__init__.py +++ b/mytoolkit/api/__init__.py @@ -3,11 +3,15 @@ from .asr import transcribe_flash from .asr_correct import correct_asr_text from .image import generate_image +from .video_gen import create_task, get_task, generate_video from .voice import synthesize_tts __all__ = [ "generate_image", + "generate_video", "synthesize_tts", "transcribe_flash", "correct_asr_text", + "create_task", + "get_task", ] diff --git a/mytoolkit/api/video_gen.py b/mytoolkit/api/video_gen.py new file mode 100644 index 0000000..02b109a --- /dev/null +++ b/mytoolkit/api/video_gen.py @@ -0,0 +1,246 @@ +"""Video generation API client via Volcengine Seedance API.""" + +from __future__ import annotations + +import json +import time +import urllib.request +from collections.abc import Callable +from pathlib import Path +from typing import Any + +_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3" +_MODEL = "doubao-seedance-2-0-260128" + +_POLL_INTERVAL = 5 # seconds +_POLL_TIMEOUT = 600 # total seconds (10 min) + + +def _headers(api_key: str) -> dict[str, str]: + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + + +def _request( + api_key: str, + method: str, + path: str, + body: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Make a JSON HTTP request to the Volcengine API.""" + url = f"{_BASE_URL}{path}" + data = json.dumps(body).encode("utf-8") if body else None + req = urllib.request.Request(url, data=data, headers=_headers(api_key), method=method) + with urllib.request.urlopen(req) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def create_task( + api_key: str, + text_prompt: str, + *, + model: str = _MODEL, + image_urls: list[str] | None = None, + video_urls: list[str] | None = None, + audio_urls: list[str] | None = None, + generate_audio: bool = True, + ratio: str = "16:9", + duration: int = 11, + watermark: bool = False, +) -> dict[str, Any]: + """Create a video generation task via Volcengine Seedance API. + + Args: + api_key: Volcengine ARK API key. + text_prompt: Video description text prompt. + model: Model identifier (default: doubao-seedance-2-0-260128). + image_urls: Optional list of public image URLs for reference. + video_urls: Optional list of public video URLs for reference. + audio_urls: Optional list of public audio URLs for reference. + generate_audio: Whether to generate synchronized audio. + ratio: Aspect ratio (16:9, 9:16, 1:1, etc.). + duration: Video duration in seconds. + watermark: Whether to add the Doubao watermark. + + Returns: + Response JSON containing task information. + """ + content: list[dict[str, Any]] = [{"type": "text", "text": text_prompt}] + + for url in image_urls or []: + content.append({ + "type": "image_url", + "image_url": {"url": url}, + "role": "reference_image", + }) + + for url in video_urls or []: + content.append({ + "type": "video_url", + "video_url": {"url": url}, + "role": "reference_video", + }) + + for url in audio_urls or []: + content.append({ + "type": "audio_url", + "audio_url": {"url": url}, + "role": "reference_audio", + }) + + body = { + "model": model, + "content": content, + "generate_audio": generate_audio, + "ratio": ratio, + "duration": duration, + "watermark": watermark, + } + + return _request(api_key, "POST", "/contents/generations/tasks", body) + + +def get_task(api_key: str, task_id: str) -> dict[str, Any]: + """Query the status of a video generation task. + + Args: + api_key: Volcengine ARK API key. + task_id: Task ID returned by create_task(). + + Returns: + Response JSON with task status and result (if completed). + """ + return _request(api_key, "GET", f"/contents/generations/tasks/{task_id}") + + +def generate_video( + api_key: str, + text_prompt: str, + *, + model: str = _MODEL, + image_urls: list[str] | None = None, + video_urls: list[str] | None = None, + audio_urls: list[str] | None = None, + generate_audio: bool = True, + ratio: str = "16:9", + duration: int = 11, + watermark: bool = False, + poll: bool = True, + poll_interval: int = _POLL_INTERVAL, + poll_timeout: int = _POLL_TIMEOUT, + output: str | None = None, + progress_callback: Callable[[str, int], None] | None = None, +) -> dict[str, Any]: + """Create a video generation task and optionally poll for results. + + Args: + api_key: Volcengine ARK API key. + text_prompt: Video description text prompt. + model: Model identifier. + image_urls: Optional reference image URLs. + video_urls: Optional reference video URLs. + audio_urls: Optional reference audio URLs. + generate_audio: Whether to generate audio. + ratio: Aspect ratio. + duration: Video duration in seconds. + watermark: Whether to add watermark. + poll: Whether to poll until completion. + poll_interval: Seconds between poll attempts. + poll_timeout: Max total poll time in seconds. + output: If set, download the result video to this path (requires poll=True). + progress_callback: Optional function(msg, elapsed_secs) for progress updates. + + Returns: + Response JSON. If poll=True, contains final result. + """ + if progress_callback is None: + progress_callback = lambda msg, elapsed: None # noqa: E731 + + task = create_task( + api_key=api_key, + text_prompt=text_prompt, + model=model, + image_urls=image_urls, + video_urls=video_urls, + audio_urls=audio_urls, + generate_audio=generate_audio, + ratio=ratio, + duration=duration, + watermark=watermark, + ) + + task_id = task.get("id") + if not task_id: + return task + + progress_callback(f"Task submitted, ID: {task_id}", 0) + + if not poll: + return task + + # Poll until completed or failed + started = time.time() + deadline = started + poll_timeout + while time.time() < deadline: + elapsed = int(time.time() - started) + progress_callback("Querying task status...", elapsed) + time.sleep(poll_interval) + result = get_task(api_key, task_id) + status = result.get("status", "").lower() + progress = result.get("data", {}).get("progress", result.get("progress")) + + elapsed = int(time.time() - started) + if progress is not None: + progress_callback(f"Progress: {progress}% | Status: {status}", elapsed) + else: + progress_callback(f"Status: {status} | Elapsed: {elapsed}s", elapsed) + + if status in ("succeeded", "completed", "finished"): + # Download video if output path specified + video_url = _extract_video_url(result) + if output and video_url: + dl_start = time.time() + progress_callback(f"Downloading video to {output}...", elapsed) + _download_video(video_url, output) + dl_elapsed = int(time.time() - dl_start) + progress_callback(f"Downloaded to {Path(output).resolve()} ({dl_elapsed}s)", elapsed + dl_elapsed) + return result + + if status in ("failed", "error"): + return result + + return { + "id": task_id, + "status": "timeout", + "message": f"Polling timed out after {poll_timeout}s", + } + + +def _extract_video_url(result: dict[str, Any]) -> str | None: + """Extract the first video URL from a completed task result.""" + # Direct nested content structure: {"content": {"video_url": "..."}} + content = result.get("content") or {} + if isinstance(content, dict): + url = content.get("video_url") or content.get("url") or content.get("media") or "" + if isinstance(url, str) and url.startswith("http"): + return url + + # data / output wrapper + data = result.get("data") or result.get("output") or {} + if isinstance(data, dict): + media = data.get("media") or data.get("video") or data.get("videos") or [] + if isinstance(media, list) and media: + return media[0].get("url") or media[0].get("video_url") + if isinstance(media, str): + return media + return None + + +def _download_video(url: str, output_path: str) -> str: + """Download video from URL to local path.""" + out = Path(output_path) + out.parent.mkdir(parents=True, exist_ok=True) + urllib.request.urlretrieve(url, out) + return str(out.resolve()) diff --git a/mytoolkit/commands/video.py b/mytoolkit/commands/video.py index 7f39304..8050b11 100644 --- a/mytoolkit/commands/video.py +++ b/mytoolkit/commands/video.py @@ -1,9 +1,10 @@ -"""Video conversion commands.""" +"""Video conversion and generation commands.""" from pathlib import Path import click +from mytoolkit.config import config from mytoolkit.utils import handle_errors, run_command # Subcommands migrated from skills. @@ -13,7 +14,7 @@ from mytoolkit.commands.video_transcribe import transcribe_video_cmd @click.group() def video(): - """Video conversion commands.""" + """Video conversion and generation commands.""" pass @@ -42,3 +43,85 @@ def avi_to_mp4(files, dry_run): ["ffmpeg", "-i", str(avi_file), "-c:v", "libx264", "-c:a", "aac", str(output)], check=False, ) + + +@video.command("generate") +@click.argument("text_prompt", required=False) +@click.option( + "-f", "--file", type=click.Path(exists=True), + help="从文件读取 text prompt", +) +@click.option("-m", "--model", default="doubao-seedance-2-0-260128", help="模型名称") +@click.option("--image-url", multiple=True, help="参考图片 URL(可多次指定)") +@click.option("--video-url", multiple=True, help="参考视频 URL(可多次指定)") +@click.option("--audio-url", multiple=True, help="参考音频 URL(可多次指定)") +@click.option("--no-audio", is_flag=True, help="不生成同步音频") +@click.option("-r", "--ratio", default="16:9", help="画面比例 (如 16:9, 9:16, 1:1)") +@click.option("-d", "--duration", default=11, type=int, help="视频时长(秒)") +@click.option("--watermark", is_flag=True, help="添加豆包水印") +@click.option("-o", "--output", help="下载视频到本地路径(需等待生成完成)") +@click.option("--no-poll", is_flag=True, help="不等待,仅提交任务并返回 task_id") +@click.option("--poll-interval", default=5, type=int, help="轮询间隔(秒)") +@click.option("--poll-timeout", default=600, type=int, help="轮询超时(秒)") +@handle_errors +def generate_video_cmd( + text_prompt, file, model, + image_url, video_url, audio_url, + no_audio, ratio, duration, watermark, + output, no_poll, poll_interval, poll_timeout, +): + """使用豆包/Seedance 模型生成视频。 + + 支持参考图片、参考视频、参考音频,可指定比例和时长。 + """ + if file: + text_prompt = Path(file).read_text(encoding="utf-8").strip() + elif not text_prompt: + raise click.UsageError("必须提供 text prompt 或使用 -f/--file 从文件读取") + + api_key = config.get("secrets.api_keys.ark") + if not api_key: + click.echo( + "Error: secrets.api_keys.ark not set. Run: mytoolkit env set secrets.api_keys.ark ", + err=True, + ) + raise click.Abort() + + from mytoolkit.api.video_gen import generate_video + + def _show_progress(msg: str, elapsed: int) -> None: + mins, secs = divmod(elapsed, 60) + click.echo(f"[{mins:02d}:{secs:02d}] {msg}") + + result = generate_video( + api_key=api_key, + text_prompt=text_prompt, + model=model, + image_urls=list(image_url) if image_url else None, + video_urls=list(video_url) if video_url else None, + audio_urls=list(audio_url) if audio_url else None, + generate_audio=not no_audio, + ratio=ratio, + duration=duration, + watermark=watermark, + poll=not no_poll, + poll_interval=poll_interval, + poll_timeout=poll_timeout, + output=output, + progress_callback=_show_progress, + ) + + status = result.get("status", "unknown") + click.echo(f"\nFinal status: {status}") + if result.get("id"): + click.echo(f"Task ID: {result['id']}") + + # Print video URL if available + if status in ("succeeded", "completed", "finished"): + from mytoolkit.api.video_gen import _extract_video_url + video_url = _extract_video_url(result) + if video_url: + click.echo(f"Video URL: {video_url}") + if output: + click.echo(f"Downloaded to: {Path(output).resolve()}") + click.echo("Done.")