287 lines
8.9 KiB
Python
287 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Agnes 视频生成工具
|
||
==================
|
||
基于 agnes-video-v2.0 模型，文字描述 → 5秒短视频。
|
||
|
||
用法:
|
||
python3 agnes_video.py "一只猫在花园散步"
|
||
python3 agnes_video.py "A cat walking" --width 1280 --height 720
|
||
python3 agnes_video.py "描述" --output /tmp/my-video.mp4
|
||
python3 agnes_video.py "描述" --push # 生成后推飞书
|
||
python3 agnes_video.py status <task_id> # 查询任务状态
|
||
python3 agnes_video.py list # 列出最近任务
|
||
"""
|
||
|
||
import json, sys, os, time, subprocess
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
API_BASE = "https://apihub.agnes-ai.com/v1"
|
||
MODEL = "agnes-video-v2.0"
|
||
OUTPUT_DIR = Path.home() / ".hermes" / "video_cache"
|
||
OUTPUT_DIR.mkdir(exist_ok=True)
|
||
|
||
FEISHU_WEBHOOK = "https://open.feishu.cn/open-apis/bot/v2/hook/446db983-e392-4d2c-bfb8-f9060e5df3ad"
|
||
|
||
|
||
def get_api_key():
|
||
"""从环境变量或.env文件获取API key"""
|
||
key = os.environ.get("AGNES_API_KEY")
|
||
if key:
|
||
return key
|
||
env_file = Path.home() / ".hermes" / ".env"
|
||
if env_file.exists():
|
||
for line in env_file.read_text().splitlines():
|
||
if line.startswith("AGNES_API_KEY="):
|
||
return line.split("=", 1)[1].strip()
|
||
print("❌ 未找到 AGNES_API_KEY，请设置环境变量或写入 ~/.hermes/.env")
|
||
sys.exit(1)
|
||
|
||
|
||
def api_request(method, path, data=None, timeout=15):
|
||
"""发送 API 请求"""
|
||
key = get_api_key()
|
||
url = f"{API_BASE}{path}"
|
||
|
||
cmd = ["curl", "-s", "--max-time", str(timeout), "-X", method, url,
|
||
"-H", f"Authorization: Bearer {key}",
|
||
"-H", "Content-Type: application/json"]
|
||
if data:
|
||
cmd.extend(["-d", json.dumps(data)])
|
||
|
||
env = dict(os.environ)
|
||
for k in ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"]:
|
||
env.pop(k, None)
|
||
|
||
r = subprocess.run(cmd, capture_output=True, timeout=timeout+5, env=env)
|
||
text = r.stdout.decode("utf-8", errors="ignore")
|
||
if not text.strip():
|
||
return None
|
||
try:
|
||
return json.loads(text)
|
||
except json.JSONDecodeError:
|
||
return {"raw": text}
|
||
|
||
|
||
def submit_video(prompt, duration=5, width=1280, height=720):
|
||
"""提交视频生成任务"""
|
||
data = {
|
||
"model": MODEL,
|
||
"prompt": prompt,
|
||
"duration": duration,
|
||
"width": width,
|
||
"height": height,
|
||
}
|
||
result = api_request("POST", "/videos", data)
|
||
if not result:
|
||
print("❌ API 请求失败")
|
||
return None
|
||
|
||
task_id = result.get("task_id") or result.get("id")
|
||
status = result.get("status", "unknown")
|
||
print(f"✅ 任务已提交")
|
||
print(f" Task ID: {task_id}")
|
||
print(f" 状态: {status}")
|
||
print(f" 提示词: {prompt}")
|
||
return task_id
|
||
|
||
|
||
def poll_status(task_id, max_wait=180, interval=10):
|
||
"""轮询任务状态直到完成"""
|
||
print(f"\n⏳ 等待生成（最多{max_wait}秒）...")
|
||
start = time.time()
|
||
|
||
while time.time() - start < max_wait:
|
||
result = api_request("GET", f"/videos/{task_id}")
|
||
if not result:
|
||
time.sleep(interval)
|
||
continue
|
||
|
||
status = result.get("status", "unknown")
|
||
progress = result.get("progress", 0)
|
||
|
||
if status == "completed":
|
||
print(f"\n✅ 生成完成！({int(time.time()-start)}秒)")
|
||
return result
|
||
elif status == "failed":
|
||
print(f"\n❌ 生成失败: {result.get('error', '未知错误')}")
|
||
return None
|
||
else:
|
||
bar = "█" * (progress // 5) + "░" * (20 - progress // 5)
|
||
print(f"\r [{bar}] {progress}% {status}", end="", flush=True)
|
||
time.sleep(interval)
|
||
|
||
print(f"\n❌ 超时（{max_wait}秒）")
|
||
return None
|
||
|
||
|
||
def download_video(url, output_path=None):
|
||
"""下载视频文件"""
|
||
if not output_path:
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
output_path = OUTPUT_DIR / f"agnes_{timestamp}.mp4"
|
||
else:
|
||
output_path = Path(output_path)
|
||
|
||
cmd = ["curl", "-s", "--max-time", "60", "-o", str(output_path), url]
|
||
env = dict(os.environ)
|
||
for k in ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"]:
|
||
env.pop(k, None)
|
||
|
||
r = subprocess.run(cmd, capture_output=True, timeout=65, env=env)
|
||
if r.returncode == 0 and output_path.exists():
|
||
size_mb = output_path.stat().st_size / 1024 / 1024
|
||
print(f"📁 已保存: {output_path} ({size_mb:.1f}MB)")
|
||
return str(output_path)
|
||
else:
|
||
print("❌ 下载失败")
|
||
return None
|
||
|
||
|
||
def push_feishu(video_path, prompt):
|
||
"""推送到飞书（发送视频文件）"""
|
||
# 飞书 webhook 只支持文本/图片/富文本，不支持视频文件
|
||
# 改为发送文本通知 + 本地路径
|
||
msg = f"""🎬 Agnes 视频生成完成
|
||
|
||
提示词: {prompt}
|
||
文件: {video_path}
|
||
大小: {Path(video_path).stat().st_size / 1024 / 1024:.1f}MB
|
||
|
||
用 MEDIA:{video_path} 发送视频"""
|
||
|
||
payload = json.dumps({"msg_type": "text", "content": {"text": msg}}).encode()
|
||
cmd = ["curl", "-s", "--max-time", "10", "-X", "POST", FEISHU_WEBHOOK,
|
||
"-H", "Content-Type: application/json", "-d", payload.decode()]
|
||
env = dict(os.environ)
|
||
for k in ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"]:
|
||
env.pop(k, None)
|
||
subprocess.run(cmd, capture_output=True, timeout=12, env=env)
|
||
print("📤 已推送飞书通知")
|
||
|
||
|
||
def check_status(task_id):
|
||
"""查询任务状态"""
|
||
result = api_request("GET", f"/videos/{task_id}")
|
||
if not result:
|
||
print("❌ 查询失败")
|
||
return
|
||
|
||
print(f"Task ID: {task_id}")
|
||
print(f"状态: {result.get('status', '?')} ({result.get('progress', 0)}%)")
|
||
print(f"模型: {result.get('model', '?')}")
|
||
print(f"时长: {result.get('seconds', '?')}秒")
|
||
print(f"尺寸: {result.get('size', '?')}")
|
||
|
||
if result.get("status") == "completed":
|
||
url = result.get("metadata", {}).get("url", "")
|
||
if url:
|
||
print(f"下载: {url}")
|
||
dl = input("下载到本地？(y/n): ").strip().lower()
|
||
if dl == "y":
|
||
download_video(url)
|
||
|
||
|
||
def generate_video(prompt, duration=5, width=1280, height=720, output=None, push=False):
|
||
"""完整流程：提交 → 等待 → 下载"""
|
||
print(f"🎬 Agnes 视频生成")
|
||
print(f" 模型: {MODEL}")
|
||
print(f" 时长: {duration}秒")
|
||
print(f" 分辨率: {width}×{height}")
|
||
print(f" 提示词: {prompt}")
|
||
print()
|
||
|
||
# 1. 提交
|
||
task_id = submit_video(prompt, duration, width, height)
|
||
if not task_id:
|
||
return None
|
||
|
||
# 2. 轮询
|
||
result = poll_status(task_id)
|
||
if not result:
|
||
return None
|
||
|
||
# 3. 下载
|
||
url = result.get("metadata", {}).get("url", "")
|
||
if not url:
|
||
print("❌ 未找到视频下载链接")
|
||
print(f" 完整响应: {json.dumps(result, indent=2)[:500]}")
|
||
return None
|
||
|
||
video_path = download_video(url, output)
|
||
if not video_path:
|
||
return None
|
||
|
||
# 4. 信息
|
||
print(f"\n📊 生成信息:")
|
||
print(f" 尺寸: {result.get('size', '?')}")
|
||
print(f" 时长: {result.get('seconds', '?')}秒")
|
||
print(f" 耗时: {result.get('completed_at', 0) - result.get('created_at', 0)}秒")
|
||
|
||
# 5. 推飞书
|
||
if push:
|
||
push_feishu(video_path, prompt)
|
||
|
||
return video_path
|
||
|
||
|
||
def list_cache():
|
||
"""列出本地缓存的视频"""
|
||
videos = sorted(OUTPUT_DIR.glob("agnes_*.mp4"), reverse=True)
|
||
if not videos:
|
||
print("📭 暂无缓存视频")
|
||
return
|
||
print(f"📁 视频缓存 ({len(videos)}个):")
|
||
for v in videos[:20]:
|
||
size_mb = v.stat().st_size / 1024 / 1024
|
||
mtime = datetime.fromtimestamp(v.stat().st_mtime).strftime("%m-%d %H:%M")
|
||
print(f" {mtime} {size_mb:.1f}MB {v.name}")
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) < 2:
|
||
print("用法:")
|
||
print(' python3 agnes_video.py "提示词" [--width 1280] [--height 720] [--output x.mp4] [--push]')
|
||
print(" python3 agnes_video.py status <task_id>")
|
||
print(" python3 agnes_video.py list")
|
||
return
|
||
|
||
if sys.argv[1] == "status" and len(sys.argv) >= 3:
|
||
check_status(sys.argv[2])
|
||
return
|
||
|
||
if sys.argv[1] == "list":
|
||
list_cache()
|
||
return
|
||
|
||
# 解析参数
|
||
prompt = sys.argv[1]
|
||
width = 1280
|
||
height = 720
|
||
output = None
|
||
push = False
|
||
duration = 5
|
||
|
||
args = sys.argv[2:]
|
||
i = 0
|
||
while i < len(args):
|
||
if args[i] == "--width" and i + 1 < len(args):
|
||
width = int(args[i + 1]); i += 2
|
||
elif args[i] == "--height" and i + 1 < len(args):
|
||
height = int(args[i + 1]); i += 2
|
||
elif args[i] == "--duration" and i + 1 < len(args):
|
||
duration = int(args[i + 1]); i += 2
|
||
elif args[i] == "--output" and i + 1 < len(args):
|
||
output = args[i + 1]; i += 2
|
||
elif args[i] == "--push":
|
||
push = True; i += 1
|
||
else:
|
||
i += 1
|
||
|
||
generate_video(prompt, duration, width, height, output, push)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|