/** * Copyright (c) 2023 MERCENARIES.AI PTE. LTD. * All rights reserved. */ //@ts-check const VERSION = '0.6.0.hf.018.a'; const express = require('express'); const http = require('http'); const session = require('express-session'); const axios = require('axios'); const { spawn } = require('child_process'); const { createProxyMiddleware } = require('http-proxy-middleware'); const fs = require('fs'); const path = require('path'); const chokidar = require('chokidar'); const fsExtra = require('fs-extra'); const app = express(); const OMNITOOL_INSTALL_SCRIPT = './omnitool_init.sh'; // './omnitool_start.sh'; const CONTAINER_HOST = '127.0.0.1'; const OMNI_URL = 'http://127.0.0.1:1688'; // URL of the OMNITOOL service const PROXY_PORT_OMNITOOL = 4444; const CONTAINER_PORT_OMNITOOL = 1688; const OMNITOOL_HEALTH_PING = '/api/v1/mercenaries/ping'; const FAVICON = 'https://github.com/omnitool-ai/omnitool/raw/main/packages/omni-ui/omni-web/public/favicon.jpg' const OMNITOOL_SPACE_URL = "omnitool-ai-omnitool-on-hf.hf.space"; const HF_SPACE_URL = "https://huggingface.co/spaces/omnitool-ai/omnitool_on_hf"; const HF_SPACE_DUPLICATE_URL = "https://huggingface.co/spaces/omnitool-ai/omnitool_on_hf?duplicate=true"; const sourceDir = 'omnitool/packages/omni-server/data.local'; const targetDir = '/data';//'omnitool/data'; const DELAY_OMNITOOL_SET_TO_RUNNING = 2000; // 2 seconds const CHECK_OMNI_INTERVAL = 60000; // 1 minute const MAX_LOG_SIZE = 1000; // Set your desired maximum log size // Global variable global.OMNITOOL_RUNNING = false; global.OMNITOOL_READY = false; global.ALREADY_STARTING = false; global.LOCAL_URL = ""; global.PING_URL = ""; global.PROTOCOL = ""; global.PROXY_STARTED = false; global.logs = []; global.error_logs = []; global.OMNITOOL_PROXY = null; global.CONNECTED_TO_MASTER = false; console.log(`************ Omnitool Proxy Server v${VERSION} ************`); // HTML templates and constants const LOG_CONTAINER_HTML = `
`; const COMMON_STYLES = ``; function getScriptsHtml(url) { const scripts_html = ``; console.log(`scripts_html = ${scripts_html}`); return scripts_html; } const INTRO_MESSAGE = `

Omnitool.ai on Hugging Face

Version: ${VERSION}

Welcome to Omnitool.ai, running on Hugging Face!
For more details, including how to install omnitool on your own computer for free,
visit us at omnitool.ai

