Initial CosAE release

README.md

@@ -0,0 +1,105 @@
+# CosAE: Convolutional Harmonic Autoencoder
+CosAE is a PyTorch implementation of the Convolutional Harmonic Autoencoder (CosAE). It encodes images into learnable harmonic representations (amplitudes and phases), constructs spatial cosine bases via a Harmonic Construction Module, and decodes back to RGB images. This repository provides the core model code, a Jupyter notebook for training and evaluation, and pretrained weights in SafeTensors format.
+
+## Features
+- Convolutional encoder with residual blocks and optional global attention
+- Harmonic Construction Module (HCM) learning per-channel frequencies
+- Decoder with upsampling and optional global attention for reconstruction
+- Supports raw RGB input or augmented FFT channels (RGB + FFT)
+- Pretrained model weights available (SafeTensors)
+- Training and tracking via Jupyter notebook and Weights & Biases
+
+## Installation
+### Requirements
+```bash
+pip install -r requirements.txt
+```
+
+## Quickstart
+### Inference with Pretrained Model
+```python
+import torch
+from cosae import CosAEModel
+
+# Load pretrained CosAE (expects 9-channel input: 3 RGB + 6 FFT)
+model = CosAEModel.from_pretrained("model/model-final")
+model.eval()
+
+# Dummy input: batch of 1, 9 channels, 256×256
+x = torch.randn(1, 9, 256, 256)
+with torch.no_grad():
+    recon = model(x)  # recon.shape == [1, 3, 256, 256]
+```
+
+### Preparing Raw Images
+```python
+from PIL import Image
+from torchvision import transforms
+
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+])
+img = Image.open("path/to/image.png").convert("RGB")
+tensor = transform(img).unsqueeze(0)  # [1, 3, 256, 256]
+```
+
+### Generating FFT Channels (Optional)
+To use FFT-augmented input (in_channels=9), compute the 2D FFT per RGB channel and stack real/imaginary parts:
+```python
+import torch
+fft = torch.fft.rfft2(tensor, norm="ortho")        # complex tensor [1,3,H,W/2+1]
+real, imag = fft.real, fft.imag
+# Optionally pad or reshape to match full spatial dims
+x9 = torch.cat([tensor, real, imag], dim=1)         # [1,9,H,W]
+```
+
+## Creating a Model from Scratch
+Use the `CosAEConfig` and `CosAEModel` classes to instantiate a model with custom settings:
+```python
+from cosae.config import CosAEConfig
+from cosae.cosae import CosAEModel
+
+# 1) Define a configuration (example uses RGB input only)
+config = CosAEConfig(
+    in_channels=3,
+    hidden_dims=[64, 128, 256, 512],
+    downsample_strides=[2, 2, 2, 2],
+    use_encoder_attention=True,
+    encoder_attention_heads=8,
+    encoder_attention_layers=1,
+    bottleneck_channels=256,
+    basis_size=16,
+    decoder_hidden_dim=256,
+    decoder_upsample_strides=[2],
+    use_decoder_attention=False,
+)
+
+# 2) Instantiate the model
+model = CosAEModel(config)
+
+# 3) (Optional) Save the model and config for later reuse
+model.save_pretrained("./my_cosae_model")  # creates model-final.safetensors and config.json
+```
+
+## Training and Evaluation
+A full training and evaluation pipeline is provided in `cosine-ae.ipynb`. Launch Jupyter to run experiments and track metrics with Weights & Biases:
+```bash
+jupyter lab cosine-ae.ipynb
+```
+
+## Repository Structure
+```
+.
+├── cosae/               # Core model implementation (config, encoder, HCM, decoder)
+├── model/               # Pretrained weights and config (SafeTensors)
+├── cosine-ae.ipynb      # Notebook: training, evaluation, and demos
+├── LICENSE              # MIT License
+└── README.md            # Project overview and usage
+```
+
+## Contributing
+Contributions, issues, and feature requests are welcome. Feel free to fork the repository and submit pull requests.
+
+## License
+This project is released under the MIT License. See [LICENSE](LICENSE) for details.

config.json

@@ -0,0 +1,40 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "CosAEModel"
+  ],
+  "basis_size": 16,
+  "bottleneck_channels": 256,
+  "decoder_attention_heads": 8,
+  "decoder_attention_layers": 0,
+  "decoder_hidden_dim": 256,
+  "decoder_upsample_strides": [
+    2
+  ],
+  "downsample_strides": [
+    2,
+    2,
+    2,
+    2
+  ],
+  "encoder_attention_heads": 8,
+  "encoder_attention_layers": 1,
+  "hidden_dims": [
+    64,
+    128,
+    256,
+    512
+  ],
+  "image_size": [
+    256,
+    256
+  ],
+  "in_channels": 9,
+  "model_type": "cosae",
+  "norm_type": "gn",
+  "num_res_blocks": 2,
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "use_decoder_attention": false,
+  "use_encoder_attention": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52bfdf9130a8688360fad5fba4763aca935a358c6221e8e0b2f009fdc7930f2f
+size 72900732