From 0c4b053f560e245452e6148d53b29034075ad7cc Mon Sep 17 00:00:00 2001 From: Siddharth Venkatesan Date: Fri, 9 Dec 2022 13:09:09 -0800 Subject: [PATCH] Update SD handler to work with custom wheel --- engines/python/setup/djl_python/stable-diffusion.py | 3 +-- serving/docker/deepspeed.Dockerfile | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/engines/python/setup/djl_python/stable-diffusion.py b/engines/python/setup/djl_python/stable-diffusion.py index 4291c5082..d896b0c0d 100644 --- a/engines/python/setup/djl_python/stable-diffusion.py +++ b/engines/python/setup/djl_python/stable-diffusion.py @@ -66,11 +66,10 @@ def initialize(self, properties: dict): kwargs["torch_dtype"] = torch.float16 kwargs["revision"] = "fp16" - torch.set_grad_enabled(False) pipeline = DiffusionPipeline.from_pretrained(self.model_id, **kwargs) pipeline.to(f"cuda:{self.device}") deepspeed.init_distributed() - engine = deepspeed.init_inference(getattr(pipeline, "model", pipeline), config=self.ds_config) + engine = deepspeed.init_inference(getattr(pipeline, "model", pipeline), **self.ds_config) if hasattr(pipeline, "model"): pipeline.model = engine diff --git a/serving/docker/deepspeed.Dockerfile b/serving/docker/deepspeed.Dockerfile index c0b041c0f..4ceb4dc86 100644 --- a/serving/docker/deepspeed.Dockerfile +++ b/serving/docker/deepspeed.Dockerfile @@ -47,7 +47,7 @@ RUN apt-get update && \ pip3 install torch==${torch_version} --extra-index-url https://download.pytorch.org/whl/cu116 && \ pip3 install ${deepspeed_wheel} && \ pip3 install transformers==${transformers_version} && \ - pip3 install triton==1.0.0 mpi4py sentencepiece accelerate==${accelerate_version} bitsandbytes && \ + pip3 install triton==2.0.0.dev20221202 mpi4py sentencepiece accelerate==${accelerate_version} bitsandbytes && \ pip3 install diffusers[torch]==${diffusers_version} && \ scripts/patch_oss_dlc.sh python && \ scripts/security_patch.sh deepspeed && \