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

extensions_built_in/captioner/BaseCaptioner.py

@@ -54,6 +54,8 @@ class CaptionConfig:
         # (low-vram machines); percent is the fraction of linears offloaded
         self.layer_offloading = kwargs.get("layer_offloading", False)
         self.layer_offloading_percent = kwargs.get("layer_offloading_percent", 1.0)
+        # [{path, strength, name}] applied as sidechains, never merged
+        self.loras = kwargs.get("loras", None) or []
 
 
 class BaseCaptioner(BaseExtensionProcess):
@@ -89,6 +91,7 @@ class BaseCaptioner(BaseExtensionProcess):
 
         self.caption_config = self.caption_config_class(**self.get_conf("caption", {}))
         self.model = None
+        self.lora_stack = None
         self.processor = None
         self.model2 = None
         self.processor2 = None
@@ -103,6 +106,7 @@ class BaseCaptioner(BaseExtensionProcess):
             self.start_stop_watcher()
             self.update_status("running", "Loading Model")
             self.load_model()
+            self.load_loras()
             self.maybe_compile_models()
             self.update_status("running", "Looking for files")
             self.find_files()
@@ -202,6 +206,28 @@ class BaseCaptioner(BaseExtensionProcess):
     def load_model(self):
         raise NotImplementedError("Model loading not implemented for this captioner")
 
+    def load_loras(self):
+        """Attach the configured LoRAs to self.model as forward-hook sidechains
+        (the inference engine's hook mode): the low-rank branch runs alongside
+        the layer in the model's compute dtype and the base weights are never
+        touched, so quantized weights stay on their grid and nothing is merged."""
+        specs = [
+            l for l in self.caption_config.loras if isinstance(l, dict) and l.get("path")
+        ]
+        if not specs or self.model is None:
+            return
+        from toolkit.inference_lora import LoRAStack
+
+        # self is the holder: LoRAStack reads .model (and .text_encoder, unused here)
+        stack = LoRAStack(self, mode="hook").load(
+            specs, status_fn=self.print_and_status_update
+        )
+        stack.apply(status_fn=self.print_and_status_update)
+        self.lora_stack = stack
+        self.print_and_status_update(
+            f"Applied {len(stack.loras)} LoRA(s) as sidechains"
+        )
+
     def maybe_compile_models(self):
         if not self.caption_config.compile:
             return

extensions_built_in/captioner/Qwen25OmniCaptioner.py

