fix test.

Signed-off-by: Ye, Xinyu <[email protected]>
opea-project · Sep 27, 2024 · 618cb81 · 618cb81
1 parent d4ba8c7
commit 618cb81
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/comps/text2image/text2image.py b/comps/text2image/text2image.py
@@ -49,14 +49,14 @@ def initialize():
                         "token": args.token,
                     }
                 )
-                # if "stable-diffusion-3" in args.model_name_or_path:
-                #     from optimum.habana.diffusers import GaudiStableDiffusion3Pipeline
+                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,
-                #     )
-                if "stable-diffusion" in args.model_name_or_path.lower() or "flux" in args.model_name_or_path.lower():
+                    pipe = GaudiStableDiffusion3Pipeline.from_pretrained(
+                        args.model_name_or_path,
+                        **kwargs,
+                    )
+                elif "stable-diffusion" in args.model_name_or_path.lower() or "flux" in args.model_name_or_path.lower():
                     from optimum.habana.diffusers import AutoPipelineForText2Image
 
                     pipe = AutoPipelineForText2Image.from_pretrained(

diff --git a/tests/text2image/test_text2image.sh b/tests/text2image/test_text2image.sh
@@ -10,7 +10,7 @@ ip_address=$(hostname -I | awk '{print $1}')
 function build_docker_images() {
     cd $WORKPATH
     echo $(pwd)
-    docker build --no-cache -t opea/text2image:latest --build-arg MODEL=stabilityai/stable-diffusion-xl-base-1.0 -f comps/text2image/Dockerfile .
+    docker build --no-cache -t opea/text2image:latest -f comps/text2image/Dockerfile .
     if [ $? -ne 0 ]; then
         echo "opea/text2image built fail"
         exit 1
@@ -21,7 +21,7 @@ function build_docker_images() {
 
 function start_service() {
     unset http_proxy
-    docker run -d --name="test-comps-text2image" -e http_proxy=$http_proxy -e https_proxy=$https_proxy -p 9379:9379 --ipc=host opea/text2image:latest
+    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
     sleep 30s
 }