103 lines
3.6 KiB
Python
Executable File
103 lines
3.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
Agnes AI 可用 subprocess curl 调用（可靠版）
|
||
⚠️ 不直接用 requests 库，IPv6 会卡住
|
||
|
||
调用方式：
|
||
python3 agnes_subprocess.py image "a cute rabbit" # 生成图片
|
||
python3 agnes_subprocess.py image "prompt" /tmp/out.png # 下载到指定路径
|
||
python3 agnes_subprocess.py text "say hello" # 文本生成
|
||
python3 agnes_subprocess.py video "a cat playing piano" 5 # 视频（异步，需轮询）
|
||
"""
|
||
import subprocess, json, os, sys, time
|
||
|
||
ENV_FILE = '/home/muc/.hermes/.env'
|
||
|
||
def get_key():
|
||
with open(ENV_FILE) as f:
|
||
for line in f:
|
||
if 'AGNES_API_KEY' in line and not line.startswith('#'):
|
||
return line.strip().split('=', 1)[1]
|
||
return ''
|
||
|
||
def curl_post(path, payload):
|
||
"""用 curl POST，返回 JSON dict"""
|
||
key = get_key()
|
||
result = subprocess.run(
|
||
['curl', '-s', '-X', 'POST',
|
||
f'https://apihub.agnes-ai.com/v1{path}',
|
||
'-H', f'Authorization: Bearer *** + key,
|
||
'-H', 'Content-Type: application/json',
|
||
'-d', json.dumps(payload)],
|
||
capture_output=True, text=True, timeout=60
|
||
)
|
||
return json.loads(result.stdout)
|
||
|
||
def curl_download(url, path):
|
||
"""下载文件"""
|
||
subprocess.run(['curl', '-sL', url, '-o', path], timeout=30)
|
||
|
||
def generate_image(prompt, size='1024x1024', model='agnes-image-2.0-flash'):
|
||
d = curl_post('/images/generations', {
|
||
'model': model, 'prompt': prompt, 'n': 1, 'size': size
|
||
})
|
||
return d.get('data', [{}])[0].get('url', '')
|
||
|
||
def generate_text(prompt, model='agnes-2.5-flash', max_tokens=200):
|
||
d = curl_post('/chat/completions', {
|
||
'model': model,
|
||
'messages': [{'role': 'user', 'content': prompt}],
|
||
'max_tokens': max_tokens
|
||
})
|
||
return d.get('choices', [{}])[0].get('message', {}).get('content', '')
|
||
|
||
def generate_video(prompt, duration=5, poll_interval=5, max_wait=300):
|
||
d = curl_post('/video/generations', {
|
||
'model': 'agnes-video-v2.0', 'prompt': prompt, 'duration': duration
|
||
})
|
||
task_id = d.get('id', '')
|
||
deadline = time.time() + max_wait
|
||
while time.time() < deadline:
|
||
task_d = curl_get(f'/video/generations/{task_id}')
|
||
status = task_d.get('data', {}).get('status', '')
|
||
print(f' status={status}', flush=True)
|
||
if status == 'SUCCESS':
|
||
return task_d.get('data', {}).get('data', {}).get('video_url', '')
|
||
time.sleep(poll_interval)
|
||
return ''
|
||
|
||
def curl_get(path):
|
||
key = get_key()
|
||
result = subprocess.run(
|
||
['curl', '-s',
|
||
f'https://apihub.agnes-ai.com/v1{path}',
|
||
'-H', f'Authorization: Bearer *** + key],
|
||
capture_output=True, text=True, timeout=15
|
||
)
|
||
return json.loads(result.stdout)
|
||
|
||
if __name__ == '__main__':
|
||
if len(sys.argv) < 3:
|
||
print('Usage: agnes_subprocess.py <image|text|video> <prompt> [save_path]')
|
||
sys.exit(1)
|
||
|
||
mode, prompt = sys.argv[1], sys.argv[2]
|
||
save_path = sys.argv[3] if len(sys.argv) > 3 else None
|
||
|
||
if mode == 'image':
|
||
url = generate_image(prompt)
|
||
print(f'Image URL: {url}')
|
||
if url:
|
||
out = save_path or '/tmp/agnes_out.png'
|
||
curl_download(url, out)
|
||
print(f'Saved: {os.path.getsize(out)} bytes → {out}')
|
||
elif mode == 'text':
|
||
print(generate_text(prompt))
|
||
elif mode == 'video':
|
||
url = generate_video(prompt)
|
||
print(f'Video URL: {url}')
|
||
if url and save_path:
|
||
curl_download(url, save_path)
|
||
print(f'Saved: {os.path.getsize(save_path)} bytes → {save_path}')
|
||
else:
|
||
print(f'Unknown mode: {mode}') |