Titan Image Generator v2の新機能対応

packages/cdk/lambda/utils/models.ts

@@ -5,6 +5,7 @@ import {
   PromptTemplate,
   StableDiffusionParams,
   TitanImageParams,
+  TitanImageV2Params,
   UnrecordedMessage,
   ConverseInferenceParams,
   UsecaseConverseInferenceParams,
@@ -371,7 +372,7 @@ const createBodyImageStableDiffusion = (params: GenerateImageParams) => {
       init_image: params.initImage,
       mask_image: params.maskImage,
       mask_source:
-        params.maskMode === 'INPAINTING'
+        params.taskType === 'INPAINTING'
           ? 'MASK_IMAGE_BLACK'
           : 'MASK_IMAGE_WHITE',
     };
@@ -390,7 +391,7 @@ const createBodyImageTitanImage = (params: GenerateImageParams) => {
     seed: params.seed % 214783648, // max for titan image
   };
   let body: Partial<TitanImageParams> = {};
-  if (params.initImage && params.maskMode === undefined) {
+  if (params.initImage && params.taskType === 'IMAGE_VARIATION') {
     body = {
       taskType: 'IMAGE_VARIATION',
       imageVariationParams: {
@@ -404,7 +405,7 @@ const createBodyImageTitanImage = (params: GenerateImageParams) => {
       },
       imageGenerationConfig: imageGenerationConfig,
     };
-  } else if (params.initImage && params.maskMode === 'INPAINTING') {
+  } else if (params.initImage && params.taskType === 'INPAINTING') {
     body = {
       taskType: 'INPAINTING',
       inPaintingParams: {
@@ -419,7 +420,7 @@ const createBodyImageTitanImage = (params: GenerateImageParams) => {
       },
       imageGenerationConfig: imageGenerationConfig,
     };
-  } else if (params.initImage && params.maskMode === 'OUTPAINTING') {
+  } else if (params.initImage && params.taskType === 'OUTPAINTING') {
     body = {
       taskType: 'OUTPAINTING',
       outPaintingParams: {
@@ -435,7 +436,7 @@ const createBodyImageTitanImage = (params: GenerateImageParams) => {
       },
       imageGenerationConfig: imageGenerationConfig,
     };
-  } else {
+  } else if (params.taskType === 'TEXT_IMAGE') {
     body = {
       taskType: 'TEXT_IMAGE',
       textToImageParams: {
@@ -447,7 +448,51 @@ const createBodyImageTitanImage = (params: GenerateImageParams) => {
       },
       imageGenerationConfig: imageGenerationConfig,
     };
+  } else {
+    body = {
+      imageGenerationConfig: imageGenerationConfig,
+    };
+  }
+  return JSON.stringify(body);
+};
+
+const createBodyImageTitanImageV2 = (params: GenerateImageParams) => {
+  // 既存の関数を呼び出して基本的なボディを取得
+  const baseBody = JSON.parse(createBodyImageTitanImage(params));
+
+  let body: Partial<TitanImageV2Params> = {
+    ...baseBody,
+  };
+
+  // 新しいタスクタイプの処理
+  if (params.taskType === 'COLOR_GUIDED_GENERATION') {
+    body = {
+      taskType: 'COLOR_GUIDED_GENERATION',
+      colorGuidedGenerationParams: {
+        text: params.textPrompt.find((x) => x.weight > 0)?.text || '',
+        negativeText: params.textPrompt.find((x) => x.weight < 0)?.text,
+        referenceImage: params.initImage,
+        colors: params.colors!,
+      },
+      imageGenerationConfig: body.imageGenerationConfig,
+    };
+  } else if (params.taskType === 'BACKGROUND_REMOVAL') {
+    body = {
+      taskType: 'BACKGROUND_REMOVAL',
+      backgroundRemovalParams: {
+        image: params.initImage!,
+      },
+    };
+  } else if (body.textToImageParams) {
+    // TEXT_IMAGE タスクタイプの拡張(Image Conditioning)
+    body.textToImageParams = {
+      ...body.textToImageParams,
+      conditionImage: params.initImage,
+      controlMode: params.controlMode,
+      controlStrength: params.controlStrength,
+    };
   }
+
   return JSON.stringify(body);
 };
 
@@ -753,7 +798,7 @@ export const BEDROCK_IMAGE_GEN_MODELS: {
     extractOutputImage: extractOutputImageTitanImage,
   },
   'amazon.titan-image-generator-v2:0': {
-    createBodyImage: createBodyImageTitanImage,
+    createBodyImage: createBodyImageTitanImageV2,
     extractOutputImage: extractOutputImageTitanImage,
   },
 };

packages/types/src/image.d.ts

@@ -1,5 +1,20 @@
+export type ControlMode = 'CANNY_EDGE' | 'SEGMENTATION';
+export type BaseGenerationMode =
+  | 'TEXT_IMAGE'
+  | 'IMAGE_VARIATION'
+  | 'INPAINTING'
+  | 'OUTPAINTING';
+export type TitanImageV2GenerationMode =
+  | 'IMAGE_CONDITIONING'
+  | 'COLOR_GUIDED_GENERATION'
+  | 'BACKGROUND_REMOVAL';
+export type GenerationMode = BaseGenerationMode | TitanImageV2GenerationMode;
 // 標準化したパラメータ
 export type GenerateImageParams = {
+  taskType?:
+    | BaseGenerationMode
+    | 'COLOR_GUIDED_GENERATION'
+    | 'BACKGROUND_REMOVAL';
   textPrompt: {
     text: string;
     weight: number;
@@ -16,7 +31,11 @@ export type GenerateImageParams = {
   // Inpaint / Outpaint
   maskImage?: string;
   maskPrompt?: string;
-  maskMode?: 'INPAINTING' | 'OUTPAINTING';
+  // Color Guided Generation
+  colors?: string[];
+  // Image Conditioning
+  controlStrength?: number;
+  controlMode?: ControlMode;
 };
 
 // Stable Diffusion
@@ -47,7 +66,7 @@ export type StableDiffusionParams = {
 // Titan Image
 // https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-image.html
 export type TitanImageParams = {
-  taskType: 'TEXT_IMAGE' | 'INPAINTING' | 'OUTPAINTING' | 'IMAGE_VARIATION';
+  taskType: BaseGenerationMode;
   textToImageParams?: {
     text: string;
     negativeText?: string;
@@ -83,6 +102,27 @@ export type TitanImageParams = {
   };
 };
 
+export type TitanImageV2Params = Omit<TitanImageParams, 'taskType'> & {
+  taskType:
+    | BaseGenerationMode
+    | 'COLOR_GUIDED_GENERATION'
+    | 'BACKGROUND_REMOVAL';
+  textToImageParams?: TitanImageParams['textToImageParams'] & {
+    conditionImage?: string; // base64 encoded image
+    controlMode?: ControlMode;
+    controlStrength?: number;
+  };
+  colorGuidedGenerationParams?: {
+    text: string;
+    negativeText?: string;
+    referenceImage?: string; // base64 encoded image
+    colors: string[]; // list of color hex codes
+  };
+  backgroundRemovalParams?: {
+    image: string; // base64 encoded image
+  };
+};
+
 export type BedrockImageGenerationResponse = {
   result: string;
   artifacts: {

packages/web/src/components/Select.tsx

@@ -3,6 +3,7 @@ import { Listbox, Transition } from '@headlessui/react';
 import { PiCaretUpDown, PiCheck, PiX } from 'react-icons/pi';
 import RowItem, { RowItemProps } from './RowItem';
 import ButtonIcon from './ButtonIcon';
+import Help from './Help';
 
 type Props = RowItemProps & {
   label?: string;
@@ -11,34 +12,71 @@ type Props = RowItemProps & {
     value: string;
     label: string;
   }[];
+  help?: string;
   clearable?: boolean;
   fullWidth?: boolean;
+  colorchip?: boolean;
   onChange: (value: string) => void;
 };
 
 const Select: React.FC<Props> = (props) => {
-  const selectedLabel = useMemo(() => {
+  const selectedOption = useMemo(() => {
     return props.value === ''
-      ? ''
-      : props.options.filter((o) => o.value === props.value)[0].label;
+      ? null
+      : props.options.find((o) => o.value === props.value);
   }, [props.options, props.value]);
 
   const onClear = useCallback(() => {
     props.onChange('');
   }, [props]);
 
+  const renderColorChips = (value: string) => {
+    const colors = value.split(',').map((color) => color.trim());
+    return (
+      <div className="flex items-center">
+        {colors.map((color, index) => (
+          <div
+            key={index}
+            style={{
+              backgroundColor: color,
+              width: '16px',
+              height: '16px',
+              marginRight: '4px',
+              border: '1px solid #ccc',
+            }}
+          />
+        ))}
+      </div>
+    );
+  };
+
+  const renderOptionContent = (option: { value: string; label: string }) => {
+    if (props.colorchip) {
+      return (
+        <div className="flex items-center">
+          {renderColorChips(option.value)}
+          <span className="mr-2">{option.label}</span>
+        </div>
+      );
+    }
+    return option.label;
+  };
+
   return (
     <RowItem notItem={props.notItem} className="relative">
       {props.label && (
-        <div>
+        <div className="flex items-center">
           <span className="text-sm">{props.label}</span>
+          {props.help && <Help className="ml-1" message={props.help} />}
         </div>
       )}
       <Listbox value={props.value} onChange={props.onChange}>
         <div className="relative">
           <Listbox.Button
             className={`relative h-8 cursor-pointer rounded border border-black/30 bg-white pl-3 pr-10 text-left focus:outline-none ${props.fullWidth ? 'w-full' : 'w-fit'}`}>
-            <span className="block truncate">{selectedLabel}</span>
+            <span className="block truncate">
+              {selectedOption ? renderOptionContent(selectedOption) : ''}
+            </span>
 
             <span className="pointer-events-none absolute inset-y-0 right-0 flex items-center pr-2">
               <PiCaretUpDown className="text-sm" />
@@ -73,7 +111,7 @@ const Select: React.FC<Props> = (props) => {
                       className={`block truncate ${
                         selected ? 'font-medium' : 'font-normal'
                       }`}>
-                      {option.label}
+                      {renderOptionContent(option)}
                     </span>
                     {selected ? (
                       <span className="text-aws-smile absolute inset-y-0 left-0 flex items-center pl-3">