`; const VIDEO_PLAYER_MESSAGE = '
' async function startOmnitoolServer() { if (global.ALREADY_STARTING) return; global.ALREADY_STARTING = true; console.log('Starting Omnitool Server...'); return new Promise((resolve, reject) => { const omnitoolStartProcess = spawn(OMNITOOL_INSTALL_SCRIPT); omnitoolStartProcess.stdout.on('data', (data) => { if (global.logs.length >= MAX_LOG_SIZE) { global.logs.shift(); // Remove the oldest log entry } global.logs.push(data.toString()); console.log(`[log] ${data}`); if (!global.OMNITOOL_RUNNING) { if (data.toString().includes(`Server has started and is ready to accept connections`)) { console.log('Omnitool server started successfully'); setOmnitoolRunning(true); } } }); omnitoolStartProcess.stderr.on('data', (data) => { if (global.error_logs.length >= MAX_LOG_SIZE) { global.error_logs.shift(); // Remove the oldest log entry } console.error(`[stderr] ${data}`); global.error_logs.push(data.toString()); }); omnitoolStartProcess.on('close', (code) => { const message = `Omnitool server process exited with code: ${code}`; console.log(message); global.logs.push(message); global.ALREADY_STARTING = false; if (code === 0) { //@ts-ignore resolve(); } else { reject(`Omnitool server did not start properly.`); } }); }); } function setOmnitoolRunning(set_running) { if (set_running === true) { if (global.OMNITOOL_READY === false) { global.OMNITOOL_READY = true; console.log('Omnitool server is READY! '); setTimeout(() => { console.log('Omnitool server is RUNNING! '); global.OMNITOOL_RUNNING = true; }, DELAY_OMNITOOL_SET_TO_RUNNING); // Delay by 2 second // start mirroring to /data when the server is running // startMirrorToDataDir(); #DISABLED! } } else { global.OMNITOOL_READY = false; global.OMNITOOL_RUNNING = false; } } // Function to check the status of the external service async function checkOmnitoolStatus() { try { //@ts-ignore const response = await axios.get('http://127.0.0.1:1688/api/v1/mercenaries/ping'); if (response.data && response.data.ping === 'pong' && Object.keys(response.data.payload).length === 0) { setOmnitoolRunning(true); } else { setOmnitoolRunning(false); } } catch (error) { setOmnitoolRunning(false); } } async function handleGetStartOmnitoolServer(req, res) { console.log(`Omnitool Server:ALREADY_STARTING = ${global.ALREADY_STARTING}`); if (global.ALREADY_STARTING) { res.writeHead(200, { 'Content-Type': 'text/html' }); res.end("Omnitool server already starting"); return; } try { await startOmnitoolServer(); res.writeHead(200, { 'Content-Type': 'text/html' }); res.end("Omnitool server started successfully"); } catch (error) { console.error(error); global.ALREADY_STARTING = false; res.writeHead(500, { 'Content-Type': 'text/html' }); res.end(`Error starting Omnitool server: ${error}`); } } async function handleGetOmnitoolLogs(req, res) { res.writeHead(200, { 'Content-Type': 'application/json' }); const reply = { logs: global.logs, error_logs: global.error_logs, ready: global.OMNITOOL_READY }; res.end(JSON.stringify(reply)); return; } async function handleGetBurstIframe(req, res) { req.session.isVisited = true; const reply = { burst_iframe: true }; res.writeHead(200, { 'Content-Type': 'application/json' }); res.end(JSON.stringify(reply)); return; } async function proxyRequest(req, res) { console.log('Proxying request'); if (global.PROXY_STARTED) return; global.PROXY_STARTED = true; // Proxy logic... const options = { hostname: CONTAINER_HOST, port: CONTAINER_PORT_OMNITOOL, path: req.url, method: req.method, headers: req.headers, }; const proxy = http.request(options, (proxyRes) => { res.writeHead(proxyRes.statusCode, proxyRes.headers); proxyRes.pipe(res, { end: true }); }); req.pipe(proxy, { end: true }); return; } async function handleGetRoot(req, res) { setGlobals(req); if (req.query.isVisited) req.session.isVisited = true; console.log(`req.session.isVisited = ${req.session.isVisited}`); console.log(`global.OMNITOOL_RUNNING = ${global.OMNITOOL_RUNNING}`); if (!req.session.isVisited || !global.OMNITOOL_RUNNING) { await checkOmnitoolStatus(); } if (global.CONNECTED_TO_MASTER && !req.session.roayl_override) { let page_message =""; page_message += '
This is the REFERENCE Omnitool Space on Huggingface.
'; page_message += '

'; page_message += '
To access Omnitool, you must first DUPLICATE this SPACE to make it your own.
'; page_message += '
Duplicating the Space will help secure your keys, recipes and outputs.
'; page_message += '
Please set your Space as PRIVATE!
'; page_message += '
DUPLICATING and using your own SPACE is FREE if you use the lowest tier of CPU. No GPU is required.
'; page_message += `
DUPLICATING this space is done by pressing this button: DUPLICATE SPACE
`; page_message += "
However, you can try using the 20 Gig Storage option (~ 5 USD/month) to persist keys, recipes and outputs between server restarts. This is still a BETA feature!
"; page_message += '

'; page_message += '
Some users have reported getting a 404 error when pressing the yellow START button after duplicating the space. If that happens to you follow the following instructions
'; page_message += '
- go to settings and set the space to PUBLIC
'; page_message += '
- confirm that the START button now works
'; page_message += '
- go to settings and set the Space back to PRIVATE
'; page_message += '
- in settings, FACTORY REBUILD your Space
'; page_message += '

'; const SCRIPTS = getScriptsHtml(); const page_html = ` Omnitool ${COMMON_STYLES} ${INTRO_MESSAGE} ${page_message}

