E-RayZer implementation

.gitignore

@@ -184,4 +184,16 @@ app.log
 
 # data
 /data/
-*.png
+*.png
+
+# machine-specific run / inspection scripts (private paths, accounts, node names)
+scripts/*.sbatch
+scripts/export_pointcloud.py
+scripts/preview_pointcloud.py
+scripts/smoke_erayzer_load.py
+scripts/viz_erayzer.py
+
+# fine-tune config holds a machine-specific init_checkpoint path (kept local)
+configs/multiview/erayzer_cheese3d_ft.yaml
+# local model checkpoints (machine-specific, large)
+checkpoints/

CONTRIBUTING.md

@@ -15,7 +15,8 @@ See our [code of conduct](CODE_OF_CONDUCT.md) for more information.
 install in editable mode with dev dependencies:
 
 ```bash
-pip install -e ".[dev]"
+pip install lightning poetry-core
+pip install -e ".[dev]" --no-build-isolation
 pre-commit install
 ```
 

README.md

@@ -45,15 +45,23 @@ For Github cloning:
 ```commandline
 git clone https://github.com/paninski-lab/beast
 cd beast
-pip install -e .
+pip install lightning poetry-core
+pip install -e . --no-build-isolation
 ```
 
 For installation through PyPI:
 
 ```commandline
-pip install beast-backbones
+pip install lightning poetry-core
+pip install beast-backbones --no-build-isolation
 ```
 
+> **Note:** `beast` depends on a custom fork of
+> [gsplat](https://github.com/QitaoZhao/gsplat) that must be compiled from source. The
+> `gsplat` build requires `torch` (provided by `lightning`) and the build backend requires
+> `poetry-core`. Installing these first and using `--no-build-isolation` lets the build find
+> them in your environment.
+
 ## Usage
 
 > The commands below are for the **single-view BEAST model**.

beast/api/model.py

@@ -10,14 +10,12 @@
 
 import torch
 
-from beast.config import BeastConfig
+from beast.config import get_beast_config_class
 from beast.inference import predict_images, predict_video
 from beast.io import load_config
 from beast.logging import log_step
 from beast.models.base import BaseLightningModel
-from beast.models.resnets import ResnetAutoencoder
-from beast.models.vits import VisionTransformer
-from beast.train import train
+from beast.models.registry import MODEL_REGISTRY, TRAIN_REGISTRY
 
 _logger = logging.getLogger(__name__)
 
@@ -46,12 +44,6 @@ class Model:
     This class manages both the model and the training/inference processes.
     """
 
-    MODEL_REGISTRY = {
-        'vit': VisionTransformer,
-        'resnet': ResnetAutoencoder,
-        # Add more models as needed
-    }
-
     def __init__(
         self,
         model: BaseLightningModel,
@@ -83,11 +75,11 @@ def from_dir(cls, model_dir: str | Path) -> 'Model':
         config = load_config(config_path)
 
         model_type = config['model'].get('model_class', '').lower()
-        if model_type not in cls.MODEL_REGISTRY:
+        if model_type not in MODEL_REGISTRY:
             raise ValueError(f'Unknown model type: {model_type}')
 
         # Initialize the LightningModule
-        model_class = cls.MODEL_REGISTRY[model_type]
+        model_class = MODEL_REGISTRY[model_type]
         model = model_class(config)
 
         _logger.info(f'Loaded a {model_class} model')
@@ -116,14 +108,15 @@ def from_config(cls, config_path: str | Path | dict) -> 'Model':
         if not isinstance(config_path, dict):
             config = load_config(config_path)
         else:
-            config = BeastConfig.model_validate(config_path).model_dump()
+            model_class = (config_path.get('model') or {}).get('model_class', '')
+            config = get_beast_config_class(model_class).model_validate(config_path).model_dump()
 
         model_type = config['model'].get('model_class', '').lower()
-        if model_type not in cls.MODEL_REGISTRY:
+        if model_type not in MODEL_REGISTRY:
             raise ValueError(f'Unknown model type: {model_type}')
 
         # Initialize the LightningModule
-        model_class = cls.MODEL_REGISTRY[model_type]
+        model_class = MODEL_REGISTRY[model_type]
         log_step(f"Creating {model_type} model instance", level='debug')
         log_step(
             f"About to call {model_class.__name__}.__init__() - this may take several"
@@ -148,8 +141,10 @@ def train(self, output_dir: str | Path = 'runs/default') -> None:
 
         """
         self.model_dir = Path(output_dir)
+        model_type = self.config['model'].get('model_class', '').lower()
+        train_fn = TRAIN_REGISTRY[model_type]
         with chdir(self.model_dir):
-            self.model = train(self.config, self.model, output_dir=self.model_dir)
+            self.model = train_fn(self.config, self.model, output_dir=self.model_dir)
 
     def predict_images(
         self,

beast/config.py

@@ -1,4 +1,32 @@
-"""Pydantic models for BEAST configuration validation."""
+"""Pydantic schemas for BEAST configuration validation and per-model dispatch.
+
+Top-level config classes
+------------------------
+BeastConfig
+    Used for models that share the standard training loop (resnet, vit).
+    Contains TrainingConfig, OptimizerConfig, and DataConfig.
+
+ERayZerBeastConfig
+    Used for ERayZer and its subclasses (e.g. BEAST3D if they share ERayZer's
+    training schema).  Contains ERayZerTrainingConfig and ERayZerOptimizerConfig.
+
+Shared schemas (reusable for new models)
+-----------------------------------------
+TrainingConfig     — standard epoch-based training fields (batch sizes, epochs, imgaug, …)
+OptimizerConfig    — AdamW/Adam + cosine/step scheduler fields
+DataConfig         — data_dir only; extend or replace for richer data configs
+
+Per-model dispatch
+------------------
+get_beast_config_class(model_class) returns the correct top-level config class for a
+given model_class string.  Both beast.io.load_config and beast.api.model.Model.from_config
+use this function so that YAML files and raw dicts are validated against the right schema.
+
+To add a new model with a divergent training schema:
+  1. Define <Model>TrainingConfig and <Model>OptimizerConfig in the model's *_config.py
+  2. Define <Model>BeastConfig here (avoids circular imports since DataConfig lives here)
+  3. Add 'model_class': <Model>BeastConfig to _MODEL_CONFIG_CLASSES
+"""
 
 from __future__ import annotations
 
@@ -7,68 +35,21 @@
 
 from pydantic import BaseModel, Field
 
+from beast.models.beast_resnet.beast_resnet_config import ResnetModelConfig
+from beast.models.beast_vit.beast_vit_config import VitModelConfig
+from beast.models.erayzer.erayzer_config import (
+    ERayZerModelConfig,
+    ERayZerOptimizerConfig,
+    ERayZerTrainingConfig,
+)
+
 
 class BeastConfig(BaseModel):
     model: ModelConfig
     training: TrainingConfig
     optimizer: OptimizerConfig
     data: DataConfig
-
-
-class ResnetModelConfig(BaseModel):
-    model_class: Literal['resnet']
-    model_params: ResnetModelParams
-    seed: int = 0
-    checkpoint: str | None = None
-
-
-class ResnetModelParams(BaseModel):
-    backbone: Literal[
-        'resnet18',
-        'resnet34',
-        'resnet50',
-        'resnet101',
-        'resnet152',
-    ] = 'resnet18'
-    num_latents: int | None = None
-    image_size: int = 224
-    num_channels: int = 3
-
-
-class VitModelConfig(BaseModel):
-    model_class: Literal['vit']
-    model_params: VitModelParams
-    seed: int = 0
-    checkpoint: str | None = None
-
-
-class VitModelParams(BaseModel):
-    hidden_size: int = 768
-    num_hidden_layers: int = 12
-    num_attention_heads: int = 12
-    intermediate_size: int = 3072
-    hidden_act: str = 'gelu'
-    hidden_dropout_prob: float = 0.0
-    attention_probs_dropout_prob: float = 0.0
-    initializer_range: float = 0.02
-    layer_norm_eps: float = 1e-12
-    image_size: int = 224
-    patch_size: int = 16
-    num_channels: int = 3
-    qkv_bias: bool = True
-    decoder_num_attention_heads: int = 16
-    decoder_hidden_size: int = 512
-    decoder_num_hidden_layers: int = 8
-    decoder_intermediate_size: int = 2048
-    mask_ratio: float = 0.75
-    norm_pix_loss: bool = False
-    embed_size: int = 768
-    temp_scale: bool = False
-    random_init: bool = False
-    use_infoNCE: bool = False
-    infoNCE_weight: float = 0.03
-    use_perceptual_loss: bool = False
-    lambda_perceptual: float = 10.0
+    inference: bool = False
 
 
 ModelConfig = Annotated[
@@ -113,4 +94,34 @@ class DataConfig(BaseModel):
     data_dir: str | Path
 
 
+class ERayZerBeastConfig(BaseModel):
+    """Complete top-level config for ERayZer training runs."""
+
+    model: ERayZerModelConfig
+    training: ERayZerTrainingConfig
+    optimizer: ERayZerOptimizerConfig
+    data: DataConfig
+    inference: bool = False
+
+
+_MODEL_CONFIG_CLASSES: dict[str, type[BaseModel]] = {
+    'erayzer': ERayZerBeastConfig,
+}
+
+
+def get_beast_config_class(model_class: str) -> type[BaseModel]:
+    """Return the top-level Pydantic config class for the given model_class identifier.
+
+    Parameters
+    ----------
+    model_class: model type identifier string (e.g., 'resnet', 'vit', 'erayzer')
+
+    Returns
+    -------
+    Pydantic model class for the full top-level config
+
+    """
+    return _MODEL_CONFIG_CLASSES.get(model_class, BeastConfig)
+
+
 BeastConfig.model_rebuild()