benfinklea-ComfyUI-ModelSync/js/modelsync-panel.js

/**
 * ModelSync UI Panel for ComfyUI
 *
 * Provides a user interface for:
 * - Viewing tracked models and their update status
 * - Configuring auto-update settings
 * - Managing sync targets
 * - Triggering manual syncs
 */

const { app } = window.comfyAPI ? window.comfyAPI.app : { app: null };

class ModelSyncPanel {
  constructor() {
    this.element = null;
    this.models = [];
    this.targets = [];
    this.config = {};
    this.currentTab = "models";
    this.searchQuery = "";
    this.filterStatus = "all"; // all, updates, auto
  }

  async init() {
    // Create menu button
    this.createMenuButton();

    // Load initial data
    await this.refreshData();
  }

  async createMenuButton() {
    const self = this;

    // Try to add new-style button first (ComfyUI 2024+)
    try {
      const { ComfyButton } = await import("../../scripts/ui/components/button.js");
      const { ComfyButtonGroup } = await import("../../scripts/ui/components/buttonGroup.js");

      // Wait for app.menu to be available
      const waitForMenu = () => {
        return new Promise((resolve) => {
          const check = () => {
            if (window.app?.menu?.settingsGroup) {
              resolve(true);
            } else {
              setTimeout(check, 100);
            }
          };
          check();
          // Timeout after 10 seconds
          setTimeout(() => resolve(false), 10000);
        });
      };

      const menuReady = await waitForMenu();
      if (menuReady && window.app?.menu?.settingsGroup) {
        const syncButton = new ComfyButton({
          icon: "sync",
          action: () => self.toggle(),
          tooltip: "Model Sync - Track and sync model updates",
          content: "Model Sync",
          classList: "comfyui-button comfyui-menu-mobile-collapse"
        });

        const syncGroup = new ComfyButtonGroup(syncButton.element);
        window.app.menu.settingsGroup.element.before(syncGroup.element);
        console.log("[ModelSync] New-style menu button added");
        return;
      }
    } catch (e) {
      console.log("[ModelSync] New-style menu not available, trying legacy:", e.message);
    }

    // Fall back to old-style button
    const tryCreateButton = () => {
      const menu = document.querySelector(".comfy-menu") || document.querySelector(".comfyui-menu");
      if (!menu) {
        setTimeout(tryCreateButton, 500);
        return;
      }

      // Check if button already exists
      if (document.querySelector(".modelsync-menu-btn")) {
        return;
      }

      const button = document.createElement("button");
      button.textContent = "Model Sync";
      button.classList.add("modelsync-menu-btn");
      button.style.cssText = "margin: 4px; padding: 4px 8px;";
      button.addEventListener("click", () => this.toggle());

      // Insert into menu
      menu.appendChild(button);
      console.log("[ModelSync] Legacy menu button added");
    };

    tryCreateButton();
  }

  async refreshData() {
    try {
      const [modelsRes, targetsRes, configRes] = await Promise.all([
        fetch("/modelsync/models"),
        fetch("/modelsync/targets"),
        fetch("/modelsync/config"),
      ]);

      if (modelsRes.ok) {
        const data = await modelsRes.json();
        this.models = data.models || [];
      }

      if (targetsRes.ok) {
        const data = await targetsRes.json();
        this.targets = data.targets || [];
      }

      if (configRes.ok) {
        this.config = await configRes.json();
      }

      this.render();
    } catch (error) {
      console.error("[ModelSync] Error refreshing data:", error);
    }
  }

  toggle() {
    if (this.element) {
      this.close();
    } else {
      this.open();
    }
  }

  open() {
    this.element = document.createElement("div");
    this.element.classList.add("modelsync-panel");
    this.element.innerHTML = this.getTemplate();

    document.body.appendChild(this.element);
    this.attachEventListeners();
    this.refreshData();
  }

  close() {
    if (this.element) {
      this.element.remove();
      this.element = null;
    }
  }

