example.js

/*

Requires Node.js 18.x to be installed.

Create an example.sh file with the following content and execute it from the command line using ./example.sh:

USEAPI_TOKEN="…" NGROK_AUTHTOKEN="…" node ./example.js

*/

import fs from 'fs';
import { promises as fsp } from "fs"
import { Readable } from 'stream';
import { finished } from 'stream/promises';
import ngrok from '@ngrok/ngrok';
import http from 'http';
import { AsyncFunctionQueue } from './query.js';
import data from './prompts.json' assert { type: 'json' };

// console.log('process.env', process.env);

/* 
  Configure your
    ✔️ PixVerse accounts https://useapi.net/docs/api-pika-v1/post-pixverse-account-channel
*/
const { USEAPI_TOKEN } = process.env;

const rootPixVerse = 'https://api.useapi.net/v1/pixverse';

const ngrok_port = 8081;

// To track the number of successfully submitted jobs
let jobsSubmitted = 0;
let jobsCompleted = 0;

// Source image/face, change to any other file name of your choice
const sourceFileName = './source.jpg';

const sleep = (ms = 0) => new Promise(resolve => setTimeout(resolve, ms));

const dateAsString = () => new Date().toISOString();

const saveToFile = async (filePath, data) => {
    await fs.writeFile(filePath, JSON.stringify(data, null, 2), (err) => {
        if (err)
            console.error('Error writing to file:', err);
    });
}

// Extract png from https://cdn.discordapp.com/attachments/server_id/channed_id/filename.png?ex=
const getFileExtensionFromUrl = (url) => {
    const matches = url.match(/\.([^.?]+)(?=\?|$)/);
    return matches ? matches[1] : '';
}

const downloadFile = async (url, fileName) => {
    const { body } = await fetch(url);
    const localPath = `./${fileName}.${getFileExtensionFromUrl(url)}`;
    const stream = fs.createWriteStream(localPath);
    await finished(Readable.fromWeb(body).pipe(stream));
    return localPath;
};

const findByJobid = (jobid, jsonObject = data) => {
    let foundObject = null;

    const traverse = (node) => {
        if (foundObject) return;
        if (node.jobid === jobid) {
            foundObject = node;
            return;
        }
        if (node.buttons && typeof node.buttons === 'object')
            for (let key in node.buttons)
                if (key[0] != '_') // Used for debugging
                    traverse(node.buttons[key]);
    }

    Object.values(jsonObject).forEach(obj => traverse(obj));

    return foundObject;
}

const getJobsCount = (obj = data) => {
    let count = 0;

    function recursiveCount(innerObj) {
        if (typeof innerObj === 'object' && innerObj !== null)
            for (const key in innerObj) {
                if (innerObj.hasOwnProperty(key)) {
                    if (key === 'params') {
                        count++;
                        continue;
                    }
                    if (key === 'buttons' && typeof innerObj[key] === 'object') {
                        count += Object.keys(innerObj[key]).length;
                        continue;
                    }
                    recursiveCount(innerObj[key]);
                }
            }
    }

    recursiveCount(obj);

    return count;
}

// https://github.com/ngrok/ngrok-javascript
const listener = await ngrok.forward({ addr: ngrok_port, authtoken_from_env: true });

console.info('Webhook', listener.url());

// Create server to receive webhook events
http
    .createServer(function (req, res) {
        switch (req.method) {
            case 'GET':
                res.writeHead(200);
                res.write("Hello from ngrok server");
                res.end();
                break;
            case 'POST':
                let body = '';
                req.on('data', chunk => body += chunk.toString());
                req.on('end', async () => {
                    const job = JSON.parse(body);
                    const { verb, status, jobid, content, attachments, button } = job;

                    console.log(`${dateAsString()} ⁝ webhook ${jobid} ${verb} ${status}`, content?.substring(0, 10) + '…');

                    // PixVerse
                    if (['pixverse-create', 'pixverse-animate', 'pixverse-create_single', 'pixverse-meme_face', 'pixverse-button']
                        .includes(verb)) {
                        switch (status) {
                            case 'completed':
                            case 'moderated':
                            case 'failed':
                            case 'cancelled':
                                jobsCompleted++;

                                console.log(`${dateAsString()} ⁝ jobs completed`, jobsCompleted);
                                const node = findByJobid(jobid);

                                node.status = status;
                                node.content = content;
                                node.completed = true;

                                let _submitted = false;

                                if (status == 'completed') {
                                    if (attachments && Array.isArray(attachments) && attachments.length) {
                                        for (let ind = 0; ind < attachments.length; ind++) {
                                            const index = attachments.length > 1 ? `-${ind}` : '';
                                            const fileName = `${jobid.split('-')[0]}-${verb.split('-')[1]}${button ? '-' + button : ''}${index}`;
                                            node.attachmentFileName = await downloadFile(attachments[ind].url, fileName);
                                        }
                                    }

                                    if (node.buttons)
                                        for (let button in node.buttons)
                                            if (button[0] != '_') {
                                                queuePixVerse.enqueue(post_pixverse_button, button, jobid, node.buttons[button]);
                                                _submitted++;
                                            }
                                }

                                saveToFile('./result.json', data);

                                if (!_submitted)
                                    queuePixVerse.enqueue(console.log, `👉 ${job.verb} ${jobid} ${status} ${content?.substring(0, 10) + '…'}`);

                                break;
                        }
                    } else {
                        console.error(`Unknown verb ${verb}`);
                    }

                    res.end('ok');
                });
                break;
        }
    })
    .listen(ngrok_port);

