Fixed minor bug in shape validation in SongUNet

[CharlelieLrt]

PhysicsNeMo Pull Request

Description

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

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Overview

Greptile Summary

Fixed an off-by-one error in the shape validation logic for SongUNet. The validation now correctly requires input dimensions to be multiples of 2^(N-1) instead of 2^N, matching the actual downsampling behavior of the U-Net architecture.

• Updated _input_shape_mult calculation from 2**self._num_levels to 2 ** (self._num_levels - 1) in physicsnemo/models/diffusion/song_unet.py:340
• Updated documentation and error messages to reflect the correct formula
• The bug was introduced in PR Added input shape validation for SongUnet #1181 and is now corrected

The fix is mathematically correct: with N levels and level 0 having no downsampling, the architecture performs (N-1) downsampling operations, resulting in a total downsampling factor of 2^(N-1).

Confidence Score: 5/5

• This PR is safe to merge with high confidence
• The change is a straightforward mathematical correction that fixes an off-by-one error in shape validation. The fix aligns the validation logic with the actual U-Net architecture implementation, which performs (N-1) downsampling operations for N levels. The changes are minimal, well-contained, and affect only the validation logic without modifying the core architecture behavior.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
physicsnemo/models/diffusion/song_unet.py	5/5	corrected shape validation from 2^N to 2^(N-1) to match actual downsampling in U-Net architecture

Sequence Diagram

sequenceDiagram
    participant User
    participant SongUNet
    participant Validation
    participant Encoder
    participant Decoder

    User->>SongUNet: forward(x, noise_labels, class_labels)
    SongUNet->>Validation: Check input shape dimensions
    
    alt Level 0 (no downsampling)
        Note over Validation: Resolution = H
    end
    
    alt Level 1
        Note over Validation: Resolution = H/2 (1 downsample)
    end
    
    alt Level N-1
        Note over Validation: Resolution = H/2^(N-1)
    end
    
    Validation->>Validation: Verify d is power of 2 OR multiple of 2^(N-1)
    
    alt Valid dimensions
        Validation->>Encoder: Process through encoder levels
        Encoder->>Decoder: Pass through N levels with skip connections
        Decoder->>SongUNet: Return denoised output
        SongUNet->>User: Output tensor
    else Invalid dimensions
        Validation->>User: ValueError: dimensions must be powers of 2 or multiples of 2^(N-1)
    end

Loading

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[coreyjadams]

LGTM

[CharlelieLrt]

/blossom-ci