mirror of
https://github.com/netbirdio/gvisor.git
synced 2026-05-22 17:12:49 -07:00
Automated rollback of changelist 614015218
PiperOrigin-RevId: 614026488
This commit is contained in:
Etienne Perot
committed by
gVisor bot
gVisor bot
parent
5f6dade838
commit
558a78de55
@@ -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
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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"
|
||||
@@ -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}')
|
||||
@@ -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"],
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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",
|
||||
],
|
||||
)
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user