test: cover autotag and lora integration
This commit is contained in:
Vito Sansevero
+6
-8
@@ -8,6 +8,7 @@ Refactored from standalone_tagger.py for integration with PromptManager API.
|
||||
import errno
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
from PIL import Image
|
||||
@@ -705,14 +706,11 @@ class AutoTagService:
|
||||
self._current_model_type = None
|
||||
gc.collect()
|
||||
|
||||
# Try to clear CUDA cache if available
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except ImportError:
|
||||
pass
|
||||
# Clear the CUDA cache only if torch is already loaded; importing
|
||||
# it here would pull in a heavy dependency the backend never used.
|
||||
torch = sys.modules.get("torch")
|
||||
if torch is not None and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
self.logger.info("Model unloaded, memory freed")
|
||||
|
||||
|
||||
@@ -532,5 +532,298 @@ class TestAutotagBatch(AutotagAPITestCase):
|
||||
self.assertEqual(self.service.unload_calls, 1)
|
||||
|
||||
|
||||
class TestAutotagModelsAndDownloadErrors(AutotagAPITestCase):
|
||||
|
||||
async def test_models_error_is_500(self):
|
||||
def boom():
|
||||
raise RuntimeError("/srv/models unreadable")
|
||||
|
||||
self.service.get_models_status = boom
|
||||
resp = await self.client.request("GET", "/prompt_manager/autotag/models")
|
||||
self.assertEqual(resp.status, 500)
|
||||
self.assertFalse((await resp.json())["success"])
|
||||
|
||||
async def test_download_invalid_model_type_is_error_event(self):
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/download/nope"
|
||||
)
|
||||
events = self._sse_events(await resp.text())
|
||||
self.assertEqual(events[-1]["type"], "error")
|
||||
self.assertEqual(self.service.download_calls, [])
|
||||
|
||||
async def test_download_failure_is_error_event(self):
|
||||
self.service.download_result = False
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/download/gguf"
|
||||
)
|
||||
events = self._sse_events(await resp.text())
|
||||
self.assertEqual(events[-1], {"type": "error", "message": "Download failed"})
|
||||
|
||||
async def test_download_exception_is_generic_error_event(self):
|
||||
def boom(model_type, progress_callback=None):
|
||||
raise OSError("/srv/models is full")
|
||||
|
||||
self.service.download_model = boom
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/download/gguf"
|
||||
)
|
||||
events = self._sse_events(await resp.text())
|
||||
self.assertEqual(events[-1]["type"], "error")
|
||||
self.assertNotIn("/srv/models", events[-1]["message"])
|
||||
|
||||
|
||||
class TestAutotagSingleEdgeCases(AutotagAPITestCase):
|
||||
|
||||
async def test_missing_file_error_hides_the_server_path(self):
|
||||
import errno
|
||||
|
||||
image = self._make_image(self.output_dir)
|
||||
self.service.generate_error = FileNotFoundError(
|
||||
errno.ENOENT, "Image not found", str(image)
|
||||
)
|
||||
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single", {"path": str(image)}
|
||||
)
|
||||
|
||||
self.assertEqual(resp.status, 500)
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["success"])
|
||||
self.assertNotIn(str(self.output_dir), data["error"])
|
||||
self.assertIn("img.png", data["error"])
|
||||
|
||||
async def test_generic_error_is_500(self):
|
||||
image = self._make_image(self.output_dir)
|
||||
self.service.generate_error = RuntimeError("model crashed")
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single", {"path": str(image)}
|
||||
)
|
||||
self.assertEqual(resp.status, 500)
|
||||
self.assertEqual((await resp.json())["error"], "model crashed")
|
||||
|
||||
async def test_oserror_without_filename_uses_strerror(self):
|
||||
from py.api.autotag_routes import _public_error
|
||||
|
||||
self.assertEqual(_public_error(OSError(5, "I/O error")), "I/O error")
|
||||
self.assertEqual(_public_error(ValueError("plain")), "plain")
|
||||
|
||||
async def test_path_with_nul_byte_is_forbidden(self):
|
||||
from py.api.autotag_routes import path_is_within
|
||||
|
||||
self.assertFalse(path_is_within("bad\x00name.png", self.output_dir))
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single", {"path": str(self.output_dir / "a\x00b")}
|
||||
)
|
||||
self.assertEqual(resp.status, 403)
|
||||
|
||||
async def test_boolean_image_id_is_400(self):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single", {"image_id": True}
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_numeric_string_image_id_is_accepted(self):
|
||||
prompt_id = self._save_prompt()
|
||||
image_id = self._link_image(prompt_id, self._make_image(self.output_dir))
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single", {"image_id": str(image_id)}
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
|
||||
async def test_non_object_body_is_400(self):
|
||||
resp = await self._post_json("/prompt_manager/autotag/single", [1, 2])
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_invalid_json_is_400(self):
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/autotag/single",
|
||||
data="{nope",
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_non_numeric_threshold_is_400(self):
|
||||
image = self._make_image(self.output_dir)
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single",
|
||||
{"path": str(image), "general_threshold": "high"},
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_prompt_lookup_failure_still_returns_tags(self):
|
||||
image = self._make_image(self.output_dir)
|
||||
with patch.object(
|
||||
self.api.db, "get_prompt_id_for_image", side_effect=RuntimeError("db")
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single", {"path": str(image)}
|
||||
)
|
||||
data = await resp.json()
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertIsNone(data["prompt_id"])
|
||||
|
||||
async def test_prompt_is_found_via_the_real_path(self):
|
||||
prompt_id = self._save_prompt()
|
||||
image = self._make_image(self.output_dir)
|
||||
self._link_image(prompt_id, image)
|
||||
dotted = os.path.join(str(self.output_dir), ".", "img.png")
|
||||
|
||||
resp = await self._post_json("/prompt_manager/autotag/single", {"path": dotted})
|
||||
|
||||
self.assertEqual((await resp.json())["prompt_id"], prompt_id)
|
||||
|
||||
async def test_custom_prompt_and_model_switch(self):
|
||||
image = self._make_image(self.output_dir)
|
||||
self.service._loaded_type = "gguf"
|
||||
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single",
|
||||
{"path": str(image), "model_type": "wd14-vit", "prompt": "terse"},
|
||||
)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(self.service.load_calls, [("wd14-vit", True)])
|
||||
self.assertEqual(self.service.custom_prompt, "terse")
|
||||
|
||||
async def test_image_id_without_stored_path_is_404(self):
|
||||
prompt_id = self._save_prompt()
|
||||
image_id = self._link_image(prompt_id, self._make_image(self.output_dir))
|
||||
with patch.object(
|
||||
self.api.db,
|
||||
"get_image_by_id",
|
||||
return_value={"id": image_id, "image_path": ""},
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single", {"image_id": image_id}
|
||||
)
|
||||
self.assertEqual(resp.status, 404)
|
||||
|
||||
|
||||
class TestApplyAutotag(AutotagAPITestCase):
|
||||
|
||||
async def test_missing_prompt_id_is_400(self):
|
||||
resp = await self._post_json("/prompt_manager/autotag/apply", {"tags": ["a"]})
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_no_tags_is_a_noop(self):
|
||||
resp = await self._post_json("/prompt_manager/autotag/apply", {"prompt_id": 1})
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["message"], "No tags to apply")
|
||||
|
||||
async def test_unknown_prompt_is_404(self):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply", {"prompt_id": 999, "tags": ["a"]}
|
||||
)
|
||||
self.assertEqual(resp.status, 404)
|
||||
|
||||
async def test_tags_are_merged_without_duplicates(self):
|
||||
prompt_id = self._save_prompt(tags=["cat"])
|
||||
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply",
|
||||
{"prompt_id": prompt_id, "tags": ["cat", "dog"]},
|
||||
)
|
||||
|
||||
data = await resp.json()
|
||||
self.assertEqual(data["added_tags"], ["dog"])
|
||||
self.assertEqual(data["total_tags"], 2)
|
||||
self.assertEqual(
|
||||
self.api.db.get_prompt_by_id(prompt_id)["tags"], ["cat", "dog"]
|
||||
)
|
||||
|
||||
async def test_string_tags_from_db_are_split(self):
|
||||
prompt_id = self._save_prompt()
|
||||
with patch.object(
|
||||
self.api.db,
|
||||
"get_prompt_by_id",
|
||||
return_value={"id": prompt_id, "tags": "cat, dog"},
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply",
|
||||
{"prompt_id": prompt_id, "tags": ["dog", "owl"]},
|
||||
)
|
||||
data = await resp.json()
|
||||
self.assertEqual(data["added_tags"], ["owl"])
|
||||
self.assertEqual(data["total_tags"], 3)
|
||||
|
||||
async def test_db_error_is_500(self):
|
||||
with patch.object(
|
||||
self.api.db, "get_prompt_by_id", side_effect=RuntimeError("db")
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply", {"prompt_id": 1, "tags": ["a"]}
|
||||
)
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
|
||||
class TestUnloadAutotag(AutotagAPITestCase):
|
||||
|
||||
async def test_nothing_loaded(self):
|
||||
resp = await self.client.request("POST", "/prompt_manager/autotag/unload")
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["message"], "No model was loaded")
|
||||
|
||||
async def test_unloads_loaded_model(self):
|
||||
self.service._loaded_type = "gguf"
|
||||
resp = await self.client.request("POST", "/prompt_manager/autotag/unload")
|
||||
data = await resp.json()
|
||||
self.assertEqual(data["message"], "GGUF model unloaded successfully")
|
||||
self.assertFalse(data["model_loaded"])
|
||||
self.assertEqual(self.service.unload_calls, 1)
|
||||
|
||||
async def test_error_is_500(self):
|
||||
def boom():
|
||||
raise RuntimeError("x")
|
||||
|
||||
self.service.is_model_loaded = boom
|
||||
resp = await self.client.request("POST", "/prompt_manager/autotag/unload")
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
|
||||
class TestScanOutputDir(AutotagAPITestCase):
|
||||
|
||||
async def test_lists_images_with_thumbnails(self):
|
||||
sub = self.output_dir / "sub"
|
||||
sub.mkdir()
|
||||
self._make_image(self.output_dir, "a.png")
|
||||
self._make_image(sub, "b.JPG")
|
||||
thumbs = self.output_dir / "thumbnails" / "sub"
|
||||
thumbs.mkdir(parents=True)
|
||||
(thumbs / "b_thumb.JPG").write_bytes(b"t")
|
||||
self._make_image(self.output_dir / "thumbnails", "ignored.png")
|
||||
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
|
||||
|
||||
resp = await self.client.request("GET", "/prompt_manager/scan_output_dir")
|
||||
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual([i["filename"] for i in data["images"]], ["a.png", "b.JPG"])
|
||||
self.assertEqual(data["count"], 2)
|
||||
self.assertIsNone(data["images"][0]["thumbnail_url"])
|
||||
self.assertEqual(
|
||||
data["images"][1]["thumbnail_url"],
|
||||
"/prompt_manager/images/serve/thumbnails/sub/b_thumb.JPG",
|
||||
)
|
||||
self.assertEqual(
|
||||
data["images"][1]["url"], "/prompt_manager/images/serve/sub/b.JPG"
|
||||
)
|
||||
|
||||
async def test_missing_output_dir_is_404(self):
|
||||
self.api._find_comfyui_output_dir = lambda: None
|
||||
resp = await self.client.request("GET", "/prompt_manager/scan_output_dir")
|
||||
self.assertEqual(resp.status, 404)
|
||||
|
||||
async def test_error_is_500(self):
|
||||
def boom():
|
||||
raise RuntimeError("x")
|
||||
|
||||
self.api._find_comfyui_output_dir = boom
|
||||
resp = await self.client.request("GET", "/prompt_manager/scan_output_dir")
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
+339
-1
@@ -34,8 +34,10 @@ from aiohttp.test_utils import AioHTTPTestCase
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from py.api import PromptManagerAPI
|
||||
from py.config import IntegrationConfig
|
||||
from py import lora_utils
|
||||
from py.config import IntegrationConfig, PromptManagerConfig
|
||||
from py.lora_utils import get_trigger_cache
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
SECRET_KEY = "sk-civitai-super-secret-0123456789"
|
||||
|
||||
@@ -60,6 +62,8 @@ class LoraAPITestCase(AioHTTPTestCase):
|
||||
"py.lora_utils.detect_lora_manager", return_value=None
|
||||
)
|
||||
self.detect_mock = self._detect_patch.start()
|
||||
self._save_patch = patch.object(PromptManagerConfig, "save_to_file")
|
||||
self.save_mock = self._save_patch.start()
|
||||
get_trigger_cache().clear()
|
||||
|
||||
app = web.Application()
|
||||
@@ -71,6 +75,7 @@ class LoraAPITestCase(AioHTTPTestCase):
|
||||
return app
|
||||
|
||||
async def tearDownAsync(self):
|
||||
self._save_patch.stop()
|
||||
self._detect_patch.stop()
|
||||
self._config_patch.stop()
|
||||
get_trigger_cache().clear()
|
||||
@@ -115,5 +120,338 @@ class TestLoraStatusNeverReturnsKey(LoraAPITestCase):
|
||||
self.assertNotIn("civitai_api_key", data)
|
||||
|
||||
|
||||
def _png_bytes():
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image
|
||||
|
||||
buf = BytesIO()
|
||||
Image.new("RGB", (2, 2), (1, 2, 3)).save(buf, "PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, body):
|
||||
self._body = body
|
||||
self._pos = 0
|
||||
self.headers = {}
|
||||
|
||||
def read(self, amt=-1):
|
||||
if amt is None or amt < 0:
|
||||
amt = len(self._body)
|
||||
chunk = self._body[self._pos : self._pos + amt]
|
||||
self._pos += len(chunk)
|
||||
return chunk
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
def _make_lora_manager(where):
|
||||
lm = where / "ComfyUI-Lora-Manager"
|
||||
(lm / "py").mkdir(parents=True)
|
||||
(lm / "__init__.py").write_text("")
|
||||
return lm
|
||||
|
||||
|
||||
def _write_metadata(directory, stem, words=None, images=None, prompt=None):
|
||||
meta = {"file_name": f"{stem}.safetensors", "model_name": stem, "civitai": {}}
|
||||
if words is not None:
|
||||
meta["civitai"]["trainedWords"] = words
|
||||
if images is not None:
|
||||
meta["civitai"]["images"] = [
|
||||
{"url": url, "meta": {"prompt": prompt} if prompt else {}} for url in images
|
||||
]
|
||||
path = directory / f"{stem}.safetensors.metadata.json"
|
||||
path.write_text(json.dumps(meta))
|
||||
return path
|
||||
|
||||
|
||||
class TestLoraDetectRoute(LoraAPITestCase):
|
||||
|
||||
async def test_detected(self):
|
||||
self.detect_mock.return_value = str(self.tmp_root)
|
||||
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertTrue(data["detected"])
|
||||
self.assertEqual(data["path"], str(self.tmp_root))
|
||||
|
||||
async def test_not_detected(self):
|
||||
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
|
||||
data = await resp.json()
|
||||
self.assertFalse(data["detected"])
|
||||
self.assertEqual(data["path"], "")
|
||||
|
||||
async def test_detection_error_is_500(self):
|
||||
self.detect_mock.side_effect = RuntimeError("disk gone")
|
||||
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
|
||||
self.assertEqual(resp.status, 500)
|
||||
self.assertFalse((await resp.json())["success"])
|
||||
|
||||
|
||||
class TestLoraStatusRoute(LoraAPITestCase):
|
||||
|
||||
async def test_status_reports_detection_and_cache(self):
|
||||
self.detect_mock.return_value = str(self.tmp_root)
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
|
||||
|
||||
data = await (
|
||||
await self.client.request("GET", "/prompt_manager/lora/status")
|
||||
).json()
|
||||
|
||||
self.assertTrue(data["enabled"])
|
||||
self.assertTrue(data["detected"])
|
||||
self.assertEqual(data["detected_path"], str(self.tmp_root))
|
||||
self.assertFalse(data["trigger_cache_loaded"])
|
||||
|
||||
async def test_status_error_is_500(self):
|
||||
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
|
||||
resp = await self.client.request("GET", "/prompt_manager/lora/status")
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
|
||||
class TestLoraEnableRoute(LoraAPITestCase):
|
||||
|
||||
async def test_enable_with_unknown_path_is_400(self):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/lora/enable", {"enabled": True, "path": "/nowhere"}
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
|
||||
self.save_mock.assert_not_called()
|
||||
|
||||
async def test_enable_persists_and_loads_trigger_cache(self):
|
||||
lm = _make_lora_manager(self.tmp_root)
|
||||
loras = self.tmp_root / "loras"
|
||||
loras.mkdir()
|
||||
_write_metadata(loras, "neon", words=["glow"])
|
||||
self.detect_mock.return_value = str(lm)
|
||||
|
||||
with patch.object(
|
||||
lora_utils, "find_lora_directories", return_value=[str(loras)]
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/lora/enable",
|
||||
{
|
||||
"enabled": True,
|
||||
"path": str(lm),
|
||||
"trigger_words_enabled": True,
|
||||
"civitai_api_key": SECRET_KEY,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
body = await resp.text()
|
||||
data = json.loads(body)
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["path"], str(lm))
|
||||
self.assertNotIn(SECRET_KEY, body)
|
||||
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
|
||||
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
|
||||
self.save_mock.assert_called_once()
|
||||
self.assertTrue(get_trigger_cache().is_loaded)
|
||||
self.assertEqual(get_trigger_cache().get_trigger_words("neon"), ["glow"])
|
||||
|
||||
async def test_disable_clears_cache(self):
|
||||
get_trigger_cache()._cache = {"x": ["y"]}
|
||||
get_trigger_cache()._loaded = True
|
||||
|
||||
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertFalse(get_trigger_cache().is_loaded)
|
||||
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
|
||||
|
||||
async def test_persist_failure_is_500(self):
|
||||
self.save_mock.side_effect = OSError("read-only")
|
||||
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
|
||||
class TestLoraScanRoute(LoraAPITestCase):
|
||||
|
||||
def _enable(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
|
||||
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
|
||||
|
||||
@staticmethod
|
||||
def _sse_events(text):
|
||||
return [
|
||||
json.loads(line[len("data: ") :])
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data: ")
|
||||
]
|
||||
|
||||
async def _scan(self, loras):
|
||||
opener_calls = []
|
||||
|
||||
def fake_open(req, timeout=None):
|
||||
opener_calls.append(req)
|
||||
return _FakeResponse(_png_bytes())
|
||||
|
||||
cache_dir = self.tmp_root / "img_cache"
|
||||
with (
|
||||
patch.object(
|
||||
lora_utils, "find_lora_directories", return_value=[str(loras)]
|
||||
),
|
||||
patch.object(
|
||||
lora_utils, "get_lora_image_cache_dir", return_value=cache_dir
|
||||
),
|
||||
patch.object(lora_utils, "_open_url", fake_open),
|
||||
):
|
||||
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
|
||||
text = await resp.text()
|
||||
return resp, self._sse_events(text), opener_calls
|
||||
|
||||
async def test_not_enabled_is_400(self):
|
||||
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_no_path_is_400(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_scan_imports_loras_with_images_and_links_existing(self):
|
||||
self._enable()
|
||||
loras = self.tmp_root / "loras"
|
||||
loras.mkdir()
|
||||
_write_metadata(
|
||||
loras,
|
||||
"neon",
|
||||
words=["glow"],
|
||||
images=["https://civitai.com/a.png", "https://evil.example/b.png"],
|
||||
prompt="neon street",
|
||||
)
|
||||
(loras / "neon.png").write_bytes(_png_bytes())
|
||||
_write_metadata(loras, "plain")
|
||||
(loras / "broken.metadata.json").write_text("{{{")
|
||||
existing_id = self.api.db.save_prompt(
|
||||
text="plain",
|
||||
category=None,
|
||||
tags=[],
|
||||
rating=None,
|
||||
prompt_hash=generate_prompt_hash("plain"),
|
||||
)
|
||||
|
||||
resp, events, opener_calls = await self._scan(loras)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
done = events[-1]
|
||||
self.assertEqual(done["type"], "complete")
|
||||
self.assertEqual(done["total"], 3)
|
||||
self.assertEqual(done["imported"], 1)
|
||||
self.assertEqual(done["skipped"], 2)
|
||||
self.assertEqual(
|
||||
[r.full_url for r in opener_calls], ["https://civitai.com/a.png"]
|
||||
)
|
||||
self.assertEqual(
|
||||
opener_calls[0].get_header("Authorization"), f"Bearer {SECRET_KEY}"
|
||||
)
|
||||
imported = self.api.db.search_prompts(category="lora-manager")
|
||||
self.assertEqual(len(imported), 1)
|
||||
self.assertEqual(imported[0]["text"], "neon street")
|
||||
self.assertIn("lora:neon", imported[0]["tags"])
|
||||
self.assertIn("glow", imported[0]["tags"])
|
||||
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 2)
|
||||
self.assertEqual(len(self.api.db.get_prompt_images(existing_id)), 0)
|
||||
|
||||
async def test_save_failure_is_counted_as_skipped(self):
|
||||
self._enable()
|
||||
loras = self.tmp_root / "loras"
|
||||
loras.mkdir()
|
||||
_write_metadata(loras, "neon", words=["glow"])
|
||||
|
||||
with patch.object(self.api.db, "save_prompt", side_effect=RuntimeError("db")):
|
||||
resp, events, _ = await self._scan(loras)
|
||||
|
||||
self.assertEqual(events[-1]["imported"], 0)
|
||||
self.assertEqual(events[-1]["skipped"], 1)
|
||||
|
||||
|
||||
class TestLoraTriggerWordRoutes(LoraAPITestCase):
|
||||
|
||||
async def test_trigger_words_not_enabled_is_400(self):
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/lora/trigger-words?name=x"
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_trigger_words_missing_name_is_400(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
resp = await self.client.request("GET", "/prompt_manager/lora/trigger-words")
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_trigger_words_loads_cache_on_demand(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
|
||||
loras = self.tmp_root / "loras"
|
||||
loras.mkdir()
|
||||
_write_metadata(loras, "neon", words=["glow", "bright"])
|
||||
|
||||
with patch.object(
|
||||
lora_utils, "find_lora_directories", return_value=[str(loras)]
|
||||
):
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/lora/trigger-words?name=NEON"
|
||||
)
|
||||
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["trigger_words"], ["glow", "bright"])
|
||||
self.assertTrue(get_trigger_cache().is_loaded)
|
||||
|
||||
async def test_trigger_words_error_is_500(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/lora/trigger-words?name=x"
|
||||
)
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
async def test_refresh_not_enabled_is_400(self):
|
||||
resp = await self.client.request("POST", "/prompt_manager/lora/refresh-cache")
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_refresh_without_path_is_400(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
resp = await self.client.request("POST", "/prompt_manager/lora/refresh-cache")
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_refresh_reloads_cache(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
|
||||
loras = self.tmp_root / "loras"
|
||||
loras.mkdir()
|
||||
_write_metadata(loras, "neon", words=["glow"])
|
||||
|
||||
with patch.object(
|
||||
lora_utils, "find_lora_directories", return_value=[str(loras)]
|
||||
):
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/lora/refresh-cache"
|
||||
)
|
||||
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
# keyed by both the file_name stem and the metadata-file stem
|
||||
self.assertEqual(data["loras_with_trigger_words"], 2)
|
||||
|
||||
async def test_refresh_error_is_500(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
|
||||
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/lora/refresh-cache"
|
||||
)
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -4,20 +4,47 @@ No real ML model is loaded; ``model_loader`` returns small fakes that mimic
|
||||
the onnxruntime session (WD14) and the llama.cpp chat model (GGUF).
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import csv
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from py import autotag
|
||||
from py.autotag import MODELS, AutoTagService
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _patch_modules(mapping):
|
||||
"""Temporarily set entries of sys.modules, restoring only those keys.
|
||||
|
||||
unittest.mock.patch.dict restores by clearing the whole dict, which evicts
|
||||
any heavy module first imported inside the window (torch cannot be
|
||||
re-imported in-process), so this patcher touches only the named keys.
|
||||
"""
|
||||
missing = object()
|
||||
saved = {name: sys.modules.get(name, missing) for name in mapping}
|
||||
try:
|
||||
for name, module in mapping.items():
|
||||
sys.modules[name] = module
|
||||
yield
|
||||
finally:
|
||||
for name, previous in saved.items():
|
||||
if previous is missing:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = previous
|
||||
|
||||
|
||||
class FakeInput:
|
||||
def __init__(self, name="input", shape=(1, 8, 8, 3)):
|
||||
self.name = name
|
||||
@@ -327,5 +354,363 @@ class TestModelsStatus(EngineTestCase):
|
||||
)
|
||||
|
||||
|
||||
# ── Model directory resolution and singleton ──────────────────────────
|
||||
|
||||
|
||||
class TestModelsDirResolution(unittest.TestCase):
|
||||
|
||||
def test_uses_folder_paths_models_dir(self):
|
||||
fake = types.ModuleType("folder_paths")
|
||||
fake.models_dir = os.path.join(tempfile.gettempdir(), "fp_models")
|
||||
with _patch_modules({"folder_paths": fake}):
|
||||
service = AutoTagService()
|
||||
self.assertEqual(service.models_dir, Path(fake.models_dir) / "LLM")
|
||||
|
||||
def test_falls_back_without_folder_paths(self):
|
||||
with _patch_modules({"folder_paths": None}):
|
||||
service = AutoTagService()
|
||||
self.assertEqual(service.models_dir.name, "LLM")
|
||||
|
||||
def test_singleton_is_created_once(self):
|
||||
with patch.object(autotag, "_service_instance", None):
|
||||
first = autotag.get_autotag_service()
|
||||
second = autotag.get_autotag_service()
|
||||
self.assertIs(first, second)
|
||||
self.assertIsInstance(first, AutoTagService)
|
||||
|
||||
|
||||
# ── HuggingFace cache lookup ──────────────────────────────────────────
|
||||
|
||||
|
||||
def _fake_hub(snapshot_dir=None, repo_id=None, raise_scan=False):
|
||||
hub = types.ModuleType("huggingface_hub")
|
||||
hub.HFCacheInfo = object
|
||||
|
||||
def scan_cache_dir():
|
||||
if raise_scan:
|
||||
raise OSError("no cache")
|
||||
revision = types.SimpleNamespace(snapshot_path=str(snapshot_dir))
|
||||
repo = types.SimpleNamespace(
|
||||
repo_id=repo_id, repo_type="model", revisions=[revision]
|
||||
)
|
||||
other = types.SimpleNamespace(
|
||||
repo_id="someone/else", repo_type="dataset", revisions=[]
|
||||
)
|
||||
return types.SimpleNamespace(repos=[other, repo])
|
||||
|
||||
hub.scan_cache_dir = scan_cache_dir
|
||||
return hub
|
||||
|
||||
|
||||
class TestHfCacheLookup(EngineTestCase):
|
||||
|
||||
def test_cache_snapshot_with_config_is_found(self):
|
||||
snapshot = self.tmp / "snap"
|
||||
snapshot.mkdir()
|
||||
(snapshot / "config.json").write_text("{}")
|
||||
hub = _fake_hub(snapshot, MODELS["hf"]["repo"])
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
|
||||
with _patch_modules({"huggingface_hub": hub}):
|
||||
self.assertTrue(service._check_hf_model())
|
||||
status = service.get_models_status()
|
||||
|
||||
self.assertTrue(status["hf"]["downloaded"])
|
||||
self.assertEqual(Path(status["hf"]["model_path"]), snapshot)
|
||||
|
||||
def test_snapshot_without_config_is_ignored(self):
|
||||
snapshot = self.tmp / "snap"
|
||||
snapshot.mkdir()
|
||||
hub = _fake_hub(snapshot, MODELS["hf"]["repo"])
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
with _patch_modules({"huggingface_hub": hub}):
|
||||
self.assertIsNone(service._get_hf_model_path())
|
||||
|
||||
def test_scan_errors_and_missing_hub_give_none(self):
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
with _patch_modules({"huggingface_hub": _fake_hub(raise_scan=True)}):
|
||||
self.assertIsNone(service._get_hf_cache_path("x/y"))
|
||||
with _patch_modules({"huggingface_hub": None}):
|
||||
self.assertIsNone(service._get_hf_cache_path("x/y"))
|
||||
|
||||
|
||||
# ── Downloads (hub patched, no network) ───────────────────────────────
|
||||
|
||||
|
||||
class _FakeDownloadHub:
|
||||
"""huggingface_hub stand-in that materialises files instead of fetching."""
|
||||
|
||||
def __init__(self):
|
||||
self.module = types.ModuleType("huggingface_hub")
|
||||
self.module.hf_hub_download = self.hf_hub_download
|
||||
self.module.snapshot_download = self.snapshot_download
|
||||
self.calls = []
|
||||
|
||||
def hf_hub_download(self, repo_id, filename, local_dir, **kwargs):
|
||||
self.calls.append(("file", repo_id, filename))
|
||||
Path(local_dir).mkdir(parents=True, exist_ok=True)
|
||||
(Path(local_dir) / filename).write_bytes(b"model")
|
||||
return str(Path(local_dir) / filename)
|
||||
|
||||
def snapshot_download(self, repo_id, local_dir, **kwargs):
|
||||
self.calls.append(("snapshot", repo_id))
|
||||
Path(local_dir).mkdir(parents=True, exist_ok=True)
|
||||
(Path(local_dir) / "config.json").write_text("{}")
|
||||
return local_dir
|
||||
|
||||
|
||||
class TestDownloadModel(EngineTestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.hub = _FakeDownloadHub()
|
||||
self.service = AutoTagService(models_dir=self.models_dir)
|
||||
self.progress = []
|
||||
|
||||
def _cb(self, status, progress):
|
||||
self.progress.append((status, progress))
|
||||
|
||||
def test_invalid_type_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.service.download_model("nope")
|
||||
|
||||
def test_missing_hub_reports_failure(self):
|
||||
with _patch_modules({"huggingface_hub": None}):
|
||||
self.assertFalse(self.service.download_model("gguf", self._cb))
|
||||
self.assertEqual(self.progress[-1][1], 0)
|
||||
self.assertIn("huggingface_hub", self.progress[-1][0])
|
||||
|
||||
def test_wd14_downloads_each_missing_file(self):
|
||||
with _patch_modules({"huggingface_hub": self.hub.module}):
|
||||
self.assertTrue(self.service.download_model("wd14-vit", self._cb))
|
||||
self.assertTrue(self.service.download_model("wd14-vit", self._cb))
|
||||
|
||||
self.assertEqual(
|
||||
[c[2] for c in self.hub.calls], ["model.onnx", "selected_tags.csv"]
|
||||
)
|
||||
self.assertTrue(self.service.get_models_status()["wd14-vit"]["downloaded"])
|
||||
self.assertEqual(self.progress[-1], ("Download complete", 100))
|
||||
|
||||
def test_gguf_downloads_model_then_mmproj(self):
|
||||
with _patch_modules({"huggingface_hub": self.hub.module}):
|
||||
self.assertTrue(self.service.download_model("gguf", self._cb))
|
||||
|
||||
self.assertEqual(
|
||||
[c[1] for c in self.hub.calls],
|
||||
[MODELS["gguf"]["repo"], MODELS["gguf"]["mmproj_repo"]],
|
||||
)
|
||||
self.assertTrue(self.service.get_models_status()["gguf"]["downloaded"])
|
||||
self.assertIn(("Main model ready", 50), self.progress)
|
||||
|
||||
def test_gguf_skips_files_already_present(self):
|
||||
_touch_model_files(self.models_dir, "gguf")
|
||||
with _patch_modules({"huggingface_hub": self.hub.module}):
|
||||
self.assertTrue(self.service.download_model("gguf"))
|
||||
self.assertEqual(self.hub.calls, [])
|
||||
|
||||
def test_hf_snapshot_download(self):
|
||||
with _patch_modules({"huggingface_hub": self.hub.module}):
|
||||
self.assertTrue(self.service.download_model("hf", self._cb))
|
||||
self.assertTrue(self.service.download_model("hf"))
|
||||
|
||||
self.assertEqual(self.hub.calls, [("snapshot", MODELS["hf"]["repo"])])
|
||||
self.assertTrue(self.service.get_models_status()["hf"]["downloaded"])
|
||||
|
||||
def test_download_errors_are_reported_not_raised(self):
|
||||
def boom(**kwargs):
|
||||
raise OSError("disk full")
|
||||
|
||||
self.hub.module.hf_hub_download = boom
|
||||
with _patch_modules({"huggingface_hub": self.hub.module}):
|
||||
self.assertFalse(self.service.download_model("wd14-vit", self._cb))
|
||||
self.assertTrue(self.progress[-1][0].startswith("Error:"))
|
||||
|
||||
|
||||
# ── Default backend loaders with the ML libraries faked ───────────────
|
||||
|
||||
|
||||
class _FakeOrt:
|
||||
def __init__(self):
|
||||
self.module = types.ModuleType("onnxruntime")
|
||||
self.module.InferenceSession = self.InferenceSession
|
||||
self.sessions = []
|
||||
|
||||
def InferenceSession(self, path, providers):
|
||||
session = FakeOnnxSession([0.9, 0.1])
|
||||
session.path = path
|
||||
session.providers = providers
|
||||
self.sessions.append(session)
|
||||
return session
|
||||
|
||||
|
||||
class TestDefaultLoaders(EngineTestCase):
|
||||
|
||||
def test_wd14_loader_builds_session_and_reads_tags(self):
|
||||
_touch_model_files(self.models_dir, "wd14-vit")
|
||||
csv_path = self.models_dir / MODELS["wd14-vit"]["subdir"] / "selected_tags.csv"
|
||||
with open(csv_path, "w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=["tag_id", "name", "category"])
|
||||
writer.writeheader()
|
||||
writer.writerow({"tag_id": 1, "name": "1girl", "category": "0"})
|
||||
writer.writerow({"tag_id": 2, "name": "smile", "category": "0"})
|
||||
ort = _FakeOrt()
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
|
||||
with _patch_modules({"onnxruntime": ort.module}):
|
||||
service.load_model("wd14-vit", use_gpu=False)
|
||||
cpu_providers = ort.sessions[-1].providers
|
||||
service.load_model("wd14-vit", use_gpu=True)
|
||||
gpu_providers = ort.sessions[-1].providers
|
||||
|
||||
self.assertEqual(cpu_providers, ["CPUExecutionProvider"])
|
||||
self.assertEqual(gpu_providers[0], "CUDAExecutionProvider")
|
||||
self.assertEqual(service.generate_tags(str(self.image_path)), ["1girl"])
|
||||
|
||||
def test_gguf_loader_wires_llama_with_mmproj(self):
|
||||
_touch_model_files(self.models_dir, "gguf")
|
||||
llama_cpp = types.ModuleType("llama_cpp")
|
||||
chat_format = types.ModuleType("llama_cpp.llama_chat_format")
|
||||
created = {}
|
||||
|
||||
class Llama:
|
||||
def __init__(self, **kwargs):
|
||||
created.update(kwargs)
|
||||
|
||||
def create_chat_completion(self, **kwargs):
|
||||
return {"choices": [{"message": {"content": "Aa, Bb"}}]}
|
||||
|
||||
class Llava15ChatHandler:
|
||||
def __init__(self, clip_model_path):
|
||||
created["clip"] = clip_model_path
|
||||
|
||||
llama_cpp.Llama = Llama
|
||||
chat_format.Llava15ChatHandler = Llava15ChatHandler
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
|
||||
with _patch_modules(
|
||||
{"llama_cpp": llama_cpp, "llama_cpp.llama_chat_format": chat_format},
|
||||
):
|
||||
service.load_model("gguf", use_gpu=False)
|
||||
|
||||
self.assertEqual(created["n_gpu_layers"], 0)
|
||||
self.assertTrue(created["clip"].endswith(MODELS["gguf"]["mmproj_filename"]))
|
||||
self.assertEqual(service.generate_tags(str(self.image_path)), ["aa", "bb"])
|
||||
|
||||
def _fake_torch(self, cuda=False):
|
||||
torch = MagicMock(name="torch")
|
||||
torch.cuda.is_available.return_value = cuda
|
||||
torch.float16 = "f16"
|
||||
torch.bfloat16 = "bf16"
|
||||
torch.inference_mode.return_value.__enter__ = lambda s: None
|
||||
torch.inference_mode.return_value.__exit__ = lambda s, *a: False
|
||||
torch.cuda.amp.autocast.return_value.__enter__ = lambda s: None
|
||||
torch.cuda.amp.autocast.return_value.__exit__ = lambda s, *a: False
|
||||
return torch
|
||||
|
||||
def _fake_transformers(self):
|
||||
transformers = MagicMock(name="transformers")
|
||||
processor = MagicMock(name="processor")
|
||||
processor.apply_chat_template.return_value = "<convo>"
|
||||
inputs = {"input_ids": MagicMock(shape=(1, 3)), "pixel_values": MagicMock()}
|
||||
processed = MagicMock(name="inputs")
|
||||
processed.to.return_value = inputs
|
||||
processor.return_value = processed
|
||||
processor.tokenizer.decode.return_value = "Red Hat, blue sky"
|
||||
transformers.AutoProcessor.from_pretrained.return_value = processor
|
||||
model = MagicMock(name="model")
|
||||
model.generate.return_value = [list(range(10))]
|
||||
transformers.LlavaForConditionalGeneration.from_pretrained.return_value = model
|
||||
return transformers, model, processor, inputs
|
||||
|
||||
def test_hf_loader_8bit_on_cpu_and_generation(self):
|
||||
hf_dir = self.models_dir / MODELS["hf"]["subdir"]
|
||||
hf_dir.mkdir(parents=True)
|
||||
(hf_dir / "config.json").write_text("{}")
|
||||
torch = self._fake_torch(cuda=False)
|
||||
transformers, model, processor, inputs = self._fake_transformers()
|
||||
pixel = inputs["pixel_values"]
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
|
||||
with _patch_modules({"torch": torch, "transformers": transformers}):
|
||||
service.load_model("hf")
|
||||
tags = service.generate_tags(str(self.image_path))
|
||||
|
||||
self.assertEqual(tags, ["red_hat", "blue_sky"])
|
||||
kwargs = (
|
||||
transformers.LlavaForConditionalGeneration.from_pretrained.call_args.kwargs
|
||||
)
|
||||
self.assertEqual(kwargs["device_map"], "cpu")
|
||||
self.assertIn("quantization_config", kwargs)
|
||||
model.eval.assert_called_once()
|
||||
pixel.to.assert_called_with("f16")
|
||||
processor.tokenizer.decode.assert_called_once()
|
||||
decoded = processor.tokenizer.decode.call_args.args[0]
|
||||
self.assertEqual(decoded, list(range(3, 10)))
|
||||
|
||||
def test_hf_loader_without_quantization_uses_bf16_and_cuda(self):
|
||||
hf_dir = self.models_dir / MODELS["hf"]["subdir"]
|
||||
hf_dir.mkdir(parents=True)
|
||||
(hf_dir / "config.json").write_text("{}")
|
||||
torch = self._fake_torch(cuda=True)
|
||||
transformers, model, processor, inputs = self._fake_transformers()
|
||||
pixel = inputs["pixel_values"]
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
|
||||
with _patch_modules({"torch": torch, "transformers": transformers}):
|
||||
kind, (loaded, _, device, dtype) = service._load_hf_tagger(
|
||||
quantization="none"
|
||||
)
|
||||
service._tagger = (kind, (loaded, processor, device, dtype))
|
||||
service._current_model_type = "hf"
|
||||
service.generate_tags(str(self.image_path))
|
||||
service.unload_model()
|
||||
|
||||
self.assertEqual((kind, device, dtype), ("hf", "cuda", "bf16"))
|
||||
kwargs = (
|
||||
transformers.LlavaForConditionalGeneration.from_pretrained.call_args.kwargs
|
||||
)
|
||||
self.assertNotIn("quantization_config", kwargs)
|
||||
pixel.to.assert_called_with("bf16")
|
||||
torch.cuda.empty_cache.assert_called_once()
|
||||
|
||||
def test_hf_loader_without_model_raises(self):
|
||||
service = AutoTagService(models_dir=self.models_dir)
|
||||
with _patch_modules(
|
||||
{"torch": self._fake_torch(), "transformers": self._fake_transformers()[0]},
|
||||
):
|
||||
with self.assertRaises(RuntimeError):
|
||||
service._load_hf_tagger()
|
||||
|
||||
|
||||
class TestUnloadDoesNotImportTorch(EngineTestCase):
|
||||
|
||||
def test_unload_skips_torch_when_not_imported(self):
|
||||
_touch_model_files(self.models_dir, "gguf")
|
||||
service = self._service(("gguf", FakeGgufModel("x")))
|
||||
service.load_model("gguf")
|
||||
|
||||
with _patch_modules({"torch": None}):
|
||||
with patch("builtins.__import__", side_effect=ImportError("no torch")):
|
||||
service.unload_model()
|
||||
|
||||
self.assertFalse(service.is_model_loaded())
|
||||
|
||||
|
||||
class TestImageConversion(EngineTestCase):
|
||||
|
||||
def test_grayscale_images_are_converted_to_rgb(self):
|
||||
_touch_model_files(self.models_dir, "wd14-vit")
|
||||
gray = self.tmp / "gray.png"
|
||||
Image.new("L", (3, 3), 128).save(gray, "PNG")
|
||||
session = FakeOnnxSession([0.9] * 6)
|
||||
service = self._service(("wd14", (session, WD14_TAGS)))
|
||||
service.load_model("wd14-vit")
|
||||
|
||||
tags = service.generate_tags(str(gray))
|
||||
|
||||
self.assertIn("1girl", tags)
|
||||
self.assertEqual(tuple(session.run_calls[0]["input"].shape), (1, 8, 8, 3))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
+299
-2
@@ -406,11 +406,15 @@ class _FakeHTTPSHandler(urllib.request.HTTPSHandler):
|
||||
host = (urllib.parse.urlsplit(req.full_url).hostname or "").lower()
|
||||
if host in CIVITAI_HOSTS:
|
||||
headers = email.message_from_string(f"Location: {self.redirect_to}\n")
|
||||
return urllib.response.addinfourl(BytesIO(b""), headers, req.full_url, 302)
|
||||
resp = urllib.response.addinfourl(BytesIO(b""), headers, req.full_url, 302)
|
||||
resp.msg = "Found" # urllib's error processor reads response.msg
|
||||
return resp
|
||||
headers = email.message_from_string("Content-Type: image/png\n")
|
||||
return urllib.response.addinfourl(
|
||||
resp = urllib.response.addinfourl(
|
||||
BytesIO(_png_bytes()), headers, req.full_url, 200
|
||||
)
|
||||
resp.msg = "OK"
|
||||
return resp
|
||||
|
||||
|
||||
class TestRedirectsAreRefused(unittest.TestCase):
|
||||
@@ -586,5 +590,298 @@ class TestTriggerWordCache(unittest.TestCase):
|
||||
self.assertEqual(errors, [])
|
||||
|
||||
|
||||
# ── Detection, directories and metadata helpers ───────────────────────
|
||||
|
||||
|
||||
class _FakeFolderPaths:
|
||||
def __init__(self, base_path, lora_dirs=None):
|
||||
self.base_path = base_path
|
||||
self._lora_dirs = lora_dirs or []
|
||||
|
||||
def get_folder_paths(self, name):
|
||||
return list(self._lora_dirs)
|
||||
|
||||
|
||||
def _make_comfy_root(tmp):
|
||||
"""Create a fake ComfyUI tree: main.py + custom_nodes/ + models/loras."""
|
||||
root = Path(tmp) / "ComfyUI"
|
||||
(root / "custom_nodes").mkdir(parents=True)
|
||||
(root / "models" / "loras").mkdir(parents=True)
|
||||
(root / "main.py").write_text("")
|
||||
return root
|
||||
|
||||
|
||||
def _make_lora_manager(custom_nodes, name="ComfyUI-Lora-Manager"):
|
||||
lm = custom_nodes / name
|
||||
(lm / "py").mkdir(parents=True)
|
||||
(lm / "__init__.py").write_text("")
|
||||
return lm
|
||||
|
||||
|
||||
class TestFindComfyuiRoot(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
|
||||
def test_folder_paths_base_path_wins(self):
|
||||
root = _make_comfy_root(self.tmp)
|
||||
with patch.dict(sys.modules, {"folder_paths": _FakeFolderPaths(str(root))}):
|
||||
self.assertEqual(lora_utils.find_comfyui_root(), root)
|
||||
|
||||
def test_walks_up_from_module_file_without_folder_paths(self):
|
||||
root = _make_comfy_root(self.tmp)
|
||||
fake_file = root / "custom_nodes" / "ext" / "py" / "lora_utils.py"
|
||||
fake_file.parent.mkdir(parents=True)
|
||||
fake_file.write_text("")
|
||||
with (
|
||||
patch.dict(sys.modules, {"folder_paths": None}),
|
||||
patch.object(lora_utils, "__file__", str(fake_file)),
|
||||
):
|
||||
self.assertEqual(lora_utils.find_comfyui_root(), root.resolve())
|
||||
|
||||
def test_returns_none_when_no_root_found(self):
|
||||
fake_file = self.tmp / "lonely" / "lora_utils.py"
|
||||
fake_file.parent.mkdir()
|
||||
fake_file.write_text("")
|
||||
with (
|
||||
patch.dict(sys.modules, {"folder_paths": None}),
|
||||
patch.object(lora_utils, "__file__", str(fake_file)),
|
||||
):
|
||||
self.assertIsNone(lora_utils.find_comfyui_root())
|
||||
|
||||
|
||||
class TestDetectLoraManager(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
|
||||
def test_custom_path_is_used_when_it_looks_right(self):
|
||||
lm = _make_lora_manager(self.tmp)
|
||||
with patch.object(lora_utils, "find_comfyui_root", return_value=None):
|
||||
self.assertEqual(lora_utils.detect_lora_manager(str(lm)), str(lm.resolve()))
|
||||
|
||||
def test_custom_path_without_init_is_ignored(self):
|
||||
bogus = self.tmp / "nope"
|
||||
bogus.mkdir()
|
||||
with patch.object(lora_utils, "find_comfyui_root", return_value=None):
|
||||
self.assertIsNone(lora_utils.detect_lora_manager(str(bogus)))
|
||||
|
||||
def test_auto_detects_under_custom_nodes_case_insensitively(self):
|
||||
root = _make_comfy_root(self.tmp)
|
||||
(root / "custom_nodes" / "other").mkdir()
|
||||
lm = _make_lora_manager(root / "custom_nodes", "comfyui-LORA-manager")
|
||||
with patch.object(lora_utils, "find_comfyui_root", return_value=root):
|
||||
self.assertEqual(lora_utils.detect_lora_manager(), str(lm.resolve()))
|
||||
|
||||
def test_lora_manager_dir_variant_is_accepted(self):
|
||||
lm = self.tmp / "LoraManager"
|
||||
(lm / "lora_manager").mkdir(parents=True)
|
||||
(lm / "__init__.py").write_text("")
|
||||
self.assertTrue(lora_utils._looks_like_lora_manager(lm))
|
||||
|
||||
def test_nothing_found(self):
|
||||
with patch.object(lora_utils, "find_comfyui_root", return_value=None):
|
||||
self.assertIsNone(lora_utils.detect_lora_manager())
|
||||
|
||||
|
||||
class TestFindLoraDirectories(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
self.root = _make_comfy_root(self.tmp)
|
||||
self.lm = _make_lora_manager(self.root / "custom_nodes")
|
||||
|
||||
def test_collects_default_extra_runtime_and_metadata_dirs(self):
|
||||
extra_abs = self.tmp / "extra_abs"
|
||||
extra_abs.mkdir()
|
||||
(self.root / "rel_loras").mkdir()
|
||||
(self.root / "extra_model_paths.yaml").write_text(
|
||||
"a:\n base_path: %s\n loras: |\n rel_loras\n %s\n\n"
|
||||
"b: notadict\nc:\n base_path: x\n"
|
||||
% (self.root.as_posix(), extra_abs.as_posix())
|
||||
)
|
||||
runtime = self.tmp / "runtime"
|
||||
runtime.mkdir()
|
||||
nested = self.lm / "user" / "loras"
|
||||
nested.mkdir(parents=True)
|
||||
(nested / "x.metadata.json").write_text("{}")
|
||||
|
||||
with (
|
||||
patch.object(lora_utils, "find_comfyui_root", return_value=self.root),
|
||||
patch.dict(
|
||||
sys.modules,
|
||||
{"folder_paths": _FakeFolderPaths(str(self.root), [str(runtime)])},
|
||||
),
|
||||
):
|
||||
dirs = lora_utils.find_lora_directories(str(self.lm))
|
||||
|
||||
expected = {
|
||||
str((self.root / "models" / "loras").resolve()),
|
||||
str((self.root / "rel_loras").resolve()),
|
||||
str(extra_abs.resolve()),
|
||||
str(runtime.resolve()),
|
||||
str(nested.resolve()),
|
||||
}
|
||||
self.assertEqual(set(dirs), expected)
|
||||
self.assertEqual(dirs, sorted(dirs))
|
||||
|
||||
def test_non_dict_yaml_and_parse_errors_are_ignored(self):
|
||||
(self.root / "extra_model_paths.yml").write_text("- just\n- a list\n")
|
||||
self.assertEqual(lora_utils._get_extra_lora_paths(self.root), [])
|
||||
(self.root / "extra_model_paths.yml").write_text("a: [unclosed\n")
|
||||
self.assertEqual(lora_utils._get_extra_lora_paths(self.root), [])
|
||||
|
||||
def test_without_root_only_metadata_dirs_are_found(self):
|
||||
with (
|
||||
patch.object(lora_utils, "find_comfyui_root", return_value=None),
|
||||
patch.dict(sys.modules, {"folder_paths": None}),
|
||||
):
|
||||
self.assertEqual(lora_utils.find_lora_directories(str(self.lm)), [])
|
||||
|
||||
|
||||
class TestModelNameAndPreviews(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
|
||||
def test_model_name_precedence(self):
|
||||
get_name = lora_utils.get_model_name_from_metadata
|
||||
self.assertEqual(get_name({"model_name": "top"}), "top")
|
||||
self.assertEqual(get_name({"civitai": {"model": {"name": "civ"}}}), "civ")
|
||||
self.assertEqual(get_name({"file_name": "f.safetensors"}), "f.safetensors")
|
||||
self.assertEqual(get_name({}), "unknown")
|
||||
|
||||
def test_preview_images_found_by_stem(self):
|
||||
meta_path = self.tmp / "mylora.safetensors.metadata.json"
|
||||
meta_path.write_text("{}")
|
||||
(self.tmp / "mylora.png").write_bytes(b"x")
|
||||
(self.tmp / "mylora.preview.jpg").write_bytes(b"x")
|
||||
|
||||
found = lora_utils.get_preview_images_from_metadata({}, meta_path)
|
||||
|
||||
self.assertEqual(
|
||||
sorted(os.path.basename(f) for f in found),
|
||||
["mylora.png", "mylora.preview.jpg"],
|
||||
)
|
||||
self.assertEqual(
|
||||
lora_utils.get_preview_image_from_metadata({}, meta_path), found[0]
|
||||
)
|
||||
|
||||
def test_preview_uses_file_name_from_metadata(self):
|
||||
meta_path = self.tmp / "other.metadata.json"
|
||||
(self.tmp / "named.webp").write_bytes(b"x")
|
||||
found = lora_utils.get_preview_images_from_metadata(
|
||||
{"file_name": "named.safetensors"}, meta_path
|
||||
)
|
||||
self.assertEqual([os.path.basename(f) for f in found], ["named.webp"])
|
||||
self.assertIsNone(lora_utils.get_preview_image_from_metadata({}, meta_path))
|
||||
|
||||
def test_example_images_dir(self):
|
||||
self.assertIsNone(lora_utils.get_example_images_dir(str(self.tmp)))
|
||||
nested = self.tmp / "user" / "example_images"
|
||||
nested.mkdir(parents=True)
|
||||
self.assertEqual(
|
||||
lora_utils.get_example_images_dir(str(self.tmp)), str(nested.resolve())
|
||||
)
|
||||
direct = self.tmp / "example_images"
|
||||
direct.mkdir()
|
||||
self.assertEqual(
|
||||
lora_utils.get_example_images_dir(str(self.tmp)), str(direct.resolve())
|
||||
)
|
||||
|
||||
|
||||
class TestDownloadEdgeCases(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
|
||||
def test_malformed_url_is_refused(self):
|
||||
self.assertFalse(is_civitai_url("https://[invalid"))
|
||||
|
||||
def test_non_numeric_content_length_is_ignored(self):
|
||||
resp = _FakeResponse(b"abc")
|
||||
resp.headers["Content-Length"] = "lots"
|
||||
self.assertEqual(lora_utils._read_capped(resp, 10), b"abc")
|
||||
|
||||
def test_download_uses_metadata_filename_when_file_name_missing(self):
|
||||
opener = _FakeOpener(_FakeResponse(_png_bytes()))
|
||||
meta_path = self.tmp / "stemmed.safetensors.metadata.json"
|
||||
metadata = {"civitai": {"images": [{"url": "https://civitai.com/a.png"}]}}
|
||||
|
||||
with patch.object(lora_utils, "_open_url", opener):
|
||||
paths = download_civitai_images(metadata, meta_path, self.tmp / "cache")
|
||||
|
||||
self.assertEqual(len(paths), 1)
|
||||
self.assertIn("stemmed", paths[0])
|
||||
|
||||
def test_failed_downloads_are_dropped(self):
|
||||
opener = _FakeOpener(_FakeResponse(b"not an image"))
|
||||
metadata = {"civitai": {"images": [{"url": "https://civitai.com/a.png"}]}}
|
||||
|
||||
with patch.object(lora_utils, "_open_url", opener):
|
||||
paths = download_civitai_images(
|
||||
metadata, self.tmp / "x.metadata.json", self.tmp / "cache"
|
||||
)
|
||||
|
||||
self.assertEqual(paths, [])
|
||||
|
||||
|
||||
class TestInjectTriggerWords(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.cache = TriggerWordCache()
|
||||
self.cache._cache = {"style": ["neon", "glow"], "empty": []}
|
||||
self.cache._loaded = True
|
||||
|
||||
def test_not_loaded_returns_unchanged(self):
|
||||
cache = TriggerWordCache()
|
||||
self.assertEqual(
|
||||
lora_utils.inject_trigger_words("<lora:style:1>", cache),
|
||||
("<lora:style:1>", []),
|
||||
)
|
||||
|
||||
def test_no_lora_tags_returns_unchanged(self):
|
||||
self.assertEqual(
|
||||
lora_utils.inject_trigger_words("plain prompt", self.cache),
|
||||
("plain prompt", []),
|
||||
)
|
||||
|
||||
def test_words_are_appended_once(self):
|
||||
text = "a <lora:Style:0.8> b <lora:style:1>"
|
||||
self.assertEqual(
|
||||
lora_utils.inject_trigger_words(text, self.cache),
|
||||
(f"{text}, neon, glow", ["neon", "glow"]),
|
||||
)
|
||||
|
||||
def test_words_already_present_are_not_repeated(self):
|
||||
text = "NEON city <lora:style:1>"
|
||||
self.assertEqual(
|
||||
lora_utils.inject_trigger_words(text, self.cache),
|
||||
(f"{text}, glow", ["glow"]),
|
||||
)
|
||||
|
||||
def test_lora_without_words_leaves_text_alone(self):
|
||||
text = "x <lora:empty:1> <lora:unknown:1>"
|
||||
self.assertEqual(lora_utils.inject_trigger_words(text, self.cache), (text, []))
|
||||
|
||||
def test_cache_load_skips_bad_and_wordless_metadata(self):
|
||||
tmp = Path(tempfile.mkdtemp())
|
||||
self.addCleanup(shutil.rmtree, tmp, True)
|
||||
(tmp / "bad.metadata.json").write_text("{{{")
|
||||
(tmp / "nowords.metadata.json").write_text("{}")
|
||||
(tmp / "named.metadata.json").write_text(
|
||||
json.dumps({"civitai": {"trainedWords": ["w"]}})
|
||||
)
|
||||
cache = TriggerWordCache()
|
||||
with patch.object(lora_utils, "find_lora_directories", return_value=[str(tmp)]):
|
||||
self.assertEqual(cache.load(str(tmp)), 1)
|
||||
self.assertEqual(cache.get_trigger_words("named"), ["w"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user