Simplified inference

app.py

@@ -39,7 +39,7 @@ torch.backends.cudnn.allow_tf32 = True
 # Global no_grad
 torch.set_grad_enabled(False)
 
-
+K_KEEP_LIST = [1, 2, 4, 8, 16, 32, 64, 128, 256]
 MAX_SEED = np.iinfo(np.int32).max
 
 MODEL_ID = 'EPFL-VILAB/flextok_d18_d28_dfn'
@@ -79,12 +79,10 @@ def infer(img_path, seed=0, randomize_seed=False, timesteps=20, cfg_scale=7.5, p
 
     # Tokenize images once
     with get_bf16_context(enable_bf16):
-        tokens_list = flextok_model.tokenize(imgs)
+        tokens = flextok_model.tokenize(imgs)[0] # 1x256 
 
     # Create all token subsequences
-    k_keep_list = [1, 2, 4, 8, 16, 32, 64, 128, 256]
-    tokens_list = tokens_list*len(k_keep_list)
-    subseq_list = [seq[:,:k_keep].clone() for seq, k_keep in zip(tokens_list, k_keep_list)]
+    subseq_list = [tokens[:,:k_keep].clone() for k_keep in K_KEEP_LIST] # [1x1, 1x2, 1x4, ..., 1x256]
 
     # Detokenize various subsequences in parallel. Batch size is 9.
     with get_bf16_context(enable_bf16):
@@ -98,7 +96,7 @@ def infer(img_path, seed=0, randomize_seed=False, timesteps=20, cfg_scale=7.5, p
     # Transform to PIL images
     all_images = [
         (TF.to_pil_image(denormalize(reconst_k).clamp(0,1)), f'{k_keep} tokens') 
-        for reconst_k, k_keep in zip(all_reconst, k_keep_list)
+        for reconst_k, k_keep in zip(all_reconst, K_KEEP_LIST)
     ]
 
     return all_images