Bug with recover_dilate_module

[zwishenzug]

make_dilate_model calls recover_dilate_module(model)

That looks at all the named submodules in order to restore the original padding/dilation

However, when doing the restore, it's actually looking at all parts of the model, not just the UNET. So it's also finding some 3x3 convs in the VAE and setting those back to a padding of (1, 1)

But the problem is some of them were originally (0,0)

For example, if I put some debugging code:

        print(f"recover_dilate_module: {submodule.dilation}|{submodule.padding}|{name}")
        submodule.dilation = (1, 1)
        submodule.padding = (1, 1)

I see this

recover_dilate_module: (1, 1)|(0, 0)|first_stage_model.encoder.down.1.downsample.conv

and there are others too.

Probably when you call

recover_dilate_module(model)

You should really be calling

recover_dilate_module(model.model.diffusion_model)

Although I'll let you decide the best fix.

[YingqingHe]

Thanks for pointing it out! It should be recover_dilate_module(model.model.diffusion_model) which is a bug during refactoring. The VAE indeed has original convolutions with padding (0,0). Fortunately, they are only in the Encoder part, which is not used in our sampling, so the results should not be affected. I have fixed this bug. Thank you!