diffsynth/diffusion/base_pipeline.py

from PIL import Image
import torch
import numpy as np
from einops import repeat, reduce
from typing import Union
from ..core import AutoTorchModule, AutoWrappedLinear, load_state_dict, ModelConfig
from ..utils.lora import GeneralLoRALoader
from ..models.model_loader import ModelPool
from ..utils.controlnet import ControlNetInput
from typing import Optional
import torchvision.transforms as T
import torchvision.transforms.functional as TF
import torch.nn.functional as F
import cv2
import random

class MotionBlur:
    """Apply motion blur to simulate camera/object motion"""
    def __init__(self, kernel_size=15, angle_range=(-45, 45)):
        self.kernel_size = kernel_size
        self.angle_range = angle_range
    
    def __call__(self, img):
        """
        Args:
            img: PIL Image
        Returns:
            PIL Image with motion blur applied
        """
        import random
        import math
        
        # Convert to tensor
        img_tensor = T.ToTensor()(img).unsqueeze(0)  # (1, C, H, W)
        
        # Random angle for motion direction
        angle = random.uniform(*self.angle_range)
        
        # Create motion blur kernel
        kernel = torch.zeros((self.kernel_size, self.kernel_size))
        center = self.kernel_size // 2
        
        # Calculate motion direction
        angle_rad = math.radians(angle)
        cos_angle = math.cos(angle_rad)
        sin_angle = math.sin(angle_rad)
        
        # Draw a line in the kernel (motion blur direction)
        for i in range(self.kernel_size):
            offset = i - center
            x = int(center + offset * cos_angle)
            y = int(center + offset * sin_angle)
            if 0 <= x < self.kernel_size and 0 <= y < self.kernel_size:
                kernel[y, x] = 1.0
        
        # Normalize kernel
        kernel = kernel / kernel.sum()
        
        # Expand kernel for each channel
        kernel = kernel.view(1, 1, self.kernel_size, self.kernel_size)
        kernel = kernel.repeat(img_tensor.shape[1], 1, 1, 1)  # (C, 1, K, K)
        
        # Apply convolution (motion blur)
        padding = self.kernel_size // 2
        blurred = F.conv2d(img_tensor, kernel, padding=padding, groups=img_tensor.shape[1])
        
        # Convert back to PIL
        blurred = blurred.squeeze(0).clamp(0, 1)
        blurred = T.ToPILImage()(blurred)
        
        return blurred


class PipelineUnit:
    def __init__(
        self,
        seperate_cfg: bool = False,
        take_over: bool = False,
        input_params: tuple[str] = None,
        output_params: tuple[str] = None,
        input_params_posi: dict[str, str] = None,
        input_params_nega: dict[str, str] = None,
        onload_model_names: tuple[str] = None
    ):
        self.seperate_cfg = seperate_cfg
        self.take_over = take_over
        self.input_params = input_params
        self.output_params = output_params
        self.input_params_posi = input_params_posi
        self.input_params_nega = input_params_nega
        self.onload_model_names = onload_model_names

    def fetch_input_params(self):
        params = []
        if self.input_params is not None:
            for param in self.input_params:
                params.append(param)
        if self.input_params_posi is not None:
            for _, param in self.input_params_posi.items():
                params.append(param)
        if self.input_params_nega is not None:
            for _, param in self.input_params_nega.items():
                params.append(param)
        params = sorted(list(set(params)))
        return params
    
    def fetch_output_params(self):
        params = []
        if self.output_params is not None:
            for param in self.output_params:
                params.append(param)
        return params

    def process(self, pipe, **kwargs) -> dict:
        return {}
    
    def post_process(self, pipe, **kwargs) -> dict:
        return {}


