fix: use tempdir

app.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import shlex
-import shutil
 import subprocess
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from textwrap import dedent
 
 import numpy as np
@@ -38,15 +38,11 @@ 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_temp = TemporaryDirectory(prefix="dataset_", dir=".")
+dataset_path = Path(dataset_temp.name)
+output_temp = TemporaryDirectory(prefix="output_", dir=".")
+output_path = Path(output_temp.name)
 
-dataset_path.mkdir()
-output_path.mkdir()
 img_path = dataset_path / f"{emb_name}.png"
 Image.fromarray(img_array).save(img_path)
 
@@ -73,7 +69,7 @@ if num_added_tokens == 0:
 cmd = """
 accelerate launch textual_inversion.py \
   --pretrained_model_name_or_path={model_name} \
-  --train_data_dir="dataset" \
+  --train_data_dir={dataset_path} \
   --learnable_property="style" \
   --placeholder_token="{emb_name}" \
   --initializer_token="{init}" \
@@ -83,16 +79,18 @@ accelerate launch textual_inversion.py \
   --gradient_accumulation_steps=1 \
   --max_train_steps={steps} \
   --learning_rate={lr} \
-  --output_dir="output" \
+  --output_dir={output_path} \
   --only_save_embeds
 """.strip()
 
 cmd = dedent(cmd).format(
     model_name=model_name,
+    dataset_path=dataset_path.as_posix(),
     emb_name=emb_name,
     init=init_token,
-    lr=learning_rate,
     steps=steps,
+    lr=learning_rate,
+    output_path=output_path.as_posix(),
 )
 cmd = shlex.split(cmd)
 
@@ -125,3 +123,6 @@ torch.save(trained_emb, result_path)
 file = result_path.read_bytes()
 download_button.download_button(f"Download {emb_name}.pt", file, f"{emb_name}.pt")
 st.download_button(f"Download {emb_name}.pt ", file, f"{emb_name}.pt")
+
+dataset_temp.cleanup()
+output_temp.cleanup()