flashstudio 0.0.1__py3-none-any.whl

flashstudio/pages/model/utils/config.py

@@ -0,0 +1,27 @@
+"""FlashStudio — Config building/applying helpers."""
+
+import streamlit as st
+from flashstudio.utils.config import apply_training_config
+from flashstudio.constants import (
+    ARCH_FAMILIES, DEFAULT_MODEL_ARCH,
+    TRAIN_EPOCHS, TRAIN_BATCH_SIZE, TRAIN_LR, TRAIN_IMG_SIZE,
+)
+
+
+def _build_cfg():
+    return {
+        "model": {"family": st.session_state.get("arch_family", ARCH_FAMILIES[0]), "variant": st.session_state.get("model_arch", DEFAULT_MODEL_ARCH)},
+        "training": {"epochs": st.session_state.get("epochs", TRAIN_EPOCHS), "batch_size": st.session_state.get("batch_size", TRAIN_BATCH_SIZE),
+                     "lr": st.session_state.get("lr", TRAIN_LR), "img_size": st.session_state.get("img_size", TRAIN_IMG_SIZE)},
+        "augmentation": {"mosaic": st.session_state.get("aug_mosaic", True), "mixup": st.session_state.get("aug_mixup", False)},
+        "advanced": {"amp": st.session_state.get("amp", True), "compile": st.session_state.get("compile_model", False)},
+    }
+
+
+def _apply_cfg(config):
+    if "model" in config:
+        m = config["model"]
+        for k, sk in [("family", "arch_family"), ("variant", "model_arch"), ("pretrained", "pretrain_option")]:
+            if k in m:
+                st.session_state[sk] = m[k]
+    apply_training_config(config)

flashstudio/pages/training/__init__.py

@@ -0,0 +1,3 @@
+"""FlashStudio — Training page package."""
+
+from flashstudio.pages.training.page import render_training_page  # noqa: F401

flashstudio/pages/training/_common.py

@@ -0,0 +1,24 @@
+"""Shared helpers for the training page."""
+
+import os
+import glob as glob_module
+import streamlit as st
+from flashstudio.constants import TRAINING_LOG_GLOB
+
+
+def _get_save_dir():
+    """Single source of truth: Model > Advanced > Save Dir. Used across the whole project."""
+    from flashstudio.utils import get_state
+    val = get_state("save_dir")
+    if not val:
+        from flashstudio.utils import DEFAULTS
+        val = DEFAULTS["save_dir"]
+    return val
+
+
+def _find_log_file(run_dir: str):
+    """Find the best training log file — prefer the largest (stdout-redirected has more data)."""
+    logs = glob_module.glob(os.path.join(run_dir, TRAINING_LOG_GLOB))
+    if not logs:
+        return None
+    return max(logs, key=lambda p: os.path.getsize(p))

flashstudio/pages/training/launch/__init__.py

@@ -0,0 +1,3 @@
+"""Training launch sub-package (tab, preflight, runner, controls, dialogs)."""
+
+from flashstudio.pages.training.launch.tab import _render_start_tab  # noqa: F401

flashstudio/pages/training/launch/controls.py

