save changes
This commit is contained in:
272
src/piterest_styletransfer_video.ts
Normal file
272
src/piterest_styletransfer_video.ts
Normal file
@ -0,0 +1,272 @@
|
||||
import { downloadImagesFromPinterestPin } from './lib/downloader';
|
||||
import { callOpenAIWithFile } from './lib/openai';
|
||||
import { generateStyledVideo } from './lib/video-generator-styled';
|
||||
import { logger } from './lib/logger';
|
||||
import * as fs from 'fs/promises';
|
||||
import dotenv from 'dotenv';
|
||||
import path from 'path';
|
||||
import puppeteer from 'puppeteer';
|
||||
import { VideoModel } from './lib/db/video';
|
||||
|
||||
dotenv.config();
|
||||
|
||||
const servers = [
|
||||
{
|
||||
baseUrl: process.env.SERVER1_COMFY_BASE_URL,
|
||||
outputDir: process.env.SERVER1_COMFY_OUTPUT_DIR,
|
||||
},
|
||||
/*
|
||||
{
|
||||
baseUrl: process.env.SERVER2_COMFY_BASE_URL,
|
||||
outputDir: process.env.SERVER2_COMFY_OUTPUT_DIR,
|
||||
},
|
||||
*/
|
||||
].filter((s): s is { baseUrl: string; outputDir: string } => !!s.baseUrl && !!s.outputDir);
|
||||
|
||||
interface GenerationTask {
|
||||
pinUrl: string;
|
||||
imagePrompt: string;
|
||||
videoPrompt: string;
|
||||
imageFileName: string;
|
||||
renamedImagePath: string;
|
||||
genre: string;
|
||||
subGenre: string;
|
||||
scene: string;
|
||||
action: string;
|
||||
camera: string;
|
||||
}
|
||||
|
||||
async function getPromptsForImage(imagePath: string, pinUrl: string, genre: string, subGenre: string): Promise<GenerationTask | null> {
|
||||
const pinId = pinUrl.split('/').filter(Boolean).pop() || `pin_${Date.now()}`;
|
||||
const timestamp = new Date().getTime();
|
||||
const imageFileName = `${pinId}_${timestamp}.png`;
|
||||
const renamedImagePath = path.join(path.dirname(imagePath), imageFileName);
|
||||
|
||||
try {
|
||||
await fs.rename(imagePath, renamedImagePath);
|
||||
logger.debug(`Renamed ${imagePath} to ${renamedImagePath}`);
|
||||
|
||||
const promptResponse = await callOpenAIWithFile(renamedImagePath,
|
||||
`Analyze the provided image and generate the following:
|
||||
1. 'scene': A description of the image's environment.
|
||||
2. 'action': A description of the main action occurring in the image.
|
||||
3. 'camera': A description of the camera shot (e.g., 'close-up', 'wide-angle').
|
||||
4. 'image_prompt': A highly detailed, creative, and artistic prompt for an image generation model, inspired by the original. This prompt should be around 200 words.
|
||||
5. 'video_prompt': A prompt for an 8-second video, describing a creative and subtle movement of the main object. The camera should be static, but slight panning is acceptable.
|
||||
|
||||
Output should be in this JSON format:
|
||||
---
|
||||
{
|
||||
"scene": "{result comes here}",
|
||||
"action": "{result comes here}",
|
||||
"camera": "{result comes here}",
|
||||
"image_prompt": "Ultra detailed illustration, {result comes here}",
|
||||
"video_prompt": "{result comes here}"
|
||||
}
|
||||
---
|
||||
`);
|
||||
const { scene, action, camera, image_prompt: imagePrompt, video_prompt: videoPrompt } = promptResponse;
|
||||
logger.info(`Image prompt for ${renamedImagePath}:`, imagePrompt);
|
||||
logger.info(`Video prompt for ${renamedImagePath}:`, videoPrompt);
|
||||
|
||||
return { pinUrl, imagePrompt, videoPrompt, imageFileName, renamedImagePath, genre, subGenre, scene, action, camera };
|
||||
} catch (error) {
|
||||
logger.error(`Failed to get prompts for ${renamedImagePath}:`, error);
|
||||
try {
|
||||
await fs.unlink(renamedImagePath);
|
||||
} catch (cleanupError) {
|
||||
// ignore
|
||||
}
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async function generateImageAndVideo(task: GenerationTask, server: { baseUrl: string; outputDir: string; }): Promise<{ imagePath: string; videoPath: string; } | null> {
|
||||
const { imagePrompt, videoPrompt, imageFileName, renamedImagePath } = task;
|
||||
const { baseUrl, outputDir } = server;
|
||||
const inputDir = outputDir.replace("output", "input");
|
||||
|
||||
try {
|
||||
const destPath = path.join(inputDir, imageFileName);
|
||||
await fs.copyFile(renamedImagePath, destPath);
|
||||
logger.info(`Copied ${renamedImagePath} to ${destPath}`);
|
||||
|
||||
const videoFileName = imageFileName.replace('.png', '.mp4');
|
||||
await generateStyledVideo(
|
||||
imagePrompt,
|
||||
videoPrompt,
|
||||
imageFileName,
|
||||
videoFileName,
|
||||
baseUrl,
|
||||
outputDir,
|
||||
{ width: 320, height: 640 }
|
||||
);
|
||||
const videoPath = path.join(outputDir, videoFileName);
|
||||
return { imagePath: destPath, videoPath };
|
||||
} catch (error) {
|
||||
logger.error(`Failed to generate styled video for ${imageFileName} on server ${baseUrl}:`, error);
|
||||
return null;
|
||||
} finally {
|
||||
try {
|
||||
await fs.unlink(renamedImagePath);
|
||||
logger.debug(`Deleted renamed source image: ${renamedImagePath}`);
|
||||
} catch (error) {
|
||||
logger.error(`Failed to delete renamed source image ${renamedImagePath}:`, error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function worker(id: number, server: { baseUrl: string; outputDir: string; }, taskQueue: GenerationTask[]) {
|
||||
logger.info(`Worker ${id} started for server ${server.baseUrl}`);
|
||||
while (taskQueue.length > 0) {
|
||||
const task = taskQueue.shift();
|
||||
if (task) {
|
||||
logger.info(`Worker ${id} processing task: ${task.imageFileName}`);
|
||||
const result = await generateImageAndVideo(task, server);
|
||||
if (result) {
|
||||
try {
|
||||
const videoData = {
|
||||
genre: task.genre,
|
||||
sub_genre: task.subGenre,
|
||||
scene: task.scene,
|
||||
action: task.action,
|
||||
camera: task.camera,
|
||||
image_prompt: task.imagePrompt,
|
||||
video_prompt: task.videoPrompt,
|
||||
image_path: result.imagePath,
|
||||
video_path: result.videoPath,
|
||||
};
|
||||
const videoId = await VideoModel.create(videoData);
|
||||
logger.info(`Successfully saved video record to database with ID: ${videoId}`);
|
||||
|
||||
const newImageName = `${videoId}_${task.genre}_${task.subGenre}${path.extname(result.imagePath)}`;
|
||||
const newVideoName = `${videoId}_${task.genre}_${task.subGenre}${path.extname(result.videoPath)}`;
|
||||
|
||||
const newImagePath = path.join(path.dirname(result.imagePath), newImageName);
|
||||
const newVideoPath = path.join(path.dirname(result.videoPath), newVideoName);
|
||||
|
||||
await fs.rename(result.imagePath, newImagePath);
|
||||
await fs.rename(result.videoPath, newVideoPath);
|
||||
|
||||
await VideoModel.update(videoId, {
|
||||
image_path: newImagePath,
|
||||
video_path: newVideoPath,
|
||||
});
|
||||
logger.info(`Renamed files and updated database record for video ID: ${videoId}`);
|
||||
} catch (error) {
|
||||
logger.error('Failed to save video record to database or rename files:', error);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
logger.info(`Worker ${id} finished.`);
|
||||
}
|
||||
|
||||
async function getPinUrlFromPinterest(keyword: string): Promise<string | null> {
|
||||
const browser = await puppeteer.launch({ headless: true });
|
||||
const page = await browser.newPage();
|
||||
await page.setUserAgent('Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36');
|
||||
await page.setViewport({ width: 1920, height: 1080 });
|
||||
try {
|
||||
const searchUrl = `https://www.pinterest.com/search/pins/?q=${encodeURIComponent(keyword)}`;
|
||||
await page.goto(searchUrl, { waitUntil: 'networkidle2' });
|
||||
|
||||
const scrollCount = Math.floor(Math.random() * 5) + 1;
|
||||
logger.info(`Scrolling ${scrollCount} times...`);
|
||||
for (let i = 0; i < scrollCount; i++) {
|
||||
await page.evaluate('window.scrollTo(0, document.body.scrollHeight)');
|
||||
await new Promise(resolve => setTimeout(resolve, Math.random() * 1000 + 1000));
|
||||
}
|
||||
|
||||
const pinLinks = await page.$$eval('a', (anchors) =>
|
||||
anchors.map((a) => a.href).filter((href) => href.includes('/pin/'))
|
||||
);
|
||||
|
||||
if (pinLinks.length > 0) {
|
||||
return pinLinks[Math.floor(Math.random() * pinLinks.length)];
|
||||
}
|
||||
return null;
|
||||
} catch (error) {
|
||||
logger.error('Error while getting pin URL from Pinterest:', error);
|
||||
return null;
|
||||
} finally {
|
||||
await browser.close();
|
||||
}
|
||||
}
|
||||
|
||||
(async () => {
|
||||
const keywords = [
|
||||
{ genre: "fantasy", subGenre: "ethereal" },
|
||||
{ genre: "fantasy", subGenre: "academia" },
|
||||
{ genre: "fantasy", subGenre: "darkacademia" },
|
||||
{ genre: "fantasy", subGenre: "illumication" },
|
||||
{ genre: "fantasy", subGenre: "aesthetic" },
|
||||
{ genre: "abstract", subGenre: "particle" },
|
||||
{ genre: "abstract", subGenre: "space" },
|
||||
{ genre: "abstract", subGenre: "science" },
|
||||
{ genre: "abstract", subGenre: "sphere" },
|
||||
{ genre: "abstract", subGenre: "cubic" },
|
||||
];
|
||||
|
||||
if (servers.length === 0) {
|
||||
logger.error("No servers configured. Please check your .env file.");
|
||||
return;
|
||||
}
|
||||
|
||||
while (true) {
|
||||
|
||||
for (const genreSubGenre of keywords) {
|
||||
|
||||
const { genre, subGenre } = genreSubGenre;
|
||||
|
||||
const keyword = `${genre} ${subGenre}`;
|
||||
logger.info(`Searching for a pin with keyword: ${keyword}`);
|
||||
const pin = await getPinUrlFromPinterest(keyword);
|
||||
|
||||
if (!pin) {
|
||||
logger.warn(`Could not find a pin for keyword: ${keyword}. Skipping.`);
|
||||
continue;
|
||||
}
|
||||
|
||||
const generationTasks: GenerationTask[] = [];
|
||||
|
||||
logger.info(`--- Starting processing for pin: ${pin} ---`);
|
||||
const downloadedImagePaths = await downloadImagesFromPinterestPin(pin);
|
||||
if (downloadedImagePaths.length === 0) {
|
||||
logger.warn(`No images were downloaded for pin ${pin}. Skipping.`);
|
||||
continue;
|
||||
}
|
||||
|
||||
const selectedImages = downloadedImagePaths.sort(() => 0.5 - Math.random()).slice(0, 1);
|
||||
logger.info(`--- Randomly selected ${selectedImages.length} images for processing ---`);
|
||||
|
||||
for (const imagePath of selectedImages) {
|
||||
const pinId = path.basename(imagePath, path.extname(imagePath)).split('_related_')[0];
|
||||
const pinUrl = `https://www.pinterest.com/pin/${pinId}/`;
|
||||
const task = await getPromptsForImage(imagePath, pinUrl, genre, subGenre);
|
||||
if (task) {
|
||||
generationTasks.push(task);
|
||||
}
|
||||
}
|
||||
|
||||
const unselectedImages = downloadedImagePaths.filter(p => !selectedImages.includes(p));
|
||||
for (const imagePath of unselectedImages) {
|
||||
try {
|
||||
await fs.unlink(imagePath);
|
||||
logger.debug(`Deleted unselected image: ${imagePath}`);
|
||||
} catch (error) {
|
||||
logger.error(`Failed to delete unselected image ${imagePath}:`, error);
|
||||
}
|
||||
}
|
||||
|
||||
if (generationTasks.length > 0) {
|
||||
logger.info(`--- Starting parallel generation of ${generationTasks.length} tasks across ${servers.length} servers ---`);
|
||||
const workers = servers.map((server, index) => worker(index + 1, server, generationTasks));
|
||||
await Promise.all(workers);
|
||||
logger.info("--- Finished parallel generation ---");
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
})();
|
||||
Reference in New Issue
Block a user