control_generator.py

import gc
import math
import os
import torch
from typing import Literal
from PIL import Image, ImageFilter, ImageOps
from PIL.ImageOps import exif_transpose
from tqdm import tqdm

from torchvision import transforms
from toolkit import device_utils

# supress all warnings
import warnings

warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)


def flush(garbage_collect=True):
    device_utils.empty_cache()
    if garbage_collect:
        gc.collect()


ControlTypes = Literal['depth', 'pose', 'line', 'inpaint', 'mask']

img_ext_list = ['.jpg', '.jpeg', '.png', '.webp']


class ControlGenerator:
    def __init__(self, device, sd=None):
        self.device = device
        self.sd = sd  # optional. It will unload the model if not None
        self.has_unloaded = False
        self.control_depth_model = None
        self.control_pose_model = None
        self.control_line_model = None
        self.control_bg_remover = None
        self.debug = False
        self.regen = False

    def get_control_path(self, img_path, control_type: ControlTypes):
        if self.regen:
            return self._generate_control(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]
        file_name_no_ext_control = f"{file_name_no_ext}.{control_type}"
        for ext in img_ext_list:
            possible_path = os.path.join(
                coltrols_folder, file_name_no_ext_control + ext)
            if os.path.exists(possible_path):
                return possible_path
        # if we get here, we need to generate the control
        return self._generate_control(img_path, control_type)

    def debug_print(self, *args, **kwargs):
        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
        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)

            # resize to a max of 1mp
            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)

        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:
                from transformers import pipeline
                self.control_depth_model = pipeline(
                    task="depth-estimation",
                    model="depth-anything/Depth-Anything-V2-Large-hf",
                    device=device,
                    torch_dtype=torch.float16
                )
            img = image.copy()
            in_size = img.size
            output = self.control_depth_model(img)
            out_tensor = output["predicted_depth"]  # shape (1, H, W) 0 - 255
            out_tensor = out_tensor.clamp(0, 255)
            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
        elif control_type == 'pose':
            self.debug_print("Generating pose control")
            if self.control_pose_model is None:
                try:
                    import onnxruntime
                    onnxruntime.set_default_logger_severity(3)
                except ImportError:
                    raise ImportError(
                        "onnxruntime is not installed. Please install it with pip install onnxruntime or onnxruntime-gpu")
                try:
                    from easy_dwpose import DWposeDetector
                    self.control_pose_model = DWposeDetector(
                        device=str(device))
                except ImportError:
                    raise ImportError(
                        "easy-dwpose is not installed. Please install it with pip install git+https://github.com/jaretburkett/easy_dwpose.git")
            img = image.copy()

            detect_res = int(math.sqrt(img.size[0] * img.size[1]))
            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

        elif control_type == 'line':
            self.debug_print("Generating line control")
            if self.control_line_model is None:
                from controlnet_aux import TEEDdetector
                self.control_line_model = TEEDdetector.from_pretrained(
                    "fal-ai/teed", filename="5_model.pth").to(device)
            img = image.copy()
            img = self.control_line_model(img, detect_resolution=1024)
            # apply threshold
            # 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
        elif control_type == 'inpaint' or control_type == '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="595e212b3eaa6a1beaad56cee49749b1e00b1596",
                    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])
            ])

            # Assuming self.device is correct
            input_images = transform_image(img).unsqueeze(
                0).to(device).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
        else:
            raise Exception(f"Error: unknown control type {control_type}")

    def cleanup(self):
        if self.control_depth_model is not None:
            self.control_depth_model = None
        if self.control_pose_model is not None:
            self.control_pose_model = None
        if self.control_line_model is not None:
            self.control_line_model = None
        if self.control_bg_remover is not None:
            self.control_bg_remover = None
        if self.sd is not None and self.has_unloaded:
            self.sd.restore_device_state()
        self.has_unloaded = False

        flush()


if __name__ == "__main__":
    import sys
    import argparse
    import time
    import transformers
    transformers.logging.set_verbosity_error()

    control_times = {
        'depth': 0,
        'pose': 0,
        'line': 0,
        'inpaint': 0,
        'mask': 0
    }

    controls = control_times.keys()

    parser = argparse.ArgumentParser(description="Generate control images")
    parser.add_argument("img_dir", type=str, help="Path to image directory")
    parser.add_argument('--debug', action='store_true',
                        help="Enable debug mode")
    parser.add_argument('--regen', action='store_true',
                        help="Regenerate all controls")

    args = parser.parse_args()
    img_dir = args.img_dir
    if not os.path.exists(img_dir):
        print(f"Error: {img_dir} does not exist")
        exit()
    if not os.path.isdir(img_dir):
        print(f"Error: {img_dir} is not a directory")
        exit()

    # find images
    img_list = []
    for root, dirs, files in os.walk(img_dir):
        for file in files:
            if "_controls" in root:
                continue
            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}")
        exit()

    # load model
    idx = 0
    for img_path in tqdm(img_list):
        for control in controls:
            start = time.time()
            control_gen = ControlGenerator(device_utils.get_device())
            control_gen.debug = args.debug
            control_gen.regen = args.regen
            control_path = control_gen.get_control_path(img_path, control)
            end = time.time()
            # dont track for first 2 images
            if idx < 2:
                continue
            control_times[control] += end - start
        idx += 1

    # determine avgt time
    for control in controls:
        control_times[control] /= (idx - 2)
        print(
            f"Avg time for {control} control: {control_times[control]:.2f} seconds")

    print("Done")