save changes

src/product/clean_background.ts

@@ -0,0 +1,75 @@
+import { convertImage } from '../lib/image-converter';
+import * as fs from 'fs-extra';
+import * as path from 'path';
+import dotenv from 'dotenv';
+
+dotenv.config();
+
+const inputDir = path.join(__dirname, '../../input');
+const outputDir = path.join(__dirname, '../../generated/clearned');
+
+const comfyUrl = process.env.SERVER1_COMFY_BASE_URL;
+const comfyOutputDir = process.env.SERVER1_COMFY_OUTPUT_DIR;
+
+if (!comfyUrl || !comfyOutputDir) {
+    console.error("ComfyUI URL or Output Directory is not set in environment variables.");
+    process.exit(1);
+}
+
+const comfyInputDir = comfyOutputDir.replace("output", "input");
+
+async function processImages() {
+    await fs.ensureDir(outputDir);
+
+    const files = await fs.readdir(inputDir);
+    let index = 1;
+
+    for (const file of files) {
+        const sourceFilePath = path.join(inputDir, file);
+        const stats = await fs.stat(sourceFilePath);
+
+        if (stats.isFile()) {
+            console.log(`Processing ${file}...`);
+
+            const comfyInputPath = path.join(comfyInputDir, file);
+
+            try {
+                // 1. Copy file to ComfyUI input directory
+                await fs.copy(sourceFilePath, comfyInputPath);
+                console.log(`Copied ${file} to ComfyUI input.`);
+
+                const prompt = "请从图1中提取主要主体，把背景设置为浅灰色，并让主体正面朝向，制作成产品照片。";
+
+                // 2. Call convertImage with correct parameters
+                const generatedFilePath = await convertImage(prompt, file, comfyUrl!, comfyOutputDir!);
+
+                if (generatedFilePath && await fs.pathExists(generatedFilePath)) {
+                    const outputFilename = `clearned_${index}.png`;
+                    const finalOutputPath = path.join(outputDir, outputFilename);
+
+                    // 3. Move the generated file to the final destination
+                    await fs.move(generatedFilePath, finalOutputPath, { overwrite: true });
+                    console.log(`Saved cleaned image to ${finalOutputPath}`);
+                    index++;
+
+                    // 4. Delete the original file from the script's input directory
+                    await fs.unlink(sourceFilePath);
+                    console.log(`Deleted original file: ${file}`);
+                }
+
+                // 5. Clean up the file from ComfyUI input directory
+                await fs.unlink(comfyInputPath);
+                console.log(`Cleaned up ${file} from ComfyUI input.`);
+
+            } catch (error) {
+                console.error(`Failed to process ${file}:`, error);
+                // If something fails, make sure to clean up the copied file if it exists
+                if (await fs.pathExists(comfyInputPath)) {
+                    await fs.unlink(comfyInputPath);
+                }
+            }
+        }
+    }
+}
+
+processImages().catch(console.error);