feat(tools): add image generation via Google AI Imagen 3
output: /mnt/work/tmp/image-gen/ triggered by !pic <prompt> in discord
This commit is contained in:
parent
bc0ae7566e
commit
5de4135ad2
70
src/tools/image-gen.ts
Normal file
70
src/tools/image-gen.ts
Normal file
@ -0,0 +1,70 @@
|
||||
/**
|
||||
* Image generation tool using Google AI (Imagen 3).
|
||||
* Usage: bun run src/tools/image-gen.ts "prompt here"
|
||||
* Outputs the path to the generated image.
|
||||
*/
|
||||
|
||||
const API_KEY = process.env.GOOGLE_AI_API_KEY;
|
||||
const OUTPUT_DIR = '/mnt/work/tmp/image-gen';
|
||||
const MODEL = 'imagen-3.0-generate-002';
|
||||
|
||||
async function generateImage(prompt: string): Promise<string> {
|
||||
if (!API_KEY) throw new Error('Missing GOOGLE_AI_API_KEY in .env');
|
||||
|
||||
console.error(`Generating: "${prompt}"...`);
|
||||
|
||||
const res = await fetch(
|
||||
`https://generativelanguage.googleapis.com/v1beta/models/${MODEL}:predict?key=${API_KEY}`,
|
||||
{
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
instances: [{ prompt }],
|
||||
parameters: {
|
||||
sampleCount: 1,
|
||||
aspectRatio: '1:1',
|
||||
safetyFilterLevel: 'block_few',
|
||||
},
|
||||
}),
|
||||
signal: AbortSignal.timeout(60_000),
|
||||
}
|
||||
);
|
||||
|
||||
if (!res.ok) {
|
||||
const text = await res.text();
|
||||
throw new Error(`API error ${res.status}: ${text}`);
|
||||
}
|
||||
|
||||
const data = await res.json();
|
||||
const b64 = data.predictions?.[0]?.bytesBase64Encoded;
|
||||
|
||||
if (!b64) {
|
||||
throw new Error(`No image in response: ${JSON.stringify(data).slice(0, 500)}`);
|
||||
}
|
||||
|
||||
const { mkdirSync, writeFileSync } = await import('fs');
|
||||
mkdirSync(OUTPUT_DIR, { recursive: true });
|
||||
|
||||
const filename = `img_${Date.now()}.png`;
|
||||
const outPath = `${OUTPUT_DIR}/${filename}`;
|
||||
writeFileSync(outPath, Buffer.from(b64, 'base64'));
|
||||
|
||||
console.error(`Saved: ${outPath}`);
|
||||
return outPath;
|
||||
}
|
||||
|
||||
async function main() {
|
||||
const prompt = process.argv.slice(2).join(' ');
|
||||
if (!prompt) {
|
||||
console.error('Usage: bun run src/tools/image-gen.ts "prompt here"');
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const path = await generateImage(prompt);
|
||||
console.log(path);
|
||||
}
|
||||
|
||||
main().catch(e => {
|
||||
console.error(e);
|
||||
process.exit(1);
|
||||
});
|
||||
Loading…
x
Reference in New Issue
Block a user