update README

condition/midas/midas/vit.py

@@ -61,7 +61,8 @@ class Transpose(nn.Module):
 def forward_vit(pretrained, x):
     b, c, h, w = x.shape
 
-    glob = pretrained.model.forward_flex(x)
+    # glob = pretrained.model.forward_flex(x)
+    glob = flat_forward_flex(pretrained.model, x)
 
     # layer_1 = pretrained.activations["1"]
     # layer_2 = pretrained.activations["2"]
@@ -127,6 +128,44 @@ def _resize_pos_embed(self, posemb, gs_h, gs_w):
     return posemb
 
 
+def flat_forward_flex(model, x):
+    b, c, h, w = x.shape
+
+    pos_embed = model._resize_pos_embed(model.pos_embed,
+                                        h // model.patch_size[1],
+                                        w // model.patch_size[0])
+
+    B = x.shape[0]
+
+    if hasattr(model.patch_embed, "backbone"):
+        x = model.patch_embed.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[
+                -1]  # last feature if backbone outputs list/tuple of features
+
+    x = model.patch_embed.proj(x).flatten(2).transpose(1, 2)
+
+    if getattr(model, "dist_token", None) is not None:
+        cls_tokens = model.cls_token.expand(
+            B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        dist_token = model.dist_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, dist_token, x), dim=1)
+    else:
+        cls_tokens = model.cls_token.expand(
+            B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+
+    x = x + pos_embed
+    x = model.pos_drop(x)
+
+    for blk in model.blocks:
+        x = blk(x)
+
+    x = model.norm(x)
+
+    return x
+
+
 def forward_flex(self, x):
     b, c, h, w = x.shape
 
@@ -174,6 +213,7 @@ def get_activation(name):
 
     return hook
 
+
 def hook_act(module, input, output):
     activations.append(output)
 
@@ -212,7 +252,7 @@ def _make_vit_b16_backbone(
     pretrained.model.blocks[hooks[1]].register_forward_hook(hook_act)
     pretrained.model.blocks[hooks[2]].register_forward_hook(hook_act)
     pretrained.model.blocks[hooks[3]].register_forward_hook(hook_act)
-    
+
     # pretrained.model.blocks[hooks[0]].register_forward_hook(
     #     get_activation("1"))
     # pretrained.model.blocks[hooks[1]].register_forward_hook(
@@ -386,20 +426,16 @@ def _make_vit_b_rn50_backbone(
     activations = []
 
     if use_vit_only == True:
-        pretrained.model.blocks[hooks[0]].register_forward_hook(
-            hook_act)
-        pretrained.model.blocks[hooks[1]].register_forward_hook(
-            hook_act)
+        pretrained.model.blocks[hooks[0]].register_forward_hook(hook_act)
+        pretrained.model.blocks[hooks[1]].register_forward_hook(hook_act)
     else:
         pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
             hook_act)
         pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
             hook_act)
 
-    pretrained.model.blocks[hooks[2]].register_forward_hook(
-        hook_act)
-    pretrained.model.blocks[hooks[3]].register_forward_hook(
-       hook_act)
+    pretrained.model.blocks[hooks[2]].register_forward_hook(hook_act)
+    pretrained.model.blocks[hooks[3]].register_forward_hook(hook_act)
 
     # if use_vit_only == True:
     #     pretrained.model.blocks[hooks[0]].register_forward_hook(