diff --git a/scripts/add_mask_dataset.py b/scripts/add_mask_dataset.py new file mode 100644 index 00000000..8dda4863 --- /dev/null +++ b/scripts/add_mask_dataset.py @@ -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() diff --git a/toolkit/control_generator.py b/toolkit/control_generator.py index ea9dd5d3..88c03ee3 100644 --- a/toolkit/control_generator.py +++ b/toolkit/control_generator.py @@ -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