← back to Whatsmystyle

scripts/tryon-worker.js

218 lines

/**
 * Try-on worker — drains the tryon_jobs queue.
 *
 * mode='avatar'      → render the catalog garment onto the user's canonical avatar.
 * mode='photo_swap'  → replace the outfit in the user's uploaded photo with the garment.
 *
 * Provider order:
 *   1) Gemini 2.5 Flash Image (if GEMINI_API_KEY)         — best image-edit quality
 *   2) IDM-VTON (if IDM_VTON_URL set)                     — open-source SOTA virtual try-on
 *   3) OOTDiffusion (if OOTD_URL set)                     — open-source alt
 *   4) Stub — composites garment image over avatar via sharp; ugly but proves the pipe.
 */
require('dotenv').config({ path: require('path').join(__dirname, '..', '.env') });
const Database = require('better-sqlite3');
const path = require('path');
const fs = require('fs');
const fetch = require('node-fetch');

const DOPPLE_KEY = process.env.DOPPLE_API_KEY;
const FASHN_KEY  = process.env.FASHN_API_KEY;
const OPENAI_KEY = process.env.OPENAI_API_KEY;
const GEMINI_KEY = process.env.GEMINI_API_KEY;
const IDM_URL    = process.env.IDM_VTON_URL;
const OOTD_URL   = process.env.OOTD_URL;
const dopple     = require('./dopple');
const fashn      = require('./fashn');
const openaiImg  = require('./openai-image');

async function fetchImageAsBuffer(urlOrPath) {
  if (urlOrPath.startsWith('http')) {
    const r = await fetch(urlOrPath);
    if (!r.ok) throw new Error(`fetch ${urlOrPath} ${r.status}`);
    return Buffer.from(await r.arrayBuffer());
  }
  return fs.readFileSync(urlOrPath);
}

async function gemini(personImgBuf, garmentImgBuf, instruction) {
  const body = {
    contents: [{
      parts: [
        { inline_data: { mime_type: 'image/jpeg', data: personImgBuf.toString('base64') } },
        { inline_data: { mime_type: 'image/jpeg', data: garmentImgBuf.toString('base64') } },
        { text: instruction },
      ],
    }],
    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: 'tryon-worker', 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) {
    // Log the actual response so we can see WHY (safety refusal, quota, etc.)
    const text = j.candidates?.[0]?.content?.parts?.find(p => p.text)?.text;
    const finish = j.candidates?.[0]?.finishReason;
    const safety = j.candidates?.[0]?.safetyRatings;
    const err = new Error(`Gemini returned no image (finishReason=${finish}, text="${(text || '').slice(0, 200)}", safety=${JSON.stringify(safety || [])})`);
    err.gemini_response = j;
    throw err;
  }
  return Buffer.from(data, 'base64');
}

