0dot77-comfyui-annotations/web/modules/persistence.js

/**
 * Persistence Module
 *
 * Handles saving and loading annotations with ComfyUI workflows.
 * Uses the workflow.extra[EXTENSION_KEY] pattern for storage.
 *
 * ## Storage Strategy
 *
 * Annotations are stored in the workflow's `extra` field:
 * ```json
 * {
 *   "extra": {
 *     "annotations": {
 *       "version": 1,
 *       "strokes": [...]
 *     }
 *   }
 * }
 * ```
 *
 * ## Versioning
 *
 * The version field allows for future migration if the stroke
 * format changes. Current version: 1
 *
 * ## Coordinate System
 *
 * Strokes are stored in canvas (world) coordinates, not screen
 * coordinates. This ensures annotations stay in the correct
 * position regardless of zoom/pan state when saved.
 */

const EXTENSION_KEY = 'annotations';
const CURRENT_VERSION = 1;

/**
 * Serializes strokes for storage in workflow.
 * @param {Array} strokes - Array of stroke objects
 * @returns {Object} Serialized data with version
 */
export function serializeAnnotations(strokes) {
    return {
        version: CURRENT_VERSION,
        strokes: strokes.map(stroke => ({
            id: stroke.id,
            points: stroke.points.map(p => ({ x: p.x, y: p.y })),
            color: stroke.color,
            size: stroke.size,
            opacity: stroke.opacity
        }))
    };
}

/**
 * Deserializes strokes from workflow data.
 * Handles version migration if needed.
 * @param {Object} data - Serialized annotation data
 * @returns {Array} Array of stroke objects
 */
export function deserializeAnnotations(data) {
    if (!data || !data.strokes) {
        console.log('[Persistence] No annotation data to deserialize');
        return [];
    }

    const version = data.version || 1;
    console.log(`[Persistence] Deserializing annotations (version ${version})`);

    // Version migration if needed
    let strokes = data.strokes;
    if (version < CURRENT_VERSION) {
        strokes = migrateStrokes(strokes, version);
    }

    // Validate and reconstruct strokes
    return strokes
        .filter(stroke => stroke && stroke.points && stroke.points.length > 0)
        .map(stroke => ({
            id: stroke.id || generateStrokeId(),
            points: stroke.points.map(p => ({ x: p.x, y: p.y })),
            color: stroke.color || '#FF0000',
            size: stroke.size || 4,
            opacity: stroke.opacity !== undefined ? stroke.opacity : 1.0
        }));
}

/**
 * Migrates strokes from older versions to current format.
 * @param {Array} strokes - Old format strokes
 * @param {number} fromVersion - Source version
 * @returns {Array} Migrated strokes
 */
function migrateStrokes(strokes, fromVersion) {
    console.log(`[Persistence] Migrating strokes from v${fromVersion} to v${CURRENT_VERSION}`);

    // Currently only version 1 exists, so no migration needed
    // Future migrations would go here:
    // if (fromVersion === 1) { ... migrate to v2 ... }

    return strokes;
}

/**
 * Generates a unique stroke ID.
 * @returns {string}
 */
function generateStrokeId() {
    return `stroke_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
}

/**
 * Gets the extension key used for storage.
 * @returns {string}
 */
export function getExtensionKey() {
    return EXTENSION_KEY;
}

/**
 * Creates hooks for ComfyUI serialization/deserialization.
 * Call this once during extension initialization.
 *
 * @param {Object} options
 * @param {() => Array} options.getStrokes - Function to get current strokes
 * @param {(strokes: Array) => void} options.setStrokes - Function to set strokes
 * @param {() => void} options.onLoad - Callback after annotations are loaded
 */
export function setupPersistence(options) {
    const { getStrokes, setStrokes, onLoad } = options;

    // Hook into LGraph serialization
    const LGraph = window.LGraph;
    if (LGraph && LGraph.prototype.serialize) {
        const originalSerialize = LGraph.prototype.serialize;
        LGraph.prototype.serialize = function() {
            const data = originalSerialize.apply(this, arguments);

            // Add annotations to extra field
            if (!data.extra) {
                data.extra = {};
            }

            const strokes = getStrokes();
            if (strokes.length > 0) {
                data.extra[EXTENSION_KEY] = serializeAnnotations(strokes);
                console.log(`[Persistence] Saved ${strokes.length} strokes to workflow`);
            } else {
                // Remove annotations key if no strokes
                delete data.extra[EXTENSION_KEY];
            }

            return data;
        };
        console.log('[Persistence] Hooked LGraph.serialize');
    } else {
        console.warn('[Persistence] LGraph.prototype.serialize not found');
    }

    // Hook into app.loadGraphData
    // This is called when loading a workflow
    const app = window.app;
    if (app) {
        const originalLoadGraphData = app.loadGraphData.bind(app);
        app.loadGraphData = async function(graphData, ...args) {
            // Clear existing annotations first
            setStrokes([]);

            // Call original load
            const result = await originalLoadGraphData(graphData, ...args);

            // Load annotations after a short delay to ensure graph is ready
            setTimeout(() => {
                if (graphData?.extra?.[EXTENSION_KEY]) {
                    const annotationData = graphData.extra[EXTENSION_KEY];
                    const strokes = deserializeAnnotations(annotationData);
                    if (strokes.length > 0) {
                        setStrokes(strokes);
                        console.log(`[Persistence] Loaded ${strokes.length} strokes from workflow`);
                        if (onLoad) {
                            onLoad();
                        }
                    }
                } else {
                    console.log('[Persistence] No annotations in workflow');
                }
            }, 100);

            return result;
        };
        console.log('[Persistence] Hooked app.loadGraphData');
    } else {
        console.warn('[Persistence] app object not found');
    }
}

/**
 * Checks if a workflow has saved annotations.
 * @param {Object} graphData - Workflow data
 * @returns {boolean}
 */
export function hasAnnotations(graphData) {
    return !!(graphData?.extra?.[EXTENSION_KEY]?.strokes?.length > 0);
}

/**
 * Gets annotation count from workflow data.
 * @param {Object} graphData - Workflow data
 * @returns {number}
 */
export function getAnnotationCount(graphData) {
    return graphData?.extra?.[EXTENSION_KEY]?.strokes?.length || 0;
}