feat: main feature

app.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+from textwrap import dedent
+
+import torch
+import streamlit as st
+import numpy as np
+from PIL import Image
+from transformers import CLIPTokenizer
+
+
+def hex_to_rgb(s: str) -> tuple[int, int, int]:
+    value = s.lstrip("#")
+    return (int(value[:2], 16), int(value[2:4], 16), int(value[4:6], 16))
+
+
+col1, col2 = st.columns([15, 85])
+color = col1.color_picker("Pick a color", "#00f900")
+col2.text_input("", color, disabled=True)
+
+emb_name = st.text_input("Embedding name", color.lstrip("#").upper())
+rgb = hex_to_rgb(color)
+
+img_array = np.zeros((128, 128, 3), dtype=np.uint8)
+for i in range(3):
+    img_array[..., i] = rgb[i]
+
+dataset_path = Path("dataset")
+output_path = Path("output")
+if dataset_path.exists():
+    shutil.rmtree(dataset_path)
+if output_path.exists():
+    shutil.rmtree(output_path)
+
+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")
+    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)
+if len(token) > 1:
+    st.warning("init_text must be a single token")
+    st.stop()
+
+# case 2: init_text 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}")
+    st.stop()
+
+cmd = """
+accelerate launch textual_inversion.py \
+  --pretrained_model_name_or_path="Linaqruf/anything-v3.0" \
+  --train_data_dir="dataset" \
+  --learnable_property="style" \
+  --placeholder_token="{emb_name}" \
+  --initializer_token="{init}" \
+  --resolution=128 \
+  --train_batch_size=1 \
+  --repeats=1 \
+  --gradient_accumulation_steps=1 \
+  --max_train_steps={steps} \
+  --learning_rate={lr} \
+  --output_dir="output" \
+  --only_save_embeds
+""".strip()
+
+cmd = dedent(cmd).format(
+    emb_name=emb_name, init=init_text, lr=learning_rate, steps=steps
+)
+
+if st.button("Start"):
+    with st.spinner("Training..."):
+        subprocess.run(cmd, shell=True)
+
+result_path = Path("output") / "learned_embeds.bin"
+if not result_path.exists():
+    st.stop()
+
+# fix unknown error
+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")