Controlnet with a conditioning strength mask

[jschoormans]

Hi everyone,

Inspired by https://github.com/Kosinkadink/ComfyUI-Advanced-ControlNet, I wanted to add a mask to Controlnets in Diffusers. I could not find it implemented yet. Please do tell me if this was already possible in Diffusers!

The approach is simply weighting the controlnet control image with a predefined mask, for instance in this case where I want to preserve the structure of a room with controlnet depth - but no furniture.

The implementation is simply done in Step 5. Controlnet_blocks in controlnet.py

        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
            down_block_res_sample = controlnet_block(down_block_res_sample)
            # *** APPLY MASK HERE ***
            if mask_optional is not None:
                mask = F.interpolate(mask_optional, size=(down_block_res_sample.shape[2], down_block_res_sample.shape[3]), mode="bilinear", align_corners=False)
                # Ensure mask has the same number of channels as down_block_res_sample
                mask = mask.repeat(down_block_res_sample.shape[0], down_block_res_sample.shape[1], 1, 1)
                down_block_res_sample = down_block_res_sample * mask

            
            controlnet_down_block_res_samples = controlnet_down_block_res_samples + (down_block_res_sample,)



        down_block_res_samples = controlnet_down_block_res_samples

        mid_block_res_sample = self.controlnet_mid_block(sample)
        if mask_optional is not None:
            mask = F.interpolate(mask_optional, size=mid_block_res_sample.shape[-2:], mode="bilinear", align_corners=False)            
            mid_block_res_sample = mid_block_res_sample * mask