Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x | import axios from 'axios';
import fs from 'fs';
import parse from 'json-templates';
import sanitize from 'sanitize-filename';
import * as stream from 'stream';
import { promisify } from 'util';
import { v4 as uuidv4 } from 'uuid';
import WebSocket from 'ws';
import dataSource from "./data-source";
import Prompt from './entity/prompt';
const comfyApi = 'https://ai.qtk.io';
const comfyWs = 'wss://ai.qtk.io';
async function generate (prompt: Prompt): Promise<string[]> {
const template = parse(JSON.parse(fs.readFileSync(`workflows/${prompt.workflow}.json`).toString()));
const seed = Math.floor(Math.random() * 4294967294);
const workflow = template({ seed: seed, prompt: prompt.prompt, negativePrompt: prompt.negativePrompt, batchSize: prompt.batchSize });
const clientId = uuidv4();
const timeStart = Date.now();
console.log(`Generating with workflow ${prompt.workflow} for prompt "${prompt.prompt}" and negative prompt "${prompt.negativePrompt}"`);
let promptId: string;
await axios.post(`${comfyApi}/prompt`, {
prompt: workflow,
client_id: clientId,
})
.then((res) => {
if (!res.data.prompt_id) return;
promptId = res.data.prompt_id;
console.log(`Got prompt ID: ${promptId}`);
});
const ws = new WebSocket(`${comfyWs}/ws?clientId=${clientId}`);
const confirmGeneration = () => new Promise((resolve, reject) => {
ws.on('error', reject);
ws.on('message', (data) => {
const message = JSON.parse(data.toString());
if (message.type == 'executing' && message.data.prompt_id == promptId && message.data.node == null) {
resolve(message);
}
});
});
await confirmGeneration();
prompt.runtime = (Date.now() - timeStart) / 1000.0;
console.log(`Completed generation in ${prompt.runtime}s`);
const urls = [];
await axios.get(`${comfyApi}/history/${promptId}`).then((res) => {
if (!res.data[promptId].outputs) return;
const nodeKey = Object.keys(res.data[promptId].outputs)[0];
const images = res.data[promptId].outputs[nodeKey].images;
for (let i = 0; i < images.length; i += 1) {
urls.push(`${comfyApi}/view?filename=${images[i].filename}&subfolder=${images[i].subfolder}&type=${images[i].type}`);
}
});
const files = [];
const finishedDownload = promisify(stream.finished);
for (let i = 0; i < urls.length; i += 1) {
const res = await axios({
url: urls[i],
method: 'GET',
responseType: 'stream',
});
const regExpFilename = /filename="(?<filename>.*)"/;
const origFilename = regExpFilename.exec(res.headers['content-disposition'])?.groups?.filename ?? `${i}.png`;
const strippedPrompt = prompt.prompt.replace('.', '').substring(0, 127);
const filename = sanitize(`${Date.now()}_${strippedPrompt}_${origFilename}`);
const path = `images/${filename}`;
const writer = fs.createWriteStream(path);
res.data.pipe(writer);
await finishedDownload(writer);
files.push(path);
console.log(`Wrote generated image: ${filename}`);
prompt.imageFilename = filename;
}
dataSource.manager.save(prompt);
return files;
}
export { generate }
|