  getTemplate() {
    return `
      <div class="modelsync-overlay" data-action="close"></div>
      <div class="modelsync-content">
        <div class="modelsync-header">
          <h2>Model Sync</h2>
          <button class="modelsync-close" data-action="close">&times;</button>
        </div>

        <div class="modelsync-tabs">
          <button class="modelsync-tab active" data-tab="models">Models</button>
          <button class="modelsync-tab" data-tab="targets">Sync Targets</button>
          <button class="modelsync-tab" data-tab="settings">Settings</button>
        </div>

        <div class="modelsync-body">
          <div class="modelsync-tab-content" id="modelsync-models">
            <div class="modelsync-toolbar">
              <input type="text" class="modelsync-search" placeholder="Search models..." />
              <select class="modelsync-filter">
                <option value="all">All Models</option>
                <option value="updates">Updates Available</option>
                <option value="auto">Auto-Update Enabled</option>
              </select>
              <button class="modelsync-btn" data-action="scan">Scan</button>
              <button class="modelsync-btn primary" data-action="lookup">Lookup Sources</button>
              <button class="modelsync-btn" data-action="stop-lookup" style="display:none;background:#c00;">Stop</button>
              <button class="modelsync-btn" data-action="check">Check Updates</button>
            </div>
            <div class="modelsync-models-list"></div>
          </div>

          <div class="modelsync-tab-content hidden" id="modelsync-targets">
            <div class="modelsync-toolbar">
              <button class="modelsync-btn primary" data-action="add-target">Add Target</button>
              <button class="modelsync-btn" data-action="sync-all">Sync All</button>
            </div>
            <div class="modelsync-targets-list"></div>
          </div>

          <div class="modelsync-tab-content hidden" id="modelsync-settings">
            <div class="modelsync-settings-form">
              <div class="modelsync-setting">
                <label>
                  <input type="checkbox" id="auto-update-enabled" />
                  Enable Auto-Update
                </label>
                <p class="modelsync-hint">Automatically download updates for enabled models</p>
              </div>

              <div class="modelsync-setting">
                <label for="check-interval">Check Interval (hours)</label>
                <input type="number" id="check-interval" min="1" max="168" value="24" />
              </div>

              <div class="modelsync-setting">
                <label>
                  <input type="checkbox" id="verify-downloads" checked />
                  Verify Downloads
                </label>
                <p class="modelsync-hint">Verify file hash after downloading</p>
              </div>

              <div class="modelsync-setting">
                <label>
                  <input type="checkbox" id="backup-before-update" checked />
                  Backup Before Update
                </label>
                <p class="modelsync-hint">Create .bak file before replacing model</p>
              </div>

              <button class="modelsync-btn primary" data-action="save-settings">Save Settings</button>
            </div>
          </div>
        </div>

        <div class="modelsync-footer">
          <span class="modelsync-status"></span>
        </div>
      </div>
    `;
  }

  attachEventListeners() {
    if (!this.element) return;

    // Close button and overlay
    this.element.querySelectorAll("[data-action='close']").forEach((el) => {
      el.addEventListener("click", () => this.close());
    });

    // Tabs
    this.element.querySelectorAll(".modelsync-tab").forEach((tab) => {
      tab.addEventListener("click", () => {
        this.currentTab = tab.dataset.tab;
        this.updateTabs();
      });
    });

    // Search
    const search = this.element.querySelector(".modelsync-search");
    if (search) {
      search.addEventListener("input", (e) => {
        this.searchQuery = e.target.value.toLowerCase();
        this.renderModelsList();
      });
    }

    // Filter
    const filter = this.element.querySelector(".modelsync-filter");
    if (filter) {
      filter.addEventListener("change", (e) => {
        this.filterStatus = e.target.value;
        this.renderModelsList();
      });
    }

    // Action buttons
    this.element.querySelectorAll("[data-action]").forEach((btn) => {
      btn.addEventListener("click", () => this.handleAction(btn.dataset.action));
    });
  }

  async handleAction(action) {
    switch (action) {
      case "scan":
        await this.scanModels();
        break;
      case "lookup":
        await this.lookupSources();
        break;
      case "stop-lookup":
        this.stopLookup();
        break;
      case "check":
        await this.checkUpdates();
        break;
      case "add-target":
        this.showAddTargetDialog();
        break;
      case "sync-all":
        await this.syncAll();
        break;
      case "save-settings":
        await this.saveSettings();
        break;
    }
  }

  async scanModels() {
    this.setStatus("Scanning models...");
    try {
      const res = await fetch("/modelsync/models/scan", { method: "POST" });
      const data = await res.json();
      this.setStatus(`Scanned ${data.scanned} models. New: ${data.new}, Updated: ${data.updated}`);
      await this.refreshData();
    } catch (error) {
      this.setStatus("Scan failed: " + error.message);
    }
  }