@@ -0,0 +1,428 @@
+"""Generic Qwen2.5-Omni captioner: same prompt-driven image / audio / video
+flow as Qwen3OmniCaptioner, on the Qwen2.5-Omni thinker from
+extensions_built_in/llm_models (single-file convrot8 checkpoints or an HF repo).
+Unlike AceStepCaptioner (fixed transcribe + describe prompts, ACE-Step/YuE2
+layouts), the caption is whatever caption_prompt asks for."""
+
+from collections import OrderedDict
+
+import os
+import torch
+
+from transformers import AutoConfig, AutoProcessor, Qwen2_5OmniForConditionalGeneration
+
+from toolkit.basic import flush
+
+from .BaseCaptioner import BaseCaptioner
+import logging
+import traceback
+import warnings
+
+warnings.filterwarnings("ignore")
+logging.disable(logging.WARNING)
+
+from extensions_built_in.llm_models.src.thinker import (  # noqa: E402  shared fast paths
+    load_thinker_single_file,
+    prepare_thinker,
+)
+
+# frame sampling rate for video captioning
+VIDEO_FPS = 2
+
+# still-image files caption through the image pipeline (no audio, no frames)
+IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"}
+# audio-only files caption through the audio pipeline (no frames)
+AUDIO_EXTENSIONS = {"mp3", "wav", "flac", "ogg", "m4a", "aac"}
+
+TARGET_SAMPLE_RATE = 16000
+
+# fixed generation ceiling under compiled decode: a constant max_length keeps
+# the static kv cache (and so the compiled decode graph) at one shape for every
+# file; the real per-caption budget is enforced by a stopping criterion
+STATIC_MAX_LENGTH = 8704
+
+# thinker hidden size -> config/processor repo; single-file checkpoints carry no config
+BASE_REPO_BY_HIDDEN = {3584: "Qwen/Qwen2.5-Omni-7B", 2048: "Qwen/Qwen2.5-Omni-3B"}
+DEFAULT_BASE_REPO = "Qwen/Qwen2.5-Omni-7B"
+
+
+class Qwen25OmniCaptioner(BaseCaptioner):
+    """Captions images, audio and video with the Qwen2.5-Omni thinker."""
+
+    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
+        super(Qwen25OmniCaptioner, self).__init__(process_id, job, config, **kwargs)
+
+    def _config_source(self, ckpt: str) -> str:
+        """Single-file checkpoints carry no config; pick the base repo (config +
+        processor) from the thinker hidden size."""
+        from safetensors import safe_open
+
+        with safe_open(ckpt, "pt") as f:
+            hidden = f.get_slice("model.embed_tokens.weight").get_shape()[1]
+        try:
+            return BASE_REPO_BY_HIDDEN[hidden]
+        except KeyError:
+            raise RuntimeError(
+                f"Unknown Qwen2.5-Omni thinker (hidden size {hidden}); "
+                "use a folder/repo that ships its own config instead"
+            ) from None
+
+    def load_model(self):
+        name_or_path = self.caption_config.model_name_or_path
+        want_quant = self.caption_config.quantize
+
+        if name_or_path.endswith(".safetensors"):
+            from toolkit.models.v2.resolver import resolve_component_file
+
+            # thinker checkpoints live in the comfy text_encoders/ folder
+            ckpt = resolve_component_file(
+                name_or_path,
+                folder="text_encoders",
+                component="qwen2.5-omni thinker",
+                status_fn=self.print_and_status_update,
+            )
+            self.print_and_status_update(
+                f"Loading Qwen2.5-Omni thinker from {os.path.basename(ckpt)}"
+            )
+            cfg_src = self._config_source(ckpt)
+            config = AutoConfig.from_pretrained(cfg_src)
+            model, prequantized = load_thinker_single_file(
+                ckpt, config.thinker_config, self.torch_dtype
+            )
+            if prequantized and want_quant:
+                print(
+                    "[AITK] Checkpoint is pre-quantized (convrot8); the quantize "
+                    "setting is ignored."
+                )
+                want_quant = False
+        else:
+            self.print_and_status_update(f"Loading Qwen2.5-Omni from {name_or_path}")
+            cfg_src = name_or_path
+            full = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+                name_or_path, dtype=self.torch_dtype, device_map="cpu"
+            )
+            model = prepare_thinker(full)
+            del full
+
+        model.eval()
+        model.to(self.device_torch)
+        if want_quant:
+            from optimum.quanto import freeze
+
+            from toolkit.util.quantize import get_qtype, quantize
+
+            self.print_and_status_update(
+                f"Quantizing thinker ({self.caption_config.qtype})"
+            )
+            quantize(model, weights=get_qtype(self.caption_config.qtype))
+            freeze(model)
+            flush()
+
+        # built from config on the single-file path, so no sampling defaults were
+        # loaded; greedy decode falls into repetition loops on long captions
+        model.generation_config.do_sample = True
+        model.generation_config.temperature = 0.7
+        model.generation_config.top_p = 0.8
+        model.generation_config.top_k = 20
+        model.generation_config.repetition_penalty = 1.05
+
+        self.model = model
+        self.processor = AutoProcessor.from_pretrained(cfg_src)
+        if self.caption_config.low_vram:
+            self.model.to("cpu")
+        flush()
+
+    def maybe_compile_models(self):
+        """CUDA-graph decode: static kv cache + compiled decode (see
+        Qwen3OmniCaptioner). Eager HF decode on the 7B thinker is launch-bound."""
+        if not self.caption_config.compile:
+            return
+        if self.caption_config.low_vram:
+            print("[AITK] low_vram is on; skipping compiled decode.")
+            return
+        import importlib.util
+
+        if importlib.util.find_spec("triton") is None:
+            print("[AITK] compile requested but triton is not installed, skipping.")
+            return
+        self.model.generation_config.cache_implementation = "static"
+        print(
+            "[AITK] Compiled decode enabled (static cache + cuda graphs). "
+            "The first file compiles (~2 min cold, faster once cached)."
+        )
+
+    @staticmethod
+    def _is_image_file(file_path: str) -> bool:
+        return os.path.splitext(file_path)[1].lower().lstrip(".") in IMAGE_EXTENSIONS
+
+    @staticmethod
+    def _is_audio_file(file_path: str) -> bool:
+        return os.path.splitext(file_path)[1].lower().lstrip(".") in AUDIO_EXTENSIONS
+
+    def _build_messages(self, _file_path: str):
+        if self._is_image_file(_file_path):
+            media = {"type": "image", "image": _file_path}
+        elif self._is_audio_file(_file_path):
+            media = {"type": "audio", "audio": _file_path}
+        else:
+            media = {"type": "video", "video": _file_path}
+        return [
+            {
+                "role": "user",
+                "content": [
+                    media,
+                    {"type": "text", "text": self.caption_config.caption_prompt},
+                ],
+            }
+        ]
+
+    def _size_kwargs(self):
+        max_pixels = self.caption_config.max_res * self.caption_config.max_res
+        # Qwen2.5-Omni's image processor takes total pixel budgets, not edges
+        return {"min_pixels": min(64 * 28 * 28, max_pixels), "max_pixels": max_pixels}
+
+    @staticmethod
+    def _load_audio(file_path: str):
+        import torchaudio
+
+        waveform, sr = torchaudio.load(file_path)
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(dim=0, keepdim=True)
+        if sr != TARGET_SAMPLE_RATE:
+            waveform = torchaudio.functional.resample(waveform, sr, TARGET_SAMPLE_RATE)
+        return waveform.squeeze(0).numpy()
+
+    def _prep_media(self, file_path: str):
+        """CPU side of one file, safe to run in a worker thread: decode +
+        subsample frames (or load the image), extract the audio track, render
+        the chat text. At batch size 1 the full processor (tokenize, resize,
+        mel) runs here too, so the main thread only moves tensors and
+        generates."""
+        if self._is_image_file(file_path):
+            from PIL import Image
+
+            image = Image.open(file_path).convert("RGB")
+            item = {"file": file_path, "kind": "image", "image": image, "audio": None}
+        elif self._is_audio_file(file_path):
+            item = {
+                "file": file_path,
+                "kind": "audio",
+                "audio": self._load_audio(file_path),
+            }
+        else:
+            from transformers.video_utils import load_video
+
+            frames = load_video(file_path, fps=VIDEO_FPS)
+            if isinstance(frames, tuple):
+                frames = frames[0]
+            audio = None
+            try:
+                a = self._load_audio(file_path)
+                if a is not None and a.size > 0:
+                    audio = a
+            except Exception:
+                pass
+            item = {
+                "file": file_path,
+                "kind": "video_audio" if audio is not None else "video_silent",
+                "frames": frames,
+                "audio": audio,
+            }
+        item["text"] = self.processor.apply_chat_template(
+            self._build_messages(file_path),
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        if self.caption_config.batch_size <= 1:
+            item["inputs"] = self._process_items([item])
+        return item
+
+    def _process_items(self, items):
+        kind = items[0]["kind"]
+        if kind == "image":
+            return self.processor(
+                text=[it["text"] for it in items],
+                images=[it["image"] for it in items],
+                return_tensors="pt",
+                padding=True,
+                **self._size_kwargs(),
+            )
+        if kind == "audio":
+            return self.processor(
+                text=[it["text"] for it in items],
+                audio=[it["audio"] for it in items],
+                return_tensors="pt",
+                padding=True,
+                sampling_rate=TARGET_SAMPLE_RATE,
+            )
+        use_audio = kind == "video_audio"
+        return self.processor(
+            text=[it["text"] for it in items],
+            audio=[it["audio"] for it in items] if use_audio else None,
+            videos=[it["frames"] for it in items],
+            return_tensors="pt",
+            padding=True,
+            use_audio_in_video=use_audio,
+            sampling_rate=TARGET_SAMPLE_RATE,
+            fps=VIDEO_FPS,
+            **self._size_kwargs(),
+        )
+
+    def _gen_kwargs(self, input_len: int) -> dict:
+        max_new = self.caption_config.max_new_tokens
+        gen_kwargs = {"max_new_tokens": max_new}
+        if self.model.generation_config.cache_implementation == "static":
+            if input_len + 16 < STATIC_MAX_LENGTH:
+                from transformers.generation import (
+                    MaxLengthCriteria,
+                    StoppingCriteriaList,
+                )
+
+                gen_kwargs = {
+                    "max_length": STATIC_MAX_LENGTH,
+                    "stopping_criteria": StoppingCriteriaList(
+                        [
+                            MaxLengthCriteria(
+                                max_length=min(input_len + max_new, STATIC_MAX_LENGTH)
+                            )
+                        ]
+                    ),
+                }
+            else:
+                # prompt would not fit the fixed cache; eager for this file
+                gen_kwargs["cache_implementation"] = "dynamic"
+        return gen_kwargs
+
+    def _caption_batch(self, items):
+        """Batched generate over preprocessed items (all the same kind: image,
+        audio, video with audio, or silent video). Returns captions in item order."""
+        use_audio = items[0]["kind"] == "video_audio"
+        if len(items) == 1 and "inputs" in items[0]:
+            inputs = items[0]["inputs"]
+        else:
+            inputs = self._process_items(items)
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        inputs = inputs.to(self.device_torch).to(self.torch_dtype)
+        # a generate that dies between static-cache creation and its first
+        # forward leaves model._cache with uninitialized layers; transformers
+        # then raises AttributeError reading cache.max_batch_size on every
+        # later call, masking the original error — drop the stale cache
+        stale_cache = getattr(self.model, "_cache", None)
+        if stale_cache is not None and not stale_cache.is_initialized:
+            del self.model._cache
+        # under static cache, generate hands the forward a prepared 4D mask;
+        # the true 2D padding mask is needed for the prefill rope index
+        self.model._pad_mask_2d = inputs.get("attention_mask", None)
+        input_len = inputs["input_ids"].shape[1]
+        generated_ids = self.model.generate(
+            **inputs,
+            use_audio_in_video=use_audio,
+            **self._gen_kwargs(input_len),
+        )
+        captions = self.processor.batch_decode(
+            generated_ids[:, input_len:],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )
+        return [c.strip() for c in captions]
+
+    def run_caption_loop(self):
+        """Batched pipeline: CPU worker threads decode/preprocess files ahead of
+        the GPU, files are grouped by kind into batches, and each batch runs one
+        model.generate call so decode work is wide enough to saturate the GPU."""
+        import concurrent.futures
+        from collections import deque
+
+        import tqdm as tqdm_mod
+
+        batch_size = max(1, int(self.caption_config.batch_size))
+        # smoothing near 1 weights recent files heavily, so the rate estimate
+        # recovers quickly after the slow compile-warmup files
+        pbar = tqdm_mod.tqdm(
+            total=len(self.file_paths),
+            desc="Captioning files",
+            unit="file",
+            smoothing=0.9,
+        )
+
+        def finish(file_path, caption):
+            if caption is not None:
+                self.save_caption_for_file(file_path, caption)
+            self.step_num += 1
+            self.update_step()
+            pbar.update(1)
+
+        def flush_bucket(bucket):
+            if len(bucket) == 0:
+                return
+            items = list(bucket)
+            bucket.clear()
+            n_real = len(items)
+            # keep the batch shape constant for the compiled decode graph:
+            # pad a final partial bucket by repeating the last item
+            if (
+                self.model.generation_config.cache_implementation == "static"
+                and 1 < n_real < batch_size
+            ):
+                items = items + [items[-1]] * (batch_size - n_real)
+            try:
+                captions = self._caption_batch(items)[:n_real]
+                for it, cap in zip(items[:n_real], captions):
+                    finish(it["file"], cap)
+            except Exception as e:
+                print(f"Batch failed ({e}); retrying files individually")
+                traceback.print_exc()
+                for it in items[:n_real]:
+                    finish(it["file"], self.get_caption_for_file(it["file"]))
+
+        executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=max(1, int(self.caption_config.num_workers))
+        )
+        try:
+            futures = deque()
+            file_iter = iter(self.file_paths)
+            # keep a couple of batches of decode work in flight ahead of the GPU
+            lookahead = batch_size * 2 + 2
+            for _ in range(lookahead):
+                path = next(file_iter, None)
+                if path is None:
+                    break
+                futures.append((path, executor.submit(self._prep_media, path)))
+
+            # batches must be homogeneous: the processor call differs per kind
+            buckets = {"image": [], "audio": [], "video_audio": [], "video_silent": []}
+            while futures:
+                if self.is_ui_captioner:
+                    self.maybe_stop()
+                    if self.is_stopping:
+                        break
+                path, fut = futures.popleft()
+                nxt = next(file_iter, None)
+                if nxt is not None:
+                    futures.append((nxt, executor.submit(self._prep_media, nxt)))
+                try:
+                    item = fut.result()
+                except Exception as e:
+                    print(f"Error preprocessing {path}: {e}")
+                    finish(path, None)
+                    continue
+                bucket = buckets[item["kind"]]
+                bucket.append(item)
+                if len(bucket) >= batch_size:
+                    flush_bucket(bucket)
+            for bucket in buckets.values():
+                flush_bucket(bucket)
+        finally:
+            executor.shutdown(wait=False, cancel_futures=True)
+            pbar.close()
+
+    def get_caption_for_file(self, file_path: str) -> str:
+        # single-file path (and the per-file fallback when a batch fails):
+        # same prep + generate flow as the batched loop, for one item
+        try:
+            return self._caption_batch([self._prep_media(file_path)])[0]
+        except Exception as e:
+            print(f"Error processing {file_path}: {e}")
+            traceback.print_exc()
+            return None

