ok

Author: johndpope

model.py

@@ -478,16 +478,9 @@ def process_tokens(self, t_c, t_r):
         return m_c, m_r
 
 
-class IMFEncoder(nn.Module):
-    def __init__(self, model):
-        super(IMFEncoder, self).__init__()
-        self.model = model
 
-    def forward(self, x_current, x_reference):
-        f_r = self.model.dense_feature_encoder(x_reference)
-        t_r = self.model.latent_token_encoder(x_reference)
-        t_c = self.model.latent_token_encoder(x_current)
-        return f_r, t_r, t_c
+
+            
 
 class MappingNetwork(nn.Module):
     def __init__(self, latent_dim, w_dim, depth):

onnxconv.py

@@ -5,6 +5,17 @@
 from PIL import Image
 from torchvision import transforms
 
+
+class IMFDecoder(nn.Module):
+    def __init__(self, model):
+        super(IMFDecoder, self).__init__()
+        self.model = model
+
+    def decode_latent_tokens(self,f_r,t_r,t_c):
+        return  self.model.decode_latent_tokens(f_r,t_r,t_c)
+
+
+
 # Define the IMFEncoder class
 class IMFEncoder(nn.Module):
     def __init__(self, model):