add toggle button

app.py

@@ -1,6 +1,7 @@
 from typing import Optional
 import streamlit as st
 from generate import ttf_to_image
+from threading import Thread
 from PIL import Image
 import os
 
@@ -8,6 +9,23 @@ LOADED_TTF_KEY = "loaded_ttf"
 SET_IMG_KEY = "set_img"
 OUTPUT_IMG_KEY = "output_img"
 
+# For multithreading toggle (prevent function from running too many time)
+process_runnning = False
+process_thread = None
+
+def toggle_process(process_running, process_thread, run_process):
+    if process_runnning:
+        # Toggle off
+        process_running = False
+        st.write("Cancled")
+        if process_thread: # Kill Thread
+            process_thread.join()
+    else:
+        # Toggle on
+        process_running = True
+        process_thread = Thread(target=run_process)
+        process_thread.start()
+
 def get_ttf(key: str) -> Optional[any]:
     if key in st.session_state:
         return st.session_state[key]
@@ -45,19 +63,10 @@ def generate_button(prefix, file_input, version, **kwargs):
         value="1,2,3,4,5,6,7,8",
         key=f"{prefix}-ref_char_ids",
         )
-    enable_attention_slicing = st.checkbox(
-        "Enable attention slicing (enables higher resolutions but is slower)",
-        key=f"{prefix}-attention-slicing",
-    )
-    enable_cpu_offload = st.checkbox(
-        "Enable CPU offload (if you run out of memory, e.g. for XL model)",
-        key=f"{prefix}-cpu-offload",
-        value=False,
-    )
 
     if st.button("Generate image", key=f"{prefix}-btn"):
         with st.spinner("⏳ Generating image..."):
-            image = ttf_to_image(file_input, n_samples, ref_char_ids, version)
+            image = toggle_process( ttf_to_image(file_input, n_samples, ref_char_ids, version) )
             set_img(OUTPUT_IMG_KEY, image.copy())
         st.image(image)
 

models/transformers.py

@@ -196,7 +196,7 @@ class Transformer_decoder(nn.Module):
         self.decoder_layers_parallel = clones(DecoderLayer(512, c(attn), c(attn), c(ff), dropout=0.0), 1)
         self.decoder_norm_parallel = nn.LayerNorm(512)
         if opts.ref_nshot == 52:
-            self.cls_embedding = nn.Embedding(92,512)
+            self.cls_embedding = nn.Embedding(96,512)
         else:
             self.cls_embedding = nn.Embedding(52,512)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, 512))