From ce64e65c39e71b7e86d06ca44183a492e1a15418 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 16 Sep 2026 10:36:19 -0600 Subject: [PATCH] Add support for Qwen 2.5 Omni and add lora loading for captioning models --- .../captioner/BaseCaptioner.py | 26 ++ .../captioner/Qwen25OmniCaptioner.py | 428 ++++++++++++++++++ extensions_built_in/captioner/__init__.py | 13 + extensions_built_in/llm_models/qwen25_omni.py | 9 +- toolkit/inference_lora.py | 18 +- toolkit/models/registry.py | 2 +- toolkit/models/v2/resolver.py | 103 +++++ ui/src/app/jobs/new/options.tsx | 2 +- ui/src/components/CaptionSimpleJob.tsx | 67 ++- ui/src/components/DatasetImageCard.tsx | 8 +- .../components/generate/LoraBrowserModal.tsx | 95 +++- ui/src/helpers/captionOptions.ts | 47 +- ui/src/types.ts | 16 + version.py | 2 +- 14 files changed, 808 insertions(+), 28 deletions(-) create mode 100644 extensions_built_in/captioner/Qwen25OmniCaptioner.py diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index 0333623..a89a6ac 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/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 diff --git a/extensions_built_in/captioner/Qwen25OmniCaptioner.py b/extensions_built_in/captioner/Qwen25OmniCaptioner.py new file mode 100644 index 0000000..d1d218a --- /dev/null +++ b/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 diff --git a/extensions_built_in/captioner/__init__.py b/extensions_built_in/captioner/__init__.py index 30a407d..e838757 100644 --- a/extensions_built_in/captioner/__init__.py +++ b/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, ] diff --git a/extensions_built_in/llm_models/qwen25_omni.py b/extensions_built_in/llm_models/qwen25_omni.py index cbf29a1..0fa6e01 100644 --- a/extensions_built_in/llm_models/qwen25_omni.py +++ b/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) diff --git a/toolkit/inference_lora.py b/toolkit/inference_lora.py index f406ab2..caf8e90 100644 --- a/toolkit/inference_lora.py +++ b/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" diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index 0133426..a2357c2 100644 --- a/toolkit/models/registry.py +++ b/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}}, } diff --git a/toolkit/models/v2/resolver.py b/toolkit/models/v2/resolver.py index 29e9cea..46556cd 100644 --- a/toolkit/models/v2/resolver.py +++ b/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/, then + under MODELS_PATH, then a recursive search for the filename in and + finally in all of MODELS_PATH. A miss is downloaded into MODELS_PATH/ + 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", diff --git a/ui/src/app/jobs/new/options.tsx b/ui/src/app/jobs/new/options.tsx index 78844fb..638fe9b 100644 --- a/ui/src/app/jobs/new/options.tsx +++ b/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], diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index f268044..a3941df 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/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 = ({ 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 (
@@ -120,6 +129,60 @@ const CaptionSimpleJob: React.FC = ({ jobConfig, setJobConfig, gpuIDs, se />
)} + {selectedCaptionOption?.supportsLoras && ( +
+
LoRAs
+
+ {loras.map((l, i) => ( +
+
+ + {l.name} + + { + 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" + /> + +
+
+ ))} + +
+
+ Applied as sidechains at caption time; the model weights are never merged. +
+ setLoraModalOpen(false)} + onPick={addLora} + cloudLoras={selectedCaptionOption?.cloudLoras} + /> +
+ )}
= ({ 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 = () => { diff --git a/ui/src/components/generate/LoraBrowserModal.tsx b/ui/src/components/generate/LoraBrowserModal.tsx index 118bde8..14c938c 100644 --- a/ui/src/components/generate/LoraBrowserModal.tsx +++ b/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(hasCloud ? 'cloud' : 'jobs'); const [loading, setLoading] = useState(false); const [jobs, setJobs] = useState([]); const [models, setModels] = useState([]); @@ -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 (
- + )} + -
@@ -100,14 +133,42 @@ export default function LoraBrowserModal({ isOpen, onClose, onPick }: Props) { Loading…
)} + {tab === 'cloud' && ( +
+
+ Downloaded to the loras folder on first use +
+ {filteredCloud.length === 0 && ( +
No published LoRAs for this model.
+ )} +
+ {filteredCloud.map(c => ( + + ))} +
+
+ )} {!loading && tab === 'jobs' && (
- {filteredJobs.length === 0 &&
No training jobs with saved LoRAs.
} + {filteredJobs.length === 0 && ( +
No training jobs with saved LoRAs.
+ )} {filteredJobs.map(j => { const isOpen = open.has(j.id) || !!q; return (