save current state

package-lock.json

@@ -12,6 +12,7 @@
                 "@types/axios": "^0.14.4",
                 "axios": "^1.11.0",
                 "dotenv": "^17.2.1",
+                "open": "^10.2.0",
                 "puppeteer": "^24.16.2"
             },
             "devDependencies": {
@@ -241,6 +242,20 @@
                 "node": "*"
             }
         },
+        "node_modules/bundle-name": {
+            "version": "4.1.0",
+            "resolved": "https://registry.npmjs.org/bundle-name/-/bundle-name-4.1.0.tgz",
+            "integrity": "sha512-tjwM5exMg6BGRI+kNmTntNsvdZS1X8BFYS6tnJ2hdH0kVxM6/eVZ2xy+FqStSWvYmtfFMDLIxurorHwDKfDz5Q==",
+            "dependencies": {
+                "run-applescript": "^7.0.0"
+            },
+            "engines": {
+                "node": ">=18"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
         "node_modules/call-bind-apply-helpers": {
             "version": "1.0.2",
             "resolved": "https://registry.npmjs.org/call-bind-apply-helpers/-/call-bind-apply-helpers-1.0.2.tgz",
@@ -362,6 +377,43 @@
                 }
             }
         },
+        "node_modules/default-browser": {
+            "version": "5.2.1",
+            "resolved": "https://registry.npmjs.org/default-browser/-/default-browser-5.2.1.tgz",
+            "integrity": "sha512-WY/3TUME0x3KPYdRRxEJJvXRHV4PyPoUsxtZa78lwItwRQRHhd2U9xOscaT/YTf8uCXIAjeJOFBVEh/7FtD8Xg==",
+            "dependencies": {
+                "bundle-name": "^4.1.0",
+                "default-browser-id": "^5.0.0"
+            },
+            "engines": {
+                "node": ">=18"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
+        "node_modules/default-browser-id": {
+            "version": "5.0.0",
+            "resolved": "https://registry.npmjs.org/default-browser-id/-/default-browser-id-5.0.0.tgz",
+            "integrity": "sha512-A6p/pu/6fyBcA1TRz/GqWYPViplrftcW2gZC9q79ngNCKAeR/X3gcEdXQHl4KNXV+3wgIJ1CPkJQ3IHM6lcsyA==",
+            "engines": {
+                "node": ">=18"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
+        "node_modules/define-lazy-prop": {
+            "version": "3.0.0",
+            "resolved": "https://registry.npmjs.org/define-lazy-prop/-/define-lazy-prop-3.0.0.tgz",
+            "integrity": "sha512-N+MeXYoqr3pOgn8xfyRPREN7gHakLYjhsHhWGT3fWAiL4IkAt0iDw14QiiEm2bE30c5XX5q0FtAA3CK5f9/BUg==",
+            "engines": {
+                "node": ">=12"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
         "node_modules/degenerator": {
             "version": "5.0.1",
             "resolved": "https://registry.npmjs.org/degenerator/-/degenerator-5.0.1.tgz",
@@ -781,6 +833,20 @@
             "resolved": "https://registry.npmjs.org/is-arrayish/-/is-arrayish-0.2.1.tgz",
             "integrity": "sha512-zz06S8t0ozoDXMG+ube26zeCTNXcKIPJZJi8hBrF4idCLms4CG9QtK7qBl1boi5ODzFpjswb5JPmHCbMpjaYzg=="
         },
+        "node_modules/is-docker": {
+            "version": "3.0.0",
+            "resolved": "https://registry.npmjs.org/is-docker/-/is-docker-3.0.0.tgz",
+            "integrity": "sha512-eljcgEDlEns/7AXFosB5K/2nCM4P7FQPkGc/DWLy5rmFEWvZayGrik1d9/QIY5nJ4f9YsVvBkA6kJpHn9rISdQ==",
+            "bin": {
+                "is-docker": "cli.js"
+            },
+            "engines": {
+                "node": "^12.20.0 || ^14.13.1 || >=16.0.0"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
         "node_modules/is-fullwidth-code-point": {
             "version": "3.0.0",
             "resolved": "https://registry.npmjs.org/is-fullwidth-code-point/-/is-fullwidth-code-point-3.0.0.tgz",
@@ -789,6 +855,37 @@
                 "node": ">=8"
             }
         },
+        "node_modules/is-inside-container": {
+            "version": "1.0.0",
+            "resolved": "https://registry.npmjs.org/is-inside-container/-/is-inside-container-1.0.0.tgz",
+            "integrity": "sha512-KIYLCCJghfHZxqjYBE7rEy0OBuTd5xCHS7tHVgvCLkx7StIoaxwNW3hCALgEUjFfeRk+MG/Qxmp/vtETEF3tRA==",
+            "dependencies": {
+                "is-docker": "^3.0.0"
+            },
+            "bin": {
+                "is-inside-container": "cli.js"
+            },
+            "engines": {
+                "node": ">=14.16"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
+        "node_modules/is-wsl": {
+            "version": "3.1.0",
+            "resolved": "https://registry.npmjs.org/is-wsl/-/is-wsl-3.1.0.tgz",
+            "integrity": "sha512-UcVfVfaK4Sc4m7X3dUSoHoozQGBEFeDC+zVo06t98xe8CzHSZZBekNXH+tu0NalHolcJ/QAGqS46Hef7QXBIMw==",
+            "dependencies": {
+                "is-inside-container": "^1.0.0"
+            },
+            "engines": {
+                "node": ">=16"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
         "node_modules/js-tokens": {
             "version": "4.0.0",
             "resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-4.0.0.tgz",
@@ -876,6 +973,23 @@
                 "wrappy": "1"
             }
         },
+        "node_modules/open": {
+            "version": "10.2.0",
+            "resolved": "https://registry.npmjs.org/open/-/open-10.2.0.tgz",
+            "integrity": "sha512-YgBpdJHPyQ2UE5x+hlSXcnejzAvD0b22U2OuAP+8OnlJT+PjWPxtgmGqKKc+RgTM63U9gN0YzrYc71R2WT/hTA==",
+            "dependencies": {
+                "default-browser": "^5.2.1",
+                "define-lazy-prop": "^3.0.0",
+                "is-inside-container": "^1.0.0",
+                "wsl-utils": "^0.1.0"
+            },
+            "engines": {
+                "node": ">=18"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
         "node_modules/pac-proxy-agent": {
             "version": "7.2.0",
             "resolved": "https://registry.npmjs.org/pac-proxy-agent/-/pac-proxy-agent-7.2.0.tgz",
@@ -1036,6 +1150,17 @@
                 "node": ">=4"
             }
         },
+        "node_modules/run-applescript": {
+            "version": "7.0.0",
+            "resolved": "https://registry.npmjs.org/run-applescript/-/run-applescript-7.0.0.tgz",
+            "integrity": "sha512-9by4Ij99JUr/MCFBUkDKLWK3G9HVXmabKz9U5MlIAIuvuzkiOicRYs8XJLxX+xahD+mLiiCYDqF9dKAgtzKP1A==",
+            "engines": {
+                "node": ">=18"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
         "node_modules/semver": {
             "version": "7.7.2",
             "resolved": "https://registry.npmjs.org/semver/-/semver-7.7.2.tgz",
@@ -1228,6 +1353,20 @@
                 }
             }
         },
+        "node_modules/wsl-utils": {
+            "version": "0.1.0",
+            "resolved": "https://registry.npmjs.org/wsl-utils/-/wsl-utils-0.1.0.tgz",
+            "integrity": "sha512-h3Fbisa2nKGPxCpm89Hk33lBLsnaGBvctQopaBSOW/uIs6FTe1ATyAnKFJrzVs9vpGdsTe73WF3V4lIsk4Gacw==",
+            "dependencies": {
+                "is-wsl": "^3.1.0"
+            },
+            "engines": {
+                "node": ">=18"
+            },
+            "funding": {
+                "url": "https://github.com/sponsors/sindresorhus"
+            }
+        },
         "node_modules/y18n": {
             "version": "5.0.8",
             "resolved": "https://registry.npmjs.org/y18n/-/y18n-5.0.8.tgz",

package.json

@@ -19,6 +19,7 @@
         "@types/axios": "^0.14.4",
         "axios": "^1.11.0",
         "dotenv": "^17.2.1",
+        "open": "^10.2.0",
         "puppeteer": "^24.16.2"
     }
 }

src/lib/downloader_civitai.ts

@@ -0,0 +1,94 @@
+import puppeteer from 'puppeteer';
+import axios from 'axios';
+import fs from 'fs';
+import dotenv from 'dotenv';
+
+dotenv.config();
+import path from 'path';
+import { logger } from './logger';
+
+const DOWNLOAD_DIR = path.join(process.cwd(), 'download');
+
+if (!fs.existsSync(DOWNLOAD_DIR)) {
+    fs.mkdirSync(DOWNLOAD_DIR, { recursive: true });
+}
+
+async function downloadImage(url: string, keyword: string): Promise<string | null> {
+    try {
+        const response = await axios.get(url, { responseType: 'arraybuffer' });
+        const fileName = `${keyword.replace(/\s/g, '_')}_${new Date().getTime()}${path.extname(url)}`;
+        const filePath = path.join(DOWNLOAD_DIR, fileName);
+        fs.writeFileSync(filePath, response.data);
+        logger.debug(`Image downloaded successfully: ${filePath}`);
+        return filePath;
+    } catch (error) {
+        logger.error(`Failed to download image from ${url}:`, error);
+        return null;
+    }
+}
+
+export async function downloadCivitaiImages(keyword: string, numberOfImages: number = 10): Promise<string[]> {
+    const browser = await puppeteer.launch({ headless: process.env.OPEN_BROWSER !== 'true' });
+    const page = await browser.newPage();
+    const downloadedImagePaths: string[] = [];
+
+    try {
+        const searchUrl = `https://civitai.com/search/images?query=${encodeURIComponent(keyword)}`;
+        await page.goto(searchUrl, { waitUntil: 'networkidle2' });
+        logger.debug(`Navigated to ${searchUrl}`);
+
+        // Wait for the privacy pop-up and click the accept button
+        try {
+            const acceptButtonSelector = 'button.mantine-Button-filled';
+            await page.waitForSelector(acceptButtonSelector, { timeout: 10000 });
+            await page.click(acceptButtonSelector);
+            logger.debug('Accepted privacy settings.');
+        } catch (error) {
+            logger.warn('Privacy pop-up not found or could not be dismissed, continuing anyway.');
+        }
+
+        // Click the filter button to open the filter options
+        const filterButtonSelector = 'button[aria-label="Filters"]';
+        await page.waitForSelector(filterButtonSelector);
+        await page.click(filterButtonSelector);
+        logger.debug('Clicked the filter button.');
+
+        // Wait for the model filter to be available and type in the model name
+        const modelFilterSelector = 'input[placeholder="Search models"]';
+        await page.waitForSelector(modelFilterSelector);
+        await page.type(modelFilterSelector, 'Flux.1 D');
+        logger.debug('Typed "Flux.1 D" into the model filter.');
+
+        // Wait for the search results to update and click the correct model
+        const modelResultSelector = 'div[aria-label="Flux.1 D"]';
+        await page.waitForSelector(modelResultSelector);
+        await page.click(modelResultSelector);
+        logger.debug('Selected the "Flux.1 D" model.');
+
+        let imageCount = 0;
+        while (imageCount < numberOfImages) {
+            const imageUrls = await page.$$eval('img', imgs => imgs.map(img => img.src));
+            for (const url of imageUrls) {
+                if (imageCount >= numberOfImages) break;
+                if (url.startsWith('https://image.civitai.com')) {
+                    const imagePath = await downloadImage(url, keyword);
+                    if (imagePath) {
+                        downloadedImagePaths.push(imagePath);
+                        imageCount++;
+                    }
+                }
+            }
+
+            if (imageCount < numberOfImages) {
+                await page.evaluate('window.scrollTo(0, document.body.scrollHeight)');
+                await new Promise(resolve => setTimeout(resolve, 2000)); // Wait for new images to load
+            }
+        }
+    } catch (error) {
+        logger.error('An error occurred while downloading images from Civitai:', error);
+    } finally {
+        await browser.close();
+    }
+
+    return downloadedImagePaths;
+}