Add support for Qwen 2.5 Omni and add lora loading for captioning models

toolkit/inference_lora.py

@@ -147,6 +147,7 @@ class LoRAEntry:
 class InferenceLoRA:
     def __init__(self, path: str, strength: float = 1.0, name: Optional[str] = None):
         self.path = path
+        self.local_path: Optional[str] = None  # set by load(): path resolved on disk
         self.strength = float(strength)
         self.name = name or os.path.splitext(os.path.basename(path))[0]
         self.entries: List[LoRAEntry] = []
@@ -154,10 +155,15 @@ class InferenceLoRA:
         self._hooks: List = []
 
     # ---- loading / key resolution ----
-    def load(self, holder) -> "InferenceLoRA":
-        if not os.path.isfile(self.path):
-            raise FileNotFoundError(f"LoRA not found: {self.path}")
-        sd = load_file(self.path)
+    def load(self, holder, status_fn=None) -> "InferenceLoRA":
+        from toolkit.models.v2.resolver import resolve_lora_file
+
+        # a spec path may be a hub reference (org/repo/file.safetensors); it is
+        # searched for under the models folder before anything is downloaded.
+        # self.path stays as given: it is the spec identity the engine keys its
+        # loaded stack on.
+        self.local_path = resolve_lora_file(self.path, status_fn=status_fn)
+        sd = load_file(self.local_path)
         convert = getattr(holder, "convert_lora_weights_before_load", None)
         if callable(convert):
             try:
@@ -395,7 +401,9 @@ class LoRAStack:
 
     def load(self, specs: List[dict], status_fn=None):
         for spec in specs:
-            lora = InferenceLoRA(spec["path"], spec.get("strength", 1.0), spec.get("name")).load(self.holder)
+            lora = InferenceLoRA(spec["path"], spec.get("strength", 1.0), spec.get("name")).load(
+                self.holder, status_fn=status_fn
+            )
             if status_fn:
                 status_fn(
                     f"LoRA {lora.name}: {len(lora.entries)} modules"

toolkit/models/registry.py

@@ -55,7 +55,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
     "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
     "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}},
-    "qwen25_omni": {"modality": "text", "model": {"name_or_path": "ostris/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 1, "guidance_scale": 1.0}},
+    "qwen25_omni": {"modality": "text", "model": {"name_or_path": "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 1, "guidance_scale": 1.0}},
     "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},
 }
 

toolkit/models/v2/resolver.py

@@ -188,6 +188,109 @@ def resolve_comfy_file(
     )
 
 
+def resolve_component_file(
+    path: str,
+    folder: str,
+    component: str = "model",
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    """Resolve a single-file component that lives in one comfy-layout folder
+    (e.g. text_encoders/) to a local file.
+
+    Accepts a local path or a hub reference 'org/repo/path/to/file.safetensors'.
+    Anything on disk wins: the file's own spot under MODELS_PATH/<folder>, then
+    under MODELS_PATH, then a recursive search for the filename in <folder> and
+    finally in all of MODELS_PATH. A miss is downloaded into MODELS_PATH/<folder>
+    at its repo-relative path.
+    """
+    if os.path.isfile(path):
+        return path
+    root = os.path.join(MODELS_PATH, folder)
+    splits = path.split("/")
+    # 'org/repo/rel/file.safetensors' -> repo 'org/repo', file 'rel/file.safetensors'
+    rel_path = "/".join(splits[2:]) if len(splits) >= 3 else splits[-1]
+    filename = splits[-1]
+    for candidate in (
+        os.path.join(root, rel_path),
+        os.path.join(root, filename),
+        os.path.join(MODELS_PATH, rel_path),
+        os.path.join(MODELS_PATH, filename),
+        os.path.join(MODELS_PATH, path),
+    ):
+        if os.path.isfile(candidate):
+            return candidate
+    for search_root in (root, MODELS_PATH):
+        found = find_file_recursive(search_root, filename)
+        if found is not None:
+            return found
+
+    if len(splits) < 3:
+        raise FileNotFoundError(
+            f"{component} not found: {path}. Must be a local file or "
+            "'org/repo/filename.safetensors' to download from the Hugging Face Hub."
+        )
+
+    import huggingface_hub
+
+    repo_id = "/".join(splits[:2])
+    if status_fn is not None:
+        status_fn(f"Downloading {rel_path} from {repo_id} into {root}")
+    return huggingface_hub.hf_hub_download(
+        repo_id=repo_id, filename=rel_path, token=hf_token, local_dir=root
+    )
+
+
+def resolve_lora_file(
+    path: str,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    """Resolve a LoRA reference to a local file.
+
+    Accepts a local path or a hub reference 'org/repo/path/to/file.safetensors'.
+    Anything already on disk wins: the exact relative spot under
+    MODELS_PATH/loras and MODELS_PATH, then a recursive search for the filename
+    under loras/ and finally under all of MODELS_PATH (a LoRA picked from a
+    training job's output folder lives outside loras/). Only a file found
+    nowhere is downloaded, into MODELS_PATH/loras at its repo-relative path.
+    """
+    if os.path.isfile(path):
+        return path
+    loras_root = os.path.join(MODELS_PATH, "loras")
+    splits = path.split("/")
+    # 'org/repo/rel/file.safetensors' -> repo 'org/repo', file 'rel/file.safetensors'
+    rel_path = "/".join(splits[2:]) if len(splits) >= 3 else splits[-1]
+    filename = splits[-1]
+    for candidate in (
+        os.path.join(loras_root, rel_path),
+        os.path.join(MODELS_PATH, rel_path),
+        os.path.join(loras_root, filename),
+        os.path.join(MODELS_PATH, path),
+    ):
+        if os.path.isfile(candidate):
+            return candidate
+    for root in (loras_root, MODELS_PATH):
+        found = find_file_recursive(root, filename)
+        if found is not None:
+            return found
+
+    if len(splits) < 3:
+        raise FileNotFoundError(
+            f"LoRA not found: {path}. Must be a local file or "
+            "'org/repo/filename.safetensors' to download from the Hugging Face Hub."
+        )
+
+    import huggingface_hub
+
+    repo_id = "/".join(splits[:2])
+    if status_fn is not None:
+        status_fn(f"Downloading LoRA {rel_path} from {repo_id} into {loras_root}")
+    return huggingface_hub.hf_hub_download(
+        repo_id=repo_id, filename=rel_path, token=hf_token, local_dir=loras_root
+    )
+
+
 def resolve_named_file(
     path: str,
     component: str = "model",