@@ -0,0 +1,84 @@
+"""Training process controls — stop, pause, resume."""
+
+import os
+import signal
+import streamlit as st
+
+
+def _stop_training():
+    """Stop the training subprocess by killing its process tree."""
+    from flashstudio.utils import flash
+
+    pid = st.session_state.get("training_pid")
+    if pid:
+        try:
+            os.kill(pid, signal.SIGTERM)
+            flash(f"Training stopped (PID {pid})", "success")
+        except ProcessLookupError:
+            flash("Training process already finished", "info")
+        except PermissionError:
+            try:
+                os.kill(pid, signal.SIGKILL)
+                flash(f"Training force-killed (PID {pid})", "warning")
+            except Exception as e:
+                flash(f"Could not kill process {pid}: {e}", "error")
+        except Exception as e:
+            flash(f"Error stopping training: {e}", "error")
+    else:
+        flash("No training process to stop", "warning")
+
+    st.session_state["training_active"] = False
+    st.session_state["training_status"] = "Stopped"
+    st.session_state["training_paused"] = False
+    st.session_state["training_pid"] = None
+    fh = st.session_state.pop("_log_file_handle", None)
+    if fh and not fh.closed:
+        try:
+            fh.close()
+        except Exception:
+            pass
+    st.rerun()
+
+
+def _pause_training():
+    """Pause the training subprocess by sending SIGSTOP."""
+    from flashstudio.utils import flash
+
+    pid = st.session_state.get("training_pid")
+    if pid:
+        try:
+            os.kill(pid, signal.SIGSTOP)
+            st.session_state["training_paused"] = True
+            st.session_state["training_status"] = "Paused"
+            flash(f"Training paused (PID {pid}). GPU memory still held.", "success")
+        except ProcessLookupError:
+            flash("Training process already finished", "info")
+            st.session_state["training_active"] = False
+            st.session_state["training_paused"] = False
+        except Exception as e:
+            flash(f"Could not pause training: {e}", "error")
+    else:
+        flash("No training process found to pause", "warning")
+    st.rerun()
+
+
+def _resume_active_training():
+    """Resume a paused training subprocess by sending SIGCONT."""
+    from flashstudio.utils import flash
+
+    pid = st.session_state.get("training_pid")
+    if pid:
+        try:
+            os.kill(pid, signal.SIGCONT)
+            st.session_state["training_paused"] = False
+            st.session_state["training_status"] = "Running"
+            flash(f"Training resumed (PID {pid})", "success")
+        except ProcessLookupError:
+            flash("Training process already finished", "info")
+            st.session_state["training_active"] = False
+            st.session_state["training_paused"] = False
+        except Exception as e:
+            flash(f"Could not resume training: {e}", "error")
+    else:
+        flash("No training process found to resume", "warning")
+    st.rerun()

flashstudio/pages/training/launch/dialogs.py

