added a dummy to not have to write pythreejs html to disk

app.py

@@ -7,16 +7,25 @@ import requests
 import numpy as np
 import numpy.typing as npt
 from dcgan import DCGAN3D_G
-import os
 import pathlib
 pv.start_xvfb()
 
+
+class DummyWriteable(object):
+    def __init__(self):
+        self.html = None
+
+    def write(self, html):
+        self.html = html
+
+
 STREAMLIT_STATIC_PATH = pathlib.Path(st.__path__[0]) / 'static'
 
 DOWNLOADS_PATH = (STREAMLIT_STATIC_PATH / "downloads")
 if not DOWNLOADS_PATH.is_dir():
     DOWNLOADS_PATH.mkdir()
 
+
 def download_checkpoint(url: str, path: str) -> None:
     resp = requests.get(url)
 
@@ -24,15 +33,22 @@ def download_checkpoint(url: str, path: str) -> None:
         f.write(resp.content)
 
 
-def generate_image(path: str,
+@st.cache(persist=True, allow_output_mutation=True)
+def load_model(path: str,
                    image_size: int = 64,
                    z_dim: int = 512,
                    n_channels: int = 1,
                    n_features: int = 32,
-                   ngpu: int = 1,
-                   latent_size: int = 3) -> npt.ArrayLike:
+                   ngpu: int = 1,) -> torch.nn.Module:
     netG = DCGAN3D_G(image_size, z_dim, n_channels, n_features, ngpu)
     netG.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
+    return netG
+
+
+@st.cache()
+def generate_image(netG: torch.nn.Module,
+                   z_dim: int = 512,
+                   latent_size: int = 3) -> npt.ArrayLike:
     z = torch.randn(1, z_dim, latent_size, latent_size, latent_size)
     with torch.no_grad():
         X = netG(z)
@@ -68,6 +84,7 @@ def create_matplotlib_figure(img: npt.ArrayLike, midpoint: int):
         a.set_axis_off()
     return fig
 
+
 def main():
     st.title("Generating Porous Media with GANs")
 
@@ -106,34 +123,37 @@ def main():
 
     if not (DOWNLOADS_PATH / model_fname).exists():
         download_checkpoint(checkpoint_url, (DOWNLOADS_PATH / model_fname))
+    netG = load_model((DOWNLOADS_PATH / model_fname))
 
     latent_size = st.slider("Latent Space Size z", min_value=1, max_value=5, step=1)
-    img = generate_image((DOWNLOADS_PATH / model_fname), latent_size=latent_size)
+    img = generate_image(netG, latent_size=latent_size)
     slices, mesh, dist = create_uniform_mesh_marching_cubes(img)
 
     pv.set_plot_theme("document")
     pl = pv.Plotter(shape=(1, 1),
                          window_size=(view_width, view_height))
     _ = pl.add_mesh(slices, cmap="gray")
-    pl.export_html((DOWNLOADS_PATH / 'slices.html'))
+
+    slices_html = DummyWriteable()
+    pl.export_html(slices_html)
 
     pl = pv.Plotter(shape=(1, 1),
                          window_size=(view_width, view_height))
     _ = pl.add_mesh(mesh, scalars=dist)
-    pl.export_html((DOWNLOADS_PATH / 'mesh.html'))
+    mesh_html = DummyWriteable()
+    pl.export_html(mesh_html)
 
     st.header("2D Cross-Section of Generated Volume")
     fig = create_matplotlib_figure(img, img.shape[0]//2)
     st.pyplot(fig=fig)
 
-    HtmlFile = open((DOWNLOADS_PATH / 'slices.html'), 'r', encoding='utf-8')
-    source_code = HtmlFile.read()
+    source_code = slices_html.html
     st.header("3D Intersections")
     components.html(source_code, width=view_width, height=view_height)
     st.markdown("_Click and drag to spin, right click to shift._")
 
-    HtmlFile = open((DOWNLOADS_PATH / 'mesh.html'), 'r', encoding='utf-8')
-    source_code = HtmlFile.read()
+
+    source_code = mesh_html.html
     st.header("3D Pore Space Mesh")
     components.html(source_code, width=view_width, height=view_height)
     st.markdown("_Click and drag to spin, right click to shift._")