From 6dd9550f361c0c19be4a43bd32fe38e6ea8d0123 Mon Sep 17 00:00:00 2001 From: Nicholas Cook Date: Fri, 20 Sep 2024 09:00:31 -0700 Subject: [PATCH] feat: add a generate image sample and test for Imagen (#3861) * feat: add a generate image sample and test for Imagen * fix lint * trigger build * trigger build --- ai-platform/snippets/imagen/generateImage.js | 106 ++++++++++++++++++ ai-platform/snippets/package.json | 2 +- .../test/imagen/generateImage.test.js | 34 ++++++ 3 files changed, 141 insertions(+), 1 deletion(-) create mode 100644 ai-platform/snippets/imagen/generateImage.js create mode 100644 ai-platform/snippets/test/imagen/generateImage.test.js diff --git a/ai-platform/snippets/imagen/generateImage.js b/ai-platform/snippets/imagen/generateImage.js new file mode 100644 index 0000000000..546b72b744 --- /dev/null +++ b/ai-platform/snippets/imagen/generateImage.js @@ -0,0 +1,106 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +'use strict'; + +const main = async () => { + // [START generativeaionvertexai_imagen_generate_image] + /** + * TODO(developer): Update these variables before running the sample. + */ + const projectId = process.env.CAIP_PROJECT_ID; + const outputFile = 'my-output'; + const prompt = 'a dog reading a newspaper'; // The text prompt describing what you want to see + const location = 'us-central1'; + + const aiplatform = require('@google-cloud/aiplatform'); + + // Imports the Google Cloud Prediction Service Client library + const {PredictionServiceClient} = aiplatform.v1; + + // Import the helper module for converting arbitrary protobuf.Value objects + const {helpers} = aiplatform; + + // Specifies the location of the api endpoint + const clientOptions = { + apiEndpoint: `${location}-aiplatform.googleapis.com`, + }; + + // Instantiates a client + const predictionServiceClient = new PredictionServiceClient(clientOptions); + + const fs = require('fs'); + const util = require('util'); + + async function generateImage() { + // Configure the parent resource + const endpoint = `projects/${projectId}/locations/${location}/publishers/google/models/imagen-3.0-generate-001`; + + const promptText = { + prompt: prompt, + }; + const instanceValue = helpers.toValue(promptText); + const instances = [instanceValue]; + + const parameter = { + sampleCount: 1, + // You can't use a seed value and watermark at the same time. + // seed: 100, + // addWatermark: false, + aspectRatio: '1:1', + safetyFilterLevel: 'block_some', + personGeneration: 'allow_adult', + }; + const parameters = helpers.toValue(parameter); + + const request = { + endpoint, + instances, + parameters, + }; + + // Predict request + const [response] = await predictionServiceClient.predict(request); + const predictions = response.predictions; + if (predictions.length === 0) { + console.log( + 'No image was generated. Check the request parameters and prompt.' + ); + } else { + let i = 1; + for (const prediction of predictions) { + const buff = Buffer.from( + prediction.structValue.fields.bytesBase64Encoded.stringValue, + 'base64' + ); + // Write image content to the output file + const writeFile = util.promisify(fs.writeFile); + const filename = `${outputFile}${i}.png`; + await writeFile(filename, buff); + console.log(`Saved image ${filename}`); + i++; + } + } + } + await generateImage(); + // [END generativeaionvertexai_imagen_generate_image] +}; + +// node generateImage.js +main().catch(err => { + console.log(err); + process.exitCode = 1; +}); diff --git a/ai-platform/snippets/package.json b/ai-platform/snippets/package.json index 4d7ad0964f..49c852c5fd 100644 --- a/ai-platform/snippets/package.json +++ b/ai-platform/snippets/package.json @@ -10,7 +10,7 @@ "*.js" ], "scripts": { - "test": "c8 mocha -p -j 2 --timeout 2400000 test/*.js" + "test": "c8 mocha -p -j 2 --timeout 2400000 test/*.js test/**/*.test.js" }, "dependencies": { "@google-cloud/aiplatform": "^3.0.0", diff --git a/ai-platform/snippets/test/imagen/generateImage.test.js b/ai-platform/snippets/test/imagen/generateImage.test.js new file mode 100644 index 0000000000..3e460a95a5 --- /dev/null +++ b/ai-platform/snippets/test/imagen/generateImage.test.js @@ -0,0 +1,34 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +'use strict'; + +const path = require('path'); +const {assert} = require('chai'); +const {describe, it} = require('mocha'); + +const cp = require('child_process'); +const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'}); +const cwd = path.join(__dirname, '..'); + +describe('AI platform generate image', () => { + it('should generate an image using the Imagen model', async () => { + const stdout = execSync('node ./imagen/generateImage.js', { + cwd, + }); + assert.match(stdout, /Saved image my-output1.png/); + }); +});