← back to Whatsmystyle
scripts/avatar-build.js
179 lines
/**
* Avatar builder — turns 3-8 user photos into a single canonical full-body
* avatar image usable as the base for virtual try-on.
*
* Strategy (tiered, picks best available):
*
* 1) GEMINI 2.5 FLASH IMAGE (Google's "Banana") — single API, edits + composites
* photos. Best for "swap outfit in old photo" workflows. Set GEMINI_API_KEY.
*
* 2) IDM-VTON (CVPR'24) / OOTDiffusion — open-source virtual try-on. Needs a
* local Python service. Set IDM_VTON_URL=http://localhost:7860 if running.
*
* 3) SAM2 + ControlNet local stack — segments the user, generates a clean
* front-facing full-body neutral-pose render. Set SAM2_URL.
*
* 4) Stub — picks the largest input photo and writes it through as the
* canonical avatar. Lets the rest of the system work end-to-end while
* the heavy lift is being wired.
*
* Body measurements (shoulder width, torso height, etc.) are extracted via
* llava-13b vision pass — same Ollama service the closet worker uses.
*/
const Database = require('better-sqlite3');
const path = require('path');
const fs = require('fs');
const fetch = require('node-fetch');
const OLLAMA = process.env.OLLAMA_URL || 'http://127.0.0.1:11434';
const VISION = process.env.OLLAMA_VISION_MODEL || 'llava:13b';
const DOPPLE_KEY = process.env.DOPPLE_API_KEY;
const GEMINI_KEY = process.env.GEMINI_API_KEY;
const IDM_URL = process.env.IDM_VTON_URL;
const SAM2_URL = process.env.SAM2_URL;
const dopple = require('./dopple');
async function measureBody(photoPath) {
try {
const b64 = fs.readFileSync(photoPath).toString('base64');
const prompt = `You see a full-body photograph of a person. Estimate proportions normalized to image height = 1.0.
Return STRICT JSON only: {"shoulder_w":0.0,"torso_h":0.0,"hip_w":0.0,"leg_h":0.0,"head_h":0.0,"pose":"front|side|3-4","clothing_visible":"top|bottom|both|partial"}`;
const r = await fetch(`${OLLAMA}/api/generate`, {
method: 'POST',
headers: { 'content-type': 'application/json' },
body: JSON.stringify({ model: VISION, prompt, images: [b64], stream: false, options: { temperature: 0.2 } }),
});
const j = await r.json();
const m = (j.response || '').match(/\{[\s\S]*\}/);
return m ? JSON.parse(m[0]) : null;
} catch (e) {
console.error('[avatar] vision err:', e.message);
return null;
}
}
async function buildViaGemini(photoPaths, outPath) {
// Gemini 2.5 Flash Image (formerly "Nano Banana") can ingest multiple images
// and composite a single canonical render. We ask for a clean full-body
// neutral-pose photo of the same person.
const images = photoPaths.slice(0, 6).map(p => ({
inline_data: { mime_type: 'image/jpeg', data: fs.readFileSync(p).toString('base64') }
}));
const body = {
contents: [{
parts: [
...images,
{ text: `Looking at these photos of the same person, produce ONE clean full-body photograph of that exact person, standing front-facing, neutral pose, plain pale-grey studio background, casual neutral outfit (white tee + grey pants). Preserve facial identity, hair, skin tone, body proportions. Photographic, not illustrated.` }
]
}],
generationConfig: { responseModalities: ['IMAGE'] }
};
const r = await fetch(`https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent?key=${GEMINI_KEY}`, {
method: 'POST',
headers: { 'content-type': 'application/json' },
body: JSON.stringify(body),
});
if (!r.ok) throw new Error(`Gemini ${r.status}: ${await r.text()}`);
const j = await r.json();
try {
const { logGemini } = require(require('os').homedir() + '/.claude/skills/cost-tracker/scripts/log-gemini.js');
logGemini(j, { app: 'whatsmystyle', note: 'avatar-build', model: 'gemini-2.5-flash-image' });
} catch {}
const part = j.candidates?.[0]?.content?.parts?.find(p => p.inline_data || p.inlineData);
const data = part?.inline_data?.data || part?.inlineData?.data;
if (!data) throw new Error('Gemini returned no image');
fs.writeFileSync(outPath, Buffer.from(data, 'base64'));
return 'gemini-2.5-flash-image';
}
function pickStubCanonical(photoPaths, outPath) {
// largest file → most likely the highest-res shot
let best = null, bestSize = 0;
for (const p of photoPaths) {
try {
const s = fs.statSync(p).size;
if (s > bestSize) { bestSize = s; best = p; }
} catch {}
}
if (!best) return null;
fs.copyFileSync(best, outPath);
return 'stub';
}
async function buildOne(db, av) {
const photos = JSON.parse(av.source_photos || '[]').filter(p => fs.existsSync(p));
if (!photos.length) {
db.prepare("UPDATE user_avatars SET status='failed' WHERE id=?").run(av.id);
return;
}
const outDir = path.join(path.dirname(photos[0]), '..', 'canonical');
fs.mkdirSync(outDir, { recursive: true });
const outPath = path.join(outDir, `avatar-${av.id}.png`);
let provider = null;
let dopple_twin_id = null;
try {
if (DOPPLE_KEY) {
// Dopple-first: create a persistent twin, then ask Dopple for a canonical render
dopple_twin_id = await dopple.createTwin(photos);
await dopple.pollTwin(dopple_twin_id);
// Render the twin in neutral pose by doing a try-on with a "neutral basics" prompt
const { buffer } = await dopple.tryOn({
twinId: dopple_twin_id,
garmentUrl: null,
options: { pose: 'neutral_front', background: 'studio_grey', outfit: 'plain_white_tee_grey_pants' },
}).catch(async () => {
// If the neutral-render path isn't supported, just save the first photo as canonical;
// the twin_id is what powers all subsequent try-ons.
return { buffer: fs.readFileSync(photos[0]) };
});
fs.writeFileSync(outPath, buffer);
provider = 'dopple';
} else if (GEMINI_KEY) provider = await buildViaGemini(photos, outPath);
else if (IDM_URL) {
// POST to local IDM-VTON service
const fd = new (require('form-data'))();
photos.slice(0, 4).forEach((p, i) => fd.append(`photo${i}`, fs.createReadStream(p)));
const r = await fetch(`${IDM_URL}/build`, { method: 'POST', body: fd });
if (!r.ok) throw new Error(`IDM-VTON ${r.status}`);
const buf = Buffer.from(await r.arrayBuffer());
fs.writeFileSync(outPath, buf);
provider = 'idm-vton';
} else if (SAM2_URL) {
const fd = new (require('form-data'))();
photos.slice(0, 4).forEach((p, i) => fd.append(`photo${i}`, fs.createReadStream(p)));
const r = await fetch(`${SAM2_URL}/avatar`, { method: 'POST', body: fd });
if (!r.ok) throw new Error(`SAM2 ${r.status}`);
const buf = Buffer.from(await r.arrayBuffer());
fs.writeFileSync(outPath, buf);
provider = 'sam2-controlnet';
} else {
provider = pickStubCanonical(photos, outPath);
}
} catch (e) {
console.error('[avatar] build err:', e.message);
pickStubCanonical(photos, outPath);
provider = 'stub (after error)';
}
const measurements = await measureBody(outPath);
const meta = { ...(measurements || {}), _provider: provider };
if (dopple_twin_id) meta._dopple_twin_id = dopple_twin_id;
db.prepare(`UPDATE user_avatars SET status='ready', canonical_path=?, body_measurements=?, ready_at=datetime('now') WHERE id=?`)
.run(outPath, JSON.stringify(meta), av.id);
console.log(`[avatar] built av#${av.id} via ${provider} → ${outPath}${dopple_twin_id ? ' (twin '+dopple_twin_id+')' : ''}`);
}
async function main() {
const dbPath = path.join(__dirname, '..', 'data', 'whatsmystyle.db');
const db = new Database(dbPath);
const rows = db.prepare("SELECT * FROM user_avatars WHERE status IN ('pending','building') LIMIT 10").all();
console.log(`[avatar] ${rows.length} pending`);
for (const r of rows) {
try { await buildOne(db, r); }
catch (e) { console.error('[avatar] fatal', r.id, e.message); db.prepare("UPDATE user_avatars SET status='failed' WHERE id=?").run(r.id); }
}
}
if (require.main === module) main().catch(e => { console.error(e); process.exit(1); });
module.exports = { buildOne };