upscael AE output if needed (body ae)

models/base.py

@@ -72,11 +72,12 @@ class AEEncoder(tf.keras.Model):
 
 class Decoder(tf.keras.Model):
     
-    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False):
+    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
         super(Decoder, self).__init__()
 
         self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
         self.reshape = tf.keras.layers.Reshape(input_size)
+        self.upscale = upscale
         
         if not is_label:
             self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
@@ -94,6 +95,10 @@ class Decoder(tf.keras.Model):
         x = self.reshape(x)
         x = self.decoder(x, training=training)
 
+        if self.upscale:
+            x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+
+
         if len(inputs) == 2:
             if len(mask.shape) == 2:
                 x = x * mask[tf.newaxis,:,:,tf.newaxis]