real app

app.py

@@ -1,4 +1,82 @@
+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
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+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
+HF_REPO_NAME_DIFFUSION = 'nostalgebraist/nostalgebraist-autoresponder-diffusion'
+model_path_diffusion = 'nostalgebraist-autoresponder-diffusion'
+timestep_respacing_sres1 = '20' # '90,60,60,20,20'
+timestep_respacing_sres2 = '20' # '250'
+
+DIFFUSION_DEFAULTS = dict(
+    batch_size=1,
+    n_samples=1,
+    clf_free_guidance=False,
+    clf_free_guidance_sres=False,
+    guidance_scale=1,
+    guidance_scale_sres=0,
+    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
+    )
+    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
+)
+
+pipeline = improved_diffusion.pipeline.SamplingPipeline(sampling_model_sres1, sampling_model_sres2)
+
+def handler(text, ts1, ts2, gs1):
+    # a = np.random.randint(0, 255, (128, 128, 3)).astype(np.uint8)
+    data = {'text': text[:380], 'guidance_scale': gs1}
+    args = {k: v for k, v in DIFFUSION_DEFAULTS.items()}
+    args.update(data)
+
+    print(f"running: {args}")
+
+    pipeline.base_model.set_timestep_respacing(str(ts1))
+    pipeline.super_res_model.set_timestep_respacing(str(ts2))
+
+    gen = pipeline.sample(**args)
+
+    for sample, pred_xstart in
+
+    return result
+
+text = st.text_area('asdf')
+
+if st.button('rweerew'):
+    handler(text, 20, 20, 0)
+
+# x = st.slider('Select a value')
+# st.write(x, 'squared is', x * x)