Add stable diffusion microservice

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

@@ -0,0 +1,18 @@
+# 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: text2image,sd,sd-gaudi
+services:
+  text2image:
+    build:
+      dockerfile: comps/text2image/Dockerfile
+    image: ${REGISTRY:-opea}/text2image:${TAG:-latest}
+  sd:
+    build:
+      dockerfile: comps/text2image/dependency/Dockerfile
+    image: ${REGISTRY:-opea}/sd:${TAG:-latest}
+  sd-gaudi:
+    build:
+      dockerfile: comps/text2image/dependency/Dockerfile.intel_hpu
+    image: ${REGISTRY:-opea}/sd-gaudi:${TAG:-latest}

comps/__init__.py

@@ -26,6 +26,8 @@
     ImagesPath,
     VideoPath,
     ImageDoc,
+    SDInputs,
+    SDOutputs,
     TextImageDoc,
     MultimodalDoc,
     EmbedMultimodalDoc,

comps/cores/proto/docarray.py

@@ -264,6 +264,15 @@ class LVMVideoDoc(BaseDoc):
     max_new_tokens: conint(ge=0, le=1024) = 512
 
 
+class SDInputs(BaseDoc):
+    prompt: str
+    num_images_per_prompt: int = 1
+
+
+class SDOutputs(BaseDoc):
+    images: list
+
+
 class ImagePath(BaseDoc):
     image_path: str
 

comps/text2image/Dockerfile

@@ -0,0 +1,18 @@
+# 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
+
+COPY comps /home/comps
+
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r /home/comps/text2image/requirements.txt
+
+ENV PYTHONPATH=$PYTHONPATH:/home
+
+WORKDIR /home/comps/text2image
+
+ENTRYPOINT ["python", "text2image.py"]

comps/text2image/README.md

