save changes

This commit is contained in:
2025-10-02 07:34:10 +02:00
parent 4bce4388f9
commit ad8ca7b997
5 changed files with 965 additions and 50 deletions

View File

@ -5,45 +5,79 @@ import * as dotenv from 'dotenv';
dotenv.config();
const clothesDir = 'C:\\Users\\ken\\Desktop\\VTON\\clothes';
const modelPath = 'C:\\Users\\ken\\Desktop\\VTON\\models\\Jessica_body.png';
const posesDir = 'C:\\Users\\ken\\Desktop\\VTON\\poses';
const modelsBodyDir = 'D:\\CatsEye\\long videos\\vton-demo\\VTON\\models_body';
const clothesDir = 'D:\\CatsEye\\long videos\\vton-demo\\VTON\\clothes';
const posesDir = 'D:\\CatsEye\\long videos\\vton-demo\\VTON\\poses';
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 files = fs.readdirSync(directory);
const vtonFiles = files.filter(file => file.startsWith('vton_') && file.endsWith('.png'));
if (vtonFiles.length === 0) {
return 0;
}
const indices = vtonFiles.map(file => {
const match = file.match(/vton_(\d+)\.png/);
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.");
}
const clothesFiles = fs.readdirSync(clothesDir).filter(file => /\.(jpg|png|jpeg)$/i.test(file));
const poseFiles = fs.readdirSync(posesDir).filter(file => /\.(jpg|png|jpeg)$/i.test(file));
let index = getNextIndex(outputDir);
if (!fs.existsSync(outputDir)) {
fs.mkdirSync(outputDir);
}
const comfyInputDir = comfyOutputDir.replace("output", "input");
for (let i = 0; i < clothesFiles.length; i++) {
const clothFile = clothesFiles[i];
const clothPath = path.join(clothesDir, clothFile);
while (true) { // Infinite loop
try {
const personFilePath = getRandomFile(modelsBodyDir);
const clothFilePath = getRandomFile(clothesDir);
const poseFilePath = getRandomFile(posesDir);
const randomPoseFile = poseFiles[Math.floor(Math.random() * poseFiles.length)];
const posePath = path.join(posesDir, randomPoseFile);
const personFileName = path.basename(personFilePath);
const clothFileName = path.basename(clothFilePath);
const poseFileName = path.basename(poseFilePath);
console.log(`Processing cloth: ${clothFile} with pose: ${randomPoseFile}`);
fs.copyFileSync(personFilePath, path.join(comfyInputDir, personFileName));
fs.copyFileSync(clothFilePath, path.join(comfyInputDir, clothFileName));
fs.copyFileSync(poseFilePath, path.join(comfyInputDir, poseFileName));
const files = [modelPath, clothPath, posePath];
const prompt = "change clothes of image1 with image2";
const outputFilename = `model_${i}.png`;
console.log(`Processing person: ${personFileName}, cloth: ${clothFileName}, pose: ${poseFileName}`);
const generatedImagePath = await convertImageVton(files, outputFilename, comfyBaseUrl, comfyOutputDir, { width: 720, height: 1280 });
const outputFilename = `vton_${index}.png`;
if (generatedImagePath) {
console.log(`Generated image saved to ${generatedImagePath}`);
} else {
console.error(`Failed to generate image for ${clothFile}`);
const generatedImagePath = await convertImageVton(personFileName, clothFileName, poseFileName, outputFilename, comfyBaseUrl, comfyOutputDir, { width: 720, height: 1280 });
if (generatedImagePath) {
console.log(`Generated image saved to ${generatedImagePath}`);
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));
}
}
}