DemonGatanjieuandClaude Opus 5.5
bffa239903
Saves move from the page to a store (api/workflow_saves.py), a group for each canvas, found by a key in the workflow's extra or by its file, and go into backups. The bar follows the open tab. A save's runs are named to the server, which finds their pictures in ComfyUI's history, so a reload during a run loses nothing. A save opens above the bar with what it made and what changed; removing one goes to the Recycle Bin, and a group past its most takes no more and drops nothing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
249 lines
13 KiB
Python
249 lines
13 KiB
Python
"""Backups (Settings -> Backup): exporting the user's data and putting it back on another
|
|
computer. Two fake computers in temporary folders; nothing outside them is touched."""
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
import zipfile
|
|
from contextlib import ExitStack
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])]
|
|
|
|
import isolated_layout # noqa: E402,F401 - side files beside the model, whatever this computer chose
|
|
from api import backup # noqa: E402
|
|
|
|
PARTS = ("recipes", "combos", "materials", "parameters", "saves", "comfy_workflows")
|
|
|
|
|
|
class Computer:
|
|
"""A user folder, a models folder and settings of one computer; `hashes` stands for the scan."""
|
|
|
|
def __init__(self, root):
|
|
self.root = root
|
|
self.stores = {part: os.path.join(root, "user", part) for part in PARTS}
|
|
self.models = os.path.join(root, "models", "loras")
|
|
self.config = os.path.join(root, "config.json")
|
|
self.download_settings = os.path.join(root, "download_settings.json")
|
|
self.temp = os.path.join(root, "temp")
|
|
self.trash = os.path.join(root, "trash")
|
|
for path in [*self.stores.values(), self.models, self.temp, self.trash]:
|
|
os.makedirs(path, exist_ok=True)
|
|
self.hashes = {}
|
|
|
|
def write(self, path, data):
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
with open(path, "wb" if isinstance(data, bytes) else "w", **({} if isinstance(data, bytes) else {"encoding": "utf-8"})) as handle:
|
|
handle.write(data)
|
|
return path
|
|
|
|
def model(self, rel, data=b"weights", sha=""):
|
|
path = self.write(os.path.join(self.models, rel), data)
|
|
if sha:
|
|
self.hashes[os.path.realpath(path)] = sha
|
|
return os.path.splitext(path)[0]
|
|
|
|
def trashed(self):
|
|
return sorted(os.listdir(self.trash))
|
|
|
|
def patches(self):
|
|
def to_trash(path):
|
|
shutil.move(path, os.path.join(self.trash, os.path.basename(path)))
|
|
|
|
stack = ExitStack()
|
|
stack.enter_context(mock.patch.object(backup, "part_dirs", return_value=self.stores))
|
|
stack.enter_context(mock.patch.object(backup, "work_dir", return_value=self.temp))
|
|
stack.enter_context(mock.patch.object(backup, "config_path", return_value=self.config))
|
|
stack.enter_context(mock.patch.object(backup, "download_settings_path", return_value=self.download_settings))
|
|
stack.enter_context(mock.patch.object(backup.folder_paths, "get_folder_paths", side_effect=lambda t: [self.models] if t == "loras" else []))
|
|
stack.enter_context(mock.patch.object(backup, "get_metadata", side_effect=lambda p: {"hash": self.hashes.get(os.path.realpath(p), "")}))
|
|
stack.enter_context(mock.patch.object(backup, "move_to_trash", side_effect=to_trash))
|
|
stack.enter_context(mock.patch.object(backup, "add_entry"))
|
|
stack.enter_context(mock.patch.object(backup, "_version", return_value="v-test"))
|
|
return stack
|
|
|
|
|
|
class BackupTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.a = Computer(os.path.join(self.tmp.name, "a"))
|
|
self.b = Computer(os.path.join(self.tmp.name, "b"))
|
|
a = self.a
|
|
a.write(os.path.join(a.stores["recipes"], "recipe_1.json"), '{"name": "portrait"}')
|
|
a.write(os.path.join(a.stores["materials"], ".assets", "m1", "preview.webp"), b"img")
|
|
a.write(os.path.join(a.stores["combos"], "Flux.json"), '{"name": "Flux"}')
|
|
a.write(os.path.join(a.stores["combos"], ".legacy_imported.json"), "{}")
|
|
a.write(os.path.join(a.stores["comfy_workflows"], "mine.json"), "{}")
|
|
a.write(a.config, json.dumps({"CIVITAI_API_KEY": "secret-a", "folder_view_mode": "abstract"}))
|
|
base = a.model("Style/cute.safetensors", sha="ab" * 32)
|
|
a.write(base + ".anomalous.json", json.dumps({"format": 1, "custom_name": "Cute", "custom_notes": "trigger: cute"}))
|
|
a.write(base + ".preview.png", b"my own cover")
|
|
civitai = a.model("plain.safetensors", sha="cd" * 32)
|
|
a.write(civitai + ".preview.png", b"civitai image")
|
|
a.write(civitai + ".civitai_bak.png", b"civitai image")
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def export(self, **options):
|
|
with self.a.patches():
|
|
path, manifest = backup.build_backup(**options)
|
|
return path, manifest
|
|
|
|
def put_back(self, path, **options):
|
|
with self.b.patches():
|
|
summary = backup.inspect_backup(path)
|
|
result = backup.apply_backup(path, **{"parts": PARTS, **options})
|
|
return summary, result
|
|
|
|
def test_export_keeps_stores_and_the_users_own_model_data(self):
|
|
path, manifest = self.export()
|
|
names = zipfile.ZipFile(path).namelist()
|
|
self.assertIn("stores/recipes/recipe_1.json", names)
|
|
self.assertIn("stores/materials/.assets/m1/preview.webp", names)
|
|
self.assertNotIn("stores/combos/.legacy_imported.json", names) # this computer's own marker
|
|
self.assertEqual(manifest["parts"]["comfy_workflows"], 1)
|
|
# Only the model with the user's name, notes and cover; a Civitai cover is no user data.
|
|
self.assertEqual([m["rel"] for m in manifest["models"]], ["Style/cute.safetensors"])
|
|
self.assertEqual(manifest["models"][0]["cover"], ".preview.png")
|
|
settings = json.loads(zipfile.ZipFile(path).read("settings/config.json"))
|
|
self.assertNotIn("CIVITAI_API_KEY", settings) # the key never leaves the computer
|
|
|
|
def test_scan_results_and_civitai_covers_come_back_where_missing(self):
|
|
civitai = os.path.splitext(os.path.join(self.a.models, "plain.safetensors"))[0]
|
|
self.a.write(civitai + ".info", '{"id": 7, "name": "Plain"}')
|
|
path, manifest = self.export(include_civitai_covers=True)
|
|
plain = [m for m in manifest["models"] if m["rel"] == "plain.safetensors"][0]
|
|
self.assertEqual((plain["info"], plain["civitai_cover"], plain["civitai_bak"]), (".info", ".preview.png", ".civitai_bak.png"))
|
|
base = self.b.model("plain.safetensors", sha="cd" * 32)
|
|
_summary, result = self.put_back(path)
|
|
self.assertEqual((result["model_infos"], result["model_civitai_covers"]), (1, 1))
|
|
restored = json.loads(Path(base + ".info").read_text(encoding="utf-8"))
|
|
self.assertEqual((restored["id"], restored["name"]), (7, "Plain"))
|
|
self.assertEqual(Path(base + ".preview.png").read_bytes(), b"civitai image")
|
|
self.assertIn("anomalous_cover_sha256", restored) # still Civitai's, not the user's: noted, not kept twice
|
|
self.assertFalse(os.path.exists(base + ".civitai_bak.png"))
|
|
self.b.write(base + ".info", '{"id": 7, "name": "Newer scan"}')
|
|
_summary, kept = self.put_back(path)
|
|
self.assertEqual(kept["model_infos"], 0)
|
|
self.assertIn("Newer scan", Path(base + ".info").read_text(encoding="utf-8"))
|
|
_path, small = self.export() # Civitai's covers only when asked
|
|
self.assertNotIn("civitai_cover", [m for m in small["models"] if m["rel"] == "plain.safetensors"][0])
|
|
|
|
def test_the_model_list_names_what_a_new_computer_lacks(self):
|
|
path, manifest = self.export()
|
|
self.assertEqual(sorted(m["rel"] for m in manifest["library"]), ["Style/cute.safetensors", "plain.safetensors"])
|
|
self.a.model("copy/cute_again.safetensors", sha="ab" * 32) # a second copy: one download
|
|
path, manifest = self.export()
|
|
self.b.model("somewhere/plain_renamed.safetensors", sha="cd" * 32) # this one is here under another name
|
|
with self.b.patches():
|
|
summary = backup.inspect_backup(path)
|
|
self.assertEqual(summary["library"]["total"], 3)
|
|
self.assertEqual([(m["rel"], m["sha256"]) for m in summary["library"]["missing"]], [("Style/cute.safetensors", "ab" * 32)])
|
|
_path, without = self.export(include_library=False)
|
|
self.assertNotIn("library", without)
|
|
|
|
def test_comfy_workflows_can_be_left_out(self):
|
|
_path, manifest = self.export(include_comfy=False)
|
|
self.assertNotIn("comfy_workflows", manifest["parts"])
|
|
|
|
def test_a_new_computer_gets_everything_and_models_are_found_by_hash(self):
|
|
path, _ = self.export()
|
|
b = self.b
|
|
base = b.model("Moved/elsewhere/cute_renamed.safetensors", sha="ab" * 32) # moved and renamed
|
|
b.write(b.config, json.dumps({"CIVITAI_API_KEY": "secret-b"}))
|
|
summary, result = self.put_back(path, settings=True)
|
|
self.assertEqual(summary["parts"]["recipes"], {"new": 1, "same": 0, "differs": 0})
|
|
self.assertEqual(summary["models"]["matched"], 1)
|
|
self.assertEqual(result["added"], 4)
|
|
self.assertTrue(os.path.isfile(os.path.join(b.stores["materials"], ".assets", "m1", "preview.webp")))
|
|
user = json.loads(Path(base + ".anomalous.json").read_text(encoding="utf-8"))
|
|
self.assertEqual(user["custom_name"], "Cute")
|
|
self.assertEqual(Path(base + ".preview.png").read_bytes(), b"my own cover")
|
|
config = json.loads(Path(b.config).read_text(encoding="utf-8"))
|
|
self.assertEqual(config["CIVITAI_API_KEY"], "secret-b") # this computer's key stays
|
|
self.assertEqual(config["folder_view_mode"], "abstract")
|
|
|
|
def test_differing_files_stay_unless_replace_and_then_go_to_the_recycle_bin(self):
|
|
path, _ = self.export()
|
|
b = self.b
|
|
mine = b.write(os.path.join(b.stores["recipes"], "recipe_1.json"), '{"name": "changed here"}')
|
|
_summary, kept = self.put_back(path)
|
|
self.assertEqual(Path(mine).read_text(encoding="utf-8"), '{"name": "changed here"}')
|
|
self.assertEqual(kept["skipped"], 1)
|
|
_summary, replaced = self.put_back(path, replace=True)
|
|
self.assertEqual(replaced["replaced"], 1)
|
|
self.assertEqual(Path(mine).read_text(encoding="utf-8"), '{"name": "portrait"}')
|
|
self.assertIn("recipe_1.json", b.trashed())
|
|
_summary, again = self.put_back(path, replace=True)
|
|
self.assertEqual((again["added"], again["replaced"]), (0, 0)) # all the same now
|
|
|
|
def test_user_fields_are_filled_where_empty_and_replaced_only_on_request(self):
|
|
path, _ = self.export()
|
|
base = self.b.model("cute.safetensors", sha="ab" * 32)
|
|
self.b.write(base + ".anomalous.json", json.dumps({"format": 1, "custom_name": "Mine"}))
|
|
self.put_back(path)
|
|
user = json.loads(Path(base + ".anomalous.json").read_text(encoding="utf-8"))
|
|
self.assertEqual((user["custom_name"], user["custom_notes"]), ("Mine", "trigger: cute"))
|
|
self.put_back(path, replace=True)
|
|
user = json.loads(Path(base + ".anomalous.json").read_text(encoding="utf-8"))
|
|
self.assertEqual(user["custom_name"], "Cute")
|
|
self.assertIn("cute.anomalous.json", self.b.trashed())
|
|
|
|
def test_the_users_cover_replaces_civitais_but_not_another_user_cover(self):
|
|
path, _ = self.export()
|
|
b = self.b
|
|
base = b.model("cute.safetensors", sha="ab" * 32)
|
|
b.write(base + ".png", b"civitai here")
|
|
b.write(base + ".civitai_bak.png", b"civitai here")
|
|
self.put_back(path)
|
|
self.assertEqual(Path(base + ".preview.png").read_bytes(), b"my own cover")
|
|
self.assertFalse(os.path.exists(base + ".png"))
|
|
self.assertTrue(os.path.exists(base + ".civitai_bak.png")) # Civitai's copy stays
|
|
b.write(base + ".preview.png", b"another of mine")
|
|
_summary, result = self.put_back(path)
|
|
self.assertEqual(result["covers_kept"], 1)
|
|
self.assertEqual(Path(base + ".preview.png").read_bytes(), b"another of mine")
|
|
|
|
def test_without_a_hash_a_model_needs_the_same_path_and_size(self):
|
|
path, _ = self.export()
|
|
with zipfile.ZipFile(path) as archive:
|
|
manifest = json.loads(archive.read("manifest.json"))
|
|
manifest["models"][0]["sha256"] = ""
|
|
rewritten = os.path.join(self.b.temp, "nohash.zip")
|
|
with zipfile.ZipFile(path) as source, zipfile.ZipFile(rewritten, "w") as target:
|
|
for info in source.infolist():
|
|
data = json.dumps(manifest) if info.filename == "manifest.json" else source.read(info)
|
|
target.writestr(info.filename, data)
|
|
self.b.model("Style/cute.safetensors", data=b"other size!")
|
|
summary, result = self.put_back(rewritten)
|
|
self.assertEqual((summary["models"]["matched"], result["models_missing"]), (0, 1))
|
|
self.b.model("Style/cute.safetensors", data=b"weights")
|
|
summary, _ = self.put_back(rewritten)
|
|
self.assertEqual(summary["models"]["matched"], 1)
|
|
|
|
def test_unsafe_or_foreign_files_are_refused(self):
|
|
evil = os.path.join(self.b.temp, "evil.zip")
|
|
with zipfile.ZipFile(evil, "w") as archive:
|
|
archive.writestr("manifest.json", json.dumps({"app": backup.APP, "format": backup.FORMAT}))
|
|
archive.writestr("stores/recipes/../../escape.json", "{}")
|
|
with self.b.patches(), self.assertRaises(backup.BackupError):
|
|
backup.inspect_backup(evil)
|
|
foreign = os.path.join(self.b.temp, "foreign.zip")
|
|
with zipfile.ZipFile(foreign, "w") as archive:
|
|
archive.writestr("readme.txt", "hello")
|
|
with self.b.patches(), self.assertRaises(backup.BackupError):
|
|
backup.inspect_backup(foreign)
|
|
text = os.path.join(self.b.temp, "text.zip")
|
|
Path(text).write_text("not a zip", encoding="utf-8")
|
|
with self.b.patches(), self.assertRaises(backup.BackupError):
|
|
backup.inspect_backup(text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|