class BasePipeline(torch.nn.Module):

    def __init__(
        self,
        device="cuda", torch_dtype=torch.float16,
        height_division_factor=64, width_division_factor=64,
        time_division_factor=None, time_division_remainder=None,
    ):
        super().__init__()
        # The device and torch_dtype is used for the storage of intermediate variables, not models.
        self.device = device
        self.torch_dtype = torch_dtype
        # The following parameters are used for shape check.
        self.height_division_factor = height_division_factor
        self.width_division_factor = width_division_factor
        self.time_division_factor = time_division_factor
        self.time_division_remainder = time_division_remainder
        # VRAM management
        self.vram_management_enabled = False
        # Pipeline Unit Runner
        self.unit_runner = PipelineUnitRunner()
        # LoRA Loader
        self.lora_loader = GeneralLoRALoader
        
        
    def to(self, *args, **kwargs):
        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
        if device is not None:
            self.device = device
        if dtype is not None:
            self.torch_dtype = dtype
        super().to(*args, **kwargs)
        return self


    def check_resize_height_width(self, height, width, num_frames=None):
        # Shape check
        if height % self.height_division_factor != 0:
            height = (height + self.height_division_factor - 1) // self.height_division_factor * self.height_division_factor
            print(f"height % {self.height_division_factor} != 0. We round it up to {height}.")
        if width % self.width_division_factor != 0:
            width = (width + self.width_division_factor - 1) // self.width_division_factor * self.width_division_factor
            print(f"width % {self.width_division_factor} != 0. We round it up to {width}.")
        if num_frames is None:
            return height, width
        else:
            if num_frames % self.time_division_factor != self.time_division_remainder:
                num_frames = (num_frames + self.time_division_factor - 1) // self.time_division_factor * self.time_division_factor + self.time_division_remainder
                print(f"num_frames % {self.time_division_factor} != {self.time_division_remainder}. We round it up to {num_frames}.")
            return height, width, num_frames



    def apply_augmentation_to_images_condition(self, images, same_augmentation=True, k=50, augmentation_strength=1.0):
        def sample_params():
            return {
                "apply_color_cast": random.random() < 0.8,
                "color_cast": [
                    random.uniform(-0.1, 0.1),
                    random.uniform(-0.1, 0.1),
                    random.uniform(-0.1, 0.1),
                ],
                "apply_blur": random.random() < 0.5,
                "blur_kernel_size": random.choice([3, 5]),
                "blur_sigma": random.uniform(0.15, 0.8),
            }

        def apply_color_cast(img_tensor, params, strength):
            if not params.get("apply_color_cast", False) or strength <= 0:
                return img_tensor

            # Mild RGB bias to simulate gentle color shift.
            cast = torch.tensor(
                params["color_cast"], dtype=img_tensor.dtype, device=img_tensor.device
            ).view(3, 1, 1)
            cast = cast * strength
            return (img_tensor + cast).clamp(0, 1)

        def apply_blur(img_tensor, params, strength):
            if not params.get("apply_blur", False) or strength <= 0:
                return img_tensor
            kernel_size = params["blur_kernel_size"]
            sigma = max(1e-6, params["blur_sigma"] * strength)
            return TF.gaussian_blur(
                img_tensor,
                kernel_size=[kernel_size, kernel_size],
                sigma=[sigma, sigma],
            )

        def apply_params(img, params, strength):
            img_tensor = TF.to_tensor(img)
            try:
                if img_tensor.ndim == 3 and img_tensor.shape[0] == 3:
                    img_tensor = apply_color_cast(img_tensor, params, strength)
                    img_tensor = apply_blur(img_tensor, params, strength)
            except Exception as e:
                print(f"Video augmentation failed: {e}")
            return img_tensor.clamp(0, 1)

        shared_params = sample_params() if same_augmentation else None

        augmented_images = []
        start_augmentation_idx = max(0, int(k))
        strength = float(max(0.0, min(1.0, augmentation_strength)))
        for idx, image in enumerate(images):
            if idx >= start_augmentation_idx:
                params = shared_params if same_augmentation else sample_params()
                augmented_tensor = apply_params(image, params, strength)
            else:
                augmented_tensor = TF.to_tensor(image)
            tensor_image = augmented_tensor.mul(2.0).sub(1.0).unsqueeze(0)
            augmented_images.append(tensor_image)

        return augmented_images


    def apply_augmentation_to_images(self, images, same_augmentation=True, k=50, augmentation_strength=1.0):
        def sample_patch_specs(params):
            min_ratio, max_ratio = params["erase_patch_ratio_range"]
            min_downsample, max_downsample = params["erase_downsample_factor_range"]
            specs = []
            for _ in range(params["erase_patch_count"]):
                specs.append({
                    "height_ratio": random.uniform(min_ratio, max_ratio),
                    "width_ratio": random.uniform(min_ratio, max_ratio),
                    "top_ratio": random.uniform(0.0, 1.0),
                    "left_ratio": random.uniform(0.0, 1.0),
                    "downsample_factor": random.randint(min_downsample, max_downsample),
                })
            return specs

        def sample_params():
            params = {
                "apply_clahe": random.random() < 0.9,
                "clahe_clip_limit": random.uniform(1.0, 64.0),
                "clahe_tile_grid_size": random.randint(1, 32),
                "apply_oversaturation": random.random() < 0.3,
                "saturation_boost": random.uniform(1.3, 3.0),
                "apply_brightness": random.random() < 0.1,
                "brightness_scale": random.uniform(0.8, 1.2),

                "apply_blur": random.random() < 0.0,
                "blur_kernel_size": random.choice([3, 5, 7]),
                "blur_sigma": random.uniform(0.2, 1.2),

                "apply_detail_erase": random.random() < 0.9,
                "erase_patch_count": random.randint(1, 4),
                "erase_patch_ratio_range": (0.1, 0.3),
                "erase_downsample_factor_range": (8, 32),
            }
            params["erase_patch_specs"] = sample_patch_specs(params)
            return params

        def apply_clahe(img_tensor, params):
            if not params.get("apply_clahe", False):
                return img_tensor
            img_np = (img_tensor.permute(1, 2, 0).clamp(0, 1).cpu().numpy() * 255.0).astype(np.uint8)
            lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
            l, a, b = cv2.split(lab)
            tile_grid_size = params["clahe_tile_grid_size"]
            clahe = cv2.createCLAHE(
                clipLimit=params["clahe_clip_limit"],
                tileGridSize=(tile_grid_size, tile_grid_size),
            )
            l = clahe.apply(l)
            lab = cv2.merge((l, a, b))
            out_np = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
            out_tensor = torch.from_numpy(out_np).to(dtype=torch.float32) / 255.0
            return out_tensor.permute(2, 0, 1)

        def apply_color_jitter(img_tensor, params, strength):
            apply_oversaturation = params.get("apply_oversaturation", False)
            apply_brightness = params.get("apply_brightness", False)
            if not apply_oversaturation and not apply_brightness:
                return img_tensor
            out = img_tensor
            if apply_oversaturation:
                saturation_factor = 1.0 + (params["saturation_boost"] - 1.0) * strength
                out = TF.adjust_saturation(out, saturation_factor)
            if apply_brightness:
                brightness_factor = 1.0 + (params["brightness_scale"] - 1.0) * strength
                out = TF.adjust_brightness(out, brightness_factor)
            return out.clamp(0, 1)

        def apply_blur(img_tensor, params, strength):
            if not params.get("apply_blur", False) or strength <= 0:
                return img_tensor
            kernel_size = params["blur_kernel_size"]
            sigma = max(1e-6, params["blur_sigma"] * strength)
            return TF.gaussian_blur(img_tensor, kernel_size=[kernel_size, kernel_size], sigma=[sigma, sigma])

        def apply_detail_erase(img_tensor, params, strength):
            if not params.get("apply_detail_erase", False):
                return img_tensor
            if strength <= 0:
                return img_tensor
            out = img_tensor.clone()
            _, h, w = out.shape
            patch_specs = params["erase_patch_specs"]
            patch_count = max(1, int(round(params["erase_patch_count"] * strength)))
            for patch_spec in patch_specs[:patch_count]:
                patch_h = max(8, int(h * patch_spec["height_ratio"]))
                patch_w = max(8, int(w * patch_spec["width_ratio"]))
                patch_h = min(patch_h, h)
                patch_w = min(patch_w, w)
                max_top = max(0, h - patch_h)
                max_left = max(0, w - patch_w)
                y1 = min(max_top, int(max_top * patch_spec["top_ratio"]))
                x1 = min(max_left, int(max_left * patch_spec["left_ratio"]))
                y2 = y1 + patch_h
                x2 = x1 + patch_w
                patch = out[:, y1:y2, x1:x2].unsqueeze(0)
                downsample_factor = patch_spec["downsample_factor"]
                small_h = max(1, patch_h // downsample_factor)
                small_w = max(1, patch_w // downsample_factor)
                patch_small = torch.nn.functional.interpolate(patch, size=(small_h, small_w), mode="area")
                patch_restore = torch.nn.functional.interpolate(patch_small, size=(patch_h, patch_w), mode="bilinear", align_corners=False)
                out[:, y1:y2, x1:x2] = patch_restore.squeeze(0)
            return out.clamp(0, 1)

        def apply_params(img, params, strength):
            img_tensor = TF.to_tensor(img)
            try:
                if img_tensor.ndim == 3 and img_tensor.shape[0] == 3:
                    img_tensor = apply_clahe(img_tensor, params)
                    img_tensor = apply_color_jitter(img_tensor, params, strength)
                    img_tensor = apply_blur(img_tensor, params, strength)
                    img_tensor = apply_detail_erase(img_tensor, params, strength)
            except Exception as e:
                print(f"Video augmentation failed: {e}")
            return img_tensor.clamp(0, 1)

        shared_params = sample_params() if same_augmentation else None

        augmented_images = []
        start_augmentation_idx = max(0, int(k))
        strength = float(max(0.0, min(1.0, augmentation_strength)))
        for idx, image in enumerate(images):
            if idx >= start_augmentation_idx:
                params = shared_params if same_augmentation else sample_params()
                augmented_tensor = apply_params(image, params, strength)
            else:
                augmented_tensor = TF.to_tensor(image)
            tensor_image = augmented_tensor.mul(2.0).sub(1.0).unsqueeze(0)
            augmented_images.append(tensor_image)

        return augmented_images



    def preprocess_image(self, image, torch_dtype=None, device=None, pattern="B C H W", min_value=-1, max_value=1):
        # Transform a PIL.Image to torch.Tensor
        image = torch.Tensor(np.array(image, dtype=np.float32))
        image = image.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device)
        image = image * ((max_value - min_value) / 255) + min_value
        image = repeat(image, f"H W C -> {pattern}", **({"B": 1} if "B" in pattern else {}))
        return image


    def preprocess_video(self, video, torch_dtype=None, device=None, pattern="B C T H W", min_value=-1, max_value=1):
        # Transform a list of PIL.Image to torch.Tensor
        video = [self.preprocess_image(image, torch_dtype=torch_dtype, device=device, min_value=min_value, max_value=max_value) for image in video]
        video = torch.stack(video, dim=pattern.index("T") // 2)
        return video


    def vae_output_to_image(self, vae_output, pattern="B C H W", min_value=-1, max_value=1):
        # Transform a torch.Tensor to PIL.Image
        if pattern != "H W C":
            vae_output = reduce(vae_output, f"{pattern} -> H W C", reduction="mean")
        image = ((vae_output - min_value) * (255 / (max_value - min_value))).clip(0, 255)
        image = image.to(device="cpu", dtype=torch.uint8)
        image = Image.fromarray(image.numpy())
        return image


    def vae_output_to_video(self, vae_output, pattern="B C T H W", min_value=-1, max_value=1):
        # Transform a torch.Tensor to list of PIL.Image
        if pattern != "T H W C":
            vae_output = reduce(vae_output, f"{pattern} -> T H W C", reduction="mean")
        video = [self.vae_output_to_image(image, pattern="H W C", min_value=min_value, max_value=max_value) for image in vae_output]
        return video


    def load_models_to_device(self, model_names):
        if self.vram_management_enabled:
            # offload models
            for name, model in self.named_children():
                if name not in model_names:
                    if hasattr(model, "vram_management_enabled") and model.vram_management_enabled:
                        if hasattr(model, "offload"):
                            model.offload()
                        else:
                            for module in model.modules():
                                if hasattr(module, "offload"):
                                    module.offload()
            torch.cuda.empty_cache()
            # onload models
            for name, model in self.named_children():
                if name in model_names:
                    if hasattr(model, "vram_management_enabled") and model.vram_management_enabled:
                        if hasattr(model, "onload"):
                            model.onload()
                        else:
                            for module in model.modules():
                                if hasattr(module, "onload"):
                                    module.onload()


    def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None):
        # Initialize Gaussian noise
        generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed)
        noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype)
        noise = noise.to(dtype=torch_dtype or self.torch_dtype, device=device or self.device)
        return noise

        
    def get_vram(self):
        return torch.cuda.mem_get_info(self.device)[1] / (1024 ** 3)
    
    def get_module(self, model, name):
        if "." in name:
            name, suffix = name[:name.index(".")], name[name.index(".") + 1:]
            if name.isdigit():
                return self.get_module(model[int(name)], suffix)
            else:
                return self.get_module(getattr(model, name), suffix)
        else:
            return getattr(model, name)
    
    def freeze_except(self, model_names):
        self.eval()
        self.requires_grad_(False)
        for name in model_names:
            module = self.get_module(self, name)
            if module is None:
                print(f"No {name} models in the pipeline. We cannot enable training on the model. If this occurs during the data processing stage, it is normal.")
                continue
            module.train()
            module.requires_grad_(True)
                
    
    def blend_with_mask(self, base, addition, mask):
        return base * (1 - mask) + addition * mask
    
    
    def step(self, scheduler, latents, progress_id, noise_pred, input_latents=None, inpaint_mask=None, **kwargs):
        timestep = scheduler.timesteps[progress_id]
        if inpaint_mask is not None:
            noise_pred_expected = scheduler.return_to_timestep(scheduler.timesteps[progress_id], latents, input_latents)
            noise_pred = self.blend_with_mask(noise_pred_expected, noise_pred, inpaint_mask)
        latents_next = scheduler.step(noise_pred, timestep, latents)
        return latents_next
    
    
    def split_pipeline_units(self, model_names: list[str]):
        return PipelineUnitGraph().split_pipeline_units(self.units, model_names)
    
    
    def flush_vram_management_device(self, device):
        for module in self.modules():
            if isinstance(module, AutoTorchModule):
                module.offload_device = device
                module.onload_device = device
                module.preparing_device = device
                module.computation_device = device
                
    
    def load_lora(
        self,
        module: torch.nn.Module,
        lora_config: Union[ModelConfig, str] = None,
        alpha=1,
        hotload=None,
        state_dict=None,
    ):
        if state_dict is None:
            if isinstance(lora_config, str):
                lora = load_state_dict(lora_config, torch_dtype=self.torch_dtype, device=self.device)
            else:
                lora_config.download_if_necessary()
                lora = load_state_dict(lora_config.path, torch_dtype=self.torch_dtype, device=self.device)
        else:
            lora = state_dict
        lora_loader = self.lora_loader(torch_dtype=self.torch_dtype, device=self.device)
        lora = lora_loader.convert_state_dict(lora)
        if hotload is None:
            hotload = hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled")
        if hotload:
            if not (hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled")):
                raise ValueError("VRAM Management is not enabled. LoRA hotloading is not supported.")
            updated_num = 0
            for _, module in module.named_modules():
                if isinstance(module, AutoWrappedLinear):
                    name = module.name
                    lora_a_name = f'{name}.lora_A.weight'
                    lora_b_name = f'{name}.lora_B.weight'
                    if lora_a_name in lora and lora_b_name in lora:
                        updated_num += 1
                        module.lora_A_weights.append(lora[lora_a_name] * alpha)
                        module.lora_B_weights.append(lora[lora_b_name])
            print(f"{updated_num} tensors are patched by LoRA. You can use `pipe.clear_lora()` to clear all LoRA layers.")
        else:
            lora_loader.fuse_lora_to_base_model(module, lora, alpha=alpha)
            
            
    def clear_lora(self):
        cleared_num = 0
        for name, module in self.named_modules():
            if isinstance(module, AutoWrappedLinear):
                if hasattr(module, "lora_A_weights"):
                    if len(module.lora_A_weights) > 0:
                        cleared_num += 1
                    module.lora_A_weights.clear()
                if hasattr(module, "lora_B_weights"):
                    module.lora_B_weights.clear()
        print(f"{cleared_num} LoRA layers are cleared.")
        
    
    def download_and_load_models(self, model_configs: list[ModelConfig] = [], vram_limit: float = None):
        model_pool = ModelPool()
        for model_config in model_configs:
            model_config.download_if_necessary()
            vram_config = model_config.vram_config()
            vram_config["computation_dtype"] = vram_config["computation_dtype"] or self.torch_dtype
            vram_config["computation_device"] = vram_config["computation_device"] or self.device
            model_pool.auto_load_model(
                model_config.path,
                vram_config=vram_config,
                vram_limit=vram_limit,
                clear_parameters=model_config.clear_parameters,
            )
        return model_pool
    
    
    def check_vram_management_state(self):
        vram_management_enabled = False
        for module in self.children():
            if hasattr(module, "vram_management_enabled") and getattr(module, "vram_management_enabled"):
                vram_management_enabled = True
        return vram_management_enabled
    
    
    def cfg_guided_model_fn(self, model_fn, cfg_scale, inputs_shared, inputs_posi, inputs_nega, **inputs_others):
        noise_pred_posi = model_fn(**inputs_posi, **inputs_shared, **inputs_others)
        if cfg_scale != 1.0:
            noise_pred_nega = model_fn(**inputs_nega, **inputs_shared, **inputs_others)
            noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega)
        else:
            noise_pred = noise_pred_posi
        return noise_pred


