Reference from ScaleCrafter[https://arxiv.org/abs/2310.07702]

dilate_conv.py

# https://arxiv.org/abs/2310.07702

import comfy.ops
ops = comfy.ops.disable_weight_init

class DilateConv:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL", ),
                "start": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1, "display": "number"}),
                "end": ("INT", {"default": 300, "min": 0, "max": 1000, "step": 1, "display": "number"}),
            },
        }

    RETURN_TYPES = ("MODEL", )
    FUNCTION = "apply"
    CATEGORY = "loaders"

    def apply(self, model, start, end):
        new_model = model.clone()
        self.replaced_conv2ds = []

        # unet計算前後のパッチ
        def apply_dilate(model_function, kwargs):
            sigmas = kwargs["timestep"]
            t = new_model.model.model_sampling.timestep(sigmas)
            if t[0] < (1000 - end) or t[0] > (1000 - start):
                return model_function(kwargs["input"], kwargs["timestep"], **kwargs["c"])
            
            self.replace_conv2d(new_model)
            retval = model_function(kwargs["input"], kwargs["timestep"], **kwargs["c"])
            self.restore_conv2d(new_model)
            return retval

        new_model.set_model_unet_function_wrapper(apply_dilate)

        return (new_model, )
    
    def replace_conv2d(self, model):
        for name, module in model.model.diffusion_model.named_modules():
            if isinstance(module, ops.Conv2d):
                if module.kernel_size == (3, 3) and module.stride == (1, 1) and module.padding == (1, 1):
                    self.replaced_conv2ds.append(name)
                    module.dilation = (2, 2)
                    module.padding = (2, 2)
    
    def restore_conv2d(self, model):
        for name, module in model.model.diffusion_model.named_modules():
            if name in self.replaced_conv2ds:
                module.dilation = (1, 1)
                module.padding = (1, 1)
        self.replaced_conv2ds = []

NODE_CLASS_MAPPINGS = {
	"DilateConv": DilateConv,
}