some boilerplate for continuous gaussian diffusion with sigmoid noise schedule formulated in log snr form

Author: lucidrains

README.md

@@ -8,6 +8,74 @@ Implementation of <a href="https://arxiv.org/abs/2212.11972">Recurrent Interface
 
 The big surprise is that the generations can reach this level of fidelity. Will need to verify this on my own machine
 
+## Install
+
+```bash
+$ pip install rin-pytorch
+```
+
+## Usage
+
+```python
+from rin_pytorch import RIN, Trainer, GaussianDiffusion
+
+model = RIN(
+    dim = 32,
+    channels = 3,
+    dim_mults = (1, 2, 4, 8),
+).cuda()
+
+diffusion = GaussianDiffusion(
+    model,
+    image_size = 128,
+    timesteps = 100,
+    use_ddim = True              # use ddim
+).cuda()
+
+trainer = Trainer(
+    diffusion,
+    '/path/to/your/data',             # path to your folder of images
+    results_folder = './results',     # where to save results
+    num_samples = 16,                 # number of samples
+    train_batch_size = 4,             # training batch size
+    gradient_accumulate_every = 4,    # gradient accumulation
+    train_lr = 1e-4,                  # learning rate
+    save_and_sample_every = 1000,     # how often to save and sample
+    train_num_steps = 700000,         # total training steps
+    ema_decay = 0.995,                # exponential moving average decay
+)
+
+trainer.train()
+```
+
+Results will be saved periodically to the `./results` folder
+
+If you would like to experiment with the `RIN` and `GaussianDiffusion` class outside the `Trainer`
+
+```python
+import torch
+from rin_pytorch import RIN, GaussianDiffusion
+
+model = RIN(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+)
+
+diffusion = GaussianDiffusion(
+    model,
+    image_size = 128,
+    timesteps = 1000
+)
+
+training_images = torch.randn(8, 3, 128, 128) # images are normalized from 0 to 1
+loss = diffusion(training_images)
+loss.backward()
+# after a lot of training
+
+sampled_images = diffusion.sample(batch_size = 4)
+sampled_images.shape # (4, 3, 128, 128)
+```
+
 ## Citations
 
 ```bibtex

rin_pytorch/__init__.py

@@ -0,0 +1 @@
+from rin_pytorch.rin_pytorch import GaussianDiffusion, RIN, Trainer