commit all

models/cond_encoder.py

@@ -1,5 +1,5 @@
 import tensorflow as tf
-from models.base import AEEncoder, Decoder
+from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
 
 
 class AutoEncoder(tf.keras.Model):
@@ -11,10 +11,28 @@ class AutoEncoder(tf.keras.Model):
         self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
 
 
-    def call(self, inputs, training=None):
+    def call(self, inputs, training):
         x = inputs
 
         x = self.encoder(x, training=training)
         x = self.decoder(x, training=training)
 
         return x
+
+
+class AutoEncoderNew(tf.keras.Model):
+    
+    def __init__(self, config):
+        super(AutoEncoderNew, self).__init__()
+
+        self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
+        self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
+
+
+    def call(self, inputs, training):
+        x = inputs
+
+        x = self.encoder(x, training=training)
+        x = self.decoder(x, training=training)
+
+        return x