Files
dj/integration-work/verify_gpu1_image_live.py

34 lines
1.8 KiB
Python

import json,sys,time,urllib.request
from pathlib import Path
from PIL import Image
root=Path(__file__).resolve().parent
state=root/'gpu1-image-verification-20261109.json'
cfg=json.loads((root.parent/'model-bridge/client.json').read_text(encoding='utf-8-sig'))
base=cfg['url'].rstrip('/')
def request(path,payload=None):
data=None if payload is None else json.dumps(payload).encode()
req=urllib.request.Request(base+path,data=data,headers={'Authorization':'Bearer '+cfg['token'],'Content-Type':'application/json'})
with urllib.request.urlopen(req,timeout=45) as response:return json.load(response)
if sys.argv[1]=='submit':
assert not state.exists(),'Verification task already exists; poll it instead.'
result=request('/v1/images/generations',{'model':'qwen_image_2_1','prompt':'A studio photograph of a red ceramic cup on a wooden table, soft daylight, no writing.','size':'512x512','gpu':'gpu1','n':1})
assert result['provider_id']=='gpu1_8189'
state.write_text(json.dumps(result,indent=2))
print(json.dumps(result))
else:
saved=json.loads(state.read_text());jid=saved['task_id']
for _ in range(8):
result=request('/v1/images/task/'+jid)
if result['status']=='completed':
path=root/'gpu1-8189-image-verification.png'
with urllib.request.urlopen(result['image_url'],timeout=30) as response:path.write_bytes(response.read())
with Image.open(path) as im:
im.verify()
with Image.open(path) as im:size=im.size
assert size==(512,512)
saved.update(status='completed',provider_id=result['provider_id'],dimensions=list(size),output=str(path))
state.write_text(json.dumps(saved,indent=2));print(json.dumps(saved));break
if result['status']=='failed':raise RuntimeError(result)
print(json.dumps({'status':result['status'],'provider':result.get('provider_id')}),flush=True)
time.sleep(5)