extensions_built_in/captioner/__init__.py

@@ -37,6 +37,18 @@ class Qwen3OmniCaptionerExtension(Extension):
         return Qwen3OmniCaptioner
 
 
+class Qwen25OmniCaptionerExtension(Extension):
+    uid = "Qwen25OmniCaptioner"
+    name = "Qwen 2.5 Omni Captioner"
+
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .Qwen25OmniCaptioner import Qwen25OmniCaptioner
+
+        return Qwen25OmniCaptioner
+
+
 class Ideogram4CaptionerExtension(Extension):
     uid = "Ideogram4Captioner"
     name = "Ideogram4 Captioner"
@@ -53,5 +65,6 @@ AI_TOOLKIT_EXTENSIONS = [
     AceStepCaptionerExtension,
     Qwen3VLCaptionerExtension,
     Qwen3OmniCaptionerExtension,
+    Qwen25OmniCaptionerExtension,
     Ideogram4CaptionerExtension,
 ]

extensions_built_in/llm_models/qwen25_omni.py

@@ -36,7 +36,7 @@ from toolkit.basic import flush
 from toolkit.config_modules import GenerateImageConfig, ModelConfig
 from toolkit.dto import DTO
 from toolkit.models.base_model import BaseModel
-from toolkit.models.v2.resolver import resolve_named_file
+from toolkit.models.v2.resolver import resolve_component_file
 from toolkit.print import print_acc
 
 from .src.thinker import attach_fast_paths, load_thinker_single_file, prepare_thinker
