From 6fed367c0aa3f2899785401db1f5e6f3a6bb5481 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sat, 17 Oct 2026 18:54:26 +0000
Subject: [PATCH] feat: show rated results as pictures to pick from in the
 training Data step

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P
---
 app/ui/features/training/data.js          | 33 ++++++++++++++++-------
 app/ui/features/training/index.js         |  2 +-
 app/ui/features/training/training.test.js | 23 ++++++++++++++++
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/app/ui/features/training/data.js b/app/ui/features/training/data.js
index 00cc2166..e5c15b99 100644
--- a/app/ui/features/training/data.js
+++ b/app/ui/features/training/data.js
@@ -1,7 +1,7 @@
 // Step 1, Training data: pick or make a set, then put media in it from any mix of sources.
 // Skippable: a set that already has what you want goes straight on to tags or to training.
 import { composer as c } from "../../composer/composer.js";
-import { LABEL_WORDS } from "./state.js";
+import { LABEL_WORDS, stillOf } from "./state.js";
 
 export const RATED = [
   { value: "liked", label: "Liked results" },
@@ -73,20 +73,35 @@ export function dataStep({ app, session, go, ui }) {
   } });
 
   const want = () => (ui.ratedKind === "all" ? ALL_RATED : [ui.ratedKind]);
-  const ratedCount = () => want().reduce((a, k) => a + ((ui.rated.counts || {})[k] || 0), 0);
+  const shownRated = () => (ui.rated.results || []).filter((r) => want().includes(r.label));
+  const keyOf = (r) => `${r.media.type}:${r.media.subfolder}/${r.media.filename}`;
+  const ratedPicked = ui.ratedPicked;
   const ratedNote = c.text.sm({ text: "" });
-  const addRated = c.button.sm({ label: "Add these results", tone: "primary", disabled: true, onClick: async () => {
-    await added((ui.rated.results || []).filter((r) => want().includes(r.label)).map(ratedEntry), "your results");
+  const ratedGallery = c.gallery.adaptive({ id: "training-rated", items: [], selection: [...ratedPicked], empty: "None found yet. Rate results in your projects and they appear here.", onActivate: (item) => {
+    if (ratedPicked.has(item.id)) ratedPicked.delete(item.id); else ratedPicked.add(item.id);
+    ratedGallery.setValue([...ratedPicked]);
+    showRated();
   } });
-  const showRated = () => {
-    const n = ratedCount();
-    ratedNote.setText(n ? `${n} rated result${n > 1 ? "s" : ""} found in your projects.` : "None found yet. Rate results in your projects and they appear here.");
+  const addRated = c.button.sm({ label: "Add these results", tone: "primary", disabled: true, onClick: async () => {
+    const all = shownRated(), chosen = all.filter((r) => ratedPicked.has(keyOf(r)));
+    await added((chosen.length ? chosen : all).map(ratedEntry), "your results");
+    ratedPicked.clear();
+    showRated();
+  } });
+  function showRated() {
+    const all = shownRated();
+    for (const k of [...ratedPicked]) if (!all.some((r) => keyOf(r) === k)) ratedPicked.delete(k);
+    ratedGallery.setItems(all.map((r) => ({ id: keyOf(r), label: r.media.filename, thumb: stillOf({ src: { kind: "output", ...r.media }, media_kind: r.media.filename.match(/\.(mp4|webm|mov|mkv|avi)$/i) ? "video" : "image" }), badge: r.label })));
+    ratedGallery.setValue([...ratedPicked]);
+    const n = all.length, p = ratedPicked.size;
+    ratedNote.setText(n ? `${n} rated result${n > 1 ? "s" : ""} found in your projects. ${p ? `${p} picked: only those are added.` : "None picked: all are added."}` : "None found yet. Rate results in your projects and they appear here.");
     addRated.node.disabled = !n || !open;
-  };
+  }
   const include = c.field.default({ label: "Captions include", hint: "What goes into the caption from each result's prompt. Shortcuts are expanded. Variables are filled in only when ticked, otherwise $name stays as written. Nothing ticked adds them with no caption.",
     control: c.checklist.default({ label: "Captions include", items: INCLUDE, values: ui.include, onChange: async (v) => {
       ui.include = v;
       try { ui.rated = await app.api.trainRated(v.length ? v : ["none"]); } catch { /* the old captions stay */ }
+      showRated();
     } }) });
   const ratedSel = c.segmented.sm({ label: "Which results", options: RATED, value: ui.ratedKind, onChange: (v) => { ui.ratedKind = v; showRated(); } });
   showRated();
@@ -100,7 +115,7 @@ export function dataStep({ app, session, go, ui }) {
     status,
     c.group.default({ label: "Drop files", rows: [dropzone] }),
     c.group.default({ label: "Choose in your Bin", rows: [gallery, c.toolbar.default({ items: [pickedNote], trailing: [addBin] })] }),
-    c.group.default({ label: "Choose your own results", hint: "Results you rated in your projects. Liked ones teach what you want; disliked ones can teach what to avoid.", rows: [ratedSel, include, c.toolbar.default({ items: [ratedNote], trailing: [addRated] })] }),
+    c.group.default({ label: "Choose your own results", hint: "Results you rated in your projects. Liked ones teach what you want; disliked ones can teach what to avoid.", rows: [ratedSel, ratedGallery, include, c.toolbar.default({ items: [ratedNote], trailing: [addRated] })] }),
     next,
   ] });
 }
diff --git a/app/ui/features/training/index.js b/app/ui/features/training/index.js
index fdbb7c50..f353607e 100644
--- a/app/ui/features/training/index.js
+++ b/app/ui/features/training/index.js
@@ -26,7 +26,7 @@ export default {
     const api = app.api;
     const session = makeSession(api);
     const ui = {
-      tab: "data", picked: new Set(), bin: [], rated: { results: [], counts: {} }, ratedKind: "liked", include: ["main"], selected: new Set(), filter: "all", tag: "", limit: 48,
+      tab: "data", picked: new Set(), ratedPicked: new Set(), bin: [], rated: { results: [], counts: {} }, ratedKind: "liked", include: ["main"], selected: new Set(), filter: "all", tag: "", limit: 48,
       tagger: { id: "wd_vit_v3", list: [], job: null, general: 0.35, limit: 0, characters: false, mode: "add", scope: "untagged" },
       advanced: false, runs: [], chosenRun: "", record: null,
       run: { shown: false, active: false, state: "idle", loss: [], smooth: [], step: 0, total: 0, elapsed: 0, message: "", error: "", lora: "", before: "" },
diff --git a/app/ui/features/training/training.test.js b/app/ui/features/training/training.test.js
index 5e73aca0..ec53460f 100644
--- a/app/ui/features/training/training.test.js
+++ b/app/ui/features/training/training.test.js
@@ -250,3 +250,26 @@ test("Tags step: Remove selected takes the ticked items out of the set and saves
   assert.equal(saved[2].items[0].id, "cccccccccccc");
   off();
 });
+
+test("rated results are shown as pictures, and only the picked ones are added", async () => {
+  const f = fake();
+  f.api.trainRated = async () => ({ results: [
+    { label: "liked", media: { type: "output", subfolder: "", filename: "x.mp4" }, caption: "a" },
+    { label: "liked", media: { type: "output", subfolder: "", filename: "y.png" }, caption: "b" }], counts: { liked: 2 } });
+  const added = [];
+  f.api.trainAdd = async (name, entries) => { added.push(entries); return { set: f.doc, added: entries.length }; };
+  const { button, off } = await open(f);
+  button.click();
+  await wait(60);
+  assert.ok(text().includes("2 rated results found"));
+  const cells = [...document.querySelectorAll(".cx-cell")];
+  assert.equal(cells.length, 3, "one bin item and the two rated results are drawn as pictures");
+  (cells[2].querySelector("button") || cells[2]).click();
+  await wait(20);
+  assert.ok(text().includes("1 picked: only those are added"));
+  click("Add these results").click();
+  await wait(40);
+  assert.equal(added[0].length, 1);
+  assert.equal(added[0][0].src.filename, "y.png");
+  off();
+});