tests/test_core.py

from __future__ import annotations

import os
import sys
import threading
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
from safetensors.numpy import save_file

os.environ["SDXL_SKIP_BOOTSTRAP"] = "1"

from lora_cache import LoraCache, MAX_FILE_BYTES
from runtime import LOGGER, MAIN_SAMPLER, Runtime
from weights import WEIGHTS


def test_headless_nodes_do_not_require_prompt_server_instance():
    root = Path(__file__).resolve().parents[1]
    sampler = (root / "vendor/kyanite/kyanite_ksampler.py").read_text()
    impact = (root / "vendor/kyanite/impact/core.py").read_text()
    assert 'getattr(PromptServer, "instance", None)' in sampler
    assert 'getattr(PromptServer, "instance", None)' in impact


def test_detailer_detaches_generated_tensor_before_numpy_conversion():
    root = Path(__file__).resolve().parents[1]
    detailer = (root / "vendor/kyanite/kyanite_face_detailer.py").read_text()
    assert "image.detach().cpu().numpy().squeeze(0)" in detailer


def test_manifest_is_fully_pinned_and_unique():
    assert len(WEIGHTS) == 11
    assert len({item.key for item in WEIGHTS}) == len(WEIGHTS)
    assert all(len(item.revision) == 40 for item in WEIGHTS)
    assert all(len(item.sha256) == 64 and item.size > 0 for item in WEIGHTS)


def test_per_lora_upload_limit_is_500_mib():
    assert MAX_FILE_BYTES == 500 * 1024 * 1024


def test_safe_prompt_policy():
    positive, negative = Runtime.prompts("a portrait", "blur", True)
    assert positive == "rating_safe, safe, sfw, a portrait"
    assert negative.endswith("rating_explicit, rating_questionable, nsfw, nude, nudity, sex, explicit")
    assert Runtime.prompts(" a ", " b ", False) == ("a", "b")


def test_eye_prompt_extracts_only_eye_tags_with_fallback():
    assert Runtime.eye_prompt("black hair, blue eyes, black jacket, closed eyes") == "blue eyes, closed eyes"
    assert Runtime.eye_prompt("single eye, portrait") == "single eye"
    assert Runtime.eye_prompt("black hair, black jacket") == "detailed eyes"


def test_lora_cache_validates_and_deduplicates(tmp_path: Path):
    upload = tmp_path / "style.safetensors"
    save_file({"lora_unet_test.lora_down.weight": __import__("numpy").zeros((2, 2), dtype="float32")}, upload)
    cache = LoraCache(tmp_path / "cache")
    first = cache.add(upload)
    second = cache.add(upload)
    assert first == second
    assert len(list((tmp_path / "cache").glob("*.safetensors"))) == 1


def test_lora_cache_reports_new_and_existing_uploads(tmp_path: Path):
    upload = tmp_path / "style.safetensors"
    save_file({"lora_unet_test.lora_down.weight": __import__("numpy").zeros((2, 2), dtype="float32")}, upload)
    cache = LoraCache(tmp_path / "cache")
    first, first_present = cache.add_with_status(upload)
    second, second_present = cache.add_with_status(upload)
    assert first == second
    assert first_present is False
    assert second_present is True


def test_lora_upload_status_uses_fast_name_and_size_lookup(tmp_path: Path, monkeypatch):
    upload = tmp_path / "style.safetensors"
    save_file({"lora_unet_test.lora_down.weight": __import__("numpy").zeros((2, 2), dtype="float32")}, upload)
    cache = LoraCache(tmp_path / "cache")
    first, first_present = cache.prepare_with_status(upload)
    monkeypatch.setattr(cache, "_digest", lambda _path: pytest.fail("duplicate upload was rehashed"))
    second, second_present = cache.prepare_with_status(upload)
    assert first == second
    assert first_present is False
    assert second_present is True


def test_lora_cache_rejects_wrong_extension(tmp_path: Path):
    upload = tmp_path / "style.bin"
    upload.write_bytes(b"not a model")
    with pytest.raises(ValueError, match="safetensors"):
        LoraCache(tmp_path / "cache").add(upload)


def test_lora_cache_rejects_non_lora_safetensors(tmp_path: Path):
    upload = tmp_path / "checkpoint.safetensors"
    save_file({"model.diffusion.weight": __import__("numpy").zeros((2, 2), dtype="float32")}, upload)
    with pytest.raises(ValueError, match="LoRA"):
        LoraCache(tmp_path / "cache").add(upload)


