|
| 1 | +import * as os from 'node:os'; |
| 2 | +import * as path from 'node:path'; |
| 3 | +import { cli, Strategy } from '../../registry.js'; |
| 4 | +import type { IPage } from '../../types.js'; |
| 5 | +import { saveBase64ToFile } from '../../utils.js'; |
| 6 | +import { GEMINI_DOMAIN, exportGeminiImages, getGeminiVisibleImageUrls, sendGeminiMessage, startNewGeminiChat, waitForGeminiImages } from './utils.js'; |
| 7 | + |
| 8 | +function extFromMime(mime: string): string { |
| 9 | + if (mime.includes('png')) return '.png'; |
| 10 | + if (mime.includes('webp')) return '.webp'; |
| 11 | + if (mime.includes('gif')) return '.gif'; |
| 12 | + return '.jpg'; |
| 13 | +} |
| 14 | + |
| 15 | +function normalizeBooleanFlag(value: unknown): boolean { |
| 16 | + if (typeof value === 'boolean') return value; |
| 17 | + const normalized = String(value ?? '').trim().toLowerCase(); |
| 18 | + return normalized === 'true' || normalized === '1' || normalized === 'yes' || normalized === 'on'; |
| 19 | +} |
| 20 | + |
| 21 | +function displayPath(filePath: string): string { |
| 22 | + const home = os.homedir(); |
| 23 | + return filePath.startsWith(home) ? `~${filePath.slice(home.length)}` : filePath; |
| 24 | +} |
| 25 | + |
| 26 | +function buildImagePrompt(prompt: string, options: { |
| 27 | + ratio?: string; |
| 28 | + style?: string; |
| 29 | +}): string { |
| 30 | + const extras: string[] = []; |
| 31 | + if (options.ratio) extras.push(`aspect ratio ${options.ratio}`); |
| 32 | + if (options.style) extras.push(`style ${options.style}`); |
| 33 | + if (extras.length === 0) return prompt; |
| 34 | + return `${prompt} |
| 35 | +
|
| 36 | +Image requirements: ${extras.join(', ')}.`; |
| 37 | +} |
| 38 | + |
| 39 | +function normalizeRatio(value: string): string { |
| 40 | + const normalized = value.trim(); |
| 41 | + const allowed = new Set(['1:1', '16:9', '9:16', '4:3', '3:4', '3:2', '2:3']); |
| 42 | + return allowed.has(normalized) ? normalized : '1:1'; |
| 43 | +} |
| 44 | +async function currentGeminiLink(page: IPage): Promise<string> { |
| 45 | + const url = await page.evaluate('window.location.href').catch(() => ''); |
| 46 | + return typeof url === 'string' && url ? url : 'https://gemini.google.com/app'; |
| 47 | +} |
| 48 | + |
| 49 | +export const imageCommand = cli({ |
| 50 | + site: 'gemini', |
| 51 | + name: 'image', |
| 52 | + description: 'Generate images with Gemini web and save them locally', |
| 53 | + domain: GEMINI_DOMAIN, |
| 54 | + strategy: Strategy.COOKIE, |
| 55 | + browser: true, |
| 56 | + navigateBefore: false, |
| 57 | + defaultFormat: 'plain', |
| 58 | + timeoutSeconds: 240, |
| 59 | + args: [ |
| 60 | + { name: 'prompt', positional: true, required: true, help: 'Image prompt to send to Gemini' }, |
| 61 | + { name: 'rt', default: '1:1', help: 'Ratio shorthand for aspect ratio (1:1, 16:9, 9:16, 4:3, 3:4, 3:2, 2:3)' }, |
| 62 | + { name: 'st', default: '', help: 'Style shorthand, e.g. anime, icon, watercolor' }, |
| 63 | + { name: 'op', default: path.join(os.homedir(), 'tmp', 'gemini-images'), help: 'Output directory shorthand' }, |
| 64 | + { name: 'sd', type: 'boolean', default: false, help: 'Skip download shorthand; only show Gemini page link' }, |
| 65 | + ], |
| 66 | + columns: ['status', 'file', 'link'], |
| 67 | + func: async (page: IPage, kwargs: any) => { |
| 68 | + const prompt = kwargs.prompt as string; |
| 69 | + const ratio = normalizeRatio(String(kwargs.rt ?? '1:1')); |
| 70 | + const style = String(kwargs.st ?? '').trim(); |
| 71 | + const outputDir = (kwargs.op as string) || path.join(os.homedir(), 'tmp', 'gemini-images'); |
| 72 | + const timeout = 120; |
| 73 | + const startFresh = true; |
| 74 | + const skipDownloadRaw = kwargs.sd; |
| 75 | + const skipDownload = skipDownloadRaw === '' || skipDownloadRaw === true || normalizeBooleanFlag(skipDownloadRaw); |
| 76 | + |
| 77 | + const effectivePrompt = buildImagePrompt(prompt, { |
| 78 | + ratio, |
| 79 | + style: style || undefined, |
| 80 | + }); |
| 81 | + |
| 82 | + if (startFresh) await startNewGeminiChat(page); |
| 83 | + |
| 84 | + const beforeUrls = await getGeminiVisibleImageUrls(page); |
| 85 | + await sendGeminiMessage(page, effectivePrompt); |
| 86 | + const urls = await waitForGeminiImages(page, beforeUrls, timeout); |
| 87 | + const link = await currentGeminiLink(page); |
| 88 | + |
| 89 | + if (!urls.length) { |
| 90 | + return [{ status: '⚠️ no-images', file: '📁 -', link: `🔗 ${link}` }]; |
| 91 | + } |
| 92 | + |
| 93 | + if (skipDownload) { |
| 94 | + return [{ status: '🎨 generated', file: '📁 -', link: `🔗 ${link}` }]; |
| 95 | + } |
| 96 | + |
| 97 | + const assets = await exportGeminiImages(page, urls); |
| 98 | + if (!assets.length) { |
| 99 | + return [{ status: '⚠️ export-failed', file: '📁 -', link: `🔗 ${link}` }]; |
| 100 | + } |
| 101 | + |
| 102 | + const stamp = Date.now(); |
| 103 | + const results = []; |
| 104 | + for (let index = 0; index < assets.length; index += 1) { |
| 105 | + const asset = assets[index]; |
| 106 | + const base64 = asset.dataUrl.replace(/^data:[^;]+;base64,/, ''); |
| 107 | + const suffix = assets.length > 1 ? `_${index + 1}` : ''; |
| 108 | + const filePath = path.join(outputDir, `gemini_${stamp}${suffix}${extFromMime(asset.mimeType)}`); |
| 109 | + await saveBase64ToFile(base64, filePath); |
| 110 | + results.push({ status: '✅ saved', file: `📁 ${displayPath(filePath)}`, link: `🔗 ${link}` }); |
| 111 | + } |
| 112 | + |
| 113 | + return results; |
| 114 | + }, |
| 115 | +}); |
0 commit comments