YAML Metadata Warning:empty or missing yaml metadata in repo card
Check out the documentation for more information.
LiquidFlow: Liquid Neural Network × Mamba-2 SSD Image Generator
A lightweight, physics-informed image generator combining Liquid Neural Networks (CfC) with Mamba-2 State Space Duality — trainable on Google Colab free tier, deployable on mobile devices.
Architecture
┌────────────────────────────────┐
Image [128×128] → │ TAESD VAE │ → Latent [16×16×4]
│ (< 1M params) │
└────────────────────────────────┘
↓
┌────────────────────────────────┐
│ LiquidFlow Backbone │
│ │
│ ┌──────────────────────────┐ │
│ │ LiquidMamba Block (×N) │ │
│ │ │ │
│ │ Input → CfC Gate │ │
│ │ ↓ │ │
│ │ Mamba-2 SSD │ │
│ │ (multi-dir scan) │ │
│ │ ↓ │ │
│ │ CfC Gate → Output │ │
│ └──────────────────────────┘ │
│ │
│ + Physics-Informed Loss │
│ (TV + Spectral + Gradient) │
└────────────────────────────────┘
↓
Predicted Noise
Core Innovations
CfC (Closed-form Continuous-time) Liquid Neural Networks
h(t) = σ(-f(x,I)·t) ⊙ g(x,I) + (1-σ(-f(x,I)·t)) ⊙ h(x,I)- No ODE solving — 100× faster than Neural ODEs
- Time-continuous adaptive gating mechanism
- From: Hasani et al., Nature Machine Intelligence (2022)
Mamba-2 SSD (State Space Duality)
h_t = A_t·h_{t-1} + B_t·x_t,y_t = C_t^T·h_t- O(N) linear complexity (vs O(N²) attention)
- Fully parallelizable via associative scan
- Pure PyTorch — no CUDA kernels needed
- From: Dao & Gu, "Transformers are SSMs" (2024)
Physics-Informed Regularization
- Total Variation + Spectral + Gradient constraints
- Training-only regularizer — zero inference cost
- Pattern from: Bastek & Sun, ICLR 2025
TAESD VAE
- < 1M parameters — 84× smaller than SD VAE
- Near-instant encoding/decoding
- From: madebyollin/taesd
Model Variants
| Variant | Parameters | Hidden Dim | Stages | Blocks/Stage | T4 VRAM |
|---|---|---|---|---|---|
| Tiny | ~2M | 128 | 2 | 2 | < 2 GB |
| Small | ~8M | 256 | 4 | 4 | ~4 GB |
| Base | ~30M | 384 | 6 | 6 | ~8 GB |
Quick Start
Google Colab
- Open the notebook
- Runtime → Change runtime type → GPU (T4)
- Run all cells
Local Training
# Clone
git clone https://huggingface.co/krystv/LiquidFlow-Gen
cd LiquidFlow-Gen
# Install
pip install torch torchvision diffusers tqdm pillow numpy
# Train (small model, 128px, CIFAR-10)
python train.py \
--dataset cifar10 \
--image_size 128 \
--variant small \
--batch_size 32 \
--epochs 100 \
--lr 2e-4
# Train (base model, 512px)
python train.py \
--dataset cifar10 \
--image_size 512 \
--variant base \
--batch_size 8 \
--epochs 200 \
--lr 1e-4
Generate Samples
from liquid_flow.generator import create_liquidflow
from liquid_flow.vae_wrapper import TAESDWrapper
# Load model
model = create_liquidflow(variant='small', image_size=128)
model.load_state_dict(torch.load('best_model.pt'))
model = model.cuda().eval()
# Load VAE
vae = TAESDWrapper.load('cuda')
# Generate
latents = model.sample(batch_size=16, steps=50, ddim=True)
images = TAESDWrapper.decode(vae, latents)
Training Details
Default Hyperparameters
- Optimizer: AdamW (β₁=0.9, β₂=0.999)
- LR: 2×10⁻⁴ (tiny/small), 1×10⁻⁴ (base)
- Weight Decay: 10⁻⁴
- LR Schedule: Cosine annealing
- Gradient Clipping: 1.0
- AMP: Enabled (when CUDA available)
Physics Regularization Weights
- TV (Total Variation): 0.01
- Conservation of Intensity: 0.001
- Spectral Regularizer: 0.01
- Gradient Penalty: 0.001
Datasets Supported
- CIFAR-10, CIFAR-100, STL-10
- CelebA, LSUN (requires download)
- ImageNet (provide path)
Mobile Deployment
LiquidFlow uses pure PyTorch — no custom CUDA kernels:
# Export to ONNX
torch.onnx.export(model, (x, t), 'liquidflow.onnx',
input_names=['noisy_latent', 'timestep'],
output_names=['predicted_noise'],
opset_version=14)
# Convert to CoreML (iOS)
# coremltools.converters.onnx.convert(model='liquidflow.onnx')
# Convert to TFLite (Android)
# onnx-tf convert -i liquidflow.onnx -o liquidflow.pb
Why This Works (Research Validation)
DiMSUM (NeurIPS 2024)
Mamba-based diffusion beats DiT transformers on ImageNet generation (FID 2.11 vs 2.27). Mamba's O(N) complexity enables 3× faster convergence than attention-based models.
PINNMamba (ICML 2025)
SSM + Physics constraints are compatible and synergistic. Mamba's selective scan naturally handles the spatio-temporal nature of PDE residuals.
LiteVAE / TAESD
Wavelet-based and tiny VAEs provide sufficient latent quality for diffusion at < 1% of the parameter count of standard VAEs. TAESD is used by 100+ real-time diffusion demos on HF Spaces.
DeepSeek V3 Insights
- Auxiliary-loss-free training (apply to physics weights)
- Multi-head architecture for efficiency
- DualPipe for overlapping computation
Repository Structure
LiquidFlow-Gen/
├── liquid_flow/
│ ├── __init__.py # Package init
│ ├── cfc_cell.py # CfC Liquid NN implementation
│ ├── mamba2_ssd.py # Mamba-2 SSD implementation
│ ├── liquid_flow_block.py # Hybrid CfC+Mamba block
│ ├── generator.py # Full diffusion generator
│ ├── vae_wrapper.py # VAE interfaces
│ ├── physics_loss.py # Physics regularizers
│ └── trainer.py # Training utilities
├── train.py # CLI training script
├── LiquidFlow_Colab.ipynb # Colab notebook
└── README.md # This file
Citations
@article{hasani2022cfc,
title={Closed-form continuous-time neural networks},
author={Hasani, Ramin and Lechner, Mathias and Amini, Alexander and others},
journal={Nature Machine Intelligence},
year={2022}
}
@article{dao2024mamba2,
title={Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
author={Dao, Tri and Gu, Albert},
journal={arXiv:2405.21060},
year={2024}
}
@inproceedings{bastek2025physics,
title={Physics-Informed Diffusion Models},
author={Bastek, Jan-Hendrik and Sun, WaiChing},
booktitle={ICLR},
year={2025}
}
@article{pham2024dimsum,
title={DiMSUM: Diffusion Mamba -- A Scalable and Unified Spatial-Frequency Method for Image Generation},
author={Pham, Hao and others},
journal={NeurIPS},
year={2024}
}
License
MIT