ape2e/tools/retrain_instructors_v2.py

"""Reproducible v2 instructor generation; reports contain no credentials."""
from pathlib import Path
import argparse
import gzip
import hashlib
import json
import re
import shutil
import requests

ROOT = Path(__file__).resolve().parents[3]
REPORT = ROOT / 'ai_person/ape2e/reports/instructors-v2-20260914'
BUNDLES = ROOT / 'ai_person/ai_person_service/builtin_avatars'
SERVICE = 'http://1.14.103.234:8123'
API = 'http://127.0.0.1:8001'
FILES = ('01.mp4', 'combined_data.json.gz', 'thumbnail.jpg', 'manifest.json')
IDS = [f'a100000000000000000000000000000{i}' for i in (5, 6, 7)]
HTTP = requests.Session()
HTTP.trust_env = False


def save(name, value):
    REPORT.mkdir(parents=True, exist_ok=True)
    (REPORT / name).write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding='utf-8')


def api_session():
    values = {}
    for line in (ROOT / 'ai_person/账号.MD').read_text(encoding='utf-8-sig').splitlines():
        match = re.match(r'^\s*([^:：]+)\s*[:：]\s*(.*?)\s*$', line)
        if match:
            values[match[1].strip()] = match[2].strip()
    session = requests.Session()
    session.trust_env = False
    response = session.post(API + '/api/auth/v1/auth/login', json={
        'roleCode': 'admin', 'username': values.get('账号') or values.get('用户名'),
        'password': values['密码'], 'rememberMe': False,
    }, timeout=30)
    response.raise_for_status()
    data = response.json()['data']
    token = data.get('accessToken') or data['tokens']['accessToken']
    session.headers['Authorization'] = 'Bearer ' + token
    return session


def get_json(url):
    response = HTTP.get(url, timeout=30)
    response.raise_for_status()
    return response.json()


def snapshot():
    if (REPORT / 'before-platform.json').exists():
        raise RuntimeError('Backup already exists; refusing to overwrite')
    session = api_session()
    response = session.get(API + '/api/v1/avatars', params={'pageSize': 200}, timeout=30)
    response.raise_for_status()
    save('before-platform.json', response.json())
    save('before-service.json', [get_json(SERVICE + '/api/v1/avatars/' + ident) for ident in IDS])
    for ident in IDS:
        shutil.copytree(BUNDLES / ident, REPORT / 'backup-local' / ident)
    print('Backed up platform catalog, service metadata and all three local bundles')


def submit():
    if not (REPORT / 'before-platform.json').exists():
        raise RuntimeError('Run snapshot first')
    state = json.loads((REPORT / 'jobs.json').read_text(encoding='utf-8')) if (REPORT / 'jobs.json').exists() else []
    for number, ident in enumerate(IDS, 1):
        if any(item['stable_id'] == ident for item in state):
            continue
        source = ROOT / '数字人素材' / f'教员{number}_v2.mp4'
        with source.open('rb') as stream:
            response = HTTP.post(SERVICE + '/api/v1/avatars',
                files={'video': (source.name, stream, 'video/mp4')},
                data={'name': f'教员{number} v2 重训 20260914', 'resize': 'true', 'matting': 'off'}, timeout=180)
        response.raise_for_status()
        item = {'number': number, 'stable_id': ident, 'source': source.name,
                'source_sha256': hashlib.sha256(source.read_bytes()).hexdigest(), **response.json()}
        state.append(item)
        save('jobs.json', state)
        print(json.dumps(item, ensure_ascii=False), flush=True)


def status():
    values = []
    for item in json.loads((REPORT / 'jobs.json').read_text(encoding='utf-8')):
        job = get_json(SERVICE + item['job_url'])
        values.append({'number': item['number'], **job})
    save('job-status.json', values)
    print(json.dumps(values, ensure_ascii=False, indent=2))


def download():
    values = []
    for item in json.loads((REPORT / 'jobs.json').read_text(encoding='utf-8')):
        avatar = get_json(SERVICE + '/api/v1/avatars/' + item['avatar_id'])
        if avatar['status'] != 'ready':
            raise RuntimeError(f"Instructor {item['number']} not ready")
        target = REPORT / 'generated' / item['stable_id'] / 'assets'
        target.mkdir(parents=True, exist_ok=True)
        for filename in FILES:
            response = HTTP.get(SERVICE + f"/api/v1/avatars/{item['avatar_id']}/assets/{filename}", timeout=120)
            response.raise_for_status()
            (target / filename).write_bytes(response.content)
        manifest = json.loads((target / 'manifest.json').read_text(encoding='utf-8'))
        for filename, length in manifest['files'].items():
            assert (target / filename).stat().st_size == length, filename
        with gzip.open(target / 'combined_data.json.gz', 'rt', encoding='utf-8') as stream:
            geometry = json.load(stream)
        values.append({'number': item['number'], 'stable_id': item['stable_id'], 'generated_id': item['avatar_id'],
                       'manifest': manifest, 'geometry_type': type(geometry).__name__,
                       'sha256': {name: hashlib.sha256((target / name).read_bytes()).hexdigest() for name in FILES}})
    save('generated-validation.json', values)
    print(json.dumps(values, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('phase', choices=['snapshot', 'submit', 'status', 'download'])
    globals()[parser.parse_args().phase]()