@@ -0,0 +1,247 @@
+"""Training launch dialogs — clean workspace, resume, folder management, config."""
+
+import os
+import yaml
+import streamlit as st
+
+from flashstudio.utils.filesystem import dir_size_str
+from flashstudio.utils.config import build_training_config, apply_training_config
+from flashstudio.constants import (
+    CKPT_BEST, CKPT_BEST_INFERENCE, CKPT_BEST_FP16,
+    CKPT_FINAL_INFERENCE, CKPT_FINAL_FP16, CKPT_LAST, CKPT_LAST_INFERENCE,
+    TRAINING_LOG_CSV,
+)
+from flashstudio.pages.training._common import _get_save_dir
+
+
+_build_config_dict = build_training_config
+_apply_config_dict = apply_training_config
+_get_folder_size_str = dir_size_str
+_get_dir_size = dir_size_str
+
+
+def _render_folder_dialogs(runs_root: str, existing_folders: list):
+    """Render rename/delete dialogs for folder management."""
+    import shutil
+    from flashstudio.utils import flash
+
+    if st.session_state.get("wf_rename_target"):
+        target = st.session_state["wf_rename_target"]
+        with st.container(border=True):
+            rename_val = st.text_input("New name", value=target, key="rename_wf_input")
+            rc1, rc2 = st.columns(2)
+            with rc1:
+                if st.button("Rename", type="primary", key="do_rename_wf"):
+                    if rename_val and rename_val != target:
+                        old_path = os.path.join(runs_root, target)
+                        new_path = os.path.join(runs_root, rename_val)
+                        if os.path.exists(new_path):
+                            flash(f"Rename failed: `{rename_val}` already exists", "error")
+                        else:
+                            os.rename(old_path, new_path)
+                            if st.session_state.get("active_workflow_folder") == target:
+                                st.session_state["active_workflow_folder"] = rename_val
+                                st.session_state["run_name"] = rename_val
+                            flash(f"Renamed `{target}` → `{rename_val}`", "success")
+                            st.session_state.pop("wf_rename_target", None)
+                            st.rerun()
+            with rc2:
+                if st.button("Cancel", key="cancel_rename_wf"):
+                    st.session_state.pop("wf_rename_target", None)
+                    st.rerun()
+
+    if st.session_state.get("wf_delete_target"):
+        target = st.session_state["wf_delete_target"]
+        with st.container(border=True):
+            st.error(f"Delete **{target}** and all contents?")
+            dc1, dc2 = st.columns(2)
+            with dc1:
+                if st.button("Delete", type="primary", key="confirm_delete_wf"):
+                    shutil.rmtree(os.path.join(runs_root, target), ignore_errors=True)
+                    if st.session_state.get("active_workflow_folder") == target:
+                        st.session_state.pop("active_workflow_folder", None)
+                    st.session_state.pop("wf_delete_target", None)
+                    flash(f"Deleted `{target}`", "success")
+                    st.rerun()
+            with dc2:
+                if st.button("Cancel", key="cancel_delete_wf"):
+                    st.session_state.pop("wf_delete_target", None)
+                    st.rerun()
+
+    if st.session_state.get("show_delete_all_wf_folders"):
+        with st.container(border=True):
+            st.error(f"Delete ALL {len(existing_folders)} folders? This is permanent!")
+            dc1, dc2 = st.columns(2)
+            with dc1:
+                if st.button("DELETE ALL", type="primary", key="confirm_delete_all_wf"):
+                    for f in existing_folders:
+                        shutil.rmtree(os.path.join(runs_root, f), ignore_errors=True)
+                    st.session_state.pop("active_workflow_folder", None)
+                    st.session_state.pop("show_delete_all_wf_folders", None)
+                    st.rerun()
+            with dc2:
+                if st.button("Cancel", key="cancel_delete_all_wf"):
+                    st.session_state.pop("show_delete_all_wf_folders", None)
+                    st.rerun()
+
+
+def _render_config_file_dialog(run_path: str):
+    """Compact save/load config dialog."""
+    with st.container(border=True):
+        tab_save, tab_load = st.tabs(["Save", "Load"])
+        with tab_save:
+            config = _build_config_dict()
+            yaml_str = yaml.dump(config, default_flow_style=False, sort_keys=False)
+            with st.expander("YAML Preview"):
+                st.code(yaml_str, language="yaml")
+            sc1, sc2 = st.columns(2)
+            with sc1:
+                st.download_button("Download", yaml_str, file_name="config.yaml", mime="text/yaml",
+                                   use_container_width=True, type="primary")
+            with sc2:
+                if st.button("Save to Run", use_container_width=True, key="save_config_to_run"):
+                    os.makedirs(run_path, exist_ok=True)
+                    with open(os.path.join(run_path, "config.yaml"), "w") as f:
+                        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
+                    st.success("Saved")
+        with tab_load:
+            load_method = st.radio("From", ["Upload", "Path", "Runs"], key="config_load_method", horizontal=True)
+            config_to_load = None
+            if load_method == "Upload":
+                uploaded = st.file_uploader("YAML", type=["yaml", "yml"], key="config_file_upload")
+                if uploaded:
+                    try:
+                        config_to_load = yaml.safe_load(uploaded.read().decode("utf-8"))
+                    except Exception as e:
+                        st.error(str(e)[:40])
+            elif load_method == "Path":
+                cp = st.text_input("Path", placeholder="/path/config.yaml", key="config_path_input")
+                if cp and os.path.isfile(cp):
+                    try:
+                        with open(cp) as f:
+                            config_to_load = yaml.safe_load(f.read())
+                    except Exception as e:
+                        st.error(str(e)[:40])
+            else:
+                sd = _get_save_dir()
+                parent = sd if os.path.isdir(sd) else ""
+                found = [(d, os.path.join(parent, d, "config.yaml"))
+                         for d in (os.listdir(parent) if parent and os.path.isdir(parent) else [])
+                         if os.path.isfile(os.path.join(parent, d, "config.yaml"))]
+                if found:
+                    sel = st.selectbox("Run", range(len(found)), format_func=lambda i: found[i][0], key="config_from_run_select")
+                    with open(found[sel][1]) as f:
+                        config_to_load = yaml.safe_load(f.read())
+            if config_to_load:
+                if st.button("Apply", type="primary", key="apply_loaded_config", use_container_width=True):
+                    _apply_config_dict(config_to_load)
+                    st.rerun()
+
+
+def _render_clean_workspace_dialog():
+    """Compact workspace cleanup."""
+    with st.container(border=True):
+        parent_dir = _get_save_dir()
+
+        if not os.path.isdir(parent_dir):
+            st.info("Nothing to clean.")
+            if st.button("Close", key="close_clean"):
+                st.session_state["show_clean_dialog"] = False
+                st.rerun()
+            return
+
+        runs = [{"name": d, "path": os.path.join(parent_dir, d), "size": _get_dir_size(os.path.join(parent_dir, d)),
+                 "has_best": (os.path.isfile(os.path.join(parent_dir, d, CKPT_BEST))
+                              or os.path.isfile(os.path.join(parent_dir, d, CKPT_BEST_INFERENCE)))}
+                for d in sorted(os.listdir(parent_dir)) if os.path.isdir(os.path.join(parent_dir, d))]
+        if not runs:
+            st.info("No runs."); return
+
+        clean_mode = st.radio("Mode", ["Selected", "Keep best only", "Delete incomplete", "Full reset"],
+                              key="clean_mode", horizontal=True)
+
+        if clean_mode == "Selected":
+            selected = st.multiselect("Runs", [r["name"] for r in runs], key="runs_to_delete")
+            if selected and st.button("Delete", type="primary", key="do_delete_selected"):
+                import shutil
+                for n in selected:
+                    shutil.rmtree(os.path.join(parent_dir, n), ignore_errors=True)
+                st.session_state["show_clean_dialog"] = False; st.rerun()
+        elif clean_mode == "Keep best only":
+            if st.button("Clean", type="primary", key="do_keep_best"):
+                for r in runs:
+                    _cleanup_run_keep_best(r["path"])
+                st.session_state["show_clean_dialog"] = False; st.rerun()
+        elif clean_mode == "Delete incomplete":
+            incomplete = [r for r in runs if not r["has_best"]]
+            st.caption(f"{len(incomplete)} incomplete")
+            if incomplete and st.button("Delete", type="primary", key="do_delete_incomplete"):
+                import shutil
+                for r in incomplete:
+                    shutil.rmtree(r["path"])
+                st.session_state["show_clean_dialog"] = False; st.rerun()
+        elif clean_mode == "Full reset":
+            if st.button("DELETE ALL", type="primary", key="do_delete_all"):
+                import shutil
+                shutil.rmtree(parent_dir); os.makedirs(parent_dir, exist_ok=True)
+                st.session_state["show_clean_dialog"] = False; st.rerun()
+
+        if st.button("Cancel", key="cancel_clean"):
+            st.session_state["show_clean_dialog"] = False; st.rerun()
+
+
+def _render_resume_dialog():
+    """Compact resume dialog."""
+    with st.container(border=True):
+        parent_dir = _get_save_dir()
+        resumable = []
+        if os.path.isdir(parent_dir):
+            for d in sorted(os.listdir(parent_dir), reverse=True):
+                full = os.path.join(parent_dir, d)
+                ckpt = os.path.join(full, CKPT_LAST)
+                if not os.path.isfile(ckpt):
+                    ckpt = os.path.join(full, CKPT_LAST_INFERENCE)
+                if os.path.isdir(full) and os.path.isfile(ckpt):
+                    resumable.append({"name": d, "path": full, "ckpt": ckpt})
+        if not resumable:
+            st.info("No resumable runs.")
+            if st.button("Close", key="close_resume"):
+                st.session_state["show_resume_dialog"] = False; st.rerun()
+            return
+        rc1, rc2, rc3 = st.columns([4, 1, 1])
+        with rc1:
+            selected = st.selectbox("Run", [r["name"] for r in resumable], key="resume_run_select", label_visibility="collapsed")
+        with rc2:
+            if st.button("Resume", type="primary", key="do_resume"):
+                from flashstudio.pages.training.launch.runner import _start_training
+                for r in resumable:
+                    if r["name"] == selected:
+                        st.session_state.update({"resume_training": True, "resume_path": r["ckpt"],
+                                                 "active_run_path": r["path"], "training_active": True,
+                                                 "training_status": "Resuming", "show_resume_dialog": False})
+                        _start_training(); break
+        with rc3:
+            if st.button("Cancel", key="cancel_resume"):
+                st.session_state["show_resume_dialog"] = False; st.rerun()
+
+
+def _cleanup_run_keep_best(run_dir: str) -> int:
+    """Remove non-essential files from a run, keeping final/best models and results."""
+    from flashstudio.constants import ONNX_MODEL_FILE, ONNX_DATA_FILE, VIS_DIR_NAMES, GT_VERIFICATION_DIR
+    keep = {
+        CKPT_BEST, TRAINING_LOG_CSV,
+        CKPT_BEST_INFERENCE, CKPT_BEST_FP16,
+        CKPT_FINAL_INFERENCE, CKPT_FINAL_FP16,
+        ONNX_MODEL_FILE, ONNX_DATA_FILE,
+    }
+    removed = 0
+    for f in os.listdir(run_dir):
+        fpath = os.path.join(run_dir, f)
+        if os.path.isfile(fpath) and f not in keep and not f.endswith(".log"):
+            os.remove(fpath)
+            removed += 1
+        elif os.path.isdir(fpath) and f in (*VIS_DIR_NAMES, GT_VERIFICATION_DIR):
+            import shutil
+            shutil.rmtree(fpath)
+            removed += 1
+    return removed

