From 21eca59dfb0c56c4b3f4a274313858f19ca4bb01 Mon Sep 17 00:00:00 2001 From: Jeff lotthammer Date: Thu, 16 Jul 2026 11:26:00 -0500 Subject: [PATCH 1/4] Remove orphaned exploratory model modules (VQ-VAE, original ResNets, 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. --- docs/autosummary/starling.models.ema.EMA.rst | 28 -- docs/autosummary/starling.models.ema.rst | 32 -- ...rling.models.quantize.VectorQuantizer2.rst | 86 ------ docs/autosummary/starling.models.quantize.rst | 32 -- ...snets_original.ResNet_Decoder_Original.rst | 83 ------ ...snets_original.ResNet_Encoder_Original.rst | 83 ------ ...els.resnets_original.Resnet101_Decoder.rst | 6 - ...els.resnets_original.Resnet101_Encoder.rst | 6 - ...els.resnets_original.Resnet152_Decoder.rst | 6 - ...els.resnets_original.Resnet152_Encoder.rst | 6 - ...dels.resnets_original.Resnet18_Decoder.rst | 6 - ...dels.resnets_original.Resnet18_Encoder.rst | 6 - ...dels.resnets_original.Resnet34_Decoder.rst | 6 - ...dels.resnets_original.Resnet34_Encoder.rst | 6 - ...dels.resnets_original.Resnet50_Decoder.rst | 6 - ...dels.resnets_original.Resnet50_Encoder.rst | 6 - .../starling.models.resnets_original.rst | 50 ---- starling/models/ema.py | 32 -- starling/models/quantize.py | 148 ---------- starling/models/resnets_original.py | 276 ------------------ 20 files changed, 910 deletions(-) delete mode 100644 docs/autosummary/starling.models.ema.EMA.rst delete mode 100644 docs/autosummary/starling.models.ema.rst delete mode 100644 docs/autosummary/starling.models.quantize.VectorQuantizer2.rst delete mode 100644 docs/autosummary/starling.models.quantize.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.ResNet_Decoder_Original.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.ResNet_Encoder_Original.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet101_Decoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet101_Encoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet152_Decoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet152_Encoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet18_Decoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet18_Encoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet34_Decoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet34_Encoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet50_Decoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.Resnet50_Encoder.rst delete mode 100644 docs/autosummary/starling.models.resnets_original.rst delete mode 100644 starling/models/ema.py delete mode 100644 starling/models/quantize.py delete mode 100644 starling/models/resnets_original.py diff --git a/docs/autosummary/starling.models.ema.EMA.rst b/docs/autosummary/starling.models.ema.EMA.rst deleted file mode 100644 index bfae538..0000000 --- a/docs/autosummary/starling.models.ema.EMA.rst +++ /dev/null @@ -1,28 +0,0 @@ -starling.models.ema.EMA -======================= - -.. currentmodule:: starling.models.ema - -.. autoclass:: EMA - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~EMA.__init__ - ~EMA.reset_parameters - ~EMA.step_ema - ~EMA.update_average - ~EMA.update_model_average - - - - - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.ema.rst b/docs/autosummary/starling.models.ema.rst deleted file mode 100644 index 09719b3..0000000 --- a/docs/autosummary/starling.models.ema.rst +++ /dev/null @@ -1,32 +0,0 @@ -starling.models.ema -=================== - -.. automodule:: starling.models.ema - :no-members: - - - - - - - - - - - - .. rubric:: Classes - - .. autosummary:: - :toctree: - :nosignatures: - - EMA - - - - - - - - - diff --git a/docs/autosummary/starling.models.quantize.VectorQuantizer2.rst b/docs/autosummary/starling.models.quantize.VectorQuantizer2.rst deleted file mode 100644 index b9f1725..0000000 --- a/docs/autosummary/starling.models.quantize.VectorQuantizer2.rst +++ /dev/null @@ -1,86 +0,0 @@ -starling.models.quantize.VectorQuantizer2 -========================================= - -.. currentmodule:: starling.models.quantize - -.. autoclass:: VectorQuantizer2 - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~VectorQuantizer2.__init__ - ~VectorQuantizer2.add_module - ~VectorQuantizer2.apply - ~VectorQuantizer2.bfloat16 - ~VectorQuantizer2.buffers - ~VectorQuantizer2.children - ~VectorQuantizer2.compile - ~VectorQuantizer2.cpu - ~VectorQuantizer2.cuda - ~VectorQuantizer2.double - ~VectorQuantizer2.eval - ~VectorQuantizer2.extra_repr - ~VectorQuantizer2.float - ~VectorQuantizer2.forward - ~VectorQuantizer2.get_buffer - ~VectorQuantizer2.get_codebook_entry - ~VectorQuantizer2.get_extra_state - ~VectorQuantizer2.get_parameter - ~VectorQuantizer2.get_submodule - ~VectorQuantizer2.half - ~VectorQuantizer2.ipu - ~VectorQuantizer2.load_state_dict - ~VectorQuantizer2.modules - ~VectorQuantizer2.mtia - ~VectorQuantizer2.named_buffers - ~VectorQuantizer2.named_children - ~VectorQuantizer2.named_modules - ~VectorQuantizer2.named_parameters - ~VectorQuantizer2.parameters - ~VectorQuantizer2.register_backward_hook - ~VectorQuantizer2.register_buffer - ~VectorQuantizer2.register_forward_hook - ~VectorQuantizer2.register_forward_pre_hook - ~VectorQuantizer2.register_full_backward_hook - ~VectorQuantizer2.register_full_backward_pre_hook - ~VectorQuantizer2.register_load_state_dict_post_hook - ~VectorQuantizer2.register_load_state_dict_pre_hook - ~VectorQuantizer2.register_module - ~VectorQuantizer2.register_parameter - ~VectorQuantizer2.register_state_dict_post_hook - ~VectorQuantizer2.register_state_dict_pre_hook - ~VectorQuantizer2.remap_to_used - ~VectorQuantizer2.requires_grad_ - ~VectorQuantizer2.set_extra_state - ~VectorQuantizer2.set_submodule - ~VectorQuantizer2.share_memory - ~VectorQuantizer2.state_dict - ~VectorQuantizer2.to - ~VectorQuantizer2.to_empty - ~VectorQuantizer2.train - ~VectorQuantizer2.type - ~VectorQuantizer2.unmap_to_all - ~VectorQuantizer2.xpu - ~VectorQuantizer2.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~VectorQuantizer2.T_destination - ~VectorQuantizer2.call_super_init - ~VectorQuantizer2.dump_patches - ~VectorQuantizer2.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.quantize.rst b/docs/autosummary/starling.models.quantize.rst deleted file mode 100644 index aca9f80..0000000 --- a/docs/autosummary/starling.models.quantize.rst +++ /dev/null @@ -1,32 +0,0 @@ -starling.models.quantize -======================== - -.. automodule:: starling.models.quantize - :no-members: - - - - - - - - - - - - .. rubric:: Classes - - .. autosummary:: - :toctree: - :nosignatures: - - VectorQuantizer2 - - - - - - - - - diff --git a/docs/autosummary/starling.models.resnets_original.ResNet_Decoder_Original.rst b/docs/autosummary/starling.models.resnets_original.ResNet_Decoder_Original.rst deleted file mode 100644 index 7bb45b2..0000000 --- a/docs/autosummary/starling.models.resnets_original.ResNet_Decoder_Original.rst +++ /dev/null @@ -1,83 +0,0 @@ -starling.models.resnets\_original.ResNet\_Decoder\_Original -=========================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autoclass:: ResNet_Decoder_Original - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~ResNet_Decoder_Original.__init__ - ~ResNet_Decoder_Original.add_module - ~ResNet_Decoder_Original.apply - ~ResNet_Decoder_Original.bfloat16 - ~ResNet_Decoder_Original.buffers - ~ResNet_Decoder_Original.children - ~ResNet_Decoder_Original.compile - ~ResNet_Decoder_Original.cpu - ~ResNet_Decoder_Original.cuda - ~ResNet_Decoder_Original.double - ~ResNet_Decoder_Original.eval - ~ResNet_Decoder_Original.extra_repr - ~ResNet_Decoder_Original.float - ~ResNet_Decoder_Original.forward - ~ResNet_Decoder_Original.get_buffer - ~ResNet_Decoder_Original.get_extra_state - ~ResNet_Decoder_Original.get_parameter - ~ResNet_Decoder_Original.get_submodule - ~ResNet_Decoder_Original.half - ~ResNet_Decoder_Original.ipu - ~ResNet_Decoder_Original.load_state_dict - ~ResNet_Decoder_Original.modules - ~ResNet_Decoder_Original.mtia - ~ResNet_Decoder_Original.named_buffers - ~ResNet_Decoder_Original.named_children - ~ResNet_Decoder_Original.named_modules - ~ResNet_Decoder_Original.named_parameters - ~ResNet_Decoder_Original.parameters - ~ResNet_Decoder_Original.register_backward_hook - ~ResNet_Decoder_Original.register_buffer - ~ResNet_Decoder_Original.register_forward_hook - ~ResNet_Decoder_Original.register_forward_pre_hook - ~ResNet_Decoder_Original.register_full_backward_hook - ~ResNet_Decoder_Original.register_full_backward_pre_hook - ~ResNet_Decoder_Original.register_load_state_dict_post_hook - ~ResNet_Decoder_Original.register_load_state_dict_pre_hook - ~ResNet_Decoder_Original.register_module - ~ResNet_Decoder_Original.register_parameter - ~ResNet_Decoder_Original.register_state_dict_post_hook - ~ResNet_Decoder_Original.register_state_dict_pre_hook - ~ResNet_Decoder_Original.requires_grad_ - ~ResNet_Decoder_Original.set_extra_state - ~ResNet_Decoder_Original.set_submodule - ~ResNet_Decoder_Original.share_memory - ~ResNet_Decoder_Original.state_dict - ~ResNet_Decoder_Original.to - ~ResNet_Decoder_Original.to_empty - ~ResNet_Decoder_Original.train - ~ResNet_Decoder_Original.type - ~ResNet_Decoder_Original.xpu - ~ResNet_Decoder_Original.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~ResNet_Decoder_Original.T_destination - ~ResNet_Decoder_Original.call_super_init - ~ResNet_Decoder_Original.dump_patches - ~ResNet_Decoder_Original.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.ResNet_Encoder_Original.rst b/docs/autosummary/starling.models.resnets_original.ResNet_Encoder_Original.rst deleted file mode 100644 index 20924ab..0000000 --- a/docs/autosummary/starling.models.resnets_original.ResNet_Encoder_Original.rst +++ /dev/null @@ -1,83 +0,0 @@ -starling.models.resnets\_original.ResNet\_Encoder\_Original -=========================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autoclass:: ResNet_Encoder_Original - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~ResNet_Encoder_Original.__init__ - ~ResNet_Encoder_Original.add_module - ~ResNet_Encoder_Original.apply - ~ResNet_Encoder_Original.bfloat16 - ~ResNet_Encoder_Original.buffers - ~ResNet_Encoder_Original.children - ~ResNet_Encoder_Original.compile - ~ResNet_Encoder_Original.cpu - ~ResNet_Encoder_Original.cuda - ~ResNet_Encoder_Original.double - ~ResNet_Encoder_Original.eval - ~ResNet_Encoder_Original.extra_repr - ~ResNet_Encoder_Original.float - ~ResNet_Encoder_Original.forward - ~ResNet_Encoder_Original.get_buffer - ~ResNet_Encoder_Original.get_extra_state - ~ResNet_Encoder_Original.get_parameter - ~ResNet_Encoder_Original.get_submodule - ~ResNet_Encoder_Original.half - ~ResNet_Encoder_Original.ipu - ~ResNet_Encoder_Original.load_state_dict - ~ResNet_Encoder_Original.modules - ~ResNet_Encoder_Original.mtia - ~ResNet_Encoder_Original.named_buffers - ~ResNet_Encoder_Original.named_children - ~ResNet_Encoder_Original.named_modules - ~ResNet_Encoder_Original.named_parameters - ~ResNet_Encoder_Original.parameters - ~ResNet_Encoder_Original.register_backward_hook - ~ResNet_Encoder_Original.register_buffer - ~ResNet_Encoder_Original.register_forward_hook - ~ResNet_Encoder_Original.register_forward_pre_hook - ~ResNet_Encoder_Original.register_full_backward_hook - ~ResNet_Encoder_Original.register_full_backward_pre_hook - ~ResNet_Encoder_Original.register_load_state_dict_post_hook - ~ResNet_Encoder_Original.register_load_state_dict_pre_hook - ~ResNet_Encoder_Original.register_module - ~ResNet_Encoder_Original.register_parameter - ~ResNet_Encoder_Original.register_state_dict_post_hook - ~ResNet_Encoder_Original.register_state_dict_pre_hook - ~ResNet_Encoder_Original.requires_grad_ - ~ResNet_Encoder_Original.set_extra_state - ~ResNet_Encoder_Original.set_submodule - ~ResNet_Encoder_Original.share_memory - ~ResNet_Encoder_Original.state_dict - ~ResNet_Encoder_Original.to - ~ResNet_Encoder_Original.to_empty - ~ResNet_Encoder_Original.train - ~ResNet_Encoder_Original.type - ~ResNet_Encoder_Original.xpu - ~ResNet_Encoder_Original.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~ResNet_Encoder_Original.T_destination - ~ResNet_Encoder_Original.call_super_init - ~ResNet_Encoder_Original.dump_patches - ~ResNet_Encoder_Original.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet101_Decoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet101_Decoder.rst deleted file mode 100644 index cf4edb4..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet101_Decoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet101\_Decoder -==================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet101_Decoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet101_Encoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet101_Encoder.rst deleted file mode 100644 index f2ea49c..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet101_Encoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet101\_Encoder -==================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet101_Encoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet152_Decoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet152_Decoder.rst deleted file mode 100644 index 04b9b62..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet152_Decoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet152\_Decoder -==================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet152_Decoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet152_Encoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet152_Encoder.rst deleted file mode 100644 index a97a747..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet152_Encoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet152\_Encoder -==================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet152_Encoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet18_Decoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet18_Decoder.rst deleted file mode 100644 index 3dcbd46..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet18_Decoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet18\_Decoder -=================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet18_Decoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet18_Encoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet18_Encoder.rst deleted file mode 100644 index b4099a2..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet18_Encoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet18\_Encoder -=================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet18_Encoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet34_Decoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet34_Decoder.rst deleted file mode 100644 index 148a526..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet34_Decoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet34\_Decoder -=================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet34_Decoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet34_Encoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet34_Encoder.rst deleted file mode 100644 index e70c9d8..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet34_Encoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet34\_Encoder -=================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet34_Encoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet50_Decoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet50_Decoder.rst deleted file mode 100644 index d46e3c1..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet50_Decoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet50\_Decoder -=================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet50_Decoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.Resnet50_Encoder.rst b/docs/autosummary/starling.models.resnets_original.Resnet50_Encoder.rst deleted file mode 100644 index e3b3c0f..0000000 --- a/docs/autosummary/starling.models.resnets_original.Resnet50_Encoder.rst +++ /dev/null @@ -1,6 +0,0 @@ -starling.models.resnets\_original.Resnet50\_Encoder -=================================================== - -.. currentmodule:: starling.models.resnets_original - -.. autofunction:: Resnet50_Encoder \ No newline at end of file diff --git a/docs/autosummary/starling.models.resnets_original.rst b/docs/autosummary/starling.models.resnets_original.rst deleted file mode 100644 index 2b02351..0000000 --- a/docs/autosummary/starling.models.resnets_original.rst +++ /dev/null @@ -1,50 +0,0 @@ -starling.models.resnets\_original -================================= - -.. automodule:: starling.models.resnets_original - :no-members: - - - - - - - - .. rubric:: Functions - - .. autosummary:: - :toctree: - :nosignatures: - - Resnet101_Decoder - Resnet101_Encoder - Resnet152_Decoder - Resnet152_Encoder - Resnet18_Decoder - Resnet18_Encoder - Resnet34_Decoder - Resnet34_Encoder - Resnet50_Decoder - Resnet50_Encoder - - - - - - .. rubric:: Classes - - .. autosummary:: - :toctree: - :nosignatures: - - ResNet_Decoder_Original - ResNet_Encoder_Original - - - - - - - - - diff --git a/starling/models/ema.py b/starling/models/ema.py deleted file mode 100644 index d022215..0000000 --- a/starling/models/ema.py +++ /dev/null @@ -1,32 +0,0 @@ -import torch -import torch.nn as nn - - -class EMA: - def __init__(self, beta: float = 0.999): - self.beta = beta - self.step = 0 - - def update_model_average(self, ma_model, current_model): - for current_params, ma_params in zip( - current_model.parameters(), ma_model.parameters() - ): - ema_model_weights, current_weighs = ma_params.data, current_params.data - ma_params.data = self.update_average(ema_model_weights, current_weighs) - - def update_average(self, ema_model_params, current_model_params): - if ema_model_params is None: - return current_model_params - - return ema_model_params * self.beta + (1 - self.beta) * current_model_params - - def step_ema(self, ema_model, model, step_start_ema=2000): - if self.step < step_start_ema: - self.reset_parameters(ema_model, model) - self.step += 1 - return - self.update_model_average(ema_model, model) - self.step += 1 - - def reset_parameters(self, ema_model, model): - ema_model.load_state_dict(model.state_dict()) diff --git a/starling/models/quantize.py b/starling/models/quantize.py deleted file mode 100644 index 90ef48c..0000000 --- a/starling/models/quantize.py +++ /dev/null @@ -1,148 +0,0 @@ -import numpy as np -import torch -from einops import rearrange -from torch import nn - -# Adapted from https://github.com/CompVis/taming-transformers/blob/master/taming/modules/vqvae/quantize.py - - -class VectorQuantizer2(nn.Module): - """ - Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly - avoids costly matrix multiplications and allows for post-hoc remapping of indices. - """ - - # NOTE: due to a bug the beta term was applied to the wrong term. for - # backwards compatibility we use the buggy version by default, but you can - # specify legacy=False to fix it. - def __init__( - self, - n_e, - e_dim, - beta, - remap=None, - unknown_index="random", - sane_index_shape=False, - legacy=True, - ): - super().__init__() - self.n_e = n_e - self.e_dim = e_dim - self.beta = beta - self.legacy = legacy - - self.embedding = nn.Embedding(self.n_e, self.e_dim) - self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) - - self.remap = remap - if self.remap is not None: - self.register_buffer("used", torch.tensor(np.load(self.remap))) - self.re_embed = self.used.shape[0] - self.unknown_index = unknown_index # "random" or "extra" or integer - if self.unknown_index == "extra": - self.unknown_index = self.re_embed - self.re_embed = self.re_embed + 1 - print( - f"Remapping {self.n_e} indices to {self.re_embed} indices. " - f"Using {self.unknown_index} for unknown indices." - ) - else: - self.re_embed = n_e - - self.sane_index_shape = sane_index_shape - - def remap_to_used(self, inds): - ishape = inds.shape - assert len(ishape) > 1 - inds = inds.reshape(ishape[0], -1) - used = self.used.to(inds) - match = (inds[:, :, None] == used[None, None, ...]).long() - new = match.argmax(-1) - unknown = match.sum(2) < 1 - if self.unknown_index == "random": - new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to( - device=new.device - ) - else: - new[unknown] = self.unknown_index - return new.reshape(ishape) - - def unmap_to_all(self, inds): - ishape = inds.shape - assert len(ishape) > 1 - inds = inds.reshape(ishape[0], -1) - used = self.used.to(inds) - if self.re_embed > self.used.shape[0]: # extra token - inds[inds >= self.used.shape[0]] = 0 # simply set to zero - back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) - return back.reshape(ishape) - - def forward(self, z, temp=None, rescale_logits=False, return_logits=False): - assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" - assert rescale_logits == False, "Only for interface compatible with Gumbel" - assert return_logits == False, "Only for interface compatible with Gumbel" - # reshape z -> (batch, height, width, channel) and flatten - z = rearrange(z, "b c h w -> b h w c").contiguous() - z_flattened = z.view(-1, self.e_dim) - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - - d = ( - torch.sum(z_flattened**2, dim=1, keepdim=True) - + torch.sum(self.embedding.weight**2, dim=1) - - 2 - * torch.einsum( - "bd,dn->bn", z_flattened, rearrange(self.embedding.weight, "n d -> d n") - ) - ) - - min_encoding_indices = torch.argmin(d, dim=1) - z_q = self.embedding(min_encoding_indices).view(z.shape) - perplexity = None - min_encodings = None - - # compute loss for embedding - if not self.legacy: - loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean( - (z_q - z.detach()) ** 2 - ) - else: - loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean( - (z_q - z.detach()) ** 2 - ) - - # preserve gradients - z_q = z + (z_q - z).detach() - - # reshape back to match original input shape - z_q = rearrange(z_q, "b h w c -> b c h w").contiguous() - - if self.remap is not None: - min_encoding_indices = min_encoding_indices.reshape( - z.shape[0], -1 - ) # add batch axis - min_encoding_indices = self.remap_to_used(min_encoding_indices) - min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten - - if self.sane_index_shape: - min_encoding_indices = min_encoding_indices.reshape( - z_q.shape[0], z_q.shape[2], z_q.shape[3] - ) - - return z_q, loss, (perplexity, min_encodings, min_encoding_indices) - - def get_codebook_entry(self, indices, shape): - # shape specifying (batch, height, width, channel) - if self.remap is not None: - indices = indices.reshape(shape[0], -1) # add batch axis - indices = self.unmap_to_all(indices) - indices = indices.reshape(-1) # flatten again - - # get quantized latent vectors - z_q = self.embedding(indices) - - if shape is not None: - z_q = z_q.view(shape) - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return z_q diff --git a/starling/models/resnets_original.py b/starling/models/resnets_original.py deleted file mode 100644 index 04b2fa5..0000000 --- a/starling/models/resnets_original.py +++ /dev/null @@ -1,276 +0,0 @@ -import torch.nn.functional as F -from torch import nn - -from starling.models.blocks import ( - ResBlockDecBasic, - ResBlockDecBottleneck, - ResBlockEncBasic, - ResBlockEncBottleneck, -) - - -class ResNet_Encoder_Original(nn.Module): - def __init__( - self, - in_channels, - num_blocks, - kernel_size=None, - dimension=None, - block_type=ResBlockEncBasic, - base=64, - ) -> None: - super().__init__() - - self.block_type = block_type - - # First convolution of the ResNet Encoder reduction in the spatial dimensions / 2 - # with kernel=7 and stride=2 AvgPool2d reduces spatial dimensions by / 2 - self.in_channels = 64 - self.first_conv = nn.Sequential( - nn.Conv2d( - in_channels=in_channels, - out_channels=self.in_channels, - kernel_size=7, - stride=2, - padding=3, - ), - nn.BatchNorm2d(self.in_channels), - nn.MaxPool2d(kernel_size=3, stride=2, padding=1), - ) - - self.layer1 = self._make_layer(self.block_type, 64, num_blocks[0], stride=1) - self.layer2 = self._make_layer(self.block_type, 128, num_blocks[1], stride=2) - self.layer3 = self._make_layer(self.block_type, 256, num_blocks[2], stride=2) - self.layer4 = self._make_layer(self.block_type, 512, num_blocks[3], stride=2) - - self.average_pool = nn.AdaptiveAvgPool2d((1, 1)) - - def _make_layer(self, block, out_channels, blocks, stride=1): - layers = [] - layers.append(block(self.in_channels, out_channels, stride)) - self.in_channels = out_channels * block.expansion - for _ in range(1, blocks): - layers.append(block(self.in_channels, out_channels, stride=1)) - return nn.Sequential(*layers) - - def forward(self, data): - data = self.first_conv(data) - # for layer in self.layers: - # data = layer(data) - # The final adaptive average can also be done through convolution - data = self.layer1(data) - data = self.layer2(data) - data = self.layer3(data) - data = self.layer4(data) - data = self.average_pool(data) - return data - - -class ResNet_Decoder_Original(nn.Module): - def __init__( - self, - out_channels, - num_blocks, - kernel_size, - dimension, - block_type=ResBlockDecBasic, - base=64, - ) -> None: - super().__init__() - - # Calculate the input channels from the encoder, assuming - # symmetric encoder and decoder setup - self.block_type = block_type - if self.block_type == ResBlockDecBasic: - self.in_channels = 512 - else: - self.in_channels = 4096 - - self.interpolate = int(dimension / (2 ** (len(num_blocks) + 1))) - - # This part can be done in many ways, this is just one of them - # It adds some number of parameters - # self.resize_conv = ResizeConv2d( - # in_channels=self.in_channels, - # out_channels=self.in_channels, - # kernel_size=kernel_size, - # size=(self.interpolate, self.interpolate), - # mode="nearest", - # ) - - self.layers = nn.ModuleList() - - self.layer1 = self._make_layer(self.block_type, 512, num_blocks[0], stride=2) - self.layer2 = self._make_layer(self.block_type, 256, num_blocks[1], stride=2) - self.layer3 = self._make_layer(self.block_type, 128, num_blocks[2], stride=2) - self.layer4 = self._make_layer( - self.block_type, 64, num_blocks[3], stride=1, last_layer=True - ) - - # for layer in [self.layer1, self.layer2, self.layer3, self.layer4]: - # self.layers.append(layer) - - # # This part could be done through interpolation (analogous to MaxPool) - self.reshaping_conv = nn.Sequential( - nn.ConvTranspose2d( - in_channels=64, - out_channels=64, - kernel_size=kernel_size, - stride=2, - padding=1, - output_padding=1, - ), - nn.BatchNorm2d(64), - # nn.LayerNorm([64, int(dimension / 2), int(dimension / 2)]), - nn.ReLU(inplace=True), - ) - - # Final output layer that looks similar to the first layer of - # the ResNet Encoder - self.output_layer = nn.Sequential( - nn.ConvTranspose2d( - in_channels=64, - out_channels=out_channels, - kernel_size=7, - stride=2, - padding=3, - output_padding=1, - ), - nn.ReLU(inplace=True), - ) - - def _make_layer(self, block, out_channels, blocks, stride=1, last_layer=False): - layers = [] - self.in_channels = out_channels * block.contraction - for _ in range(1, blocks): - layers.append(block(self.in_channels, out_channels, stride=1)) - if stride > 1 and block == ResBlockDecBasic: - out_channels = int(out_channels / 2) - layers.append( - block(self.in_channels, out_channels, stride, last_layer=last_layer) - ) - return nn.Sequential(*layers) - - def forward(self, data): - # data = self.resize_conv(data) - data = F.interpolate(data, size=(self.interpolate, self.interpolate)) - data = self.layer1(data) - data = self.layer2(data) - data = self.layer3(data) - data = self.layer4(data) - data = self.reshaping_conv(data) - data = self.output_layer(data) - return data - - -# Current implementations of ResNets - - -def Resnet18_Encoder(in_channels, kernel_size, dimension, base): - return ResNet_Encoder_Original( - block_type=ResBlockEncBasic, - in_channels=in_channels, - num_blocks=[2, 2, 2, 2], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet18_Decoder(out_channels, kernel_size, dimension, base): - return ResNet_Decoder_Original( - block_type=ResBlockDecBasic, - out_channels=out_channels, - num_blocks=[2, 2, 2, 2], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet34_Encoder(in_channels, kernel_size, dimension, base): - return ResNet_Encoder_Original( - block_type=ResBlockEncBasic, - in_channels=in_channels, - num_blocks=[3, 4, 6, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet34_Decoder(out_channels, kernel_size, dimension, base): - return ResNet_Decoder_Original( - block_type=ResBlockDecBasic, - out_channels=out_channels, - num_blocks=[3, 6, 4, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet50_Encoder(in_channels, kernel_size, dimension, base): - return ResNet_Encoder_Original( - block_type=ResBlockEncBottleneck, - in_channels=in_channels, - num_blocks=[3, 4, 6, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet50_Decoder(out_channels, kernel_size, dimension, base): - return ResNet_Decoder_Original( - block_type=ResBlockDecBottleneck, - out_channels=out_channels, - num_blocks=[3, 6, 4, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet101_Encoder(in_channels, kernel_size, dimension, base): - return ResNet_Encoder_Original( - block_type=ResBlockEncBottleneck, - in_channels=in_channels, - num_blocks=[3, 4, 23, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet101_Decoder(out_channels, kernel_size, dimension, base): - return ResNet_Decoder_Original( - block_type=ResBlockDecBottleneck, - out_channels=out_channels, - num_blocks=[3, 23, 4, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet152_Encoder(in_channels, kernel_size, dimension, base): - return ResNet_Encoder_Original( - block_type=ResBlockEncBottleneck, - in_channels=in_channels, - num_blocks=[3, 8, 36, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) - - -def Resnet152_Decoder(out_channels, kernel_size, dimension, base): - return ResNet_Decoder_Original( - block_type=ResBlockDecBottleneck, - out_channels=out_channels, - num_blocks=[3, 36, 8, 3], - kernel_size=kernel_size, - dimension=dimension, - base=base, - ) From 38fe7cab5740af445dc4d8d442910e17f6301a55 Mon Sep 17 00:00:00 2001 From: Jeff lotthammer Date: Thu, 16 Jul 2026 11:28:51 -0500 Subject: [PATCH 2/4] Retire pre-transformer UNet backbone 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. --- ...ling.models.unet.ConditionalSequential.rst | 87 ---- ....models.unet.CrossAttentionResnetLayer.rst | 83 ---- .../starling.models.unet.Downsample.rst | 83 ---- .../starling.models.unet.ResnetLayer.rst | 83 ---- .../starling.models.unet.SinusoidalPosEmb.rst | 83 ---- .../starling.models.unet.UNetConditional.rst | 83 ---- docs/autosummary/starling.models.unet.rst | 37 -- starling/configs/unet/unet.yaml | 7 - starling/models/unet.py | 464 ------------------ 9 files changed, 1010 deletions(-) delete mode 100644 docs/autosummary/starling.models.unet.ConditionalSequential.rst delete mode 100644 docs/autosummary/starling.models.unet.CrossAttentionResnetLayer.rst delete mode 100644 docs/autosummary/starling.models.unet.Downsample.rst delete mode 100644 docs/autosummary/starling.models.unet.ResnetLayer.rst delete mode 100644 docs/autosummary/starling.models.unet.SinusoidalPosEmb.rst delete mode 100644 docs/autosummary/starling.models.unet.UNetConditional.rst delete mode 100644 docs/autosummary/starling.models.unet.rst delete mode 100644 starling/configs/unet/unet.yaml delete mode 100644 starling/models/unet.py diff --git a/docs/autosummary/starling.models.unet.ConditionalSequential.rst b/docs/autosummary/starling.models.unet.ConditionalSequential.rst deleted file mode 100644 index 6b5a856..0000000 --- a/docs/autosummary/starling.models.unet.ConditionalSequential.rst +++ /dev/null @@ -1,87 +0,0 @@ -starling.models.unet.ConditionalSequential -========================================== - -.. currentmodule:: starling.models.unet - -.. autoclass:: ConditionalSequential - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~ConditionalSequential.__init__ - ~ConditionalSequential.add_module - ~ConditionalSequential.append - ~ConditionalSequential.apply - ~ConditionalSequential.bfloat16 - ~ConditionalSequential.buffers - ~ConditionalSequential.children - ~ConditionalSequential.compile - ~ConditionalSequential.cpu - ~ConditionalSequential.cuda - ~ConditionalSequential.double - ~ConditionalSequential.eval - ~ConditionalSequential.extend - ~ConditionalSequential.extra_repr - ~ConditionalSequential.float - ~ConditionalSequential.forward - ~ConditionalSequential.get_buffer - ~ConditionalSequential.get_extra_state - ~ConditionalSequential.get_parameter - ~ConditionalSequential.get_submodule - ~ConditionalSequential.half - ~ConditionalSequential.insert - ~ConditionalSequential.ipu - ~ConditionalSequential.load_state_dict - ~ConditionalSequential.modules - ~ConditionalSequential.mtia - ~ConditionalSequential.named_buffers - ~ConditionalSequential.named_children - ~ConditionalSequential.named_modules - ~ConditionalSequential.named_parameters - ~ConditionalSequential.parameters - ~ConditionalSequential.pop - ~ConditionalSequential.register_backward_hook - ~ConditionalSequential.register_buffer - ~ConditionalSequential.register_forward_hook - ~ConditionalSequential.register_forward_pre_hook - ~ConditionalSequential.register_full_backward_hook - ~ConditionalSequential.register_full_backward_pre_hook - ~ConditionalSequential.register_load_state_dict_post_hook - ~ConditionalSequential.register_load_state_dict_pre_hook - ~ConditionalSequential.register_module - ~ConditionalSequential.register_parameter - ~ConditionalSequential.register_state_dict_post_hook - ~ConditionalSequential.register_state_dict_pre_hook - ~ConditionalSequential.requires_grad_ - ~ConditionalSequential.set_extra_state - ~ConditionalSequential.set_submodule - ~ConditionalSequential.share_memory - ~ConditionalSequential.state_dict - ~ConditionalSequential.to - ~ConditionalSequential.to_empty - ~ConditionalSequential.train - ~ConditionalSequential.type - ~ConditionalSequential.xpu - ~ConditionalSequential.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~ConditionalSequential.T_destination - ~ConditionalSequential.call_super_init - ~ConditionalSequential.dump_patches - ~ConditionalSequential.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.unet.CrossAttentionResnetLayer.rst b/docs/autosummary/starling.models.unet.CrossAttentionResnetLayer.rst deleted file mode 100644 index b6c6023..0000000 --- a/docs/autosummary/starling.models.unet.CrossAttentionResnetLayer.rst +++ /dev/null @@ -1,83 +0,0 @@ -starling.models.unet.CrossAttentionResnetLayer -============================================== - -.. currentmodule:: starling.models.unet - -.. autoclass:: CrossAttentionResnetLayer - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~CrossAttentionResnetLayer.__init__ - ~CrossAttentionResnetLayer.add_module - ~CrossAttentionResnetLayer.apply - ~CrossAttentionResnetLayer.bfloat16 - ~CrossAttentionResnetLayer.buffers - ~CrossAttentionResnetLayer.children - ~CrossAttentionResnetLayer.compile - ~CrossAttentionResnetLayer.cpu - ~CrossAttentionResnetLayer.cuda - ~CrossAttentionResnetLayer.double - ~CrossAttentionResnetLayer.eval - ~CrossAttentionResnetLayer.extra_repr - ~CrossAttentionResnetLayer.float - ~CrossAttentionResnetLayer.forward - ~CrossAttentionResnetLayer.get_buffer - ~CrossAttentionResnetLayer.get_extra_state - ~CrossAttentionResnetLayer.get_parameter - ~CrossAttentionResnetLayer.get_submodule - ~CrossAttentionResnetLayer.half - ~CrossAttentionResnetLayer.ipu - ~CrossAttentionResnetLayer.load_state_dict - ~CrossAttentionResnetLayer.modules - ~CrossAttentionResnetLayer.mtia - ~CrossAttentionResnetLayer.named_buffers - ~CrossAttentionResnetLayer.named_children - ~CrossAttentionResnetLayer.named_modules - ~CrossAttentionResnetLayer.named_parameters - ~CrossAttentionResnetLayer.parameters - ~CrossAttentionResnetLayer.register_backward_hook - ~CrossAttentionResnetLayer.register_buffer - ~CrossAttentionResnetLayer.register_forward_hook - ~CrossAttentionResnetLayer.register_forward_pre_hook - ~CrossAttentionResnetLayer.register_full_backward_hook - ~CrossAttentionResnetLayer.register_full_backward_pre_hook - ~CrossAttentionResnetLayer.register_load_state_dict_post_hook - ~CrossAttentionResnetLayer.register_load_state_dict_pre_hook - ~CrossAttentionResnetLayer.register_module - ~CrossAttentionResnetLayer.register_parameter - ~CrossAttentionResnetLayer.register_state_dict_post_hook - ~CrossAttentionResnetLayer.register_state_dict_pre_hook - ~CrossAttentionResnetLayer.requires_grad_ - ~CrossAttentionResnetLayer.set_extra_state - ~CrossAttentionResnetLayer.set_submodule - ~CrossAttentionResnetLayer.share_memory - ~CrossAttentionResnetLayer.state_dict - ~CrossAttentionResnetLayer.to - ~CrossAttentionResnetLayer.to_empty - ~CrossAttentionResnetLayer.train - ~CrossAttentionResnetLayer.type - ~CrossAttentionResnetLayer.xpu - ~CrossAttentionResnetLayer.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~CrossAttentionResnetLayer.T_destination - ~CrossAttentionResnetLayer.call_super_init - ~CrossAttentionResnetLayer.dump_patches - ~CrossAttentionResnetLayer.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.unet.Downsample.rst b/docs/autosummary/starling.models.unet.Downsample.rst deleted file mode 100644 index 750280f..0000000 --- a/docs/autosummary/starling.models.unet.Downsample.rst +++ /dev/null @@ -1,83 +0,0 @@ -starling.models.unet.Downsample -=============================== - -.. currentmodule:: starling.models.unet - -.. autoclass:: Downsample - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~Downsample.__init__ - ~Downsample.add_module - ~Downsample.apply - ~Downsample.bfloat16 - ~Downsample.buffers - ~Downsample.children - ~Downsample.compile - ~Downsample.cpu - ~Downsample.cuda - ~Downsample.double - ~Downsample.eval - ~Downsample.extra_repr - ~Downsample.float - ~Downsample.forward - ~Downsample.get_buffer - ~Downsample.get_extra_state - ~Downsample.get_parameter - ~Downsample.get_submodule - ~Downsample.half - ~Downsample.ipu - ~Downsample.load_state_dict - ~Downsample.modules - ~Downsample.mtia - ~Downsample.named_buffers - ~Downsample.named_children - ~Downsample.named_modules - ~Downsample.named_parameters - ~Downsample.parameters - ~Downsample.register_backward_hook - ~Downsample.register_buffer - ~Downsample.register_forward_hook - ~Downsample.register_forward_pre_hook - ~Downsample.register_full_backward_hook - ~Downsample.register_full_backward_pre_hook - ~Downsample.register_load_state_dict_post_hook - ~Downsample.register_load_state_dict_pre_hook - ~Downsample.register_module - ~Downsample.register_parameter - ~Downsample.register_state_dict_post_hook - ~Downsample.register_state_dict_pre_hook - ~Downsample.requires_grad_ - ~Downsample.set_extra_state - ~Downsample.set_submodule - ~Downsample.share_memory - ~Downsample.state_dict - ~Downsample.to - ~Downsample.to_empty - ~Downsample.train - ~Downsample.type - ~Downsample.xpu - ~Downsample.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~Downsample.T_destination - ~Downsample.call_super_init - ~Downsample.dump_patches - ~Downsample.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.unet.ResnetLayer.rst b/docs/autosummary/starling.models.unet.ResnetLayer.rst deleted file mode 100644 index b139243..0000000 --- a/docs/autosummary/starling.models.unet.ResnetLayer.rst +++ /dev/null @@ -1,83 +0,0 @@ -starling.models.unet.ResnetLayer -================================ - -.. currentmodule:: starling.models.unet - -.. autoclass:: ResnetLayer - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~ResnetLayer.__init__ - ~ResnetLayer.add_module - ~ResnetLayer.apply - ~ResnetLayer.bfloat16 - ~ResnetLayer.buffers - ~ResnetLayer.children - ~ResnetLayer.compile - ~ResnetLayer.cpu - ~ResnetLayer.cuda - ~ResnetLayer.double - ~ResnetLayer.eval - ~ResnetLayer.extra_repr - ~ResnetLayer.float - ~ResnetLayer.forward - ~ResnetLayer.get_buffer - ~ResnetLayer.get_extra_state - ~ResnetLayer.get_parameter - ~ResnetLayer.get_submodule - ~ResnetLayer.half - ~ResnetLayer.ipu - ~ResnetLayer.load_state_dict - ~ResnetLayer.modules - ~ResnetLayer.mtia - ~ResnetLayer.named_buffers - ~ResnetLayer.named_children - ~ResnetLayer.named_modules - ~ResnetLayer.named_parameters - ~ResnetLayer.parameters - ~ResnetLayer.register_backward_hook - ~ResnetLayer.register_buffer - ~ResnetLayer.register_forward_hook - ~ResnetLayer.register_forward_pre_hook - ~ResnetLayer.register_full_backward_hook - ~ResnetLayer.register_full_backward_pre_hook - ~ResnetLayer.register_load_state_dict_post_hook - ~ResnetLayer.register_load_state_dict_pre_hook - ~ResnetLayer.register_module - ~ResnetLayer.register_parameter - ~ResnetLayer.register_state_dict_post_hook - ~ResnetLayer.register_state_dict_pre_hook - ~ResnetLayer.requires_grad_ - ~ResnetLayer.set_extra_state - ~ResnetLayer.set_submodule - ~ResnetLayer.share_memory - ~ResnetLayer.state_dict - ~ResnetLayer.to - ~ResnetLayer.to_empty - ~ResnetLayer.train - ~ResnetLayer.type - ~ResnetLayer.xpu - ~ResnetLayer.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~ResnetLayer.T_destination - ~ResnetLayer.call_super_init - ~ResnetLayer.dump_patches - ~ResnetLayer.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.unet.SinusoidalPosEmb.rst b/docs/autosummary/starling.models.unet.SinusoidalPosEmb.rst deleted file mode 100644 index 4be7707..0000000 --- a/docs/autosummary/starling.models.unet.SinusoidalPosEmb.rst +++ /dev/null @@ -1,83 +0,0 @@ -starling.models.unet.SinusoidalPosEmb -===================================== - -.. currentmodule:: starling.models.unet - -.. autoclass:: SinusoidalPosEmb - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~SinusoidalPosEmb.__init__ - ~SinusoidalPosEmb.add_module - ~SinusoidalPosEmb.apply - ~SinusoidalPosEmb.bfloat16 - ~SinusoidalPosEmb.buffers - ~SinusoidalPosEmb.children - ~SinusoidalPosEmb.compile - ~SinusoidalPosEmb.cpu - ~SinusoidalPosEmb.cuda - ~SinusoidalPosEmb.double - ~SinusoidalPosEmb.eval - ~SinusoidalPosEmb.extra_repr - ~SinusoidalPosEmb.float - ~SinusoidalPosEmb.forward - ~SinusoidalPosEmb.get_buffer - ~SinusoidalPosEmb.get_extra_state - ~SinusoidalPosEmb.get_parameter - ~SinusoidalPosEmb.get_submodule - ~SinusoidalPosEmb.half - ~SinusoidalPosEmb.ipu - ~SinusoidalPosEmb.load_state_dict - ~SinusoidalPosEmb.modules - ~SinusoidalPosEmb.mtia - ~SinusoidalPosEmb.named_buffers - ~SinusoidalPosEmb.named_children - ~SinusoidalPosEmb.named_modules - ~SinusoidalPosEmb.named_parameters - ~SinusoidalPosEmb.parameters - ~SinusoidalPosEmb.register_backward_hook - ~SinusoidalPosEmb.register_buffer - ~SinusoidalPosEmb.register_forward_hook - ~SinusoidalPosEmb.register_forward_pre_hook - ~SinusoidalPosEmb.register_full_backward_hook - ~SinusoidalPosEmb.register_full_backward_pre_hook - ~SinusoidalPosEmb.register_load_state_dict_post_hook - ~SinusoidalPosEmb.register_load_state_dict_pre_hook - ~SinusoidalPosEmb.register_module - ~SinusoidalPosEmb.register_parameter - ~SinusoidalPosEmb.register_state_dict_post_hook - ~SinusoidalPosEmb.register_state_dict_pre_hook - ~SinusoidalPosEmb.requires_grad_ - ~SinusoidalPosEmb.set_extra_state - ~SinusoidalPosEmb.set_submodule - ~SinusoidalPosEmb.share_memory - ~SinusoidalPosEmb.state_dict - ~SinusoidalPosEmb.to - ~SinusoidalPosEmb.to_empty - ~SinusoidalPosEmb.train - ~SinusoidalPosEmb.type - ~SinusoidalPosEmb.xpu - ~SinusoidalPosEmb.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~SinusoidalPosEmb.T_destination - ~SinusoidalPosEmb.call_super_init - ~SinusoidalPosEmb.dump_patches - ~SinusoidalPosEmb.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.unet.UNetConditional.rst b/docs/autosummary/starling.models.unet.UNetConditional.rst deleted file mode 100644 index ea5cfb1..0000000 --- a/docs/autosummary/starling.models.unet.UNetConditional.rst +++ /dev/null @@ -1,83 +0,0 @@ -starling.models.unet.UNetConditional -==================================== - -.. currentmodule:: starling.models.unet - -.. autoclass:: UNetConditional - :members: - :show-inheritance: - :special-members: __init__, __call__ - - - - .. rubric:: Methods - - .. autosummary:: - :nosignatures: - - ~UNetConditional.__init__ - ~UNetConditional.add_module - ~UNetConditional.apply - ~UNetConditional.bfloat16 - ~UNetConditional.buffers - ~UNetConditional.children - ~UNetConditional.compile - ~UNetConditional.cpu - ~UNetConditional.cuda - ~UNetConditional.double - ~UNetConditional.eval - ~UNetConditional.extra_repr - ~UNetConditional.float - ~UNetConditional.forward - ~UNetConditional.get_buffer - ~UNetConditional.get_extra_state - ~UNetConditional.get_parameter - ~UNetConditional.get_submodule - ~UNetConditional.half - ~UNetConditional.ipu - ~UNetConditional.load_state_dict - ~UNetConditional.modules - ~UNetConditional.mtia - ~UNetConditional.named_buffers - ~UNetConditional.named_children - ~UNetConditional.named_modules - ~UNetConditional.named_parameters - ~UNetConditional.parameters - ~UNetConditional.register_backward_hook - ~UNetConditional.register_buffer - ~UNetConditional.register_forward_hook - ~UNetConditional.register_forward_pre_hook - ~UNetConditional.register_full_backward_hook - ~UNetConditional.register_full_backward_pre_hook - ~UNetConditional.register_load_state_dict_post_hook - ~UNetConditional.register_load_state_dict_pre_hook - ~UNetConditional.register_module - ~UNetConditional.register_parameter - ~UNetConditional.register_state_dict_post_hook - ~UNetConditional.register_state_dict_pre_hook - ~UNetConditional.requires_grad_ - ~UNetConditional.set_extra_state - ~UNetConditional.set_submodule - ~UNetConditional.share_memory - ~UNetConditional.state_dict - ~UNetConditional.to - ~UNetConditional.to_empty - ~UNetConditional.train - ~UNetConditional.type - ~UNetConditional.xpu - ~UNetConditional.zero_grad - - - - - - .. rubric:: Attributes - - .. autosummary:: - - ~UNetConditional.T_destination - ~UNetConditional.call_super_init - ~UNetConditional.dump_patches - ~UNetConditional.training - - \ No newline at end of file diff --git a/docs/autosummary/starling.models.unet.rst b/docs/autosummary/starling.models.unet.rst deleted file mode 100644 index 7356e1d..0000000 --- a/docs/autosummary/starling.models.unet.rst +++ /dev/null @@ -1,37 +0,0 @@ -starling.models.unet -==================== - -.. automodule:: starling.models.unet - :no-members: - - - - - - - - - - - - .. rubric:: Classes - - .. autosummary:: - :toctree: - :nosignatures: - - ConditionalSequential - CrossAttentionResnetLayer - Downsample - ResnetLayer - SinusoidalPosEmb - UNetConditional - - - - - - - - - diff --git a/starling/configs/unet/unet.yaml b/starling/configs/unet/unet.yaml deleted file mode 100644 index 8373304..0000000 --- a/starling/configs/unet/unet.yaml +++ /dev/null @@ -1,7 +0,0 @@ - in_channels: 1 - out_channels: 1 - base: 64 - norm: group - blocks: [2, 2, 2] - middle_blocks: 2 - labels_dim: 512 \ No newline at end of file diff --git a/starling/models/unet.py b/starling/models/unet.py deleted file mode 100644 index 11564b8..0000000 --- a/starling/models/unet.py +++ /dev/null @@ -1,464 +0,0 @@ -import math -from typing import List - -import torch -from torch import nn - -from starling.models.blocks import ResBlockEncBasic, ResizeConv2d -from starling.models.normalization import RMSNorm -from starling.models.transformer import SpatialTransformer - - -class SinusoidalPosEmb(nn.Module): - def __init__(self, dim: int, theta: int = 10000): - """ - Generates sinusoidal positional embeddings that are used in the denoising-diffusion - models to encode the timestep information. The positional embeddings are generated - using sine and cosine functions. It takes in time in the shape of (batch_size, 1) - and returns the positional embeddings in the shape of (batch_size, dim). The positional - encodings are later used in each of the ResNet blocks to encode the timestep information. - - Parameters - ---------- - dim : int - Dimension of the input data. - theta : int, optional - A scaling factor for the positional embeddings. The default value is 10000. - """ - super().__init__() - self.dim = dim - self.theta = theta - - def forward(self, time: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the positional (timestep) embeddings. - - Parameters - ---------- - time : torch.Tensor - Timestep information in the shape of (batch_size, 1). - - Returns - ------- - torch.Tensor - Positional (timestep) embeddings in the shape of (batch_size, dim). - """ - device = time.device - - # The number of unique frequencies in the positional embeddings, half - # will be used for sine and the other half for cosine functions - half_dim = self.dim // 2 - emb = math.log(self.theta) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device) * -emb) - emb = time[:, None] * emb[None, :] - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class ConditionalSequential(nn.Sequential): - def forward(self, x, condition): - for module in self._modules.values(): - x = module(x, condition) - return x - - -class Downsample(nn.Module): - def __init__(self, in_channels: int, out_channels: int, norm: str): - """ - A convolutional block that reduces the spatial dimensions of the input tensor by a factor of 2. - The block consists of a convolutional layer with a kernel size of 3, stride of 2, and padding of 1. - The convolutional layer is followed by a normalization layer and a ReLU activation function. - - Parameters - ---------- - in_channels : int - The number of features in the input tensor. - out_channels : int - The number of features in the output tensor. - norm : str - The normalization layer to be used in the block. Choose from batch, instance, rms, or group. - """ - super().__init__() - - normalization = { - "batch": nn.BatchNorm2d, - "instance": nn.InstanceNorm2d, - "rms": RMSNorm, - "group": nn.GroupNorm, - } - - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1), - normalization[norm](out_channels) - if norm != "group" - else normalization[norm](32, out_channels), - nn.ReLU(inplace=True), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.conv(x) - return x - - -class ResnetLayer(nn.Module): - def __init__( - self, - in_channels, - out_channels, - norm, - num_blocks, - timestep_dim, - class_dim=None, - ): - super().__init__() - - self.layer = nn.ModuleList() - - self.in_channels = in_channels - - for block in range(num_blocks): - self.layer.append( - ResBlockEncBasic( - self.in_channels, out_channels, 1, norm, timestep_dim, class_dim - ) - ) - - self.in_channels = out_channels - - def forward(self, x, time): - for layer in self.layer: - x = layer(x, time) - return x - - -class CrossAttentionResnetLayer(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - norm: str, - num_blocks: int, - attention_heads: int, - timestep_dim: int, - label_dim: int, - ): - """ - A combination of ResNet blocks followed by spatial transformer blocks. The ResNet block - processes the input tensor and the spatial transformer block captures the relationships - between the input tensor and the context data (protein sequences). - - Parameters - ---------- - in_channels : int - The number of features in the input tensor. - out_channels : int - The number of features in the output tensor. - norm : str - The normalization layer to be used in the block. Choose from batch, instance, rms, or group. - num_blocks : int - The number of ResNet + spatial transformer blocks in the layer. - attention_heads : int - The number of heads in the multi-head attention layer. - timestep_dim : int - The dimension of the timestep embeddings. - label_dim : int - The dimension of the context data (protein sequences). - """ - super().__init__() - - self.layer = nn.ModuleList() - self.transformer = nn.ModuleList() - - self.in_channels = in_channels - - for block in range(num_blocks): - self.layer.append( - ResBlockEncBasic(self.in_channels, out_channels, 1, norm, timestep_dim) - ) - self.transformer.append( - SpatialTransformer(out_channels, attention_heads, label_dim), - ) - - self.in_channels = out_channels - - def forward( - self, - x: torch.Tensor, - time: torch.Tensor, - sequence_label: torch.Tensor, - sequence_mask: torch.Tensor, - ) -> torch.Tensor: - """ - Forward pass of the ResNet + spatial transformer blocks. - - Parameters - ---------- - x : torch.Tensor - Input tensor to be processed by the ResNet + spatial transformer blocks. - time : torch.Tensor - Timestep embeddings to be used by the network. - sequence_label : torch.Tensor - Context data (protein sequences) to guide the prediction. - - Returns - ------- - torch.Tensor - Output of the ResNet + spatial transformer blocks. - """ - for layer, transformer in zip(self.layer, self.transformer): - x = layer(x, time) - x = transformer(x, context=sequence_label, mask=sequence_mask) - return x - - -class UNetConditional(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - base: int, - norm: str, - blocks: List = [2, 2, 2], - middle_blocks: int = 2, - labels_dim: int = 512, - sinusoidal_pos_emb_theta: int = 10000, - ): - """ - A U-Net architecture that uses ResNet blocks with spatial transformer blocks to process the input - tensor and the context data (protein sequences). The U-Net architecture consists of an encoder, - a middle section, and a decoder. The spatial transformer blocks are used to capture the relationships - between the input tensor and the context data (in our case protein sequences). - - Parameters - ---------- - in_channels : int - The number of features in the input tensor. - out_channels : int - The number of features in the output tensor. - base : int - The base number of features in the U-Net architecture. - norm : str - The normalization layer to be used in the block. Choose from batch, instance, rms, or group - blocks : List, optional - The number of ResNet + spatial transformer blocks in each section of the U-Net architecture, by default [2, 2, 2] - middle_blocks : int, optional - The number of ResNet + spatial transformer blocks in the middle section of the U-Net architecture, by default 2 - labels_dim : int, optional - The dimension of the context data (i.e., protein sequences), by default 512 - sinusoidal_pos_emb_theta : int, optional - A scaling factor for the positional (timestep) embeddings, by default 10000 - """ - super().__init__() - - normalization = { - "batch": nn.BatchNorm2d, - "instance": nn.InstanceNorm2d, - "rms": RMSNorm, - "group": nn.GroupNorm, - } - - self.norm = norm - self.in_channels = in_channels - self.out_channels = out_channels - self.time_dim = base * 4 - self.base = base - self.labels_dim = labels_dim - - # Time embeddings - self.time_emb = SinusoidalPosEmb(self.base, theta=sinusoidal_pos_emb_theta) - self.time_mlp = nn.Sequential( - self.time_emb, - nn.Linear(self.base, self.time_dim), - nn.SiLU(inplace=False), - nn.Linear(self.time_dim, self.time_dim), - ) - - all_in_channels = [base * (2**i) for i in range(len(blocks) + 1)] - - # Encoder part of UNet - - self.conv_in = CrossAttentionResnetLayer( - in_channels, - all_in_channels[0], - self.norm, - blocks[0], - 8, - self.time_dim, - self.labels_dim, - ) - - self.encoder_layer1 = CrossAttentionResnetLayer( - all_in_channels[0], - all_in_channels[0], - self.norm, - blocks[0], - 8, - self.time_dim, - self.labels_dim, - ) - - self.downsample1 = Downsample(all_in_channels[0], all_in_channels[1], norm) - - self.encoder_layer2 = CrossAttentionResnetLayer( - all_in_channels[1], - all_in_channels[1], - self.norm, - blocks[1], - 8, - self.time_dim, - self.labels_dim, - ) - - self.downsample2 = Downsample(all_in_channels[1], all_in_channels[2], norm) - - self.encoder_layer3 = CrossAttentionResnetLayer( - all_in_channels[2], - all_in_channels[2], - self.norm, - blocks[2], - 8, - self.time_dim, - self.labels_dim, - ) - - self.downsample3 = Downsample(all_in_channels[2], all_in_channels[3], norm) - - # Middle convolution of the UNet - - self.middle = CrossAttentionResnetLayer( - all_in_channels[3], - all_in_channels[3], - self.norm, - middle_blocks, - 8, - self.time_dim, - self.labels_dim, - ) - - # Decoder part of UNet - - self.upconv1 = ResizeConv2d( - all_in_channels[3], - all_in_channels[2], - kernel_size=3, - padding=1, - scale_factor=2, - norm=normalization[norm], - activation="relu", - ) - - self.decoder_layer1 = CrossAttentionResnetLayer( - all_in_channels[2] * 2, - all_in_channels[2], - self.norm, - blocks[2], - 8, - self.time_dim, - self.labels_dim, - ) - - self.upconv2 = ResizeConv2d( - all_in_channels[2], - all_in_channels[1], - kernel_size=3, - padding=1, - scale_factor=2, - norm=normalization[norm], - activation="relu", - ) - - self.decoder_layer2 = CrossAttentionResnetLayer( - all_in_channels[1] * 2, - all_in_channels[1], - self.norm, - blocks[1], - 8, - self.time_dim, - self.labels_dim, - ) - - self.upconv3 = ResizeConv2d( - all_in_channels[1], - all_in_channels[0], - kernel_size=3, - padding=1, - scale_factor=2, - norm=normalization[norm], - activation="relu", - ) - - self.decoder_layer3 = CrossAttentionResnetLayer( - all_in_channels[0] * 2, - all_in_channels[0], - self.norm, - blocks[1], - 8, - self.time_dim, - self.labels_dim, - ) - - self.conv_out = nn.Conv2d(all_in_channels[0], out_channels, kernel_size=1) - - def forward( - self, - x: torch.Tensor, - time: torch.Tensor, - labels: torch.Tensor, - sequence_mask: torch.Tensor, - ) -> torch.Tensor: - """ - Forward pass of the UNet architecture. - - Parameters - ---------- - x : torch.Tensor - Data to pass through the UNet architecture. - time : torch.Tensor - Timestep embeddings. - labels : torch.Tensor, optional - Context data (protein sequences) to guide the prediction, by default None - - Returns - ------- - torch.Tensor - Output of the UNet architecture. - """ - # Get the time embeddings - time = self.time_mlp(time) - - # Initial convolution - x = self.conv_in(x, time, labels, sequence_mask) - - # Encoder forward passes - x = self.encoder_layer1(x, time, labels, sequence_mask) - x_layer1 = x.clone() - x = self.downsample1(x) - - x = self.encoder_layer2(x, time, labels, sequence_mask) - x_layer2 = x.clone() - x = self.downsample2(x) - - x = self.encoder_layer3(x, time, labels, sequence_mask) - x_layer3 = x.clone() - x = self.downsample3(x) - - # Mid UNet - x = self.middle(x, time, labels, sequence_mask) - - # Decoder forward passes with skip connections from the encoder - x = self.upconv1(x) - x = torch.cat((x, x_layer3), dim=1) - x = self.decoder_layer1(x, time, labels, sequence_mask) - - x = self.upconv2(x) - x = torch.cat((x, x_layer2), dim=1) - x = self.decoder_layer2(x, time, labels, sequence_mask) - - x = self.upconv3(x) - x = torch.cat((x, x_layer1), dim=1) - x = self.decoder_layer3(x, time, labels, sequence_mask) - - # Final convolutions - x = self.conv_out(x) - - return x From 076b32dde1035a173ad4c94b833db57cfb96855b Mon Sep 17 00:00:00 2001 From: Jeff lotthammer Date: Thu, 16 Jul 2026 11:29:52 -0500 Subject: [PATCH 3/4] Fix broken entry points and dead import - 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. --- docs/api.rst | 4 ---- pyproject.toml | 2 -- starling/configs/configs.yaml | 1 - starling/models/vit.py | 3 +-- starling/training/diffusion_train.py | 12 ++++-------- 5 files changed, 5 insertions(+), 17 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index 19b5de3..fff1dde 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -60,17 +60,13 @@ Models models.vae models.diffusion - models.unet models.vit models.transformer models.attention models.blocks models.normalization - models.ema models.continuous_diffusion models.vae_components - models.quantize - models.resnets_original Samplers -------- diff --git a/pyproject.toml b/pyproject.toml index 1358f59..f315841 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,8 +49,6 @@ Documentation = "https://idptools-starling.readthedocs.io" [project.scripts] starling-vae-train = "starling.training.vae_train:train_vae" starling-ddpm-train = "starling.training.diffusion_train:train_model" -starling-sample = "starling.training.vae_generate:vae_generate" -ae-train = "starling.training.ae_train:train_ae" starling = "starling.scripts.starling_main_cli:main" starling-benchmark = "starling.scripts.starling_main_cli:starling_benchmark" diff --git a/starling/configs/configs.yaml b/starling/configs/configs.yaml index 3a367eb..941c3fe 100644 --- a/starling/configs/configs.yaml +++ b/starling/configs/configs.yaml @@ -2,7 +2,6 @@ defaults: - dataloader: dataloader.yaml - trainer: trainer.yaml - diffusion: diffusion.yaml - - unet: unet.yaml - sequence_encoder: sequence_encoder.yaml - _self_ diff --git a/starling/models/vit.py b/starling/models/vit.py index 1ee0585..a757fa4 100644 --- a/starling/models/vit.py +++ b/starling/models/vit.py @@ -3,8 +3,7 @@ from einops.layers.torch import Rearrange from torch import nn -from starling.models.transformer import DiTBlock -from starling.models.unet import SinusoidalPosEmb +from starling.models.transformer import DiTBlock, SinusoidalPosEmb class PatchEmbed(nn.Module): diff --git a/starling/training/diffusion_train.py b/starling/training/diffusion_train.py index fb52240..68fca8e 100644 --- a/starling/training/diffusion_train.py +++ b/starling/training/diffusion_train.py @@ -12,13 +12,11 @@ from pytorch_lightning.loggers import WandbLogger from pytorch_lightning.utilities.rank_zero import rank_zero_only -import starling.data.ddpm_loader as ddpm_loader from starling.data.argument_parser import get_params from starling.data.ddpm_loader_tar import DDPMDataLoader from starling.models.continuous_diffusion import ContinuousDiffusion from starling.models.diffusion import DiffusionModel from starling.models.transformer import SequenceEncoder -from starling.models.unet import UNetConditional from starling.models.vae import VAE from starling.models.vit import ViT @@ -82,14 +80,12 @@ def setup_data_module(cfg, effective_batch_size=None): def setup_models(config): - """Set up the UNet and Diffusion models.""" + """Set up the ViT backbone and Diffusion model.""" model_path = config.trainer.checkpoint diffusion_models = {"discrete": DiffusionModel, "continuous": ContinuousDiffusion} - unet_config_dict = OmegaConf.to_container(config.unet, resolve=True) seq_encoder_dict = OmegaConf.to_container(config.sequence_encoder, resolve=True) - UNet_model = UNetConditional(**unet_config_dict) vit = ViT(12, 512, 8, 512) sequence_encoder = SequenceEncoder(**seq_encoder_dict) @@ -118,7 +114,7 @@ def setup_models(config): **diffusion_config_dict, ) - return UNet_model, diffusion_model + return diffusion_model def setup_logger(config, diffusion_model): @@ -173,11 +169,11 @@ def train_model(cfg: DictConfig): dataset = setup_data_module(cfg, effective_batch_size=effective_batch_size) # Setup models - UNet_model, diffusion_model = setup_models(cfg) + diffusion_model = setup_models(cfg) # Save model architecture with open(f"{output_path}/model_architecture.txt", "w") as f: - f.write(str(UNet_model)) + f.write(str(diffusion_model)) # Setup logger wandb_logger = setup_logger(cfg, diffusion_model) From 9aa37636a7f003b09d5eecc3e87165ce7330483c Mon Sep 17 00:00:00 2001 From: Borna Novak Date: Wed, 12 Aug 2026 08:57:28 -0500 Subject: [PATCH 4/4] removed continuous-time diffusion model formulation --- starling/models/continuous_diffusion.py | 415 ------------------------ 1 file changed, 415 deletions(-) delete mode 100644 starling/models/continuous_diffusion.py diff --git a/starling/models/continuous_diffusion.py b/starling/models/continuous_diffusion.py deleted file mode 100644 index fd6fdf7..0000000 --- a/starling/models/continuous_diffusion.py +++ /dev/null @@ -1,415 +0,0 @@ -import math -from typing import List - -import pytorch_lightning as pl -import torch -import torch.nn.functional as F -from einops import reduce, repeat -from torch import nn, sqrt -from torch.amp import autocast -from torch.optim.lr_scheduler import ( - CosineAnnealingLR, - CosineAnnealingWarmRestarts, - LambdaLR, - OneCycleLR, -) -from torch.special import expm1 - -# Adapted from https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py - -# helpers - - -def exists(val): - return val is not None - - -def default(val, d): - if exists(val): - return val - return d() if callable(d) else d - - -# diffusion helpers - - -def right_pad_dims_to(x, t): - padding_dims = x.ndim - t.ndim - if padding_dims <= 0: - return t - return t.view(*t.shape, *((1,) * padding_dims)) - - -# continuous schedules - -# equations are taken from https://openreview.net/attachment?id=2LdBqxc1Yv&name=supplementary_material -# @crowsonkb Katherine's repository also helped here https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/utils.py - -# log(snr) that approximates the original linear schedule - - -def log(t, eps=1e-20): - return torch.log(t.clamp(min=eps)) - - -def beta_linear_log_snr(t): - return -log(expm1(1e-4 + 10 * (t**2))) - - -def alpha_cosine_log_snr(t, s=0.008): - return -log((torch.cos((t + s) / (1 + s) * math.pi * 0.5) ** -2) - 1, eps=1e-5) - - -# From paper https://arxiv.org/abs/2206.00364; equation 5 -def karras_log_snr(t, sigma_min=0.002, sigma_max=80.0, rho=7.0): - """ - Implements the noise schedule from Karras et al. (2022) - "Elucidating the Design Space of Diffusion-Based Generative Models" - """ - # Convert t from [0,1] to the sigma space - inverse_rho = 1.0 / rho - sigma = sigma_min**inverse_rho + t * ( - sigma_max**inverse_rho - sigma_min**inverse_rho - ) - sigma = sigma**rho - - # Convert sigma to log(SNR) - return -2 * torch.log(sigma) - - -class ContinuousDiffusion(pl.LightningModule): - def __init__( - self, - model, - set_lr, - config_scheduler, - noise_schedule="karras", - min_snr_loss_weight=False, - min_snr_gamma=5, - ): - super().__init__() - - # Save the hyperparameters of the model but ignore the encoder_model and the U-Net model - self.save_hyperparameters(ignore=["model"]) - - self.model = model - - self.set_lr = set_lr - self.config_scheduler = config_scheduler - - self.monitor = "epoch_val_loss" - - # continuous noise schedule related stuff - - if noise_schedule == "linear": - self.log_snr = beta_linear_log_snr - elif noise_schedule == "cosine": - self.log_snr = alpha_cosine_log_snr - elif noise_schedule == "karras": - self.log_snr = karras_log_snr - else: - raise ValueError(f"unknown noise schedule {noise_schedule}") - - # proposed https://arxiv.org/abs/2303.09556 - # can converge 3.4 times faster than baseline if used - - self.min_snr_loss_weight = min_snr_loss_weight - self.min_snr_gamma = min_snr_gamma - - self.sequence_embedding = nn.Embedding(21, self.model.labels_dim) - - latent_space_scaling_factor = torch.tensor(1.0, dtype=torch.float32) - - # Register the buffer - self.register_buffer("latent_space_scaling_factor", latent_space_scaling_factor) - - @property - def device(self): - return next(self.model.parameters()).device - - # training related functions - noise prediction - - def sequence2labels(self, sequences: List) -> torch.Tensor: - """ - Converts sequences to labels based on user defined models, - - Parameters - ---------- - sequences : List - A list of sequences to convert to labels - - Returns - ------- - torch.Tensor - Returns the labels for the decoder - - Raises - ------ - ValueError - If the labels are not one of the three options - """ - - encoded = self.sequence_embedding(sequences) - - return encoded - - @autocast("cuda", enabled=False) - def q_sample(self, x_start, times, masks=None, noise=None): - noise = default(noise, lambda: torch.randn_like(x_start)) - - log_snr = self.log_snr(times) - - log_snr_padded = right_pad_dims_to(x_start, log_snr) - alpha, sigma = sqrt(log_snr_padded.sigmoid()), sqrt((-log_snr_padded).sigmoid()) - x_noised = x_start * alpha + noise * sigma - - if masks is not None: - x_noised = x_noised * masks + x_start * (1 - masks) - - return x_noised, log_snr - - def random_times(self, batch_size): - # times are now uniform from 0 to 1 - return torch.zeros((batch_size,), device=self.device).float().uniform_(0, 1) - - def p_losses( - self, - x_start: torch.Tensor, - t: torch.Tensor, - labels: torch.Tensor = None, - noise: torch.Tensor = None, - masks: torch.Tensor = None, - ) -> torch.Tensor: - """ - Calculate model loss based on predicted vs actual noise. - - Parameters - ---------- - x_start : torch.Tensor - The starting tensor to denoise - t : torch.Tensor - Timesteps along the denoising-diffusion process - labels : torch.Tensor, optional - Condition labels for the model - noise : torch.Tensor, optional - Optional pre-defined noise, otherwise sampled from N(0,I) - masks : torch.Tensor, optional - Optional masks for conditional generation - - Returns - ------- - torch.Tensor - Mean MSE loss between predicted and actual noise - """ - # Use standard normal distribution if no noise provided - noise = torch.randn_like(x_start) if noise is None else noise - - # Apply noise according to timestep - noised_input, log_snr = self.q_sample(x_start=x_start, times=t, noise=noise) - - # Prepare condition labels - condition_labels = self.sequence2labels(labels) - - # Predict the noise - predicted_noise = self.model(noised_input, log_snr, condition_labels) - - # Calculate per-element loss and reduce to per-batch loss - per_element_loss = F.mse_loss(predicted_noise, noise, reduction="none") - per_batch_loss = reduce(per_element_loss, "b ... -> b", "mean") - - # Apply minimum SNR loss weighting if enabled - if self.min_snr_loss_weight: - snr = log_snr.exp() - loss_weight = snr.clamp(min=self.min_snr_gamma) / snr - per_batch_loss = per_batch_loss * loss_weight - - return per_batch_loss.mean() - - def forward( - self, x: torch.Tensor, labels: torch.Tensor, masks: torch.Tensor = None - ) -> torch.Tensor: - """ - Forward pass that samples random timesteps and calculates loss. - - Parameters - ---------- - x : torch.Tensor - Input tensor - labels : torch.Tensor - Condition labels - masks : torch.Tensor, optional - Optional masks - - Returns - ------- - torch.Tensor - Loss value - """ - batch_size = x.shape[0] - random_timesteps = self.random_times(batch_size) - - return self.p_losses(x, random_timesteps, labels, masks=masks) - - def _initialize_latent_scaling(self, latent_encoding: torch.Tensor) -> None: - """ - Initialize the latent space scaling factor using the first batch. - - Parameters - ---------- - latent_encoding : torch.Tensor - Batch of encoded latent vectors - """ - # Calculate local standard deviation - local_std = latent_encoding.std() - - # Gather from all processes and compute global standard deviation - gathered_std = self.all_gather(local_std) - mean_std = gathered_std.mean() - - # Set consistent scaling factor across all GPUs - scaling_factor = 1 / mean_std - self.latent_space_scaling_factor = scaling_factor.float().to(self.device) - - def training_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: - """ - Training step that encodes inputs and calculates diffusion loss. - - Parameters - ---------- - batch : torch.Tensor - Batch containing data and sequence labels - batch_idx : int - Index of the current batch - - Returns - ------- - torch.Tensor - Training loss - """ - latent_encoding, sequences = batch - - # Initialize scaling factor on first batch - if self.global_step == 0 and batch_idx == 0: - self._initialize_latent_scaling(latent_encoding) - - # Scale latent vectors to have unit standard deviation - normalized_latents = self.latent_space_scaling_factor * latent_encoding - - # Calculate diffusion loss - loss = self.forward(normalized_latents, labels=sequences) - - # Log the training loss - self.log("train_loss", loss, prog_bar=True, batch_size=latent_encoding.size(0)) - - return loss - - def validation_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor: - """Validation step that evaluates diffusion loss on a batch.""" - latent_encoding, sequences = batch - - # Scale the latent encoding to have unit std - latent_encoding = self.latent_space_scaling_factor * latent_encoding - - loss = self.forward(latent_encoding, labels=sequences) - - self.log( - "epoch_val_loss", - loss, - prog_bar=True, - sync_dist=True, - batch_size=latent_encoding.size(0), - ) - - return loss - - def configure_optimizers(self): - """ - Configure the optimizer and the learning rate scheduler for the model. - Here I am using NVIDIA suggested settings for learning rate and weight - decay. For ResNet50 they have seen best performance with CosineAnnealingLR, - initial learning rate of 0.256 for batch size of 256 and linearly scaling - it down/up for other batch sizes. The weight decay is set to 1/32768 for all - parameters except the batch normalization layers. For further information check: - https://catalog.ngc.nvidia.com/orgs/nvidia/resources/resnet_50_v1_5_for_pytorch - - Returns - ------- - List - Returns the optimizer and the learning rate scheduler - - Raises - ------ - ValueError - If the scheduler is not implemented - """ - optimizer = torch.optim.AdamW( - self.model.parameters(), - lr=self.set_lr, - betas=(0.9, 0.999), - eps=1e-08, - weight_decay=0.01, - amsgrad=False, - ) - - if self.config_scheduler == "CosineAnnealingWarmRestarts": - lr_scheduler = { - "scheduler": CosineAnnealingWarmRestarts( - optimizer, T_0=5, eta_min=1e-4 - ), - "monitor": self.monitor, - "interval": "epoch", - } - - elif self.config_scheduler == "OneCycleLR": - lr_scheduler = { - "scheduler": OneCycleLR( - optimizer, - max_lr=0.01, - total_steps=self.trainer.estimated_stepping_batches, - ), - "monitor": self.monitor, - "interval": "step", - } - - elif self.config_scheduler == "CosineAnnealingLR": - num_epochs = self.trainer.max_epochs - lr_scheduler = { - "scheduler": CosineAnnealingLR( - optimizer, - T_max=num_epochs, - eta_min=1e-8, - ), - "monitor": self.monitor, - "interval": "epoch", - } - elif self.config_scheduler == "LinearWarmupCosineAnnealingLR": - num_epochs = self.trainer.max_epochs - total_steps = self.trainer.estimated_stepping_batches - steps_per_epoch = total_steps // num_epochs - # Warmup for 5% of the total steps - warmup_steps = steps_per_epoch * int(num_epochs * 0.05) - - def lr_lambda(current_step): - if current_step < warmup_steps: - # Linear warmup phase - return current_step / max(1, warmup_steps) - else: - # Cosine annealing phase - eta_min = 1e-8 - remaining_steps = current_step - warmup_steps - current_epoch = remaining_steps // steps_per_epoch - cosine_factor = 0.5 * ( - 1 + math.cos(math.pi * current_epoch / num_epochs) - ) - return eta_min + (1 - eta_min) * cosine_factor - - lr_scheduler = { - "scheduler": LambdaLR(optimizer, lr_lambda=lr_lambda), - "monitor": self.monitor, - "interval": "step", - } - - else: - raise ValueError(f"{self.config_scheduler} lr_scheduler is not implemented") - - return [optimizer], [lr_scheduler]