RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add a control generation script.

This commit is contained in:
Jaret Burkett 2026-06-25 10:14:24 -06:00
parent d14f6e567a
commit 4eb0707639
2 changed files with 366 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

256
scripts/add_mask_dataset.py Normal file
View File

@ -0,0 +1,256 @@
import os
import sys
import queue
import random
import argparse
import threading
import traceback
import torch
from tqdm import tqdm
# allow importing from project root
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from toolkit.control_generator import ControlGenerator, img_ext_list
def control_exists(img_path, control_type):
# mirrors the lookup in ControlGenerator.get_control_path so we can skip
# images another instance has already finished
controls_folder = os.path.join(os.path.dirname(img_path), "_controls")
file_name_no_ext = os.path.splitext(os.path.basename(img_path))[0]
file_name_no_ext_control = f"{file_name_no_ext}.{control_type}"
for ext in img_ext_list:
if os.path.exists(os.path.join(controls_folder, file_name_no_ext_control + ext)):
return True
return False
# sentinel pushed through the queues to tell workers to stop
_DONE = object()
def run_pipeline(control_gen, img_list, control_type, regen, n_load, n_save):
# Three-stage pipeline so the GPU never waits on disk/CPU work:
# loaders (N threads) -> read + exif + resize + preprocess to a CPU tensor
# gpu worker (1 thread) -> model forward only (kept single so VRAM is bounded)
# savers (M threads) -> postprocess (resize/alpha) + encode + write
# The heavy CPU work (resize/normalize on input, resize/convert on output) is
# pushed onto the loader/saver threads so the GPU thread does almost nothing
# but the forward pass. Bounded queues apply backpressure so we don't load the
# whole dataset into RAM.
infer_q = queue.Queue(maxsize=n_load * 2)
save_q = queue.Queue(maxsize=n_save * 2)
path_q = queue.Queue()
for img_path in img_list:
path_q.put(img_path)
# miniters=1 disables tqdm's dynamic-miniters heuristic (which otherwise
# raises the redraw threshold after a fast burst and makes the bar look
# frozen); mininterval keeps redraws time-based and cheap.
pbar = tqdm(total=len(img_list), desc=f"Generating {control_type}",
miniters=1, mininterval=0.5)
pbar_lock = threading.Lock()
# set on completion OR on Ctrl-C; every blocking call below uses a timeout and
# re-checks this so the worker threads can actually be shut down.
stop_event = threading.Event()
def put(q, item):
# interruptible put: blocks until there's room, but wakes periodically so
# a stop request (or KeyboardInterrupt on the main thread) is honored.
while not stop_event.is_set():
try:
q.put(item, timeout=0.2)
return
except queue.Full:
continue
def loader():
while not stop_event.is_set():
try:
img_path = path_q.get_nowait()
except queue.Empty:
break
try:
if not regen and control_exists(img_path, control_type):
# another instance (or a previous run) already did it
with pbar_lock:
pbar.update(1)
continue
image = control_gen.load_image(img_path)
payload = control_gen.preprocess(image, control_type)
put(infer_q, (img_path, image, payload))
except Exception:
traceback.print_exc()
with pbar_lock:
pbar.update(1)
def saver():
while not stop_event.is_set():
try:
item = save_q.get(timeout=0.2)
except queue.Empty:
continue
if item is _DONE:
break
img_path, image, result = item
try:
out_image = control_gen.postprocess(result, image, control_type)
save_path = control_gen.control_save_path(img_path, control_type)
control_gen.save_control(out_image, save_path)
except Exception:
traceback.print_exc()
finally:
with pbar_lock:
pbar.update(1)
loaders = [threading.Thread(target=loader, daemon=True) for _ in range(n_load)]
savers = [threading.Thread(target=saver, daemon=True) for _ in range(n_save)]
for t in loaders + savers:
t.start()
# GPU stage runs on the main thread: pull preprocessed tensors, run the
# forward pass, hand the raw result off to the savers. We stop once every
# loader has exited and nothing is left queued for inference.
interrupted = False
try:
while not stop_event.is_set():
if not any(t.is_alive() for t in loaders) and infer_q.empty():
break
try:
img_path, image, payload = infer_q.get(timeout=0.1)
except queue.Empty:
continue
try:
result = control_gen.run_inference(payload, control_type)
put(save_q, (img_path, image, result))
except Exception:
traceback.print_exc()
with pbar_lock:
pbar.update(1)
except KeyboardInterrupt:
interrupted = True
print("\nInterrupted, shutting down...")
if interrupted:
# abort: tell every worker to stop; pending items are dropped
stop_event.set()
else:
# normal finish: let savers drain whatever is still queued, then stop
for _ in savers:
save_q.put(_DONE)
# join with a timeout so a stuck worker can never wedge shutdown; threads are
# daemons, so anything still alive is torn down when we return.
for t in savers:
t.join(timeout=5)
pbar.close()
if interrupted:
raise KeyboardInterrupt
def main():
parser = argparse.ArgumentParser(
description="Generate masks for a dataset using the ControlGenerator"
)
parser.add_argument("img_dir", type=str, help="Path to image directory")
parser.add_argument(
"--control",
type=str,
default="mask",
choices=["mask", "inpaint", "depth", "pose", "line", "sapiens2_mask"],
help="Control type to generate (default: mask)",
)
parser.add_argument(
"--device", type=str, default="cuda", help="Device to run on (default: cuda)"
)
parser.add_argument("--debug", action="store_true", help="Enable debug mode")
parser.add_argument(
"--regen",
action="store_true",
help="Regenerate controls even if they already exist",
)
parser.add_argument(
"--shuffle",
action="store_true",
help="Shuffle image order so multiple instances on the same dataset "
"don't chase the same images",
)
parser.add_argument(
"--load-workers",
type=int,
default=16,
help="Number of threads for loading/resizing images (default: 4)",
)
parser.add_argument(
"--save-workers",
type=int,
default=16,
help="Number of threads for saving controls (default: 4)",
)
args = parser.parse_args()
img_dir = args.img_dir
if not os.path.isdir(img_dir):
print(f"Error: {img_dir} is not a directory")
sys.exit(1)
# find images, skipping existing _controls folders and hidden files
img_list = []
for root, dirs, files in os.walk(img_dir):
if "_controls" in root:
continue
for file in files:
if file.startswith("."):
continue
if file.lower().endswith(tuple(img_ext_list)):
img_list.append(os.path.join(root, file))
if len(img_list) == 0:
print(f"Error: no images found in {img_dir}")
sys.exit(1)
# filter out images that already have controls up front so the progress bar
# reflects only real work (otherwise it races through thousands of skips and
# the rate/ETA are meaningless). The loader still re-checks just before
# processing to handle the multi-instance race.
if not args.regen:
total = len(img_list)
img_list = [p for p in img_list if not control_exists(p, args.control)]
skipped = total - len(img_list)
if skipped:
print(f"Skipping {skipped} images that already have '{args.control}' controls")
if len(img_list) == 0:
print("All images already have controls. Nothing to do.")
return
if args.shuffle:
random.shuffle(img_list)
control_gen = ControlGenerator(torch.device(args.device))
control_gen.debug = args.debug
control_gen.regen = args.regen
interrupted = False
try:
run_pipeline(
control_gen,
img_list,
args.control,
args.regen,
max(1, args.load_workers),
max(1, args.save_workers),
)
except KeyboardInterrupt:
interrupted = True
finally:
control_gen.cleanup()
if interrupted:
sys.exit(130)
print("Done")
if __name__ == "__main__":
main()

