Fix corr-diff conditional scaling

[nbren12]

Modulus Pull Request

Description

This PR fixes the scaling for conditional inputs in corrdiff. Previously the conditional inputs were concatenated with the noisy target BEFORE scaling. This means that for large noise the conditional inputs will also be scaled down to near 0. This is undesirable since the large-noise limits corresponds to deterministic prediction---we expect the conditional inputs to matter most in this case! This PR fixes this by scaling the noise image BEFORE concatenating it with the conditioning.

I wrote this PR with unit testing in mind, so I wrote a modular class _ConditionalPrecond which is testable in isolation. I added a public API EDMPrecondSRV2 to provide a similar interface as EDMPrecondSR that also builds the UNet.

Closes #229.

Checklist

• [ ] I am familiar with the Contributing Guidelines.
• [ ] New or existing tests cover these changes.
• [ ] The documentation is up to date with these changes.
• [ ] The CHANGELOG.md is up to date with these changes.
• [ ] An issue is linked to this pull request.

Dependencies

[nbren12]

BTW, I've used this code preconditioner in some other projects with good results.

[MortezaMardani]

makes sense to me.

[yairchn]

/blossom-ci

[nbren12]

/blossom-ci

[ktangsali]

/blossom-ci

[ktangsali]

/blossom-ci