{"version":3,"sources":["../../src/vertexai/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 { ActionMetadata, modelActionMetadata, z } from 'genkit';\nimport {\n  GenerationCommonConfigSchema,\n  ModelAction,\n  ModelInfo,\n  ModelReference,\n  modelRef,\n} from 'genkit/model';\nimport { model as pluginModel } from 'genkit/plugin';\nimport { isKnownKey } from '../common/utils.mjs';\nimport { imagenPredict } from './client.mjs';\nimport { fromImagenResponse, toImagenPredictRequest } from './converters.mjs';\nimport { ClientOptions, Model, VertexPluginOptions } from './types.mjs';\nimport {\n  calculateRequestOptions,\n  checkModelName,\n  extractVersion,\n  modelName,\n} from './utils.mjs';\n\n/**\n * See https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api.\n */\nexport const 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        .passthrough()\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    .passthrough()\n    .optional(),\n}).passthrough();\nexport type ImagenConfigSchemaType = typeof ImagenConfigSchema;\nexport type ImagenConfig = z.infer<ImagenConfigSchemaType>;\n\nexport const ImagenTryOnConfigSchema = z\n  .object({\n    sampleCount: z\n      .number()\n      .int()\n      .min(1)\n      .optional()\n      .describe('Number of images to generate.'),\n    seed: z\n      .number()\n      .int()\n      .optional()\n      .describe('Random seed for the image generation.'),\n    baseSteps: z\n      .number()\n      .int()\n      .min(1)\n      .max(100)\n      .optional()\n      .describe('Number of denoising steps.'),\n    personGeneration: z\n      .enum(['dont_allow', 'allow_adult', 'allow_all'])\n      .optional(),\n    safetySetting: z\n      .enum(['block_most', 'block_some', 'block_few', 'block_fewest'])\n      .optional(),\n    storageUri: z\n      .string()\n      .optional()\n      .describe('Cloud Storage URI to store the generated images.'),\n    outputOptions: z\n      .object({\n        mimeType: z.string().optional(),\n        compressionQuality: z.number().int().optional(),\n      })\n      .optional(),\n  })\n  .passthrough();\nexport type ImagenTryOnConfigSchemaType = typeof ImagenTryOnConfigSchema;\nexport type ImagenTryOnConfig = z.infer<ImagenTryOnConfigSchemaType>;\n\nexport type ConfigSchemaType =\n  | ImagenConfigSchemaType\n  | ImagenTryOnConfigSchemaType;\n\nfunction commonRef(\n  name: string,\n  info?: ModelInfo,\n  configSchema: ConfigSchemaType = ImagenConfigSchema\n): ModelReference<ConfigSchemaType> {\n  return modelRef({\n    name: `vertexai/${name}`,\n    configSchema,\n    info: info ?? {\n      supports: {\n        media: true,\n        multiturn: false,\n        tools: false,\n        toolChoice: false,\n        systemRole: false,\n        output: ['media'],\n      },\n    },\n  });\n}\n\n// There are no more future imagen models, set these to the commonRef values.\nconst GENERIC_MODEL = commonRef('imagen', {\n  supports: {\n    media: true,\n    multiturn: false,\n    tools: false,\n    toolChoice: false,\n    systemRole: false,\n    output: ['media'],\n  },\n});\n\n// The 'regular' imagen models are deprecated, but not shutdown yet.\n// Shutdown is June 30, 2026 - until then, we need to keep the same\n// logic as before in the definitions so we are backwards compatible.\nexport const KNOWN_MODELS = {\n  'virtual-try-on-001': commonRef(\n    'virtual-try-on-001',\n    {\n      supports: {\n        media: true,\n        multiturn: false,\n        tools: false,\n        systemRole: true,\n        output: ['media'],\n      },\n    },\n    ImagenTryOnConfigSchema\n  ),\n} as const;\nexport type KnownModels = keyof typeof KNOWN_MODELS;\n\n// Models we no longer suggest, but that still work: they still resolve and\n// are still registered as actions, but aren't in the KnownModels type, so\n// they don't show up in autocomplete. Set `stage` on each one (keeping the\n// default `supports`). Remove a model once the server returns 404 for it.\nconst HIDDEN_MODELS = {} as const;\n\nconst ALL_MODELS = {\n  ...HIDDEN_MODELS,\n  ...KNOWN_MODELS,\n};\n\nexport type ImagenModelName =\n  | KnownModels\n  | `imagen-${string}`\n  | `virtual-try-on-${string}`;\nexport function isImagenModelName(value?: string): value is ImagenModelName {\n  return (\n    !!value &&\n    (value.startsWith('imagen-') ||\n      value.startsWith('virtual-try-on-') ||\n      value in ALL_MODELS)\n  );\n}\n\nexport function model(\n  version: string,\n  config: ImagenConfig = {}\n): ModelReference<ConfigSchemaType> {\n  const name = checkModelName(version);\n\n  if (isKnownKey(name, ALL_MODELS)) {\n    return ALL_MODELS[name].withConfig(config);\n  }\n\n  return modelRef({\n    name: `vertexai/${name}`,\n    config,\n    configSchema: ImagenConfigSchema,\n    info: {\n      ...GENERIC_MODEL.info,\n    },\n  });\n}\n\nexport function listActions(models: Model[]): ActionMetadata[] {\n  return models\n    .filter((m: Model) => isImagenModelName(modelName(m.name)))\n    .map((m: Model) => {\n      const ref = model(m.name);\n      return modelActionMetadata({\n        name: ref.name,\n        info: ref.info,\n        configSchema: ref.configSchema,\n      });\n    });\n}\n\nexport function listKnownModels(\n  clientOptions: ClientOptions,\n  pluginOptions?: VertexPluginOptions\n) {\n  return Object.keys(ALL_MODELS).map((name: string) =>\n    defineModel(name, clientOptions, pluginOptions)\n  );\n}\n\nexport function defineModel(\n  name: string,\n  clientOptions: ClientOptions,\n  pluginOptions?: VertexPluginOptions\n): ModelAction {\n  const ref = model(name);\n\n  return pluginModel(\n    {\n      name: ref.name,\n      ...ref.info,\n      configSchema: ref.configSchema,\n    },\n    async (request, { abortSignal }) => {\n      const clientOpt = calculateRequestOptions(\n        { ...clientOptions, signal: abortSignal },\n        request.config\n      );\n\n      const imagenPredictRequest = toImagenPredictRequest(request);\n\n      const response = await imagenPredict(\n        extractVersion(ref),\n        imagenPredictRequest,\n        clientOpt\n      );\n\n      if (!response.predictions || response.predictions.length == 0) {\n        throw new Error(\n          'Model returned no predictions. Possibly due to content filters.'\n        );\n      }\n\n      return fromImagenResponse(response, request);\n    }\n  );\n}\n\nexport const TEST_ONLY = {\n  GENERIC_MODEL,\n  KNOWN_MODELS,\n  HIDDEN_MODELS,\n  ALL_MODELS,\n};\n"],"mappings":"AAgBA,SAAyB,qBAAqB,SAAS;AACvD;AAAA,EACE;AAAA,EAIA;AAAA,OACK;AACP,SAAS,SAAS,mBAAmB;AACrC,SAAS,kBAAkB;AAC3B,SAAS,qBAAqB;AAC9B,SAAS,oBAAoB,8BAA8B;AAE3D;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,OACK;AAKA,MAAM,qBAAqB,6BAA6B,OAAO;AAAA;AAAA;AAAA;AAAA,EAIpE,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,YAAY,EACZ,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,YAAY,EACZ,SAAS;AACd,CAAC,EAAE,YAAY;AAIR,MAAM,0BAA0B,EACpC,OAAO;AAAA,EACN,aAAa,EACV,OAAO,EACP,IAAI,EACJ,IAAI,CAAC,EACL,SAAS,EACT,SAAS,+BAA+B;AAAA,EAC3C,MAAM,EACH,OAAO,EACP,IAAI,EACJ,SAAS,EACT,SAAS,uCAAuC;AAAA,EACnD,WAAW,EACR,OAAO,EACP,IAAI,EACJ,IAAI,CAAC,EACL,IAAI,GAAG,EACP,SAAS,EACT,SAAS,4BAA4B;AAAA,EACxC,kBAAkB,EACf,KAAK,CAAC,cAAc,eAAe,WAAW,CAAC,EAC/C,SAAS;AAAA,EACZ,eAAe,EACZ,KAAK,CAAC,cAAc,cAAc,aAAa,cAAc,CAAC,EAC9D,SAAS;AAAA,EACZ,YAAY,EACT,OAAO,EACP,SAAS,EACT,SAAS,kDAAkD;AAAA,EAC9D,eAAe,EACZ,OAAO;AAAA,IACN,UAAU,EAAE,OAAO,EAAE,SAAS;AAAA,IAC9B,oBAAoB,EAAE,OAAO,EAAE,IAAI,EAAE,SAAS;AAAA,EAChD,CAAC,EACA,SAAS;AACd,CAAC,EACA,YAAY;AAQf,SAAS,UACP,MACA,MACA,eAAiC,oBACC;AAClC,SAAO,SAAS;AAAA,IACd,MAAM,YAAY,IAAI;AAAA,IACtB;AAAA,IACA,MAAM,QAAQ;AAAA,MACZ,UAAU;AAAA,QACR,OAAO;AAAA,QACP,WAAW;AAAA,QACX,OAAO;AAAA,QACP,YAAY;AAAA,QACZ,YAAY;AAAA,QACZ,QAAQ,CAAC,OAAO;AAAA,MAClB;AAAA,IACF;AAAA,EACF,CAAC;AACH;AAGA,MAAM,gBAAgB,UAAU,UAAU;AAAA,EACxC,UAAU;AAAA,IACR,OAAO;AAAA,IACP,WAAW;AAAA,IACX,OAAO;AAAA,IACP,YAAY;AAAA,IACZ,YAAY;AAAA,IACZ,QAAQ,CAAC,OAAO;AAAA,EAClB;AACF,CAAC;AAKM,MAAM,eAAe;AAAA,EAC1B,sBAAsB;AAAA,IACpB;AAAA,IACA;AAAA,MACE,UAAU;AAAA,QACR,OAAO;AAAA,QACP,WAAW;AAAA,QACX,OAAO;AAAA,QACP,YAAY;AAAA,QACZ,QAAQ,CAAC,OAAO;AAAA,MAClB;AAAA,IACF;AAAA,IACA;AAAA,EACF;AACF;AAOA,MAAM,gBAAgB,CAAC;AAEvB,MAAM,aAAa;AAAA,EACjB,GAAG;AAAA,EACH,GAAG;AACL;AAMO,SAAS,kBAAkB,OAA0C;AAC1E,SACE,CAAC,CAAC,UACD,MAAM,WAAW,SAAS,KACzB,MAAM,WAAW,iBAAiB,KAClC,SAAS;AAEf;AAEO,SAAS,MACd,SACA,SAAuB,CAAC,GACU;AAClC,QAAM,OAAO,eAAe,OAAO;AAEnC,MAAI,WAAW,MAAM,UAAU,GAAG;AAChC,WAAO,WAAW,IAAI,EAAE,WAAW,MAAM;AAAA,EAC3C;AAEA,SAAO,SAAS;AAAA,IACd,MAAM,YAAY,IAAI;AAAA,IACtB;AAAA,IACA,cAAc;AAAA,IACd,MAAM;AAAA,MACJ,GAAG,cAAc;AAAA,IACnB;AAAA,EACF,CAAC;AACH;AAEO,SAAS,YAAY,QAAmC;AAC7D,SAAO,OACJ,OAAO,CAAC,MAAa,kBAAkB,UAAU,EAAE,IAAI,CAAC,CAAC,EACzD,IAAI,CAAC,MAAa;AACjB,UAAM,MAAM,MAAM,EAAE,IAAI;AACxB,WAAO,oBAAoB;AAAA,MACzB,MAAM,IAAI;AAAA,MACV,MAAM,IAAI;AAAA,MACV,cAAc,IAAI;AAAA,IACpB,CAAC;AAAA,EACH,CAAC;AACL;AAEO,SAAS,gBACd,eACA,eACA;AACA,SAAO,OAAO,KAAK,UAAU,EAAE;AAAA,IAAI,CAAC,SAClC,YAAY,MAAM,eAAe,aAAa;AAAA,EAChD;AACF;AAEO,SAAS,YACd,MACA,eACA,eACa;AACb,QAAM,MAAM,MAAM,IAAI;AAEtB,SAAO;AAAA,IACL;AAAA,MACE,MAAM,IAAI;AAAA,MACV,GAAG,IAAI;AAAA,MACP,cAAc,IAAI;AAAA,IACpB;AAAA,IACA,OAAO,SAAS,EAAE,YAAY,MAAM;AAClC,YAAM,YAAY;AAAA,QAChB,EAAE,GAAG,eAAe,QAAQ,YAAY;AAAA,QACxC,QAAQ;AAAA,MACV;AAEA,YAAM,uBAAuB,uBAAuB,OAAO;AAE3D,YAAM,WAAW,MAAM;AAAA,QACrB,eAAe,GAAG;AAAA,QAClB;AAAA,QACA;AAAA,MACF;AAEA,UAAI,CAAC,SAAS,eAAe,SAAS,YAAY,UAAU,GAAG;AAC7D,cAAM,IAAI;AAAA,UACR;AAAA,QACF;AAAA,MACF;AAEA,aAAO,mBAAmB,UAAU,OAAO;AAAA,IAC7C;AAAA,EACF;AACF;AAEO,MAAM,YAAY;AAAA,EACvB;AAAA,EACA;AAAA,EACA;AAAA,EACA;AACF;","names":[]}