From d8713fc4cc414bc2ecb3e6ec73c0a190e34c1724 Mon Sep 17 00:00:00 2001
From: Jason Hoku <jlevien808@gmail.com>
Date: Mon, 18 May 2026 01:24:59 -0700
Subject: [PATCH] =?UTF-8?q?feat:=20florence2=5Fhires=20=E2=80=94=20model?=
 =?UTF-8?q?=20loading=20with=20module-level=20cache?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 florence2_hires.py                    | 35 +++++++++++
 tests/test_florence2_model_loading.py | 83 +++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 tests/test_florence2_model_loading.py

diff --git a/florence2_hires.py b/florence2_hires.py
index a1f08fe..2a2a3c8 100644
--- a/florence2_hires.py
+++ b/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
diff --git a/tests/test_florence2_model_loading.py b/tests/test_florence2_model_loading.py
new file mode 100644
index 0000000..f1ccd07
--- /dev/null
+++ b/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