flashstudio/pages/training/launch/preflight.py

@@ -0,0 +1,77 @@
+"""Pre-flight validation checks before training."""
+
+import os
+import streamlit as st
+
+from flashstudio.utils.device import has_cuda
+from flashstudio.pages.training._common import _get_save_dir
+
+
+def _run_preflight_checks() -> list:
+    """Run pre-flight validation checks before training. Returns list of (label, ok, msg)."""
+    checks = []
+
+    train_path = st.session_state.get("train_img_path", "")
+    if train_path and os.path.isdir(train_path):
+        checks.append(("Dataset", True, ""))
+    else:
+        checks.append(("Dataset", False, "No train data path"))
+
+    if train_path and os.path.isdir(train_path):
+        ann_found = False
+        # Check in train_img_path itself
+        ann_file = os.path.join(train_path, "_annotations.coco.json")
+        json_files = [f for f in os.listdir(train_path) if f.endswith(".json")]
+        if os.path.isfile(ann_file) or json_files:
+            ann_found = True
+        # Check parent directory (common layout: train/images/ + train/_annotations.coco.json)
+        if not ann_found:
+            parent = os.path.dirname(train_path)
+            parent_ann = os.path.join(parent, "_annotations.coco.json")
+            parent_jsons = [f for f in os.listdir(parent) if f.endswith(".json")] if os.path.isdir(parent) else []
+            if os.path.isfile(parent_ann) or parent_jsons:
+                ann_found = True
+        # Check session state annotation path
+        if not ann_found and st.session_state.get("train_ann_path"):
+            if os.path.isfile(st.session_state["train_ann_path"]):
+                ann_found = True
+        checks.append(("Annotations", ann_found, "" if ann_found else "No COCO JSON found"))
+    else:
+        checks.append(("Annotations", False, "No data path"))
+
+    nc = st.session_state.get("num_classes", 0)
+    cls_names = st.session_state.get("class_names", "")
+    if isinstance(cls_names, list):
+        cls_names = "\n".join(cls_names)
+    if nc and nc > 0:
+        checks.append(("Classes", True, f"{nc} classes"))
+    elif isinstance(cls_names, str) and cls_names.strip():
+        n = len([c for c in cls_names.strip().split("\n") if c.strip()])
+        checks.append(("Classes", True, f"{n} classes"))
+    else:
+        checks.append(("Classes", False, "No classes — go to Data → Upload"))
+
+    from flashstudio.utils import get_state
+    arch = get_state("model_arch") or ""
+    if arch:
+        checks.append(("Model", True, ""))
+    else:
+        checks.append(("Model", False, "No model selected"))
+
+    if has_cuda():
+        checks.append(("GPU", True, "CUDA"))
+    else:
+        checks.append(("GPU", True, "CPU (slow)"))
+
+    save_dir = _get_save_dir()
+    try:
+        stat = os.statvfs(os.path.dirname(save_dir) if not os.path.isdir(save_dir) else save_dir)
+        free_gb = (stat.f_bavail * stat.f_frsize) / (1024**3)
+        if free_gb > 1:
+            checks.append(("Disk", True, ""))
+        else:
+            checks.append(("Disk", False, f"Only {free_gb:.1f} GB free"))
+    except (OSError, AttributeError):
+        checks.append(("Disk", True, ""))
+
+    return checks