champions-of-luneraria/tools/meshy/meshy.py

"""Small Meshy API client for the U5.5 figure pilot (standard library only).

The API key is read from ~/.config/luneraria/meshy.key (or the file named by LUNE_MESHY_KEY) and
is never printed, logged or written anywhere. Every created task is appended to a ledger
(tools/out/meshy/ledger.jsonl: time, endpoint, task id, credits) so the spend stays traceable.
Downloads go into a fresh folder per task; signed URLs are stored without their query string.

  meshy.py balance
  meshy.py library [search]                       animation library (free)
  meshy.py create <endpoint> <request.json>       prints the task id
  meshy.py wait <endpoint> <task id>              polls until the task ends, prints the task
  meshy.py fetch <endpoint> <task id> <folder>    downloads every result file of a finished task

<endpoint> is the API path below /openapi/, e.g. v1/image-to-image, v1/image-to-3d, v1/rigging,
v1/animations, v1/text-to-motion. Request files may name local images as "@path/to/image.png";
they are sent as data URIs.
"""
import base64, json, os, sys, time, urllib.error, urllib.parse, urllib.request

API = 'https://api.meshy.ai/openapi/'
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
LEDGER = os.path.join(ROOT, 'tools', 'out', 'meshy', 'ledger.jsonl')


def key():
    path = os.environ.get('LUNE_MESHY_KEY', os.path.expanduser('~/.config/luneraria/meshy.key'))
    with open(path) as handle:
        return handle.read().strip()


def call(method, path, body=None):
    data = json.dumps(body).encode() if body is not None else None
    request = urllib.request.Request(API + path, data=data, method=method)
    request.add_header('Authorization', 'Bearer ' + key())
    if data is not None:
        request.add_header('Content-Type', 'application/json')
    try:
        with urllib.request.urlopen(request, timeout=120) as response:
            return json.loads(response.read().decode() or 'null')
    except urllib.error.HTTPError as error:
        detail = error.read().decode(errors='replace')[:800]
        sys.exit('Meshy %s %s: HTTP %d %s' % (method, path, error.code, detail))


def inline(value):
    """"@file" strings become data URIs (recursively in lists and objects)."""
    if isinstance(value, str) and value.startswith('@'):
        path = value[1:]
        mime = 'image/png' if path.lower().endswith('.png') else 'image/jpeg'
        with open(path, 'rb') as handle:
            return 'data:%s;base64,%s' % (mime, base64.b64encode(handle.read()).decode())
    if isinstance(value, list):
        return [inline(v) for v in value]
    if isinstance(value, dict):
        return {k: inline(v) for k, v in value.items()}
    return value


def unsigned(value):
    """Task objects without signatures in their URLs and without inline images."""
    if isinstance(value, str):
        if value.startswith('data:'):
            return value[:40] + '...'
        if value.startswith('http'):
            return value.split('?', 1)[0]
        return value
    if isinstance(value, list):
        return [unsigned(v) for v in value]
    if isinstance(value, dict):
        return {k: unsigned(v) for k, v in value.items()}
    return value


def ledger(entry):
    os.makedirs(os.path.dirname(LEDGER), exist_ok=True)
    entry['time'] = time.strftime('%Y-%m-%dT%H:%M:%S')
    with open(LEDGER, 'a') as handle:
        handle.write(json.dumps(entry) + '\n')


def urls(value, trail=''):
    """(name, url) of every downloadable result in a task object."""
    found = []
    if isinstance(value, dict):
        for k, v in value.items():
            found += urls(v, (trail + '.' if trail else '') + k)
    elif isinstance(value, list):
        for i, v in enumerate(value):
            found += urls(v, '%s.%d' % (trail, i))
    elif isinstance(value, str) and value.startswith('https://') and 'meshy' in urllib.parse.urlparse(value).netloc:
        found.append((trail, value))
    return found


def main(args):
    if not args:
        sys.exit(__doc__)
    command = args[0]
    if command == 'balance':
        print(json.dumps(call('GET', 'v1/balance')))
    elif command == 'library':
        query = '?search=' + urllib.parse.quote(args[1]) if len(args) > 1 else ''
        for item in call('GET', 'v1/animations/library' + query):
            print('%5s  %-14s %-18s %-40s %s' % (item.get('action_id'), item.get('category'), item.get('sub_category'), item.get('name'), item.get('key')))
    elif command == 'create':
        endpoint, request_file = args[1], args[2]
        with open(request_file) as handle:
            body = json.load(handle)
        result = call('POST', endpoint, inline(body))
        task = result.get('result')
        ledger({'endpoint': endpoint, 'task': task, 'request': unsigned(body), 'request_file': os.path.relpath(request_file, ROOT)})
        print(task)
    elif command == 'wait':
        endpoint, task = args[1], args[2]
        while True:
            obj = call('GET', endpoint + '/' + task)
            status = obj.get('status')
            if status in ('SUCCEEDED', 'FAILED', 'CANCELED'):
                ledger({'endpoint': endpoint, 'task': task, 'status': status, 'credits': obj.get('consumed_credits'), 'error': (obj.get('task_error') or {}).get('message')})
                print(json.dumps(unsigned(obj), indent=1))
                return 0 if status == 'SUCCEEDED' else 1
            sys.stderr.write('%s %s %s%%\n' % (task, status, obj.get('progress')))
            time.sleep(10)
    elif command == 'fetch':
        endpoint, task, folder = args[1], args[2], args[3]
        obj = call('GET', endpoint + '/' + task)
        if obj.get('status') != 'SUCCEEDED':
            sys.exit('task %s is %s' % (task, obj.get('status')))
        os.makedirs(folder, exist_ok=True)
        with open(os.path.join(folder, 'task.json'), 'w') as handle:
            json.dump(unsigned(obj), handle, indent=1)
        for name, url in urls(obj):
            if name in ('texture_image_url',) or name.startswith('request'):
                continue
            extension = os.path.splitext(urllib.parse.urlparse(url).path)[1] or '.bin'
            target = os.path.join(folder, name.replace('.', '-') + extension)
            with urllib.request.urlopen(url, timeout=300) as response, open(target, 'wb') as handle:
                handle.write(response.read())
            print('%9d  %s' % (os.path.getsize(target), os.path.relpath(target, ROOT)))
    else:
        sys.exit(__doc__)
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))