Skip to content

feat: Add model skeleton package with stubs for U-Net, noise schedule…#21

Open
KrishanYadav333 wants to merge 1 commit intoML4SCI:mainfrom
KrishanYadav333:feat/model-skeleton
Open

feat: Add model skeleton package with stubs for U-Net, noise schedule…#21
KrishanYadav333 wants to merge 1 commit intoML4SCI:mainfrom
KrishanYadav333:feat/model-skeleton

Conversation

@KrishanYadav333
Copy link

Pre-GSoC groundwork for the "Denoising Astronomical Observations
of Protoplanetary Disks" project (GSoC 2026).

Following up on the config system PR, this lays out the model
module structure before any real implementation starts.
The idea is to define clean interfaces first so the next PRs (actual
forward pass logic, training loop, etc.) have a clear contract
to fill in rather than figuring out the structure mid-way.

All forward methods raise NotImplementedError for now — that's
intentional, not a bug.

What's included:

  • UNet — constructor signature with in_channels, base_channels,
    channel_multipliers, attention_levels, dropout
  • NoiseScheduler — q_sample, p_sample, p_sample_loop signatures
    with full docstrings explaining what each method should do
  • DDPM — wrapper combining UNet + NoiseScheduler, exposes
    training_loss and sample as the two main entry points
  • blocks.py — ResidualBlock, AttentionBlock, SinusoidalTimeEmbedding
    stubs with documented args and return shapes
  • 20 unit tests covering imports, instantiation, and interface contracts

Files changed:

  • src/models/__init__.py
  • src/models/blocks.py
  • src/models/unet.py
  • src/models/noise_scheduler.py
  • src/models/ddpm.py
  • tests/test_models_skeleton.py

Let me know if the interface design needs any changes before I
start filling in the implementations!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant