257 lines
8.9 KiB
Python
257 lines
8.9 KiB
Python
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()
|