Skip to content

Remove exploratory/branched implementations from main (VQ-VAE, original ResNets, EMA, pre-transformer UNet) - #12

Open
jlotthammer wants to merge 3 commits into
mainfrom
cleanup/remove-exploratory-code
Open

Remove exploratory/branched implementations from main (VQ-VAE, original ResNets, EMA, pre-transformer UNet)#12
jlotthammer wants to merge 3 commits into
mainfrom
cleanup/remove-exploratory-code

Conversation

@jlotthammer

Copy link
Copy Markdown
Contributor

Summary

Removes exploratory / branched implementations from main that were not part of the deployed STARLING system, leaving only the production model code (VAE distance-map autoencoder + ViT/DiT latent diffusion) and the training code used to produce it.

Full pre-cleanup state is preserved on the archive/pre-cleanup branch (tag archive/pre-cleanup-2026-07-16).

Net change: −1937 lines, no production code paths altered.

What was removed

Orphaned model modules (zero importers in production or training):

  • models/quantize.py — VQ-VAE VectorQuantizer2
  • models/resnets_original.py — pre-vae_components ResNet enc/dec
  • models/ema.pyEMA, never wired into training or diffusion

Pre-transformer UNet backbone (superseded by the ViT/DiT transformer):

  • models/unet.py, configs/unet/unet.yaml (+ its configs.yaml defaults entry)
  • Its only remaining tie to production was vit.py importing SinusoidalPosEmb; that class is byte-identical to the one already in transformer.py, so the import was redirected there.
  • Dropped the unused UNetConditional construction in diffusion_train.py (the real backbone passed to the diffusion model was already the ViT).

Broken references / bit-rot:

  • pyproject.toml: removed starling-sample and ae-train console scripts (targeted non-existent modules training/vae_generate.py, training/ae_train.py).
  • diffusion_train.py: removed import starling.data.ddpm_loader (module does not exist — only ddpm_loader_tar is present; this broke starling-ddpm-train at import time).

Docs: removed the Sphinx autosummary stubs and api.rst entries for all deleted modules.

What was deliberately kept

  • The core VAE (vae.py, vae_components.py, distributions.py) — this is the production distance-map autoencoder, loaded at inference. Only the VQ-VAE variant was exploratory.
  • evaluate_vae.py + benchmark_mds.py — wired to the live starling-benchmark CLI.
  • devtools/scripts/ — left untouched.

Verification

  • All tracked .py byte-compile; import starling + all affected modules import at runtime; ViT(12,512,8,512) constructs.
  • All 14 remaining pyproject console scripts resolve to real module:function targets.
  • No dangling references to any removed module remain.
  • Fast test suite: 138 passed, 2 skipped.

Out of scope (pre-existing, not caused by this PR)

8 stochastic generation/reconstruction tests in test_starling.py fail because ~19/100 generated conformations contain NaN on this environment (torch 2.13.0). Confirmed identical failures on a clean main worktree, so this is unrelated to the cleanup and tracked separately.

…EMA)

These modules had zero importers in the production or training paths:
- models/quantize.py (VQ-VAE VectorQuantizer2)
- models/resnets_original.py (pre-vae_components ResNet enc/dec)
- models/ema.py (EMA, never wired into training/diffusion)

Also drops their Sphinx autosummary stubs and api.rst entries.
Preserved on archive/pre-cleanup.
The production diffusion backbone is the ViT/DiT transformer; the UNet was
the pre-transformer implementation. Its only remaining tie to production was
vit.py importing SinusoidalPosEmb, which is an identical class already defined
in transformer.py.

- vit.py: import SinusoidalPosEmb from transformer instead of unet
- diffusion_train.py: drop the unused UNetConditional construction (the real
  backbone passed to the diffusion model was already the ViT); dump the actual
  diffusion model architecture to model_architecture.txt
- delete models/unet.py, configs/unet/unet.yaml (and its defaults entry)
- drop unet Sphinx stubs and api.rst entry

Preserved on archive/pre-cleanup.
- pyproject: remove starling-sample and ae-train console scripts; both
  pointed at modules that do not exist (training/vae_generate.py,
  training/ae_train.py)
- diffusion_train.py: remove 'import starling.data.ddpm_loader' (module does
  not exist; only ddpm_loader_tar is present), which broke starling-ddpm-train
  at import time. The alias was never referenced.
@jlotthammer
jlotthammer marked this pull request as ready for review July 16, 2026 17:52
@jlotthammer
jlotthammer requested a review from bnovakborna July 16, 2026 17:55
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.

2 participants