From 011168c3b881798007e3d72407f9cba44de9f0a3 Mon Sep 17 00:00:00 2001 From: iamdoubz <> Date: Mon, 29 Jun 2026 13:07:43 -0500 Subject: [PATCH] Add progress bar for ai gen image (est) --- src/app/admin/AdminCreations.tsx | 7 ++++++- src/app/api/create/[id]/route.ts | 14 ++++++++++++-- src/app/create/CreateForm.tsx | 24 +++++++++++++++++++++-- src/lib/comfyui.ts | 33 +++++++++++++++++++++++++++++--- src/lib/createdLessons.ts | 15 ++++++++++++++- 5 files changed, 84 insertions(+), 9 deletions(-) diff --git a/src/app/admin/AdminCreations.tsx b/src/app/admin/AdminCreations.tsx index ffc1d9f..6d99421 100644 --- a/src/app/admin/AdminCreations.tsx +++ b/src/app/admin/AdminCreations.tsx @@ -18,8 +18,12 @@ export default function AdminCreations({ review, ready }: { review: Row[]; ready const [busy, setBusy] = useState(null); const [lightbox, setLightbox] = useState(null); const [editing, setEditing] = useState(null); + const [prog, setProg] = useState>({}); const editVal = useRef(""); + const etaText = (secs: number | null) => + secs === null ? "estimating…" : secs <= 0 ? "almost done…" : secs < 60 ? `about ${secs}s left` : `about ${Math.ceil(secs / 60)} min left`; + // Auto-update: while anything is generating, poll its status and refresh when it changes // (so a Regenerate's new image appears on screen without a manual reload). useEffect(() => { @@ -30,6 +34,7 @@ export default function AdminCreations({ review, ready }: { review: Row[]; ready try { const res = await fetch(`/api/create/${r.id}`); const d = await res.json(); + if (d.progress) setProg((prev) => ({ ...prev, [r.id]: { percent: d.progress.percent, etaSeconds: d.progress.etaSeconds } })); if (d.status && d.status !== r.status) { router.refresh(); break; @@ -97,7 +102,7 @@ export default function AdminCreations({ review, ready }: { review: Row[]; ready )}
- {working ? "🎨 Drawing… (updates automatically)" : r.status === "review" ? "🔎 Awaiting approval" : r.status === "needs_review" ? "⏳ Needs generation" : r.status === "failed" ? `⚠️ Failed: ${r.error || "unknown error"}` : r.status === "ready" ? "✓ Ready" : r.status} + {working ? `🎨 Drawing… ${prog[r.id] ? `${prog[r.id].percent}% · ${etaText(prog[r.id].etaSeconds)}` : "(updates automatically)"}` : r.status === "review" ? "🔎 Awaiting approval" : r.status === "needs_review" ? "⏳ Needs generation" : r.status === "failed" ? `⚠️ Failed: ${r.error || "unknown error"}` : r.status === "ready" ? "✓ Ready" : r.status}
diff --git a/src/app/api/create/[id]/route.ts b/src/app/api/create/[id]/route.ts index 88b588b..82ded6f 100644 --- a/src/app/api/create/[id]/route.ts +++ b/src/app/api/create/[id]/route.ts @@ -1,6 +1,6 @@ import { NextResponse } from "next/server"; import { getCurrentUser } from "@/lib/session"; -import { getCreatedById } from "@/lib/createdLessons"; +import { getCreatedById, getCreatedProgress } from "@/lib/createdLessons"; // Poll the status of a creation (owner or admin). export async function GET(_req: Request, ctx: { params: Promise<{ id: string }> }) { @@ -13,5 +13,15 @@ export async function GET(_req: Request, ctx: { params: Promise<{ id: string }> if (row.user_id !== user.id && user.role !== "admin") return NextResponse.json({ error: "Not authorized." }, { status: 403 }); - return NextResponse.json({ id: row.id, slug: row.slug, status: row.status, error: row.error, subject: row.subject }); + // Best-effort generation progress + ETA (from ComfyUI's WebSocket step events). + let progress: { value: number; max: number; percent: number; etaSeconds: number | null } | null = null; + const p = getCreatedProgress(row.id); + if (p && p.max > 0) { + const elapsed = Date.now() - p.startedAt; + const perStep = p.value > 0 ? elapsed / p.value : 0; + const etaSeconds = perStep > 0 ? Math.round((Math.max(0, p.max - p.value) * perStep) / 1000) : null; + progress = { value: p.value, max: p.max, percent: Math.round((p.value / p.max) * 100), etaSeconds }; + } + + return NextResponse.json({ id: row.id, slug: row.slug, status: row.status, error: row.error, subject: row.subject, progress }); } diff --git a/src/app/create/CreateForm.tsx b/src/app/create/CreateForm.tsx index 8968a3b..04488ce 100644 --- a/src/app/create/CreateForm.tsx +++ b/src/app/create/CreateForm.tsx @@ -10,6 +10,7 @@ export default function CreateForm({ level: initialLevel, levels, enabled }: { l const [phase, setPhase] = useState("idle"); const [msg, setMsg] = useState(""); const [slug, setSlug] = useState(""); + const [progress, setProgress] = useState<{ percent: number; etaSeconds: number | null } | null>(null); const poll = useRef | undefined>(undefined); useEffect(() => () => clearInterval(poll.current), []); @@ -50,12 +51,19 @@ export default function CreateForm({ level: initialLevel, levels, enabled }: { l try { const r = await fetch(`/api/create/${id}`); const d = await r.json(); + if (d.progress) setProgress({ percent: d.progress.percent, etaSeconds: d.progress.etaSeconds }); if (d.status === "ready") { clearInterval(poll.current); + setProgress(null); setSlug(baseSlug); setPhase("ready"); + } else if (d.status === "review") { + clearInterval(poll.current); + setProgress(null); + setPhase("review"); } else if (d.status === "failed" || d.status === "blocked") { clearInterval(poll.current); + setProgress(null); setPhase("failed"); setMsg(d.error || "The drawing couldn't be made. Try a different subject."); } @@ -65,6 +73,13 @@ export default function CreateForm({ level: initialLevel, levels, enabled }: { l }, 2500); } + function etaText(secs: number | null): string { + if (secs === null) return "estimating…"; + if (secs <= 0) return "almost done…"; + if (secs < 60) return `about ${secs}s left`; + return `about ${Math.ceil(secs / 60)} min left`; + } + if (!enabled) { return (
@@ -107,8 +122,13 @@ export default function CreateForm({ level: initialLevel, levels, enabled }: { l {phase === "generating" && (
-

🎨 Drawing your {subject || "picture"}… this can take a little while.

-

You can stay on this page.

+

🎨 Drawing your {subject || "picture"}…

+
+
+
+

+ {progress ? `${progress.percent}% · ${etaText(progress.etaSeconds)}` : "Starting up…"} +

)} {phase === "review" && ( diff --git a/src/lib/comfyui.ts b/src/lib/comfyui.ts index 2742c3a..dfae882 100644 --- a/src/lib/comfyui.ts +++ b/src/lib/comfyui.ts @@ -54,9 +54,12 @@ async function sleep(ms: number) { /** * Generate coloring-book line art for the given subject. Returns the PNG bytes. * Injects the subject into the workflow's prompt placeholder + a fresh seed. + * `onProgress` (best-effort, via ComfyUI's WebSocket) reports sampler step value/max for an ETA. * Throws on timeout / ComfyUI errors (caller marks the lesson 'failed'). */ -export async function generateLineArt(subject: string): Promise { +export type ProgressFn = (p: { value: number; max: number }) => void; + +export async function generateLineArt(subject: string, onProgress?: ProgressFn): Promise { if (!isCreateEnabled()) throw new Error("Create is disabled (COMFYUI_URL not set)."); const workflow = loadWorkflow(); @@ -70,16 +73,39 @@ export async function generateLineArt(subject: string): Promise { const clientId = `drawit-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + // Open a best-effort WebSocket for live step progress (ComfyUI has no "time left" endpoint). + let ws: WebSocket | undefined; + if (onProgress && typeof WebSocket !== "undefined") { + try { + const wsUrl = `${URL_BASE.replace(/^http/, "ws")}/ws?clientId=${clientId}`; + ws = new WebSocket(wsUrl); + ws.onmessage = (ev: MessageEvent) => { + if (typeof ev.data !== "string") return; // ignore binary preview frames + try { + const msg = JSON.parse(ev.data); + if (msg?.type === "progress" && msg.data && typeof msg.data.max === "number") { + onProgress({ value: Number(msg.data.value) || 0, max: Number(msg.data.max) || 0 }); + } + } catch { + /* ignore */ + } + }; + ws.onerror = () => { /* best-effort; ignore */ }; + } catch { + ws = undefined; + } + } + // 1) Queue the workflow. const queueRes = await fetch(`${URL_BASE}/prompt`, { method: "POST", headers: { "Content-Type": "application/json", ...authHeaders() }, body: JSON.stringify({ prompt: workflow, client_id: clientId }), }); - if (!queueRes.ok) throw new Error(`ComfyUI /prompt failed: ${queueRes.status} ${await queueRes.text()}`); + if (!queueRes.ok) { try { ws?.close(); } catch { /* */ } throw new Error(`ComfyUI /prompt failed: ${queueRes.status} ${await queueRes.text()}`); } const queued = (await queueRes.json()) as { prompt_id?: string; error?: unknown }; const promptId = queued.prompt_id; - if (!promptId) throw new Error(`ComfyUI did not return a prompt_id: ${JSON.stringify(queued)}`); + if (!promptId) { try { ws?.close(); } catch { /* */ } throw new Error(`ComfyUI did not return a prompt_id: ${JSON.stringify(queued)}`); } // 2) Poll history until the outputs appear (or timeout). const deadline = Date.now() + TIMEOUT_MS; @@ -97,6 +123,7 @@ export async function generateLineArt(subject: string): Promise { } if (image) break; } + try { ws?.close(); } catch { /* */ } if (!image) throw new Error("ComfyUI timed out before producing an image."); // 3) Fetch the image bytes. diff --git a/src/lib/createdLessons.ts b/src/lib/createdLessons.ts index 6b34112..d65000a 100644 --- a/src/lib/createdLessons.ts +++ b/src/lib/createdLessons.ts @@ -122,6 +122,13 @@ export function promoteCreated(id: number) { getDb().prepare("UPDATE created_lessons SET promoted = 1, updated_at = datetime('now') WHERE id = ?").run(id); } +// Live generation progress (in-memory; single self-hosted node). Cleared when a run finishes. +export interface CreatedProgress { value: number; max: number; startedAt: number; updatedAt: number } +const progressStore = new Map(); +export function getCreatedProgress(id: number): CreatedProgress | undefined { + return progressStore.get(id); +} + /** * Run generation for a row: ComfyUI -> vectorize -> ready (or failed). Fire-and-forget from the API * route; safe on a long-lived self-hosted node. Heavy deps are imported lazily so the data layer stays @@ -132,9 +139,13 @@ export async function processCreation(id: number): Promise { if (!row) return; try { updateCreatedStatus(id, "generating"); + progressStore.set(id, { value: 0, max: 0, startedAt: Date.now(), updatedAt: Date.now() }); const { generateLineArt } = await import("./comfyui"); const { pngToLessonArt } = await import("./vectorize"); - const png = await generateLineArt(row.subject); + const png = await generateLineArt(row.subject, (p) => { + const prev = progressStore.get(id); + progressStore.set(id, { value: p.value, max: p.max, startedAt: prev?.startedAt ?? Date.now(), updatedAt: Date.now() }); + }); const art = await pngToLessonArt(png); // { outline, details, full } // Auto-approved (allowlist) or admin-approved → live; anything else is quarantined for review. const fresh = getCreatedById(id); @@ -142,6 +153,8 @@ export async function processCreation(id: number): Promise { setGenerated(id, JSON.stringify(art), `data:image/png;base64,${png.toString("base64")}`, needsReview ? "review" : "ready"); } catch (e) { updateCreatedStatus(id, "failed", e instanceof Error ? e.message : String(e)); + } finally { + progressStore.delete(id); } } -- 2.34.1