@@ -0,0 +1,92 @@
+# Text-to-Image Microservice
+
+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.
+
+# 🚀1. Start Microservice with Python (Option 1)
+
+## 1.1 Install Requirements
+
+```bash
+pip install -r requirements.txt
+pip install -r dependency/requirements.txt
+```
+
+## 1.2 Start SD Service
+
+```bash
+# Start SD service
+cd dependency/
+python sd_server.py --token $HF_TOKEN
+```
+
+## 1.3 Start Text-to-Image Microservice
+
+```bash
+cd ..
+# Start the OPEA Microservice
+python text2image.py
+```
+
+# 🚀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
+# SD3
+export MODEL=stabilityai/stable-diffusion-3-medium-diffusers
+# SDXL
+export MODEL=stabilityai/stable-diffusion-xl-base-1.0
+```
+
+### 2.1.1 SD Server Image
+
+Build SD server image on Xeon with below command:
+
+```bash
+cd ../..
+docker build -t opea/sd:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy --build-arg MODEL=$MODEL -f comps/text2image/dependency/Dockerfile .
+```
+
+Build SD server image on Gaudi with below command:
+
+```bash
+cd ../..
+docker build -t opea/sd-gaudi:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy --build-arg MODEL=$MODEL -f comps/text2image/dependency/Dockerfile.intel_hpu .
+```
+
+### 2.1.2 Text-to-Image Service Image
+
+```bash
+docker build -t opea/text2image:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy -f comps/text2image/Dockerfile .
+```
+
+## 2.2 Start SD and Text-to-Image Service
+
+### 2.2.1 Start SD server
+
+Start SD server on Xeon with below command:
+
+```bash
+docker run --ipc=host -p 9378:9378 -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e HF_TOKEN=$HF_TOKEN opea/sd:latest
+```
+
+Start SD server on Gaudi with below command:
+
+```bash
+docker run -p 9378:9378 --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 opea/sd-gaudi:latest
+```
+
+### 2.2.2 Start Text-to-Image service
+
+```bash
+ip_address=$(hostname -I | awk '{print $1}')
+docker run -p 9379:9379 --ipc=host -e http_proxy=$http_proxy -e https_proxy=$https_proxy -e SD_ENDPOINT=http://$ip_address:9378 opea/text2image:latest
+```
+
+### 2.2.3 Test
+
+```bash
+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'
+```

comps/text2image/__init__.py

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

comps/text2image/dependency/Dockerfile

@@ -0,0 +1,25 @@
+# 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"
+
+ARG MODEL
+
+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/text2image/dependency/requirements.txt
+
+ENV PYTHONPATH=$PYTHONPATH:/home
+
+WORKDIR /home/comps/text2image/dependency
+
+RUN echo python sd_server.py --model_name_or_path $MODEL >> run.sh
+
+CMD bash run.sh

comps/text2image/dependency/Dockerfile.intel_hpu

@@ -0,0 +1,34 @@
+# 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/text2image
+
+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/:/optimum-habana
+
+ARG MODEL
+
+# Install requirements and optimum habana
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r /home/user/comps/text2image/dependency/requirements.txt && \
+    pip install --no-cache-dir optimum[habana]
+
+ENV PYTHONPATH=$PYTHONPATH:/home/user
+
+WORKDIR /home/user/comps/text2image/dependency
+
+RUN echo python sd_server.py --device hpu --bf16 --model_name_or_path $MODEL >> run.sh
+
+CMD bash run.sh

comps/text2image/dependency/requirements.txt

@@ -0,0 +1,6 @@
+accelerate
+diffusers
+fastapi
+torch
+transformers
+uvicorn

comps/text2image/dependency/sd_server.py

@@ -0,0 +1,97 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+"""Stand-alone Stable Diffusion FastAPI Server."""
+
+import argparse
+import base64
+import os
+import time
+
+import torch
+import uvicorn
+from diffusers import DiffusionPipeline
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, Response
+
+app = FastAPI()
+
+
+@app.post("/generate")
+async def generate(request: Request) -> Response:
+    print("SD generation begin.")
+    request_dict = await request.json()
+    prompt = request_dict.pop("prompt")
+    num_images_per_prompt = request_dict.pop("num_images_per_prompt", 1)
+
+    start = time.time()
+    generator = torch.manual_seed(args.seed)
+    images = pipe(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)
+    end = time.time()
+    print(f"SD Images output in {image_path}, time = {end-start}s")
+    return JSONResponse({"images": results})
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=9378)
+    parser.add_argument("--model_name_or_path", type=str, default="stabilityai/stable-diffusion-3-medium-diffusers")
+    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()
+    if not args.token:
+        args.token = os.getenv("HF_TOKEN")
+    if args.device == "hpu":
+        kwargs = {
+            "use_habana": True,
+            "use_hpu_graphs": args.use_hpu_graphs,
+            "gaudi_config": "Habana/stable-diffusion",
+            "token": args.token,
+        }
+        if args.bf16:
+            kwargs["torch_dtype"] = torch.bfloat16
+        if "stable-diffusion-3" in args.model_name_or_path:
+            from optimum.habana.diffusers import GaudiStableDiffusion3Pipeline
+
+            pipe = GaudiStableDiffusion3Pipeline.from_pretrained(
+                args.model_name_or_path,
+                **kwargs,
+            )
+        elif "stable-diffusion-xl" in args.model_name_or_path:
+            from optimum.habana.diffusers import GaudiStableDiffusionXLPipeline
+
+            pipe = GaudiStableDiffusionXLPipeline.from_pretrained(
+                args.model_name_or_path,
+                **kwargs,
+            )
+        else:
+            raise NotImplementedError(
+                "Only support stable-diffusion-3 and stable-diffusion-xl now, "
+                + f"model {args.model_name_or_path} not supported."
+            )
+    elif args.device == "cpu":
+        pipe = DiffusionPipeline.from_pretrained(args.model_name_or_path, token=args.token)
+    else:
+        raise NotImplementedError(f"Only support cpu and hpu device now, device {args.device} not supported.")
+    print("Stable Diffusion model initialized.")
+
+    uvicorn.run(
+        app,
+        host=args.host,
+        port=args.port,
+        log_level="debug",
+    )

comps/text2image/requirements.txt

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

comps/text2image/text2image.py

@@ -0,0 +1,46 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+
+import json
+import os
+import time
+
+import requests
+
+from comps import (
+    SDInputs,
+    SDOutputs,
+    ServiceType,
+    opea_microservices,
+    register_microservice,
+    register_statistics,
+    statistics_dict,
+)
+
+
+@register_microservice(
+    name="opea_service@text2image",
+    service_type=ServiceType.TEXT2IMAGE,
+    endpoint="/v1/text2image",
+    host="0.0.0.0",
+    port=9379,
+    input_datatype=SDInputs,
+    output_datatype=SDOutputs,
+)
+@register_statistics(names=["opea_service@text2image"])
+async def text2image(input: SDInputs):
+    start = time.time()
+    inputs = {"prompt": input.prompt, "num_images_per_prompt": input.num_images_per_prompt}
+    images = requests.post(url=f"{sd_endpoint}/generate", data=json.dumps(inputs), proxies={"http": None}).json()[
+        "images"
+    ]
+
+    statistics_dict["opea_service@text2image"].append_latency(time.time() - start, None)
+    return SDOutputs(images=images)
+
+
+if __name__ == "__main__":
+    sd_endpoint = os.getenv("SD_ENDPOINT", "http://localhost:9378")
+    print("Text2image server started.")
+    opea_microservices["opea_service@text2image"].start()