feat: move a flake's external data into its folder on save (#392)

flake_io.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import logging
 import os
 import posixpath
 import re
+import shutil
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -590,12 +592,190 @@ def save_flake(
                     os.replace(old_cover, new_cover)
                 except OSError:
                     pass
+            # The owned folder follows the yaml; its `./<old_base>/` references
+            # and owned cover take the new base name.
+            old_base = old_name.rsplit("/", 1)[-1]
+            new_base = os.path.basename(canonical)
+            old_owned = os.path.join(os.path.dirname(old_yaml), old_base)
+            owned = os.path.join(os.path.dirname(target), new_base)
+            if os.path.isdir(old_owned) and not os.path.exists(owned):
+                try:
+                    _move_file(old_owned, owned)
+                    data = _rewrite_prefix(data, f"./{old_base}/", f"./{new_base}/")
+                    for ext in _COVER_EXTENSIONS:
+                        if old_base != new_base and os.path.isfile(os.path.join(owned, old_base + ext)):
+                            os.replace(os.path.join(owned, old_base + ext), os.path.join(owned, new_base + ext))
+                except OSError as exc:
+                    logging.warning("[flakes] could not move %s to %s: %s", old_owned, owned, exc)
+
+    _take_ownership(data, target)
 
     with open(target, "w", encoding="utf-8") as f:
         yaml.safe_dump(data, f, sort_keys=False, allow_unicode=True)
     return canonical
 
 
+_CN_BATCH_RE = re.compile(r"^(.+)_\d+_[^_]+$")
+
+
+def _file_safe(value: Any) -> str:
+    return re.sub(r"[^A-Za-z0-9_\-]+", "_", str(value or "")).strip("_")
+
+
+def _move_file(src: str, dest: str) -> None:
+    """Move *src* (file or folder) to *dest*; copy it, keeping the origin, when
+    *dest* is on another device."""
+    os.makedirs(os.path.dirname(dest), exist_ok=True)
+    if os.stat(src).st_dev == os.stat(os.path.dirname(dest)).st_dev:
+        os.replace(src, dest)
+    elif os.path.isdir(src):
+        shutil.copytree(src, dest)
+    else:
+        shutil.copy2(src, dest)
+
+
+def _rewrite_prefix(value: Any, old: str, new: str) -> Any:
+    if isinstance(value, dict):
+        return {k: _rewrite_prefix(v, old, new) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_rewrite_prefix(v, old, new) for v in value]
+    if isinstance(value, str) and value.startswith(old):
+        return new + value[len(old):]
+    return value
+
+
+def _cn_batch(path: str) -> list[str]:
+    """*path* plus the files beside it from the same ControlNet batch: stems
+    ``<prefix>_<n>_<token>`` sharing the selected image's ``<prefix>``."""
+    match = _CN_BATCH_RE.match(os.path.splitext(os.path.basename(path))[0])
+    if not match:
+        return [path]
+    folder = os.path.dirname(path)
+    batch = []
+    for fn in os.listdir(folder):
+        other = _CN_BATCH_RE.match(os.path.splitext(fn)[0])
+        if other and other.group(1) == match.group(1) and os.path.isfile(os.path.join(folder, fn)):
+            batch.append(os.path.join(folder, fn))
+    return batch
+
+
+def _external_image(value: Any, moved: dict[str, str]) -> str | None:
+    """Absolute path of the image a non-``./`` *value* refers to: input-relative
+    (extension optional, like ``_load_cn_image``), absolute, or under the LoRA
+    tree. Files already in *moved* count as existing."""
+    if not isinstance(value, str) or not value or value.startswith("./"):
+        return None
+    direct = os.path.join(folder_paths.get_input_directory(), value)
+    for path in (direct, *(direct + e for e in (".png", ".jpg", ".jpeg", ".webp")), folder_paths.get_full_path("loras", value)):
+        if not path or os.path.splitext(path)[1].lower() not in _COVER_MIME_MAP:
+            continue
+        path = os.path.abspath(path)
+        if os.path.isfile(path) or os.path.normcase(path) in moved:
+            return path
+    return None
+
+
+def _take_ownership(data: dict[str, Any], target: str) -> None:
+    """Move the external files *data* references (LoRAs + sidecars, ControlNet
+    image batches, variant images, the cover) into the flake's owned folder
+    ``<dirname(target)>/<base>`` and rewrite them to ``./<base>/...`` (#392).
+    Flake links, ``./`` values and missing files are left alone."""
+    flake_dir = os.path.dirname(target)
+    base = os.path.splitext(os.path.basename(target))[0]
+    owned = os.path.join(flake_dir, base)
+    moved: dict[str, str] = {}
+    lora_refs: dict[str, str] = {}
+    lora_stems: set[str] = set()
+
+    def own(src: str, dest: str) -> str | None:
+        """Move *src* to *dest* once; its ``./`` reference, or None on failure."""
+        key = os.path.normcase(src)
+        if key not in moved:
+            try:
+                _move_file(src, dest)
+            except OSError as exc:
+                logging.warning("[flakes] could not move %s into %s: %s", src, owned, exc)
+                return None
+            moved[key] = dest
+        return "./" + os.path.relpath(moved[key], flake_dir).replace(os.sep, "/")
+
+    def own_lora(path: Any, display: Any) -> Any:
+        if not isinstance(path, str) or not path or path.startswith("./"):
+            return path
+        if path in lora_refs:
+            return lora_refs[path]
+        src = folder_paths.get_full_path("loras", _resolve_model_name("loras", path))
+        if not src or not os.path.isfile(src):
+            return path
+        src = os.path.abspath(src)
+        src_dir, src_name = os.path.split(src)
+        src_stem, ext = os.path.splitext(src_name)
+        stem = f"{base}_{_file_safe(display)}" if _file_safe(display) else base
+        new_stem, n = stem, 1
+        while new_stem in lora_stems or os.path.exists(os.path.join(owned, new_stem + ext)):
+            n += 1
+            new_stem = f"{stem}_{n}"
+        lora_stems.add(new_stem)
+        # Sidecars (previews, .civitai.info, .json) follow under the new stem.
+        for fn in os.listdir(src_dir):
+            if fn.startswith(src_stem + ".") and not fn.endswith(ext):
+                own(os.path.join(src_dir, fn), os.path.join(owned, new_stem + fn[len(src_stem):]))
+        ref = own(src, os.path.join(owned, new_stem + ext))
+        if ref:
+            lora_refs[path] = ref
+        return ref or path
+
+    def own_entries(loras: Any, cns: Any) -> None:
+        for lora in loras if isinstance(loras, list) else ():
+            if isinstance(lora, dict):
+                lora["path"] = own_lora(lora.get("path"), lora.get("name"))
+        for cn in cns if isinstance(cns, list) else ():
+            src = _external_image(cn.get("image"), moved) if isinstance(cn, dict) else None
+            if not src:
+                continue
+            if os.path.isfile(src):
+                for f in _cn_batch(src):
+                    if f != src:
+                        own(f, os.path.join(owned, "cnets", os.path.basename(f)))
+            cn["image"] = own(src, os.path.join(owned, "cnets", os.path.basename(src))) or cn["image"]
+
+    # Resolved before anything moves: the cover may be a LoRA preview or a CN
+    # batch image that is about to land in the owned folder.
+    cover_src = _external_image(data.get("cover_image"), moved)
+    if "cover_image" not in data and not any(os.path.isfile(os.path.join(owned, base + e)) for e in _COVER_EXTENSIONS):
+        cover_src = next((p for e in _COVER_EXTENSIONS if os.path.isfile(p := os.path.join(flake_dir, base + e))), None)
+
+    own_entries(data.get("loras"), data.get("controlnets"))
+    if "loras" not in data and data.get("path"):
+        data["path"] = own_lora(data["path"], None)
+    variants = data.get("variants")
+    for group, choices in variants.items() if isinstance(variants, dict) else ():
+        for choice, choice_data in choices.items() if isinstance(choices, dict) else ():
+            if not isinstance(choice_data, dict):
+                continue
+            own_entries(choice_data.get("loras"), choice_data.get("controlnets"))
+            src = _external_image(choice_data.get("image"), moved)
+            if src:
+                dest = os.path.join(owned, f"{base}_{_file_safe(group)}_{_file_safe(choice)}{os.path.splitext(src)[1]}")
+                choice_data["image"] = own(src, dest) or choice_data["image"]
+
+    if cover_src:
+        src = moved.get(os.path.normcase(cover_src), cover_src)
+        dest = os.path.join(owned, base + os.path.splitext(src)[1].lower())
+        try:
+            if not os.path.normcase(src).startswith(os.path.normcase(owned) + os.sep):
+                _move_file(src, dest)
+            elif os.path.normcase(src) != os.path.normcase(dest):
+                shutil.copy2(src, dest)
+            for e in _COVER_EXTENSIONS:
+                for old in (os.path.join(owned, base + e), os.path.join(flake_dir, base + e)):
+                    if os.path.normcase(old) != os.path.normcase(dest) and os.path.isfile(old):
+                        os.remove(old)
+            data.pop("cover_image", None)
+        except OSError as exc:
+            logging.warning("[flakes] could not move cover %s into %s: %s", src, owned, exc)
+
+
 def delete_flake(name: str) -> None:
     path = _resolve_file(name)
     os.remove(path)
@@ -1037,10 +1217,8 @@ def save_cover(name: str, ext: str, data: bytes) -> None:
     # Remove existing covers (other extension or the legacy sibling)
     _delete_cover(name)
 
-    # Place in primary root
-    root = _primary_root()
-    path = os.path.join(root, f"{_owned_cover_stem(name)}{ext_lower}")
-    _ensure_inside(path, root)
+    base = name.rsplit("/", 1)[-1]
+    path = os.path.join(os.path.dirname(_resolve_file(name)), base, f"{base}{ext_lower}")
     os.makedirs(os.path.dirname(path), exist_ok=True)
     with open(path, "wb") as f:
         f.write(data)

tests/test_portable_flakes.py

@@ -1,5 +1,5 @@
-"""Tests for `./` paths relative to the flake yaml and the owned cover
-location (#391). Same loading approach as test_variant_extras.py, with the
+"""Tests for `./` paths relative to the flake yaml, the owned cover
+location (#391) and taking ownership of external data on save (#392). Same loading approach as test_variant_extras.py, with the
 stub `folder_paths` pointing the flakes root at a temp dir."""
 import contextlib
 import importlib.util
@@ -168,6 +168,148 @@ def test_list_dirs_skips_owned_folders():
         assert fio.list_dirs() == ["chars", "chars/loose"]
 
 
+@contextlib.contextmanager
+def _save_env():
+    """Flakes root plus temp `input/` and LoRA dirs for the ownership pass (#392)."""
+    with _flakes_root() as root, tempfile.TemporaryDirectory() as ext:
+        inp, loras = os.path.join(ext, "input"), os.path.join(ext, "loras")
+        fp.get_input_directory = lambda: inp
+        fp.get_full_path = lambda c, n: os.path.join(loras, n) if c == "loras" and os.path.isfile(os.path.join(loras, n)) else None
+        try:
+            yield root, inp, loras, ext
+        finally:
+            fp.get_full_path = lambda category, name: None
+
+
+def _exists(*parts):
+    return os.path.isfile(os.path.join(*parts))
+
+
+def test_cn_batch_detection():
+    with tempfile.TemporaryDirectory() as tmp:
+        for fn in ("001_x_1_hed.png", "001_x_2_depth.png", "001_x_3_cover.png", "001_x_extra.png", "002_x_1_hed.png", "plain.png"):
+            _write(os.path.join(tmp, fn), PNG)
+        batch = sorted(os.path.basename(p) for p in fio._cn_batch(os.path.join(tmp, "001_x_2_depth.png")))
+        assert batch == ["001_x_1_hed.png", "001_x_2_depth.png", "001_x_3_cover.png"], batch
+        assert fio._cn_batch(os.path.join(tmp, "001_x_extra.png")) == [os.path.join(tmp, "001_x_extra.png")]
+        assert fio._cn_batch(os.path.join(tmp, "plain.png")) == [os.path.join(tmp, "plain.png")]
+
+
+def test_save_moves_loras_with_sidecars():
+    with _save_env() as (root, _, loras, _ext):
+        for fn in ("hero.safetensors", "hero.preview.png", "hero.civitai.info", "hero.v2.safetensors", "b.safetensors"):
+            _write(os.path.join(loras, "chars", fn), b"w")
+        data = {
+            "loras": [
+                {"name": "Hero!", "path": "chars/hero.safetensors"},
+                {"name": "Hero", "path": "chars/b.safetensors"},
+                {"path": "./musashi/kept.safetensors"},
+                {"path": "missing.safetensors"},
+            ],
+            "variants": {"outfit": {"armor": {"loras": [{"name": "x", "path": "chars/hero.safetensors"}]}}},
+        }
+        fio.save_flake("chars/musashi", data)
+        raw = fio.read_flake_raw("chars/musashi")
+        assert [lr["path"] for lr in raw["loras"]] == [
+            "./musashi/musashi_Hero.safetensors",
+            "./musashi/musashi_Hero_2.safetensors",
+            "./musashi/kept.safetensors",
+            "missing.safetensors",
+        ], raw["loras"]
+        assert raw["variants"]["outfit"]["armor"]["loras"][0]["path"] == "./musashi/musashi_Hero.safetensors"
+        owned = os.path.join(root, "chars", "musashi")
+        for fn in ("musashi_Hero.safetensors", "musashi_Hero.preview.png", "musashi_Hero.civitai.info", "musashi_Hero_2.safetensors"):
+            assert _exists(owned, fn), fn
+        assert not _exists(loras, "chars", "hero.safetensors")
+        assert _exists(loras, "chars", "hero.v2.safetensors")
+
+
+def test_save_unnamed_legacy_lora_path():
+    with _save_env() as (root, _, loras, _ext):
+        _write(os.path.join(loras, "solo.safetensors"), b"w")
+        fio.save_flake("chars/musashi", {"path": "solo.safetensors", "strength": 0.5})
+        assert fio.read_flake_raw("chars/musashi")["path"] == "./musashi/musashi.safetensors"
+        assert _exists(root, "chars", "musashi", "musashi.safetensors")
+
+
+def test_save_moves_cn_batches_and_variant_images():
+    with _save_env() as (root, inp, _, _ext):
+        for fn in ("001_x_1_hed.png", "001_x_2_depth.png", "other.png", "unrelated.png"):
+            _write(os.path.join(inp, "poses", fn), PNG)
+        data = {
+            "controlnets": [{"type": "openpose", "image": "poses/001_x_1_hed"}],
+            "variants": {"outfit": {
+                "Red Armor": {"image": "poses/other.png", "controlnets": [{"type": "depth", "image": "poses/001_x_2_depth.png"}]},
+                "plain": {"image": "poses/001_x_1_hed.png"},
+            }},
+        }
+        fio.save_flake("chars/musashi", data)
+        raw = fio.read_flake_raw("chars/musashi")
+        outfit = raw["variants"]["outfit"]
+        assert raw["controlnets"][0]["image"] == "./musashi/cnets/001_x_1_hed.png"
+        assert outfit["Red Armor"]["controlnets"][0]["image"] == "./musashi/cnets/001_x_2_depth.png"
+        assert outfit["Red Armor"]["image"] == "./musashi/musashi_outfit_Red_Armor.png"
+        assert outfit["plain"]["image"] == "./musashi/cnets/001_x_1_hed.png"
+        assert _exists(root, "chars", "musashi", "cnets", "001_x_2_depth.png")
+        assert _exists(root, "chars", "musashi", "musashi_outfit_Red_Armor.png")
+        assert _exists(inp, "poses", "unrelated.png")
+        assert not _exists(inp, "poses", "001_x_1_hed.png")
+
+
+def test_save_moves_external_cover():
+    with _save_env() as (root, _, _l, ext):
+        source = os.path.join(ext, "pics", "c.JPG")
+        _write(source, b"jpg")
+        _write(os.path.join(root, "chars", "musashi.png"), PNG)
+        fio.save_flake("chars/musashi", {"cover_image": source})
+        assert "cover_image" not in fio.read_flake_raw("chars/musashi")
+        assert not os.path.exists(source)
+        assert not _exists(root, "chars", "musashi.png")
+        assert fio.read_cover("chars/musashi") == (b"jpg", "image/jpeg")
+
+
+def test_save_copies_cover_already_owned():
+    with _save_env() as (root, inp, _, _ext):
+        _write(os.path.join(inp, "001_x_1_hed.png"), b"hed")
+        _write(os.path.join(inp, "001_x_2_cover.png"), PNG)
+        fio.save_flake("chars/musashi", {"controlnets": [{"type": "openpose", "image": "001_x_1_hed.png"}], "cover_image": "001_x_2_cover.png"})
+        assert _exists(root, "chars", "musashi", "cnets", "001_x_2_cover.png")
+        assert _exists(root, "chars", "musashi", "musashi.png")
+        assert "cover_image" not in fio.read_flake_raw("chars/musashi")
+
+
+def test_save_moves_legacy_sibling_cover():
+    with _save_env() as (root, _, _l, _ext):
+        _write(os.path.join(root, "chars", "musashi.png"), PNG)
+        fio.save_flake("chars/musashi", {})
+        assert not _exists(root, "chars", "musashi.png")
+        assert _exists(root, "chars", "musashi", "musashi.png")
+
+
+def test_rename_moves_owned_folder():
+    with _save_env() as (root, _, _l, _ext):
+        _write(os.path.join(root, "chars", "musashi", "musashi_hero.safetensors"), b"w")
+        _write(os.path.join(root, "chars", "musashi", "musashi.png"), PNG)
+        data = {"loras": [{"name": "hero", "path": "./musashi/musashi_hero.safetensors"}], "flake_links": [{"target": "chars/other"}]}
+        assert fio.save_flake("solo/miyamoto", data, old_name="chars/musashi") == "solo/miyamoto"
+        assert not os.path.exists(os.path.join(root, "chars", "musashi.yaml"))
+        assert not os.path.exists(os.path.join(root, "chars", "musashi"))
+        raw = fio.read_flake_raw("solo/miyamoto")
+        assert raw["loras"][0]["path"] == "./miyamoto/musashi_hero.safetensors"
+        assert raw["flake_links"] == [{"target": "chars/other"}]
+        assert _exists(root, "solo", "miyamoto", "musashi_hero.safetensors")
+        assert fio.read_cover("solo/miyamoto") == (PNG, "image/png")
+
+
+def test_save_cover_next_to_yaml_in_other_root():
+    with _flakes_root() as root, tempfile.TemporaryDirectory() as second:
+        fp.folder_names_and_paths["flakes"] = ([root, second], {".yaml", ".yml"})
+        _write(os.path.join(second, "solo.yaml"), b"{}")
+        fio.save_cover("solo", ".png", PNG)
+        assert _exists(second, "solo", "solo.png")
+        assert not os.path.exists(os.path.join(root, "solo"))
+
+
 if __name__ == "__main__":
     failures = 0
     for name, fn in sorted(globals().items()):