# 🌀Spatial Reasoners
**A Python package for spatial reasoning over continuous variables with generative denoising models.**
[](LICENSE)
[]()
[](https://pypi.org/project/spatialreasoners/)
[](https://spatialreasoners.github.io)
## Overview
<p align="center">
<img src="assets/overview.png" alt="Spatial Reasoners Overview"/>
</p>
🌀Spatial Reasoners is a Python package for spatial reasoning over continuous variables with generative denoising models. Denoising generative models have become the de-facto standard for image generation, due to their effectiveness in sampling from complex, high-dimensional distributions. Recently, they have started being explored in the context of reasoning over multiple continuous variables.
Our package provides a comprehensive framework to facilitate research in this area, offering easy-to-use interfaces to control:
* **Variable Mapping:** Seamlessly map variables from arbitrary data domains.
* **Generative Model Paradigms:** Flexibly work with a wide range of denoising formulations.
* **Samplers & Inference Strategies:** Implement and experiment with diverse samplers and inference techniques.
🌀Spatial Reasoners is a generalization of [Spatial Reasoning Models (SRMs)](https://geometric-rl.mpi-inf.mpg.de/srm/) to new domains, packaged as a reusable library for the research community.
## 🛠️ Installation
### Quick Install (Recommended)
Install Spatial Reasoners directly from PyPI:
```bash
pip install spatialreasoners
```
### Development Install
For development or to use the latest features:
```bash
git clone https://github.com/spatialreasoners/spatialreasoners.git
cd spatialreasoners
pip install -e .
```
### Requirements
- Python 3.11+ (Recommended: 3.13)
- PyTorch 1.13+
- PyTorch Lightning 2.0+
## 🚀 Quick Start
### Basic Usage with predefined experiments
```python
import spatialreasoners as sr
# 🚀 One-line training with sensible defaults
sr.run_training()
# 🔍 With enhanced type checking for better error messages
sr.run_training(enable_beartype=True)
# ⚙️ Customize training parameters
sr.run_training(overrides=[
"experiment=mnist_sudoku", # Use specific experiment
"trainer.max_epochs=50", # Train for 50 epochs
"data_loader.train.batch_size=32" # Adjust batch size
])
# 🔧 Advanced usage with different model architectures
sr.run_training(overrides=[
"denoising_model.denoiser=dit_l_2", # Use large DiT model
"denoising_model.flow=cosine", # Use cosine flow
"variable_mapper=image" # Image variable mapping
])
```
## 🏗️ Custom Projects & Training
Spatial Reasoners provides two clean approaches for creating custom research projects with your own datasets and models.
### Method 1: @sr.config_main Decorator (Recommended)
The cleanest interface for most use cases - similar to `@hydra.main` but with automatic config merging.
**Create your training script** (`training.py`):
```python
#!/usr/bin/env python3
import spatialreasoners as sr
# Import your custom components to auto-register them
import src # This imports and registers all your custom components
@sr.config_main(config_path="configs", config_name="main")
def main(cfg):
"""Main training function with full control over the process."""
# Create components from the loaded config
lightning_module = sr.create_lightning_module(cfg)
data_module = sr.create_data_module(cfg)
trainer = sr.create_trainer(cfg)
# Full control - add custom callbacks, modify trainer, etc.
trainer.fit(lightning_module, datamodule=data_module)
if __name__ == "__main__":
main()
```
**CLI Usage:**
```bash
# Basic training with your experiment
python training.py experiment=my_experiment
# Customize any parameter via CLI
python training.py experiment=my_experiment trainer.max_epochs=100
# Multiple overrides
python training.py experiment=my_experiment trainer.max_epochs=50 dataset.subset_size=15000
# Enable dynamic type checking
python training.py experiment=my_experiment --enable-beartype
# Get help and examples
python training.py --help
```
**Advantages:**
- ✅ **Cleanest interface** - just like `@hydra.main`
- ✅ **Automatic config merging** (local + embedded configs)
- ✅ **No boilerplate code** - just import, decorate, and run
- ✅ **Full control** - inspect and modify config before training
- ✅ **Easy type-safe debugging** via `--enable-beartype` flag
### Method 2: Programmatic Configuration
For automation, notebooks, or when you need to generate configurations dynamically.
```python
#!/usr/bin/env python3
import spatialreasoners as sr
# Import your custom components to auto-register them
import src
def main():
"""Programmatic training configuration."""
# Define overrides as needed
overrides = [
"experiment=my_experiment",
"trainer.max_epochs=100",
"dataset.subset_size=20000"
]
# Run training with the overrides
sr.run_training(
config_name="main",
config_path="configs",
overrides=overrides,
enable_beartype=True,
)
if __name__ == "__main__":
main()
```
**Advantages:**
- ✅ **Programmatic control** - generate configs dynamically
- ✅ **Easy integration** into larger Python programs
- ✅ **Good for automation** - scripts, pipelines, notebooks
- ✅ **No CLI complexity** - simple function calls
### Configuration Structure
Organize your project with this recommended structure:
```
your_project/
├── training.py # Your main training script
├── src/ # Custom components
│ ├── __init__.py # Auto-register components
│ ├── dataset.py # Custom datasets
│ ├── denoiser.py # Custom models
│ ├── variable_mapper.py # Custom variable mappers
│ └── tokenizer.py # Custom tokenizers
└── configs/ # Configuration files
├── main.yaml # Main config (references experiments)
├── experiment/ # Experiment-specific configs
│ └── my_experiment.yaml
├── dataset/ # Custom dataset configs
└── variable_mapper/ # Custom mapper configs
```
**Example main config** (`configs/main.yaml`):
```yaml
defaults:
- experiment: null # Users specify via CLI
- time_sampler: mean_beta
- optimizer: default
# Your project-specific defaults
trainer:
max_steps: 3000
val_check_interval: 1000
data_loader:
train:
batch_size: 128
num_workers: 16
```
**Example experiment config** (`configs/experiment/my_experiment.yaml`):
```yaml
# @package _global_
defaults:
- /dataset: my_custom_dataset # Your custom dataset
# Mix local and embedded components
variable_mapper:
name: my_custom_mapper
denoising_model:
flow: rectified # From embedded configs
denoiser:
name: my_custom_model
```
### Register Custom Components
Define custom components in your `src/` directory and auto-register them:
```python
# src/dataset.py
import spatialreasoners as sr
from spatialreasoners.dataset import register_dataset, DatasetCfg
from dataclasses import dataclass
@dataclass
class MyDatasetCfg(DatasetCfg):
name: str = "my_dataset"
data_path: str = "data/"
subset_size: int = 10000
@register_dataset("my_dataset", MyDatasetCfg)
class MyDataset(sr.Dataset):
def __init__(self, cfg: MyDatasetCfg):
# Your dataset implementation
pass
```
```python
# src/denoiser.py
import spatialreasoners as sr
from spatialreasoners.denoising_model.denoiser import register_denoiser, DenoiserCfg
@dataclass
class MyModelCfg(DenoiserCfg):
name: str = "my_model"
hidden_dim: int = 256
@register_denoiser("my_model", MyModelCfg)
class MyModel(sr.Denoiser):
def __init__(self, cfg: MyModelCfg, tokenizer, num_classes=None):
# Your model implementation
pass
```
```python
from src import * # Register all components
from . import dataset
from . import denoiser
from . import variable_mapper
from . import tokenizer
# Add other component imports as needed
```
### Config Merging
🌀Spatial Reasoners automatically merges your local configs with embedded configurations:
- **Local configs take precedence** - your custom components override built-in ones
- **Built-in components remain accessible** - use `dataset=cifar10`, `denoising_model.flow=rectified`, etc.
- **Seamless composition** - mix and match local and embedded components freely
### Quick Comparison
| Method | Interface | CLI Support | Setup | Best For |
|--------|-----------|-------------|-------|----------|
| `@sr.config_main` | Decorator | ✅ Automatic | Minimal | General use, research, experimentation |
| Programmatic | Function | ❌ None | Minimal | Automation, notebooks, production |
**Recommendation:** Start with Method 1 (`@sr.config_main`) for most use cases. Use Method 2 for automation or when generating configurations dynamically.
## 📖 Documentation & Examples
### Example Projects
Check out the `example_project/` directory for a complete working example that demonstrates:
- **Two training approaches**: `@sr.config_main` decorator and programmatic configuration
- **Custom component organization**: Structured `src/` directory with auto-registration on import
- **Config composition**: Local configs that reference embedded Spatial Reasoners components
- **Professional workflows**: Proper project structure for research projects
The example implements a spiral dataset where the model learns to generate points along a spiral pattern, showcasing:
- Custom dataset, variable mapper, tokenizer, and denoiser implementations
- Clean configuration management with experiment-specific configs
- Visualization and evaluation during training
**Run the example:**
```bash
cd example_project
# Method 1: @sr.config_main decorator (recommended)
python training_decorator.py experiment=spiral_training
# Method 2: Programmatic configuration
python training_programmatic.py
```
### Configuration System
Spatial Reasoners uses Hydra for flexible configuration management with automatic merging between your local configs and embedded components.
**Key Configuration Concepts:**
- **Main Config** (`configs/main.yaml`): Project-wide defaults and structure (eg. wandb setups, float precision, etc.)
- **Experiments** (`configs/experiment/`): Complete run-specific configurations
- **Component Configs**: Modular configs for datasets, models, etc.
- **Embedded Components**: Built-in configs from Spatial Reasoners (datasets, flows, optimizers, configurable denoisers)
You can inspect the embedded configs here in the `config` directory. For most datasets, you will need to provide the `root` directory -- the place where your images/videos are.
**Advanced Configuration Loading:**
```python
# Multiple ways to load and customize configs
config = sr.load_default_config() # Built-in api_default experiment
config = sr.load_config_from_yaml(overrides=["experiment=mnist_sudoku"])
config = sr.load_config_from_yaml("./configs", "main", ["experiment=custom"])
# Programmatic config modification
config.trainer.max_epochs = 100
config.data_loader.train.batch_size = 32
```
**CLI Configuration:**
```bash
# Use embedded experiments
python training.py experiment=mnist_sudoku
# Override any nested parameter
python training.py experiment=mnist_sudoku trainer.max_epochs=100 data_loader.train.batch_size=64
# Mix local and embedded components
python training.py experiment=my_experiment denoising_model.flow=cosine optimizer=adamw
```
## 💾 Datasets & Checkpoints
### Datasets
We provide datasets from the original SRM project. Download them from the [SRM releases](https://github.com/Chrixtar/SRM/releases):
```bash
# Extract datasets.zip to your data directory
mkdir -p data
cd data
wget https://github.com/Chrixtar/SRM/releases/download/v1.0/datasets.zip
unzip datasets.zip
```
For FFHQ-based datasets, download [FFHQ](https://github.com/NVlabs/ffhq-dataset) and update the path in your dataset config.
### Pretrained Models
Download pretrained checkpoints from the [SRM releases](https://github.com/Chrixtar/SRM/releases):
```bash
mkdir -p checkpoints
cd checkpoints
wget https://github.com/Chrixtar/SRM/releases/download/v1.0/checkpoints.zip
unzip checkpoints.zip
```
## 📊 Research & Benchmarks
### Running Benchmarks
Evaluate models on standard benchmarks:
```python
import spatialreasoners as sr
@sr.config_main(config_path="configs", config_name="main_spiral")
def demo_spiral_testing(cfg):
# Create components from the loaded config
lightning_module = sr.create_lightning_module(cfg)
data_module = sr.create_data_module(cfg)
trainer = sr.create_trainer(cfg)
checkpoint_path = cfg.checkpointing.load
trainer.test(lightning_module, datamodule=data_module, ckpt_path=checkpoint_path)
print("🎉 Testing complete!")
```
## 🏗️ Architecture
Spatial Reasoners is built with modularity and extensibility in mind:
```
spatialreasoners/
├── api/ # High-level API
├── dataset/ # Data loading and processing
├── denoising_model/ # Model implementations
│ ├── denoiser/ # Denoiser architectures (UNet, DiT, MAR, etc.)
│ ├── flow/ # Flow variants (rectified, cosine, etc.)
│ └── tokenizer/ # Tokenizers of variables for the denoiser
├── training/ # Training infrastructure
├── variable_mapper/ # Variable mapping logic
├── benchmark/ # Evaluation framework
└── configs/ # Embedded default configs
```
### Key Components
- **Variable Mappers**: Split data sample into a set of variables, that could have different noise levels
- **Denoising Models**: Combined objects containing of a Denoiser (predictive object), Tokenizer (transforming the variables to model's inputs), Flow (eg. rectified, cosine or continuous diffusion) and parametrization ($u_t$, $\epsilon$, $x_0$, $v$)
- **Training System**: PyTorch Lightning-based training with full configurability
- **Benchmark Sets**: Each benchmark is a pair of a Dataset and an Evaluation
## 🔬 Research Applications
Spatial Reasoners are a generalization of the idea of diffusion models that allows using different noise levels within a sample. Before SRMs this approach has been explored by for example MAR, xAR, Rolling Diffusion and Diffusion Forcing -- Spatial Reasoners allows you to build similiar setups. For some architectures (such as Unet, DiT, xAR's variant of DiT or History Guided Diffusion's U-ViT-pose) you can just specify the denoiser config and directly start training.
In some domains starting your work could be even faster due to already implemented Variable Mappers and some evaluations -- this is true for tasks like:
- **Sudoku generation** Our MNIST Sudoku dataset
- **Image generation** With prepared dataset implementations for ImageNet, CIFAR10, CelebA, SRM's Counting Stars and many others
- **Video generation** Where a variable is a single frame -- as in [Diffusion Forcing](https://www.boyuan.space/diffusion-forcing/)
We also highly encourage you to take Spatial Reasoners to completely new domains -- see our [example project](https://github.com/spatialreasoners/spatialreasoners/tree/rc/sr-0.1.4-c/example_project) to see how to train new models in your domain!
### Citation
If you use Spatial Reasoners in your research, please cite:
```bibtex
@software{pogodzinski25spatialreasoners,
title={Spatial Reasoners for Continuous Variables in Any Domain},
author={Pogodzinski, Bart and Wewer, Christopher and Lenssen, Jan Eric and Schiele, Bernt},
year={2025},
url={https://spatialreasoners.github.io}
}
@inproceedings{wewer25srm,
title = {Spatial Reasoning with Denoising Models},
author = {Wewer, Christopher and Pogodzinski, Bartlomiej and Schiele, Bernt and Lenssen, Jan Eric},
booktitle = {International Conference on Machine Learning ({ICML})},
year = {2025},
}
```
## 🤝 Contributing
We welcome contributions from the research community! Here's how you can help:
### Ways to Contribute
- **New Models**: Implement novel denoising architectures
- **Datasets**: Add support for new spatial reasoning tasks
- **Benchmarks**: Contribute evaluation protocols
- **Documentation**: Improve docs and examples
- **Bug Reports**: Report issues and suggest improvements
### Development Setup
```bash
git clone https://github.com/spatialreasoners/spatialreasoners.git
cd spatialreasoners
pip install -e ".[dev]"
```
### Running Tests
```bash
pytest tests/
```
## 📜 License
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
## 🙋 Support & Community
- **Documentation**: [spatialreasoners.github.io](https://spatialreasoners.github.io)
- **Issues**: [GitHub Issues](https://github.com/spatialreasoners/spatialreasoners/issues)
- **Email**: bpogodzi@mpi-inf.mpg.de
Raw data
{
"_id": null,
"home_page": null,
"name": "spatialreasoners",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.11",
"maintainer_email": null,
"keywords": "machine learning, deep learning, spatial reasoning, pytorch, pytorch lightning, diffusion models, Spatial Reasoners, SRM",
"author": null,
"author_email": "Bart Pogodzinski <bpogodzi@mpi-inf.mpg.de>, Christopher Wewer <cwewer@mpi-inf.mpg.de>, Jan Eric Lenssen <jlenssen@mpi-inf.mpg.de>, Bernt Schiele <schiele@mpi-inf.mpg.de>",
"download_url": "https://files.pythonhosted.org/packages/6b/18/1cab78e12e93c4e620e4c4a3f0d1f7b9a907ede34d3ae9eb9fdd2b2eb1dc/spatialreasoners-0.1.4.tar.gz",
"platform": null,
"description": "# \ud83c\udf00Spatial Reasoners \n\n**A Python package for spatial reasoning over continuous variables with generative denoising models.**\n\n[](LICENSE)\n[]()\n[](https://pypi.org/project/spatialreasoners/)\n[](https://spatialreasoners.github.io)\n\n## Overview\n\n<p align=\"center\">\n <img src=\"assets/overview.png\" alt=\"Spatial Reasoners Overview\"/>\n</p>\n\n\ud83c\udf00Spatial Reasoners is a Python package for spatial reasoning over continuous variables with generative denoising models. Denoising generative models have become the de-facto standard for image generation, due to their effectiveness in sampling from complex, high-dimensional distributions. Recently, they have started being explored in the context of reasoning over multiple continuous variables.\n\nOur package provides a comprehensive framework to facilitate research in this area, offering easy-to-use interfaces to control:\n\n* **Variable Mapping:** Seamlessly map variables from arbitrary data domains.\n* **Generative Model Paradigms:** Flexibly work with a wide range of denoising formulations.\n* **Samplers & Inference Strategies:** Implement and experiment with diverse samplers and inference techniques.\n\n\ud83c\udf00Spatial Reasoners is a generalization of [Spatial Reasoning Models (SRMs)](https://geometric-rl.mpi-inf.mpg.de/srm/) to new domains, packaged as a reusable library for the research community.\n\n## \ud83d\udee0\ufe0f Installation\n\n### Quick Install (Recommended)\n\nInstall Spatial Reasoners directly from PyPI:\n\n```bash\npip install spatialreasoners\n```\n\n### Development Install\n\nFor development or to use the latest features:\n\n```bash\ngit clone https://github.com/spatialreasoners/spatialreasoners.git\ncd spatialreasoners\npip install -e .\n```\n\n### Requirements\n\n- Python 3.11+ (Recommended: 3.13)\n- PyTorch 1.13+\n- PyTorch Lightning 2.0+\n\n## \ud83d\ude80 Quick Start\n\n### Basic Usage with predefined experiments\n\n```python\nimport spatialreasoners as sr\n\n# \ud83d\ude80 One-line training with sensible defaults\nsr.run_training()\n\n# \ud83d\udd0d With enhanced type checking for better error messages\nsr.run_training(enable_beartype=True)\n\n# \u2699\ufe0f Customize training parameters\nsr.run_training(overrides=[\n \"experiment=mnist_sudoku\", # Use specific experiment\n \"trainer.max_epochs=50\", # Train for 50 epochs\n \"data_loader.train.batch_size=32\" # Adjust batch size\n])\n\n# \ud83d\udd27 Advanced usage with different model architectures\nsr.run_training(overrides=[\n \"denoising_model.denoiser=dit_l_2\", # Use large DiT model\n \"denoising_model.flow=cosine\", # Use cosine flow\n \"variable_mapper=image\" # Image variable mapping\n])\n```\n\n## \ud83c\udfd7\ufe0f Custom Projects & Training\n\nSpatial Reasoners provides two clean approaches for creating custom research projects with your own datasets and models.\n\n### Method 1: @sr.config_main Decorator (Recommended)\n\nThe cleanest interface for most use cases - similar to `@hydra.main` but with automatic config merging.\n\n**Create your training script** (`training.py`):\n```python\n#!/usr/bin/env python3\nimport spatialreasoners as sr\n\n# Import your custom components to auto-register them\nimport src # This imports and registers all your custom components\n\n@sr.config_main(config_path=\"configs\", config_name=\"main\")\ndef main(cfg):\n \"\"\"Main training function with full control over the process.\"\"\"\n \n # Create components from the loaded config\n lightning_module = sr.create_lightning_module(cfg)\n data_module = sr.create_data_module(cfg)\n trainer = sr.create_trainer(cfg)\n \n # Full control - add custom callbacks, modify trainer, etc.\n trainer.fit(lightning_module, datamodule=data_module)\n\nif __name__ == \"__main__\":\n main()\n```\n\n**CLI Usage:**\n```bash\n# Basic training with your experiment\npython training.py experiment=my_experiment\n\n# Customize any parameter via CLI\npython training.py experiment=my_experiment trainer.max_epochs=100\n\n# Multiple overrides\npython training.py experiment=my_experiment trainer.max_epochs=50 dataset.subset_size=15000\n\n# Enable dynamic type checking\npython training.py experiment=my_experiment --enable-beartype\n\n# Get help and examples\npython training.py --help\n```\n\n**Advantages:**\n- \u2705 **Cleanest interface** - just like `@hydra.main`\n- \u2705 **Automatic config merging** (local + embedded configs)\n- \u2705 **No boilerplate code** - just import, decorate, and run\n- \u2705 **Full control** - inspect and modify config before training\n- \u2705 **Easy type-safe debugging** via `--enable-beartype` flag\n\n### Method 2: Programmatic Configuration\n\nFor automation, notebooks, or when you need to generate configurations dynamically.\n\n```python\n#!/usr/bin/env python3\nimport spatialreasoners as sr\n\n# Import your custom components to auto-register them\nimport src\n\ndef main():\n \"\"\"Programmatic training configuration.\"\"\"\n \n # Define overrides as needed\n overrides = [\n \"experiment=my_experiment\", \n \"trainer.max_epochs=100\",\n \"dataset.subset_size=20000\"\n ]\n \n # Run training with the overrides\n sr.run_training(\n config_name=\"main\",\n config_path=\"configs\",\n overrides=overrides,\n enable_beartype=True,\n )\n\nif __name__ == \"__main__\":\n main()\n```\n\n**Advantages:**\n- \u2705 **Programmatic control** - generate configs dynamically\n- \u2705 **Easy integration** into larger Python programs\n- \u2705 **Good for automation** - scripts, pipelines, notebooks\n- \u2705 **No CLI complexity** - simple function calls\n\n### Configuration Structure\n\nOrganize your project with this recommended structure:\n\n```\nyour_project/\n\u251c\u2500\u2500 training.py # Your main training script\n\u251c\u2500\u2500 src/ # Custom components\n\u2502 \u251c\u2500\u2500 __init__.py # Auto-register components\n\u2502 \u251c\u2500\u2500 dataset.py # Custom datasets\n\u2502 \u251c\u2500\u2500 denoiser.py # Custom models \n\u2502 \u251c\u2500\u2500 variable_mapper.py # Custom variable mappers\n\u2502 \u2514\u2500\u2500 tokenizer.py # Custom tokenizers\n\u2514\u2500\u2500 configs/ # Configuration files\n \u251c\u2500\u2500 main.yaml # Main config (references experiments)\n \u251c\u2500\u2500 experiment/ # Experiment-specific configs\n \u2502 \u2514\u2500\u2500 my_experiment.yaml\n \u251c\u2500\u2500 dataset/ # Custom dataset configs\n \u2514\u2500\u2500 variable_mapper/ # Custom mapper configs\n```\n\n**Example main config** (`configs/main.yaml`):\n```yaml\ndefaults:\n - experiment: null # Users specify via CLI\n - time_sampler: mean_beta\n - optimizer: default\n\n# Your project-specific defaults\ntrainer:\n max_steps: 3000\n val_check_interval: 1000\n \ndata_loader:\n train:\n batch_size: 128\n num_workers: 16\n```\n\n**Example experiment config** (`configs/experiment/my_experiment.yaml`):\n```yaml\n# @package _global_\ndefaults:\n - /dataset: my_custom_dataset # Your custom dataset\n\n# Mix local and embedded components\nvariable_mapper:\n name: my_custom_mapper\n\ndenoising_model:\n flow: rectified # From embedded configs\n denoiser:\n name: my_custom_model\n```\n\n### Register Custom Components\n\nDefine custom components in your `src/` directory and auto-register them:\n\n```python\n# src/dataset.py\nimport spatialreasoners as sr\nfrom spatialreasoners.dataset import register_dataset, DatasetCfg\nfrom dataclasses import dataclass\n\n@dataclass \nclass MyDatasetCfg(DatasetCfg):\n name: str = \"my_dataset\"\n data_path: str = \"data/\"\n subset_size: int = 10000\n\n@register_dataset(\"my_dataset\", MyDatasetCfg)\nclass MyDataset(sr.Dataset):\n def __init__(self, cfg: MyDatasetCfg):\n # Your dataset implementation\n pass\n```\n\n```python\n# src/denoiser.py \nimport spatialreasoners as sr\nfrom spatialreasoners.denoising_model.denoiser import register_denoiser, DenoiserCfg\n\n@dataclass\nclass MyModelCfg(DenoiserCfg):\n name: str = \"my_model\"\n hidden_dim: int = 256\n\n@register_denoiser(\"my_model\", MyModelCfg)\nclass MyModel(sr.Denoiser):\n def __init__(self, cfg: MyModelCfg, tokenizer, num_classes=None):\n # Your model implementation\n pass\n```\n\n```python\nfrom src import * # Register all components\n\nfrom . import dataset\nfrom . import denoiser\nfrom . import variable_mapper\nfrom . import tokenizer\n# Add other component imports as needed\n```\n\n### Config Merging\n\n\ud83c\udf00Spatial Reasoners automatically merges your local configs with embedded configurations:\n\n- **Local configs take precedence** - your custom components override built-in ones\n- **Built-in components remain accessible** - use `dataset=cifar10`, `denoising_model.flow=rectified`, etc.\n- **Seamless composition** - mix and match local and embedded components freely\n\n### Quick Comparison\n\n| Method | Interface | CLI Support | Setup | Best For |\n|--------|-----------|-------------|-------|----------|\n| `@sr.config_main` | Decorator | \u2705 Automatic | Minimal | General use, research, experimentation |\n| Programmatic | Function | \u274c None | Minimal | Automation, notebooks, production |\n\n**Recommendation:** Start with Method 1 (`@sr.config_main`) for most use cases. Use Method 2 for automation or when generating configurations dynamically.\n\n## \ud83d\udcd6 Documentation & Examples\n\n### Example Projects\n\nCheck out the `example_project/` directory for a complete working example that demonstrates:\n\n- **Two training approaches**: `@sr.config_main` decorator and programmatic configuration\n- **Custom component organization**: Structured `src/` directory with auto-registration on import \n- **Config composition**: Local configs that reference embedded Spatial Reasoners components\n- **Professional workflows**: Proper project structure for research projects\n\nThe example implements a spiral dataset where the model learns to generate points along a spiral pattern, showcasing:\n- Custom dataset, variable mapper, tokenizer, and denoiser implementations\n- Clean configuration management with experiment-specific configs\n- Visualization and evaluation during training\n\n**Run the example:**\n```bash\ncd example_project\n\n# Method 1: @sr.config_main decorator (recommended)\npython training_decorator.py experiment=spiral_training\n\n# Method 2: Programmatic configuration \npython training_programmatic.py\n```\n\n### Configuration System\n\nSpatial Reasoners uses Hydra for flexible configuration management with automatic merging between your local configs and embedded components.\n\n**Key Configuration Concepts:**\n\n- **Main Config** (`configs/main.yaml`): Project-wide defaults and structure (eg. wandb setups, float precision, etc.)\n- **Experiments** (`configs/experiment/`): Complete run-specific configurations\n- **Component Configs**: Modular configs for datasets, models, etc.\n- **Embedded Components**: Built-in configs from Spatial Reasoners (datasets, flows, optimizers, configurable denoisers)\n\nYou can inspect the embedded configs here in the `config` directory. For most datasets, you will need to provide the `root` directory -- the place where your images/videos are. \n\n**Advanced Configuration Loading:**\n```python\n# Multiple ways to load and customize configs\nconfig = sr.load_default_config() # Built-in api_default experiment\nconfig = sr.load_config_from_yaml(overrides=[\"experiment=mnist_sudoku\"])\nconfig = sr.load_config_from_yaml(\"./configs\", \"main\", [\"experiment=custom\"])\n\n# Programmatic config modification\nconfig.trainer.max_epochs = 100\nconfig.data_loader.train.batch_size = 32\n```\n\n**CLI Configuration:**\n```bash\n# Use embedded experiments\npython training.py experiment=mnist_sudoku\n\n# Override any nested parameter\npython training.py experiment=mnist_sudoku trainer.max_epochs=100 data_loader.train.batch_size=64\n\n# Mix local and embedded components \npython training.py experiment=my_experiment denoising_model.flow=cosine optimizer=adamw\n```\n\n## \ud83d\udcbe Datasets & Checkpoints\n\n### Datasets\nWe provide datasets from the original SRM project. Download them from the [SRM releases](https://github.com/Chrixtar/SRM/releases):\n\n```bash\n# Extract datasets.zip to your data directory\nmkdir -p data\ncd data\nwget https://github.com/Chrixtar/SRM/releases/download/v1.0/datasets.zip\nunzip datasets.zip\n```\n\nFor FFHQ-based datasets, download [FFHQ](https://github.com/NVlabs/ffhq-dataset) and update the path in your dataset config.\n\n### Pretrained Models\nDownload pretrained checkpoints from the [SRM releases](https://github.com/Chrixtar/SRM/releases):\n\n```bash\nmkdir -p checkpoints\ncd checkpoints\nwget https://github.com/Chrixtar/SRM/releases/download/v1.0/checkpoints.zip\nunzip checkpoints.zip\n```\n\n## \ud83d\udcca Research & Benchmarks\n\n### Running Benchmarks\n\nEvaluate models on standard benchmarks:\n\n```python\nimport spatialreasoners as sr\n\n@sr.config_main(config_path=\"configs\", config_name=\"main_spiral\")\ndef demo_spiral_testing(cfg):\n # Create components from the loaded config\n lightning_module = sr.create_lightning_module(cfg)\n data_module = sr.create_data_module(cfg)\n trainer = sr.create_trainer(cfg)\n \n checkpoint_path = cfg.checkpointing.load\n trainer.test(lightning_module, datamodule=data_module, ckpt_path=checkpoint_path)\n \n print(\"\ud83c\udf89 Testing complete!\")\n```\n\n## \ud83c\udfd7\ufe0f Architecture\n\nSpatial Reasoners is built with modularity and extensibility in mind:\n\n```\nspatialreasoners/\n\u251c\u2500\u2500 api/ # High-level API\n\u251c\u2500\u2500 dataset/ # Data loading and processing\n\u251c\u2500\u2500 denoising_model/ # Model implementations\n\u2502 \u251c\u2500\u2500 denoiser/ # Denoiser architectures (UNet, DiT, MAR, etc.)\n\u2502 \u251c\u2500\u2500 flow/ # Flow variants (rectified, cosine, etc.)\n\u2502 \u2514\u2500\u2500 tokenizer/ # Tokenizers of variables for the denoiser\n\u251c\u2500\u2500 training/ # Training infrastructure\n\u251c\u2500\u2500 variable_mapper/ # Variable mapping logic\n\u251c\u2500\u2500 benchmark/ # Evaluation framework\n\u2514\u2500\u2500 configs/ # Embedded default configs\n```\n\n### Key Components\n\n- **Variable Mappers**: Split data sample into a set of variables, that could have different noise levels\n- **Denoising Models**: Combined objects containing of a Denoiser (predictive object), Tokenizer (transforming the variables to model's inputs), Flow (eg. rectified, cosine or continuous diffusion) and parametrization ($u_t$, $\\epsilon$, $x_0$, $v$)\n- **Training System**: PyTorch Lightning-based training with full configurability\n- **Benchmark Sets**: Each benchmark is a pair of a Dataset and an Evaluation\n\n## \ud83d\udd2c Research Applications\n\nSpatial Reasoners are a generalization of the idea of diffusion models that allows using different noise levels within a sample. Before SRMs this approach has been explored by for example MAR, xAR, Rolling Diffusion and Diffusion Forcing -- Spatial Reasoners allows you to build similiar setups. For some architectures (such as Unet, DiT, xAR's variant of DiT or History Guided Diffusion's U-ViT-pose) you can just specify the denoiser config and directly start training. \n\nIn some domains starting your work could be even faster due to already implemented Variable Mappers and some evaluations -- this is true for tasks like:\n\n- **Sudoku generation** Our MNIST Sudoku dataset\n- **Image generation** With prepared dataset implementations for ImageNet, CIFAR10, CelebA, SRM's Counting Stars and many others\n- **Video generation** Where a variable is a single frame -- as in [Diffusion Forcing](https://www.boyuan.space/diffusion-forcing/)\n\nWe also highly encourage you to take Spatial Reasoners to completely new domains -- see our [example project](https://github.com/spatialreasoners/spatialreasoners/tree/rc/sr-0.1.4-c/example_project) to see how to train new models in your domain!\n\n\n\n### Citation\n\nIf you use Spatial Reasoners in your research, please cite:\n\n```bibtex\n@software{pogodzinski25spatialreasoners,\n title={Spatial Reasoners for Continuous Variables in Any Domain},\n author={Pogodzinski, Bart and Wewer, Christopher and Lenssen, Jan Eric and Schiele, Bernt},\n year={2025},\n url={https://spatialreasoners.github.io}\n}\n\n@inproceedings{wewer25srm,\n title = {Spatial Reasoning with Denoising Models},\n author = {Wewer, Christopher and Pogodzinski, Bartlomiej and Schiele, Bernt and Lenssen, Jan Eric},\n booktitle = {International Conference on Machine Learning ({ICML})},\n year = {2025},\n}\n```\n\n## \ud83e\udd1d Contributing\n\nWe welcome contributions from the research community! Here's how you can help:\n\n### Ways to Contribute\n\n- **New Models**: Implement novel denoising architectures\n- **Datasets**: Add support for new spatial reasoning tasks\n- **Benchmarks**: Contribute evaluation protocols\n- **Documentation**: Improve docs and examples\n- **Bug Reports**: Report issues and suggest improvements\n\n### Development Setup\n\n```bash\ngit clone https://github.com/spatialreasoners/spatialreasoners.git\ncd spatialreasoners\npip install -e \".[dev]\"\n```\n\n### Running Tests\n\n```bash\npytest tests/\n```\n\n## \ud83d\udcdc License\n\nThis project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.\n\n## \ud83d\ude4b Support & Community\n\n- **Documentation**: [spatialreasoners.github.io](https://spatialreasoners.github.io)\n- **Issues**: [GitHub Issues](https://github.com/spatialreasoners/spatialreasoners/issues)\n- **Email**: bpogodzi@mpi-inf.mpg.de\n",
"bugtrack_url": null,
"license": null,
"summary": "SpatialReasoners: A framework for training Spatial Reasoning Models in any domain",
"version": "0.1.4",
"project_urls": {
"Documentation": "https://spatialreasoners.github.io/",
"Homepage": "https://github.com/spatialreasoners/spatialreasoners",
"Issues": "https://github.com/spatialreasoners/spatialreasoners/issues",
"Repository": "https://github.com/spatialreasoners/spatialreasoners"
},
"split_keywords": [
"machine learning",
" deep learning",
" spatial reasoning",
" pytorch",
" pytorch lightning",
" diffusion models",
" spatial reasoners",
" srm"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "2d9753c7059a272ec27afb3a0a58e226ee7926af102db14fe3e1848ccaca68e1",
"md5": "a960147071142a26435bb08cc772be1b",
"sha256": "3c6c3c1040f2c9e1b62775f58d494b3098720de3b278e05d532df73e29b200f5"
},
"downloads": -1,
"filename": "spatialreasoners-0.1.4-py3-none-any.whl",
"has_sig": false,
"md5_digest": "a960147071142a26435bb08cc772be1b",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.11",
"size": 258172,
"upload_time": "2025-07-13T04:43:14",
"upload_time_iso_8601": "2025-07-13T04:43:14.060896Z",
"url": "https://files.pythonhosted.org/packages/2d/97/53c7059a272ec27afb3a0a58e226ee7926af102db14fe3e1848ccaca68e1/spatialreasoners-0.1.4-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": null,
"digests": {
"blake2b_256": "6b181cab78e12e93c4e620e4c4a3f0d1f7b9a907ede34d3ae9eb9fdd2b2eb1dc",
"md5": "08b13a7ca6034203f6d204396b02feb2",
"sha256": "19dd103fd48dc8f25db86d99acd63659b51e54268d4ac033d5e022d9ba2ef10d"
},
"downloads": -1,
"filename": "spatialreasoners-0.1.4.tar.gz",
"has_sig": false,
"md5_digest": "08b13a7ca6034203f6d204396b02feb2",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.11",
"size": 178492,
"upload_time": "2025-07-13T04:43:15",
"upload_time_iso_8601": "2025-07-13T04:43:15.594389Z",
"url": "https://files.pythonhosted.org/packages/6b/18/1cab78e12e93c4e620e4c4a3f0d1f7b9a907ede34d3ae9eb9fdd2b2eb1dc/spatialreasoners-0.1.4.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-07-13 04:43:15",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "spatialreasoners",
"github_project": "spatialreasoners",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [
{
"name": "accelerate",
"specs": []
},
{
"name": "av",
"specs": []
},
{
"name": "beartype",
"specs": []
},
{
"name": "colorama",
"specs": []
},
{
"name": "colorspacious",
"specs": []
},
{
"name": "dacite",
"specs": []
},
{
"name": "diffusers",
"specs": []
},
{
"name": "einops",
"specs": []
},
{
"name": "gdown",
"specs": []
},
{
"name": "hydra-core",
"specs": []
},
{
"name": "jaxtyping",
"specs": []
},
{
"name": "lightning",
"specs": []
},
{
"name": "matplotlib",
"specs": []
},
{
"name": "pandas",
"specs": []
},
{
"name": "pytest",
"specs": []
},
{
"name": "pytest-beartype",
"specs": []
},
{
"name": "scipy",
"specs": []
},
{
"name": "roma",
"specs": []
},
{
"name": "rotary_embedding_torch",
"specs": []
},
{
"name": "timm",
"specs": []
},
{
"name": "torch",
"specs": []
},
{
"name": "torchaudio",
"specs": []
},
{
"name": "torchvision",
"specs": []
},
{
"name": "wandb",
"specs": []
},
{
"name": "wandb",
"specs": []
}
],
"lcname": "spatialreasoners"
}
JSON
