RandomVideoMaker/src/tools/vton_generator.ts

import * as fs from 'fs';
import * as path from 'path';
import { convertImageVton, convertImage } from '../lib/image-converter';
import * as dotenv from 'dotenv';
import sharp from 'sharp';

dotenv.config();

const clothesDir = 'D:\\projects\\random_video_maker\\input';
const outputDir = 'generated';

const comfyBaseUrl = process.env.SERVER1_COMFY_BASE_URL;
const comfyOutputDir = process.env.SERVER1_COMFY_OUTPUT_DIR;

function getNextIndex(directory: string): number {
    if (!fs.existsSync(directory)) {
        fs.mkdirSync(directory, { recursive: true });
        return 0;
    }
    const dirs = fs.readdirSync(directory, { withFileTypes: true })
        .filter(dirent => dirent.isDirectory())
        .map(dirent => dirent.name);

    const vtonDirs = dirs.filter(dir => dir.startsWith('vton_'));
    if (vtonDirs.length === 0) {
        return 0;
    }
    const indices = vtonDirs.map(dir => {
        const match = dir.match(/vton_(\d+)/);
        return match ? parseInt(match[1], 10) : -1;
    });
    return Math.max(...indices) + 1;
}

function getRandomFile(directory: string): string {
    const files = fs.readdirSync(directory).filter(file => /\.(jpg|png|jpeg)$/i.test(file));
    if (files.length === 0) {
        throw new Error(`No image files found in directory: ${directory}`);
    }
    const randomFile = files[Math.floor(Math.random() * files.length)];
    return path.join(directory, randomFile);
}

async function generateVtonImages() {
    if (!comfyBaseUrl || !comfyOutputDir) {
        throw new Error("ComfyUI URL or Output Directory is not set in environment variables.");
    }

    let index = getNextIndex(outputDir);

    const comfyInputDir = comfyOutputDir.replace("output", "input");

    while (true) { // Infinite loop
        const iterationDir = path.join(outputDir, `vton_${index}`);
        fs.mkdirSync(iterationDir, { recursive: true });

        try {
            const personOrigPath = getRandomFile(clothesDir);
            const clothOrigPath = getRandomFile(clothesDir);

            fs.copyFileSync(personOrigPath, path.join(iterationDir, '1-personOrig.png'));
            fs.copyFileSync(clothOrigPath, path.join(iterationDir, '3-clothOrig.png'));

            const personOrigFileName = path.basename(personOrigPath);
            const clothOrigFileName = path.basename(clothOrigPath);

            fs.copyFileSync(personOrigPath, path.join(comfyInputDir, personOrigFileName));
            fs.copyFileSync(clothOrigPath, path.join(comfyInputDir, clothOrigFileName));

            console.log(`Processing person: ${personOrigPath}, cloth: ${clothOrigPath}`);

            const cleanePersonImagePath = await convertImage("请把姿势改成站立的，转换成全身照片。去掉衣服，只保留白色运动文胸和白色短裤。双脚保持赤脚。背景为浅灰色。", personOrigFileName, comfyBaseUrl, comfyOutputDir, { width: 720, height: 1280 });
            fs.copyFileSync(cleanePersonImagePath, path.join(iterationDir, '2-personCleaned.png'));
            const cleanedPersonFileName = path.basename(cleanePersonImagePath);
            fs.copyFileSync(cleanePersonImagePath, path.join(comfyInputDir, cleanedPersonFileName));

            const cleanedClothImagePath = await convertImage("请将图1中的上衣、下装和配饰分别提取出来，放到同一个浅灰色的背景上。", clothOrigFileName, comfyBaseUrl, comfyOutputDir, { width: 720, height: 1280 });
            fs.copyFileSync(cleanedClothImagePath, path.join(iterationDir, '4-clothCleaned.png'));
            const cleanedClothFileName = path.basename(cleanedClothImagePath);
            fs.copyFileSync(cleanedClothImagePath, path.join(comfyInputDir, cleanedClothFileName));

            const outputFilename = `vton_final_${index}.png`;
            const generatedImagePath = await convertImageVton(cleanedPersonFileName, cleanedClothFileName, outputFilename, comfyBaseUrl, comfyOutputDir, { width: 720, height: 1280 });

            if (generatedImagePath) {
                fs.copyFileSync(generatedImagePath, path.join(iterationDir, '5-finalResult.png'));
                console.log(`Generated image saved to ${generatedImagePath}`);

                // --- Create composite image ---
                const imagePaths = [
                    path.join(iterationDir, '1-personOrig.png'),
                    path.join(iterationDir, '3-clothOrig.png'),
                    path.join(iterationDir, '2-personCleaned.png'),
                    path.join(iterationDir, '4-clothCleaned.png'),
                    path.join(iterationDir, '5-finalResult.png')
                ];

                const resizedImages = [];
                let totalWidth = 10; // Initial left margin
                const resizedHeight = 720;

                for (const imagePath of imagePaths) {
                    const image = sharp(imagePath);
                    const metadata = await image.metadata();
                    if (!metadata.width || !metadata.height) {
                        throw new Error(`Could not get metadata for image ${imagePath}`);
                    }
                    const resizedWidth = Math.round((metadata.width / metadata.height) * resizedHeight);

                    const resizedImageBuffer = await image.resize(resizedWidth, resizedHeight).toBuffer();
                    resizedImages.push({
                        buffer: resizedImageBuffer,
                        width: resizedWidth
                    });
                    totalWidth += resizedWidth + 10; // Add image width and right margin
                }

                const compositeOps = [];
                let currentLeft = 10; // Start with left margin

                for (const img of resizedImages) {
                    compositeOps.push({
                        input: img.buffer,
                        top: 10, // 10px top margin
                        left: currentLeft
                    });
                    currentLeft += img.width + 10; // Move to the next position
                }

                await sharp({
                    create: {
                        width: totalWidth,
                        height: 740,
                        channels: 4,
                        background: { r: 255, g: 255, b: 255, alpha: 1 }
                    }
                })
                    .composite(compositeOps)
                    .toFile(path.join(iterationDir, 'process.png'));

                console.log(`Generated composite image process.png in ${iterationDir}`);
                // --- End of composite image creation ---

                index++;
            } else {
                console.error(`Failed to generate image for index ${index}`);
            }
        } catch (error) {
            console.error("An error occurred during image generation:", error);
            // Optional: wait for a bit before retrying to avoid spamming errors
            await new Promise(resolve => setTimeout(resolve, 5000));
        }
    }
}

generateVtonImages().catch(console.error);