lib/diffusion.js

/*
//
// Dibuat oleh Kaze 
// https://github.com/KazeDevID
// https://whatsapp.com/channel/0029VaFNnRTHLHQR6G0fC01O
//
*/

import WebSocket from "ws";


function createHash(length) {
  var chars = 'abcdefghijklmnopqrstuvwxyz0123456789';
  var text = '';

  length = length || 6;

  for (var i = 0; i < length; i++) {
    text += chars.charAt(Math.floor(Math.random() * chars.length));
  }

  return text;
}

const generateHash = () => {
  let m = createHash(12);
  return {
    session_hash: m,
    fn_index: 2
  }
}

async function stablediff(prompt) {
  return new Promise((resolve, reject) => {
    let timerCounter = setTimeout(async () => {
      reject(new Error('Permintaan Anda telah habis waktu. Silakan coba lagi'));
    }, 129000)
    const ws = new WebSocket('wss://stabilityai-stable-diffusion.hf.space/queue/join');
    const hash = generateHash();
    
    ws.on('open', () => {});
    
    ws.on('message', async (message) => {
      try {
        const msg = JSON.parse(`${message}`);
      
        if (msg.msg === 'send_hash') {
          ws.send(JSON.stringify(hash));
        } else if (msg.msg === 'send_data') {
          const data = {
            data: [prompt,"",10],
            ...hash,
          };
          ws.send(JSON.stringify(data));
        } else if (msg.msg === 'estimation') {
        } else if (msg.msg === 'process_completed') {
          clearTimeout(timerCounter)
          try {
            const resultsArr = msg.output.data[0];
            const imgArr = [];
            const resultsStr = [resultsArr].toString();
            for (let i = 0; i < resultsArr.length; i++) {
              const imgData = resultsArr[i].split(',')[1];
              const buffer = Buffer.from(imgData, 'base64');
              imgArr.push(buffer);
            }
            resolve(imgArr);
          } catch (error) {
            console.error(error);
          } finally {
            ws.removeAllListeners();
            ws.close();
          }
        } else if (msg.msg === 'queue_full') {
          try {
            reject(new Error('Antrean penuh.'));
          }
          catch (error) {
            console.error(error);
            reject(new Error('Terjadi kesalahan saat menghasilkan gambar'));
          }
        }
      } catch (error) {
        reject(error);
      };
    });
    ws.on('error', async (error) => {
      console.error(error)
    }); 
  });    
}

export { stablediff }