Skip to content

Commit 7d08835

Browse files
XinyuYe-Intelpre-commit-ci[bot]
authored and
ctao456
committed
Add stable diffusion microservice (opea-project#729)
* add stable diffusion microservice. Signed-off-by: Ye, Xinyu <xinyu.ye@intel.com> * added test. Signed-off-by: Ye, Xinyu <xinyu.ye@intel.com> * changed output to images bytes data Signed-off-by: Ye, Xinyu <xinyu.ye@intel.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * unified inference and wrapper into one microservice. Signed-off-by: Ye, Xinyu <xinyu.ye@intel.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix test. Signed-off-by: Ye, Xinyu <xinyu.ye@intel.com> --------- Signed-off-by: Ye, Xinyu <xinyu.ye@intel.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Chun Tao <chun.tao@intel.com>
1 parent fbb493b commit 7d08835

10 files changed

Lines changed: 374 additions & 0 deletions

File tree

‎.github/workflows/docker/compose/text2image-compose-cd.yaml‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# Copyright (C) 2024 Intel Corporation
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
# this file should be run in the root of the repo
5+
# images used by GenAIExamples: text2image,text2image-gaudi
6+
services:
7+
text2image:
8+
build:
9+
dockerfile: comps/text2image/Dockerfile
10+
image: ${REGISTRY:-opea}/text2image:${TAG:-latest}
11+
text2image-gaudi:
12+
build:
13+
dockerfile: comps/text2image/Dockerfile.intel_hpu
14+
image: ${REGISTRY:-opea}/text2image-gaudi:${TAG:-latest}

‎comps/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
ImagesPath,
2727
VideoPath,
2828
ImageDoc,
29+
SDInputs,
30+
SDOutputs,
2931
TextImageDoc,
3032
MultimodalDoc,
3133
EmbedMultimodalDoc,

‎comps/cores/proto/docarray.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,15 @@ class LVMVideoDoc(BaseDoc):
284284
max_new_tokens: conint(ge=0, le=1024) = 512
285285

286286

287+
class SDInputs(BaseDoc):
288+
prompt: str
289+
num_images_per_prompt: int = 1
290+
291+
292+
class SDOutputs(BaseDoc):
293+
images: list
294+
295+
287296
class ImagePath(BaseDoc):
288297
image_path: str
289298

‎comps/text2image/Dockerfile‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Copyright (C) 2024 Intel Corporation
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
FROM python:3.11-slim
5+
6+
# Set environment variables
7+
ENV LANG=en_US.UTF-8
8+
9+
ARG ARCH="cpu"
10+
11+
COPY comps /home/comps
12+
13+
RUN pip install --no-cache-dir --upgrade pip && \
14+
if [ ${ARCH} = "cpu" ]; then pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu; fi && \
15+
pip install --no-cache-dir -r /home/comps/text2image/requirements.txt
16+
17+
ENV PYTHONPATH=$PYTHONPATH:/home
18+
19+
WORKDIR /home/comps/text2image
20+
21+
RUN echo python text2image.py --bf16 >> run.sh
22+
23+
CMD bash run.sh

‎comps/text2image/Dockerfile.intel_hpu‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# Copyright (C) 2024 Intel Corporation
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
# HABANA environment
5+
# FROM vault.habana.ai/gaudi-docker/1.16.1/ubuntu22.04/habanalabs/pytorch-installer-2.2.2:latest as hpu
6+
FROM opea/habanalabs:1.16.1-pytorch-installer-2.2.2 as hpu
7+
RUN useradd -m -s /bin/bash user && \
8+
mkdir -p /home/user && \
9+
chown -R user /home/user/
10+
11+
COPY comps /home/user/comps
12+
13+
RUN chown -R user /home/user/comps/text2image
14+
15+
RUN rm -rf /etc/ssh/ssh_host*
16+
USER user
17+
# Set environment variables
18+
ENV LANG=en_US.UTF-8
19+
ENV PYTHONPATH=/home/user:/usr/lib/habanalabs/:/home/user/optimum-habana
20+
21+
# Install requirements and optimum habana
22+
RUN pip install --no-cache-dir --upgrade pip && \
23+
pip install --no-cache-dir -r /home/user/comps/text2image/requirements.txt && \
24+
pip install --no-cache-dir optimum[habana]
25+
26+
WORKDIR /home/user/comps/text2image
27+
28+
RUN echo python text2image.py --device hpu --use_hpu_graphs --bf16 >> run.sh
29+
30+
CMD bash run.sh

‎comps/text2image/README.md‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
# Text-to-Image Microservice
2+
3+
Text-to-Image is a task that generate image conditioning on the provided text. This microservice supports text-to-image task by using Stable Diffusion (SD) model.
4+
5+
# 🚀1. Start Microservice with Python (Option 1)
6+
7+
## 1.1 Install Requirements
8+
9+
```bash
10+
pip install -r requirements.txt
11+
```
12+
13+
## 1.2 Start Text-to-Image Microservice
14+
15+
Select Stable Diffusion (SD) model and assign its name to a environment variable as below:
16+
17+
```bash
18+
# SD1.5
19+
export MODEL=stable-diffusion-v1-5/stable-diffusion-v1-5
20+
# SD2.1
21+
export MODEL=stabilityai/stable-diffusion-2-1
22+
# SDXL
23+
export MODEL=stabilityai/stable-diffusion-xl-base-1.0
24+
# SD3
25+
export MODEL=stabilityai/stable-diffusion-3-medium-diffusers
26+
```
27+
28+
Set huggingface token:
29+
30+
```bash
31+
export HF_TOKEN=<your huggingface token>
32+
```
33+
34+
Start the OPEA Microservice:
35+
36+
```bash
37+
python text2image.py --bf16 --model_name_or_path $MODEL --token $HF_TOKEN
38+
```
39+
40+
# 🚀2. Start Microservice with Docker (Option 2)
41+
42+
## 2.1 Build Images
43+
44+
Select Stable Diffusion (SD) model and assign its name to a environment variable as below:
45+
46+
```bash
47+
# SD1.5
48+
export MODEL=stable-diffusion-v1-5/stable-diffusion-v1-5
49+
# SD2.1
50+
export MODEL=stabilityai/stable-diffusion-2-1
51+
# SDXL
52+
export MODEL=stabilityai/stable-diffusion-xl-base-1.0
53+
# SD3
54+
export MODEL=stabilityai/stable-diffusion-3-medium-diffusers
55+
```
56+
57+
### 2.1.1 Text-to-Image Service Image on Xeon
58+
59+
Build text-to-image service image on Xeon with below command:
60+
61+
```bash
62+
cd ../..
63+
docker build -t opea/text2image:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/text2image/Dockerfile .
64+
```
65+
66+
### 2.1.2 Text-to-Image Service Image on Gaudi
67+
68+
Build text-to-image service image on Gaudi with below command:
69+
70+
```bash
71+
cd ../..
72+
docker build -t opea/text2image-gaudi:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/text2image/Dockerfile.intel_hpu .
73+
```
74+
75+
## 2.2 Start Text-to-Image Service
76+
77+
### 2.2.1 Start Text-to-Image Service on Xeon
78+
79+
Start text-to-image service on Xeon with below command:
80+
81+
```bash
82+
docker run --ipc=host -p 9379:9379 -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e HF_TOKEN=$HF_TOKEN -e MODEL=$MODEL opea/text2image:latest
83+
```
84+
85+
### 2.2.2 Start Text-to-Image Service on Gaudi
86+
87+
Start text-to-image service on Gaudi with below command:
88+
89+
```bash
90+
docker run -p 9379:9379 --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/text2image-gaudi:latest
91+
```
92+
93+
# 3 Test Text-to-Image Service
94+
95+
```bash
96+
http_proxy="" curl http://localhost:9379/v1/text2image -XPOST -d '{"prompt":"An astronaut riding a green horse", "num_images_per_prompt":1}' -H 'Content-Type: application/json'
97+
```

‎comps/text2image/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
# Copyright (C) 2024 Intel Corporation
2+
# SPDX-License-Identifier: Apache-2.0

‎comps/text2image/requirements.txt‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
accelerate
2+
datasets
3+
diffusers
4+
docarray[full]
5+
fastapi
6+
opentelemetry-api
7+
opentelemetry-exporter-otlp
8+
opentelemetry-sdk
9+
prometheus-fastapi-instrumentator
10+
pydantic==2.7.2
11+
pydub
12+
shortuuid
13+
torch
14+
transformers
15+
uvicorn

‎comps/text2image/text2image.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
# Copyright (C) 2024 Intel Corporation
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
import argparse
5+
import base64
6+
import os
7+
import threading
8+
import time
9+
10+
import torch
11+
from diffusers import DiffusionPipeline
12+
13+
from comps import (
14+
CustomLogger,
15+
SDInputs,
16+
SDOutputs,
17+
ServiceType,
18+
opea_microservices,
19+
register_microservice,
20+
register_statistics,
21+
statistics_dict,
22+
)
23+
24+
logger = CustomLogger("text2image")
25+
pipe = None
26+
args = None
27+
initialization_lock = threading.Lock()
28+
initialized = False
29+
30+
31+
def initialize():
32+
global pipe, args, initialized
33+
with initialization_lock:
34+
if not initialized:
35+
# initialize model and tokenizer
36+
if os.getenv("MODEL", None):
37+
args.model_name_or_path = os.getenv("MODEL")
38+
kwargs = {}
39+
if args.bf16:
40+
kwargs["torch_dtype"] = torch.bfloat16
41+
if not args.token:
42+
args.token = os.getenv("HF_TOKEN")
43+
if args.device == "hpu":
44+
kwargs.update(
45+
{
46+
"use_habana": True,
47+
"use_hpu_graphs": args.use_hpu_graphs,
48+
"gaudi_config": "Habana/stable-diffusion",
49+
"token": args.token,
50+
}
51+
)
52+
if "stable-diffusion-3" in args.model_name_or_path:
53+
from optimum.habana.diffusers import GaudiStableDiffusion3Pipeline
54+
55+
pipe = GaudiStableDiffusion3Pipeline.from_pretrained(
56+
args.model_name_or_path,
57+
**kwargs,
58+
)
59+
elif "stable-diffusion" in args.model_name_or_path.lower() or "flux" in args.model_name_or_path.lower():
60+
from optimum.habana.diffusers import AutoPipelineForText2Image
61+
62+
pipe = AutoPipelineForText2Image.from_pretrained(
63+
args.model_name_or_path,
64+
**kwargs,
65+
)
66+
else:
67+
raise NotImplementedError(
68+
"Only support stable-diffusion, stable-diffusion-xl, stable-diffusion-3 and flux now, "
69+
+ f"model {args.model_name_or_path} not supported."
70+
)
71+
elif args.device == "cpu":
72+
pipe = DiffusionPipeline.from_pretrained(args.model_name_or_path, token=args.token, **kwargs)
73+
else:
74+
raise NotImplementedError(f"Only support cpu and hpu device now, device {args.device} not supported.")
75+
logger.info("Stable Diffusion model initialized.")
76+
initialized = True
77+
78+
79+
@register_microservice(
80+
name="opea_service@text2image",
81+
service_type=ServiceType.TEXT2IMAGE,
82+
endpoint="/v1/text2image",
83+
host="0.0.0.0",
84+
port=9379,
85+
input_datatype=SDInputs,
86+
output_datatype=SDOutputs,
87+
)
88+
@register_statistics(names=["opea_service@text2image"])
89+
def text2image(input: SDInputs):
90+
initialize()
91+
start = time.time()
92+
prompt = input.prompt
93+
num_images_per_prompt = input.num_images_per_prompt
94+
95+
generator = torch.manual_seed(args.seed)
96+
images = pipe(prompt, generator=generator, num_images_per_prompt=num_images_per_prompt).images
97+
image_path = os.path.join(os.getcwd(), prompt.strip().replace(" ", "_").replace("/", ""))
98+
os.makedirs(image_path, exist_ok=True)
99+
results = []
100+
for i, image in enumerate(images):
101+
save_path = os.path.join(image_path, f"image_{i+1}.png")
102+
image.save(save_path)
103+
with open(save_path, "rb") as f:
104+
bytes = f.read()
105+
b64_str = base64.b64encode(bytes).decode()
106+
results.append(b64_str)
107+
statistics_dict["opea_service@text2image"].append_latency(time.time() - start, None)
108+
return SDOutputs(images=results)
109+
110+
111+
if __name__ == "__main__":
112+
parser = argparse.ArgumentParser()
113+
parser.add_argument("--model_name_or_path", type=str, default="stabilityai/stable-diffusion-3-medium-diffusers")
114+
parser.add_argument("--use_hpu_graphs", default=False, action="store_true")
115+
parser.add_argument("--device", type=str, default="cpu")
116+
parser.add_argument("--token", type=str, default=None)
117+
parser.add_argument("--seed", type=int, default=42)
118+
parser.add_argument("--bf16", action="store_true")
119+
120+
args = parser.parse_args()
121+
122+
logger.info("Text2image server started.")
123+
opea_microservices["opea_service@text2image"].start()

‎tests/text2image/test_text2image.sh‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
#!/bin/bash
2+
# Copyright (C) 2024 Intel Corporation
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
set -x
6+
7+
WORKPATH=$(dirname "$PWD")
8+
ip_address=$(hostname -I | awk '{print $1}')
9+
10+
function build_docker_images() {
11+
cd $WORKPATH
12+
echo $(pwd)
13+
docker build --no-cache -t opea/text2image:latest -f comps/text2image/Dockerfile .
14+
if [ $? -ne 0 ]; then
15+
echo "opea/text2image built fail"
16+
exit 1
17+
else
18+
echo "opea/text2image built successful"
19+
fi
20+
}
21+
22+
function start_service() {
23+
unset http_proxy
24+
docker run -d --name="test-comps-text2image" -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e MODEL=stabilityai/stable-diffusion-xl-base-1.0 -p 9379:9379 --ipc=host opea/text2image:latest
25+
sleep 30s
26+
}
27+
28+
function validate_microservice() {
29+
result=$(http_proxy="" curl http://localhost:9379/v1/text2image -XPOST -d '{"prompt":"An astronaut riding a green horse", "num_images_per_prompt":1}' -H 'Content-Type: application/json')
30+
if [[ $result == *"images"* ]]; then
31+
echo "Result correct."
32+
else
33+
echo "Result wrong."
34+
docker logs test-comps-text2image
35+
exit 1
36+
fi
37+
38+
}
39+
40+
function stop_docker() {
41+
cid=$(docker ps -aq --filter "name=test-comps-text2image*")
42+
if [[ ! -z "$cid" ]]; then docker stop $cid && docker rm $cid && sleep 1s; fi
43+
}
44+
45+
function main() {
46+
47+
stop_docker
48+
49+
build_docker_images
50+
start_service
51+
52+
validate_microservice
53+
54+
stop_docker
55+
echo y | docker system prune
56+
57+
}
58+
59+
main

0 commit comments

Comments
 (0)