def test_detail_pass_constants_match_fixed_workflow():
    face = Runtime._detail_kwargs(7, "portrait", detector="bbox/face_yolov8m.pt", face=True)
    hand = Runtime._detail_kwargs(7, "portrait", detector="segm/PitHandDetailer-v2-Test-v9c.pt", hand=True)
    foot = Runtime._detail_kwargs(7, "portrait", detector="segm/foot-yolov8l.pt")
    common = {
        "guide_size_for": True, "max_size": 768,
        "seed": 7, "steps": 6, "cfg": 1.8, "neg_embedding": "BadPonyHD.safetensors",
        "embedding_emphasis": 1.0, "denoise": 0.5, "feather": 5,
        "noise_mask": True, "force_inpaint": True, "bbox_threshold": 0.5,
        "bbox_dilation": 10, "bbox_crop_factor": 3, "sam_detection_hint": "center-1",
        "sam_dilation": 0, "sam_threshold": 0.93, "sam_bbox_expansion": 0,
        "sam_mask_hint_threshold": 0.7, "sam_mask_hint_use_negative": "False",
        "drop_size": 10, "wildcard": "", "cycle": 1, "inpaint_model": False,
        "noise_mask_feather": 20, "tiled_encode": False, "tiled_decode": False,
        "scene_names": "", "scene_table": "",
    }
    hand_negative = "extra fingers, missing fingers, fused fingers, malformed fingers, broken fingers, deformed hand, mutated hand, bad anatomy, 6 fingers"
    foot_negative = "extra fingers, extra toes, fused fingers, poorly detailed fingers, polish, toe polish"
    assert face == common | {
        "guide_size": 512,
        "sampler_name": "dpmpp_2m", "scheduler": "sgm_uniform",
        "positive_text": "portrait", "negative_text": hand_negative,
        "bbox_model_name": "bbox/face_yolov8m.pt", "skip_keywords": "",
    }
    assert hand == common | {
        "guide_size": 512,
        "sampler_name": "lcm", "scheduler": "simple",
        "positive_text": "detailed hand, 5 fingers", "negative_text": hand_negative,
        "bbox_model_name": "segm/PitHandDetailer-v2-Test-v9c.pt",
        "skip_keywords": "flowjob, footjob, fj",
    }
    assert foot == common | {
        "guide_size": 384,
        "sampler_name": "lcm", "scheduler": "simple",
        "positive_text": "Ultra detailed feet, 5 toes", "negative_text": foot_negative,
        "bbox_model_name": "segm/foot-yolov8l.pt", "skip_keywords": "",
    }


def test_main_sampler_contract_matches_fixed_workflow():
    assert MAIN_SAMPLER == {
        "base_positive": "",
        "neg_embedding": "BadPonyHD.safetensors",
        "embedding_emphasis": 0.95,
        "width": 896,
        "height": 504,
        "batch_size": 1,
        "seed_type": "Fixed",
        "steps": 27,
        "cfg": 8,
        "sampler_name": "dpmpp_2m",
        "scheduler": "sgm_uniform",
        "denoise": 1.0,
        "ckpt_name": "None",
        "lora_name": "None",
        "lora_mode": "LoRA Model + Clip",
        "strength_model": 0,
        "strength_clip": 0,
        "clip_skip": 2,
    }


def test_runtime_stage_logger_keeps_info_timings_enabled():
    assert LOGGER.isEnabledFor(20)


def test_original_flow_uses_independent_randomized_node_seeds(monkeypatch):
    values = iter((11, 22, 33, 44, 55))
    monkeypatch.setattr("runtime.random.randrange", lambda _limit: next(values))
    seeds = Runtime.generation_seeds(None)
    assert (seeds.main, seeds.face, seeds.foot, seeds.hand, seeds.eye) == (11, 22, 33, 44, 55)


def test_original_detailer_branches_and_face_expression_are_preserved():
    source = (Path(__file__).resolve().parents[1] / "runtime.py").read_text()
    assert '"dmd2_sdxl_4step_lora.safetensors"' in source
    assert "expression_seductive=True" in source
    assert "detail_scale=1.5" in source
    assert "detail_scale=1.0" in source


def test_optional_detail_branches_and_eye_final_pass_are_present():
    source = (Path(__file__).resolve().parents[1] / "runtime.py").read_text()
    assert "if face_detail:" in source
    assert "if hands_detail:" in source
    assert source.count("if feet_detail:") == 2
    assert "if eyes_detail:" in source
    assert 'detector="segm/PitEyeDetailer-v2-seg.pt"' in source
    assert source.index("if eyes_detail:") > source.index("c.foot_apply.apply")


def test_detail_flags_skip_work_and_eyes_run_last_with_extracted_prompt(monkeypatch):
    calls = []

    class FakeImage:
        def __getitem__(self, _index):
            return self

        def detach(self):
            return self

        def cpu(self):
            return self

        def numpy(self):
            return np.zeros((8, 8, 3), dtype=np.float32)

    class FakeLoader:
        def __init__(self, label):
            self.label = label

        def load_lora_model_only(self, model, _name, _strength):
            calls.append(self.label)
            return (model,)

    class FakeDetailer:
        def doit(self, *, image, model, clip, vae, **kwargs):
            calls.append(("detail", kwargs["bbox_model_name"], kwargs["positive_text"]))
            return image, model, clip, vae, None

    class FakeFootPrepare:
        def prepare(self, *, image, model, clip, vae, **_kwargs):
            calls.append("feet-prepare")
            return image, {"patch_items": []}, None, None, model, clip, vae, None

    class FakeFootApply:
        def apply(self, image, _batch):
            calls.append("feet-apply")
            return (image,)

    class FakeSampler:
        def sample(self, *, model, clip, vae, positive_text, negative_text, seed, **_kwargs):
            return FakeImage(), model, clip, vae, positive_text, negative_text, seed

    monkeypatch.setitem(
        sys.modules,
        "vendor.kyanite.kyanite_ksampler",
        SimpleNamespace(KyaniteKSampler=FakeSampler),
    )
    runtime = Runtime.__new__(Runtime)
    runtime.lock = threading.RLock()
    runtime.loras = SimpleNamespace()
    runtime.components = SimpleNamespace(
        model=object(), clip=object(), vae=object(), user_loader=object(),
        hyper_loader=FakeLoader("hyper"), dmd_loader=FakeLoader("dmd"),
        face_detailer=FakeDetailer(), foot_prepare=FakeFootPrepare(),
        foot_apply=FakeFootApply(),
    )

    runtime.generate(
        "black hair, blue eyes, black jacket", "blur", False,
        False, False, False, False, [], [], [], 7,
    )
    assert calls == []

    runtime.generate(
        "black hair, blue eyes, black jacket", "blur", False,
        True, True, True, True, [], [], [], 7,
    )
    assert calls[-1] == (
        "detail", "segm/PitEyeDetailer-v2-seg.pt", "blue eyes"
    )
    assert calls.index("feet-apply") < len(calls) - 1