@genkit-ai/vertexai/lib/imagen.mjs.map

Genkit AI framework plugin for Google Cloud Vertex AI APIs including Gemini APIs, Imagen, and more.

{"version":3,"sources":["../src/imagen.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport { Genkit, z } from 'genkit';\nimport {\n  CandidateData,\n  GenerateRequest,\n  GenerationCommonConfigSchema,\n  ModelAction,\n  ModelReference,\n  getBasicUsageStats,\n  modelRef,\n} from 'genkit/model';\nimport { GoogleAuth } from 'google-auth-library';\nimport { PluginOptions } from './common/types.js';\nimport { PredictClient, predictModel } from './predict.js';\n\n/**\n * See https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api.\n */\nconst ImagenConfigSchema = GenerationCommonConfigSchema.extend({\n  // TODO: Remove common config schema extension since Imagen models don't support\n  // most of the common config parameters. Also, add more parameters like sampleCount\n  // from the above reference.\n  language: z\n    .enum(['auto', 'en', 'es', 'hi', 'ja', 'ko', 'pt', 'zh-TW', 'zh', 'zh-CN'])\n    .describe('Language of the prompt text.')\n    .optional(),\n  aspectRatio: z\n    .enum(['1:1', '9:16', '16:9', '3:4', '4:3'])\n    .describe('Desired aspect ratio of the output image.')\n    .optional(),\n  negativePrompt: z\n    .string()\n    .describe(\n      'A description of what to discourage in the generated images. ' +\n        'For example: \"animals\" (removes animals), \"blurry\" ' +\n        '(makes the image clearer), \"text\" (removes text), or ' +\n        '\"cropped\" (removes cropped images).'\n    )\n    .optional(),\n  seed: z\n    .number()\n    .int()\n    .min(1)\n    .max(2147483647)\n    .describe(\n      'Controls the randomization of the image generation process. Use the ' +\n        'same seed across requests to provide consistency, or change it to ' +\n        'introduce variety in the response.'\n    )\n    .optional(),\n  location: z\n    .string()\n    .describe('Google Cloud region e.g. us-central1.')\n    .optional(),\n  personGeneration: z\n    .enum(['dont_allow', 'allow_adult', 'allow_all'])\n    .describe('Control if/how images of people will be generated by the model.')\n    .optional(),\n  safetySetting: z\n    .enum(['block_most', 'block_some', 'block_few', 'block_fewest'])\n    .describe('Adds a filter level to safety filtering.')\n    .optional(),\n  addWatermark: z\n    .boolean()\n    .describe('Add an invisible watermark to the generated images.')\n    .optional(),\n  storageUri: z\n    .string()\n    .describe('Cloud Storage URI to store the generated images.')\n    .optional(),\n  mode: z\n    .enum(['upscale'])\n    .describe('Mode must be set for upscaling requests.')\n    .optional(),\n  /**\n   * Describes the editing intention for the request.\n   *\n   * See https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api#edit_images_2 for details.\n   */\n  editConfig: z\n    .object({\n      editMode: z\n        .enum([\n          'inpainting-insert',\n          'inpainting-remove',\n          'outpainting',\n          'product-image',\n        ])\n        .describe('Editing intention for the request.')\n        .optional(),\n      maskMode: z\n        .object({\n          maskType: z\n            .enum(['background', 'foreground', 'semantic'])\n            .describe(\n              '\"background\" automatically generates a mask for all ' +\n                'regions except the primary subject(s) of the image, ' +\n                '\"foreground\" automatically generates a mask for the primary ' +\n                'subjects(s) of the image. \"semantic\" segments one or more ' +\n                'of the segmentation classes using class ID.'\n            ),\n          classes: z\n            .array(z.number())\n            .describe('List of class IDs for segmentation.')\n            .length(5)\n            .optional(),\n        })\n        .describe(\n          'Prompts the model to generate a mask instead of you ' +\n            'needing to provide one. Consequently, when you provide ' +\n            'this parameter you can omit a mask object.'\n        )\n        .optional(),\n      maskDilation: z\n        .number()\n        .describe('Dilation percentage of the mask provided.')\n        .min(0.0)\n        .max(1.0)\n        .optional(),\n      guidanceScale: z\n        .number()\n        .describe(\n          'Controls how much the model adheres to the text prompt. ' +\n            'Large values increase output and prompt alignment, but may ' +\n            'compromise image quality. Suggested values are 0-9 ' +\n            '(low strength), 10-20 (medium strength), 21+ (high strength).'\n        )\n        .optional(),\n      productPosition: z\n        .enum(['reposition', 'fixed'])\n        .describe(\n          'Defines whether the product should stay fixed or be ' +\n            'repositioned.'\n        )\n        .optional(),\n    })\n    .passthrough()\n    .optional(),\n  upscaleConfig: z\n    .object({\n      upscaleFactor: z\n        .enum(['x2', 'x4'])\n        .describe('The factor to upscale the image.'),\n    })\n    .describe('Configuration for upscaling.')\n    .optional(),\n}).passthrough();\n\nexport const imagen2 = modelRef({\n  name: 'vertexai/imagen2',\n  info: {\n    label: 'Vertex AI - Imagen2',\n    versions: ['imagegeneration@006', 'imagegeneration@005'],\n    supports: {\n      media: false,\n      multiturn: false,\n      tools: false,\n      systemRole: false,\n      output: ['media'],\n    },\n  },\n  version: 'imagegeneration@006',\n  configSchema: ImagenConfigSchema,\n});\n\nexport const imagen3 = modelRef({\n  name: 'vertexai/imagen3',\n  info: {\n    label: 'Vertex AI - Imagen3',\n    versions: ['imagen-3.0-generate-001'],\n    supports: {\n      media: true,\n      multiturn: false,\n      tools: false,\n      systemRole: false,\n      output: ['media'],\n    },\n  },\n  version: 'imagen-3.0-generate-001',\n  configSchema: ImagenConfigSchema,\n});\n\nexport const imagen3Fast = modelRef({\n  name: 'vertexai/imagen3-fast',\n  info: {\n    label: 'Vertex AI - Imagen3 Fast',\n    versions: ['imagen-3.0-fast-generate-001'],\n    supports: {\n      media: false,\n      multiturn: false,\n      tools: false,\n      systemRole: false,\n      output: ['media'],\n    },\n  },\n  version: 'imagen-3.0-fast-generate-001',\n  configSchema: ImagenConfigSchema,\n});\n\nexport const SUPPORTED_IMAGEN_MODELS = {\n  imagen2: imagen2,\n  imagen3: imagen3,\n  'imagen3-fast': imagen3Fast,\n};\n\nfunction extractText(request: GenerateRequest) {\n  return request.messages\n    .at(-1)!\n    .content.map((c) => c.text || '')\n    .join('');\n}\n\ninterface ImagenParameters {\n  sampleCount?: number;\n  aspectRatio?: string;\n  negativePrompt?: string;\n  seed?: number;\n  language?: string;\n  personGeneration?: string;\n  safetySetting?: string;\n  addWatermark?: boolean;\n  storageUri?: string;\n}\n\nfunction toParameters(\n  request: GenerateRequest<typeof ImagenConfigSchema>\n): ImagenParameters {\n  const out = {\n    sampleCount: request.candidates ?? 1,\n    ...request?.config,\n  };\n\n  for (const k in out) {\n    if (!out[k]) delete out[k];\n  }\n\n  return out;\n}\n\nfunction extractMaskImage(request: GenerateRequest): string | undefined {\n  return request.messages\n    .at(-1)\n    ?.content.find((p) => !!p.media && p.metadata?.type === 'mask')\n    ?.media?.url.split(',')[1];\n}\n\nfunction extractBaseImage(request: GenerateRequest): string | undefined {\n  return request.messages\n    .at(-1)\n    ?.content.find(\n      (p) => !!p.media && (!p.metadata?.type || p.metadata?.type === 'base')\n    )\n    ?.media?.url.split(',')[1];\n}\n\ninterface ImagenPrediction {\n  bytesBase64Encoded: string;\n  mimeType: string;\n}\n\ninterface ImagenInstance {\n  prompt: string;\n  image?: { bytesBase64Encoded: string };\n  mask?: { image?: { bytesBase64Encoded: string } };\n}\n\nexport function imagenModel(\n  ai: Genkit,\n  name: string,\n  client: GoogleAuth,\n  options: PluginOptions\n): ModelAction {\n  const modelName = `vertexai/${name}`;\n  const model: ModelReference<z.ZodTypeAny> = SUPPORTED_IMAGEN_MODELS[name];\n  if (!model) throw new Error(`Unsupported model: ${name}`);\n\n  const predictClients: Record<\n    string,\n    PredictClient<ImagenInstance, ImagenPrediction, ImagenParameters>\n  > = {};\n  const predictClientFactory = (\n    request: GenerateRequest<typeof ImagenConfigSchema>\n  ): PredictClient<ImagenInstance, ImagenPrediction, ImagenParameters> => {\n    const requestLocation = request.config?.location || options.location;\n    if (!predictClients[requestLocation]) {\n      predictClients[requestLocation] = predictModel<\n        ImagenInstance,\n        ImagenPrediction,\n        ImagenParameters\n      >(\n        client,\n        {\n          ...options,\n          location: requestLocation,\n        },\n        request.config?.version || model.version || name\n      );\n    }\n    return predictClients[requestLocation];\n  };\n\n  return ai.defineModel(\n    {\n      name: modelName,\n      ...model.info,\n      configSchema: ImagenConfigSchema,\n    },\n    async (request) => {\n      const instance: ImagenInstance = {\n        prompt: extractText(request),\n      };\n      const baseImage = extractBaseImage(request);\n      if (baseImage) {\n        instance.image = { bytesBase64Encoded: baseImage };\n      }\n      const maskImage = extractMaskImage(request);\n      if (maskImage) {\n        instance.mask = {\n          image: { bytesBase64Encoded: maskImage },\n        };\n      }\n\n      const predictClient = predictClientFactory(request);\n      const response = await predictClient([instance], toParameters(request));\n\n      const candidates: CandidateData[] = response.predictions.map((p, i) => {\n        const b64data = p.bytesBase64Encoded;\n        const mimeType = p.mimeType;\n        return {\n          index: i,\n          finishReason: 'stop',\n          message: {\n            role: 'model',\n            content: [\n              {\n                media: {\n                  url: `data:${mimeType};base64,${b64data}`,\n                  contentType: mimeType,\n                },\n              },\n            ],\n          },\n        };\n      });\n      return {\n        candidates,\n        usage: {\n          ...getBasicUsageStats(request.messages, candidates),\n          custom: { generations: candidates.length },\n        },\n        custom: response,\n      };\n    }\n  );\n}\n"],"mappings":"AAgBA,SAAiB,SAAS;AAC1B;AAAA,EAGE;AAAA,EAGA;AAAA,EACA;AAAA,OACK;AAGP,SAAwB,oBAAoB;AAK5C,MAAM,qBAAqB,6BAA6B,OAAO;AAAA;AAAA;AAAA;AAAA,EAI7D,UAAU,EACP,KAAK,CAAC,QAAQ,MAAM,MAAM,MAAM,MAAM,MAAM,MAAM,SAAS,MAAM,OAAO,CAAC,EACzE,SAAS,8BAA8B,EACvC,SAAS;AAAA,EACZ,aAAa,EACV,KAAK,CAAC,OAAO,QAAQ,QAAQ,OAAO,KAAK,CAAC,EAC1C,SAAS,2CAA2C,EACpD,SAAS;AAAA,EACZ,gBAAgB,EACb,OAAO,EACP;AAAA,IACC;AAAA,EAIF,EACC,SAAS;AAAA,EACZ,MAAM,EACH,OAAO,EACP,IAAI,EACJ,IAAI,CAAC,EACL,IAAI,UAAU,EACd;AAAA,IACC;AAAA,EAGF,EACC,SAAS;AAAA,EACZ,UAAU,EACP,OAAO,EACP,SAAS,uCAAuC,EAChD,SAAS;AAAA,EACZ,kBAAkB,EACf,KAAK,CAAC,cAAc,eAAe,WAAW,CAAC,EAC/C,SAAS,iEAAiE,EAC1E,SAAS;AAAA,EACZ,eAAe,EACZ,KAAK,CAAC,cAAc,cAAc,aAAa,cAAc,CAAC,EAC9D,SAAS,0CAA0C,EACnD,SAAS;AAAA,EACZ,cAAc,EACX,QAAQ,EACR,SAAS,qDAAqD,EAC9D,SAAS;AAAA,EACZ,YAAY,EACT,OAAO,EACP,SAAS,kDAAkD,EAC3D,SAAS;AAAA,EACZ,MAAM,EACH,KAAK,CAAC,SAAS,CAAC,EAChB,SAAS,0CAA0C,EACnD,SAAS;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAMZ,YAAY,EACT,OAAO;AAAA,IACN,UAAU,EACP,KAAK;AAAA,MACJ;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,IACF,CAAC,EACA,SAAS,oCAAoC,EAC7C,SAAS;AAAA,IACZ,UAAU,EACP,OAAO;AAAA,MACN,UAAU,EACP,KAAK,CAAC,cAAc,cAAc,UAAU,CAAC,EAC7C;AAAA,QACC;AAAA,MAKF;AAAA,MACF,SAAS,EACN,MAAM,EAAE,OAAO,CAAC,EAChB,SAAS,qCAAqC,EAC9C,OAAO,CAAC,EACR,SAAS;AAAA,IACd,CAAC,EACA;AAAA,MACC;AAAA,IAGF,EACC,SAAS;AAAA,IACZ,cAAc,EACX,OAAO,EACP,SAAS,2CAA2C,EACpD,IAAI,CAAG,EACP,IAAI,CAAG,EACP,SAAS;AAAA,IACZ,eAAe,EACZ,OAAO,EACP;AAAA,MACC;AAAA,IAIF,EACC,SAAS;AAAA,IACZ,iBAAiB,EACd,KAAK,CAAC,cAAc,OAAO,CAAC,EAC5B;AAAA,MACC;AAAA,IAEF,EACC,SAAS;AAAA,EACd,CAAC,EACA,YAAY,EACZ,SAAS;AAAA,EACZ,eAAe,EACZ,OAAO;AAAA,IACN,eAAe,EACZ,KAAK,CAAC,MAAM,IAAI,CAAC,EACjB,SAAS,kCAAkC;AAAA,EAChD,CAAC,EACA,SAAS,8BAA8B,EACvC,SAAS;AACd,CAAC,EAAE,YAAY;AAER,MAAM,UAAU,SAAS;AAAA,EAC9B,MAAM;AAAA,EACN,MAAM;AAAA,IACJ,OAAO;AAAA,IACP,UAAU,CAAC,uBAAuB,qBAAqB;AAAA,IACvD,UAAU;AAAA,MACR,OAAO;AAAA,MACP,WAAW;AAAA,MACX,OAAO;AAAA,MACP,YAAY;AAAA,MACZ,QAAQ,CAAC,OAAO;AAAA,IAClB;AAAA,EACF;AAAA,EACA,SAAS;AAAA,EACT,cAAc;AAChB,CAAC;AAEM,MAAM,UAAU,SAAS;AAAA,EAC9B,MAAM;AAAA,EACN,MAAM;AAAA,IACJ,OAAO;AAAA,IACP,UAAU,CAAC,yBAAyB;AAAA,IACpC,UAAU;AAAA,MACR,OAAO;AAAA,MACP,WAAW;AAAA,MACX,OAAO;AAAA,MACP,YAAY;AAAA,MACZ,QAAQ,CAAC,OAAO;AAAA,IAClB;AAAA,EACF;AAAA,EACA,SAAS;AAAA,EACT,cAAc;AAChB,CAAC;AAEM,MAAM,cAAc,SAAS;AAAA,EAClC,MAAM;AAAA,EACN,MAAM;AAAA,IACJ,OAAO;AAAA,IACP,UAAU,CAAC,8BAA8B;AAAA,IACzC,UAAU;AAAA,MACR,OAAO;AAAA,MACP,WAAW;AAAA,MACX,OAAO;AAAA,MACP,YAAY;AAAA,MACZ,QAAQ,CAAC,OAAO;AAAA,IAClB;AAAA,EACF;AAAA,EACA,SAAS;AAAA,EACT,cAAc;AAChB,CAAC;AAEM,MAAM,0BAA0B;AAAA,EACrC;AAAA,EACA;AAAA,EACA,gBAAgB;AAClB;AAEA,SAAS,YAAY,SAA0B;AAC7C,SAAO,QAAQ,SACZ,GAAG,EAAE,EACL,QAAQ,IAAI,CAAC,MAAM,EAAE,QAAQ,EAAE,EAC/B,KAAK,EAAE;AACZ;AAcA,SAAS,aACP,SACkB;AAClB,QAAM,MAAM;AAAA,IACV,aAAa,QAAQ,cAAc;AAAA,IACnC,GAAG,SAAS;AAAA,EACd;AAEA,aAAW,KAAK,KAAK;AACnB,QAAI,CAAC,IAAI,CAAC,EAAG,QAAO,IAAI,CAAC;AAAA,EAC3B;AAEA,SAAO;AACT;AAEA,SAAS,iBAAiB,SAA8C;AACtE,SAAO,QAAQ,SACZ,GAAG,EAAE,GACJ,QAAQ,KAAK,CAAC,MAAM,CAAC,CAAC,EAAE,SAAS,EAAE,UAAU,SAAS,MAAM,GAC5D,OAAO,IAAI,MAAM,GAAG,EAAE,CAAC;AAC7B;AAEA,SAAS,iBAAiB,SAA8C;AACtE,SAAO,QAAQ,SACZ,GAAG,EAAE,GACJ,QAAQ;AAAA,IACR,CAAC,MAAM,CAAC,CAAC,EAAE,UAAU,CAAC,EAAE,UAAU,QAAQ,EAAE,UAAU,SAAS;AAAA,EACjE,GACE,OAAO,IAAI,MAAM,GAAG,EAAE,CAAC;AAC7B;AAaO,SAAS,YACd,IACA,MACA,QACA,SACa;AACb,QAAM,YAAY,YAAY,IAAI;AAClC,QAAM,QAAsC,wBAAwB,IAAI;AACxE,MAAI,CAAC,MAAO,OAAM,IAAI,MAAM,sBAAsB,IAAI,EAAE;AAExD,QAAM,iBAGF,CAAC;AACL,QAAM,uBAAuB,CAC3B,YACsE;AACtE,UAAM,kBAAkB,QAAQ,QAAQ,YAAY,QAAQ;AAC5D,QAAI,CAAC,eAAe,eAAe,GAAG;AACpC,qBAAe,eAAe,IAAI;AAAA,QAKhC;AAAA,QACA;AAAA,UACE,GAAG;AAAA,UACH,UAAU;AAAA,QACZ;AAAA,QACA,QAAQ,QAAQ,WAAW,MAAM,WAAW;AAAA,MAC9C;AAAA,IACF;AACA,WAAO,eAAe,eAAe;AAAA,EACvC;AAEA,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM;AAAA,MACN,GAAG,MAAM;AAAA,MACT,cAAc;AAAA,IAChB;AAAA,IACA,OAAO,YAAY;AACjB,YAAM,WAA2B;AAAA,QAC/B,QAAQ,YAAY,OAAO;AAAA,MAC7B;AACA,YAAM,YAAY,iBAAiB,OAAO;AAC1C,UAAI,WAAW;AACb,iBAAS,QAAQ,EAAE,oBAAoB,UAAU;AAAA,MACnD;AACA,YAAM,YAAY,iBAAiB,OAAO;AAC1C,UAAI,WAAW;AACb,iBAAS,OAAO;AAAA,UACd,OAAO,EAAE,oBAAoB,UAAU;AAAA,QACzC;AAAA,MACF;AAEA,YAAM,gBAAgB,qBAAqB,OAAO;AAClD,YAAM,WAAW,MAAM,cAAc,CAAC,QAAQ,GAAG,aAAa,OAAO,CAAC;AAEtE,YAAM,aAA8B,SAAS,YAAY,IAAI,CAAC,GAAG,MAAM;AACrE,cAAM,UAAU,EAAE;AAClB,cAAM,WAAW,EAAE;AACnB,eAAO;AAAA,UACL,OAAO;AAAA,UACP,cAAc;AAAA,UACd,SAAS;AAAA,YACP,MAAM;AAAA,YACN,SAAS;AAAA,cACP;AAAA,gBACE,OAAO;AAAA,kBACL,KAAK,QAAQ,QAAQ,WAAW,OAAO;AAAA,kBACvC,aAAa;AAAA,gBACf;AAAA,cACF;AAAA,YACF;AAAA,UACF;AAAA,QACF;AAAA,MACF,CAAC;AACD,aAAO;AAAA,QACL;AAAA,QACA,OAAO;AAAA,UACL,GAAG,mBAAmB,QAAQ,UAAU,UAAU;AAAA,UAClD,QAAQ,EAAE,aAAa,WAAW,OAAO;AAAA,QAC3C;AAAA,QACA,QAAQ;AAAA,MACV;AAAA,IACF;AAAA,EACF;AACF;","names":[]}