${VIDEO_PLAYER_MESSAGE} ${SCRIPTS} `; res.writeHead(200, { 'Content-Type': 'text/html' }); res.end(page_html); return; } if (!req.session.isVisited) { let page_message =""; page_message += `
You are now connected to the space [ ${global.LOCAL_URL} ]
` page_message += `
If this is NOT your SPACE, it is HIGHLY recommended that you duplicate this space and make it private to secure your keys, recipes and outputs.
` page_message += `
DUPLICATING and using your own SPACE is FREE. You can do it by pressing this button: DUPLICATE SPACE
However, consider choosing the 20 Gig Storage option (~ 5 USD/month) to persist keys, recipes and outputs between server restarts.
`; page_message += `

`; page_message += `
To ACCESS OMNITOOL on this SPACE: press the [START] button below.
`; const gotoButtonClass = 'highlight-button'; let buttons_html = ` `; const SCRIPTS = getScriptsHtml(); const page_html = ` Omnitool ${COMMON_STYLES} ${INTRO_MESSAGE} ${page_message}

${buttons_html}

${LOG_CONTAINER_HTML} ${SCRIPTS} `; res.writeHead(200, { 'Content-Type': 'text/html' }); res.end(page_html); return; } if (!global.OMNITOOL_RUNNING) { let page_message =""; page_message += `
It looks like Omnitool is not running on this Space. Let's fix that.
`; console.log('Omnitool server is not running'); const startButtonClass = 'highlight-button-green'; let buttonsHTML = ` `; const SCRIPTS = getScriptsHtml(); const page_html = ` Omnitool ${COMMON_STYLES} ${INTRO_MESSAGE} ${page_message}

