diff --git a/.github/workflows/docker/compose/image2image-compose-cd.yaml b/.github/workflows/docker/compose/image2image-compose-cd.yaml new file mode 100644 index 000000000..c2db4d40a --- /dev/null +++ b/.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} diff --git a/comps/__init__.py b/comps/__init__.py index a6ef41f65..34cfe8d18 100644 --- a/comps/__init__.py +++ b/comps/__init__.py @@ -27,6 +27,7 @@ VideoPath, ImageDoc, SDInputs, + SDImg2ImgInputs, SDOutputs, TextImageDoc, MultimodalDoc, diff --git a/comps/cores/mega/constants.py b/comps/cores/mega/constants.py index 0917bd46c..a0523daba 100644 --- a/comps/cores/mega/constants.py +++ b/comps/cores/mega/constants.py @@ -32,6 +32,7 @@ class ServiceType(Enum): IMAGE2VIDEO = 15 TEXT2IMAGE = 16 ANIMATION = 17 + IMAGE2IMAGE = 18 class MegaServiceEndpoint(Enum): diff --git a/comps/cores/proto/docarray.py b/comps/cores/proto/docarray.py index ee06017a5..71b6f15ec 100644 --- a/comps/cores/proto/docarray.py +++ b/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 diff --git a/comps/image2image/Dockerfile b/comps/image2image/Dockerfile new file mode 100644 index 000000000..3e2cce5e3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/comps/image2image/Dockerfile.intel_hpu b/comps/image2image/Dockerfile.intel_hpu new file mode 100644 index 000000000..a47d9dfb0 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/comps/image2image/README.md b/comps/image2image/README.md new file mode 100644 index 000000000..79d20949b --- /dev/null +++ b/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= +``` + +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' +``` diff --git a/comps/image2image/__init__.py b/comps/image2image/__init__.py new file mode 100644 index 000000000..916f3a44b --- /dev/null +++ b/comps/image2image/__init__.py @@ -0,0 +1,2 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 diff --git a/comps/image2image/image2image.py b/comps/image2image/image2image.py new file mode 100644 index 000000000..36e0cbd4f --- /dev/null +++ b/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() diff --git a/comps/image2image/requirements.txt b/comps/image2image/requirements.txt new file mode 100644 index 000000000..f639b4ba1 --- /dev/null +++ b/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 diff --git a/tests/image2image/test_image2image.sh b/tests/image2image/test_image2image.sh new file mode 100644 index 000000000..2b8883f11 --- /dev/null +++ b/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