  async lookupSources() {
    this.lookupRunning = true;
    let foundCount = 0;
    let notFoundCount = 0;
    let processedCount = 0;

    // Show stop button, hide lookup button
    const lookupBtn = this.element?.querySelector('[data-action="lookup"]');
    const stopBtn = this.element?.querySelector('[data-action="stop-lookup"]');
    if (lookupBtn) lookupBtn.style.display = 'none';
    if (stopBtn) stopBtn.style.display = 'inline-block';

    const processNext = async () => {
      if (!this.lookupRunning) {
        if (lookupBtn) lookupBtn.style.display = 'inline-block';
        if (stopBtn) stopBtn.style.display = 'none';
        this.setStatus(`Lookup stopped. Found ${foundCount}, not found ${notFoundCount}.`);
        await this.refreshData();
        return;
      }

      try {
        const res = await fetch("/modelsync/lookup-sources", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({}),
        });
        const data = await res.json();

        if (data.status === "complete") {
          this.lookupRunning = false;
          if (lookupBtn) lookupBtn.style.display = 'inline-block';
          if (stopBtn) stopBtn.style.display = 'none';
          this.setStatus(`Lookup complete. Found ${foundCount} on CivitAI, ${notFoundCount} not found.`);
          await this.refreshData();
          return;
        }

        if (data.status === "found") {
          foundCount++;
          processedCount++;
        } else if (data.status === "not_found") {
          notFoundCount++;
          processedCount++;
        } else if (data.status === "hashing") {
          // Just hashed, will lookup next call
        }

        const progress = data.total > 0 ? Math.round(((data.total - data.remaining) / data.total) * 100) : 0;
        this.setStatus(`[${progress}%] ${data.message} (${foundCount} found, ${data.remaining} remaining)`);

        // Update UI periodically
        if (processedCount % 5 === 0) {
          await this.refreshData();
        }

        // Continue to next model
        setTimeout(processNext, 100);

      } catch (error) {
        this.lookupRunning = false;
        if (lookupBtn) lookupBtn.style.display = 'inline-block';
        if (stopBtn) stopBtn.style.display = 'none';
        this.setStatus("Lookup error: " + error.message);
        await this.refreshData();
      }
    };