${buttonsHTML} ${LOG_CONTAINER_HTML} ${SCRIPTS} `; res.writeHead(200, { 'Content-Type': 'text/html' }); res.end(page_html); return; } proxyRequest(req, res); } function setGlobals(req) { console.log(`Setting global using req.headers['host'] = ${req.headers['host']}`); global.LOCAL_URL = req.headers['host']; if (req.protocol === "https") global.PROTOCOL = 'https'; else global.PROTOCOL = 'http'; console.log(`global.LOCAL_URL = ${global.LOCAL_URL}\nglobal.PROTOCOL = ${global.PROTOCOL}`); if (!global.LOCAL_URL) throw new Error('No host header found in request' + JSON.stringify(req.headers, null, 2)); const hostname = global.LOCAL_URL.split(':')[0]; const newUrl = `${global.PROTOCOL}://${hostname}`;//:${PROXY_PORT_OMNITOOL}`; global.PING_URL = `${newUrl}${OMNITOOL_HEALTH_PING}`; console.log(`global.PING_URL = ${global.PING_URL}`); // New logic to set CONNECTED_TO_MASTER global.CONNECTED_TO_MASTER = false; if (global.LOCAL_URL === OMNITOOL_SPACE_URL) global.CONNECTED_TO_MASTER = true; console.log(`global.CONNECTED_TO_MASTER = ${global.CONNECTED_TO_MASTER}`); } async function startHuggingfaceServer() { // Start server http.createServer(app).listen(PROXY_PORT_OMNITOOL, () => { console.log(`Server running on port ${PROXY_PORT_OMNITOOL}`); }); } function omnitoolProxyMiddleware(req, res, next) { if (global.OMNITOOL_RUNNING && req.session.isVisited) { // Proxy all requests to OMNITOOL when conditions are met return createProxyMiddleware({ target: OMNI_URL, changeOrigin: true, ws: true // if you need WebSocket support })(req, res, next); } else { // Continue with normal processing for other cases next(); } } async function main(app) { if (fs.existsSync(targetDir)) { fsExtra.copySync(targetDir, sourceDir); console.log('Copied files from', targetDir, 'to', sourceDir); } await checkOmnitoolStatus(); // Configure session middleware app.use(session({ secret: 'your-secret-key', resave: false, saveUninitialized: true })); try { await startHuggingfaceServer(); } catch (error) { console.error(`There was an error starting the server: ${error}`); return; } app.use(omnitoolProxyMiddleware); app.get('/favicon.ico', (req, res) => res.status(204)); app.get('/start-omnitool-server', (req, res) => handleGetStartOmnitoolServer(req, res)); app.get('/omnitool-logs', (req, res) => handleGetOmnitoolLogs(req, res)); app.get('/burst-iframe', (req, res) => handleGetBurstIframe(req, res)); app.get('/', (req, res) => handleGetRoot(req, res)); setInterval(async () => { await checkOmnitoolStatus(); }, CHECK_OMNI_INTERVAL); } // Define the proxy middleware outside the function global.OMNITOOL_PROXY = createProxyMiddleware({ target: OMNI_URL, changeOrigin: true, ws: true // if you need WebSocket support }); function omnitoolProxyMiddleware(req, res, next) { if (global.OMNITOOL_RUNNING && req.session.isVisited) { // Use the predefined proxy middleware return global.OMNITOOL_PROXY(req, res, next); } else { // Continue with normal processing for other cases next(); } } main(app); function startMirrorToDataDir() { if (!fs.existsSync(targetDir)) { console.warn(`---------> Target directory ${targetDir} does not exist. Not starting mirroring.`); return; } else { console.warn(`---------> Target directory ${targetDir} exists. Starting mirroring.`); } const watcher = chokidar.watch(sourceDir, { ignored: /[\/\\]\./, persistent: true }); watcher .on('add', function(file_path) { //if (file_path.includes('-journal') || file_path.includes('httpd.db')) return; if (file_path.includes('-journal')) return; const targetFile = file_path.replace(sourceDir, targetDir); //console.warn(`TargetFile = ${targetFile}. File ${file_path} added and copied from ${sourceDir} to ${targetDir}`); try { fsExtra.ensureDirSync(path.dirname(targetFile)); fs.copyFile(file_path, targetFile, (err) => { if (err) { //throw err; const error_message = `ERROR when doing: File ${file_path} added -> copied to ${targetFile}\n`; console.log(error_message); global.error_logs.push(error_message); }}); } catch(error) { console.error(`WATCHER: error: ${error}`); } }) .on('change', function(file_path) { //if (file_path.includes('-journal') || file_path.includes('httpd.db')) return; if (file_path.includes('-journal')) return; const targetFile = file_path.replace(sourceDir, targetDir); //console.warn(`TargetFile = ${targetFile}. File ${file_path} changed and copied from ${sourceDir} to ${targetDir}`); try{ fsExtra.ensureDirSync(path.dirname(targetFile)); fs.copyFile(file_path, targetFile, (err) => { if (err) { //throw err; const error_message = `ERROR when doing: File ${file_path} changed -> copied to ${targetFile}\n`; console.log(error_message); global.error_logs.push(error_message); } }); } catch(error) { console.error(`WATCHER: error: ${error}`); } }) .on('unlink', function(file_path) { //if (file_path.includes('-journal') || file_path.includes('httpd.db')) return; if (file_path.includes('-journal')) return; const targetFile = file_path.replace(sourceDir, targetDir); //console.warn(`TargetFile = ${targetFile}. File ${file_path} removed from ${targetDir}`); try { fsExtra.ensureDirSync(path.dirname(targetFile)); fs.unlink(targetFile, (err) => { if (err) { // throw err; const error_message = `ERROR when doing: File ${file_path} has been removed from ${targetDir}\n`; console.log(error_message); global.error_logs.push(error_message); } }); } catch(error) { console.error(`WATCHER: error: ${error}`); } }); return watcher; }