feat: florence2_hires — model loading with module-level cache

florence2_hires.py

@@ -278,3 +278,38 @@ def preflight_florence2():
     install hint if either required node is missing.
     """
     get_florence2_node_classes()
+
+
+# Module-level cache for Florence2 models. Keyed by model name string.
+# Cleared on Comfy restart; survives across jobs in the same session.
+_FLORENCE2_MODEL_CACHE = {}
+
+
+def load_florence2_model(model_name):
+    """Load a Florence2 model, caching by name across jobs.
+
+    Args:
+        model_name: HF Hub id, e.g. "microsoft/Florence-2-base"
+
+    Returns:
+        The loaded model handle (whatever the loader node returns at index 0).
+    """
+    if model_name in _FLORENCE2_MODEL_CACHE:
+        return _FLORENCE2_MODEL_CACHE[model_name]
+
+    from ltx_video_generation import _call_node, _unwrap
+
+    classes = get_florence2_node_classes()
+    loader_cls = classes["DownloadAndLoadFlorence2Model"]
+
+    print(f"[Florence2HiResFix] Loading {model_name}...")
+    result = _call_node(
+        loader_cls,
+        model=model_name,
+        precision="fp16",
+        attention="sdpa",
+        convert_to_safetensors=False,
+    )
+    handle = _unwrap(result, 0)
+    _FLORENCE2_MODEL_CACHE[model_name] = handle
+    return handle

tests/test_florence2_model_loading.py

@@ -0,0 +1,83 @@
+"""Tests for Florence2 model loading + caching."""
+from unittest.mock import MagicMock
+
+import pytest
+
+
+def test_load_florence2_model_caches_per_name(monkeypatch):
+    """Same model name -> load called once."""
+    load_call_count = {"n": 0}
+
+    class FakeLoader:
+        FUNCTION = "loadmodel"
+
+        def loadmodel(self, **kwargs):
+            load_call_count["n"] += 1
+            return (f"loaded:{kwargs['model']}",)
+
+    import nodes
+    monkeypatch.setattr(nodes, "NODE_CLASS_MAPPINGS", {
+        "DownloadAndLoadFlorence2Model": FakeLoader,
+        "Florence2Run": object,
+    }, raising=False)
+
+    # Reset cache (module-level dict)
+    import florence2_hires
+    florence2_hires._FLORENCE2_MODEL_CACHE.clear()
+
+    m1 = florence2_hires.load_florence2_model("microsoft/Florence-2-base")
+    m2 = florence2_hires.load_florence2_model("microsoft/Florence-2-base")
+    assert m1 is m2
+    assert load_call_count["n"] == 1
+
+
+def test_load_florence2_model_different_names_load_separately(monkeypatch):
+    load_call_count = {"n": 0}
+
+    class FakeLoader:
+        FUNCTION = "loadmodel"
+
+        def loadmodel(self, **kwargs):
+            load_call_count["n"] += 1
+            return (f"loaded:{kwargs['model']}",)
+
+    import nodes
+    monkeypatch.setattr(nodes, "NODE_CLASS_MAPPINGS", {
+        "DownloadAndLoadFlorence2Model": FakeLoader,
+        "Florence2Run": object,
+    }, raising=False)
+
+    import florence2_hires
+    florence2_hires._FLORENCE2_MODEL_CACHE.clear()
+
+    m1 = florence2_hires.load_florence2_model("microsoft/Florence-2-base")
+    m2 = florence2_hires.load_florence2_model("microsoft/Florence-2-large")
+    assert m1 != m2
+    assert load_call_count["n"] == 2
+
+
+def test_load_florence2_passes_hidden_defaults(monkeypatch):
+    """Verify hidden defaults: precision=fp16, attention=sdpa, convert_to_safetensors=False."""
+    captured = {}
+
+    class FakeLoader:
+        FUNCTION = "loadmodel"
+
+        def loadmodel(self, **kwargs):
+            captured.update(kwargs)
+            return ("model_handle",)
+
+    import nodes
+    monkeypatch.setattr(nodes, "NODE_CLASS_MAPPINGS", {
+        "DownloadAndLoadFlorence2Model": FakeLoader,
+        "Florence2Run": object,
+    }, raising=False)
+
+    import florence2_hires
+    florence2_hires._FLORENCE2_MODEL_CACHE.clear()
+    florence2_hires.load_florence2_model("microsoft/Florence-2-base")
+
+    assert captured["model"] == "microsoft/Florence-2-base"
+    assert captured["precision"] == "fp16"
+    assert captured["attention"] == "sdpa"
+    assert captured["convert_to_safetensors"] is False