save current state

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;
+}