34 lines
1.8 KiB
Python
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)
|