class PipelineUnitGraph:
    def __init__(self):
        pass
    
    def build_edges(self, units: list[PipelineUnit]):
        # Establish dependencies between units
        # to search for subsequent related computation units.
        last_compute_unit_id = {}
        edges = []
        for unit_id, unit in enumerate(units):
            for input_param in unit.fetch_input_params():
                if input_param in last_compute_unit_id:
                    edges.append((last_compute_unit_id[input_param], unit_id))
            for output_param in unit.fetch_output_params():
                last_compute_unit_id[output_param] = unit_id
        return edges
    
    def build_chains(self, units: list[PipelineUnit]):
        # Establish updating chains for each variable
        # to track their computation process.
        params = sum([unit.fetch_input_params() + unit.fetch_output_params() for unit in units], [])
        params = sorted(list(set(params)))
        chains = {param: [] for param in params}
        for unit_id, unit in enumerate(units):
            for param in unit.fetch_output_params():
                chains[param].append(unit_id)
        return chains
    
    def search_direct_unit_ids(self, units: list[PipelineUnit], model_names: list[str]):
        # Search for units that directly participate in the model's computation.
        related_unit_ids = []
        for unit_id, unit in enumerate(units):
            for model_name in model_names:
                if unit.onload_model_names is not None and model_name in unit.onload_model_names:
                    related_unit_ids.append(unit_id)
                    break
        return related_unit_ids
    
    def search_related_unit_ids(self, edges, start_unit_ids, direction="target"):
        # Search for subsequent related computation units.
        related_unit_ids = [unit_id for unit_id in start_unit_ids]
        while True:
            neighbors = []
            for source, target in edges:
                if direction == "target" and source in related_unit_ids and target not in related_unit_ids:
                    neighbors.append(target)
                elif direction == "source" and source not in related_unit_ids and target in related_unit_ids:
                    neighbors.append(source)
            neighbors = sorted(list(set(neighbors)))
            if len(neighbors) == 0:
                break
            else:
                related_unit_ids.extend(neighbors)
        related_unit_ids = sorted(list(set(related_unit_ids)))
        return related_unit_ids
    
    def search_updating_unit_ids(self, units: list[PipelineUnit], chains, related_unit_ids):
        # If the input parameters of this subgraph are updated outside the subgraph,
        # search for the units where these updates occur.
        first_compute_unit_id = {}
        for unit_id in related_unit_ids:
            for param in units[unit_id].fetch_input_params():
                if param not in first_compute_unit_id:
                    first_compute_unit_id[param] = unit_id
        updating_unit_ids = []
        for param in first_compute_unit_id:
            unit_id = first_compute_unit_id[param]
            chain = chains[param]
            if unit_id in chain and chain.index(unit_id) != len(chain) - 1:
                for unit_id_ in chain[chain.index(unit_id) + 1:]:
                    if unit_id_ not in related_unit_ids:
                        updating_unit_ids.append(unit_id_)
        related_unit_ids.extend(updating_unit_ids)
        related_unit_ids = sorted(list(set(related_unit_ids)))
        return related_unit_ids
    
    def split_pipeline_units(self, units: list[PipelineUnit], model_names: list[str]):
        # Split the computation graph,
        # separating all model-related computations.
        related_unit_ids = self.search_direct_unit_ids(units, model_names)
        edges = self.build_edges(units)
        chains = self.build_chains(units)
        while True:
            num_related_unit_ids = len(related_unit_ids)
            related_unit_ids = self.search_related_unit_ids(edges, related_unit_ids, "target")
            related_unit_ids = self.search_updating_unit_ids(units, chains, related_unit_ids)
            if len(related_unit_ids) == num_related_unit_ids:
                break
            else:
                num_related_unit_ids = len(related_unit_ids)
        related_units = [units[i] for i in related_unit_ids]
        unrelated_units = [units[i] for i in range(len(units)) if i not in related_unit_ids]
        return related_units, unrelated_units