const submit_payload = async (url, payload, owner) => {
    let response;
    let retryCount = 0;

    // When using a very slow connection, fetch may fail to POST a large payload, so we will retry up to 3 times
    while (retryCount < 3)
        try {
            response = await fetch(url, payload);
            break;
        } catch (ex) {
            console.error(`fetch ${url} failed #${retryCount}`, ex);
            if (retryCount > 1)
                throw ex;
            retryCount++;
        }

    const job = await response.json();

    const { jobid, status, error, errorDetails, verb, button, executingJobs } = job;

    console.log(`${dateAsString()} ⁝ #${jobsSubmitted} ${verb ?? url} ${button ?? ''} (${owner.params?.prompt ?? ''}) HTTP ${response.status}`,
        { status, jobid, error, errorDetails });

    switch (response.status) {
        case 429:
            if (executingJobs)
                return 'full';
            else {
                // We received a 429 rate limit response from Discord. 
                // To be safe, let's wait for approximately 10 seconds before trying again.
                await sleep(10 * 1000);
                return 'retry';
            }
        case 504:
            // Query overflow (should never happen unless maxJobs is misconfigured)        
            console.error('504 query overflow detected, sleeping for 3 minutes to allow already running jobs to complete');

            await sleep(3 * 60 * 1000);
        default:
            owner.jobid = jobid;
            owner.status = status;
            owner.error = error;
            owner.errorDetails = errorDetails;
            owner.code = response.status;
            owner.completed = !!error;

            saveToFile('./result.json', data);

            jobsSubmitted++;
    }
}

/*
  Use this function to execute API v1 PixVerse pixverse/button, see https://useapi.net/docs/api-pika-v1/post-pixverse-button
  👉 button param may have unique id after which will be removed.
  Examples:
    post_pixverse_button('U1', <parent_jobid>, { params: {} } );
    post_pixverse_button('V3-456', <parent_jobid>, { params: { prompt: 'color it red' } } );
*/
const post_pixverse_button = async (button, parent_jobid, owner) => {
    // Prompt is the only optional params relevant here
    let { prompt } = owner.params ?? {};

    // button-xx -> button
    button = button.split('-')[0];

    const payload = {
        method: 'POST',
        headers: {
            'Authorization': `Bearer ${USEAPI_TOKEN}`,
            'Content-Type': 'application/json'
        },
        body: JSON.stringify({
            jobid: parent_jobid,
            button,
            prompt,
            replyUrl: listener.url()
        })
    };

    return await submit_payload(`${rootPixVerse}/button`, payload, owner);
}

/*
  Use this function to execute any of API v1 PixVerse jobs, see https://useapi.net/docs/api-pixverse-v1.
  👉 verb param may have unique id after which will be removed.
  Examples:
    post_pixverse('create', { params: { prompt: 'color it red' } } );
    post_pixverse('create_singe-456', { params: { prompt: 'color it red' } } );
*/
const post_pixverse = async (verb, owner) => {
    // Let's default to multipart/form-data to account for cases with image/face param
    const formData = new FormData();

    formData.append('replyUrl', listener.url());

    for (let [key, value] of Object.entries(owner.params)) {
        // Special case 
        if (key === 'jobid') {
            formData.append('replyRef', jobid);
            continue;
        };

        // value has file name to load
        if (['image', 'face'].includes(key))
            value = new Blob([await fsp.readFile(value)]);

        if (key === 'seed_value')
            value = getRandomInt();

        formData.append(key, value);
    }

    const payload = {
        method: 'POST',
        headers: {
            'Authorization': `Bearer ${USEAPI_TOKEN}`
        },
        body: formData
    };

    // verb-xx -> verb
    verb = verb.split('-')[0];

    return await submit_payload(`${rootPixVerse}/${verb}`, payload, owner);
}

// PixVerse seed value valid range 0-2147483647
const getRandomInt = (min = 0, max = 2147483647) => Math.floor(Math.random() * (max - min)) + min;

saveToFile('./result.json', data);

const startTime = performance.now();

const queuePixVerse = new AsyncFunctionQueue();

const jobsToRun = getJobsCount(data);

console.log(`${dateAsString()} ⁝ jobs to execute`, jobsToRun);

Object.keys(data).forEach(key => queuePixVerse.enqueue(post_pixverse, key, data[key]));

while (jobsToRun > jobsCompleted)
    await sleep(10 * 1000);

console.log(`${dateAsString()} ⁝ total elapsed time ${new Date(performance.now() - startTime).toISOString()}`);

process.exit();