75 lines
9.3 KiB
JavaScript
75 lines
9.3 KiB
JavaScript
/* Run on the application host; credentials stay in memory and reports contain no secrets. */
|
|
const fs=require('node:fs');const path=require('node:path');const {randomUUID,createHash}=require('node:crypto');
|
|
require('dotenv').config({path:'.env.local',quiet:true});require('dotenv').config({path:'.env',quiet:true});
|
|
const sharp=require('sharp');const postgres=require('postgres');const {S3Client,PutObjectCommand}=require('@aws-sdk/client-s3');
|
|
const out=path.resolve(process.env.PROBE_DIR||'tmp/vip98k-matrix-0911');fs.mkdirSync(out,{recursive:true});
|
|
const delay=ms=>new Promise(r=>setTimeout(r,ms));const sanitize=s=>String(s).replace(/https?:\/\/[^\s"<]+/g,'[url]').slice(0,1800);
|
|
const ratios=(process.env.PROBE_RATIOS||'1:1').split(',');const tiers=(process.env.PROBE_TIERS||'1K,2K,4K').split(',');
|
|
function urls(x){if(!x||typeof x!=='object')return [];if(Array.isArray(x))return x.flatMap(urls);return Object.entries(x).flatMap(([k,v])=>['url','image_url','file_download_url','download_url','output_url'].includes(k)&&typeof v==='string'&&v.startsWith('http')?[v]:typeof v==='object'?urls(v):[]);}
|
|
function b64(x){if(!x||typeof x!=='object')return null;if(x.b64_json)return x.b64_json;for(const v of Object.values(x)){const s=b64(v);if(s)return s;}return null;}
|
|
(async()=>{
|
|
const sql=postgres(process.env.POSTGRES_URL,{max:1});const [p]=await sql`select "apiKey","baseUrl" from "AiProvider" where type='vip98k' and "isEnabled"=true limit 1`;
|
|
const [s]=await sql`select * from "StorageConfig" where "isDefault"=true and "isEnabled"=true limit 1`;await sql.end();
|
|
const base=p.baseUrl.replace(/\/+$/,'').replace(/\/v1$/,'');
|
|
async function request(url,body,headers={},method='POST'){
|
|
const r=await fetch(url,{method,headers:{Authorization:'Bearer '+p.apiKey,...(body&&! (body instanceof FormData)?{'Content-Type':'application/json'}:{}),...headers},body:body instanceof FormData?body:body?JSON.stringify(body):undefined,signal:AbortSignal.timeout(240000)});
|
|
const txt=await r.text();let data;try{data=JSON.parse(txt)}catch{data={message:txt}};
|
|
if(!r.ok){const err=new Error(sanitize(data.error?.message||data.message||txt));err.status=r.status;err.requestId=r.headers.get('x-request-id');throw err;}
|
|
return {data,requestId:r.headers.get('x-request-id')};
|
|
}
|
|
const catalog=(await request(base+'/api/v2/catalog/models',null,{},'GET')).data;
|
|
const auth=(await request(base+'/v1/models',null,{},'GET')).data;
|
|
const allowed=new Set(auth.data.map(m=>m.id));let models=catalog.data.sections.flatMap(x=>x.models||[]).filter(m=>allowed.has(m.id)&&(/^(gpt-image|nano-banana)/.test(m.id)));
|
|
if(process.env.PROBE_MODELS){const ids=new Set(process.env.PROBE_MODELS.split(','));models=models.filter(m=>ids.has(m.id));}
|
|
const client=new S3Client({endpoint:s.endpoint,region:s.region||'auto',forcePathStyle:true,credentials:{accessKeyId:s.accessKeyId,secretAccessKey:s.secretAccessKey}});
|
|
const refs=[];const refBuffers=[];
|
|
for(const [color,label] of [['#175bdd','A'],['#df2828','B']]){
|
|
const buf=await sharp(Buffer.from(`<svg width="512" height="512"><rect width="512" height="512" fill="${color}"/><text x="256" y="345" text-anchor="middle" fill="white" font-family="sans-serif" font-size="260">${label}</text></svg>`)).png().toBuffer();
|
|
const key='uploads/vip98k-test-0911/'+createHash('sha256').update(buf).digest('hex')+'.png';
|
|
await client.send(new PutObjectCommand({Bucket:s.bucket,Key:key,Body:buf,ContentType:'image/png'}));refs.push(s.publicUrl.replace(/\/+$/,'')+'/'+key);refBuffers.push(buf);
|
|
}
|
|
fs.writeFileSync(path.join(out,'catalog-summary.json'),JSON.stringify(models.map(m=>({id:m.id,resolutions:m.supported_resolutions,invocations:m.invocations,invocation:m.invocation})),null,2));
|
|
let resolver; if(process.env.PROBE_VARIANT==='production'){if(process.env.PROBE_RESOLVER_MODULE)resolver=require(path.resolve(process.env.PROBE_RESOLVER_MODULE));else{require('tsx/cjs');resolver=require('../lib/ai/vip98k-task-invocation.ts');}}
|
|
const jobs=[];for(const m of models)for(const tier of tiers)for(const ratio of ratios)jobs.push({m,tier,ratio,variant:process.env.PROBE_VARIANT||'pixel'});
|
|
console.log(JSON.stringify({event:'start',models:models.map(m=>m.id),jobs:jobs.length,out}));
|
|
let cursor=0;const results=[];
|
|
async function run(job){const {m,tier,ratio,variant}=job;const started=Date.now();const record={model:m.id,tier,ratio,variant,startedAt:new Date().toISOString()};
|
|
try{
|
|
const list=m.invocations||[m.invocation];let inv=list.find(v=>v.path==='/extensions/v1/tasks')||list.find(v=>v.path==='/v1/images/edits')||m.invocation||list[0];
|
|
const headers={};const chosen=m.invocation?.path===inv.path?m.invocation:inv;for(const [k,v]of Object.entries(chosen.headers||{}))if(k.toLowerCase().startsWith('x-relay-'))headers[k]=v;
|
|
const [w,h]=ratio.split(':').map(Number);const edge=tier==='1K'?1024:tier==='2K'?2048:m.id.includes('2.5')?2880:4096;
|
|
const dims=ratio==='1:1'?[edge,edge]:[Math.round(edge*Math.sqrt(w/h)/16)*16,Math.round(edge/Math.sqrt(w/h)/16)*16];
|
|
let size=dims.join('x');let body={...(inv.request_example||{}),model:m.id,prompt:`Use both image references. Put the blue A card on the left and the red B card on the right on a neutral studio background. Preserve both letters and colors. Output aspect ratio ${ratio}.`};
|
|
if(inv.path==='/extensions/v1/tasks'){body.file_urls=refs;body.resolution=tier;body.aspect_ratio=ratio;body.size='auto';}
|
|
else {body.image=refs;body.size=variant==='bucket'?tier:size;if(variant==='bucket'||variant==='pixel-ratio')body.aspect_ratio=ratio;}
|
|
if(variant==='production'&&m.id.startsWith('gpt-image'))body.size=resolver.resolveVip98kImageTaskSize({imageSize:tier,aspectRatio:ratio,useGptImage25AspectSizing:true,preserveResolutionBucket:inv.path==='/v1/images/generations'});
|
|
if(process.env.PROBE_TEXT_ONLY==='1'){delete body.image;delete body.file_urls;body.prompt=`Two cards, a blue A on the left and a red B on the right, on a neutral studio background. Aspect ratio ${ratio}.`;}
|
|
record.path=inv.path;record.requestedSize=body.size;record.referenceCount=process.env.PROBE_TEXT_ONLY==='1'?0:2;record.route=headers['X-Relay-Route-Option']||null;
|
|
if(inv.content_type==='multipart/form-data'){
|
|
const form=new FormData();for(const [k,v]of Object.entries(body))if(!['image','image[]','images'].includes(k)&&v!==null&&typeof v!=='object')form.append(k,String(v));
|
|
refBuffers.forEach((buf,i)=>form.append('image[]',new Blob([buf],{type:'image/png'}),`ref-${i}.png`));body=form;
|
|
}
|
|
let response=await request(base+inv.path,body,{...headers,'Idempotency-Key':randomUUID()});record.requestId=response.requestId;let data=response.data;record.submitMs=Date.now()-started;
|
|
let imageUrl=urls(data)[0];let encoded=b64(data);let task=data.data?.task_id||data.task_id||data.data?.id||data.id;
|
|
const pollStart=Date.now();while(!imageUrl&&!encoded&&task&&Date.now()-pollStart<900000){
|
|
await delay(5000);const rp=inv.result_path||data.result_path||'/extensions/v1/tasks/{task_id}';const poll=rp.replaceAll('{task_id}',encodeURIComponent(task));
|
|
data=(await request(poll.startsWith('http')?poll:base+poll,null,{},'GET')).data;
|
|
const state=String(data.data?.status||data.data?.state||data.status||'').toLowerCase();if(['failed','error','cancelled'].includes(state))throw new Error(sanitize(JSON.stringify(data)));
|
|
imageUrl=urls(data)[0];encoded=b64(data);
|
|
if(!imageUrl&&!encoded&&['succeeded','completed','success'].includes(state)){
|
|
const delivery=(await request(base+'/api/external/task/result?taskId='+encodeURIComponent(task),null,{},'GET')).data;imageUrl=urls(delivery)[0];encoded=b64(delivery);
|
|
}
|
|
}
|
|
if(!imageUrl&&!encoded)throw new Error('No output image after polling');
|
|
record.generationMs=Date.now()-started;
|
|
let buf;if(encoded)buf=Buffer.from(encoded,'base64');else{const r=await fetch(imageUrl,{signal:AbortSignal.timeout(120000)});if(!r.ok)throw new Error('Output download HTTP '+r.status);buf=Buffer.from(await r.arrayBuffer());}
|
|
const meta=await sharp(buf).metadata();const {data:pix,info}=await sharp(buf).removeAlpha().resize(256,256,{fit:'fill'}).raw().toBuffer({resolveWithObject:true});let red=0,blue=0;for(let i=0;i<pix.length;i+=info.channels){const [r,g,b]=[pix[i],pix[i+1],pix[i+2]];if(r>100&&r>b*1.25&&r>g*1.15)red++;if(b>100&&b>r*1.25&&b>g*1.1)blue++;}
|
|
record.actual={width:meta.width,height:meta.height,bytes:buf.length,redShare:red/65536,blueShare:blue/65536};record.ratioMatches=Math.abs(meta.width/meta.height-w/h)/(w/h)<0.025;record.bothReferences=red/65536>.03&&blue/65536>.03;
|
|
record.status='generated';record.file=`${m.id}-${tier}-${ratio.replace(':','x')}-${variant}.${meta.format==='jpeg'?'jpg':meta.format}`;fs.writeFileSync(path.join(out,record.file),buf);
|
|
}catch(e){record.status='failed';record.error=sanitize(e.message);record.httpStatus=e.status;record.requestId=e.requestId||record.requestId;}
|
|
record.durationMs=Date.now()-started;results.push(record);fs.appendFileSync(path.join(out,'results.jsonl'),JSON.stringify(record)+'\n');console.log(JSON.stringify(record));
|
|
}
|
|
await Promise.all(Array.from({length:Number(process.env.PROBE_CONCURRENCY||3)},async()=>{while(cursor<jobs.length){const job=jobs[cursor++];await run(job);}}));
|
|
fs.writeFileSync(path.join(out,'completed.json'),JSON.stringify({finishedAt:new Date().toISOString(),results},null,2));console.log(JSON.stringify({event:'complete',count:results.length}));
|
|
})().catch(e=>{console.error(sanitize(e.message));process.exit(1)});
|