Fixed dim_z to 120, modified path to be dynamic and cleaned VAE constructor of unnecessary parameters

model.py

@@ -1,4 +1,6 @@
 import tensorflow as tf
+import os
+import inspect
 
 
 _CAP = 3501 # Cap for the number of notes
@@ -112,16 +114,24 @@ class VAECost:
 
 class VAE(tf.keras.Model):
 
-  def __init__(self, dim_z=120, seed=2000, analytic_kl=True, name="autoencoder", **kwargs):
+  def __init__(self, **kwargs):
     super(VAE, self).__init__(name=name, **kwargs)
     self.dim_x = (3, _CAP, 1)
-    self.dim_z = dim_z
-    self.seed = seed
-    self.analytic_kl = analytic_kl
-    self.encoder = Encoder_Z(dim_z=self.dim_z).build()
-    self.decoder = Decoder_X(dim_z=self.dim_z).build()
+    self.encoder = Encoder_Z(dim_z=120).build()
+    self.decoder = Decoder_X(dim_z=120).build()
     self.cost_func = VAECost(self)
-    self.load_weights("./weights/")
+
+    # Get the path of the script that defines this method
+    script_path = inspect.getfile(inspect.currentframe())
+
+    # Get the directory containing the script
+    script_dir = os.path.dirname(os.path.abspath(script_path))
+
+    # Construct the path to the weights folder
+    weights_dir = os.path.join(script_dir, 'weights') + os.sep
+    
+    # Load pretrained weights
+    self.load_weights(weights_dir)
 
   @tf.function()
   def train_step(self, data):
@@ -143,12 +153,12 @@ class VAE(tf.keras.Model):
       
     mu, rho = tf.split(self.encoder(x_input), num_or_size_splits=2, axis=1)
     sd = tf.math.log(1 + tf.math.exp(rho))
-    z_sample = mu + sd * tf.random.normal(shape=(self.dim_z,))
+    z_sample = mu + sd * tf.random.normal(shape=(120,))
     return z_sample, mu, sd
 
   def generate(self, z_sample=None):
     # Decode a latent representation of a song, which is provided or sampled
       
     if z_sample == None:
-      z_sample = tf.expand_dims(tf.random.normal(shape=(self.dim_z,)), axis=0)
+      z_sample = tf.expand_dims(tf.random.normal(shape=(120,)), axis=0)
     return self.decoder(z_sample)