diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index 3c8f298b..08724419 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/extensions_built_in/captioner/BaseCaptioner.py @@ -45,6 +45,11 @@ class CaptionConfig: "caption_prompt", "Describe this image in detail." ) self.compile = kwargs.get("compile", False) + # batched captioners: files generated per model.generate call, and CPU + # preprocessing threads that keep the GPU fed. Default 1 for VRAM + # safety; raise it to saturate a large GPU. + self.batch_size = kwargs.get("batch_size", 1) + self.num_workers = kwargs.get("num_workers", 3) class BaseCaptioner(BaseExtensionProcess): diff --git a/extensions_built_in/captioner/Qwen3OmniCaptioner.py b/extensions_built_in/captioner/Qwen3OmniCaptioner.py new file mode 100644 index 00000000..05fa7e2a --- /dev/null +++ b/extensions_built_in/captioner/Qwen3OmniCaptioner.py @@ -0,0 +1,695 @@ +from transformers import AutoConfig, AutoProcessor +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeThinkerForConditionalGeneration, +) +from collections import OrderedDict + +import os +import torch +import torch.nn.functional as F + +from toolkit.basic import flush +from toolkit.util.comfy_quant_import import ( + import_comfy_quantized_layers, + parse_comfy_quant_blob, +) +from toolkit.util.convrot_quant import rotate + +from .BaseCaptioner import BaseCaptioner +from .Qwen3VLCaptioner import patch_qwen_vl_patch_embed +import logging +import traceback +import warnings + +warnings.filterwarnings("ignore") +logging.disable(logging.WARNING) + +# frame sampling rate for video captioning +VIDEO_FPS = 2 + +# single-file comfy-format checkpoint (thinker only, convrot8 int8) produced by +# scripts/convert_vllm_to_comfy.py. This is always what we load — never the +# original bf16 shards. +CONVROT_FILENAME = "qwen3_omni_30b_a3b_instruct_thinker_convrot8.safetensors" +# config + processor (tokenizer, feature extractors) come from the original repo +BASE_REPO = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + + +class OstrisQwen3OmniThinker(Qwen3OmniMoeThinkerForConditionalGeneration): + """Thinker with static-cache-safe MRoPE handling. + + Upstream breaks under ``cache_implementation="static"``: generate passes a + prepared 4D bool attention mask, but the forward's rope-delta block does + ``1 - attention_mask`` and ``get_rope_index`` assumes a 2D long padding + mask. We compute position_ids ourselves — prefill from the true 2D mask + (stashed by the caller before generate), decode from cache_position with + no data-dependent ops — so the upstream block (which only runs when + position_ids is None) is skipped entirely. Also required for CUDA-graph + decode: the decode branch is sync-free and shape-static.""" + + _pad_mask_2d = None + + def forward( + self, + input_ids=None, + attention_mask=None, + position_ids=None, + past_key_values=None, + cache_position=None, + **kwargs, + ): + if position_ids is None and input_ids is not None: + if input_ids.shape[1] > 1 or self.rope_deltas is None: + # prefill: replicate the upstream math with a valid 2D mask + mask2d = ( + attention_mask + if attention_mask is not None and attention_mask.dim() == 2 + else self._pad_mask_2d + ) + if mask2d is None: + mask2d = torch.ones_like(input_ids) + mask2d = mask2d.long() + if mask2d.shape[1] != input_ids.shape[1]: + # static cache pads the mask out to max_cache_len + mask2d = mask2d[:, : input_ids.shape[1]] + feature_attention_mask = kwargs.get("feature_attention_mask", None) + if feature_attention_mask is not None: + audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) + else: + audio_feature_lengths = kwargs.get("audio_feature_lengths", None) + delta0 = (1 - mask2d).sum(dim=-1).unsqueeze(1) + position_ids, rope_deltas = self.get_rope_index( + input_ids, + kwargs.get("image_grid_thw", None), + kwargs.get("video_grid_thw", None), + mask2d, + kwargs.get("use_audio_in_video", None) or False, + audio_feature_lengths, + kwargs.get("video_second_per_grid", None), + ) + self.rope_deltas = rope_deltas - delta0 + else: + # decode: continue from the cache position; sync-free + batch_size, seq_length = input_ids.shape + deltas = self.rope_deltas.to(input_ids.device) + if cache_position is not None: + pos = cache_position.view(1, -1) + deltas + else: + # get_seq_length may be a tensor (static cache); keep it on-device + past_len = ( + past_key_values.get_seq_length() + if past_key_values is not None + else 0 + ) + pos = ( + torch.arange(seq_length, device=input_ids.device).view(1, -1) + + past_len + + deltas + ) + position_ids = pos.unsqueeze(0).expand(3, batch_size, seq_length) + return super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + + +class ConvRot8Experts(torch.nn.Module): + """Drop-in replacement for Qwen3OmniMoeThinkerTextExperts that keeps the + fused expert banks in comfy convrot8 storage (regular-Hadamard rotated, + per-output-row symmetric int8). Experts are dequantized one at a time at + forward, so the full-precision banks (the bulk of the 30B) never + materialize.""" + + def __init__( + self, gate_up_q, gate_up_s, gate_up_rot, down_q, down_s, down_rot, dtype + ): + super().__init__() + self.num_experts = gate_up_q.shape[0] + self.gate_up_rot = gate_up_rot + self.down_rot = down_rot + self.out_dtype = dtype + self.register_buffer("gate_up_q", gate_up_q.contiguous(), persistent=False) + self.register_buffer("down_q", down_q.contiguous(), persistent=False) + # fp32 scales stored as uint8 byte views so a later .to(dtype=...) on the + # model cannot silently cast them (same convention as the cr8 backend) + self.register_buffer( + "gate_up_s", + gate_up_s.detach().float().contiguous().view(torch.uint8), + persistent=False, + ) + self.register_buffer( + "down_s", + down_s.detach().float().contiguous().view(torch.uint8), + persistent=False, + ) + + def _dequant(self, qdata, scales_u8, rot, i): + # scales are [E, out, 1]; rotation is self-inverse along the in dim + scales = scales_u8.view(torch.float32) + w = qdata[i].float() * scales[i] + return rotate(w, rot).to(self.out_dtype) + + def _dequant_batch(self, qdata, scales_u8, rot, hit, dtype): + """Dequantize the hit experts in one shot: [n_hit, out, in].""" + scales = scales_u8.view(torch.float32) + w = qdata[hit].float() * scales[hit] + return rotate(w, rot).to(dtype) + + def forward(self, hidden_states, top_k_index, top_k_weights): + """Fully batched MoE: group tokens by expert (sort + bincount), pad the + groups to a rectangle, dequantize the hit experts in one op, and run the + whole layer as two bmms — no per-expert python loop. Decode touches only + the routed experts' weights; prefill runs every expert in one launch.""" + hidden_dim = hidden_states.shape[1] + top_k = top_k_index.shape[-1] + + n_pairs = hidden_states.shape[0] * top_k + if n_pairs <= 64: + # decode-size batches: one bmm per (token, expert) pair with fixed + # shapes and NO data-dependent ops — the grouped path below needs + # nonzero()/max() which each force a GPU sync, and 2 syncs x 48 + # layers per token is exactly what stalls the GPU at small batch + flat = top_k_index.reshape(-1) + x_rep = hidden_states.repeat_interleave(top_k, dim=0).unsqueeze(1) + w_gate_up = self._dequant_batch( + self.gate_up_q, + self.gate_up_s, + self.gate_up_rot, + flat, + hidden_states.dtype, + ) + gate, up = torch.bmm(x_rep, w_gate_up.transpose(1, 2)).chunk(2, dim=-1) + del w_gate_up + h = F.silu(gate) * up + w_down = self._dequant_batch( + self.down_q, self.down_s, self.down_rot, flat, hidden_states.dtype + ) + out = torch.bmm(h, w_down.transpose(1, 2)).squeeze(1) + del w_down + out = out * top_k_weights.reshape(-1, 1) + return ( + out.view(hidden_states.shape[0], top_k, hidden_dim) + .sum(dim=1) + .to(hidden_states.dtype) + ) + device = hidden_states.device + dtype = hidden_states.dtype + + flat_expert = top_k_index.reshape(-1) # [n_tokens * top_k] + order = flat_expert.argsort() + sorted_expert = flat_expert[order] + token_of_pair = order // top_k + counts = torch.bincount(flat_expert, minlength=self.num_experts) + hit = counts.nonzero().flatten() + hit_counts = counts[hit] + group_size = int(hit_counts.max()) + # rank of each routed pair inside its expert group + group_start = (torch.cumsum(counts, 0) - counts)[sorted_expert] + rank = torch.arange(order.shape[0], device=device) - group_start + slot = torch.searchsorted(hit, sorted_expert) + + padded_x = torch.zeros( + hit.shape[0], group_size, hidden_dim, device=device, dtype=dtype + ) + padded_x[slot, rank] = hidden_states[token_of_pair] + + w_gate_up = self._dequant_batch( + self.gate_up_q, self.gate_up_s, self.gate_up_rot, hit, dtype + ) + gate, up = torch.bmm(padded_x, w_gate_up.transpose(1, 2)).chunk(2, dim=-1) + del w_gate_up + h = F.silu(gate) * up + w_down = self._dequant_batch( + self.down_q, self.down_s, self.down_rot, hit, dtype + ) + out = torch.bmm(h, w_down.transpose(1, 2)) + del w_down + + pair_out = out[slot, rank] * top_k_weights.reshape(-1)[order].unsqueeze(1) + final_hidden_states = torch.zeros_like(hidden_states) + final_hidden_states.index_add_(0, token_of_pair, pair_out.to(dtype)) + return final_hidden_states + + def _forward_dequant(self, hidden_states, top_k_index, top_k_weights): + # mirrors Qwen3OmniMoeThinkerTextExperts.forward with per-expert dequant + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + for expert_idx in expert_hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + w_gate_up = self._dequant( + self.gate_up_q, self.gate_up_s, self.gate_up_rot, expert_idx + ) + gate, up = F.linear(current_state, w_gate_up).chunk(2, dim=-1) + current_hidden_states = F.silu(gate) * up + w_down = self._dequant(self.down_q, self.down_s, self.down_rot, expert_idx) + current_hidden_states = F.linear(current_hidden_states, w_down) + current_hidden_states = ( + current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + ) + final_hidden_states.index_add_( + 0, token_idx, current_hidden_states.to(final_hidden_states.dtype) + ) + + return final_hidden_states + + +def swap_convrot_expert_banks(root, state_dict, dtype): + """Replace each MoE experts module with a ConvRot8Experts holding the + quantized banks from the checkpoint, consuming their state dict entries. + Returns (remaining_state_dict, num_swapped).""" + state_dict = dict(state_dict) + bank_paths = sorted( + { + k[: -len(".gate_up_proj.comfy_quant")] + for k in state_dict + if k.endswith(".gate_up_proj.comfy_quant") and ".experts" in k + } + ) + for experts_path in bank_paths: + tensors = {} + rots = {} + for proj in ("gate_up_proj", "down_proj"): + prefix = f"{experts_path}.{proj}" + conf = parse_comfy_quant_blob(state_dict.pop(f"{prefix}.comfy_quant")) + if conf.get("format") != "int8_tensorwise" or not conf.get("convrot"): + raise ValueError( + f"Expert bank {prefix} has unsupported quant config {conf}" + ) + tensors[proj + "_q"] = state_dict.pop(f"{prefix}.weight") + tensors[proj + "_s"] = state_dict.pop(f"{prefix}.weight_scale") + rots[proj] = int(conf.get("convrot_groupsize", 256)) + + parent_path, _, attr = experts_path.rpartition(".") + parent = root.get_submodule(parent_path) + setattr( + parent, + attr, + ConvRot8Experts( + tensors["gate_up_proj_q"], + tensors["gate_up_proj_s"], + rots["gate_up_proj"], + tensors["down_proj_q"], + tensors["down_proj_s"], + rots["down_proj"], + dtype, + ), + ) + return state_dict, len(bank_paths) + + +class Qwen3OmniCaptioner(BaseCaptioner): + """Captions videos using their audio track via the Qwen3-Omni thinker, + loaded from the pre-quantized convrot8 single-file checkpoint.""" + + def __init__(self, process_id: int, job, config: OrderedDict, **kwargs): + super(Qwen3OmniCaptioner, self).__init__(process_id, job, config, **kwargs) + + def _resolve_checkpoint(self) -> str: + """model_name_or_path can be the checkpoint file itself, a folder + holding it, or a hub repo. Known local spots under MODELS_PATH + (text_encoders/, the root, then any subfolder of text_encoders/) are + searched before downloading; downloads land in + MODELS_PATH/text_encoders.""" + from toolkit.paths import MODELS_PATH + + name_or_path = self.caption_config.model_name_or_path + if os.path.isfile(name_or_path): + return name_or_path + if os.path.isdir(name_or_path): + candidate = os.path.join(name_or_path, CONVROT_FILENAME) + if os.path.exists(candidate): + return candidate + files = [f for f in os.listdir(name_or_path) if f.endswith(".safetensors")] + if len(files) == 1: + return os.path.join(name_or_path, files[0]) + raise FileNotFoundError( + f"No {CONVROT_FILENAME} (or single .safetensors) in {name_or_path}" + ) + + te_dir = os.path.join(MODELS_PATH, "text_encoders") + for candidate in ( + os.path.join(te_dir, CONVROT_FILENAME), + os.path.join(MODELS_PATH, CONVROT_FILENAME), + ): + if os.path.exists(candidate): + return candidate + if os.path.isdir(te_dir): + for dirpath, dirnames, filenames in os.walk(te_dir): + dirnames.sort() + if CONVROT_FILENAME in filenames: + return os.path.join(dirpath, CONVROT_FILENAME) + + import huggingface_hub + + self.print_and_status_update( + f"Downloading {CONVROT_FILENAME} from {name_or_path} into {te_dir}" + ) + return huggingface_hub.hf_hub_download( + repo_id=name_or_path, filename=CONVROT_FILENAME, local_dir=te_dir + ) + + def load_model(self): + from accelerate import init_empty_weights + from safetensors.torch import load_file + + ckpt_path = self._resolve_checkpoint() + self.print_and_status_update("Loading Qwen3-Omni thinker (convrot8)") + + config = AutoConfig.from_pretrained(BASE_REPO) + with init_empty_weights(include_buffers=False): + model = OstrisQwen3OmniThinker(config.thinker_config) + model.eval() + + from transformers.utils import is_flash_attn_2_available + + if is_flash_attn_2_available(): + try: + model.set_attn_implementation("flash_attention_2") + self.print_and_status_update(" - using flash attention 2") + except Exception as e: + print(f"[AITK] flash attention 2 not usable ({e}); staying on sdpa") + + state_dict = load_file(ckpt_path) + + # MoE expert banks stay int8 in ConvRot8Experts modules + state_dict, num_banks = swap_convrot_expert_banks( + model, state_dict, self.torch_dtype + ) + # everything else quantized (attention, vision, audio linears) attaches + # to the toolkit's convrot8 backend in place — no dequantization + state_dict, num_quantized = import_comfy_quantized_layers( + model, state_dict, orig_dtype=self.torch_dtype + ) + self.print_and_status_update( + f" - attached {num_banks} expert banks and {num_quantized} ConvRot layers" + ) + result = model.load_state_dict(state_dict, assign=True, strict=False) + # the importer already attached weights (and popped + assigned biases) + # of quantized layers, so load_state_dict reports them as missing + expected_missing = set() + for name, module in model.named_modules(): + if hasattr(module, "ostris_quantizer"): + expected_missing.add(f"{name}.weight") + expected_missing.add(f"{name}.bias") + bad_missing = [k for k in result.missing_keys if k not in expected_missing] + if bad_missing or result.unexpected_keys: + raise RuntimeError( + f"Checkpoint mismatch. missing: {bad_missing[:8]} " + f"unexpected: {result.unexpected_keys[:8]}" + ) + leftover_meta = [ + n for n, p in model.named_parameters() if p.device.type == "meta" + ] + if leftover_meta: + raise RuntimeError(f"Params never loaded: {leftover_meta[:8]}") + + model.generation_config.pad_token_id = 151643 + model.generation_config.eos_token_id = [151645, 151643] + + # swap the slow bf16 Conv3d patch_embed for an equivalent fast linear + patch_qwen_vl_patch_embed(model) + + if self.caption_config.quantize: + print( + "[AITK] Qwen3-Omni loads pre-quantized (convrot8); the quantize " + "setting is ignored." + ) + + self.model = model + self.model.to(self.device_torch) + self.processor = AutoProcessor.from_pretrained(BASE_REPO) + flush() + + def _build_messages(self, _file_path: str): + return [ + { + "role": "user", + "content": [ + {"type": "video", "video": _file_path}, + {"type": "text", "text": self.caption_config.caption_prompt}, + ], + } + ] + + def _size_kwargs(self): + max_pixels = self.caption_config.max_res * self.caption_config.max_res + # shortest_edge/longest_edge are total pixel counts + # (min_pixels/max_pixels), not edge lengths + return { + "shortest_edge": min(131072, max_pixels), + "longest_edge": max_pixels, + } + + def _prep_media(self, file_path: str): + """CPU side of one video, safe to run in a worker thread: decode + + subsample frames, 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.""" + from transformers.video_utils import load_video + from transformers.audio_utils import load_audio + + frames = load_video(file_path, fps=VIDEO_FPS) + if isinstance(frames, tuple): + frames = frames[0] + audio = None + try: + a = load_audio(file_path, sampling_rate=16000) + if a is not None and a.size > 0: + audio = a + except Exception: + pass + text = self.processor.apply_chat_template( + self._build_messages(file_path), tokenize=False, add_generation_prompt=True + ) + item = {"file": file_path, "frames": frames, "audio": audio, "text": text} + if self.caption_config.batch_size <= 1: + item["inputs"] = self._process_items([item]) + return item + + def _process_items(self, items): + use_audio = items[0]["audio"] is not None + 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, + fps=VIDEO_FPS, + do_sample_frames=False, + size=self._size_kwargs(), + ) + + def _caption_batch(self, items): + """Batched generate over preprocessed items (all with audio, or all + silent). Returns captions in item order.""" + use_audio = items[0]["audio"] is not None + if len(items) == 1 and "inputs" in items[0]: + inputs = items[0]["inputs"] + else: + inputs = self._process_items(items) + inputs = inputs.to(self.device_torch).to(self.torch_dtype) + # 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) + generated_ids = self.model.generate( + **inputs, + use_audio_in_video=use_audio, + max_new_tokens=self.caption_config.max_new_tokens, + ) + trimmed = generated_ids[:, inputs["input_ids"].shape[1] :] + captions = self.processor.batch_decode( + trimmed, 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 videos ahead + of the GPU, videos are grouped (with-audio vs silent) 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 videos + 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): + if len(bucket) == 0: + return + items = list(bucket) + bucket.clear() + try: + captions = self._caption_batch(items) + for it, cap in zip(items, captions): + finish(it["file"], cap) + except Exception as e: + print(f"Batch failed ({e}); retrying files individually") + traceback.print_exc() + for it in items: + 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))) + + with_audio, 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 = with_audio if item["audio"] is not None else silent + bucket.append(item) + if len(bucket) >= batch_size: + flush(bucket) + flush(with_audio) + flush(silent) + finally: + executor.shutdown(wait=False, cancel_futures=True) + pbar.close() + + def maybe_compile_models(self): + """CUDA-graph decode: static kv cache + reduce-overhead compile of the + text model. Each decode step replays as one captured graph, removing + the per-kernel python/launch gaps that cap GPU utilization at small + batch sizes. First video per batch shape is slow (compile warmup).""" + if not self.caption_config.compile: + return + import importlib.util + + if importlib.util.find_spec("triton") is None: + print("[AITK] compile requested but triton is not installed, skipping.") + return + # a static (compileable) cache makes generate auto-compile its decode + # loop (get_compiled_call); prefill stays eager + self.model.generation_config.cache_implementation = "static" + print( + "[AITK] Compiled decode enabled (static cache + cuda graphs). " + "The first video of each batch size will be slow while it compiles." + ) + + def get_caption_for_file(self, file_path: str) -> str: + try: + messages = [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": file_path, + }, + {"type": "text", "text": self.caption_config.caption_prompt}, + ], + } + ] + + max_pixels = self.caption_config.max_res * self.caption_config.max_res + # render the chat text only; the media goes to the processor + # directly so the audio track is interleaved INTO the video block + # (use_audio_in_video) instead of forming a separate audio segment + text = self.processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + # pull the audio track out of the video file; silent videos fall + # back to frames only + from transformers.audio_utils import load_audio + + use_audio = True + try: + audio = load_audio(file_path, sampling_rate=16000) + if audio.size == 0: + use_audio = False + except Exception as audio_err: + print( + f"No audio track for {file_path} ({audio_err}); captioning frames only" + ) + use_audio = False + + inputs = self.processor( + text=text, + audio=[audio] if use_audio else None, + videos=[file_path], + return_tensors="pt", + padding=True, + use_audio_in_video=use_audio, + fps=VIDEO_FPS, + do_sample_frames=True, + # shortest_edge/longest_edge are total pixel counts + # (min_pixels/max_pixels), not edge lengths + size={ + "shortest_edge": min(131072, max_pixels), + "longest_edge": max_pixels, + }, + ) + inputs = inputs.to(self.device_torch).to(self.torch_dtype) + + generated_ids = self.model.generate( + **inputs, + use_audio_in_video=use_audio, + max_new_tokens=self.caption_config.max_new_tokens, + ) + generated_ids_trimmed = [ + out_ids[len(in_ids) :] + for in_ids, out_ids in zip(inputs.input_ids, generated_ids) + ] + output_text = self.processor.batch_decode( + generated_ids_trimmed, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ) + + return output_text[0].strip() + 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 f09cb5d1..30a407dc 100644 --- a/extensions_built_in/captioner/__init__.py +++ b/extensions_built_in/captioner/__init__.py @@ -25,6 +25,18 @@ class Qwen3VLCaptionerExtension(Extension): return Qwen3VLCaptioner +class Qwen3OmniCaptionerExtension(Extension): + uid = "Qwen3OmniCaptioner" + name = "Qwen 3 Omni Captioner" + + @classmethod + def get_process(cls): + # import your process class here so it is only loaded when needed and return it + from .Qwen3OmniCaptioner import Qwen3OmniCaptioner + + return Qwen3OmniCaptioner + + class Ideogram4CaptionerExtension(Extension): uid = "Ideogram4Captioner" name = "Ideogram4 Captioner" @@ -40,5 +52,6 @@ class Ideogram4CaptionerExtension(Extension): AI_TOOLKIT_EXTENSIONS = [ AceStepCaptionerExtension, Qwen3VLCaptionerExtension, + Qwen3OmniCaptionerExtension, Ideogram4CaptionerExtension, ] diff --git a/scripts/convert_vllm_to_comfy.py b/scripts/convert_vllm_to_comfy.py new file mode 100644 index 00000000..85c7c292 --- /dev/null +++ b/scripts/convert_vllm_to_comfy.py @@ -0,0 +1,336 @@ +"""Convert a huggingface LLM / VLM checkpoint into a single ComfyUI-format +safetensors file, quantized with convrot8 (ComfyUI's ``int8_tensorwise`` + +``convrot`` layout, which comfy_kitchen runs natively and is bit-identical to +the toolkit's convrot8 storage). + +ComfyUI ships its text-generation models (the CLIPLoader / TextGenerate path) +as one safetensors file holding the huggingface state dict, with quantized +layers marked two equivalent ways, both of which this script writes: + - a ``.comfy_quant`` uint8 JSON tensor baked into the state dict + (what comfy's _load_quantized_module and the toolkit's + comfy_quant_import read) + - a ``_quantization_metadata`` JSON blob in the safetensors file metadata + (comfy's documented checkpoint format; converted to the markers at load) + +Quantized layer storage (per-output-row symmetric int8 on regular-Hadamard +rotated weights, rotation block = min(256, largest power-of-4 divisor of +in_features)): + .weight int8 [out, in] (or [E, out, in] for MoE banks) + .weight_scale fp32 [out] (or [E, out]) + .comfy_quant {"format": "int8_tensorwise", "convrot": true, + "convrot_groupsize": G} (+ "num_experts": E for banks) + +Per-expert MoE linears (``...experts..gate_proj/up_proj/down_proj.weight``) +are fused into the 3D banks comfy's ops.MoEExperts expects: +``...experts.gate_up_proj.weight`` ([E, 2I, H], gate rows first, matching how +transformers fuses the same checkpoints) and ``...experts.down_proj.weight`` +([E, H, I]). + +Note: the toolkit's own import_comfy_quantized_layers only understands 2D +linear/embedding markers, not the 3D expert banks — the banks are for comfy. + +Usage: + python scripts/convert_vllm_to_comfy.py Qwen/Qwen3-Omni-30B-A3B-Instruct \ + /path/out/qwen3_omni_30b_thinker_convrot8.safetensors + python scripts/convert_vllm_to_comfy.py out.safetensors --no-quant +""" + +import argparse +import json +import os +import re +import sys + +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from tqdm import tqdm + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from toolkit.util.convrot_quant import ( + largest_pow4_divisor, + quantize_int8_rows, + rotate, +) + +# convrot8 eligibility (mirrors ConvRotInt8Quantizer.can_quantize) +MAX_ROT = 256 +MIN_ROT = 16 + +EXPERT_KEY_RE = re.compile( + r"^(?P.*\.experts)\.(?P\d+)\.(?Pgate_proj|up_proj|down_proj)\.weight$" +) + + +def comfy_quant_marker(conf: dict) -> torch.Tensor: + return torch.tensor(list(json.dumps(conf).encode("utf-8")), dtype=torch.uint8) + + +class ArchHandler: + """Generic model: keep every key, no expert fusion. + + Subclasses override to drop components, rename keys, and add + architecture-specific quantization excludes. + """ + + # substrings/suffixes never quantized: output head and embeddings are + # quality-critical, MoE router logits are precision-sensitive + exclude_contains = ("embed_tokens", "pos_embed", "lm_head") + exclude_suffixes = (".mlp.gate.weight",) + + def __init__(self, config: dict): + self.config = config + + def map_key(self, key: str): + """Return the output key, or None to drop the tensor.""" + return key + + def num_experts(self): + for path in ( + ("num_experts",), + ("text_config", "num_experts"), + ("thinker_config", "text_config", "num_experts"), + ): + node = self.config + for p in path: + if not isinstance(node, dict) or p not in node: + node = None + break + node = node[p] + if isinstance(node, int): + return node + return None + + def is_excluded(self, key: str) -> bool: + if any(s in key for s in self.exclude_contains): + return True + return any(key.endswith(s) for s in self.exclude_suffixes) + + +class Qwen3OmniHandler(ArchHandler): + """Qwen3-Omni: keep only the thinker (the VLM — talker and code2wav are + speech synthesis, dead weight for text generation) and strip its prefix + so keys start at model./visual./audio_tower./lm_head. like comfy's other + qwen checkpoints.""" + + def map_key(self, key: str): + if not key.startswith("thinker."): + return None + return key[len("thinker.") :] + + +ARCH_HANDLERS = { + "Qwen3OmniMoeForConditionalGeneration": Qwen3OmniHandler, +} + + +def resolve_model_dir(name_or_path: str) -> str: + if os.path.isdir(name_or_path): + return name_or_path + from huggingface_hub import snapshot_download + + return snapshot_download( + name_or_path, + allow_patterns=["*.safetensors", "model.safetensors.index.json", "config.json"], + ) + + +def shard_files(model_dir: str): + index_path = os.path.join(model_dir, "model.safetensors.index.json") + if os.path.exists(index_path): + with open(index_path, "r") as f: + index = json.load(f) + return sorted(set(index["weight_map"].values())) + single = os.path.join(model_dir, "model.safetensors") + if os.path.exists(single): + return ["model.safetensors"] + raise FileNotFoundError(f"No model.safetensors(.index.json) in {model_dir}") + + +def rot_for(in_features: int) -> int: + return min(MAX_ROT, largest_pow4_divisor(in_features)) + + +def can_quantize_2d(out_features: int, in_features: int) -> bool: + return ( + in_features % 16 == 0 + and out_features % 8 == 0 + and rot_for(in_features) >= MIN_ROT + ) + + +@torch.no_grad() +def quantize_rows(weight: torch.Tensor, rot: int, device): + """convrot8: per-row symmetric int8 on the rotated weight. Accepts + [out, in] or [E, out, in]; scales come back fp32 [out, 1] / [E, out, 1] + (the trailing 1 is comfy_kitchen's per-channel scale convention — its + int8 dequant broadcasts scale directly against [out, in]).""" + shape = weight.shape + w = weight.to(device=device, dtype=torch.float32) + q, scales = quantize_int8_rows(rotate(w, rot).reshape(-1, shape[-1])) + return q.reshape(shape).cpu(), scales.reshape(shape[:-1] + (1,)).cpu() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("model", help="hf repo id or local model directory") + parser.add_argument("output", help="output .safetensors path") + parser.add_argument( + "--no-quant", + action="store_true", + help="just repack to a single file, no quantization", + ) + parser.add_argument( + "--dtype", + default="bf16", + choices=["bf16", "fp16", "fp32"], + help="dtype for non-quantized float tensors", + ) + parser.add_argument( + "--extra-exclude", + default="", + help="comma-separated substrings of keys to keep unquantized", + ) + parser.add_argument( + "--device", + default="cuda" if torch.cuda.is_available() else "cpu", + help="device to run quantization math on", + ) + args = parser.parse_args() + + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[ + args.dtype + ] + extra_exclude = tuple(s for s in args.extra_exclude.split(",") if s.strip()) + + model_dir = resolve_model_dir(args.model) + with open(os.path.join(model_dir, "config.json"), "r") as f: + config = json.load(f) + + handler_cls = ArchHandler + for arch in config.get("architectures", []): + if arch in ARCH_HANDLERS: + handler_cls = ARCH_HANDLERS[arch] + break + handler = handler_cls(config) + print(f"Architecture handler: {handler_cls.__name__}") + + num_experts = handler.num_experts() + quantize = not args.no_quant + + out_sd = {} + quant_layers = {} + # per-(bank, proj) accumulation of expert weights until all E arrive + pending_experts = {} + + def add_quantized(layer: str, weight: torch.Tensor, conf_extra=None): + rot = rot_for(weight.shape[-1]) + q, scales = quantize_rows(weight, rot, args.device) + conf = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": rot} + if conf_extra: + conf.update(conf_extra) + out_sd[f"{layer}.weight"] = q + out_sd[f"{layer}.weight_scale"] = scales + out_sd[f"{layer}.comfy_quant"] = comfy_quant_marker(conf) + quant_layers[layer] = conf + + def flush_bank(bank: str): + """Fuse a completed expert group into comfy MoEExperts banks and + quantize them. gate/up fuse into one [E, 2I, H] bank (gate rows + first), down stays its own [E, H, I] bank.""" + gate = pending_experts.pop((bank, "gate_proj"), None) + up = pending_experts.pop((bank, "up_proj"), None) + down = pending_experts.pop((bank, "down_proj"), None) + if ( + gate is None + or up is None + or down is None + or None in gate + or None in up + or None in down + ): + raise RuntimeError(f"Incomplete expert group for {bank}") + e = len(gate) + gate_up = torch.stack([torch.cat([g, u], dim=0) for g, u in zip(gate, up)]) + down = torch.stack(down) + for name, bankw in (("gate_up_proj", gate_up), ("down_proj", down)): + layer = f"{bank}.{name}" + if quantize and can_quantize_2d(bankw.shape[1], bankw.shape[2]): + add_quantized(layer, bankw, conf_extra={"num_experts": e}) + else: + out_sd[f"{layer}.weight"] = bankw.to(dtype) + + def bank_complete(bank: str) -> bool: + counts = [ + sum(w is not None for w in pending_experts.get((bank, p), [])) + for p in ("gate_proj", "up_proj", "down_proj") + ] + return num_experts is not None and counts == [num_experts] * 3 + + shards = shard_files(model_dir) + for shard in tqdm(shards, desc="Shards"): + with safe_open( + os.path.join(model_dir, shard), framework="pt", device="cpu" + ) as f: + for key in f.keys(): + out_key = handler.map_key(key) + if out_key is None: + continue + tensor = f.get_tensor(key) + + m = EXPERT_KEY_RE.match(out_key) + if m is not None and quantize: + bank, idx, proj = ( + m.group("bank"), + int(m.group("idx")), + m.group("proj"), + ) + slot = pending_experts.setdefault( + (bank, proj), [None] * (num_experts or idx + 1) + ) + if idx >= len(slot): + slot.extend([None] * (idx + 1 - len(slot))) + slot[idx] = tensor + if bank_complete(bank): + flush_bank(bank) + continue + + excluded = handler.is_excluded(out_key) or any( + s in out_key for s in extra_exclude + ) + if ( + quantize + and not excluded + and out_key.endswith(".weight") + and tensor.ndim == 2 + and can_quantize_2d(*tensor.shape) + ): + add_quantized(out_key[: -len(".weight")], tensor) + elif tensor.is_floating_point(): + out_sd[out_key] = tensor.to(dtype) + else: + out_sd[out_key] = tensor + + # experts that never completed during streaming (unknown count or + # shard-straddling groups): fuse whatever is fully populated now + for bank in sorted({b for (b, _p) in pending_experts}): + flush_bank(bank) + + metadata = { + "_quantization_metadata": json.dumps( + {"format_version": "1.0", "layers": quant_layers} + ) + } + total_bytes = sum(t.numel() * t.element_size() for t in out_sd.values()) + print( + f"Saving {len(out_sd)} tensors ({total_bytes / 1e9:.2f} GB, " + f"{len(quant_layers)} quantized layers) to {args.output}" + ) + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + save_file(out_sd, args.output, metadata=metadata) + print("Done") + + +if __name__ == "__main__": + main() diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index 269d4b53..c89eea61 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/ui/src/components/CaptionSimpleJob.tsx @@ -10,6 +10,7 @@ import { import { CaptionJobConfig } from '@/types'; import { handleCaptionerTypeChange } from '@/helpers/captionJobConfig'; import { + batchSizeOptions, captionerTypes, defaultQtype, groupedCaptionerTypes, @@ -165,6 +166,21 @@ const CaptionSimpleJob: React.FC = ({ jobConfig, setJobConfig, gpuIDs, se /> )} + {additionalSections.includes('caption.batch_size') && ( +
+ { + const intVal = parseInt(value); + if (!isNaN(intVal)) { + setJobConfig(intVal, 'config.process[0].caption.batch_size'); + } + }} + options={batchSizeOptions} + /> +
+ )}
diff --git a/ui/src/helpers/captionJobConfig.ts b/ui/src/helpers/captionJobConfig.ts index 19832ea1..0d881fe8 100644 --- a/ui/src/helpers/captionJobConfig.ts +++ b/ui/src/helpers/captionJobConfig.ts @@ -16,7 +16,7 @@ export const defaultCaptionJobConfig: CaptionJobConfig = { model_name_or_path2: "ACE-Step/acestep-captioner", dtype: 'bf16', quantize: true, - qtype: 'float8', + qtype: 'convrot8', low_vram: true, extensions: ['mp3', 'wav', 'flac', 'ogg'], path_to_caption: '', diff --git a/ui/src/helpers/captionOptions.ts b/ui/src/helpers/captionOptions.ts index f0ef61c6..7b0c3964 100644 --- a/ui/src/helpers/captionOptions.ts +++ b/ui/src/helpers/captionOptions.ts @@ -1,7 +1,7 @@ import { GroupedSelectOption, SelectOption } from "@/types"; type CaptionGroup = 'image' | 'music' | 'video'; -type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.thinking'; +type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.thinking' | 'caption.batch_size'; export interface CaptionOption { name: string; @@ -91,6 +91,8 @@ export const captionerTypes: CaptionOption[] = [ 'config.process[0].caption.caption_prompt': [defaultVideoCaptionPrompt, undefined], 'config.process[0].caption.max_res': [512, undefined], 'config.process[0].caption.max_new_tokens': [256, undefined], + 'config.process[0].caption.batch_size': [1, undefined], + 'config.process[0].caption.compile': [true, false], }, name_or_path_options: [ { value: 'ostris/Qwen3-Omni-30B-A3B-Instruct', label: 'ostris/Qwen3-Omni-30B-A3B-Instruct' }, @@ -99,6 +101,7 @@ export const captionerTypes: CaptionOption[] = [ 'caption.caption_prompt', 'caption.max_res', 'caption.max_new_tokens', + 'caption.batch_size', ], }, { @@ -166,6 +169,17 @@ export const quantizationOptions: SelectOption[] = [ { value: 'uint2', label: '2 bit' }, ]; +export const batchSizeOptions: SelectOption[] = [ + { value: '1', label: '1 (default)' }, + { value: '2', label: '2' }, + { value: '4', label: '4' }, + { value: '8', label: '8' }, + { value: '12', label: '12' }, + { value: '16', label: '16' }, + { value: '24', label: '24' }, + { value: '32', label: '32' }, +]; + export const maxResOptions: SelectOption[] = [ { value: '256', label: '256' }, { value: '512', label: '512 (default)' }, @@ -184,4 +198,4 @@ export const maxNewTokensOptions: SelectOption[] = [ { value: '8192', label: '8192' }, ]; -export const defaultQtype = 'float8'; \ No newline at end of file +export const defaultQtype = 'convrot8'; \ No newline at end of file diff --git a/ui/src/types.ts b/ui/src/types.ts index 554dbc08..eb287473 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -298,6 +298,7 @@ export interface CaptionProcessConfig { fixed_caption?: string; caption_extension?: string; thinking?: boolean; + batch_size?: number; } } diff --git a/version.py b/version.py index 26add729..05275330 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.12.15" +VERSION = "0.12.16"