latest output update

app.py

@@ -65,8 +65,8 @@ def generate_button(prefix, file_input, version, **kwargs):
         )
 
     if st.button("Generate image", key=f"{prefix}-btn"):
-        with st.spinner(f"⏳ Generating image... ({n_samples*5} minutes estimated time)"):
-            image = toggle_process( ttf_to_image(file_input, n_samples, ref_char_ids, version) )
+        with st.spinner(f"⏳ Generating image (5 minutes per n_sample estimated time)"):
+            image = toggle_process( ttf_to_image(file_input, OUTPUT_IMG_KEY, n_samples, ref_char_ids, version) )
             set_img(OUTPUT_IMG_KEY, image.copy())
         st.image(image)
 
@@ -110,7 +110,7 @@ def main():
     generate_tab()
 
     with st.sidebar:
-        st.header("Latest Output Image")
+        st.header("Latest Output")
         output_image = get_img(OUTPUT_IMG_KEY)
         if output_image:
             st.image(output_image)

generate.py

@@ -123,9 +123,10 @@ def preprocessing(ttf_file) -> str:
     print("Saved at", output_path)
     return output_path
 
-def inference_model(n_samples, ref_char_ids, version):
+def inference_model(OUTPUT_IMG_KEY, n_samples, ref_char_ids, version):
     opts.n_samples = n_samples
     opts.ref_char_ids = ref_char_ids
+    opts.OUTPUT_IMG_KEY = OUTPUT_IMG_KEY
 
     # Select Model
     if version == "TH2TH":
@@ -137,9 +138,9 @@ def inference_model(n_samples, ref_char_ids, version):
 
     return test_main_model(opts)
 
-def ttf_to_image(ttf_file, n_samples=10, ref_char_ids="1,2,3,4,5,6,7,8", version="TH2TH"):
+def ttf_to_image(ttf_file, OUTPUT_IMG_KEY, n_samples=10, ref_char_ids="1,2,3,4,5,6,7,8", version="TH2TH"):
     preprocessing(ttf_file) # Make Data
-    merge_svg_img = inference_model(n_samples, ref_char_ids, version) # Inference
+    merge_svg_img = inference_model(OUTPUT_IMG_KEY, n_samples, ref_char_ids, version) # Inference
     return merge_svg_img
 
 def main():

test_few_shot.py

@@ -31,6 +31,9 @@ def test_main_model(opts):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print("Inference With Device:", device)
     if opts.streamlit:
+        def set_img(key: str, img: Image.Image):
+            st.session_state[key] = img
+
         st.write("Loading Model Weight...")
         st.write("Inference With Device:", device)
 
@@ -78,7 +81,8 @@ def test_main_model(opts):
                 if opts.streamlit:
                     st.progress((sample_idx+1)/opts.n_samples, f"Generating Font Sample {sample_idx+1} Please wait...")
                     im = Image.open(save_file_merge)
-                    st.image(im, caption='img_sample_merge')
+                    set_img(opts.OUTPUT_IMG_KEY, im)
+                    st.image(im, caption=f"sample {sample_idx+1}")
                 
                 for char_idx in tqdm(range(opts.char_num)):
                     img_gt = (1.0 - img_trg[char_idx,...]).data