async function processJob(db, job) {
  // resolve garment — either catalog item OR user's own closet item
  let garmentImageUrl = null;
  let garmentLocalPath = null;
  let garmentTitle = '', garmentCategory = '', garmentColor = '';
  if (job.item_id) {
    const item = db.prepare('SELECT id, image_url, title, category, color FROM items WHERE id=?').get(job.item_id);
    if (!item) throw new Error('catalog item missing');
    garmentImageUrl = item.image_url;
    garmentTitle = item.title; garmentCategory = item.category; garmentColor = item.color;
  } else if (job.closet_id) {
    const c = db.prepare('SELECT id, photo_path, category, color, pattern, vendor_guess FROM closet WHERE id=?').get(job.closet_id);
    if (!c) throw new Error('closet item missing');
    garmentLocalPath = c.photo_path;
    garmentTitle = c.vendor_guess || c.pattern || c.color || 'closet item';
    garmentCategory = c.category || 'top'; garmentColor = c.color || '';
  }

  // resolve "person" base — avatar canonical OR an old photo
  let personPath = null;
  let dopple_twin_id = null;
  if (job.mode === 'photo_swap' || job.mode === 'closet_to_oldphoto') {
    personPath = job.source_photo_path;
  } else {
    const av = db.prepare('SELECT canonical_path, body_measurements FROM user_avatars WHERE id=?').get(job.avatar_id);
    personPath = av?.canonical_path;
    try {
      const m = JSON.parse(av?.body_measurements || '{}');
      dopple_twin_id = m._dopple_twin_id || null;
    } catch {}
  }
  if (!personPath || !fs.existsSync(personPath)) throw new Error('person image missing');

  const personBuf  = fs.readFileSync(personPath);
  const garmentBuf = garmentLocalPath ? fs.readFileSync(garmentLocalPath) : await fetchImageAsBuffer(garmentImageUrl);

  const outDir = path.join(__dirname, '..', 'data', 'uploads', 'tryon');
  fs.mkdirSync(outDir, { recursive: true });
  const outPath = path.join(outDir, `job-${job.id}.png`);

  const occasion = job.occasion ? ` Scene context: ${job.occasion}.` : '';
  const instruction = (job.mode === 'photo_swap' || job.mode === 'closet_to_oldphoto')
    ? `Edit the FIRST image: replace whatever clothing the wearer (the subject) is currently wearing with the garment shown in the SECOND image. Preserve the wearer's face, hair, body shape, pose, lighting, the background, and any other people exactly as they are. The new garment is: ${garmentTitle} (${garmentCategory}, ${garmentColor}).${occasion} Photographic realism.`
    : `Render the person in the FIRST image wearing the garment shown in the SECOND image. The new garment is: ${garmentTitle} (${garmentCategory}, ${garmentColor}). Keep face, hair, skin, pose. Clean studio background. Photographic.`;

  let provider = job.provider || 'stub';
  try {
    if (DOPPLE_KEY) {
      // Dopple path — use twin if we have one, otherwise create on the fly from the person image.
      let twinId = dopple_twin_id;
      if (!twinId) {
        twinId = await dopple.createTwin([personPath]);
        await dopple.pollTwin(twinId);
      }
      let result;
      if (job.mode === 'photo_swap' || job.mode === 'closet_to_oldphoto') {
        result = await dopple.editScene({
          twinId,
          sourceImagePath: personPath,
          garmentUrl: garmentImageUrl,
          garmentBuffer: garmentLocalPath ? garmentBuf : null,
          occasion: job.occasion,
        });
      } else {
        result = await dopple.tryOn({
          twinId,
          garmentUrl: garmentImageUrl,
          garmentBuffer: garmentLocalPath ? garmentBuf : null,
          options: { category: garmentCategory },
        });
      }
      fs.writeFileSync(outPath, result.buffer);
      provider = 'dopple';
    } else if (FASHN_KEY) {
      const result = await fashn.tryOn({ personBuf, garmentBuf, category: garmentCategory });
      fs.writeFileSync(outPath, result.buffer);
      provider = 'fashn-ai';
      try {
        db.prepare("INSERT INTO provider_costs (provider, kind, cost_cents, user_id, ref_job_id) VALUES (?, ?, ?, ?, ?)")
          .run(provider, job.mode || 'tryon', result.cost_cents, job.user_id, job.id);
      } catch (e) { console.error('[tryon] cost record err:', e.message); }
    } else if (OPENAI_KEY) {
      const result = await openaiImg.tryOn({ personBuf, garmentBuf, instruction });
      fs.writeFileSync(outPath, result.buffer);
      provider = 'gpt-image-1';
      try {
        db.prepare("INSERT INTO provider_costs (provider, kind, cost_cents, user_id, ref_job_id) VALUES (?, ?, ?, ?, ?)")
          .run(provider, job.mode || 'tryon', result.cost_cents, job.user_id, job.id);
      } catch (e) { console.error('[tryon] cost record err:', e.message); }
    } else if (GEMINI_KEY) {
      const out = await gemini(personBuf, garmentBuf, instruction);
      fs.writeFileSync(outPath, out);
      provider = 'gemini-2.5-flash-image';
      try {
        db.prepare("INSERT INTO provider_costs (provider, kind, cost_cents, user_id, ref_job_id) VALUES (?, ?, ?, ?, ?)")
          .run(provider, job.mode || 'tryon', 3.9, job.user_id, job.id);
      } catch (e) { console.error('[tryon] cost record err:', e.message); }
    } else if (IDM_URL) {
      const FormData = require('form-data');
      const fd = new FormData();
      fd.append('person', personBuf, { filename: 'person.jpg' });
      fd.append('garment', garmentBuf, { filename: 'garment.jpg' });
      fd.append('category', item.category || 'top');
      const r = await fetch(`${IDM_URL}/tryon`, { method: 'POST', body: fd });
      if (!r.ok) throw new Error(`IDM-VTON ${r.status}`);
      fs.writeFileSync(outPath, Buffer.from(await r.arrayBuffer()));
      provider = 'idm-vton';
    } else if (OOTD_URL) {
      const FormData = require('form-data');
      const fd = new FormData();
      fd.append('person', personBuf, { filename: 'person.jpg' });
      fd.append('garment', garmentBuf, { filename: 'garment.jpg' });
      const r = await fetch(`${OOTD_URL}/run`, { method: 'POST', body: fd });
      if (!r.ok) throw new Error(`OOTD ${r.status}`);
      fs.writeFileSync(outPath, Buffer.from(await r.arrayBuffer()));
      provider = 'ootdiffusion';
    } else {
      // stub: just copy the person image — proves the pipe, ugly result
      fs.writeFileSync(outPath, personBuf);
      provider = 'stub';
    }
    db.prepare(`UPDATE tryon_jobs SET status='done', result_path=?, provider=?, done_at=datetime('now') WHERE id=?`)
      .run(outPath, provider, job.id);
    console.log(`[tryon] ✓ job#${job.id} via ${provider}`);
  } catch (e) {
    db.prepare(`UPDATE tryon_jobs SET status='failed', error=?, done_at=datetime('now') WHERE id=?`).run(e.message, job.id);
    console.error(`[tryon] ✗ job#${job.id}: ${e.message}`);
  }
}

async function main() {
  const dbPath = path.join(__dirname, '..', 'data', 'whatsmystyle.db');
  const db = new Database(dbPath);
  const rows = db.prepare("SELECT * FROM tryon_jobs WHERE status='queued' ORDER BY id ASC LIMIT 20").all();
  console.log(`[tryon] ${rows.length} queued`);
  for (const j of rows) {
    db.prepare("UPDATE tryon_jobs SET status='running' WHERE id=?").run(j.id);
    try { await processJob(db, j); }
    catch (e) { db.prepare(`UPDATE tryon_jobs SET status='failed', error=? WHERE id=?`).run(e.message, j.id); }
  }
}

if (require.main === module) main().catch(e => { console.error(e); process.exit(1); });
module.exports = { processJob };