class PipelineUnitRunner:
    def __init__(self):
        pass

    def __call__(self, unit: PipelineUnit, pipe: BasePipeline, inputs_shared: dict, inputs_posi: dict, inputs_nega: dict) -> tuple[dict, dict]:
        if unit.take_over:
            # Let the pipeline unit take over this function.
            inputs_shared, inputs_posi, inputs_nega = unit.process(pipe, inputs_shared=inputs_shared, inputs_posi=inputs_posi, inputs_nega=inputs_nega)
        elif unit.seperate_cfg:
            # Positive side
            processor_inputs = {name: inputs_posi.get(name_) for name, name_ in unit.input_params_posi.items()}
            if unit.input_params is not None:
                for name in unit.input_params:
                    processor_inputs[name] = inputs_shared.get(name)
            processor_outputs = unit.process(pipe, **processor_inputs)
            inputs_posi.update(processor_outputs)
            # Negative side
            if inputs_shared["cfg_scale"] != 1:
                processor_inputs = {name: inputs_nega.get(name_) for name, name_ in unit.input_params_nega.items()}
                if unit.input_params is not None:
                    for name in unit.input_params:
                        processor_inputs[name] = inputs_shared.get(name)
                processor_outputs = unit.process(pipe, **processor_inputs)
                inputs_nega.update(processor_outputs)
            else:
                inputs_nega.update(processor_outputs)
        else:
            processor_inputs = {name: inputs_shared.get(name) for name in unit.input_params}
            processor_outputs = unit.process(pipe, **processor_inputs)
            inputs_shared.update(processor_outputs)
        return inputs_shared, inputs_posi, inputs_nega