Add image2image microservice (opea-project#794)

* added image2image microservice.

Signed-off-by: Ye, Xinyu <[email protected]>

* added UT.

Signed-off-by: Ye, Xinyu <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Ye, Xinyu <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/workflows/docker/compose/image2image-compose-cd.yaml

@@ -0,0 +1,14 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# this file should be run in the root of the repo
+# images used by GenAIExamples: image2image,image2image-gaudi
+services:
+  image2image:
+    build:
+      dockerfile: comps/image2image/Dockerfile
+    image: ${REGISTRY:-opea}/image2image:${TAG:-latest}
+  image2image-gaudi:
+    build:
+      dockerfile: comps/image2image/Dockerfile.intel_hpu
+    image: ${REGISTRY:-opea}/image2image-gaudi:${TAG:-latest}

comps/__init__.py

@@ -27,6 +27,7 @@
     VideoPath,
     ImageDoc,
     SDInputs,
+    SDImg2ImgInputs,
     SDOutputs,
     TextImageDoc,
     MultimodalDoc,

comps/cores/mega/constants.py

@@ -32,6 +32,7 @@ class ServiceType(Enum):
     IMAGE2VIDEO = 15
     TEXT2IMAGE = 16
     ANIMATION = 17
+    IMAGE2IMAGE = 18
 
 
 class MegaServiceEndpoint(Enum):

comps/cores/proto/docarray.py

@@ -291,6 +291,12 @@ class SDInputs(BaseDoc):
     num_images_per_prompt: int = 1
 
 
+class SDImg2ImgInputs(BaseDoc):
+    image: str
+    prompt: str = ""
+    num_images_per_prompt: int = 1
+
+
 class SDOutputs(BaseDoc):
     images: list
 

comps/image2image/Dockerfile

@@ -0,0 +1,23 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+FROM python:3.11-slim
+
+# Set environment variables
+ENV LANG=en_US.UTF-8
+
+ARG ARCH="cpu"
+
+COPY comps /home/comps
+
+RUN pip install --no-cache-dir --upgrade pip && \
+    if [ ${ARCH} = "cpu" ]; then pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu; fi && \
+    pip install --no-cache-dir -r /home/comps/image2image/requirements.txt
+
+ENV PYTHONPATH=$PYTHONPATH:/home
+
+WORKDIR /home/comps/image2image
+
+RUN echo python image2image.py --bf16 >> run.sh
+
+CMD bash run.sh

comps/image2image/Dockerfile.intel_hpu

@@ -0,0 +1,30 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# HABANA environment
+# FROM vault.habana.ai/gaudi-docker/1.16.1/ubuntu22.04/habanalabs/pytorch-installer-2.2.2:latest as hpu
+FROM opea/habanalabs:1.16.1-pytorch-installer-2.2.2 as hpu
+RUN useradd -m -s /bin/bash user && \
+    mkdir -p /home/user && \
+    chown -R user /home/user/
+
+COPY comps /home/user/comps
+
+RUN chown -R user /home/user/comps/image2image
+
+RUN rm -rf /etc/ssh/ssh_host*
+USER user
+# Set environment variables
+ENV LANG=en_US.UTF-8
+ENV PYTHONPATH=/home/user:/usr/lib/habanalabs/:/home/user/optimum-habana
+
+# Install requirements and optimum habana
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r /home/user/comps/image2image/requirements.txt && \
+    pip install --no-cache-dir optimum[habana]
+
+WORKDIR /home/user/comps/image2image
+
+RUN echo python image2image.py --device hpu --use_hpu_graphs --bf16 >> run.sh
+
+CMD bash run.sh

comps/image2image/README.md

@@ -0,0 +1,85 @@
+# Image-to-Image Microservice
+
+Image-to-Image is a task that generate image conditioning on the provided image and text. This microservice supports image-to-image task by using Stable Diffusion (SD) model.
+
+# 🚀1. Start Microservice with Python (Option 1)
+
+## 1.1 Install Requirements
+
+```bash
+pip install -r requirements.txt
+```
+
+## 1.2 Start Image-to-Image Microservice
+
+Select Stable Diffusion (SD) model and assign its name to a environment variable as below:
+
+```bash
+# SDXL
+export MODEL=stabilityai/stable-diffusion-xl-refiner-1.0
+```
+
+Set huggingface token:
+
+```bash
+export HF_TOKEN=<your huggingface token>
+```
+
+Start the OPEA Microservice:
+
+```bash
+python image2image.py --bf16 --model_name_or_path $MODEL --token $HF_TOKEN
+```
+
+# 🚀2. Start Microservice with Docker (Option 2)
+
+## 2.1 Build Images
+
+Select Stable Diffusion (SD) model and assign its name to a environment variable as below:
+
+```bash
+# SDXL
+export MODEL=stabilityai/stable-diffusion-xl-refiner-1.0
+```
+
+### 2.1.1 Image-to-Image Service Image on Xeon
+
+Build image-to-image service image on Xeon with below command:
+
+```bash
+cd ../..
+docker build -t opea/image2image:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/image2image/Dockerfile .
+```
+
+### 2.1.2 Image-to-Image Service Image on Gaudi
+
+Build image-to-image service image on Gaudi with below command:
+
+```bash
+cd ../..
+docker build -t opea/image2image-gaudi:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/image2image/Dockerfile.intel_hpu .
+```
+
+## 2.2 Start Image-to-Image Service
+
+### 2.2.1 Start Image-to-Image Service on Xeon
+
+Start image-to-image service on Xeon with below command:
+
+```bash
+docker run --ipc=host -p 9389:9389 -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e HF_TOKEN=$HF_TOKEN -e MODEL=$MODEL opea/image2image:latest
+```
+
+### 2.2.2 Start Image-to-Image Service on Gaudi
+
+Start image-to-image service on Gaudi with below command:
+
+```bash
+docker run -p 9389:9389 --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --ipc=host -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e HF_TOKEN=$HF_TOKEN -e MODEL=$MODEL opea/image2image-gaudi:latest
+```
+
+# 3 Test Image-to-Image Service
+
+```bash
+http_proxy="" curl http://localhost:9389/v1/image2image -XPOST -d '{"image": "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png", "prompt":"a photo of an astronaut riding a horse on mars", "num_images_per_prompt":1}' -H 'Content-Type: application/json'
+```

comps/image2image/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0

comps/image2image/image2image.py

@@ -0,0 +1,117 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import base64
+import os
+import threading
+import time
+
+import torch
+from diffusers import AutoPipelineForImage2Image
+from diffusers.utils import load_image
+
+from comps import (
+    CustomLogger,
+    SDImg2ImgInputs,
+    SDOutputs,
+    ServiceType,
+    opea_microservices,
+    register_microservice,
+    register_statistics,
+    statistics_dict,
+)
+
+logger = CustomLogger("image2image")
+pipe = None
+args = None
+initialization_lock = threading.Lock()
+initialized = False
+
+
+def initialize():
+    global pipe, args, initialized
+    with initialization_lock:
+        if not initialized:
+            # initialize model and tokenizer
+            if os.getenv("MODEL", None):
+                args.model_name_or_path = os.getenv("MODEL")
+            kwargs = {}
+            if args.bf16:
+                kwargs["torch_dtype"] = torch.bfloat16
+            if not args.token:
+                args.token = os.getenv("HF_TOKEN")
+            if args.device == "hpu":
+                kwargs.update(
+                    {
+                        "use_habana": True,
+                        "use_hpu_graphs": args.use_hpu_graphs,
+                        "gaudi_config": "Habana/stable-diffusion",
+                        "token": args.token,
+                    }
+                )
+                if "stable-diffusion-xl" in args.model_name_or_path:
+                    from optimum.habana.diffusers import GaudiStableDiffusionXLImg2ImgPipeline
+
+                    pipe = GaudiStableDiffusionXLImg2ImgPipeline.from_pretrained(
+                        args.model_name_or_path,
+                        **kwargs,
+                    )
+                else:
+                    raise NotImplementedError(
+                        "Only support stable-diffusion-xl now, " + f"model {args.model_name_or_path} not supported."
+                    )
+            elif args.device == "cpu":
+                pipe = AutoPipelineForImage2Image.from_pretrained(args.model_name_or_path, token=args.token, **kwargs)
+            else:
+                raise NotImplementedError(f"Only support cpu and hpu device now, device {args.device} not supported.")
+            logger.info("Stable Diffusion model initialized.")
+            initialized = True
+
+
+@register_microservice(
+    name="opea_service@image2image",
+    service_type=ServiceType.IMAGE2IMAGE,
+    endpoint="/v1/image2image",
+    host="0.0.0.0",
+    port=9389,
+    input_datatype=SDImg2ImgInputs,
+    output_datatype=SDOutputs,
+)
+@register_statistics(names=["opea_service@image2image"])
+def image2image(input: SDImg2ImgInputs):
+    initialize()
+    start = time.time()
+    image = load_image(input.image).convert("RGB")
+    prompt = input.prompt
+    num_images_per_prompt = input.num_images_per_prompt
+
+    generator = torch.manual_seed(args.seed)
+    images = pipe(image=image, prompt=prompt, generator=generator, num_images_per_prompt=num_images_per_prompt).images
+    image_path = os.path.join(os.getcwd(), prompt.strip().replace(" ", "_").replace("/", ""))
+    os.makedirs(image_path, exist_ok=True)
+    results = []
+    for i, image in enumerate(images):
+        save_path = os.path.join(image_path, f"image_{i+1}.png")
+        image.save(save_path)
+        with open(save_path, "rb") as f:
+            bytes = f.read()
+        b64_str = base64.b64encode(bytes).decode()
+        results.append(b64_str)
+    statistics_dict["opea_service@image2image"].append_latency(time.time() - start, None)
+    return SDOutputs(images=results)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name_or_path", type=str, default="stabilityai/stable-diffusion-xl-refiner-1.0")
+    parser.add_argument("--use_hpu_graphs", default=False, action="store_true")
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument("--token", type=str, default=None)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--bf16", action="store_true")
+
+    args = parser.parse_args()
+
+    logger.info("Image2image server started.")
+    opea_microservices["opea_service@image2image"].start()

comps/image2image/requirements.txt

@@ -0,0 +1,15 @@
+accelerate
+datasets
+diffusers
+docarray[full]
+fastapi
+opentelemetry-api
+opentelemetry-exporter-otlp
+opentelemetry-sdk
+prometheus-fastapi-instrumentator
+pydantic==2.7.2
+pydub
+shortuuid
+torch
+transformers
+uvicorn

tests/image2image/test_image2image.sh

@@ -0,0 +1,59 @@
+#!/bin/bash
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+set -x
+
+WORKPATH=$(dirname "$PWD")
+ip_address=$(hostname -I | awk '{print $1}')
+
+function build_docker_images() {
+    cd $WORKPATH
+    echo $(pwd)
+    docker build --no-cache -t opea/image2image:latest -f comps/image2image/Dockerfile .
+    if [ $? -ne 0 ]; then
+        echo "opea/image2image built fail"
+        exit 1
+    else
+        echo "opea/image2image built successful"
+    fi
+}
+
+function start_service() {
+    unset http_proxy
+    docker run -d --name="test-comps-image2image" -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e MODEL=stabilityai/stable-diffusion-xl-refiner-1.0 -p 9389:9389 --ipc=host opea/image2image:latest
+    sleep 30s
+}
+
+function validate_microservice() {
+    result=$(http_proxy="" curl http://localhost:9389/v1/image2image -XPOST -d '{"image": "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png", "prompt":"a photo of an astronaut riding a horse on mars", "num_images_per_prompt":1}' -H 'Content-Type: application/json')
+    if [[ $result == *"images"* ]]; then
+        echo "Result correct."
+    else
+        echo "Result wrong."
+        docker logs test-comps-image2image
+        exit 1
+    fi
+
+}
+
+function stop_docker() {
+    cid=$(docker ps -aq --filter "name=test-comps-image2image*")
+    if [[ ! -z "$cid" ]]; then docker stop $cid && docker rm $cid && sleep 1s; fi
+}
+
+function main() {
+
+    stop_docker
+
+    build_docker_images
+    start_service
+
+    validate_microservice
+
+    stop_docker
+    echo y | docker system prune
+
+}
+
+main