fluxdev/stable-diffusion-webui-forge

1import torch2 3 4def apply_controlnet_advanced(5        unet,6        controlnet,7        image_bchw,8        strength,9        start_percent,10        end_percent,11        positive_advanced_weighting=None,12        negative_advanced_weighting=None,13        advanced_frame_weighting=None,14        advanced_sigma_weighting=None,15        advanced_mask_weighting=None16):17    """18 19    # positive_advanced_weighting or negative_advanced_weighting20 21    Unet has input, middle, output blocks, and we can give different weights to each layers in all blocks.22    Below is an example for stronger control in middle block.23    This is helpful for some high-res fix passes.24 25        positive_advanced_weighting = {26            'input': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2],27            'middle': [1.0],28            'output': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2]29        }30        negative_advanced_weighting = {31            'input': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2],32            'middle': [1.0],33            'output': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2]34        }35 36    # advanced_frame_weighting37 38    The advanced_frame_weighting is a weight applied to each image in a batch.39    The length of this list must be same with batch size40    For example, if batch size is 5, you can use advanced_frame_weighting = [0, 0.25, 0.5, 0.75, 1.0]41    If you view the 5 images as 5 frames in a video, this will lead to progressively stronger control over time.42 43    # advanced_sigma_weighting44 45    The advanced_sigma_weighting allows you to dynamically compute control46    weights given diffusion timestep (sigma).47    For example below code can softly make beginning steps stronger than ending steps.48 49        sigma_max = unet.model.model_sampling.sigma_max50        sigma_min = unet.model.model_sampling.sigma_min51        advanced_sigma_weighting = lambda s: (s - sigma_min) / (sigma_max - sigma_min)52 53    # advanced_mask_weighting54 55    A mask can be applied to control signals.56    This should be a tensor with shape B 1 H W where the H and W can be arbitrary.57    This mask will be resized automatically to match the shape of all injection layers.58 59    """60 61    cnet = controlnet.copy().set_cond_hint(image_bchw, strength, (start_percent, end_percent))62    cnet.positive_advanced_weighting = positive_advanced_weighting63    cnet.negative_advanced_weighting = negative_advanced_weighting64    cnet.advanced_frame_weighting = advanced_frame_weighting65    cnet.advanced_sigma_weighting = advanced_sigma_weighting66 67    if advanced_mask_weighting is not None:68        assert isinstance(advanced_mask_weighting, torch.Tensor)69        B, C, H, W = advanced_mask_weighting.shape70        assert B > 0 and C == 1 and H > 0 and W > 071 72    cnet.advanced_mask_weighting = advanced_mask_weighting73 74    m = unet.clone()75    m.add_patched_controlnet(cnet)76    return m77 78