@@ -45,7 +45,7 @@ BASE_REPO = "Qwen/Qwen2.5-Omni-7B"
 # thinker hidden size -> config/processor repo; single-file checkpoints carry no config
 BASE_REPO_BY_HIDDEN = {3584: "Qwen/Qwen2.5-Omni-7B", 2048: "Qwen/Qwen2.5-Omni-3B"}
 # single-file convrot8 thinker written by scripts/convert_vllm_to_comfy.py
-DEFAULT_CHECKPOINT = "ostris/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors"
+DEFAULT_CHECKPOINT = "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors"
 DEFAULT_INSTRUCTION = "Describe this in detail."
 SAMPLE_RATE = 16000
 
@@ -124,7 +124,10 @@ class Qwen25OmniLLM(BaseModel):
             raise NotImplementedError("Layer offloading is not implemented for qwen25_omni")
 
         if name_or_path.endswith(".safetensors"):
-            ckpt = resolve_named_file(name_or_path, component="qwen2.5-omni thinker")
+            # thinker checkpoints live in the comfy text_encoders/ folder
+            ckpt = resolve_component_file(
+                name_or_path, folder="text_encoders", component="qwen2.5-omni thinker"
+            )
             self.print_and_status_update(f"Loading Qwen2.5-Omni thinker from {os.path.basename(ckpt)}")
             cfg_src = self._config_source(ckpt)
             config = AutoConfig.from_pretrained(cfg_src)

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",

