cache

app.py

@@ -1,15 +1,8 @@
-import os, subprocess, sys
-os.system("git clone https://github.com/nostalgebraist/improved-diffusion.git && cd improved-diffusion && git fetch origin nbar-space && git checkout nbar-dev && pip install -e .")
-os.system("pip install tokenizers x-transformers==0.22.0 axial-positional-embedding")
-os.system("pip install einops==0.3.2")
-sys.path.append("improved-diffusion")
-
 import streamlit as st
 
 import numpy as np
 from PIL import Image
 
-import improved_diffusion.pipeline
 from transformer_utils.util.tfm_utils import get_local_path_from_huggingface_cdn
 
 # constants
@@ -28,33 +21,46 @@ DIFFUSION_DEFAULTS = dict(
     yield_intermediates=True
 )
 
-if not os.path.exists(model_path_diffusion):
-    model_tar_name = 'model.tar'
-    model_tar_path = get_local_path_from_huggingface_cdn(
-        HF_REPO_NAME_DIFFUSION, model_tar_name
+@st.cache
+def setup():
+    import os, subprocess, sys
+    os.system("git clone https://github.com/nostalgebraist/improved-diffusion.git && cd improved-diffusion && git fetch origin nbar-space && git checkout nbar-dev && pip install -e .")
+    os.system("pip install tokenizers x-transformers==0.22.0 axial-positional-embedding")
+    os.system("pip install einops==0.3.2")
+    sys.path.append("improved-diffusion")
+
+    import improved_diffusion.pipeline
+
+    if not os.path.exists(model_path_diffusion):
+        model_tar_name = 'model.tar'
+        model_tar_path = get_local_path_from_huggingface_cdn(
+            HF_REPO_NAME_DIFFUSION, model_tar_name
+        )
+        subprocess.run(f"tar -xf {model_tar_path} && rm {model_tar_path}", shell=True)
+
+    checkpoint_path_sres1 = os.path.join(model_path_diffusion, "sres1.pt")
+    config_path_sres1 = os.path.join(model_path_diffusion, "config_sres1.json")
+
+    checkpoint_path_sres2 = os.path.join(model_path_diffusion, "sres2.pt")
+    config_path_sres2 = os.path.join(model_path_diffusion, "config_sres2.json")
+
+    # load
+    sampling_model_sres1 = improved_diffusion.pipeline.SamplingModel.from_config(
+        checkpoint_path=checkpoint_path_sres1,
+        config_path=config_path_sres1,
+        timestep_respacing=timestep_respacing_sres1
     )
-    subprocess.run(f"tar -xf {model_tar_path} && rm {model_tar_path}", shell=True)
-
-checkpoint_path_sres1 = os.path.join(model_path_diffusion, "sres1.pt")
-config_path_sres1 = os.path.join(model_path_diffusion, "config_sres1.json")
-
-checkpoint_path_sres2 = os.path.join(model_path_diffusion, "sres2.pt")
-config_path_sres2 = os.path.join(model_path_diffusion, "config_sres2.json")
 
-# load
-sampling_model_sres1 = improved_diffusion.pipeline.SamplingModel.from_config(
-    checkpoint_path=checkpoint_path_sres1,
-    config_path=config_path_sres1,
-    timestep_respacing=timestep_respacing_sres1
-)
+    sampling_model_sres2 = improved_diffusion.pipeline.SamplingModel.from_config(
+        checkpoint_path=checkpoint_path_sres2,
+        config_path=config_path_sres2,
+        timestep_respacing=timestep_respacing_sres2
+    )
 
-sampling_model_sres2 = improved_diffusion.pipeline.SamplingModel.from_config(
-    checkpoint_path=checkpoint_path_sres2,
-    config_path=config_path_sres2,
-    timestep_respacing=timestep_respacing_sres2
-)
+    pipeline = improved_diffusion.pipeline.SamplingPipeline(sampling_model_sres1, sampling_model_sres2)
+    return pipeline
 
-pipeline = improved_diffusion.pipeline.SamplingPipeline(sampling_model_sres1, sampling_model_sres2)
+pipeline = setup()
 
 def handler(text, ts1, ts2, gs1):
     # a = np.random.randint(0, 255, (128, 128, 3)).astype(np.uint8)