180
toolkit/control_generator.py
View File

@ -57,38 +57,109 @@ class ControlGenerator:
if self.debug:
print(*args, **kwargs)
def _generate_control(self, img_path, control_type):
device = self.device
image: Image = None
coltrols_folder = os.path.join(os.path.dirname(img_path), '_controls')
file_name_no_ext = os.path.splitext(os.path.basename(img_path))[0]
# we need to generate the control. Unload model if not unloaded
def ensure_unloaded(self):
# unload the training model (if any) before generating controls
if not self.has_unloaded:
if self.sd is not None:
print("Unloading model to generate controls")
self.sd.set_device_state_preset('unload')
self.has_unloaded = True
if image is None:
# make sure image is loaded if we havent loaded it with another control
image = Image.open(img_path).convert('RGB')
image = exif_transpose(image)
def load_image(self, img_path):
# CPU/disk stage: read, orient, and downscale to a max of 1mp
image = Image.open(img_path).convert('RGB')
image = exif_transpose(image)
# resize to a max of 1mp
max_size = 1024 * 1024
max_size = 1024 * 1024
w, h = image.size
if w * h > max_size:
scale = math.sqrt(max_size / (w * h))
w = int(w * scale)
h = int(h * scale)
image = image.resize((w, h), Image.BICUBIC)
return image
w, h = image.size
if w * h > max_size:
scale = math.sqrt(max_size / (w * h))
w = int(w * scale)
h = int(h * scale)
image = image.resize((w, h), Image.BICUBIC)
def control_save_path(self, img_path, control_type):
coltrols_folder = os.path.join(os.path.dirname(img_path), '_controls')
file_name_no_ext = os.path.splitext(os.path.basename(img_path))[0]
# inpaint needs alpha and mask is a near-binary single channel; webp
# compresses both far smaller than jpg. The rest stay jpg.
ext = 'webp' if control_type in ('inpaint', 'mask') else 'jpg'
return os.path.join(
coltrols_folder, f"{file_name_no_ext}.{control_type}.{ext}")
def save_control(self, out_image, save_path):
# CPU/disk stage: encode and write the generated control
os.makedirs(os.path.dirname(save_path), exist_ok=True)
if save_path.lower().endswith('.webp'):
# method=6 trades CPU (already off the GPU thread) for smaller files
out_image.save(save_path, quality=80, method=6)
else:
out_image.save(save_path)
def _bg_transform(self):
return transforms.Compose([
transforms.Resize((1024, 1024)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
def _ensure_bg_remover(self):
if self.control_bg_remover is None:
from transformers import AutoModelForImageSegmentation
self.control_bg_remover = AutoModelForImageSegmentation.from_pretrained(
'ZhengPeng7/BiRefNet_HR',
trust_remote_code=True,
revision="a7a562f6fd16021180f2f4348f4de003a2d3d1e1",
dtype=torch.float16
).to(self.device)
self.control_bg_remover.eval()
def preprocess(self, image, control_type):
# CPU stage. For the bg-remover path this does the expensive resize +
# normalize and returns a ready-to-run float16 tensor, so the GPU thread
# never has to. Other control types preprocess inside their model, so we
# just pass the PIL image straight through.
if control_type in ('inpaint', 'mask'):
return self._bg_transform()(image).unsqueeze(0).to(torch.float16)
return image
def run_inference(self, payload, control_type):
# GPU stage. Returns an intermediate result for postprocess(). Models are
# lazily loaded here, so call from a single thread per generator instance.
self.ensure_unloaded()
if control_type in ('inpaint', 'mask'):
self._ensure_bg_remover()
x = payload.to(self.device).to(torch.float16)
with torch.inference_mode():
preds = self.control_bg_remover(x)[-1].sigmoid().cpu()
return preds[0].squeeze() # CPU mask tensor, 1024x1024
# everything else does preprocessing + inference together on this thread
return self.run_control(payload, control_type)
def postprocess(self, result, image, control_type):
# CPU stage. Turns the inference result into the final control image.
if control_type in ('inpaint', 'mask'):
mask = transforms.ToPILImage()(result).resize(image.size)
if control_type == 'inpaint':
# inpainting currently only supports the "erased" section to inpaint
mask = ImageOps.invert(mask)
out = image.copy()
out.putalpha(mask)
return out
# keep the mask single-channel grayscale; the loader converts as
# needed and this roughly thirds the file size vs RGB
return mask
# the fallback path already produced a finished PIL image
return result
def run_control(self, image, control_type):
# GPU stage: run inference on an already-loaded image and return the
# resulting PIL image (no disk IO). Models are lazily loaded here, so
# this must be called from a single thread per generator instance.
device = self.device
self.ensure_unloaded()
save_path = os.path.join(
coltrols_folder, f"{file_name_no_ext}.{control_type}.jpg")
os.makedirs(coltrols_folder, exist_ok=True)
if control_type == 'depth':
self.debug_print("Generating depth control")
if self.control_depth_model is None:
@ -107,8 +178,7 @@ class ControlGenerator:
out_tensor = out_tensor.squeeze(0).cpu().numpy()
img = Image.fromarray(out_tensor.astype('uint8'))
img = img.resize(in_size, Image.LANCZOS)
img.save(save_path)
return save_path
return img
elif control_type == 'pose':
self.debug_print("Generating pose control")
if self.control_pose_model is None:
@ -131,8 +201,7 @@ class ControlGenerator:
img = self.control_pose_model(
img, output_type="pil", include_hands=True, include_face=True, detect_resolution=detect_res)
img = img.convert('RGB')
img.save(save_path)
return save_path
return img
elif control_type == 'line':
self.debug_print("Generating line control")
@ -146,63 +215,34 @@ class ControlGenerator:
# img = img.filter(ImageFilter.GaussianBlur(radius=1))
img = img.point(lambda p: p > 128 and 255)
img = img.convert('RGB')
img.save(save_path)
return save_path
return img
elif control_type in ['inpaint', 'mask']:
self.debug_print("Generating inpaint/mask control")
img = image.copy()
if self.control_bg_remover is None:
from transformers import AutoModelForImageSegmentation
self.control_bg_remover = AutoModelForImageSegmentation.from_pretrained(
'ZhengPeng7/BiRefNet_HR',
trust_remote_code=True,
revision="a7a562f6fd16021180f2f4348f4de003a2d3d1e1",
torch_dtype=torch.float16
).to(device)
self.control_bg_remover.eval()
image_size = (1024, 1024)
transform_image = transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [
0.229, 0.224, 0.225])
])
input_images = transform_image(img).unsqueeze(
0).to('cuda').to(torch.float16)
# Prediction
preds = self.control_bg_remover(input_images)[-1].sigmoid().cpu()
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(img.size)
if control_type == 'inpaint':
# inpainting feature currently only supports "erased" section desired to inpaint
mask = ImageOps.invert(mask)
img.putalpha(mask)
save_path = os.path.join(
coltrols_folder, f"{file_name_no_ext}.{control_type}.webp")
else:
img = mask
img = img.convert('RGB')
img.save(save_path)
return save_path
# delegate to the staged methods so this matches the threaded path
payload = self.preprocess(image, control_type)
result = self.run_inference(payload, control_type)
return self.postprocess(result, image, control_type)
elif control_type in ['sapiens2_mask']:
self.debug_print("Generating sapiens2_mask control")
if self.control_bg_remover is None:
from toolkit.models.sapiens2 import Sapiens2Matting
self.control_bg_remover = Sapiens2Matting.from_pretrained(
device=device,
device=device,
dtype=torch.float16
)
img = image.copy()
img = self.control_bg_remover(img)
img.save(save_path)
return save_path
return img
else:
raise Exception(f"Error: unknown control type {control_type}")
def _generate_control(self, img_path, control_type):
image = self.load_image(img_path)
out_image = self.run_control(image, control_type)
save_path = self.control_save_path(img_path, control_type)
self.save_control(out_image, save_path)
return save_path
def cleanup(self):
if self.control_depth_model is not None:
self.control_depth_model = None