feat: capture progressbar

app.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
+import shlex
 import shutil
 import subprocess
 from pathlib import Path
 from textwrap import dedent
 
-import torch
-import streamlit as st
 import numpy as np
+import streamlit as st
+import torch
 from PIL import Image
 from transformers import CLIPTokenizer
 
@@ -22,6 +23,7 @@ color = col1.color_picker("Pick a color", "#00f900")
 col2.text_input("", color, disabled=True)
 
 emb_name = st.text_input("Embedding name", color.lstrip("#").upper())
+init_token = st.text_input("Initializer token", "init token name")
 rgb = hex_to_rgb(color)
 
 img_array = np.zeros((128, 128, 3), dtype=np.uint8)
@@ -38,23 +40,22 @@ if output_path.exists():
 dataset_path.mkdir()
 img_path = dataset_path / f"{emb_name}.png"
 Image.fromarray(img_array).save(img_path)
-tokenizer = CLIPTokenizer.from_pretrained(
-    "Linaqruf/anything-v3.0", subfolder="tokenizer"
-)
 
 with st.sidebar:
-    init_text = st.text_input("Initializer", "init token name")
+    model_name = st.text_input("Model name", "Linaqruf/anything-v3.0")
     steps = st.slider("Steps", 1, 100, 30, step=1)
     learning_rate = st.text_input("Learning rate", "0.005")
     learning_rate = float(learning_rate)
 
-# case 1: init_text is not a single token
-token = tokenizer.tokenize(init_text)
+tokenizer = CLIPTokenizer.from_pretrained(model_name, subfolder="tokenizer")
+
+# case 1: init_token is not a single token
+token = tokenizer.tokenize(init_token)
 if len(token) > 1:
-    st.warning("init_text must be a single token")
+    st.warning("Initializer token must be a single token")
     st.stop()
 
-# case 2: init_text already exists in the tokenizer
+# case 2: init_token already exists in the tokenizer
 num_added_tokens = tokenizer.add_tokens(emb_name)
 if num_added_tokens == 0:
     st.warning(f"The tokenizer already contains the token {emb_name}")
@@ -62,7 +63,7 @@ if num_added_tokens == 0:
 
 cmd = """
 accelerate launch textual_inversion.py \
-  --pretrained_model_name_or_path="Linaqruf/anything-v3.0" \
+  --pretrained_model_name_or_path={model_name} \
   --train_data_dir="dataset" \
   --learnable_property="style" \
   --placeholder_token="{emb_name}" \
@@ -78,22 +79,39 @@ accelerate launch textual_inversion.py \
 """.strip()
 
 cmd = dedent(cmd).format(
-    emb_name=emb_name, init=init_text, lr=learning_rate, steps=steps
+    model_name=model_name,
+    emb_name=emb_name,
+    init=init_token,
+    lr=learning_rate,
+    steps=steps,
 )
+cmd = shlex.split(cmd)
+
+result_path = output_path / "learned_embeds.bin"
+captured = ""
 
-if st.button("Start"):
+start_button = st.button("Start")
+download_button = st.empty()
+
+if start_button:
     with st.spinner("Training..."):
-        subprocess.run(cmd, shell=True)
+        placeholder = st.empty()
+        p = subprocess.Popen(
+            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf-8"
+        )
+
+        while line := p.stderr.readline():
+            captured += line
+            placeholder.code(captured, language="bash")
 
-result_path = Path("output") / "learned_embeds.bin"
 if not result_path.exists():
     st.stop()
 
-# fix unknown error
+# fix unknown file volume bug
 trained_emb = torch.load(result_path, map_location="cpu")
 for k, v in trained_emb.items():
     trained_emb[k] = torch.from_numpy(v.numpy())
 torch.save(trained_emb, result_path)
 
 file = result_path.read_bytes()
-st.download_button("Download", file, f"{emb_name}.pt")
+download_button.download_button(f"Download {emb_name}.pt", file, f"{emb_name}.pt")