    this.setStatus("Starting source lookup...");
    processNext();
  }

  stopLookup() {
    this.lookupRunning = false;
    this.setStatus("Stopping lookup...");
  }

  async checkUpdates() {
    this.setStatus("Checking for updates...");
    try {
      const res = await fetch("/modelsync/check", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({}),
      });
      const data = await res.json();
      this.setStatus(`Checked ${data.checked} models. Updates available: ${data.updates_available}`);
      await this.refreshData();
    } catch (error) {
      this.setStatus("Check failed: " + error.message);
    }
  }

  async syncAll() {
    this.setStatus("Syncing to all targets...");
    try {
      const res = await fetch("/modelsync/sync", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({}),
      });
      const data = await res.json();
      const total = data.results.reduce((sum, r) => sum + r.models_synced, 0);
      this.setStatus(`Sync complete. ${total} models synced.`);
    } catch (error) {
      this.setStatus("Sync failed: " + error.message);
    }
  }

  async saveSettings() {
    const settings = {
      auto_update_enabled: this.element.querySelector("#auto-update-enabled").checked.toString(),
      check_interval_hours: this.element.querySelector("#check-interval").value,
      verify_downloads: this.element.querySelector("#verify-downloads").checked.toString(),
      backup_before_update: this.element.querySelector("#backup-before-update").checked.toString(),
    };

    try {
      await fetch("/modelsync/config", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(settings),
      });
      this.setStatus("Settings saved");
      this.config = settings;
    } catch (error) {
      this.setStatus("Failed to save settings: " + error.message);
    }
  }

  showAddTargetDialog() {
    const name = prompt("Target name (e.g., 'frodo'):");
    if (!name) return;

    const sshHost = prompt("SSH host (e.g., 'user@host' or empty for local):");
    const modelPath = prompt("Models path on target:");
    if (!modelPath) return;

    const maxVram = prompt("Max VRAM in MB (empty for no limit):");

    this.addTarget({
      name,
      ssh_host: sshHost || null,
      model_path: modelPath,
      max_vram_mb: maxVram ? parseInt(maxVram) : null,
    });
  }

  async addTarget(target) {
    try {
      await fetch("/modelsync/targets", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(target),
      });
      this.setStatus("Target added");
      await this.refreshData();
    } catch (error) {
      this.setStatus("Failed to add target: " + error.message);
    }
  }

  async deleteTarget(targetId) {
    if (!confirm("Delete this sync target?")) return;

    try {
      await fetch(`/modelsync/targets/${targetId}`, { method: "DELETE" });
      this.setStatus("Target deleted");
      await this.refreshData();
    } catch (error) {
      this.setStatus("Failed to delete target: " + error.message);
    }
  }

  async syncToTarget(targetId) {
    this.setStatus("Syncing...");
    try {
      const res = await fetch(`/modelsync/sync/${targetId}`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({}),
      });
      const data = await res.json();
      this.setStatus(`Synced ${data.models_synced} models to ${data.target}`);
    } catch (error) {
      this.setStatus("Sync failed: " + error.message);
    }
  }

  async toggleAutoUpdate(modelId, enabled) {
    try {
      await fetch(`/modelsync/models/${modelId}`, {
        method: "PATCH",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ auto_update: enabled ? 1 : 0 }),
      });
    } catch (error) {
      console.error("[ModelSync] Failed to update model:", error);
    }
  }

  async downloadUpdate(modelId) {
    this.setStatus("Downloading update...");
    try {
      const res = await fetch(`/modelsync/download/${modelId}`, { method: "POST" });
      const data = await res.json();
      this.setStatus(data.success ? "Download complete" : "Download failed");
      await this.refreshData();
    } catch (error) {
      this.setStatus("Download failed: " + error.message);
    }
  }

  updateTabs() {
    if (!this.element) return;

    this.element.querySelectorAll(".modelsync-tab").forEach((tab) => {
      tab.classList.toggle("active", tab.dataset.tab === this.currentTab);
    });

    this.element.querySelectorAll(".modelsync-tab-content").forEach((content) => {
      content.classList.toggle("hidden", !content.id.includes(this.currentTab));
    });
  }

  render() {
    if (!this.element) return;

    this.renderModelsList();
    this.renderTargetsList();
    this.renderSettings();
  }

  renderModelsList() {
    const container = this.element?.querySelector(".modelsync-models-list");
    if (!container) return;

    let filtered = this.models;

    // Apply search filter
    if (this.searchQuery) {
      filtered = filtered.filter((m) =>
        m.filename.toLowerCase().includes(this.searchQuery)
      );
    }

    // Apply status filter
    if (this.filterStatus === "updates") {
      filtered = filtered.filter((m) => m.update_available);
    } else if (this.filterStatus === "auto") {
      filtered = filtered.filter((m) => m.auto_update);
    }

    if (filtered.length === 0) {
      container.innerHTML = '<p class="modelsync-empty">No models found. Click "Scan" to discover models.</p>';
      return;
    }

    container.innerHTML = `
      <table class="modelsync-table">
        <thead>
          <tr>
            <th>Name</th>
            <th>Type</th>
            <th>Status</th>
            <th>Size</th>
            <th>Auto</th>
            <th>Actions</th>
          </tr>
        </thead>
        <tbody>
          ${filtered.map((m) => this.renderModelRow(m)).join("")}
        </tbody>
      </table>
    `;

    // Attach row event listeners
    container.querySelectorAll("[data-model-action]").forEach((el) => {
      el.addEventListener("click", () => {
        const action = el.dataset.modelAction;
        const modelId = parseInt(el.dataset.modelId);

        if (action === "download") {
          this.downloadUpdate(modelId);
        } else if (action === "toggle-auto") {
          const model = this.models.find((m) => m.id === modelId);
          this.toggleAutoUpdate(modelId, !model.auto_update);
          model.auto_update = !model.auto_update;
          this.renderModelsList();
        }
      });
    });
  }

  renderModelRow(model) {
    const sizeStr = model.file_size
      ? (model.file_size / (1024 * 1024 * 1024)).toFixed(2) + " GB"
      : "Unknown";

    const statusClass = model.update_available
      ? "update"
      : model.last_checked
        ? "current"
        : "unknown";

    const statusText = model.update_available
      ? "Update Available"
      : model.last_checked
        ? "Up to Date"
        : "Not Checked";

    return `
      <tr>
        <td title="${model.path}">${model.filename}</td>
        <td>${model.model_type}</td>
        <td><span class="modelsync-badge ${statusClass}">${statusText}</span></td>
        <td>${sizeStr}</td>
        <td>
          <input type="checkbox" ${model.auto_update ? "checked" : ""}
            data-model-action="toggle-auto" data-model-id="${model.id}" />
        </td>
        <td>
          ${
            model.update_available
              ? `<button class="modelsync-btn small" data-model-action="download" data-model-id="${model.id}">Download</button>`
              : ""
          }
        </td>
      </tr>
    `;
  }

  renderTargetsList() {
    const container = this.element?.querySelector(".modelsync-targets-list");
    if (!container) return;

    if (this.targets.length === 0) {
      container.innerHTML = '<p class="modelsync-empty">No sync targets configured.</p>';
      return;
    }

    container.innerHTML = `
      <div class="modelsync-targets">
        ${this.targets.map((t) => this.renderTargetCard(t)).join("")}
      </div>
    `;

    // Attach target event listeners
    container.querySelectorAll("[data-target-action]").forEach((el) => {
      el.addEventListener("click", () => {
        const action = el.dataset.targetAction;
        const targetId = parseInt(el.dataset.targetId);

        if (action === "sync") {
          this.syncToTarget(targetId);
        } else if (action === "delete") {
          this.deleteTarget(targetId);
        }
      });
    });
  }

  renderTargetCard(target) {
    return `
      <div class="modelsync-target-card">
        <div class="modelsync-target-header">
          <strong>${target.name}</strong>
          <span class="modelsync-badge ${target.enabled ? "current" : "unknown"}">
            ${target.enabled ? "Enabled" : "Disabled"}
          </span>
        </div>
        <div class="modelsync-target-details">
          <p><strong>Host:</strong> ${target.ssh_host || "Local"}</p>
          <p><strong>Path:</strong> ${target.model_path}</p>
          <p><strong>Max VRAM:</strong> ${target.max_vram_mb ? target.max_vram_mb + " MB" : "No limit"}</p>
          <p><strong>Last Sync:</strong> ${target.last_sync || "Never"}</p>
        </div>
        <div class="modelsync-target-actions">
          <button class="modelsync-btn small primary" data-target-action="sync" data-target-id="${target.id}">Sync Now</button>
          <button class="modelsync-btn small" data-target-action="delete" data-target-id="${target.id}">Delete</button>
        </div>
      </div>
    `;
  }

  renderSettings() {
    if (!this.element) return;

    const autoUpdate = this.element.querySelector("#auto-update-enabled");
    const checkInterval = this.element.querySelector("#check-interval");
    const verifyDownloads = this.element.querySelector("#verify-downloads");
    const backupBefore = this.element.querySelector("#backup-before-update");

    if (autoUpdate) {
      autoUpdate.checked = this.config.auto_update_enabled === "true";
    }
    if (checkInterval) {
      checkInterval.value = this.config.check_interval_hours || "24";
    }
    if (verifyDownloads) {
      verifyDownloads.checked = this.config.verify_downloads !== "false";
    }
    if (backupBefore) {
      backupBefore.checked = this.config.backup_before_update !== "false";
    }
  }

  setStatus(message) {
    const status = this.element?.querySelector(".modelsync-status");
    if (status) {
      status.textContent = message;
    }
  }
}

// Initialize when DOM is ready
const modelSyncPanel = new ModelSyncPanel();

// Load CSS
function loadCSS() {
  if (document.querySelector('link[href*="modelsync.css"]')) return;

  const link = document.createElement("link");
  link.rel = "stylesheet";
  link.href = "/extensions/ComfyUI-ModelSync/modelsync.css";
  document.head.appendChild(link);
}

// Initialize
function init() {
  loadCSS();
  modelSyncPanel.init();
  console.log("[ModelSync] Extension initialized");
}

// Wait for page to be ready
if (document.readyState === "loading") {
  document.addEventListener("DOMContentLoaded", init);
} else {
  init();
}

// Also try after a delay for dynamic ComfyUI loading
setTimeout(init, 2000);

// Export for external access
window.ModelSyncPanel = modelSyncPanel;