From 558a78de55a40dfde9cff147c00be0fd03660108 Mon Sep 17 00:00:00 2001 From: Etienne Perot Date: Fri, 8 Mar 2024 13:18:43 -0800 Subject: [PATCH] Automated rollback of changelist 614015218 PiperOrigin-RevId: 614026488 --- Makefile | 4 +- .../gpu/stable-diffusion-xl/Dockerfile.x86_64 | 35 -- .../download_checkpoints.py | 38 --- images/gpu/stable-diffusion-xl/generate_image | 45 --- .../gpu/stable-diffusion-xl/generate_image.py | 304 ------------------ test/gpu/BUILD | 12 - test/gpu/imagegen_test.go | 57 ---- test/gpu/stablediffusion/BUILD | 17 - test/gpu/stablediffusion/stablediffusion.go | 179 ----------- 9 files changed, 1 insertion(+), 690 deletions(-) delete mode 100644 images/gpu/stable-diffusion-xl/Dockerfile.x86_64 delete mode 100644 images/gpu/stable-diffusion-xl/download_checkpoints.py delete mode 100644 images/gpu/stable-diffusion-xl/generate_image delete mode 100644 images/gpu/stable-diffusion-xl/generate_image.py delete mode 100644 test/gpu/imagegen_test.go delete mode 100644 test/gpu/stablediffusion/BUILD delete mode 100644 test/gpu/stablediffusion/stablediffusion.go diff --git a/Makefile b/Makefile index ff4579456..c12919476 100644 --- a/Makefile +++ b/Makefile @@ -298,14 +298,13 @@ 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 load-gpu_stablediffusionxl +gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_python .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 @@ -313,7 +312,6 @@ 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 deleted file mode 100644 index 9660a4c03..000000000 --- a/images/gpu/stable-diffusion-xl/Dockerfile.x86_64 +++ /dev/null @@ -1,35 +0,0 @@ -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 deleted file mode 100644 index cef81e55f..000000000 --- a/images/gpu/stable-diffusion-xl/download_checkpoints.py +++ /dev/null @@ -1,38 +0,0 @@ -#!/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 deleted file mode 100644 index eb7e08639..000000000 --- a/images/gpu/stable-diffusion-xl/generate_image +++ /dev/null @@ -1,45 +0,0 @@ -#!/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 deleted file mode 100644 index 3dbd61c0c..000000000 --- a/images/gpu/stable-diffusion-xl/generate_image.py +++ /dev/null @@ -1,304 +0,0 @@ -#!/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 747a23e7f..4b302ca3e 100644 --- a/test/gpu/BUILD +++ b/test/gpu/BUILD @@ -61,15 +61,3 @@ 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 deleted file mode 100644 index ab7995bc8..000000000 --- a/test/gpu/imagegen_test.go +++ /dev/null @@ -1,57 +0,0 @@ -// 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 deleted file mode 100644 index fb098f6ca..000000000 --- a/test/gpu/stablediffusion/BUILD +++ /dev/null @@ -1,17 +0,0 @@ -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 deleted file mode 100644 index 7d1c2c90f..000000000 --- a/test/gpu/stablediffusion/stablediffusion.go +++ /dev/null @@ -1,179 +0,0 @@ -// 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 -}