import itertools
import os
import random
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
import numpy as np
import safetensors.torch
import torch
import torch.nn.functional as F
import torchvision.transforms
import torchvision.transforms.functional as TF
from PIL import Image
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import default_collate
from transformers import (CLIPTextModel, CLIPTextModelWithProjection,
CLIPTokenizerFast)
from diffusion import (default_num_train_timesteps,
euler_ode_solver_diffusion_loop, make_sigmas)
from sdxl_models import (SDXLAdapter, SDXLControlNet, SDXLControlNetFull,
SDXLControlNetPreEncodedControlnetCond, SDXLUNet,
SDXLVae)
class SDXLTraining:
text_encoder_one: CLIPTextModel
text_encoder_two: CLIPTextModelWithProjection
vae: SDXLVae
sigmas: torch.Tensor
unet: SDXLUNet
adapter: Optional[SDXLAdapter]
controlnet: Optional[Union[SDXLControlNet, SDXLControlNetFull]]
train_unet: bool
train_unet_up_blocks: bool
mixed_precision: Optional[torch.dtype]
timestep_sampling: Literal["uniform", "cubic"]
validation_images_logged: bool
log_validation_input_images_every_time: bool
get_sdxl_conditioning_images: Callable[[Image.Image], Dict[str, Any]]
def __init__(
self,
device,
train_unet,
get_sdxl_conditioning_images,
train_unet_up_blocks=False,
unet_resume_from=None,
controlnet_cls=None,
controlnet_resume_from=None,
adapter_cls=None,
adapter_resume_from=None,
mixed_precision=None,
timestep_sampling="uniform",
log_validation_input_images_every_time=True,
):
self.text_encoder_one = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", variant="fp16", torch_dtype=torch.float16)
self.text_encoder_one.to(device=device)
self.text_encoder_one.requires_grad_(False)
self.text_encoder_one.eval()
self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2", variant="fp16", torch_dtype=torch.float16)
self.text_encoder_two.to(device=device)
self.text_encoder_two.requires_grad_(False)
self.text_encoder_two.eval()
self.vae = SDXLVae.load_fp16_fix(device=device)
self.vae.requires_grad_(False)
self.vae.eval()
self.sigmas = make_sigmas(device=device)
if train_unet:
if unet_resume_from is None:
self.unet = SDXLUNet.load_fp32(device=device)
else:
self.unet = SDXLUNet.load(unet_resume_from, device=device)
self.unet.requires_grad_(True)
self.unet.train()
self.unet = DDP(self.unet, device_ids=[device])
elif train_unet_up_blocks:
if unet_resume_from is None:
self.unet = SDXLUNet.load_fp32(device=device)
else:
self.unet = SDXLUNet.load_fp32(device=device, overrides=[unet_resume_from])
self.unet.requires_grad_(False)
self.unet.eval()
self.unet.up_blocks.requires_grad_(True)
self.unet.up_blocks.train()
self.unet = DDP(self.unet, device_ids=[device], find_unused_parameters=True)
else:
self.unet = SDXLUNet.load_fp16(device=device)
self.unet.requires_grad_(False)
self.unet.eval()
if controlnet_cls is not None:
if controlnet_resume_from is None:
self.controlnet = controlnet_cls.from_unet(self.unet)
self.controlnet.to(device)
else:
self.controlnet = controlnet_cls.load(controlnet_resume_from, device=device)
self.controlnet.train()
self.controlnet.requires_grad_(True)
# TODO add back
# controlnet.enable_gradient_checkpointing()
# TODO - should be able to remove find_unused_parameters. Comes from pre encoded controlnet
self.controlnet = DDP(self.controlnet, device_ids=[device], find_unused_parameters=True)
else:
self.controlnet = None
if adapter_cls is not None:
if adapter_resume_from is None:
self.adapter = adapter_cls()
self.adapter.to(device=device)
else:
self.adapter = adapter_cls.load(adapter_resume_from, device=device)
self.adapter.train()
self.adapter.requires_grad_(True)
self.adapter = DDP(self.adapter, device_ids=[device])
else:
self.adapter = None
self.mixed_precision = mixed_precision
self.timestep_sampling = timestep_sampling
self.validation_images_logged = False
self.log_validation_input_images_every_time = log_validation_input_images_every_time
self.get_sdxl_conditioning_images = get_sdxl_conditioning_images
self.train_unet = train_unet
self.train_unet_up_blocks = train_unet_up_blocks
def train_step(self, batch):
with torch.no_grad():
if isinstance(self.unet, DDP):
unet_dtype = self.unet.module.dtype
unet_device = self.unet.module.device
else:
unet_dtype = self.unet.dtype
unet_device = self.unet.device
micro_conditioning = batch["micro_conditioning"].to(device=unet_device)
image = batch["image"].to(self.vae.device, dtype=self.vae.dtype)
latents = self.vae.encode(image).to(dtype=unet_dtype)
text_input_ids_one = batch["text_input_ids_one"].to(self.text_encoder_one.device)
text_input_ids_two = batch["text_input_ids_two"].to(self.text_encoder_two.device)
encoder_hidden_states, pooled_encoder_hidden_states = sdxl_text_conditioning(self.text_encoder_one, self.text_encoder_two, text_input_ids_one, text_input_ids_two)
encoder_hidden_states = encoder_hidden_states.to(dtype=unet_dtype)
pooled_encoder_hidden_states = pooled_encoder_hidden_states.to(dtype=unet_dtype)
bsz = latents.shape[0]
if self.timestep_sampling == "uniform":
timesteps = torch.randint(0, default_num_train_timesteps, (bsz,), device=unet_device)
elif self.timestep_sampling == "cubic":
# Cubic sampling to sample a random timestep for each image
timesteps = torch.rand((bsz,), device=unet_device)
timesteps = (1 - timesteps**3) * default_num_train_timesteps
timesteps = timesteps.long()
timesteps = timesteps.clamp(0, default_num_train_timesteps - 1)
else:
assert False
sigmas_ = self.sigmas[timesteps].to(dtype=latents.dtype)
noise = torch.randn_like(latents)
noisy_latents = latents + noise * sigmas_
scaled_noisy_latents = noisy_latents / ((sigmas_**2 + 1) ** 0.5)
if "conditioning_image" in batch:
conditioning_image = batch["conditioning_image"].to(unet_device)
if self.controlnet is not None and isinstance(self.controlnet, SDXLControlNetPreEncodedControlnetCond):
controlnet_device = self.controlnet.module.device
controlnet_dtype = self.controlnet.module.dtype
conditioning_image = self.vae.encode(conditioning_image.to(self.vae.dtype)).to(device=controlnet_device, dtype=controlnet_dtype)
conditioning_image_mask = TF.resize(batch["conditioning_image_mask"], conditioning_image.shape[2:]).to(device=controlnet_device, dtype=controlnet_dtype)
conditioning_image = torch.concat((conditioning_image, conditioning_image_mask), dim=1)
with torch.autocast(
"cuda",
self.mixed_precision,
enabled=self.mixed_precision is not None,
):
down_block_additional_residuals = None
mid_block_additional_residual = None
add_to_down_block_inputs = None
add_to_output = None
if self.adapter is not None:
down_block_additional_residuals = self.adapter(conditioning_image)
if self.controlnet is not None:
controlnet_out = self.controlnet(
x_t=scaled_noisy_latents,
t=timesteps,
encoder_hidden_states=encoder_hidden_states,
micro_conditioning=micro_conditioning,
pooled_encoder_hidden_states=pooled_encoder_hidden_states,
controlnet_cond=conditioning_image,
)
down_block_additional_residuals = controlnet_out["down_block_res_samples"]
mid_block_additional_residual = controlnet_out["mid_block_res_sample"]
add_to_down_block_inputs = controlnet_out.get("add_to_down_block_inputs", None)
add_to_output = controlnet_out.get("add_to_output", None)
model_pred = self.unet(
x_t=scaled_noisy_latents,
t=timesteps,
encoder_hidden_states=encoder_hidden_states,
micro_conditioning=micro_conditioning,
pooled_encoder_hidden_states=pooled_encoder_hidden_states,
down_block_additional_residuals=down_block_additional_residuals,
mid_block_additional_residual=mid_block_additional_residual,
add_to_down_block_inputs=add_to_down_block_inputs,
add_to_output=add_to_output,
).sample
loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
return loss
@torch.no_grad()
def log_validation(self, step, num_validation_images: int, validation_prompts: Optional[List[str]] = None, validation_images: Optional[List[str]] = None):
import wandb
if isinstance(self.unet, DDP):
unet = self.unet.module
unet.eval()
unet_set_to_eval = True
else:
unet = self.unet
unet_set_to_eval = False
if self.adapter is not None:
adapter = self.adapter.module
adapter.eval()
else:
adapter = None
if self.controlnet is not None:
controlnet = self.controlnet.module
controlnet.eval()
else:
controlnet = None
formatted_validation_images = None
if validation_images is not None:
formatted_validation_images = []
wandb_validation_images = []
for validation_image_path in validation_images:
validation_image = Image.open(validation_image_path)
validation_image = validation_image.convert("RGB")
validation_image = validation_image.resize((1024, 1024))
conditioning_images = self.get_sdxl_conditioning_images(validation_image)
conditioning_image = conditioning_images["conditioning_image"]
if self.controlnet is not None and isinstance(self.controlnet, SDXLControlNetPreEncodedControlnetCond):
conditioning_image = self.vae.encode(conditioning_image[None, :, :, :].to(self.vae.device, dtype=self.vae.dtype))
conditionin_mask_image = TF.resize(conditioning_images["conditioning_mask_image"], conditioning_image.shape[2:]).to(conditioning_image.dtype, conditioning_image.device)
conditioning_image = torch.concat(conditioning_image, conditionin_mask_image, dim=1)
formatted_validation_images.append(conditioning_image)
wandb_validation_images.append(wandb.Image(conditioning_images["conditioning_image_as_pil"]))
if self.log_validation_input_images_every_time or not self.validation_images_logged:
wandb.log({"validation_conditioning": wandb_validation_images}, step=step)
self.validation_images_logged = True
generator = torch.Generator().manual_seed(0)
output_validation_images = []
for formatted_validation_image, validation_prompt in zip(formatted_validation_images, validation_prompts):
for _ in range(num_validation_images):
with torch.autocast("cuda"):
x_0 = sdxl_diffusion_loop(
prompts=validation_prompt,
images=formatted_validation_image,
unet=unet,
text_encoder_one=self.text_encoder_one,
text_encoder_two=self.text_encoder_two,
controlnet=controlnet,
adapter=adapter,
sigmas=self.sigmas,
generator=generator,
)
x_0 = self.vae.decode(x_0)
x_0 = self.vae.output_tensor_to_pil(x_0)[0]
output_validation_images.append(wandb.Image(x_0, caption=validation_prompt))
wandb.log({"validation": output_validation_images}, step=step)
if unet_set_to_eval:
unet.train()
if adapter is not None:
adapter.train()
if controlnet is not None:
controlnet.train()
def parameters(self):
if self.train_unet:
return self.unet.parameters()
if self.controlnet is not None and self.train_unet_up_blocks:
return itertools.chain(self.controlnet.parameters(), self.unet.up_blocks.parameters())
if self.controlnet is not None:
return self.controlnet.parameters()
if self.adapter is not None:
return self.adapter.parameters()
assert False
def save(self, save_to):
if self.train_unet:
safetensors.torch.save_file(self.unet.module.state_dict(), os.path.join(save_to, "unet.safetensors"))
if self.controlnet is not None and self.train_unet_up_blocks:
safetensors.torch.save_file(self.controlnet.module.state_dict(), os.path.join(save_to, "controlnet.safetensors"))
safetensors.torch.save_file(self.unet.module.up_blocks.state_dict(), os.path.join(save_to, "unet.safetensors"))
if self.controlnet is not None:
safetensors.torch.save_file(self.controlnet.module.state_dict(), os.path.join(save_to, "controlnet.safetensors"))
if self.adapter is not None:
safetensors.torch.save_file(self.adapter.module.state_dict(), os.path.join(save_to, "adapter.safetensors"))
def get_sdxl_dataset(train_shards: str, shuffle_buffer_size: int, batch_size: int, proportion_empty_prompts: float, get_sdxl_conditioning_images=None):
import webdataset as wds
dataset = (
wds.WebDataset(
train_shards,
resampled=True,
handler=wds.ignore_and_continue,
)
.shuffle(shuffle_buffer_size)
.decode("pil", handler=wds.ignore_and_continue)
.rename(
image="jpg;png;jpeg;webp",
text="text;txt;caption",
metadata="json",
handler=wds.warn_and_continue,
)
.map(lambda d: make_sample(d, proportion_empty_prompts=proportion_empty_prompts, get_sdxl_conditioning_images=get_sdxl_conditioning_images))
.select(lambda sample: "conditioning_image" not in sample or sample["conditioning_image"] is not None)
)
dataset = dataset.batched(batch_size, partial=False, collation_fn=default_collate)
return dataset
@torch.no_grad()
def make_sample(d, proportion_empty_prompts, get_sdxl_conditioning_images=None):
image = d["image"]
metadata = d["metadata"]
if random.random() < proportion_empty_prompts:
text = ""
else:
text = d["text"]
c_top, c_left, _, _ = get_random_crop_params([image.height, image.width], [1024, 1024])
original_width = int(metadata.get("original_width", 0.0))
original_height = int(metadata.get("original_height", 0.0))
micro_conditioning = torch.tensor([original_width, original_height, c_top, c_left, 1024, 1024])
text_input_ids_one = sdxl_tokenize_one(text)[0]
text_input_ids_two = sdxl_tokenize_two(text)[0]
image = image.convert("RGB")
image = TF.resize(
image,
1024,
interpolation=torchvision.transforms.InterpolationMode.BILINEAR,
)
image = TF.crop(
image,
c_top,
c_left,
1024,
1024,
)
sample = {
"micro_conditioning": micro_conditioning,
"text_input_ids_one": text_input_ids_one,
"text_input_ids_two": text_input_ids_two,
"image": SDXLVae.input_pil_to_tensor(image),
}
if get_sdxl_conditioning_images is not None:
conditioning_images = get_sdxl_conditioning_images(image)
sample["conditioning_image"] = conditioning_images["conditioning_image"]
if conditioning_images["conditioning_image_mask"] is not None:
sample["conditioning_image_mask"] = conditioning_images["conditioning_image_mask"]
return sample
def get_random_crop_params(input_size: Tuple[int, int], output_size: Tuple[int, int]) -> Tuple[int, int, int, int]:
h, w = input_size
th, tw = output_size
if h < th or w < tw:
raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
if w == tw and h == th:
return 0, 0, h, w
i = torch.randint(0, h - th + 1, size=(1,)).item()
j = torch.randint(0, w - tw + 1, size=(1,)).item()
return i, j, th, tw
def get_adapter_openpose_conditioning_image(image, open_pose):
resolution = image.width
conditioning_image = open_pose(image, detect_resolution=resolution, image_resolution=resolution, return_pil=False)
if (conditioning_image == 0).all():
return None, None
conditioning_image_as_pil = Image.fromarray(conditioning_image)
conditioning_image = TF.to_tensor(conditioning_image)
return dict(conditioning_image=conditioning_image, conditioning_image_as_pil=conditioning_image_as_pil)
def get_controlnet_canny_conditioning_image(image):
import cv2
conditioning_image = np.array(image)
conditioning_image = cv2.Canny(conditioning_image, 100, 200)
conditioning_image = conditioning_image[:, :, None]
conditioning_image = np.concatenate([conditioning_image, conditioning_image, conditioning_image], axis=2)
conditioning_image_as_pil = Image.fromarray(conditioning_image)
conditioning_image = TF.to_tensor(conditioning_image)
return dict(conditioning_image=conditioning_image, conditioning_image_as_pil=conditioning_image_as_pil)
def get_controlnet_pre_encoded_controlnet_inpainting_conditioning_image(image, conditioning_image_mask):
resolution = image.width
if conditioning_image_mask is None:
if random.random() <= 0.25:
conditioning_image_mask = np.ones((resolution, resolution), np.float32)
else:
conditioning_image_mask = random.choice([make_random_rectangle_mask, make_random_irregular_mask, make_outpainting_mask])(resolution, resolution)
conditioning_image_mask = torch.from_numpy(conditioning_image_mask)
conditioning_image_mask = conditioning_image_mask[None, :, :]
conditioning_image = TF.to_tensor(image)
# where mask is 1, zero out the pixels. Note that this requires mask to be concattenated
# with the mask so that the network knows the zeroed out pixels are from the mask and
# are not just zero in the original image
conditioning_image = conditioning_image * (conditioning_image_mask < 0.5)
conditioning_image_as_pil = TF.to_pil_image(conditioning_image)
conditioning_image = TF.normalize(conditioning_image, [0.5], [0.5])
return dict(conditioning_image=conditioning_image, conditioning_image_mask=conditioning_image_mask, conditioning_image_as_pil=conditioning_image_as_pil)
def get_controlnet_inpainting_conditioning_image(image, conditioning_image_mask):
resolution = image.width
if conditioning_image_mask is None:
if random.random() <= 0.25:
conditioning_image_mask = np.ones((resolution, resolution), np.float32)
else:
conditioning_image_mask = random.choice([make_random_rectangle_mask, make_random_irregular_mask, make_outpainting_mask])(resolution, resolution)
conditioning_image_mask = torch.from_numpy(conditioning_image_mask)
conditioning_image_mask = conditioning_image_mask[None, :, :]
conditioning_image = TF.to_tensor(image)
# Just zero out the pixels which will be masked
conditioning_image_as_pil = TF.to_pil_image(conditioning_image * (conditioning_image_mask < 0.5))
# where mask is set to 1, set to -1 "special" masked image pixel.
# -1 is outside of the 0-1 range that the controlnet normalized
# input is in.
conditioning_image = conditioning_image * (conditioning_image_mask < 0.5) + -1.0 * (conditioning_image_mask >= 0.5)
return dict(conditioning_image=conditioning_image, conditioning_image_mask=conditioning_image_mask, conditioning_image_as_pil=conditioning_image_as_pil)
# TODO: would be nice to just call a function from a tokenizers https://github.com/huggingface/tokenizers
# i.e. afaik tokenizing shouldn't require holding any state
tokenizer_one = CLIPTokenizerFast.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="tokenizer")
tokenizer_two = CLIPTokenizerFast.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="tokenizer_2")
def sdxl_tokenize_one(prompts):
return tokenizer_one(
prompts,
padding="max_length",
max_length=tokenizer_one.model_max_length,
truncation=True,
return_tensors="pt",
).input_ids
def sdxl_tokenize_two(prompts):
return tokenizer_two(
prompts,
padding="max_length",
max_length=tokenizer_one.model_max_length,
truncation=True,
return_tensors="pt",
).input_ids
def sdxl_text_conditioning(text_encoder_one, text_encoder_two, text_input_ids_one, text_input_ids_two):
prompt_embeds_1 = text_encoder_one(
text_input_ids_one,
output_hidden_states=True,
).hidden_states[-2]
prompt_embeds_1 = prompt_embeds_1.view(prompt_embeds_1.shape[0], prompt_embeds_1.shape[1], -1)
prompt_embeds_2 = text_encoder_two(
text_input_ids_two,
output_hidden_states=True,
)
pooled_encoder_hidden_states = prompt_embeds_2[0]
prompt_embeds_2 = prompt_embeds_2.hidden_states[-2]
prompt_embeds_2 = prompt_embeds_2.view(prompt_embeds_2.shape[0], prompt_embeds_2.shape[1], -1)
encoder_hidden_states = torch.cat((prompt_embeds_1, prompt_embeds_2), dim=-1)
return encoder_hidden_states, pooled_encoder_hidden_states
def make_random_rectangle_mask(
height,
width,
margin=10,
bbox_min_size=100,
bbox_max_size=512,
min_times=1,
max_times=2,
):
mask = np.zeros((height, width), np.float32)
bbox_max_size = min(bbox_max_size, height - margin * 2, width - margin * 2)
times = np.random.randint(min_times, max_times + 1)
for i in range(times):
box_width = np.random.randint(bbox_min_size, bbox_max_size)
box_height = np.random.randint(bbox_min_size, bbox_max_size)
start_x = np.random.randint(margin, width - margin - box_width + 1)
start_y = np.random.randint(margin, height - margin - box_height + 1)
mask[start_y : start_y + box_height, start_x : start_x + box_width] = 1
return mask
def make_random_irregular_mask(height, width, max_angle=4, max_len=60, max_width=256, min_times=1, max_times=2):
import cv2
mask = np.zeros((height, width), np.float32)
times = np.random.randint(min_times, max_times + 1)
for i in range(times):
start_x = np.random.randint(width)
start_y = np.random.randint(height)
for j in range(1 + np.random.randint(5)):
angle = 0.01 + np.random.randint(max_angle)
if i % 2 == 0:
angle = 2 * 3.1415926 - angle
length = 10 + np.random.randint(max_len)
brush_w = 5 + np.random.randint(max_width)
end_x = np.clip((start_x + length * np.sin(angle)).astype(np.int32), 0, width)
end_y = np.clip((start_y + length * np.cos(angle)).astype(np.int32), 0, height)
choice = random.randint(0, 2)
if choice == 0:
cv2.line(mask, (start_x, start_y), (end_x, end_y), 1.0, brush_w)
elif choice == 1:
cv2.circle(mask, (start_x, start_y), radius=brush_w, color=1.0, thickness=-1)
elif choice == 2:
radius = brush_w // 2
mask[
start_y - radius : start_y + radius,
start_x - radius : start_x + radius,
] = 1
else:
assert False
start_x, start_y = end_x, end_y
return mask
def make_outpainting_mask(height, width, probs=[0.5, 0.5, 0.5, 0.5]):
mask = np.zeros((height, width), np.float32)
at_least_one_mask_applied = False
coords = [
[(0, 0), (1, get_padding(height))],
[(0, 0), (get_padding(width), 1)],
[(0, 1 - get_padding(height)), (1, 1)],
[(1 - get_padding(width), 0), (1, 1)],
]
for pp, coord in zip(probs, coords):
if np.random.random() < pp:
at_least_one_mask_applied = True
mask = apply_padding(mask=mask, coord=coord)
if not at_least_one_mask_applied:
idx = np.random.choice(range(len(coords)), p=np.array(probs) / sum(probs))
mask = apply_padding(mask=mask, coord=coords[idx])
return mask
def get_padding(size, min_padding_percent=0.04, max_padding_percent=0.5):
n1 = int(min_padding_percent * size)
n2 = int(max_padding_percent * size)
return np.random.randint(n1, n2) / size
def apply_padding(mask, coord):
height, width = mask.shape
mask[
int(coord[0][0] * height) : int(coord[1][0] * height),
int(coord[0][1] * width) : int(coord[1][1] * width),
] = 1
return mask
@torch.no_grad()
def sdxl_diffusion_loop(
prompts: Union[str, List[str]],
unet,
text_encoder_one,
text_encoder_two,
images=None,
controlnet=None,
adapter=None,
sigmas=None,
timesteps=None,
x_T=None,
micro_conditioning=None,
guidance_scale=5.0,
generator=None,
negative_prompts=None,
diffusion_loop=euler_ode_solver_diffusion_loop,
):
if isinstance(prompts, str):
prompts = [prompts]
batch_size = len(prompts)
if negative_prompts is not None and guidance_scale > 1.0:
prompts += negative_prompts
encoder_hidden_states, pooled_encoder_hidden_states = sdxl_text_conditioning(
text_encoder_one,
text_encoder_two,
sdxl_tokenize_one(prompts).to(text_encoder_one.device),
sdxl_tokenize_two(prompts).to(text_encoder_two.device),
)
encoder_hidden_states = encoder_hidden_states.to(unet.dtype)
pooled_encoder_hidden_states = pooled_encoder_hidden_states.to(unet.dtype)
if guidance_scale > 1.0:
if negative_prompts is None:
negative_encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
negative_pooled_encoder_hidden_states = torch.zeros_like(pooled_encoder_hidden_states)
else:
encoder_hidden_states, negative_encoder_hidden_states = torch.chunk(encoder_hidden_states, 2)
pooled_encoder_hidden_states, negative_pooled_encoder_hidden_states = torch.chunk(pooled_encoder_hidden_states, 2)
else:
negative_encoder_hidden_states = None
negative_pooled_encoder_hidden_states = None
if sigmas is None:
sigmas = make_sigmas(device=unet.device)
if timesteps is None:
timesteps = torch.linspace(0, sigmas.numel() - 1, 50, dtype=torch.long, device=unet.device)
if x_T is None:
x_T = torch.randn((batch_size, 4, 1024 // 8, 1024 // 8), dtype=unet.dtype, device=unet.device, generator=generator)
x_T = x_T * ((sigmas[timesteps[-1]] ** 2 + 1) ** 0.5)
if micro_conditioning is None:
micro_conditioning = torch.tensor([[1024, 1024, 0, 0, 1024, 1024]], dtype=torch.long, device=unet.device)
micro_conditioning = micro_conditioning.expand(batch_size, -1)
if adapter is not None:
down_block_additional_residuals = adapter(images.to(dtype=adapter.dtype, device=adapter.device))
else:
down_block_additional_residuals = None
if controlnet is not None:
controlnet_cond = images.to(dtype=controlnet.dtype, device=controlnet.device)
else:
controlnet_cond = None
eps_theta = lambda *args, **kwargs: sdxl_eps_theta(
*args,
**kwargs,
unet=unet,
encoder_hidden_states=encoder_hidden_states,
pooled_encoder_hidden_states=pooled_encoder_hidden_states,
negative_encoder_hidden_states=negative_encoder_hidden_states,
negative_pooled_encoder_hidden_states=negative_pooled_encoder_hidden_states,
micro_conditioning=micro_conditioning,
guidance_scale=guidance_scale,
controlnet=controlnet,
controlnet_cond=controlnet_cond,
down_block_additional_residuals=down_block_additional_residuals,
)
x_0 = diffusion_loop(eps_theta=eps_theta, timesteps=timesteps, sigmas=sigmas, x_T=x_T)
return x_0
@torch.no_grad()
def sdxl_eps_theta(
x_t,
t,
sigma,
unet,
encoder_hidden_states,
pooled_encoder_hidden_states,
negative_encoder_hidden_states,
negative_pooled_encoder_hidden_states,
micro_conditioning,
guidance_scale,
controlnet=None,
controlnet_cond=None,
down_block_additional_residuals=None,
):
# TODO - how does this not effect the ode we are solving
scaled_x_t = x_t / ((sigma**2 + 1) ** 0.5)
if guidance_scale > 1.0:
scaled_x_t = torch.concat([scaled_x_t, scaled_x_t])
encoder_hidden_states = torch.concat((encoder_hidden_states, negative_encoder_hidden_states))
pooled_encoder_hidden_states = torch.concat((pooled_encoder_hidden_states, negative_pooled_encoder_hidden_states))
micro_conditioning = torch.concat([micro_conditioning, micro_conditioning])
if controlnet_cond is not None:
controlnet_cond = torch.concat([controlnet_cond, controlnet_cond])
if controlnet is not None:
controlnet_out = controlnet(
x_t=scaled_x_t.to(controlnet.dtype),
t=t,
encoder_hidden_states=encoder_hidden_states.to(controlnet.dtype),
micro_conditioning=micro_conditioning.to(controlnet.dtype),
pooled_encoder_hidden_states=pooled_encoder_hidden_states.to(controlnet.dtype),
controlnet_cond=controlnet_cond,
)
down_block_additional_residuals = [x.to(unet.dtype) for x in controlnet_out["down_block_res_samples"]]
mid_block_additional_residual = controlnet_out["mid_block_res_sample"].to(unet.dtype)
add_to_down_block_inputs = controlnet_out.get("add_to_down_block_inputs", None)
if add_to_down_block_inputs is not None:
add_to_down_block_inputs = [x.to(unet.dtype) for x in add_to_down_block_inputs]
add_to_output = controlnet_out.get("add_to_output", None)
if add_to_output is not None:
add_to_output = add_to_output.to(unet.dtype)
else:
mid_block_additional_residual = None
add_to_down_block_inputs = None
add_to_output = None
eps_hat = unet(
x_t=scaled_x_t,
t=t,
encoder_hidden_states=encoder_hidden_states,
micro_conditioning=micro_conditioning,
pooled_encoder_hidden_states=pooled_encoder_hidden_states,
down_block_additional_residuals=down_block_additional_residuals,
mid_block_additional_residual=mid_block_additional_residual,
add_to_down_block_inputs=add_to_down_block_inputs,
add_to_output=add_to_output,
)
if guidance_scale > 1.0:
eps_hat, eps_hat_uncond = eps_hat.chunk(2)
eps_hat = eps_hat_uncond + guidance_scale * (eps_hat - eps_hat_uncond)
return eps_hat
known_negative_prompt = "text, watermark, low-quality, signature, moiré pattern, downsampling, aliasing, distorted, blurry, glossy, blur, jpeg artifacts, compression artifacts, poorly drawn, low-resolution, bad, distortion, twisted, excessive, exaggerated pose, exaggerated limbs, grainy, symmetrical, duplicate, error, pattern, beginner, pixelated, fake, hyper, glitch, overexposed, high-contrast, bad-contrast"
if __name__ == "__main__":
from argparse import ArgumentParser
args = ArgumentParser()
args.add_argument("--prompts", required=True, type=str, nargs="+")
args.add_argument("--negative_prompts", required=False, type=str, nargs="+")
args.add_argument("--use_known_negative_prompt", action="store_true")
args.add_argument("--num_images_per_prompt", required=True, type=int, default=1)
args.add_argument("--num_inference_steps", required=False, type=int, default=50)
args.add_argument("--images", required=False, type=str, default=None, nargs="+")
args.add_argument("--masks", required=False, type=str, default=None, nargs="+")
args.add_argument("--controlnet_checkpoint", required=False, type=str, default=None)
args.add_argument("--controlnet", required=False, choices=["SDXLControlNet", "SDXLControlNetFull", "SDXLControNetPreEncodedControlnetCond"], default=None)
args.add_argument("--adapter_checkpoint", required=False, type=str, default=None)
args.add_argument("--device", required=False, default=None)
args.add_argument("--dtype", required=False, default="fp16", choices=["fp16", "fp32"])
args.add_argument("--guidance_scale", required=False, default=5.0, type=float)
args.add_argument("--seed", required=False, type=int)
args = args.parse_args()
if args.device is None:
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
device = "mps"
if args.dtype == "fp16":
dtype = torch.float16
text_encoder_one = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", variant="fp16", torch_dtype=torch.float16)
text_encoder_one.to(device=device)
text_encoder_two = CLIPTextModelWithProjection.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2", variant="fp16", torch_dtype=torch.float16)
text_encoder_two.to(device=device)
vae = SDXLVae.load_fp16_fix(device=device)
vae.to(torch.float16)
unet = SDXLUNet.load_fp16(device=device)
elif args.dtype == "fp32":
dtype = torch.float32
text_encoder_one = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
text_encoder_one.to(device=device)
text_encoder_two = CLIPTextModelWithProjection.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2")
text_encoder_two.to(device=device)
vae = SDXLVae.load_fp16_fix(device=device)
unet = SDXLUNet.load_fp32(device=device)
else:
assert False
if args.controlnet == "SDXLControlNet":
controlnet = SDXLControlNet.load(args.controlnet_checkpoint, device=device)
controlnet.to(dtype)
elif args.controlnet == "SDXLControlNetFull":
controlnet = SDXLControlNetFull.load(args.controlnet_checkpoint, device=device)
controlnet.to(dtype)
elif args.controlnet == "SDXLControlNetPreEncodedControlnetCond":
controlnet = SDXLControlNetPreEncodedControlnetCond.load(args.controlnet_checkpoint, device=device)
controlnet.to(dtype)
else:
controlnet = None
if args.adapter_checkpoint is not None:
adapter = SDXLAdapter.load(args.adapter_checkpoint, device=device)
adapter.to(dtype)
else:
adapter = None
sigmas = make_sigmas(device=device).to(unet.dtype)
timesteps = torch.linspace(0, sigmas.numel() - 1, args.num_inference_steps, dtype=torch.long, device=unet.device)
prompts = []
for prompt in args.prompts:
prompts += [prompt] * args.num_images_per_prompt
if args.use_known_negative_prompt:
args.negative_prompts = [known_negative_prompt]
if args.negative_prompts is None:
negative_prompts = None
elif len(args.negative_prompts) == 1:
negative_prompts = args.negative_prompts * len(prompts)
elif len(args.negative_prompts) == len(args.prompts):
negative_prompts = []
for negative_prompt in args.negative_prompts:
negative_prompts += [negative_prompt] * args.num_images_per_prompt
else:
assert False
if args.images is not None:
images = []
for image_idx, image in enumerate(args.images):
image = Image.open(image)
image = image.convert("RGB")
image = image.resize((1024, 1024))
image = TF.to_tensor(image)
if args.masks is not None:
mask = args.masks[image_idx]
mask = Image.open(mask)
mask = mask.convert("L")
mask = mask.resize((1024, 1024))
mask = TF.to_tensor(mask)
if isinstance(controlnet, SDXLControlNetPreEncodedControlnetCond):
image = image * (mask < 0.5)
image = TF.normalize(image, [0.5], [0.5])
image = vae.encode(image[None, :, :, :].to(dtype=vae.dtype, device=vae.device)).to(dtype=controlnet.dtype, device=controlnet.device)
mask = TF.resize(mask, (1024 // 8, 1024 // 8))[None, :, :, :].to(dtype=image.dtype, device=image.device)
image = torch.concat((image, mask), dim=1)
else:
image = (image * (mask < 0.5) + -1.0 * (mask >= 0.5)).to(dtype=dtype, device=device)
image = image[None, :, :, :]
images += [image] * args.num_images_per_prompt
images = torch.concat(images)
else:
images = None
if args.seed is None:
generator = None
else:
generator = torch.Generator(device).manual_seed(args.seed)
images = sdxl_diffusion_loop(
prompts=prompts,
unet=unet,
text_encoder_one=text_encoder_one,
text_encoder_two=text_encoder_two,
images=images,
controlnet=controlnet,
adapter=adapter,
sigmas=sigmas,
timesteps=timesteps,
guidance_scale=args.guidance_scale,
negative_prompts=negative_prompts,
generator=generator,
)
images = vae.output_tensor_to_pil(vae.decode(images))
for i, image in enumerate(images):
image.save(f"out_{i}.png")