"""Generate the 100 A-gritty-colour images through the official OpenMOSS-Team/MOSS-Image Space. Every image is tracked separately. State is written after every step: manifest.json one record per image (prompt, payload, every attempt, final file + hashes) attempts.jsonl append-only log of every attempt, successful or failed originals/NNN-*.png the untouched PNG returned by the Space provenance/NNN-*.json sidecar written the moment an image is saved failed/ anything downloaded that failed validation (kept, never deleted) Usage: gen.py [--only 1,2,3] [--workers N] [--max-attempts N] """ import argparse, hashlib, json, os, re, threading, time, datetime from concurrent.futures import ThreadPoolExecutor from pathlib import Path import requests from PIL import Image, ImageStat from huggingface_hub import get_token, HfApi from scenes import SCENES ROOT = Path(__file__).resolve().parent ORIG, PROV, FAILED = ROOT / 'originals', ROOT / 'provenance', ROOT / 'failed' MANIFEST, ATTEMPTS = ROOT / 'manifest.json', ROOT / 'attempts.jsonl' SPACE_ID = 'OpenMOSS-Team/MOSS-Image' H = 'https://openmoss-team-moss-image.hf.space' SEED_BASE = 70000 W, HGT = 1536, 1024 # --- Style and settings, taken verbatim from the A-gritty-colour records in source-manifest.json --- src = json.loads((ROOT / 'source-manifest.json').read_text()) src_a = [r for r in src if r['direction'] == 'A-gritty-colour'] assert len(src_a) == 5 _p0 = src_a[0]['payload']['param_1'] STYLE = _p0.split('\n\nART DIRECTION: ', 1)[1].split('\nColour motif: ', 1)[0] TAIL = 'One coherent cinematic image, atmospheric storytelling with clear visual hierarchy and authentic human emotion. No recognizable existing actors or franchise characters. No typography, captions, logos, watermarks, borders or collage.' SETTINGS = {k: v for k, v in src_a[0]['payload'].items() if k not in ('param_1', 'param_8')} def build_prompt(description, motif): # Same construction as source-generate.py for direction A (col 0 gets the colour motif). return f'{description}\n\nART DIRECTION: {STYLE}\nColour motif: {motif}. {TAIL}' # Prove the reconstruction is exact against every source A record, and that settings are shared. _src_scene = { '01-orchid-empire': 'sickly greenhouse green', '02-spin-cycle': 'a small vivid red bouquet', '03-bird-market': 'muted amber pigeon-ring boxes', '04-carwash-wake': 'a luminous orange peel', '05-fishmarket-dawn': 'a tiny golden fish'} for r in src_a: desc = r['payload']['param_1'].split('\n\nART DIRECTION: ', 1)[0] assert build_prompt(desc, _src_scene[r['scene']]) == r['payload']['param_1'], r['name'] assert {k: v for k, v in r['payload'].items() if k not in ('param_1', 'param_8')} == SETTINGS, r['name'] assert SETTINGS['selected'] == 'OpenMOSS-Team/MOSS-Image-SFT' assert (SETTINGS['param_3'], SETTINGS['param_4'], SETTINGS['param_5']) == (HGT, W, 50) lock = threading.Lock() quota_until = [0.0] # shared back-off when the ZeroGPU quota is exhausted def now(): return datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='seconds') def atomic_write(path, text): tmp = path.with_suffix(path.suffix + '.tmp') tmp.write_text(text) os.replace(tmp, path) def save_manifest(manifest): with lock: atomic_write(MANIFEST, json.dumps(manifest, indent=2, ensure_ascii=False)) def log_attempt(entry): with lock: with ATTEMPTS.open('a') as f: f.write(json.dumps(entry, ensure_ascii=False) + '\n') def init_manifest(): if MANIFEST.exists(): manifest = json.loads(MANIFEST.read_text()) else: manifest = [] have = {r['index'] for r in manifest} for i, (slug, motif, desc) in enumerate(SCENES, 1): if i in have: continue name = f'{i:03d}-{slug}' payload = dict(SETTINGS) payload['param_1'] = build_prompt(desc, motif) payload['param_8'] = SEED_BASE + i payload = {k: payload[k] for k in ['selected'] + [f'param_{n}' for n in range(1, 15)]} manifest.append({'index': i, 'name': name, 'slug': slug, 'direction': 'A-gritty-colour', 'colour_motif': motif, 'scene': desc, 'seed': payload['param_8'], 'payload': payload, 'status': 'pending', 'attempts': []}) manifest.sort(key=lambda r: r['index']) # Every stored prompt must still match the current scene list. for r in manifest: slug, motif, desc = SCENES[r['index'] - 1] assert r['slug'] == slug and r['payload']['param_1'] == build_prompt(desc, motif), r['name'] save_manifest(manifest) return manifest def validate_png(path): with Image.open(path) as im: im.verify() with Image.open(path) as im: im.load() info = {'format': im.format, 'mode': im.mode, 'width': im.width, 'height': im.height} stat = ImageStat.Stat(im.convert('RGB')) info['channel_stddev'] = [round(x, 2) for x in stat.stddev] info['channel_mean'] = [round(x, 2) for x in stat.mean] problems = [] if info['format'] != 'PNG': problems.append(f"format {info['format']}") if (info['width'], info['height']) != (W, HGT): problems.append(f"size {info['width']}x{info['height']}") if max(info['channel_stddev']) < 8: problems.append('near-blank image') return info, problems def parse_quota_wait(msg): m = re.search(r'(\d+):(\d{2}):(\d{2})', msg) if m: h, mi, s = map(int, m.groups()) return h * 3600 + mi * 60 + s return None def run_one(rec, manifest, session, max_attempts): idx, name, payload = rec['index'], rec['name'], rec['payload'] final = ORIG / f'{name}.png' failures = sum(1 for a in rec['attempts'] if a.get('error') and not a.get('quota')) while rec['status'] != 'ok': if failures >= max_attempts: rec['status'] = 'failed' save_manifest(manifest) print(f'GAVE UP {idx:03d} after {failures} failures', flush=True) return wait = quota_until[0] - time.time() if wait > 0: time.sleep(wait) a = {'attempt': len(rec['attempts']) + 1, 'started_at': now()} rec['attempts'].append(a) rec['status'] = 'running' save_manifest(manifest) t0 = time.monotonic() try: r = session.post(H + '/gradio_api/call/v2/generate', json=payload, timeout=60) r.raise_for_status() eid = r.json()['event_id'] a['event_id'] = eid save_manifest(manifest) event, result = '', None with session.get(H + '/gradio_api/call/generate/' + eid, stream=True, timeout=(30, 900)) as st: st.raise_for_status() for line in st.iter_lines(decode_unicode=True): if not line: continue if line.startswith('event:'): event = line[6:].strip() elif line.startswith('data:') and event in ('complete', 'error'): data = json.loads(line[5:].strip()) if event == 'error': raise RuntimeError(f'Space error event: {data}') result = data break if result is None: raise RuntimeError(f'SSE ended without a completed image (last event {event!r})') a['result'] = result status_text = result[1] if len(result) > 1 else '' a['space_status'] = status_text if f"seed {payload['param_8']}" not in str(status_text): raise RuntimeError(f'seed mismatch in Space status: {status_text!r}') image = result[0][0]['image'] url = image.get('url') or H + '/gradio_api/file=' + image['path'] a['image_url'] = url rr = session.get(url, timeout=180) rr.raise_for_status() part = ORIG / f'{name}.png.part' part.write_bytes(rr.content) info, problems = validate_png(part) a['validation'] = info if problems: keep = FAILED / f"{name}-attempt{a['attempt']}.png" os.replace(part, keep) a['kept_invalid_file'] = str(keep.relative_to(ROOT)) raise RuntimeError('invalid image: ' + ', '.join(problems)) os.replace(part, final) data = final.read_bytes() rec.update({ 'status': 'ok', 'file': str(final.relative_to(ROOT)), 'sha256': hashlib.sha256(data).hexdigest(), 'bytes': len(data), 'width': info['width'], 'height': info['height'], 'generated_at': now(), 'event_id': eid, 'space_status': status_text, }) a['seconds'] = round(time.monotonic() - t0, 2) a['finished_at'] = now() prov = {k: rec[k] for k in ('index', 'name', 'slug', 'direction', 'colour_motif', 'scene', 'seed', 'file', 'sha256', 'bytes', 'width', 'height', 'generated_at', 'event_id', 'space_status')} prov.update({'space': SPACE_ID, 'space_sha': SPACE_SHA, 'endpoint': '/generate', 'model': payload['selected'], 'prompt': payload['param_1'], 'negative_prompt': payload['param_2'], 'payload': payload, 'settings': {'height': payload['param_3'], 'width': payload['param_4'], 'steps': payload['param_5'], 'cfg': payload['param_6'], 'shift': payload['param_7'], 'seed': payload['param_8'], 'images': payload['param_9'], 'skip_layers': payload['param_10'], 'slg_layers': payload['param_11'], 'slg_scale': payload['param_12'], 'skip_start': payload['param_13'], 'skip_stop': payload['param_14']}, 'attempt_count': len(rec['attempts'])}) atomic_write(PROV / f'{name}.json', json.dumps(prov, indent=2, ensure_ascii=False)) save_manifest(manifest) log_attempt({'index': idx, 'name': name, **a, 'ok': True}) print(f'OK {idx:03d} {name} {a["seconds"]}s attempt {a["attempt"]}', flush=True) except Exception as e: msg = str(e) a['error'] = msg[:2000] a['seconds'] = round(time.monotonic() - t0, 2) a['finished_at'] = now() low = msg.lower() if 'quota' in low: a['quota'] = True else: failures += 1 # Persist the failed attempt before any back-off sleep. rec['status'] = 'retrying' save_manifest(manifest) log_attempt({'index': idx, 'name': name, **a, 'ok': False}) if a.get('quota'): wait = parse_quota_wait(msg) wait = (wait + 60) if wait else 900 with lock: quota_until[0] = max(quota_until[0], time.time() + wait) print(f'QUOTA {idx:03d}: waiting {wait}s :: {msg[:300]}', flush=True) else: print(f'ERROR {idx:03d} attempt {a["attempt"]}: {msg[:300]}', flush=True) time.sleep(min(30 * failures, 180)) def main(): ap = argparse.ArgumentParser() ap.add_argument('--only', default='') ap.add_argument('--workers', type=int, default=1) ap.add_argument('--max-attempts', type=int, default=6) args = ap.parse_args() for d in (ORIG, PROV, FAILED): d.mkdir(exist_ok=True) global SPACE_SHA SPACE_SHA = HfApi().space_info(SPACE_ID).sha manifest = init_manifest() # Any image whose file vanished or no longer matches its hash goes back to pending. for r in manifest: if r['status'] == 'ok': p = ROOT / r['file'] if not p.exists() or hashlib.sha256(p.read_bytes()).hexdigest() != r['sha256']: r['status'] = 'pending' elif r['status'] in ('running', 'retrying', 'failed'): r['status'] = 'pending' save_manifest(manifest) only = {int(x) for x in args.only.split(',') if x.strip()} todo = [r for r in manifest if r['status'] != 'ok' and (not only or r['index'] in only)] print(f'space sha {SPACE_SHA}; {len(todo)} to generate with {args.workers} worker(s)', flush=True) session = requests.Session() token = get_token() if token: session.headers['Authorization'] = 'Bearer ' + token with ThreadPoolExecutor(args.workers) as ex: for f in [ex.submit(run_one, r, manifest, session, args.max_attempts) for r in todo]: f.result() ok = sum(r['status'] == 'ok' for r in manifest) print(f'DONE ok={ok}/100 failed_attempts={sum(1 for r in manifest for a in r["attempts"] if a.get("error"))}', flush=True) if __name__ == '__main__': main()