smartbotic/nodes/sdcpp/sdcpp-unload.js

/**
 * @node sdcpp-unload
 * @name SD.cpp Unload
 * @category sdcpp
 * @version 1.0.0
 * @description Free the main model slot, the upscaler slot, or both
 * @icon eraser
 */

const configSchema = {
    type: 'object',
    properties: {
        serverUrl: {
            type: 'string', title: 'Server URL',
            description: 'Base address of the sdcpp-restapi server',
            default: 'http://localhost:8077'
        },
        credentialId: {
            type: 'string', title: 'Credential',
            description: 'A basic credential holding the sdcpp-restapi username and password',
            dynamicOptions: { source: 'credentials', filter: { type: ['basic'] } }
        },
        target: {
            type: 'string', title: 'Unload',
            enum: ['model', 'upscaler', 'both'],
            default: 'model',
            description: 'Which slot to free. They are independent - unloading the model leaves the upscaler in place, and the other way round'
        },
        timeout: { type: 'number', title: 'Timeout (ms)', default: 60000 }
    },
    required: []
};

const inputSchema = { type: 'object', properties: { data: { type: 'any' } } };

const outputSchema = {
    type: 'object',
    properties: {
        unloadedModel: { type: 'boolean' },
        unloadedUpscaler: { type: 'boolean' },
        previousModel: { type: 'string', description: 'What was in the main slot before' },
        previousUpscaler: { type: 'string' }
    }
};

function normalizeServer(url) {
    const value = String(url || '').trim();
    if (!value) {
        throw new Error('SD.cpp: a server URL is required, such as http://localhost:8077');
    }
    return value.replace(/\/+$/, '');
}

function readCredential(credentialId) {
    const auth = smartbotic.credentials.get(credentialId);
    if (!auth || auth.success !== true) {
        throw new Error('SD.cpp: could not read the credential: ' +
            ((auth && auth.error) || 'unknown error'));
    }

    const value = auth.headerValue || '';
    if (value.indexOf('Basic ') !== 0) {
        throw new Error('SD.cpp: the credential must be a basic one, holding the sdcpp-restapi ' +
            'username and password');
    }

    const decoded = smartbotic.utils.base64Decode(value.substring(6));
    const separator = decoded.indexOf(':');
    if (separator < 1) {
        throw new Error('SD.cpp: the credential is malformed, expected a username and a password');
    }

    return {
        username: decoded.substring(0, separator),
        password: decoded.substring(separator + 1)
    };
}

function call(options) {
    const response = smartbotic.http.request(options);

    let body = response.data;
    if (typeof body === 'string' && body.length > 0) {
        try {
            body = JSON.parse(body);
        } catch (e) {
            const snippet = body.substring(0, 200).replace(/\s+/g, ' ');
            throw new Error('SD.cpp: ' + options.what + ' returned HTTP ' + response.status +
                ' with a body that is not JSON: ' + snippet);
        }
    }

    if (response.status < 200 || response.status >= 300) {
        const detail = (body && (body.message || body.error)) || ('HTTP ' + response.status);
        throw new Error('SD.cpp: ' + options.what + ' failed: ' + detail);
    }

    return body || {};
}

function login(server, credential, timeout) {
    const session = call({
        method: 'POST',
        url: server + '/auth/login',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
            username: credential.username,
            password: credential.password
        }),
        timeout: timeout,
        what: 'signing in'
    });

    if (!session.token) {
        throw new Error('SD.cpp: the server accepted the login but returned no token');
    }
    return session.token;
}

// /health is unauthenticated, and it is the only way to find out what is
// already loaded without asking for a token first.
function readHealth(server, timeout) {
    return call({
        method: 'GET',
        url: server + '/health',
        timeout: timeout,
        what: 'reading server health'
    });
}

function putIfSet(target, key, value) {
    if (value === undefined || value === null || value === '') {
        return;
    }
    target[key] = value;
}

async function execute(config, input, context) {
    const server = normalizeServer(config.serverUrl);
    const timeout = config.timeout > 0 ? config.timeout : 60000;
    const target = config.target || 'model';

    const health = readHealth(server, Math.min(timeout, 15000));
    const hadModel = health.model_loaded === true;
    const hadUpscaler = health.upscaler_loaded === true;

    const wantModel = target === 'model' || target === 'both';
    const wantUpscaler = target === 'upscaler' || target === 'both';

    // Nothing to free is a success, not an error. A cleanup step at the end of a
    // workflow should not fail because an earlier branch already unloaded.
    if ((!wantModel || !hadModel) && (!wantUpscaler || !hadUpscaler)) {
        smartbotic.log.info('SD.cpp: nothing to unload');
        return {
            unloadedModel: false,
            unloadedUpscaler: false,
            previousModel: '',
            previousUpscaler: ''
        };
    }

    const credential = readCredential(config.credentialId);
    const token = login(server, credential, Math.min(timeout, 30000));
    const headers = { 'Content-Type': 'application/json', 'Authorization': 'Bearer ' + token };

    let unloadedModel = false;
    let unloadedUpscaler = false;

    if (wantModel && hadModel) {
        call({ method: 'POST', url: server + '/models/unload', headers: headers,
               body: '{}', timeout: timeout, what: 'unloading the model' });
        unloadedModel = true;
    }
    if (wantUpscaler && hadUpscaler) {
        call({ method: 'POST', url: server + '/upscaler/unload', headers: headers,
               body: '{}', timeout: timeout, what: 'unloading the upscaler' });
        unloadedUpscaler = true;
    }

    smartbotic.log.info('SD.cpp: unloaded' +
        (unloadedModel ? ' model ' + (health.model_name || '') : '') +
        (unloadedUpscaler ? ' upscaler ' + (health.upscaler_name || '') : ''));

    return {
        unloadedModel: unloadedModel,
        unloadedUpscaler: unloadedUpscaler,
        previousModel: unloadedModel ? (health.model_name || '') : '',
        previousUpscaler: unloadedUpscaler ? (health.upscaler_name || '') : ''
    };
}

module.exports = { configSchema, inputSchema, outputSchema, execute };