ui/src/app/jobs/new/options.tsx

@@ -1387,7 +1387,7 @@ export const modelArchs: ModelArch[] = [
     defaults: {
       // default updates when [selected, unselected] in the UI
       'config.process[0].model.name_or_path': [
-        'ostris/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors',
+        'ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors',
         defaultNameOrPath,
       ],
       'config.process[0].model.quantize': [true, false],

ui/src/components/CaptionSimpleJob.tsx

@@ -1,4 +1,6 @@
-import React from 'react';
+import React, { useState } from 'react';
+import { Plus, X } from 'lucide-react';
+import LoraBrowserModal, { LoraPick } from '@/components/generate/LoraBrowserModal';
 import {
   Checkbox,
   CreatableSelectInput,
@@ -8,7 +10,7 @@ import {
   TextAreaInput,
   TextInput,
 } from '@/components/formInputs';
-import { CaptionJobConfig } from '@/types';
+import { CaptionJobConfig, CaptionLora } from '@/types';
 import { handleCaptionerTypeChange } from '@/helpers/captionJobConfig';
 import {
   batchSizeOptions,
@@ -37,6 +39,13 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
   const promptPresetNames = Object.keys(captionPrompts);
   const minNewTokens = selectedCaptionOption?.minNewTokens ?? 0;
   const newTokensOptions = maxNewTokensOptions.filter(option => parseInt(option.value) >= minNewTokens);
+  const [loraModalOpen, setLoraModalOpen] = useState(false);
+  const loras: CaptionLora[] = jobConfig.config.process[0].caption.loras || [];
+  const setLoras = (next: CaptionLora[]) => setJobConfig(next, 'config.process[0].caption.loras');
+  const addLora = (pick: LoraPick) => {
+    if (loras.some(l => l.path === pick.path)) return;
+    setLoras([...loras, { path: pick.path, name: pick.name, strength: 1.0 }]);
+  };
 
   return (
     <div className="text-sm text-gray-400">
@@ -120,6 +129,60 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
           />
         </div>
       )}
+      {selectedCaptionOption?.supportsLoras && (
+        <div className="mt-4">
+          <div className="text-xs text-gray-300 mb-1">LoRAs</div>
+          <div className="space-y-2">
+            {loras.map((l, i) => (
+              <div key={l.path} className="bg-gray-950/60 border border-gray-800 rounded-md px-2 py-1.5">
+                <div className="flex items-center gap-2">
+                  <span className="text-xs truncate flex-1 text-gray-200" title={l.path}>
+                    {l.name}
+                  </span>
+                  <input
+                    type="number"
+                    step={0.05}
+                    min={-2}
+                    max={3}
+                    value={l.strength}
+                    onChange={e => {
+                      const v = parseFloat(e.target.value);
+                      const next = [...loras];
+                      next[i] = { ...l, strength: isNaN(v) ? 0 : v };
+                      setLoras(next);
+                    }}
+                    className="w-16 text-xs px-1.5 py-0.5 bg-gray-950 border border-gray-700 rounded text-gray-100 text-right"
+                  />
+                  <button
+                    type="button"
+                    onClick={() => setLoras(loras.filter(x => x.path !== l.path))}
+                    className="text-gray-500 hover:text-red-400"
+                    title="Remove"
+                  >
+                    <X className="w-3.5 h-3.5" />
+                  </button>
+                </div>
+              </div>
+            ))}
+            <button
+              type="button"
+              onClick={() => setLoraModalOpen(true)}
+              className="w-full px-2 py-1 rounded-md bg-gray-800 hover:bg-gray-700 text-gray-200 text-xs flex items-center justify-center gap-1"
+            >
+              <Plus className="w-3.5 h-3.5" /> Add LoRA
+            </button>
+          </div>
+          <div className="text-[11px] text-gray-500 mt-1">
+            Applied as sidechains at caption time; the model weights are never merged.
+          </div>
+          <LoraBrowserModal
+            isOpen={loraModalOpen}
+            onClose={() => setLoraModalOpen(false)}
+            onPick={addLora}
+            cloudLoras={selectedCaptionOption?.cloudLoras}
+          />
+        </div>
+      )}
       <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
         <div>
           <SelectInput

ui/src/components/DatasetImageCard.tsx

@@ -148,7 +148,13 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
   useEffect(() => {
     if (!isAutoCaptioning) return;
     const interval = setInterval(() => setPollTick(t => t + 1), 5000);
-    return () => clearInterval(interval);
+    return () => {
+      clearInterval(interval);
+      // the job writes every caption before it reports completion, so the last
+      // ones land between the final tick and polling stopping here; without one
+      // more refresh they stay stale until the card is remounted
+      setPollTick(t => t + 1);
+    };
   }, [isAutoCaptioning]);
 
   const saveCaption = () => {

ui/src/components/generate/LoraBrowserModal.tsx

@@ -1,9 +1,10 @@
 'use client';
 
 import { useEffect, useMemo, useState } from 'react';
-import { ChevronDown, ChevronRight, FolderOpen, Loader2, Search } from 'lucide-react';
+import { ChevronDown, ChevronRight, Cloud, FolderOpen, Loader2, Search } from 'lucide-react';
 import { Modal } from '@/components/Modal';
 import { apiClient } from '@/utils/api';
+import { CloudLora } from '@/types';
 
 export interface LoraPick {
   path: string;
@@ -30,13 +31,19 @@ interface Props {
   isOpen: boolean;
   onClose: () => void;
   onPick: (lora: LoraPick) => void;
+  // LoRAs published for the selected model. Picked by their hub reference; the
+  // backend finds or downloads the file when the job runs.
+  cloudLoras?: CloudLora[];
 }
 
+type Tab = 'cloud' | 'jobs' | 'models';
+
 const fmtSize = (n: number) => (n > 1e9 ? `${(n / 1e9).toFixed(2)} GB` : `${(n / 1e6).toFixed(0)} MB`);
 const fmtDate = (ms: number) => new Date(ms).toLocaleString();
 
-export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) {
-  const [tab, setTab] = useState<'jobs' | 'models'>('jobs');
+export default function LoraBrowserModal({ isOpen, onClose, onPick, cloudLoras }: Props) {
+  const hasCloud = !!cloudLoras?.length;
+  const [tab, setTab] = useState<Tab>(hasCloud ? 'cloud' : 'jobs');
   const [loading, setLoading] = useState(false);
   const [jobs, setJobs] = useState<JobEntry[]>([]);
   const [models, setModels] = useState<LoraFile[]>([]);
@@ -46,6 +53,7 @@ export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) {
 
   useEffect(() => {
     if (!isOpen) return;
+    setTab(hasCloud ? 'cloud' : 'jobs');
     setLoading(true);
     apiClient
       .get('/api/loras')
@@ -60,27 +68,52 @@ export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) {
 
   const q = filter.trim().toLowerCase();
   const filteredJobs = useMemo(
-    () => (q ? jobs.filter(j => j.name.toLowerCase().includes(q) || j.files.some(f => f.name.toLowerCase().includes(q))) : jobs),
+    () =>
+      q
+        ? jobs.filter(j => j.name.toLowerCase().includes(q) || j.files.some(f => f.name.toLowerCase().includes(q)))
+        : jobs,
     [jobs, q],
   );
-  const filteredModels = useMemo(() => (q ? models.filter(m => (m.relpath || m.name).toLowerCase().includes(q)) : models), [models, q]);
+  const filteredModels = useMemo(
+    () => (q ? models.filter(m => (m.relpath || m.name).toLowerCase().includes(q)) : models),
+    [models, q],
+  );
+  const filteredCloud = useMemo(() => {
+    const list = cloudLoras || [];
+    return q ? list.filter(c => c.name.toLowerCase().includes(q) || c.path.toLowerCase().includes(q)) : list;
+  }, [cloudLoras, q]);
 
-  const pick = (f: LoraFile, jobName?: string) => {
-    onPick({ path: f.path, name: jobName ? `${jobName} / ${f.name.replace(/\.safetensors$/, '')}` : (f.relpath || f.name).replace(/\.safetensors$/, '') });
+  const pickCloud = (c: CloudLora) => {
+    onPick({ path: c.path, name: c.name });
     onClose();
   };
 
-  const tabClass = (t: 'jobs' | 'models') =>
+  const pick = (f: LoraFile, jobName?: string) => {
+    onPick({
+      path: f.path,
+      name: jobName
+        ? `${jobName} / ${f.name.replace(/\.safetensors$/, '')}`
+        : (f.relpath || f.name).replace(/\.safetensors$/, ''),
+    });
+    onClose();
+  };
+
+  const tabClass = (t: Tab) =>
     `px-3 py-1.5 text-sm border-b-2 ${tab === t ? 'border-blue-500 text-blue-400' : 'border-transparent text-gray-400 hover:text-gray-200'}`;
 
   return (
     <Modal isOpen={isOpen} onClose={onClose} title="Add LoRA" size="lg">
       <div className="text-gray-200">
         <div className="flex items-center gap-2 border-b border-gray-800 mb-3">
-          <button className={tabClass('jobs')} onClick={() => setTab('jobs')}>
+          {hasCloud && (
+            <button type="button" className={tabClass('cloud')} onClick={() => setTab('cloud')}>
+              For this model
+            </button>
+          )}
+          <button type="button" className={tabClass('jobs')} onClick={() => setTab('jobs')}>
             Training jobs
           </button>
-          <button className={tabClass('models')} onClick={() => setTab('models')}>
+          <button type="button" className={tabClass('models')} onClick={() => setTab('models')}>
             Models folder
           </button>
           <div className="flex-1" />
@@ -100,14 +133,42 @@ export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) {
               <Loader2 className="w-4 h-4 animate-spin" /> Loading…
             </div>
           )}
+          {tab === 'cloud' && (
+            <div>
+              <div className="text-[11px] text-gray-500 mb-2 flex items-center gap-1">
+                <Cloud className="w-3.5 h-3.5" /> Downloaded to the loras folder on first use
+              </div>
+              {filteredCloud.length === 0 && (
+                <div className="text-xs text-gray-500">No published LoRAs for this model.</div>
+              )}
+              <div className="space-y-0.5">
+                {filteredCloud.map(c => (
+                  <button
+                    type="button"
+                    key={c.path}
+                    onClick={() => pickCloud(c)}
+                    className="w-full px-2 py-1 rounded text-xs text-left hover:bg-blue-900/40"
+                    title={c.path}
+                  >
+                    <div className="truncate text-gray-200">{c.name}</div>
+                    <div className="truncate font-mono text-[11px] text-gray-500">{c.path}</div>
+                    {c.description && <div className="truncate text-[11px] text-gray-500">{c.description}</div>}
+                  </button>
+                ))}
+              </div>
+            </div>
+          )}
           {!loading && tab === 'jobs' && (
             <div className="space-y-1">
-              {filteredJobs.length === 0 && <div className="text-xs text-gray-500">No training jobs with saved LoRAs.</div>}
+              {filteredJobs.length === 0 && (
+                <div className="text-xs text-gray-500">No training jobs with saved LoRAs.</div>
+              )}
               {filteredJobs.map(j => {
                 const isOpen = open.has(j.id) || !!q;
                 return (
                   <div key={j.id} className="rounded-md bg-gray-900 border border-gray-800">
                     <button
+                      type="button"
                       className="w-full flex items-center gap-2 px-3 py-2 text-sm text-left hover:bg-gray-800/60"
                       onClick={() =>
                         setOpen(prev => {
@@ -118,7 +179,11 @@ export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) {
                         })
                       }
                     >
-                      {isOpen ? <ChevronDown className="w-4 h-4 text-gray-500" /> : <ChevronRight className="w-4 h-4 text-gray-500" />}
+                      {isOpen ? (
+                        <ChevronDown className="w-4 h-4 text-gray-500" />
+                      ) : (
+                        <ChevronRight className="w-4 h-4 text-gray-500" />
+                      )}
                       <span className="truncate flex-1">{j.name}</span>
                       <span className="text-[11px] text-gray-500">
                         {j.files.length} file{j.files.length === 1 ? '' : 's'}
@@ -130,6 +195,7 @@ export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) {
                           .filter(f => !q || f.name.toLowerCase().includes(q) || j.name.toLowerCase().includes(q))
                           .map(f => (
                             <button
+                              type="button"
                               key={f.path}
                               onClick={() => pick(f, j.name)}
                               className="w-full flex items-center gap-2 px-2 py-1 rounded text-xs text-left hover:bg-blue-900/40"
@@ -152,10 +218,13 @@ export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) {
               <div className="text-[11px] text-gray-500 mb-2 flex items-center gap-1">
                 <FolderOpen className="w-3.5 h-3.5" /> {lorasRoot || 'MODELS_PATH/loras'}
               </div>
-              {filteredModels.length === 0 && <div className="text-xs text-gray-500">No .safetensors files found under the loras folder.</div>}
+              {filteredModels.length === 0 && (
+                <div className="text-xs text-gray-500">No .safetensors files found under the loras folder.</div>
+              )}
               <div className="space-y-0.5">
                 {filteredModels.map(f => (
                   <button
+                    type="button"
                     key={f.path}
                     onClick={() => pick(f)}
                     className="w-full flex items-center gap-2 px-2 py-1 rounded text-xs text-left hover:bg-blue-900/40"

ui/src/helpers/captionOptions.ts

@@ -1,4 +1,4 @@
-import { GroupedSelectOption, SelectOption } from "@/types";
+import { CloudLora, GroupedSelectOption, SelectOption } from "@/types";
 
 type CaptionGroup = 'image' | 'music' | 'video' | 'image/video/sound';
 type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
@@ -9,6 +9,12 @@ export interface CaptionOption {
     group: CaptionGroup;
     hasMultiLinePrompts?: boolean;
     minNewTokens?: number;
+    // captioner can run LoRAs on its model (applied as sidechains, never merged)
+    supportsLoras?: boolean;
+    // LoRAs published for this model, offered in the LoRA browser alongside
+    // local files. Paths are 'org/repo/path_to/file.safetensors'; the captioner
+    // searches the models folder for the file before downloading it to loras/.
+    cloudLoras?: CloudLora[];
     defaults?: { [key: string]: any };
     additionalSections?: AdditionalSections[];
     name_or_path_options?: SelectOption[];
@@ -168,6 +174,45 @@ export const captionerTypes: CaptionOption[] = [
             'caption.thinking',
         ],
     },
+    {
+        name: 'Qwen25OmniCaptioner',
+        label: 'Qwen2.5-Omni',
+        group: 'image/video/sound',
+        supportsLoras: true,
+        cloudLoras: [
+            {
+                path: 'ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_lora_caption_this_song.safetensors',
+                name: 'Caption This Song',
+            },
+        ],
+        defaults: {
+            'config.process[0].caption.loras': [[], undefined],
+            'config.process[0].caption.model_name_or_path': ['ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors', defaultNameOrPath],
+            'config.process[0].caption.extensions': [[...extensionsVideo, ...extensionsImage, ...extensionsAudio], defaultExtensions],
+            'config.process[0].caption.caption_prompt': [defaultVideoCaptionPrompt, undefined],
+            'config.process[0].caption.max_res': [512, undefined],
+            'config.process[0].caption.max_new_tokens': [512, undefined],
+            'config.process[0].caption.batch_size': [1, undefined],
+            'config.process[0].caption.compile': [true, false],
+        },
+        name_or_path_options: [
+            { value: 'ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors', label: 'ai-toolkit/Qwen2.5-Omni-7B (convrot8)' },
+            { value: 'Qwen/Qwen2.5-Omni-7B', label: 'Qwen/Qwen2.5-Omni-7B' },
+            { value: 'Qwen/Qwen2.5-Omni-3B', label: 'Qwen/Qwen2.5-Omni-3B' },
+        ],
+        captionPrompts: {
+            'General': defaultVideoCaptionPrompt,
+            'MiniMax H4 T2V': minimaxT2VCaptionPrompt,
+            'MiniMax H4 Image': minimaxImageCaptionPrompt,
+            'YuE2': yue2CaptionPrompt,
+        },
+        additionalSections: [
+            'caption.caption_prompt',
+            'caption.max_res',
+            'caption.max_new_tokens',
+            'caption.batch_size',
+        ],
+    },
     {
         name: 'Ideogram4Captioner',
         label: 'Ideogram 4 Captioner',

ui/src/types.ts

@@ -303,6 +303,21 @@ export interface JobConfig {
   meta: MetaConfig;
 }
 
+// A LoRA published on the hub, offered for a specific model option. `path` is a
+// 'org/repo/path_to/file.safetensors' reference; the backend looks for it under
+// the models folder first and downloads it into MODELS_PATH/loras if missing.
+export interface CloudLora {
+  path: string;
+  name: string;
+  description?: string;
+}
+
+export interface CaptionLora {
+  path: string;
+  name: string;
+  strength: number;
+}
+
 export interface CaptionProcessConfig {
   type: string;
   sqlite_db_path?: string;
@@ -328,6 +343,7 @@ export interface CaptionProcessConfig {
     batch_size?: number;
     layer_offloading?: boolean;
     layer_offloading_percent?: number;
+    loras?: CaptionLora[];
   }
 }
 

version.py

@@ -1 +1 @@
-VERSION = "0.13.13"
+VERSION = "0.13.14"