diff --git a/Makefile b/Makefile index 5ed3b08a3..4e4b83d17 100644 --- a/Makefile +++ b/Makefile @@ -298,13 +298,14 @@ 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_stable-diffusion-xl .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 @@ -312,6 +313,7 @@ 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 diff --git a/images/gpu/stable-diffusion-xl/Dockerfile.x86_64 b/images/gpu/stable-diffusion-xl/Dockerfile.x86_64 new file mode 100644 index 000000000..9660a4c03 --- /dev/null +++ b/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"] diff --git a/images/gpu/stable-diffusion-xl/download_checkpoints.py b/images/gpu/stable-diffusion-xl/download_checkpoints.py new file mode 100644 index 000000000..cef81e55f --- /dev/null +++ b/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", +) diff --git a/images/gpu/stable-diffusion-xl/generate_image b/images/gpu/stable-diffusion-xl/generate_image new file mode 100644 index 000000000..eb7e08639 --- /dev/null +++ b/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" diff --git a/images/gpu/stable-diffusion-xl/generate_image.py b/images/gpu/stable-diffusion-xl/generate_image.py new file mode 100644 index 000000000..3dbd61c0c --- /dev/null +++ b/images/gpu/stable-diffusion-xl/generate_image.py @@ -0,0 +1,304 @@ +#!/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. + +"""Generate image with Stable Diffusion XL. + +Images are written to stdout by wrapper script. +""" + +import argparse +import array +import base64 +import datetime +import enum +import fcntl +import io +import json +import os +import subprocess +import termios + +import diffusers +import torch + + +# Define arguments. +class Format(enum.Enum): + """Output format enum.""" + + PNG = 'PNG' + JPEG = 'JPEG' + ASCII = 'ASCII' + BRAILLE = 'BRAILLE' + PNG_BASE64 = 'PNG-BASE64' + METRICS = 'METRICS' + + @property + def is_terminal_output(self): + return self in (Format.ASCII, Format.BRAILLE, Format.METRICS) + + def __str__(self): + return self.value + + +parser = argparse.ArgumentParser( + prog='generate_image', + description='Generate an image using Stable Diffusion XL', +) + +# Arguments passed by wrapper script. +parser.add_argument('--out', required=True, type=str, help=argparse.SUPPRESS) +parser.add_argument('--terminal_pixel_width', type=str, help=argparse.SUPPRESS) +parser.add_argument('--terminal_pixel_height', type=str, help=argparse.SUPPRESS) + +parser.add_argument( + '--quiet_stderr', + action='store_true', + help=( + 'Suppress PyTorch messages to stderr; useful if stderr output is' + ' captured.' + ), +) +parser.add_argument( + '--enable_model_cpu_offload', + action='store_true', + help='Offload non-main components of model to CPU if low on GPU VRAM', +) +parser.add_argument( + '--format', + type=Format, + choices=list(Format), + default=Format.BRAILLE, + help='Output file format: ' + ', '.join(str(v) for v in Format), +) +parser.add_argument( + '--steps', default=50, type=int, help='Number of diffusion steps' +) +parser.add_argument( + '--noise_frac', default=0.8, type=float, help='Noise fraction' +) +parser.add_argument( + '--enable_refiner', + action='store_true', + help='Use the refiner model on top of the base model for better results', +) +parser.add_argument( + '--warm', + action='store_true', + help='Generate the image twice; timing metrics will measure both images', +) +parser.add_argument('prompt', type=str, help='Prompt to generate image') +args = parser.parse_args() + +# Load base model. +time_start = datetime.datetime.now(datetime.timezone.utc) +base = diffusers.DiffusionPipeline.from_pretrained( + 'stabilityai/stable-diffusion-xl-base-1.0', + torch_dtype=torch.float16, + variant='fp16', + use_safetensors=True, +) +if args.enable_model_cpu_offload: + base.enable_model_cpu_offload() +else: + base.to('cuda') +base.unet = torch.compile(base.unet, mode='reduce-overhead', fullgraph=True) + +# Load refiner model if enabled. +refiner = None +if args.enable_refiner: + 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', + ) + if args.enable_model_cpu_offload: + refiner.enable_model_cpu_offload() + else: + refiner.to('cuda') + refiner.unet = torch.compile( + refiner.unet, mode='reduce-overhead', fullgraph=True + ) + +# Set the prompt. +default_prompt = ( + 'Photorealistic image of two androids playing chess aboard a spaceship' +) +if args.format.is_terminal_output: + # If displaying in a terminal, cartoony pictures that have sharp edges will + # look much clearer than photorealistic pictures. + default_prompt = 'A boring flat corporate logo that says "gVisor"' +prompt = args.prompt or default_prompt + + +# Generate image. +def generate_image(): + """Run the base model and maybe the refiner model to generate the image.""" + + time_start_image = datetime.datetime.now(datetime.timezone.utc) + if not args.enable_refiner: + img = base( + prompt=prompt, + num_inference_steps=args.steps, + output_type='pil', + ).images[0] + time_base_done = datetime.datetime.now(datetime.timezone.utc) + time_refiner_done = None + else: + base_images = base( + prompt=prompt, + num_inference_steps=args.steps, + denoising_end=args.noise_frac, + output_type='latent', + ).images + time_base_done = datetime.datetime.now(datetime.timezone.utc) + img = refiner( + prompt=prompt, + num_inference_steps=args.steps, + denoising_start=args.noise_frac, + image=base_images, + ).images[0] + time_refiner_done = datetime.datetime.now(datetime.timezone.utc) + return img, time_start_image, time_base_done, time_refiner_done + + +image, cold_start_image, cold_base_done, cold_refiner_done = generate_image() +warm_start_image, warm_base_done, warm_refiner_done = None, None, None +if args.warm: + image, warm_start_image, warm_base_done, warm_refiner_done = generate_image() + + +def get_optimal_terminal_width(): + """Returns the width of the terminal for ASCII image display.""" + try: + terminal_width, terminal_height = os.get_terminal_size() + except OSError: # Not a TTY, return a sane default. + return 80 + if terminal_width == 0 or terminal_height == 0: # Incoherent terminal size. + return 80 + if terminal_width <= 42: + # Ridiculously small terminal, return default dimension anyway because + # whatever we do won't look nice regardless. + return 80 + # Try to find the aspect ratio of a single terminal character. + terminal_pixel_width = 0 + terminal_pixel_height = 0 + if args.terminal_pixel_width.isdigit(): + terminal_pixel_width = int(args.terminal_pixel_width) + if args.terminal_pixel_height.isdigit(): + terminal_pixel_height = int(args.terminal_pixel_height) + if terminal_pixel_width == 0 or terminal_pixel_height == 0: + termios_buf = array.array('H', [0, 0, 0, 0]) + fcntl.ioctl(1, termios.TIOCGWINSZ, termios_buf) + _, _, terminal_pixel_width, terminal_pixel_height = termios_buf + if terminal_pixel_width != 0 and terminal_pixel_height != 0: + character_width = float(terminal_pixel_width) / float(terminal_width) + character_height = float(terminal_pixel_height) / float(terminal_height) + character_aspect_ratio = character_width / character_height + else: + character_aspect_ratio = 0.5 # Just use a sane default. + adjusted_terminal_height = float(terminal_height) / float( + character_aspect_ratio + ) + image_width, image_height = image.size + width_ratio = float(image_width) / float(terminal_width) + height_ratio = float(image_height) / adjusted_terminal_height + if width_ratio > height_ratio: + # Width is determining factor. + return terminal_width + # Height is the determining factor. + final_width = int( + adjusted_terminal_height * float(image_width) / float(image_height) + ) + # Remove one just to not make it take literally the entire console, and + # in case our estimation for things like character size is wrong. + final_width -= 1 + if final_width < 8: + # Very vertical image, most likely text? So it's OK if it scrolls. + # Return a sane default width. + return 42 + return final_width + + +# Save image in desired format. +if args.format in (Format.PNG, Format.JPEG): + image.save(args.out, args.format) +else: + buf = io.BytesIO() + image.save(buf, format=Format.PNG) + image_bytes = buf.getvalue() + time_done = datetime.datetime.now(datetime.timezone.utc) + with open(args.out, 'wb') as f: + if args.format == Format.PNG_BASE64: + f.write(base64.standard_b64encode(image_bytes)) + elif args.format.is_terminal_output: + image_converter_args = [ + '/usr/bin/ascii-image-converter', + '/dev/stdin', + '--width=%d' % (get_optimal_terminal_width(),), + ] + if args.format in (Format.BRAILLE, Format.METRICS): + image_converter_args.extend(('--braille', '--dither')) + else: + image_converter_args.append('--complex') + image_ascii = subprocess.run( + image_converter_args, + input=image_bytes, + capture_output=True, + check=True, + timeout=60, + ) + if args.format == Format.METRICS: + split_lines = lambda x: [ + x[i : i + 1024] for i in range(0, len(x), 1024) + ] + results = { + 'image_ascii_base64': split_lines( + base64.standard_b64encode(image_ascii.stdout).decode('ascii') + ), + 'image_png_base64': split_lines( + base64.standard_b64encode(image_bytes).decode('ascii') + ), + } + for name, timestamp in ( + ('start', time_start), + ('cold_start_image', cold_start_image), + ('cold_base_done', cold_base_done), + ('cold_refiner_done', cold_refiner_done), + ('warm_start_image', warm_start_image), + ('warm_base_done', warm_base_done), + ('warm_refiner_done', warm_refiner_done), + ('done', time_done), + ): + results[name] = ( + timestamp.isoformat() if timestamp is not None else None + ) + # Python's `json` module always outputs strings, not bytes, so + # we cannot directly dump to `f`. Output to string instead, then + # encode. + # Also, `json.dumps` doesn't add a trailing newline, so we do. + results_json = ( + json.dumps(results, sort_keys=True, ensure_ascii=True, indent=2) + + '\n' + ) + f.write(results_json.encode('ascii')) + else: + f.write(image_ascii.stdout) + else: + raise ValueError(f'Unknown format: {args.format}') diff --git a/test/gpu/BUILD b/test/gpu/BUILD index 4b302ca3e..747a23e7f 100644 --- a/test/gpu/BUILD +++ b/test/gpu/BUILD @@ -61,3 +61,15 @@ go_test( "//pkg/test/testutil", ], ) + +go_test( + name = "imagegen_test", + srcs = ["imagegen_test.go"], + tags = [ + "local", + "noguitar", + "notap", + ], + visibility = ["//:sandbox"], + deps = ["//test/gpu/stablediffusion"], +) diff --git a/test/gpu/imagegen_test.go b/test/gpu/imagegen_test.go new file mode 100644 index 000000000..ab7995bc8 --- /dev/null +++ b/test/gpu/imagegen_test.go @@ -0,0 +1,57 @@ +// 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. + +// Package imagegen_test runs Stable Diffusion and generates images with it. +package imagegen_test + +import ( + "context" + "testing" + "time" + + "gvisor.dev/gvisor/test/gpu/stablediffusion" +) + +// TestStableDiffusionXL generates an image with Stable Diffusion XL. +func TestStableDiffusionXL(t *testing.T) { + ctx := context.Background() + sdxl := stablediffusion.NewDockerXL(t) + generateCtx, generateCancel := context.WithTimeout(ctx, 15*time.Minute) + defer generateCancel() + image, err := sdxl.Generate(generateCtx, &stablediffusion.XLPrompt{ + Query: `A boring flat corporate logo that says "gVisor"`, + AllowCPUOffload: false, + UseRefiner: false, + NoiseFraction: 0.8, + // This is just a test to make sure Stable Diffusion works at all, + // so we don't need a lot of steps here: + Steps: 8, + }) + if err != nil { + t.Fatalf("Cannot generate image with Stable Diffusion XL: %v", err) + } + img, err := image.Image() + if err != nil { + t.Fatalf("Cannot decode image: %v", err) + } + size := img.Bounds().Size() + if size.X <= 0 || size.Y <= 0 { + t.Fatalf("Generated image has invalid size: %dx%d", size.X, size.Y) + } + ascii, err := image.ASCII() + if err != nil { + t.Fatalf("Cannot convert image to ASCII: %v", err) + } + t.Logf("Generated image (size %dx%d pixels):\n%s\n", size.X, size.Y, ascii) +} diff --git a/test/gpu/stablediffusion/BUILD b/test/gpu/stablediffusion/BUILD new file mode 100644 index 000000000..fb098f6ca --- /dev/null +++ b/test/gpu/stablediffusion/BUILD @@ -0,0 +1,17 @@ +load("//tools:defs.bzl", "go_library") + +package( + default_applicable_licenses = ["//:license"], + licenses = ["notice"], +) + +go_library( + name = "stablediffusion", + testonly = 1, + srcs = ["stablediffusion.go"], + visibility = ["//:sandbox"], + deps = [ + "//pkg/test/dockerutil", + "//pkg/test/testutil", + ], +) diff --git a/test/gpu/stablediffusion/stablediffusion.go b/test/gpu/stablediffusion/stablediffusion.go new file mode 100644 index 000000000..7d1c2c90f --- /dev/null +++ b/test/gpu/stablediffusion/stablediffusion.go @@ -0,0 +1,179 @@ +// 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. + +// Package stablediffusion provides utilities to generate images with +// Stable Diffusion. +package stablediffusion + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "image" + "image/png" + "strings" + "time" + + "gvisor.dev/gvisor/pkg/test/dockerutil" + "gvisor.dev/gvisor/pkg/test/testutil" +) + +// ContainerRunner is an interface to run containers. +type ContainerRunner interface { + // Run runs a container with the given image and arguments to completion, + // and returns its combined output as a byte string. + Run(ctx context.Context, image string, argv []string) ([]byte, error) +} + +// dockerRunner runs Docker containers on the local machine. +type dockerRunner struct { + logger testutil.Logger +} + +// Run implements `ContainerRunner.Run`. +func (dr *dockerRunner) Run(ctx context.Context, image string, argv []string) ([]byte, error) { + cont := dockerutil.MakeContainer(ctx, dr.logger) + defer cont.CleanUp(ctx) + opts := dockerutil.GPURunOpts() + opts.Image = image + if err := cont.Spawn(ctx, opts, argv...); err != nil { + return nil, fmt.Errorf("could not start Stable Diffusion container: %v", err) + } + waitErr := cont.Wait(ctx) + logs, logsErr := cont.Logs(ctx) + if waitErr != nil { + if logsErr == nil { + return nil, fmt.Errorf("container exited with error: %v; logs: %v", waitErr, logs) + } + return nil, fmt.Errorf("container exited with error: %v (cannot get logs: %v)", waitErr, logsErr) + } + if logsErr != nil { + return nil, fmt.Errorf("could not get container logs: %v", logsErr) + } + return []byte(logs), nil +} + +// XL generates images using Stable Diffusion XL. +type XL struct { + image string + runner ContainerRunner +} + +// NewXL returns a new Stable Diffusion XL generator. +func NewXL(sdxlImage string, runner ContainerRunner) *XL { + return &XL{ + image: sdxlImage, + runner: runner, + } +} + +// NewDockerXL returns a new Stable Diffusion XL generator using Docker +// containers on the local machine. +func NewDockerXL(logger testutil.Logger) *XL { + return NewXL("gpu/stable-diffusion-xl", &dockerRunner{logger: logger}) +} + +// XLPrompt is the input to Stable Diffusion XL to generate an image. +type XLPrompt struct { + // Query is the text query to generate the image with. + Query string + + // AllowCPUOffload is whether to allow offloading parts of the model to CPU. + AllowCPUOffload bool + + // UseRefiner is whether to use the refiner model after the base model. + // This takes more VRAM and more time but produces a better image. + UseRefiner bool + + // NoiseFraction is the fraction of noise to seed the image with. + // Must be between 0.0 and 1.0 inclusively. + NoiseFraction float64 + + // Steps is the number of diffusion steps to run for the base and refiner + // models. More steps generally means sharper results but more time to + // generate the image. A reasonable value is between 30 and 50. + Steps int + + // Warm controls whether the image will be generated while the model is + // warm. This will double the running time, as the image will still be + // generated with a cold model first. + Warm bool +} + +// xlImageJSON is the JSON response from the Stable Diffusion XL +// container's generate_image.py. +// Warm* fields are only present when `XLPrompt.Warm` is set. +type xlImageJSON struct { + ImageASCIIBase64 []string `json:"image_ascii_base64"` + ImagePNGBase64 []string `json:"image_png_base64"` + Start time.Time `json:"start"` + ColdStartImage time.Time `json:"cold_start_image"` + ColdBaseDone time.Time `json:"cold_base_done"` + ColdRefinerDone time.Time `json:"cold_refiner_done"` + WarmStartImage time.Time `json:"warm_start_image"` + WarmBaseDone time.Time `json:"warm_base_done"` + WarmRefinerDone time.Time `json:"warm_refiner_done"` + Done time.Time `json:"done"` +} + +// XLImage is an image generated by Stable Diffusion XL. +type XLImage struct { + Prompt *XLPrompt + data xlImageJSON +} + +// ASCII returns an ASCII version of the generated image. +func (i *XLImage) ASCII() (string, error) { + ascii, err := base64.StdEncoding.DecodeString(strings.Join(i.data.ImageASCIIBase64, "")) + if err != nil { + return "", fmt.Errorf("invalid base64: %w", err) + } + return string(ascii), nil +} + +// Image returns the generated image. +func (i *XLImage) Image() (image.Image, error) { + return png.Decode(base64.NewDecoder(base64.StdEncoding, bytes.NewBufferString(strings.Join(i.data.ImagePNGBase64, "")))) +} + +// Generate generates an image with Stable Diffusion XL. +func (xl *XL) Generate(ctx context.Context, prompt *XLPrompt) (*XLImage, error) { + argv := []string{ + "--format=METRICS", + fmt.Sprintf("--steps=%d", prompt.Steps), + fmt.Sprintf("--noise_frac=%f", prompt.NoiseFraction), + "--quiet_stderr", + } + if prompt.AllowCPUOffload { + argv = append(argv, "--enable_model_cpu_offload") + } + if prompt.UseRefiner { + argv = append(argv, "--enable_refiner") + } + if prompt.Warm { + argv = append(argv, "--warm") + } + argv = append(argv, prompt.Query) + output, err := xl.runner.Run(ctx, xl.image, argv) + if err != nil { + return nil, err + } + xlImage := &XLImage{Prompt: prompt} + if err := json.Unmarshal(output, &xlImage.data); err != nil { + return nil, fmt.Errorf("malformed JSON output %q: %w", string(output), err) + } + return xlImage, nil +}