79 lines
2.3 KiB
TypeScript
79 lines
2.3 KiB
TypeScript
import * as fs from 'fs/promises';
|
|
import * as path from 'path';
|
|
import axios from 'axios';
|
|
import dotenv from 'dotenv';
|
|
|
|
dotenv.config();
|
|
|
|
interface ImageSize {
|
|
width: number;
|
|
height: number;
|
|
}
|
|
|
|
async function generateImage(
|
|
prompt: string,
|
|
faceImage: string,
|
|
styleImage: string,
|
|
newFileName: string,
|
|
comfyBaseUrl: string,
|
|
comfyOutputDir: string,
|
|
size: ImageSize = { width: 720, height: 1280 }
|
|
): Promise<string> {
|
|
const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
|
|
const COMFY_OUTPUT_DIR = comfyOutputDir;
|
|
let workflow;
|
|
|
|
workflow = JSON.parse(await fs.readFile('src/comfyworkflows/generate_image_style_faceswap.json', 'utf-8'));
|
|
workflow['14']['inputs']['text'] = prompt;
|
|
|
|
// Set image name
|
|
workflow['13']['inputs']['image'] = styleImage;
|
|
|
|
// Set image name
|
|
workflow['19']['inputs']['image'] = faceImage;
|
|
|
|
workflow['3']['inputs']['width'] = size.width;
|
|
workflow['3']['inputs']['height'] = size.height;
|
|
|
|
|
|
const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
|
|
const promptId = response.data.prompt_id;
|
|
|
|
let history;
|
|
do {
|
|
await new Promise(resolve => setTimeout(resolve, 1000));
|
|
const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
|
|
history = historyResponse.data[promptId];
|
|
} while (!history || Object.keys(history.outputs).length === 0);
|
|
|
|
const files = await fs.readdir(COMFY_OUTPUT_DIR!);
|
|
const generatedFiles = files.filter(file => file.startsWith('STYLEDVIDEOMAKER'));
|
|
|
|
const fileStats = await Promise.all(
|
|
generatedFiles.map(async (file) => {
|
|
const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
|
|
return { file, mtime: stat.mtime };
|
|
})
|
|
);
|
|
|
|
fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());
|
|
|
|
const latestFile = fileStats[0].file;
|
|
const newFilePath = path.resolve('./generated', newFileName);
|
|
|
|
await fs.mkdir('./generated', { recursive: true });
|
|
|
|
const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
|
|
try {
|
|
await fs.unlink(newFilePath);
|
|
} catch (err) {
|
|
// ignore if not exists
|
|
}
|
|
|
|
await fs.copyFile(sourcePath, newFilePath);
|
|
|
|
return newFilePath;
|
|
}
|
|
|
|
export { generateImage };
|