FillipandClaude Opus 4.6
8e16057d2d
Pixel-space diffusion text-to-image generation with REPA LoRA training. 6 nodes: Model Loader, Text Encoder Loader, Generate, Training Config, Train, LoRA Loader. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
949 lines
25 KiB
JavaScript
949 lines
25 KiB
JavaScript
import { app } from "../../../scripts/app.js";
|
|
import { api } from "../../../scripts/api.js";
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Styles — identical to AceStep Training widget visual
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const TRAINING_WIDGET_STYLES = `
|
|
.pixelgen-training-widget {
|
|
--primary: #06b6d4;
|
|
--primary-glow: rgba(6, 182, 212, 0.4);
|
|
--secondary: #8b5cf6;
|
|
--success: #22c55e;
|
|
--danger: #ef4444;
|
|
--warning: #f59e0b;
|
|
--bg-dark: #0f0f12;
|
|
--bg-card: #18181b;
|
|
--bg-elevated: #1f1f23;
|
|
--border: #27272a;
|
|
--border-hover: #3f3f46;
|
|
--text-primary: #fafafa;
|
|
--text-secondary: #a1a1aa;
|
|
--text-muted: #71717a;
|
|
|
|
background: var(--bg-card);
|
|
border-radius: 12px;
|
|
border: 1px solid var(--border);
|
|
overflow: hidden;
|
|
position: relative;
|
|
font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
|
|
color: var(--text-primary);
|
|
box-sizing: border-box;
|
|
height: 100%;
|
|
min-height: 300px;
|
|
display: flex;
|
|
flex-direction: column;
|
|
}
|
|
|
|
.pixelgen-training-widget * {
|
|
box-sizing: border-box;
|
|
}
|
|
|
|
/* Header */
|
|
.pixelgen-header {
|
|
display: flex;
|
|
justify-content: space-between;
|
|
align-items: center;
|
|
padding: 6px 10px;
|
|
background: var(--bg-elevated);
|
|
border-bottom: 1px solid var(--border);
|
|
flex-shrink: 0;
|
|
}
|
|
|
|
.pixelgen-title {
|
|
font-size: 11px;
|
|
font-weight: 600;
|
|
color: var(--text-primary);
|
|
display: flex;
|
|
align-items: center;
|
|
gap: 8px;
|
|
}
|
|
|
|
.pixelgen-badge {
|
|
padding: 2px 8px;
|
|
background: var(--primary);
|
|
color: white;
|
|
border-radius: 10px;
|
|
font-size: 10px;
|
|
font-weight: 500;
|
|
}
|
|
|
|
.pixelgen-badge.idle {
|
|
background: var(--text-muted);
|
|
}
|
|
|
|
.pixelgen-badge.training {
|
|
background: var(--success);
|
|
animation: pixelgen-pulse 2s infinite;
|
|
}
|
|
|
|
@keyframes pixelgen-pulse {
|
|
0%, 100% { opacity: 1; }
|
|
50% { opacity: 0.7; }
|
|
}
|
|
|
|
/* Content */
|
|
.pixelgen-content {
|
|
flex: 1;
|
|
display: flex;
|
|
flex-direction: column;
|
|
padding: 10px;
|
|
gap: 8px;
|
|
overflow: hidden;
|
|
}
|
|
|
|
/* Stats — compact inline row */
|
|
.pixelgen-stats {
|
|
display: flex;
|
|
gap: 12px;
|
|
justify-content: center;
|
|
align-items: baseline;
|
|
flex-shrink: 0;
|
|
}
|
|
|
|
.pixelgen-stat {
|
|
display: flex;
|
|
align-items: baseline;
|
|
gap: 4px;
|
|
}
|
|
|
|
.pixelgen-stat-label {
|
|
font-size: 9px;
|
|
color: var(--text-muted);
|
|
text-transform: uppercase;
|
|
}
|
|
|
|
.pixelgen-stat-value {
|
|
font-size: 11px;
|
|
font-weight: 600;
|
|
color: var(--primary);
|
|
font-variant-numeric: tabular-nums;
|
|
}
|
|
|
|
/* Progress Bar */
|
|
.pixelgen-progress-section {
|
|
background: var(--bg-elevated);
|
|
border-radius: 6px;
|
|
padding: 8px 10px;
|
|
flex-shrink: 0;
|
|
}
|
|
|
|
.pixelgen-progress-header {
|
|
display: flex;
|
|
justify-content: space-between;
|
|
margin-bottom: 4px;
|
|
}
|
|
|
|
.pixelgen-progress-label {
|
|
font-size: 9px;
|
|
color: var(--text-secondary);
|
|
}
|
|
|
|
.pixelgen-progress-value {
|
|
font-size: 9px;
|
|
color: var(--text-primary);
|
|
font-weight: 500;
|
|
}
|
|
|
|
.pixelgen-progress-bar {
|
|
height: 4px;
|
|
background: var(--bg-dark);
|
|
border-radius: 2px;
|
|
overflow: hidden;
|
|
}
|
|
|
|
.pixelgen-progress-fill {
|
|
height: 100%;
|
|
background: linear-gradient(90deg, var(--primary), var(--secondary));
|
|
border-radius: 2px;
|
|
transition: width 0.3s ease;
|
|
width: 0%;
|
|
}
|
|
|
|
/* Loss Chart */
|
|
.pixelgen-chart-section {
|
|
background: var(--bg-elevated);
|
|
border-radius: 6px;
|
|
padding: 8px 10px;
|
|
flex: 1;
|
|
min-height: 80px;
|
|
display: flex;
|
|
flex-direction: column;
|
|
overflow: hidden;
|
|
}
|
|
|
|
.pixelgen-chart-header {
|
|
font-size: 9px;
|
|
color: var(--text-secondary);
|
|
margin-bottom: 4px;
|
|
}
|
|
|
|
.pixelgen-chart-canvas {
|
|
width: 100%;
|
|
height: 100%;
|
|
display: block;
|
|
}
|
|
|
|
/* Status */
|
|
.pixelgen-status {
|
|
background: var(--bg-elevated);
|
|
border-radius: 6px;
|
|
padding: 6px 10px;
|
|
font-size: 10px;
|
|
color: var(--text-secondary);
|
|
text-align: center;
|
|
border-left: 3px solid var(--primary);
|
|
flex-shrink: 0;
|
|
}
|
|
|
|
.pixelgen-status.error {
|
|
border-left-color: var(--danger);
|
|
color: var(--danger);
|
|
}
|
|
|
|
.pixelgen-status.success {
|
|
border-left-color: var(--success);
|
|
color: var(--success);
|
|
}
|
|
|
|
/* Validation Carousel */
|
|
.pixelgen-preview-section {
|
|
background: var(--bg-elevated);
|
|
border-radius: 6px;
|
|
padding: 8px 10px;
|
|
flex: 1;
|
|
min-height: 80px;
|
|
display: flex;
|
|
flex-direction: column;
|
|
overflow: hidden;
|
|
}
|
|
|
|
.pixelgen-preview-header {
|
|
font-size: 9px;
|
|
color: var(--text-secondary);
|
|
margin-bottom: 6px;
|
|
flex-shrink: 0;
|
|
}
|
|
|
|
.pixelgen-preview-carousel {
|
|
display: flex;
|
|
gap: 6px;
|
|
overflow-x: auto;
|
|
overflow-y: hidden;
|
|
flex: 1;
|
|
align-items: center;
|
|
padding-bottom: 4px;
|
|
scrollbar-width: thin;
|
|
scrollbar-color: var(--border-hover) transparent;
|
|
}
|
|
|
|
.pixelgen-preview-empty {
|
|
width: 100%;
|
|
text-align: center;
|
|
font-size: 10px;
|
|
color: var(--text-muted);
|
|
}
|
|
|
|
.pixelgen-preview-carousel::-webkit-scrollbar {
|
|
height: 4px;
|
|
}
|
|
.pixelgen-preview-carousel::-webkit-scrollbar-track {
|
|
background: transparent;
|
|
}
|
|
.pixelgen-preview-carousel::-webkit-scrollbar-thumb {
|
|
background: var(--border-hover);
|
|
border-radius: 2px;
|
|
}
|
|
|
|
.pixelgen-preview-tile {
|
|
flex-shrink: 0;
|
|
display: flex;
|
|
flex-direction: column;
|
|
align-items: center;
|
|
gap: 3px;
|
|
}
|
|
|
|
.pixelgen-preview-tile img {
|
|
width: 64px;
|
|
height: 64px;
|
|
object-fit: cover;
|
|
border-radius: 4px;
|
|
border: 1px solid var(--border);
|
|
cursor: pointer;
|
|
transition: border-color 0.2s;
|
|
}
|
|
|
|
.pixelgen-preview-tile img:hover {
|
|
border-color: var(--primary);
|
|
}
|
|
|
|
.pixelgen-preview-tile .tile-label {
|
|
font-size: 9px;
|
|
color: var(--text-muted);
|
|
text-align: center;
|
|
white-space: nowrap;
|
|
}
|
|
|
|
/* Lightbox overlay */
|
|
.pixelgen-lightbox {
|
|
position: fixed;
|
|
top: 0; left: 0;
|
|
width: 100%; height: 100%;
|
|
background: rgba(0, 0, 0, 0.92);
|
|
backdrop-filter: blur(4px);
|
|
z-index: 10000;
|
|
display: flex;
|
|
align-items: center;
|
|
justify-content: center;
|
|
animation: pixelgen-fade-in 0.2s ease-out;
|
|
}
|
|
|
|
@keyframes pixelgen-fade-in {
|
|
from { opacity: 0; }
|
|
to { opacity: 1; }
|
|
}
|
|
|
|
.pixelgen-lightbox img {
|
|
max-width: 90vw;
|
|
max-height: 85vh;
|
|
object-fit: contain;
|
|
border-radius: 8px;
|
|
box-shadow: 0 8px 32px rgba(0, 0, 0, 0.8);
|
|
}
|
|
|
|
.pixelgen-lightbox-nav {
|
|
position: absolute;
|
|
top: 50%;
|
|
transform: translateY(-50%);
|
|
width: 48px; height: 48px;
|
|
border-radius: 50%;
|
|
background: rgba(0, 0, 0, 0.6);
|
|
border: 2px solid rgba(255, 255, 255, 0.3);
|
|
cursor: pointer;
|
|
display: flex;
|
|
align-items: center;
|
|
justify-content: center;
|
|
color: white;
|
|
font-size: 22px;
|
|
transition: all 0.2s ease;
|
|
}
|
|
.pixelgen-lightbox-nav:hover {
|
|
background: rgba(6, 182, 212, 0.6);
|
|
border-color: #06b6d4;
|
|
}
|
|
.pixelgen-lightbox-nav.left { left: 20px; }
|
|
.pixelgen-lightbox-nav.right { right: 20px; }
|
|
|
|
.pixelgen-lightbox-close {
|
|
position: absolute;
|
|
top: 16px; right: 16px;
|
|
width: 40px; height: 40px;
|
|
border-radius: 8px;
|
|
background: rgba(255, 255, 255, 0.08);
|
|
border: 1px solid rgba(255, 255, 255, 0.15);
|
|
cursor: pointer;
|
|
color: white;
|
|
font-size: 18px;
|
|
display: flex;
|
|
align-items: center;
|
|
justify-content: center;
|
|
transition: all 0.2s ease;
|
|
}
|
|
.pixelgen-lightbox-close:hover {
|
|
background: rgba(239, 68, 68, 0.6);
|
|
}
|
|
|
|
.pixelgen-lightbox-label {
|
|
position: absolute;
|
|
bottom: 20px;
|
|
color: rgba(255, 255, 255, 0.7);
|
|
font-size: 14px;
|
|
font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
|
|
}
|
|
`;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// TrainingWidget — 1:1 visual match with AceStep Training widget
|
|
// ---------------------------------------------------------------------------
|
|
|
|
class TrainingWidget {
|
|
constructor(options) {
|
|
this.node = options.node;
|
|
this.container = options.container;
|
|
this.element = document.createElement("div");
|
|
this.element.className = "pixelgen-training-widget";
|
|
|
|
// UI elements
|
|
this.badgeEl = null;
|
|
this.epochValueEl = null;
|
|
this.stepValueEl = null;
|
|
this.lossValueEl = null;
|
|
this.progressFillEl = null;
|
|
this.progressLabelEl = null;
|
|
this.statusEl = null;
|
|
this.canvasEl = null;
|
|
|
|
// State
|
|
this.lossHistory = [];
|
|
this.isTraining = false;
|
|
|
|
// Resize handling
|
|
this.resizeObserver = null;
|
|
this.resizeTimeout = null;
|
|
|
|
this.injectStyles();
|
|
this.createUI();
|
|
this.container.appendChild(this.element);
|
|
}
|
|
|
|
injectStyles() {
|
|
const styleId = "pixelgen-training-styles";
|
|
if (!document.getElementById(styleId)) {
|
|
const style = document.createElement("style");
|
|
style.id = styleId;
|
|
style.textContent = TRAINING_WIDGET_STYLES;
|
|
document.head.appendChild(style);
|
|
}
|
|
}
|
|
|
|
createUI() {
|
|
this.element.innerHTML = `
|
|
<div class="pixelgen-header">
|
|
<div class="pixelgen-title">
|
|
<span>PixelGen Training</span>
|
|
<span class="pixelgen-badge idle">Idle</span>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="pixelgen-content">
|
|
<div class="pixelgen-stats">
|
|
<div class="pixelgen-stat">
|
|
<span class="pixelgen-stat-label">Epoch</span>
|
|
<span class="pixelgen-stat-value" data-stat="epoch">0/0</span>
|
|
</div>
|
|
<div class="pixelgen-stat">
|
|
<span class="pixelgen-stat-label">Step</span>
|
|
<span class="pixelgen-stat-value" data-stat="step">0</span>
|
|
</div>
|
|
<div class="pixelgen-stat">
|
|
<span class="pixelgen-stat-label">Loss</span>
|
|
<span class="pixelgen-stat-value" data-stat="loss">-</span>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="pixelgen-progress-section">
|
|
<div class="pixelgen-progress-header">
|
|
<span class="pixelgen-progress-label">Training Progress</span>
|
|
<span class="pixelgen-progress-value" data-progress-label>0%</span>
|
|
</div>
|
|
<div class="pixelgen-progress-bar">
|
|
<div class="pixelgen-progress-fill" data-progress-fill></div>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="pixelgen-chart-section">
|
|
<div class="pixelgen-chart-header">Loss History</div>
|
|
<canvas class="pixelgen-chart-canvas"></canvas>
|
|
</div>
|
|
|
|
<div class="pixelgen-status">
|
|
Ready to train
|
|
</div>
|
|
|
|
<div class="pixelgen-preview-section">
|
|
<div class="pixelgen-preview-header">Validation Samples</div>
|
|
<div class="pixelgen-preview-carousel">
|
|
<div class="pixelgen-preview-empty">Samples will appear at each checkpoint</div>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
`;
|
|
|
|
// Get references
|
|
this.badgeEl = this.element.querySelector(".pixelgen-badge");
|
|
this.epochValueEl = this.element.querySelector('[data-stat="epoch"]');
|
|
this.stepValueEl = this.element.querySelector('[data-stat="step"]');
|
|
this.lossValueEl = this.element.querySelector('[data-stat="loss"]');
|
|
this.progressFillEl = this.element.querySelector("[data-progress-fill]");
|
|
this.progressLabelEl = this.element.querySelector("[data-progress-label]");
|
|
this.statusEl = this.element.querySelector(".pixelgen-status");
|
|
this.canvasEl = this.element.querySelector(".pixelgen-chart-canvas");
|
|
|
|
// Set up ResizeObserver for canvas scaling - watch container, not canvas
|
|
const chartSection = this.element.querySelector(".pixelgen-chart-section");
|
|
if (this.canvasEl && chartSection) {
|
|
this.resizeObserver = new ResizeObserver(() => {
|
|
if (this.resizeTimeout) {
|
|
clearTimeout(this.resizeTimeout);
|
|
}
|
|
this.resizeTimeout = window.setTimeout(() => {
|
|
this.drawChart();
|
|
}, 16); // ~60fps max
|
|
});
|
|
this.resizeObserver.observe(chartSection);
|
|
}
|
|
|
|
// Draw initial chart
|
|
this.drawChart();
|
|
}
|
|
|
|
updateProgress(epoch, totalEpochs, loss, step) {
|
|
this.isTraining = true;
|
|
|
|
// Update badge
|
|
if (this.badgeEl) {
|
|
this.badgeEl.textContent = "Training";
|
|
this.badgeEl.className = "pixelgen-badge training";
|
|
}
|
|
|
|
// Update stats
|
|
if (this.epochValueEl) {
|
|
this.epochValueEl.textContent = `${epoch}/${totalEpochs}`;
|
|
}
|
|
if (this.stepValueEl) {
|
|
this.stepValueEl.textContent = step.toString();
|
|
}
|
|
if (this.lossValueEl) {
|
|
this.lossValueEl.textContent = loss.toFixed(6);
|
|
}
|
|
|
|
// Update progress bar
|
|
const progress = (epoch / totalEpochs) * 100;
|
|
if (this.progressFillEl) {
|
|
this.progressFillEl.style.width = `${progress}%`;
|
|
}
|
|
if (this.progressLabelEl) {
|
|
this.progressLabelEl.textContent = `${progress.toFixed(1)}%`;
|
|
}
|
|
}
|
|
|
|
updateLossHistory(history) {
|
|
this.lossHistory = history;
|
|
this.drawChart();
|
|
}
|
|
|
|
updateStatus(message, type = "normal") {
|
|
if (this.statusEl) {
|
|
this.statusEl.textContent = message;
|
|
this.statusEl.className = "pixelgen-status";
|
|
if (type !== "normal") {
|
|
this.statusEl.classList.add(type);
|
|
}
|
|
}
|
|
}
|
|
|
|
onTrainingComplete(finalPath) {
|
|
this.isTraining = false;
|
|
|
|
if (this.badgeEl) {
|
|
this.badgeEl.textContent = "Complete";
|
|
this.badgeEl.className = "pixelgen-badge";
|
|
this.badgeEl.style.background = "#22c55e";
|
|
}
|
|
|
|
this.updateStatus(`Training complete! Saved to: ${finalPath}`, "success");
|
|
}
|
|
|
|
addValidationImage(imageBase64, checkpointPath) {
|
|
const carousel = this.element.querySelector(".pixelgen-preview-carousel");
|
|
if (!carousel) return;
|
|
|
|
// Clear placeholder on first image
|
|
const empty = carousel.querySelector(".pixelgen-preview-empty");
|
|
if (empty) empty.remove();
|
|
|
|
// Extract a short label from the checkpoint path
|
|
const parts = checkpointPath.replace(/\\/g, "/").split("/");
|
|
const label = parts.find((p) => p.startsWith("checkpoint-epoch-")) || parts[parts.length - 1];
|
|
const shortLabel = label.replace("checkpoint-epoch-", "E");
|
|
|
|
// Create tile
|
|
const tile = document.createElement("div");
|
|
tile.className = "pixelgen-preview-tile";
|
|
|
|
const img = document.createElement("img");
|
|
img.src = imageBase64;
|
|
img.title = checkpointPath;
|
|
|
|
img.addEventListener("click", () => {
|
|
const tiles = carousel.querySelectorAll(".pixelgen-preview-tile");
|
|
const images = Array.from(tiles).map((t) => ({
|
|
src: t.querySelector("img").src,
|
|
label: t.querySelector(".tile-label")?.textContent || "",
|
|
}));
|
|
const index = Array.from(tiles).indexOf(tile);
|
|
this.openLightbox(images, index >= 0 ? index : 0);
|
|
});
|
|
|
|
const labelEl = document.createElement("div");
|
|
labelEl.className = "tile-label";
|
|
labelEl.textContent = shortLabel;
|
|
|
|
tile.appendChild(img);
|
|
tile.appendChild(labelEl);
|
|
carousel.appendChild(tile);
|
|
|
|
// Auto-scroll to the newest tile
|
|
carousel.scrollLeft = carousel.scrollWidth;
|
|
}
|
|
|
|
openLightbox(images, startIndex) {
|
|
if (this._lightbox) this.closeLightbox();
|
|
if (!images.length) return;
|
|
|
|
let currentIndex = startIndex;
|
|
|
|
const overlay = document.createElement("div");
|
|
overlay.className = "pixelgen-lightbox";
|
|
|
|
const mainImg = document.createElement("img");
|
|
const label = document.createElement("div");
|
|
label.className = "pixelgen-lightbox-label";
|
|
|
|
const update = () => {
|
|
mainImg.src = images[currentIndex].src;
|
|
label.textContent = `${images[currentIndex].label} (${currentIndex + 1} / ${images.length})`;
|
|
};
|
|
|
|
// Nav arrows
|
|
const leftBtn = document.createElement("button");
|
|
leftBtn.className = "pixelgen-lightbox-nav left";
|
|
leftBtn.innerHTML = "‹";
|
|
leftBtn.onclick = (e) => { e.stopPropagation(); currentIndex = (currentIndex - 1 + images.length) % images.length; update(); };
|
|
|
|
const rightBtn = document.createElement("button");
|
|
rightBtn.className = "pixelgen-lightbox-nav right";
|
|
rightBtn.innerHTML = "›";
|
|
rightBtn.onclick = (e) => { e.stopPropagation(); currentIndex = (currentIndex + 1) % images.length; update(); };
|
|
|
|
// Close button
|
|
const closeBtn = document.createElement("button");
|
|
closeBtn.className = "pixelgen-lightbox-close";
|
|
closeBtn.innerHTML = "✕";
|
|
closeBtn.onclick = () => this.closeLightbox();
|
|
|
|
// Click backdrop to close
|
|
overlay.onclick = (e) => { if (e.target === overlay) this.closeLightbox(); };
|
|
|
|
// Keyboard navigation
|
|
this._lightboxKeyHandler = (e) => {
|
|
if (e.key === "Escape") this.closeLightbox();
|
|
if (e.key === "ArrowLeft") { currentIndex = (currentIndex - 1 + images.length) % images.length; update(); }
|
|
if (e.key === "ArrowRight") { currentIndex = (currentIndex + 1) % images.length; update(); }
|
|
};
|
|
document.addEventListener("keydown", this._lightboxKeyHandler);
|
|
|
|
overlay.append(mainImg, leftBtn, rightBtn, closeBtn, label);
|
|
document.body.appendChild(overlay);
|
|
this._lightbox = overlay;
|
|
update();
|
|
}
|
|
|
|
closeLightbox() {
|
|
if (!this._lightbox) return;
|
|
if (this._lightboxKeyHandler) {
|
|
document.removeEventListener("keydown", this._lightboxKeyHandler);
|
|
this._lightboxKeyHandler = null;
|
|
}
|
|
this._lightbox.style.animation = "pixelgen-fade-in 0.15s ease-in reverse";
|
|
const el = this._lightbox;
|
|
this._lightbox = null;
|
|
setTimeout(() => {
|
|
if (el.parentNode) el.parentNode.removeChild(el);
|
|
}, 150);
|
|
}
|
|
|
|
reset() {
|
|
this.lossHistory = [];
|
|
this.isTraining = false;
|
|
|
|
if (this.badgeEl) {
|
|
this.badgeEl.textContent = "Training";
|
|
this.badgeEl.className = "pixelgen-badge training";
|
|
}
|
|
|
|
if (this.epochValueEl) this.epochValueEl.textContent = "0/0";
|
|
if (this.stepValueEl) this.stepValueEl.textContent = "0";
|
|
if (this.lossValueEl) this.lossValueEl.textContent = "-";
|
|
|
|
if (this.progressFillEl) this.progressFillEl.style.width = "0%";
|
|
if (this.progressLabelEl) this.progressLabelEl.textContent = "0%";
|
|
|
|
this.updateStatus("Starting training...");
|
|
this.drawChart();
|
|
|
|
// Clear preview carousel and restore placeholder
|
|
const carousel = this.element.querySelector(".pixelgen-preview-carousel");
|
|
if (carousel) {
|
|
carousel.innerHTML = '<div class="pixelgen-preview-empty">Samples will appear at each checkpoint</div>';
|
|
}
|
|
}
|
|
|
|
drawChart() {
|
|
if (!this.canvasEl) return;
|
|
|
|
const canvas = this.canvasEl;
|
|
const ctx = canvas.getContext("2d");
|
|
if (!ctx) return;
|
|
|
|
// Get actual size
|
|
const rect = canvas.getBoundingClientRect();
|
|
canvas.width = rect.width * window.devicePixelRatio;
|
|
canvas.height = rect.height * window.devicePixelRatio;
|
|
ctx.scale(window.devicePixelRatio, window.devicePixelRatio);
|
|
|
|
const w = rect.width;
|
|
const h = rect.height;
|
|
|
|
// Clear
|
|
ctx.fillStyle = "#0f0f12";
|
|
ctx.fillRect(0, 0, w, h);
|
|
|
|
// No data message
|
|
if (this.lossHistory.length < 2) {
|
|
ctx.fillStyle = "#71717a";
|
|
ctx.font = "11px Inter, sans-serif";
|
|
ctx.textAlign = "center";
|
|
ctx.fillText("Waiting for training data...", w / 2, h / 2);
|
|
return;
|
|
}
|
|
|
|
// Calculate scales
|
|
const padding = { top: 20, right: 20, bottom: 25, left: 50 };
|
|
const chartW = w - padding.left - padding.right;
|
|
const chartH = h - padding.top - padding.bottom;
|
|
|
|
const maxStep = Math.max(...this.lossHistory.map((d) => d.step));
|
|
const minStep = Math.min(...this.lossHistory.map((d) => d.step));
|
|
const maxLoss = Math.max(...this.lossHistory.map((d) => d.loss));
|
|
const minLoss = Math.min(...this.lossHistory.map((d) => d.loss));
|
|
const lossRange = maxLoss - minLoss || 1;
|
|
|
|
// Draw grid
|
|
ctx.strokeStyle = "#27272a";
|
|
ctx.lineWidth = 0.5;
|
|
|
|
for (let i = 0; i <= 4; i++) {
|
|
const y = padding.top + (chartH / 4) * i;
|
|
ctx.beginPath();
|
|
ctx.moveTo(padding.left, y);
|
|
ctx.lineTo(w - padding.right, y);
|
|
ctx.stroke();
|
|
|
|
// Y-axis labels
|
|
const lossVal = maxLoss - (lossRange / 4) * i;
|
|
ctx.fillStyle = "#71717a";
|
|
ctx.font = "9px Inter, sans-serif";
|
|
ctx.textAlign = "right";
|
|
ctx.fillText(lossVal.toFixed(4), padding.left - 5, y + 3);
|
|
}
|
|
|
|
// Draw loss line
|
|
ctx.strokeStyle = "#06b6d4";
|
|
ctx.lineWidth = 2;
|
|
ctx.lineJoin = "round";
|
|
ctx.lineCap = "round";
|
|
ctx.beginPath();
|
|
|
|
this.lossHistory.forEach((point, i) => {
|
|
const x = padding.left + ((point.step - minStep) / (maxStep - minStep || 1)) * chartW;
|
|
const y = padding.top + chartH - ((point.loss - minLoss) / lossRange) * chartH;
|
|
|
|
if (i === 0) {
|
|
ctx.moveTo(x, y);
|
|
} else {
|
|
ctx.lineTo(x, y);
|
|
}
|
|
});
|
|
|
|
ctx.stroke();
|
|
|
|
// Draw gradient fill
|
|
const gradient = ctx.createLinearGradient(0, padding.top, 0, h - padding.bottom);
|
|
gradient.addColorStop(0, "rgba(6, 182, 212, 0.3)");
|
|
gradient.addColorStop(1, "rgba(6, 182, 212, 0)");
|
|
|
|
ctx.fillStyle = gradient;
|
|
ctx.beginPath();
|
|
|
|
this.lossHistory.forEach((point, i) => {
|
|
const x = padding.left + ((point.step - minStep) / (maxStep - minStep || 1)) * chartW;
|
|
const y = padding.top + chartH - ((point.loss - minLoss) / lossRange) * chartH;
|
|
|
|
if (i === 0) {
|
|
ctx.moveTo(x, y);
|
|
} else {
|
|
ctx.lineTo(x, y);
|
|
}
|
|
});
|
|
|
|
// Close the path
|
|
const lastPoint = this.lossHistory[this.lossHistory.length - 1];
|
|
const lastX = padding.left + ((lastPoint.step - minStep) / (maxStep - minStep || 1)) * chartW;
|
|
ctx.lineTo(lastX, padding.top + chartH);
|
|
ctx.lineTo(padding.left, padding.top + chartH);
|
|
ctx.closePath();
|
|
ctx.fill();
|
|
|
|
// Current loss label
|
|
ctx.fillStyle = "#fafafa";
|
|
ctx.font = "bold 10px Inter, sans-serif";
|
|
ctx.textAlign = "left";
|
|
ctx.fillText(`Loss: ${lastPoint.loss.toFixed(6)}`, padding.left + 5, padding.top + 12);
|
|
}
|
|
|
|
dispose() {
|
|
this.closeLightbox();
|
|
// Clean up ResizeObserver
|
|
if (this.resizeObserver) {
|
|
this.resizeObserver.disconnect();
|
|
this.resizeObserver = null;
|
|
}
|
|
if (this.resizeTimeout) {
|
|
clearTimeout(this.resizeTimeout);
|
|
this.resizeTimeout = null;
|
|
}
|
|
|
|
if (this.element && this.element.parentNode) {
|
|
this.element.parentNode.removeChild(this.element);
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ComfyUI Extension Registration
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const widgetInstances = new Map();
|
|
|
|
function createTrainingWidget(node) {
|
|
const container = document.createElement("div");
|
|
container.id = `pixelgen-training-widget-${node.id}`;
|
|
container.style.width = "100%";
|
|
container.style.height = "100%";
|
|
container.style.minHeight = "280px";
|
|
|
|
const widget = node.addDOMWidget(
|
|
"training_ui",
|
|
"training-widget",
|
|
container,
|
|
{
|
|
getMinHeight: () => 400,
|
|
hideOnZoom: false,
|
|
serialize: false,
|
|
}
|
|
);
|
|
|
|
setTimeout(() => {
|
|
const trainingWidget = new TrainingWidget({
|
|
node,
|
|
container,
|
|
});
|
|
widgetInstances.set(node.id, trainingWidget);
|
|
}, 100);
|
|
|
|
widget.onRemove = () => {
|
|
const instance = widgetInstances.get(node.id);
|
|
if (instance) {
|
|
instance.dispose();
|
|
widgetInstances.delete(node.id);
|
|
}
|
|
};
|
|
|
|
return { widget };
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "ComfyUI.FL_PixelGen_Training",
|
|
|
|
nodeCreated(node) {
|
|
const comfyClass = (node.constructor && node.constructor.comfyClass) || "";
|
|
|
|
// FL_ theming for all PixelGen nodes
|
|
if (comfyClass.startsWith("FL_PixelGen")) {
|
|
node.color = "#16727c";
|
|
node.bgcolor = "#4F0074";
|
|
}
|
|
|
|
// Training widget only on Train node
|
|
if (comfyClass !== "FL_PixelGen_Train") {
|
|
return;
|
|
}
|
|
|
|
const [oldWidth, oldHeight] = node.size;
|
|
node.setSize([Math.max(oldWidth, 400), Math.max(oldHeight, 600)]);
|
|
createTrainingWidget(node);
|
|
},
|
|
});
|
|
|
|
// Listen for training progress events
|
|
api.addEventListener("pixelgen.training.progress", (event) => {
|
|
const detail = event.detail;
|
|
if (!detail || !detail.node) return;
|
|
|
|
const nodeId = parseInt(detail.node, 10);
|
|
const widget = widgetInstances.get(nodeId);
|
|
if (!widget) return;
|
|
|
|
switch (detail.type) {
|
|
case "progress":
|
|
widget.updateProgress(
|
|
detail.epoch ?? 0,
|
|
detail.total_epochs ?? 0,
|
|
detail.loss ?? 0,
|
|
detail.step ?? 0
|
|
);
|
|
if (detail.loss_history) {
|
|
widget.updateLossHistory(detail.loss_history);
|
|
}
|
|
break;
|
|
|
|
case "status":
|
|
widget.updateStatus(detail.message ?? "");
|
|
break;
|
|
|
|
case "checkpoint":
|
|
if (detail.checkpoint_path) {
|
|
widget.updateStatus(`Checkpoint saved: ${detail.checkpoint_path}`, "success");
|
|
}
|
|
if (detail.validation_image && detail.checkpoint_path) {
|
|
widget.addValidationImage(detail.validation_image, detail.checkpoint_path);
|
|
}
|
|
break;
|
|
|
|
case "complete":
|
|
if (detail.final_path) {
|
|
widget.onTrainingComplete(detail.final_path);
|
|
}
|
|
break;
|
|
}
|
|
});
|
|
|
|
// Listen for node execution results
|
|
api.addEventListener("executed", (event) => {
|
|
const detail = event.detail;
|
|
if (!detail || !detail.node || !detail.output) return;
|
|
|
|
const nodeId = parseInt(detail.node, 10);
|
|
const widget = widgetInstances.get(nodeId);
|
|
if (!widget) return;
|
|
|
|
const finalPath = detail.output && detail.output.lora_path;
|
|
if (finalPath && finalPath.length > 0) {
|
|
widget.onTrainingComplete(finalPath[0]);
|
|
}
|
|
});
|
|
|
|
// Listen for execution start to reset the widget
|
|
api.addEventListener("executing", (event) => {
|
|
const detail = event.detail;
|
|
if (!detail || !detail.node) return;
|
|
|
|
const nodeId = parseInt(detail.node, 10);
|
|
const widget = widgetInstances.get(nodeId);
|
|
if (widget) {
|
|
widget.reset();
|
|
}
|
|
});
|