diff --git a/flake_io.py b/flake_io.py index ec04b19..580b4ab 100644 --- a/flake_io.py +++ b/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 `.//` 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 + ``__`` sharing the selected image's ````.""" + 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 + ``/`` and rewrite them to ``.//...`` (#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) diff --git a/tests/test_portable_flakes.py b/tests/test_portable_flakes.py index 8621d14..fed99a6 100644 --- a/tests/test_portable_flakes.py +++ b/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()):