ui + logs

app.py

@@ -1,9 +1,13 @@
 import streamlit as st
 
-import time
+import time, uuid
 import numpy as np
 from PIL import Image
 
+if 'session_id' not in st.session_state:
+    st.session_state.session_id = str(uuid.uuid4())
+    st.session_state.n_gen = 0
+
 # constants
 HF_REPO_NAME_DIFFUSION = 'nostalgebraist/nostalgebraist-autoresponder-diffusion'
 model_path_diffusion = 'nostalgebraist-autoresponder-diffusion'
@@ -63,14 +67,23 @@ def setup():
     return pipeline
 
 
-def handler(text, ts1, ts2, gs1):
+def log(msg, st_state):
+    session_id = st_state.session_id if 'session_id' in st_state else None
+    n_gen = st.session_state.n_gen if 'n_gen' in st_state else None
+
+    print(f"{session_id} ({n_gen}th gen):\n\t{msg}\n")
+
+def handler(text, ts1, ts2, gs1, st_state):
     pipeline = setup()
 
     data = {'text': text[:380], 'guidance_scale': gs1}
     args = {k: v for k, v in DIFFUSION_DEFAULTS.items()}
     args.update(data)
 
-    print(f"running: {args}")
+    log_data = {'ts1': ts2, 'ts2': ts2}
+    log_data.update(args)
+
+    log(repr(log_data), st_state))
 
     pipeline.base_model.set_timestep_respacing(str(ts1))
     pipeline.super_res_model.set_timestep_respacing(str(ts2))
@@ -121,6 +134,7 @@ low_res = st.empty()
 high_res = st.empty()
 
 if button_go:
+    st.session_state.n_gen = st.session_state.n_gen + 1
     with generating_marker.container():
         st.write('**Generating...**')
         st.write('**Prompt:**')
@@ -131,7 +145,7 @@ if button_go:
 
     t = time.time()
 
-    for s, xs in handler(text, ts1, ts2, gs1):
+    for s, xs in handler(text, ts1, ts2, gs1, st.session_state):
         s = Image.fromarray(s[0])
         xs = Image.fromarray(xs[0])
 
@@ -165,7 +179,9 @@ if button_go:
             st.write(f'{prefix} | {count:02d} / {total} frames | {rate:.2f} seconds/frame')
 
         if button_stop:
+            log('gen stopped', st_state))
             break
 
     with generating_marker.container():
+        log('gen complete', st_state))
         st.write('')