add continuation feature

README.md

@@ -86,6 +86,29 @@ recon_img = patches_to_images(top_level)
 loss = F.mse_loss(img, recon_img)
 loss.backward()
 ```
+
+You can pass in the state of the column and levels back into the model to continue where you left off (perhaps if you are processing consecutive frames of a slow video, as mentioned in the paper)
+
+```python
+import torch
+from glom_pytorch import Glom
+
+model = Glom(
+    dim = 512,
+    levels = 6,
+    image_size = 224,
+    patch_size = 14
+)
+
+img1 = torch.randn(1, 3, 224, 224)
+img2 = torch.randn(1, 3, 224, 224)
+img3 = torch.randn(1, 3, 224, 224)
+
+levels1 = model(img1, iters = 12)                   # show image 1 for 12 iterations
+levels2 = model(img2, levels = levels1, iters = 10) # image 2 for 10 iteratoins
+levels3 = model(img3, levels = levels2, iters = 6)  # image 3 for 6 iterations
+```
+
 ### Appreciation
 
 Thanks goes out to <a href="https://github.com/cfoster0">Cfoster0</a> for reviewing the code

glom_pytorch/glom_pytorch.py

@@ -108,7 +108,7 @@ def __init__(
         # consensus attention
         self.attention = ConsensusAttention(num_patches_side, attend_self = consensus_self, local_consensus_radius = local_consensus_radius)
 
-    def forward(self, img, iters = None, return_all = False):
+    def forward(self, img, iters = None, levels = None, return_all = False):
         b, h, w, _, device = *img.shape, img.device
         iters = default(iters, self.levels * 2)   # need to have twice the number of levels of iterations in order for information to propagate up and back down. can be overridden
 
@@ -121,7 +121,8 @@ def forward(self, img, iters = None, return_all = False):
         bottom_level = tokens
         bottom_level = rearrange(bottom_level, 'b n d -> b n () d')
 
-        levels = repeat(self.init_levels, 'l d -> b n l d', b = b, n = n)
+        if not exists(levels):
+            levels = repeat(self.init_levels, 'l d -> b n l d', b = b, n = n)
 
         hiddens = [levels]
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'glom-pytorch',
   packages = find_packages(),
-  version = '0.0.10',
+  version = '0.0.11',
   license='MIT',
   description = 'Glom - Pytorch',
   author = 'Phil Wang',