caching model generation

app.py

@@ -8,32 +8,33 @@ import transforms as trans
 
 P_MODEL = "model/drilling_ds_btcvae"
 
-# Decode Funktion --------------------------------------------------
-sorter = trans.LatentSorter(disvae.get_kl_dict(P_MODEL))
-vae = disvae.load_model(P_MODEL)
-scaler = trans.MinMaxScaler(_min=torch.tensor([1.3]),_max=torch.tensor([4.0]),min_norm=0.3,max_norm=0.6)
-imaging = trans.SumField()
-
-_dec = trans.sequential_function(
-    sorter.inv,
-    vae.decoder,
-    scaler.inv
-)
-
-def decode(latent):
-    with torch.no_grad():
-        return trans.np_sample(_dec)(latent)
+@st.cache_resource
+def load_decode_function():
+    sorter = trans.LatentSorter(disvae.get_kl_dict(P_MODEL))
+    vae = disvae.load_model(P_MODEL)
+    scaler = trans.MinMaxScaler(_min=torch.tensor([1.3]),_max=torch.tensor([4.0]),min_norm=0.3,max_norm=0.6)
+    imaging = trans.SumField()
+
+    _dec = trans.sequential_function(
+        sorter.inv,
+        vae.decoder,
+        scaler.inv,
+        imaging.inv
+    )
+
+    def decode(latent):
+        with torch.no_grad():
+            return trans.np_sample(_dec)(latent)
+    
+    return decode
     
-img2ts = trans.np_sample(imaging.inv)
+decode = load_decode_function()
 
-# GUI -----------------------------------------------------------
 
 latent_vector = np.array([st.slider(f"L{l}",min_value=-3.0,max_value=3.0,value=0.0) for l in range(3)])
 latent_vector = np.concatenate([latent_vector,np.zeros(7)],axis=0)
 
-value = decode(latent_vector)
-
-ts = img2ts(value)
+ts = decode(latent_vector)
 
 df = pd.DataFrame({
     "x":np.arange(len(ts)),
@@ -41,10 +42,4 @@ df = pd.DataFrame({
 }
 )
 
-st.line_chart(df,x="x",y="y")
-st.write(ts)
-# st.write(value)
-# st.image(value, use_column_width="always")
-
-# x = st.slider("Select a value")
-# st.write(x, "squared is", x * x)
+st.line_chart(df,x="x",y="y")