mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-05-06 16:36:45 -03:00
366 lines
11 KiB
Python
366 lines
11 KiB
Python
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
import piexif
|
|
from PIL import Image
|
|
|
|
from py.services.service_registry import ServiceRegistry
|
|
from py.nodes.save_image import SaveImageLM
|
|
|
|
|
|
class _DummyTensor:
|
|
def __init__(self, array):
|
|
self._array = array
|
|
self.shape = array.shape
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def numpy(self):
|
|
return self._array
|
|
|
|
|
|
def _make_image():
|
|
return _DummyTensor(
|
|
np.array(
|
|
[
|
|
[[0.0, 0.1, 0.2], [0.3, 0.4, 0.5]],
|
|
[[0.6, 0.7, 0.8], [0.9, 1.0, 0.0]],
|
|
],
|
|
dtype="float32",
|
|
)
|
|
)
|
|
|
|
|
|
def _configure_save_paths(monkeypatch, tmp_path):
|
|
monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(tmp_path), raising=False)
|
|
monkeypatch.setattr(
|
|
"folder_paths.get_save_image_path",
|
|
lambda *_args, **_kwargs: (str(tmp_path), "sample", 1, "", "sample"),
|
|
raising=False,
|
|
)
|
|
|
|
|
|
def _configure_metadata(monkeypatch, metadata_dict):
|
|
monkeypatch.setattr("py.nodes.save_image.get_metadata", lambda: {"raw": "metadata"})
|
|
monkeypatch.setattr(
|
|
"py.nodes.save_image.MetadataProcessor.to_dict",
|
|
lambda raw_metadata, node_id: metadata_dict,
|
|
)
|
|
|
|
|
|
def test_save_image_defaults_to_writing_png_metadata(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
node = SaveImageLM()
|
|
node.save_images([_make_image()], "ComfyUI", "png", id="node-1")
|
|
|
|
image_path = tmp_path / "sample_00001_.png"
|
|
with Image.open(image_path) as img:
|
|
assert img.info["parameters"] == "prompt text\nSeed: 123"
|
|
|
|
|
|
def test_save_image_skips_png_parameters_when_metadata_disabled_and_keeps_workflow(
|
|
monkeypatch, tmp_path
|
|
):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
node = SaveImageLM()
|
|
workflow = {"nodes": [{"id": 1}]}
|
|
node.save_images(
|
|
[_make_image()],
|
|
"ComfyUI",
|
|
"png",
|
|
id="node-1",
|
|
embed_workflow=True,
|
|
extra_pnginfo={"workflow": workflow},
|
|
save_with_metadata=False,
|
|
)
|
|
|
|
image_path = tmp_path / "sample_00001_.png"
|
|
with Image.open(image_path) as img:
|
|
assert "parameters" not in img.info
|
|
assert img.info["workflow"] == json.dumps(workflow)
|
|
|
|
|
|
def test_save_image_skips_jpeg_metadata_when_disabled(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
node = SaveImageLM()
|
|
node.save_images(
|
|
[_make_image()],
|
|
"ComfyUI",
|
|
"jpeg",
|
|
id="node-1",
|
|
save_with_metadata=False,
|
|
)
|
|
|
|
image_path = tmp_path / "sample_00001_.jpg"
|
|
exif_dict = piexif.load(str(image_path))
|
|
assert piexif.ExifIFD.UserComment not in exif_dict.get("Exif", {})
|
|
|
|
|
|
def test_save_image_skips_webp_metadata_when_disabled(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
node = SaveImageLM()
|
|
node.save_images(
|
|
[_make_image()],
|
|
"ComfyUI",
|
|
"webp",
|
|
id="node-1",
|
|
save_with_metadata=False,
|
|
)
|
|
|
|
image_path = tmp_path / "sample_00001_.webp"
|
|
exif_dict = piexif.load(str(image_path))
|
|
assert piexif.ExifIFD.UserComment not in exif_dict.get("Exif", {})
|
|
|
|
|
|
def test_process_image_returns_passthrough_result_and_ui_images(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
images = [_make_image()]
|
|
node = SaveImageLM()
|
|
|
|
result = node.process_image(images, id="node-1")
|
|
|
|
assert result["result"] == (images,)
|
|
assert result["ui"] == {
|
|
"images": [{"filename": "sample_00001_.png", "subfolder": "", "type": "output"}]
|
|
}
|
|
|
|
|
|
def test_process_image_returns_empty_ui_images_when_save_fails(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
def _raise_save_error(*args, **kwargs):
|
|
raise OSError("disk full")
|
|
|
|
monkeypatch.setattr(Image.Image, "save", _raise_save_error)
|
|
|
|
images = [_make_image()]
|
|
node = SaveImageLM()
|
|
|
|
result = node.process_image(images, id="node-1")
|
|
|
|
assert result["result"] == (images,)
|
|
assert result["ui"] == {"images": []}
|
|
|
|
|
|
def test_save_image_does_not_save_recipe_by_default(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
calls = []
|
|
monkeypatch.setattr(
|
|
SaveImageLM,
|
|
"_save_image_as_recipe",
|
|
lambda self, file_path, metadata_dict: calls.append((file_path, metadata_dict)),
|
|
)
|
|
|
|
node = SaveImageLM()
|
|
node.save_images([_make_image()], "ComfyUI", "png", id="node-1")
|
|
|
|
assert calls == []
|
|
|
|
|
|
def test_save_image_saves_recipe_when_enabled(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
metadata_dict = {"prompt": "prompt text", "seed": 123}
|
|
_configure_metadata(monkeypatch, metadata_dict)
|
|
|
|
calls = []
|
|
monkeypatch.setattr(
|
|
SaveImageLM,
|
|
"_save_image_as_recipe",
|
|
lambda self, file_path, metadata_dict: calls.append((file_path, metadata_dict)),
|
|
)
|
|
|
|
node = SaveImageLM()
|
|
node.save_images(
|
|
[_make_image()],
|
|
"ComfyUI",
|
|
"png",
|
|
id="node-1",
|
|
save_as_recipe=True,
|
|
)
|
|
|
|
assert calls == [(str(tmp_path / "sample_00001_.png"), metadata_dict)]
|
|
|
|
|
|
def test_save_image_saves_recipe_for_each_successful_batch_image(monkeypatch, tmp_path):
|
|
monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(tmp_path), raising=False)
|
|
monkeypatch.setattr(
|
|
"folder_paths.get_save_image_path",
|
|
lambda *_args, **_kwargs: (str(tmp_path), "sample", 7, "", "sample"),
|
|
raising=False,
|
|
)
|
|
metadata_dict = {"prompt": "prompt text", "seed": 123}
|
|
_configure_metadata(monkeypatch, metadata_dict)
|
|
|
|
calls = []
|
|
monkeypatch.setattr(
|
|
SaveImageLM,
|
|
"_save_image_as_recipe",
|
|
lambda self, file_path, metadata_dict: calls.append((file_path, metadata_dict)),
|
|
)
|
|
|
|
node = SaveImageLM()
|
|
node.save_images(
|
|
[_make_image(), _make_image()],
|
|
"ComfyUI",
|
|
"png",
|
|
id="node-1",
|
|
save_as_recipe=True,
|
|
)
|
|
|
|
assert calls == [
|
|
(str(tmp_path / "sample_00007_.png"), metadata_dict),
|
|
(str(tmp_path / "sample_00008_.png"), metadata_dict),
|
|
]
|
|
|
|
|
|
def test_save_image_does_not_save_recipe_when_image_save_fails(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
def _raise_save_error(*args, **kwargs):
|
|
raise OSError("disk full")
|
|
|
|
calls = []
|
|
monkeypatch.setattr(Image.Image, "save", _raise_save_error)
|
|
monkeypatch.setattr(
|
|
SaveImageLM,
|
|
"_save_image_as_recipe",
|
|
lambda self, file_path, metadata_dict: calls.append((file_path, metadata_dict)),
|
|
)
|
|
|
|
node = SaveImageLM()
|
|
node.save_images(
|
|
[_make_image()],
|
|
"ComfyUI",
|
|
"png",
|
|
id="node-1",
|
|
save_as_recipe=True,
|
|
)
|
|
|
|
assert calls == []
|
|
|
|
|
|
def test_process_image_keeps_image_result_when_recipe_save_fails(monkeypatch, tmp_path):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
_configure_metadata(monkeypatch, {"prompt": "prompt text", "seed": 123})
|
|
|
|
def _raise_recipe_error(*args, **kwargs):
|
|
raise RuntimeError("recipe unavailable")
|
|
|
|
monkeypatch.setattr(SaveImageLM, "_save_image_as_recipe", _raise_recipe_error)
|
|
|
|
images = [_make_image()]
|
|
node = SaveImageLM()
|
|
|
|
result = node.process_image(images, id="node-1", save_as_recipe=True)
|
|
|
|
assert result["result"] == (images,)
|
|
assert result["ui"] == {
|
|
"images": [{"filename": "sample_00001_.png", "subfolder": "", "type": "output"}]
|
|
}
|
|
|
|
|
|
def test_save_image_as_recipe_writes_recipe_without_async_scanner_calls(
|
|
monkeypatch, tmp_path
|
|
):
|
|
_configure_save_paths(monkeypatch, tmp_path)
|
|
source_image = tmp_path / "source.png"
|
|
Image.new("RGB", (16, 16), color=(10, 20, 30)).save(source_image)
|
|
recipes_dir = tmp_path / "recipes"
|
|
|
|
class _Cache:
|
|
def __init__(self, raw_data=None):
|
|
self.raw_data = raw_data or []
|
|
self.sorted_by_name = []
|
|
self.sorted_by_date = []
|
|
self.folders = []
|
|
self.folder_tree = {}
|
|
|
|
class _ModelScanner:
|
|
def __init__(self, raw_data):
|
|
self._cache = _Cache(raw_data)
|
|
|
|
class _PersistentCache:
|
|
def __init__(self):
|
|
self.updates = []
|
|
|
|
def update_recipe(self, recipe_data, json_path):
|
|
self.updates.append((recipe_data, json_path))
|
|
|
|
class _RecipeScanner:
|
|
def __init__(self):
|
|
self.recipes_dir = str(recipes_dir)
|
|
self._cache = _Cache([])
|
|
self._json_path_map = {}
|
|
self._persistent_cache = _PersistentCache()
|
|
self._lora_scanner = _ModelScanner(
|
|
[
|
|
{
|
|
"file_name": "foo",
|
|
"sha256": "ABC123",
|
|
"base_model": "SDXL",
|
|
"civitai": {
|
|
"id": 456,
|
|
"name": "Foo v1",
|
|
"model": {"name": "Foo"},
|
|
},
|
|
}
|
|
]
|
|
)
|
|
self._checkpoint_scanner = _ModelScanner([])
|
|
self.fts_updates = []
|
|
|
|
def _update_folder_metadata(self, cache):
|
|
cache.folders = [""]
|
|
cache.folder_tree = {}
|
|
|
|
def _update_fts_index_for_recipe(self, recipe_data, operation):
|
|
self.fts_updates.append((recipe_data["id"], operation))
|
|
|
|
scanner = _RecipeScanner()
|
|
monkeypatch.setitem(ServiceRegistry._services, "recipe_scanner", scanner)
|
|
|
|
node = SaveImageLM()
|
|
node._save_image_as_recipe(
|
|
str(source_image),
|
|
{
|
|
"prompt": "prompt text",
|
|
"seed": 123,
|
|
"checkpoint": "model.safetensors",
|
|
"loras": "<lora:foo:0.7>",
|
|
},
|
|
)
|
|
|
|
recipe_files = list(recipes_dir.glob("*.recipe.json"))
|
|
preview_files = list(recipes_dir.glob("*.webp"))
|
|
|
|
assert len(recipe_files) == 1
|
|
assert len(preview_files) == 1
|
|
assert len(scanner._cache.raw_data) == 1
|
|
assert len(scanner._persistent_cache.updates) == 1
|
|
|
|
recipe = json.loads(recipe_files[0].read_text(encoding="utf-8"))
|
|
assert recipe["file_path"] == os.path.normpath(str(preview_files[0]))
|
|
assert recipe["title"] == "foo-0.70"
|
|
assert recipe["base_model"] == "SDXL"
|
|
assert recipe["loras"][0]["hash"] == "abc123"
|
|
assert recipe["loras"][0]["modelVersionId"] == 456
|
|
assert recipe["gen_params"] == {"prompt": "prompt text", "seed": 123}
|
|
assert scanner._json_path_map[recipe["id"]] == os.path.normpath(str(recipe_files[0]))
|
|
assert scanner.fts_updates == [(recipe["id"], "add")]
|