save changes
This commit is contained in:
@ -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));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user