diff --git a/generative_ai/imagen/edit_image_product_image.py b/generative_ai/imagen/edit_image_product_image.py new file mode 100644 index 00000000000..212b50d6ea6 --- /dev/null +++ b/generative_ai/imagen/edit_image_product_image.py @@ -0,0 +1,104 @@ +# 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. + +"""Google Cloud Vertex AI sample for editing a product image. You can + modify the background content but preserve the product's appearance. +Example usage: + python edit_image_product_image.py --project_id \ + --location --input_file \ + --output_file --prompt +""" + +# [START generativeaionvertexai_imagen_edit_image_product_image] + +import argparse + +import vertexai +from vertexai.preview.vision_models import Image, ImageGenerationModel + + +def edit_image_product_image( + project_id: str, + location: str, + input_file: str, + output_file: str, + prompt: str, +) -> vertexai.preview.vision_models.ImageGenerationResponse: + """Edit a local image by modifying the background content. + Args: + project_id: Google Cloud project ID, used to initialize Vertex AI. + location: Google Cloud region, used to initialize Vertex AI. + input_file: Local path to the input image file. Image can be in PNG or JPEG format. + output_file: Local path to the output image file. + prompt: The text prompt describing what you want to see in the background.""" + + vertexai.init(project=project_id, location=location) + + model = ImageGenerationModel.from_pretrained("imagegeneration@006") + base_img = Image.load_from_file(location=input_file) + + images = model.edit_image( + base_image=base_img, + prompt=prompt, + edit_mode="product-image", + # Optional parameters + number_of_images=1, + # Allows for product repositioning: "reposition" (default; moves product + # to center), "fixed" (keeps product in original position) + # product_position="fixed", + ) + + images[0].save(location=output_file) + + # Optional. View the edited image in a notebook. + # images[0].show() + + print(f"Created output image using {len(images[0]._image_bytes)} bytes") + + return images + + +# [END generativeaionvertexai_imagen_edit_image_product_image] + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--project_id", help="Your Cloud project ID.", required=True) + parser.add_argument( + "--location", + help="The location in which to initialize Vertex AI.", + default="us-central1", + ) + parser.add_argument( + "--input_file", + help="The local path to the input file (e.g., 'my-input.png').", + required=True, + ) + parser.add_argument( + "--output_file", + help="The local path to the output file (e.g., 'my-output.png').", + required=True, + ) + parser.add_argument( + "--prompt", + help="The text prompt describing what you want to insert into the background.", + required=True, + ) + args = parser.parse_args() + edit_image_product_image( + args.project_id, + args.location, + args.input_file, + args.output_file, + args.prompt, + ) diff --git a/generative_ai/imagen/edit_image_product_image_test.py b/generative_ai/imagen/edit_image_product_image_test.py new file mode 100644 index 00000000000..43e2976dd21 --- /dev/null +++ b/generative_ai/imagen/edit_image_product_image_test.py @@ -0,0 +1,42 @@ +# 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. + +import os + +import backoff + +import edit_image_product_image + +from google.api_core.exceptions import ResourceExhausted + + +_RESOURCES = os.path.join(os.path.dirname(__file__), "test_resources") +_PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") +_LOCATION = "us-central1" +_INPUT_FILE = os.path.join(_RESOURCES, "pillow.png") +_OUTPUT_FILE = os.path.join(_RESOURCES, "pillow_on_beach.png") +_PROMPT = "beach" + + +@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=60) +def test_edit_image_product_image() -> None: + response = edit_image_product_image.edit_image_product_image( + _PROJECT_ID, + _LOCATION, + _INPUT_FILE, + _OUTPUT_FILE, + _PROMPT, + ) + + assert len(response[0]._image_bytes) > 1000 diff --git a/generative_ai/imagen/test_resources/cat.png b/generative_ai/imagen/test_resources/cat.png index bd7d6612073..67f2b55a6f4 100644 Binary files a/generative_ai/imagen/test_resources/cat.png and b/generative_ai/imagen/test_resources/cat.png differ diff --git a/generative_ai/imagen/test_resources/dog.png b/generative_ai/imagen/test_resources/dog.png index 3e176130099..a5040ca8f75 100644 Binary files a/generative_ai/imagen/test_resources/dog.png and b/generative_ai/imagen/test_resources/dog.png differ diff --git a/generative_ai/imagen/test_resources/dog_book.png b/generative_ai/imagen/test_resources/dog_book.png index a49a9311fcc..a71a7bbf858 100644 Binary files a/generative_ai/imagen/test_resources/dog_book.png and b/generative_ai/imagen/test_resources/dog_book.png differ diff --git a/generative_ai/imagen/test_resources/dog_newspaper.png b/generative_ai/imagen/test_resources/dog_newspaper.png index f54414076dc..ed69902f9c6 100644 Binary files a/generative_ai/imagen/test_resources/dog_newspaper.png and b/generative_ai/imagen/test_resources/dog_newspaper.png differ diff --git a/generative_ai/imagen/test_resources/dog_newspaper_mask.png b/generative_ai/imagen/test_resources/dog_newspaper_mask.png index a0edd713d75..d0cceae3a06 100644 Binary files a/generative_ai/imagen/test_resources/dog_newspaper_mask.png and b/generative_ai/imagen/test_resources/dog_newspaper_mask.png differ diff --git a/generative_ai/imagen/test_resources/pillow.png b/generative_ai/imagen/test_resources/pillow.png new file mode 100644 index 00000000000..8803f49defe Binary files /dev/null and b/generative_ai/imagen/test_resources/pillow.png differ diff --git a/generative_ai/imagen/test_resources/pillow_on_beach.png b/generative_ai/imagen/test_resources/pillow_on_beach.png new file mode 100644 index 00000000000..630d987dae9 Binary files /dev/null and b/generative_ai/imagen/test_resources/pillow_on_beach.png differ diff --git a/generative_ai/requirements.txt b/generative_ai/requirements.txt index 975c55934e1..b143b8465fb 100644 --- a/generative_ai/requirements.txt +++ b/generative_ai/requirements.txt @@ -2,6 +2,6 @@ pandas==1.3.5; python_version == '3.7' pandas==2.0.1; python_version > '3.7' pillow==9.5.0; python_version < '3.8' pillow==10.0.1; python_version >= '3.8' -google-cloud-aiplatform[pipelines]==1.42.1 +google-cloud-aiplatform[pipelines]==1.46.0 google-auth==2.17.3 anthropic[vertex]==0.21.3