Create Stable Diffusion XL Docker image and regression test.

PiperOrigin-RevId: 614015218

Author: EtiennePerot

Makefile

@@ -298,20 +298,22 @@ cos-gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN)
 # This is a superset of those needed for smoke tests.
 # It includes non-GPU images that are used as part of GPU tests,
 # e.g. busybox and python.
-gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_python
+gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_python load-gpu_stablediffusionxl
 .PHONY: gpu-images
 
 gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true --nvproxy-docker=true)
 	@$(call sudo,test/gpu:pytorch_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
+	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-all-tests
 
 cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true)
 	@$(call sudo,test/gpu:pytorch_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
+	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-all-tests
 

images/gpu/stable-diffusion-xl/Dockerfile.x86_64

@@ -0,0 +1,35 @@
+FROM nvidia/cuda:12.3.1-devel-ubuntu22.04
+
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install --yes \
+      python3 \
+      python3-distutils \
+      python3-pip \
+      clang \
+      wget \
+      vim \
+      git \
+      libgl1 \
+      libglib2.0-0 \
+      libgl1-mesa-glx \
+      golang
+
+RUN python3 -m pip install --ignore-installed \
+      diffusers \
+      transformers \
+      accelerate \
+      xformers \
+      invisible-watermark
+
+RUN go install \
+      github.com/TheZoraiz/ascii-image-converter@d05a757c5e02ab23e97b6f6fca4e1fbeb10ab559 && \
+      mv "$HOME/go/bin/ascii-image-converter" /usr/bin/
+
+COPY download_checkpoints.py /tmp
+RUN chmod +x /tmp/download_checkpoints.py && \
+      /tmp/download_checkpoints.py && \
+      rm /tmp/download_checkpoints.py
+
+COPY generate_image generate_image.py /
+RUN chmod 555 /generate_image /generate_image.py
+ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+ENTRYPOINT ["/generate_image"]

images/gpu/stable-diffusion-xl/download_checkpoints.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+# Copyright 2024 The gVisor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Download Stable Diffusion XL checkpoints from Hugging Face."""
+
+import diffusers
+import torch
+
+# Download base model.
+base = diffusers.DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True,
+)
+
+# Download refiner model.
+refiner = diffusers.DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=base.text_encoder_2,
+    vae=base.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)

images/gpu/stable-diffusion-xl/generate_image

@@ -0,0 +1,45 @@
+#!/bin/bash
+
+set -euo pipefail
+
+quiet_stderr=false
+for arg; do
+  if [[ "$arg" == '--out' ]] || echo "$arg" | grep -qE '^--out='; then
+    echo 'Cannot specify --out parameter; the image file will be written to stdout.' >&2
+    exit 1
+  fi
+  if [[ "$arg" == '--quiet_stderr' ]]; then
+    quiet_stderr=true
+  fi
+done
+
+# Try to find out pixel size of the shell.
+terminal_pixel_width=0
+terminal_pixel_height=0
+if [[ -t 1 ]]; then
+  echo -e -n '\e[14t'; IFS=';' read -rs -t 0.5 -d 't' rest height width <$(tty)
+  terminal_pixel_width="$width"
+  terminal_pixel_height="$height"
+fi
+
+out_dir="$(mktemp -d)"
+
+set +e
+  /generate_image.py \
+    --out="$out_dir/out_image" \
+    --terminal_pixel_width="$terminal_pixel_width" \
+    --terminal_pixel_height="$terminal_pixel_height" \
+    "$@" \
+    1>/dev/null \
+    2>"$out_dir/stderr"
+  return_code="$?"
+set -e
+
+if [[ "$return_code" == 0 ]]; then
+  cat "$out_dir/out_image"
+fi
+if [[ "$return_code" != 0 ]] || [[ "$quiet_stderr" == false ]]; then
+  cat "$out_dir/stderr" >&2
+fi
+rm -rf "$out_dir"
+exit "$return_code"