Add custom tag input and fix model download (zip directory)

- Add trigger word input to Data Source tab for LoRA training
- Fix model download: zip the output directory before serving

app.py

@@ -313,7 +313,7 @@ def _build_review_dataframe():
     return builder.get_samples_dataframe_data()
 
 
-def lora_download_hf(dataset_id, max_files, hf_offset, training_state):
+def lora_download_hf(dataset_id, custom_tag, max_files, hf_offset, training_state):
     """Download HuggingFace dataset batch, restore labels from HF repo, and scan."""
     try:
         if not dataset_id or not dataset_id.strip():
@@ -333,6 +333,11 @@ def lora_download_hf(dataset_id, max_files, hf_offset, training_state):
 
         builder = get_dataset_builder()
 
+        # Set trigger word for LoRA training
+        tag = custom_tag.strip() if custom_tag else ""
+        if tag:
+            builder.set_custom_tag(tag)
+
         # Restore labels/flags from dataset.json pulled from HF repo
         dataset_json_path = str(Path(local_dir) / "dataset.json")
         if Path(dataset_json_path).exists():
@@ -592,10 +597,19 @@ def lora_stop_training():
 
 
 def lora_download_model(model_path):
-    """Return model path for Gradio file download."""
-    if model_path and Path(model_path).exists():
-        return model_path
-    return None
+    """Zip the LoRA model directory and return the zip for Gradio file download."""
+    import shutil
+
+    if not model_path or not Path(model_path).exists():
+        return None
+
+    path = Path(model_path)
+    if path.is_dir():
+        zip_path = path.parent / path.name
+        shutil.make_archive(str(zip_path), "zip", root_dir=str(path.parent), base_dir=path.name)
+        return str(zip_path) + ".zip"
+
+    return model_path
 
 
 # ==================== GRADIO UI ====================
@@ -832,6 +846,10 @@ def create_ui():
                             label="Dataset ID",
                             placeholder="username/dataset-name",
                         )
+                        lora_custom_tag = gr.Textbox(
+                            label="Custom Tag (trigger word for LoRA)",
+                            placeholder="lofi, synthwave, jazz-piano…",
+                        )
                         with gr.Row():
                             lora_hf_max = gr.Slider(
                                 minimum=1, maximum=500, value=50, step=1,
@@ -980,7 +998,7 @@ def create_ui():
                 # Data Source
                 lora_hf_btn.click(
                     fn=lora_download_hf,
-                    inputs=[lora_hf_id, lora_hf_max, lora_hf_offset, training_state],
+                    inputs=[lora_hf_id, lora_custom_tag, lora_hf_max, lora_hf_offset, training_state],
                     outputs=[lora_source_status, training_state, lora_hf_offset, lora_progress],
                 )