Add support for training VLLMs, omni models, and LLMs

extensions_built_in/llm_models/__init__.py

@@ -0,0 +1,5 @@
+from .qwen25_omni import Qwen25OmniLLM
+
+AI_TOOLKIT_MODELS = [
+    Qwen25OmniLLM,
+]

extensions_built_in/llm_models/qwen25_omni.py

@@ -0,0 +1,465 @@
+"""Qwen2.5-Omni thinker as a trainable text-generation model (arch ``qwen25_omni``).
+
+Input media (audio, image or video file + a .txt caption) is encoded by the
+frozen audio / vision towers into the tower embeddings [N, H] plus the
+processor's expanded placeholder token span and the MRoPE metadata. This runs
+on the GPU per step (a few hundred ms per file); latent caching is optional and
+stores the same tensors (~54 MB per 300 s of audio). Training runs the text stack on
+``prefix + media span + instruction + assistant header + caption`` and applies
+next-token cross-entropy on the caption (LoRA on the text stack; towers and
+lm_head stay frozen). ``is_llm`` routes the trainer to ``train_llm_accumulation``,
+which calls ``get_llm_loss`` directly: no noise, scheduler, VAE or prompt encoding.
+
+Samples: ``ctrl_img`` is the media file (any of the three kinds), the sample
+prompt is the instruction; the generated text is written as a .txt.
+
+model_kwargs:
+  instruction            user-turn text used for every training item (default below)
+  system_prompt          override the chat template's default system message
+  max_pixels/min_pixels  image and video frame budget for the vision tower
+  video_fps              fps the dataset extracted video frames at (dataset fps)
+  max_new_tokens         sample generation budget (default 512)
+  sample_compile         static-cache compiled decode for samples (default false)
+  loss_chunk             positions per lm_head chunk in the CE (default 512)
+"""
+
+import os
+from typing import List, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import AutoConfig, AutoProcessor, Qwen2_5OmniForConditionalGeneration
+
+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.print import print_acc
+
+from .src.thinker import attach_fast_paths, load_thinker_single_file, prepare_thinker
+
+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/text_encoders/qwen2_5_omni_7b_convrot8.safetensors"
+DEFAULT_INSTRUCTION = "Describe this in detail."
+SAMPLE_RATE = 16000
+
+KIND_AUDIO, KIND_IMAGE, KIND_VIDEO = 0, 1, 2
+AUDIO_EXTS = (".mp3", ".wav", ".flac", ".aac", ".ogg", ".m4a")
+VIDEO_EXTS = (".mp4", ".avi", ".mov", ".webm", ".mkv", ".wmv", ".m4v", ".flv")
+
+
+def _media_kind(path: str) -> int:
+    ext = os.path.splitext(path)[1].lower()
+    if ext in AUDIO_EXTS:
+        return KIND_AUDIO
+    if ext in VIDEO_EXTS:
+        return KIND_VIDEO
+    return KIND_IMAGE
+
+
+class Qwen25OmniLLM(BaseModel):
+    arch = "qwen25_omni"
+    is_llm = True
+    # the data loader accepts audio, image and video files in one dataset
+    is_multimodal_llm = True
+    sample_rate = SAMPLE_RATE
+
+    def __init__(self, device, model_config: ModelConfig, dtype="bf16", custom_pipeline=None, noise_scheduler=None, **kwargs):
+        super().__init__(device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs)
+        self.is_transformer = True
+        # LoRA goes on the text stack only (towers are frozen, lm_head stays base)
+        self.target_lora_modules = ["Qwen2_5OmniThinkerTextModel"]
+        kw = self.model_config.model_kwargs
+        self.instruction = str(kw.get("instruction", DEFAULT_INSTRUCTION))
+        self.system_prompt: Optional[str] = kw.get("system_prompt", None)
+        self.max_pixels = int(kw.get("max_pixels", 512 * 512))
+        self.min_pixels = int(kw.get("min_pixels", 64 * 28 * 28))
+        self.video_fps = float(kw.get("video_fps", 24))
+        self.max_new_tokens = int(kw.get("max_new_tokens", 512))
+        self.sample_compile = bool(kw.get("sample_compile", False))
+        self.loss_chunk = int(kw.get("loss_chunk", 512))
+        self.debug = bool(kw.get("debug", False))
+        self.processor = None
+        self.additional_loss_logs = {}
+
+    @staticmethod
+    def get_train_scheduler():
+        # no diffusion: the trainer never samples timesteps for this model
+        return None
+
+    # ------------------------------------------------------------------
+    # loading
+    # ------------------------------------------------------------------
+    def _config_source(self, ckpt: Optional[str] = None) -> str:
+        extras = self.model_config.extras_name_or_path
+        if extras and extras != self.model_config.name_or_path and not extras.endswith(".safetensors"):
+            return extras
+        if ckpt is None:
+            return BASE_REPO
+        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}); set extras_name_or_path to its HF repo"
+            ) from None
+
+    def load_model(self):
+        dtype = self.torch_dtype
+        device = self.device_torch
+        name_or_path = self.model_config.name_or_path or DEFAULT_CHECKPOINT
+        cfg_src = self._config_source()
+        want_quant = self.model_config.quantize
+        qtype = self.model_config.qtype
+        if self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0:
+            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")
+            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)
+            thinker, prequantized = load_thinker_single_file(ckpt, config.thinker_config, dtype)
+            if prequantized:
+                thinker.aitk_is_quantized = True
+                if not want_quant or qtype != "convrot8":
+                    print_acc("Checkpoint is pre-quantized (convrot8); keeping it")
+                want_quant = False
+        else:
+            self.print_and_status_update(f"Loading Qwen2.5-Omni from {name_or_path}")
+            full = Qwen2_5OmniForConditionalGeneration.from_pretrained(name_or_path, dtype=dtype, device_map="cpu")
+            thinker = prepare_thinker(full)
+            del full
+        thinker.eval()
+        thinker.requires_grad_(False)
+        thinker.to(device)
+        if want_quant:
+            from optimum.quanto import freeze
+
+            from toolkit.util.quantize import get_qtype, quantize
+
+            self.print_and_status_update(f"Quantizing thinker ({qtype})")
+            quantize(thinker, weights=get_qtype(qtype))
+            freeze(thinker)
+            thinker.aitk_is_quantized = True
+            flush()
+        self.model = thinker
+        self.processor = AutoProcessor.from_pretrained(cfg_src)
+        self.tokenizer = self.processor.tokenizer
+        self.text_encoder = None
+        self.vae = None
+        self.pipeline = self
+        self._build_template()
+        flush()
+
+    def _build_template(self):
+        """Split the chat template once around the media span and the instruction so
+        training sequences can be assembled from cached media spans without the processor."""
+        mark = "<<AITK_INSTRUCTION>>"
+        msgs = []
+        if self.system_prompt is not None:
+            msgs.append({"role": "system", "content": self.system_prompt})
+        msgs.append({"role": "user", "content": [{"type": "audio", "audio": "x"}, {"type": "text", "text": mark}]})
+        text = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
+        span = "<|audio_bos|><|AUDIO|><|audio_eos|>"
+        if span not in text or mark not in text:
+            raise RuntimeError("Unexpected Qwen2.5-Omni chat template; cannot locate the media span")
+        before, rest = text.split(span, 1)
+        mid, after = rest.split(mark, 1)
+        enc = lambda t: self.tokenizer(t, add_special_tokens=False).input_ids
+        self._prefix_ids = enc(before)
+        self._mid_ids = enc(mid)
+        self._after_ids = enc(after)
+        # probe texts: the processor expands the single placeholder in these
+        self._probe = {
+            KIND_AUDIO: text.replace(mark, "x"),
+            KIND_IMAGE: text.replace(span, "<|vision_bos|><|IMAGE|><|vision_eos|>").replace(mark, "x"),
+            KIND_VIDEO: text.replace(span, "<|vision_bos|><|VIDEO|><|vision_eos|>").replace(mark, "x"),
+        }
+        self._span_tokens = {
+            KIND_AUDIO: (self.tokenizer.convert_tokens_to_ids("<|audio_bos|>"), self.tokenizer.convert_tokens_to_ids("<|audio_eos|>")),
+            KIND_IMAGE: (self.tokenizer.convert_tokens_to_ids("<|vision_bos|>"), self.tokenizer.convert_tokens_to_ids("<|vision_eos|>")),
+            KIND_VIDEO: (self.tokenizer.convert_tokens_to_ids("<|vision_bos|>"), self.tokenizer.convert_tokens_to_ids("<|vision_eos|>")),
+        }
+        self._placeholder = {
+            KIND_AUDIO: self.model.config.audio_token_index,
+            KIND_IMAGE: self.model.config.image_token_index,
+            KIND_VIDEO: self.model.config.video_token_index,
+        }
+        self._im_end = self.tokenizer.convert_tokens_to_ids("<|im_end|>")
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return ["model.layers"]
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    def get_bucket_divisibility(self):
+        return 16
+
+    def save_model(self, output_path, meta, save_dtype):
+        raise NotImplementedError("qwen25_omni: only LoRA training is supported (no full-model save)")
+
+    # LoRA keys are thinker-relative (model.layers.N....), the layout the
+    # ACE-Step captioner and scripts/extract_lora.py use
+    def convert_lora_weights_before_save(self, state_dict):
+        return {(k[len("transformer.") :] if k.startswith("transformer.") else k): v for k, v in state_dict.items()}
+
+    def convert_lora_weights_before_load(self, state_dict):
+        return {(k if k.startswith("transformer.") else "transformer." + k): v for k, v in state_dict.items()}
+
+    # ------------------------------------------------------------------
+    # media encoding (latent cache)
+    # ------------------------------------------------------------------
+    def _span(self, input_ids: torch.Tensor, kind: int) -> torch.Tensor:
+        bos, eos = self._span_tokens[kind]
+        ids = input_ids[0]
+        s = (ids == bos).nonzero()[0].item()
+        e = (ids == eos).nonzero()[-1].item()
+        return ids[s : e + 1].to(torch.int32)
+
+    def _pack(self, embeds, media_ids, kind, audio_seqlen=0, grid=(0, 0, 0), second_per_grid=0.0):
+        dev = embeds.device
+        return DTO(
+            embeds[None].to(self.torch_dtype),
+            media_ids=media_ids[None].to(dev),
+            kind=torch.tensor([kind], dtype=torch.int32, device=dev),
+            audio_seqlen=torch.tensor([int(audio_seqlen)], dtype=torch.int32, device=dev),
+            grid_thw=torch.tensor([list(grid)], dtype=torch.int32, device=dev),
+            second_per_grid=torch.tensor([float(second_per_grid)], dtype=torch.float32, device=dev),
+        )
+
+    @torch.no_grad()
+    def encode_images(self, image_list: torch.Tensor, device=None, dtype=None):
+        """[1, C, S] waveform (16 kHz) | [1, 3, H, W] image in [-1, 1] | [1, F, 3, H, W] frames."""
+        if image_list.shape[0] != 1:
+            raise ValueError("qwen25_omni encodes one item at a time (enable latent caching, batch_size 1)")
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        x = image_list[0]
+        if x.ndim == 2:
+            return self._encode_audio_item(x)
+        if x.ndim == 3:
+            return self._encode_image_item(x)
+        if x.ndim == 4:
+            return self._encode_video_item(x)
+        raise ValueError(f"Unexpected media tensor shape {tuple(image_list.shape)}")
+
+    def encode_audio(self, audio_data_list):
+        # the data loader offers a video's own audio track here; v1 keeps video vision-only
+        return torch.zeros(len(audio_data_list), 1)
+
+    def _log_mel(self, wav: torch.Tensor) -> torch.Tensor:
+        """Whisper log-mel on the GPU, same math as the processor's torch path
+        (n_fft 400, hop 160, 128 mels, log10, max-8 floor, (x+4)/4). [128, frames]"""
+        fe = self.processor.feature_extractor
+        if getattr(self, "_mel_filters", None) is None or self._mel_filters.device != wav.device:
+            self._mel_filters = torch.from_numpy(fe.mel_filters).to(wav.device, torch.float32)
+            self._mel_window = torch.hann_window(fe.n_fft, device=wav.device)
+        stft = torch.stft(wav, fe.n_fft, fe.hop_length, window=self._mel_window, return_complex=True)
+        magnitudes = stft[..., :-1].abs() ** 2
+        log_spec = torch.clamp(self._mel_filters.T @ magnitudes, min=1e-10).log10()
+        log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+        return (log_spec + 4.0) / 4.0
+
+    def _encode_audio_item(self, wav: torch.Tensor):
+        """Stays on the GPU: mel here, tower on the real frames only (the processor
+        pads to 300s and the tower drops the padding anyway)."""
+        fe = self.processor.feature_extractor
+        wav = wav.to(self.device_torch, torch.float32)
+        wav = wav.mean(0) if wav.shape[0] > 1 else wav[0]
+        wav = wav[: fe.n_samples]  # the processor truncates at chunk_length seconds
+        # frame count as the processor's padding mask defines it: ceil(samples / hop)
+        n_frames = -(-wav.shape[0] // fe.hop_length)
+        # the processor zero-pads the clip to chunk_length seconds and the STFT reflects past
+        # that; pad the same way (never beyond n_samples) so the tail frames match it
+        wav = F.pad(wav, (0, min(n_frames * fe.hop_length + fe.n_fft, fe.n_samples) - wav.shape[0]))
+        mel = self._log_mel(wav)[:, :n_frames]
+        feats = self.model.get_audio_features(
+            mel[None].to(self.torch_dtype),
+            feature_attention_mask=torch.ones(1, n_frames, dtype=torch.long, device=self.device_torch),
+        ).last_hidden_state
+        # placeholder count as the processor computes it from the frame count
+        n_tokens = int(((n_frames - 1) // 2 + 1 - 2) // 2 + 1)
+        bos, eos = self._span_tokens[KIND_AUDIO]
+        span = torch.tensor([bos] + [self._placeholder[KIND_AUDIO]] * n_tokens + [eos], dtype=torch.int32)
+        return self._pack(feats, span, KIND_AUDIO, audio_seqlen=n_frames)
+
+    @staticmethod
+    def _to_uint8(img: torch.Tensor) -> np.ndarray:
+        # [-1, 1] float [3, H, W] -> uint8 HWC
+        return ((img.float().clamp(-1, 1) + 1.0) * 127.5).round().to(torch.uint8).permute(1, 2, 0).cpu().numpy()
+
+    def _encode_image_item(self, img: torch.Tensor):
+        from PIL import Image
+
+        pil = Image.fromarray(self._to_uint8(img))
+        inputs = self.processor(text=self._probe[KIND_IMAGE], images=[pil], return_tensors="pt", min_pixels=self.min_pixels, max_pixels=self.max_pixels)
+        grid = inputs["image_grid_thw"]
+        feats = self.model.get_image_features(inputs["pixel_values"].to(self.device_torch, self.torch_dtype), grid.to(self.device_torch), return_dict=True).pooler_output
+        return self._pack(feats, self._span(inputs["input_ids"], KIND_IMAGE), KIND_IMAGE, grid=grid[0].tolist())
+
+    def _encode_video_item(self, frames: torch.Tensor):
+        video = np.stack([self._to_uint8(f) for f in frames])  # [F, H, W, 3]
+        inputs = self.processor(
+            text=self._probe[KIND_VIDEO], videos=[video], return_tensors="pt", fps=self.video_fps,
+            min_pixels=self.min_pixels, max_pixels=self.max_pixels,
+        )
+        grid = inputs["video_grid_thw"]
+        feats = self.model.get_video_features(inputs["pixel_values_videos"].to(self.device_torch, self.torch_dtype), grid.to(self.device_torch), return_dict=True).pooler_output
+        spg = inputs.get("video_second_per_grid", None)
+        spg = float(spg[0]) if spg is not None else 2.0 / self.video_fps
+        return self._pack(feats, self._span(inputs["input_ids"], KIND_VIDEO), KIND_VIDEO, grid=grid[0].tolist(), second_per_grid=spg)
+
+    # ------------------------------------------------------------------
+    # training
+    # ------------------------------------------------------------------
+    def _sequence(self, media_ids: torch.Tensor, caption: str):
+        """ids [L] and the index where the caption (loss span) starts."""
+        instr = self.tokenizer(self.instruction, add_special_tokens=False).input_ids
+        target = self.tokenizer(caption, add_special_tokens=False).input_ids + [self._im_end]
+        head = self._prefix_ids + media_ids.tolist() + self._mid_ids + instr + self._after_ids
+        ids = torch.tensor(head + target, dtype=torch.long, device=self.device_torch)
+        return ids, len(head)
+
+    def _rope_kwargs(self, kind: int, audio_seqlen: int, grid, second_per_grid: float):
+        dev = self.device_torch
+        kw = dict(image_grid_thw=None, video_grid_thw=None, audio_seqlens=None, second_per_grids=None)
+        if kind == KIND_AUDIO:
+            kw["audio_seqlens"] = torch.tensor([audio_seqlen], device=dev)
+        elif kind == KIND_IMAGE:
+            kw["image_grid_thw"] = torch.tensor([list(grid)], device=dev)
+        else:
+            kw["video_grid_thw"] = torch.tensor([list(grid)], device=dev)
+            kw["second_per_grids"] = torch.tensor([second_per_grid], device=dev)
+        return kw
+
+    def _lm_loss(self, ids: torch.Tensor, start: int, media: torch.Tensor, kind: int, rope_kw: dict):
+        model = self.model
+        embeds = model.get_input_embeddings()(ids[None])
+        mask = ids == self._placeholder[kind]
+        n_ph = int(mask.sum())
+        if n_ph != media.shape[0]:
+            raise ValueError(f"media placeholder count {n_ph} != cached embeddings {media.shape[0]}")
+        embeds = embeds.masked_scatter(mask[None, :, None].expand_as(embeds), media.to(embeds.dtype))
+        attention_mask = torch.ones(1, ids.shape[0], dtype=torch.long, device=ids.device)
+        position_ids, _ = model.get_rope_index(
+            ids[None], rope_kw["image_grid_thw"], rope_kw["video_grid_thw"], attention_mask, False,
+            rope_kw["audio_seqlens"], rope_kw["second_per_grids"],
+        )
+        out = model.model(inputs_embeds=embeds, position_ids=position_ids, attention_mask=attention_mask, use_cache=False)
+        hidden = out.last_hidden_state[0, start - 1 : -1]
+        targets = ids[start:]
+        lm_head = model.lm_head
+
+        def chunk_ce(h, t):
+            return F.cross_entropy(lm_head(h).float(), t, reduction="sum")
+
+        total = hidden.new_zeros((), dtype=torch.float32)
+        for s0 in range(0, targets.shape[0], self.loss_chunk):
+            h, t = hidden[s0 : s0 + self.loss_chunk], targets[s0 : s0 + self.loss_chunk]
+            if torch.is_grad_enabled():
+                total = total + torch.utils.checkpoint.checkpoint(chunk_ce, h, t, use_reentrant=False)
+            else:
+                total = total + chunk_ce(h, t)
+        return total / targets.shape[0]
+
+    def get_llm_loss(self, batch) -> torch.Tensor:
+        """Next-token cross-entropy on the caption, from the cached media embeddings.
+        Called by the trainer's train_llm_accumulation; logs ride in additional_loss_logs."""
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        lat = batch.latents
+        if not isinstance(lat, DTO) or lat.get("media_ids") is None:
+            # no latent cache: run the frozen towers on the raw media now
+            if batch.tensor is None:
+                raise ValueError("qwen25_omni got a batch with neither media tensors nor cached embeddings")
+            lat = self.encode_images(batch.tensor)
+        if lat.shape[0] != 1:
+            raise ValueError("qwen25_omni trains with batch_size 1")
+        captions = batch.get_caption_list()
+        media = lat.tensor[0].to(self.device_torch, self.torch_dtype)
+        kind = int(lat.get("kind")[0])
+        rope_kw = self._rope_kwargs(
+            kind, int(lat.get("audio_seqlen")[0]), lat.get("grid_thw")[0].tolist(), float(lat.get("second_per_grid")[0])
+        )
+        ids, start = self._sequence(lat.get("media_ids")[0].long(), captions[0])
+        ce = self._lm_loss(ids, start, media, kind, rope_kw)
+        if self.debug:
+            print_acc(
+                f"qwen25_omni: seq {ids.shape[0]} target {ids.shape[0] - start} ce {ce.item():.4f} "
+                f"peak {torch.cuda.max_memory_allocated() / 1e9:.1f} GB"
+            )
+        self.additional_loss_logs = {"loss/ce": ce.detach().float().item()}
+        return ce
+
+    # ------------------------------------------------------------------
+    # sampling: ctrl_img is the media file, prompt is the instruction, output is text
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return self
+
+    def _load_sample_media(self, path: str):
+        kind = _media_kind(path)
+        if kind == KIND_AUDIO:
+            import torchaudio
+
+            wav, sr = torchaudio.load(path)
+            wav = wav.mean(0, keepdim=True)
+            if sr != SAMPLE_RATE:
+                wav = torchaudio.functional.resample(wav, sr, SAMPLE_RATE)
+            return kind, {"type": "audio", "audio": "x"}, {"audio": [wav[0].numpy()]}
+        if kind == KIND_IMAGE:
+            from PIL import Image
+
+            return kind, {"type": "image", "image": "x"}, {"images": [Image.open(path).convert("RGB")]}
+        from transformers.video_utils import load_video
+
+        frames = load_video(path, fps=self.video_fps)
+        if isinstance(frames, tuple):
+            frames = frames[0]
+        return kind, {"type": "video", "video": "x"}, {"videos": [frames], "fps": self.video_fps}
+
+    def generate_single_image(self, pipeline, gen_config: GenerateImageConfig, conditional_embeds, unconditional_embeds, generator, extra):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        instruction = gen_config.prompt.strip() if gen_config.prompt and gen_config.prompt.strip() else self.instruction
+        content = []
+        media_kwargs = {}
+        if gen_config.ctrl_img:
+            _, entry, media_kwargs = self._load_sample_media(gen_config.ctrl_img)
+            content.append(entry)
+        content.append({"type": "text", "text": instruction})
+        msgs = ([{"role": "system", "content": self.system_prompt}] if self.system_prompt is not None else []) + [{"role": "user", "content": content}]
+        text = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
+        inputs = self.processor(
+            text=text, return_tensors="pt", padding=True, sampling_rate=SAMPLE_RATE,
+            min_pixels=self.min_pixels, max_pixels=self.max_pixels, **media_kwargs,
+        ).to(self.device_torch).to(self.torch_dtype)
+        input_len = inputs["input_ids"].shape[1]
+        self.model._pad_mask_2d = inputs.get("attention_mask", None)
+        gen_kwargs = {"max_new_tokens": self.max_new_tokens, "do_sample": False}
+        if self.sample_compile:
+            from transformers.generation import MaxLengthCriteria, StoppingCriteriaList
+
+            static_len = max(8704, ((input_len + self.max_new_tokens + 255) // 256) * 256)
+            gen_kwargs = {
+                "max_length": static_len, "do_sample": False, "cache_implementation": "static",
+                "stopping_criteria": StoppingCriteriaList([MaxLengthCriteria(max_length=input_len + self.max_new_tokens)]),
+            }
+        with torch.no_grad():
+            ids = self.model.generate(**inputs, **gen_kwargs)
+        # a str return is saved as <sample>.txt by GenerateImageConfig.save_image
+        return self.processor.batch_decode(ids[:, input_len:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0].strip()

extensions_built_in/llm_models/src/thinker.py

@@ -0,0 +1,297 @@
+"""Qwen2.5-Omni thinker with the toolkit's fast paths, shared by the ACE-Step
+captioner and the qwen25_omni trainable model:
+  - ostris_sdpa: grouped single-query decode attention (sdpa's masked kernel is
+    ~1 ms/layer against a static cache)
+  - OstrisQwen25OmniAudioEncoder: per-window batched attention instead of a dense
+    whole-clip block-diagonal mask (fp32-identical, ~150x less attention work)
+  - OstrisQwen25OmniThinker: static-cache-safe MRoPE handling for compiled decode
+"""
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import AttentionInterface, AttentionMaskInterface
+from transformers import Qwen2_5OmniForConditionalGeneration
+from transformers.integrations.sdpa_attention import sdpa_attention_forward
+from transformers.masking_utils import sdpa_mask
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniAudioEncoder,
+    Qwen2_5OmniThinkerForConditionalGeneration,
+)
+
+PAD_TOKEN_ID = 151643
+EOS_TOKEN_IDS = [151645, 151643]
+
+
+def ostris_sdpa_attention_forward(
+    module, query, key, value, attention_mask, dropout=0.0, scaling=None, is_causal=None, **kwargs
+):
+    """sdpa for prefill; for single-query decode against the static cache sdpa's
+    masked (memory-efficient) kernel takes ~1ms per layer, 27ms/token on the 7B.
+    A grouped matmul reads the cache once instead."""
+    if query.shape[2] != 1:
+        return sdpa_attention_forward(
+            module, query, key, value, attention_mask, dropout=dropout, scaling=scaling, is_causal=is_causal, **kwargs
+        )
+    b, hq, _, d = query.shape
+    hk = key.shape[1]
+    g = hq // hk
+    L = key.shape[2]
+    if scaling is None:
+        scaling = d**-0.5
+    # scores land in fp32 straight from the bf16 GEMM: Qwen2.5's sink logits reach ~1e4,
+    # where a bf16 ulp is 64, so a bf16 score tensor rounds the softmax into a coin flip
+    q = query.reshape(b * hk, g, d)
+    kT = key.reshape(b * hk, L, d).transpose(-1, -2)
+    scores = torch.bmm(q, kT, out_dtype=torch.float32).view(b, hk, g, L) * scaling
+    if attention_mask is not None:
+        mask = attention_mask[:, :, -1:, :L]
+        if mask.dtype == torch.bool:
+            scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
+        else:
+            scores = scores + mask.float()
+    probs = torch.softmax(scores, dim=-1).to(query.dtype)
+    out = torch.bmm(probs.view(b * hk, g, L), value.reshape(b * hk, L, d))  # [B*hk, g, d]
+    return out.view(b, hq, 1, d).transpose(1, 2).contiguous(), None
+
+
+AttentionInterface.register("ostris_sdpa", ostris_sdpa_attention_forward)
+AttentionMaskInterface.register("ostris_sdpa", sdpa_mask)
+
+
+class OstrisQwen25OmniAudioEncoder(Qwen2_5OmniAudioEncoder):
+    """Upstream flattens every 2s window of the whole clip into one sequence
+    and attends with a dense [S, S] block-diagonal mask (S = 15000 for 300s):
+    O(S^2) memory and time for attention that is only ever within a window.
+    Batch the windows instead: [num_windows, 100, d] with a key-padding mask
+    for the short tail window. Same math, ~150x less attention work."""
+
+    def forward(self, input_features, feature_lens=None, aftercnn_lens=None, **kwargs):
+        chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
+        chunk_lengths = torch.full(
+            (chunk_num.sum(),), self.n_window * 2, dtype=torch.long, device=feature_lens.device
+        )
+        tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
+        chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
+        chunk_lengths = torch.where(chunk_lengths == 0, self.n_window * 2, chunk_lengths)
+        chunk_list = input_features.split(chunk_lengths.tolist(), dim=1)
+        padded_feature, padded_mask, padded_mask_after_cnn = self.padded_and_mask_function(
+            chunk_list, chunk_lengths, padding_value=0, padding_side="right"
+        )
+        x = F.gelu(self.conv1(padded_feature)) * padded_mask
+        x = F.gelu(self.conv2(x)).transpose(1, 2)
+        x = x + self.positional_embedding.positional_embedding[: x.shape[1], :].unsqueeze(0).to(x.dtype)
+        key_mask = None
+        if not bool(padded_mask_after_cnn.all()):
+            key_mask = padded_mask_after_cnn[:, None, None, :]
+        n, seq_len, dim = x.shape
+        for layer in self.layers:
+            attn = layer.self_attn
+            residual = x
+            h = layer.self_attn_layer_norm(x)
+            q = attn.q_proj(h).view(n, seq_len, attn.num_heads, -1).transpose(1, 2)
+            k = attn.k_proj(h).view(n, seq_len, attn.num_heads, -1).transpose(1, 2)
+            v = attn.v_proj(h).view(n, seq_len, attn.num_heads, -1).transpose(1, 2)
+            o = F.scaled_dot_product_attention(q, k, v, attn_mask=key_mask, scale=attn.scaling)
+            x = residual + attn.out_proj(o.transpose(1, 2).reshape(n, seq_len, dim))
+            residual = x
+            h = layer.final_layer_norm(x)
+            x = residual + layer.fc2(layer.activation_fn(layer.fc1(h)))
+        hidden_states = x[padded_mask_after_cnn]
+        outs = []
+        for each in hidden_states.split(aftercnn_lens.tolist(), dim=0):
+            each = self.avg_pooler(each.transpose(0, 1)).transpose_(0, 1)
+            outs.append(self.proj(self.ln_post(each)))
+        return BaseModelOutputWithPooling(last_hidden_state=torch.cat(outs, dim=0))
+
+
+class OstrisQwen25OmniThinker(Qwen2_5OmniThinkerForConditionalGeneration):
+    """Thinker with static-cache-safe MRoPE handling (same fix as the Qwen3-Omni
+    captioner): under a compileable cache generate hands forward a 4D mask, but
+    upstream's rope block does `1 - attention_mask` and get_rope_index needs the
+    2D padding mask. Prefill uses the true 2D mask stashed by the caller; decode
+    continues from the cache length with no data-dependent ops."""
+
+    _pad_mask_2d = None
+
+    # per-file media shapes must not reach the compiled decode graph or it
+    # recompiles for every file
+    _PREFILL_ONLY_KEYS = (
+        "input_features",
+        "feature_attention_mask",
+        "audio_feature_lengths",
+        "pixel_values",
+        "pixel_values_videos",
+        "image_grid_thw",
+        "video_grid_thw",
+        "video_second_per_grid",
+    )
+
+    def enable_gradient_checkpointing(self):
+        """Trainer hook. transformers gates its own layer checkpointing on .training,
+        but the toolkit trains the frozen base in eval mode with adapters on top, so
+        wrap the text-stack decoder layers to checkpoint whenever autograd is recording."""
+        for layer in self.model.layers:
+            if getattr(layer, "_aitk_checkpointed", False):
+                continue
+            orig_forward = layer.forward
+
+            def forward(*args, _orig=orig_forward, **kwargs):
+                if torch.is_grad_enabled():
+                    return torch.utils.checkpoint.checkpoint(_orig, *args, use_reentrant=False, **kwargs)
+                return _orig(*args, **kwargs)
+
+            layer.forward = forward
+            layer._aitk_checkpointed = True
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        model_inputs = super().prepare_inputs_for_generation(*args, **kwargs)
+        ids = model_inputs.get("input_ids", None)
+        if ids is not None and ids.shape[1] == 1:
+            for key in self._PREFILL_ONLY_KEYS:
+                model_inputs.pop(key, None)
+        return model_inputs
+
+    def forward(
+        self,
+        input_ids=None,
+        input_features=None,
+        pixel_values=None,
+        pixel_values_videos=None,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        attention_mask=None,
+        feature_attention_mask=None,
+        audio_feature_lengths=None,
+        position_ids=None,
+        past_key_values=None,
+        use_audio_in_video=None,
+        video_second_per_grid=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:
+                mask2d = (
+                    attention_mask
+                    if isinstance(attention_mask, torch.Tensor) 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]:
+                    mask2d = mask2d[:, : input_ids.shape[1]]
+                audio_lens = (
+                    feature_attention_mask.sum(1)
+                    if feature_attention_mask is not None
+                    else audio_feature_lengths
+                )
+                delta0 = (1 - mask2d).sum(dim=-1).unsqueeze(1)
+                position_ids, rope_deltas = self.get_rope_index(
+                    input_ids,
+                    image_grid_thw,
+                    video_grid_thw,
+                    mask2d,
+                    use_audio_in_video or False,
+                    audio_lens,
+                    video_second_per_grid,
+                )
+                self.rope_deltas = rope_deltas - delta0
+            else:
+                batch_size, seq_length = input_ids.shape
+                # static cache length is a device tensor: stays sync-free
+                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
+                    + self.rope_deltas.to(input_ids.device)
+                )
+                position_ids = pos.unsqueeze(0).expand(3, batch_size, seq_length)
+        return super().forward(
+            input_ids=input_ids,
+            input_features=input_features,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            attention_mask=attention_mask,
+            feature_attention_mask=feature_attention_mask,
+            audio_feature_lengths=audio_feature_lengths,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_audio_in_video=use_audio_in_video,
+            video_second_per_grid=video_second_per_grid,
+            **kwargs,
+        )
+
+
+def prepare_thinker(full: Qwen2_5OmniForConditionalGeneration):
+    """Strip a loaded Qwen2.5-Omni to its thinker with the fast paths attached."""
+    # drop the talker before anything touches the GPU
+    full.disable_talker()
+    model = full.thinker
+    model.__class__ = OstrisQwen25OmniThinker
+    model.audio_tower.__class__ = OstrisQwen25OmniAudioEncoder
+    model.model.config._attn_implementation = "ostris_sdpa"
+    model.generation_config.eos_token_id = full.generation_config.eos_token_id
+    model.generation_config.pad_token_id = (
+        full.generation_config.pad_token_id
+        if full.generation_config.pad_token_id is not None
+        else 151643
+    )
+    return model
+
+
+def attach_fast_paths(thinker):
+    """Swap in the fast subclasses on an already-built thinker (any load path)."""
+    thinker.__class__ = OstrisQwen25OmniThinker
+    thinker.audio_tower.__class__ = OstrisQwen25OmniAudioEncoder
+    thinker.model.config._attn_implementation = "ostris_sdpa"
+    if thinker.generation_config.eos_token_id is None:
+        thinker.generation_config.eos_token_id = list(EOS_TOKEN_IDS)
+    if thinker.generation_config.pad_token_id is None:
+        thinker.generation_config.pad_token_id = PAD_TOKEN_ID
+    return thinker
+
+
+def load_thinker_single_file(ckpt_path: str, thinker_config, dtype: torch.dtype):
+    """Build a thinker from a single-file checkpoint in the comfy layout written
+    by scripts/convert_vllm_to_comfy.py (bf16 or convrot8-quantized). Returns
+    (thinker, is_prequantized)."""
+    from accelerate import init_empty_weights
+    from safetensors.torch import load_file
+
+    from toolkit.util.comfy_quant_import import import_comfy_quantized_layers
+
+    with init_empty_weights(include_buffers=False):
+        thinker = OstrisQwen25OmniThinker(thinker_config)
+    thinker.eval()
+    sd = load_file(ckpt_path)
+    prequantized = any(k.endswith(".comfy_quant") for k in sd)
+    if prequantized:
+        sd, _ = import_comfy_quantized_layers(thinker, sd, orig_dtype=dtype)
+    sd = {k: (v.to(dtype) if v.is_floating_point() else v) for k, v in sd.items()}
+    result = thinker.load_state_dict(sd, assign=True, strict=False)
+    allowed = set()
+    for name, module in thinker.named_modules():
+        if hasattr(module, "ostris_quantizer"):
+            allowed.add(f"{name}.weight")
+            allowed.add(f"{name}.bias")
+    missing = [k for k in result.missing_keys if k not in allowed]
+    if missing or result.unexpected_keys:
+        raise RuntimeError(
+            f"Qwen2.5-Omni checkpoint mismatch. missing: {missing[:8]} unexpected: {result.unexpected_keys[:8]}"
+        )
+    leftover = [n for n, p in thinker.named_parameters() if p.device.type == "meta"]
+    if leftover:
+        raise RuntimeError(f"Params never loaded: {leftover[:8]}")
+    # quantized layers own their buffers; cast only the plain float tensors
+    for module in thinker.modules():
+        if hasattr(module, "ostris_quantizer"):
+            continue
+        for pname, prm in list(module.named_parameters(recurse=False)):
+            if prm.is_floating_point() and prm.dtype != dtype:
+                prm.data = prm.data.to(dtype)
+    attach_fast_paths(thinker)
+    return thinker, prequantized

extensions_built_in/sd_trainer/SDTrainer.py

@@ -298,15 +298,16 @@ class SDTrainer(BaseSDTrainProcess):
             # make sure model is on cpu for this part so we don't oom.
             self.sd.unet.to('cpu')
         
-        # cache unconditional embeds (blank prompt)
-        with torch.no_grad():
-            self.unconditional_embeds = self.encode_static_prompt(
-                [self.train_config.unconditional_prompt],
-                long_prompts=self.do_long_prompts,
-            ).to(
-                self.device_torch,
-                dtype=self.sd.torch_dtype
-            ).detach()
+        # cache unconditional embeds (blank prompt); text-generating models have no text encoder
+        if not getattr(self.sd, 'is_llm', False):
+            with torch.no_grad():
+                self.unconditional_embeds = self.encode_static_prompt(
+                    [self.train_config.unconditional_prompt],
+                    long_prompts=self.do_long_prompts,
+                ).to(
+                    self.device_torch,
+                    dtype=self.sd.torch_dtype
+                ).detach()
         
         if self.train_config.do_prior_divergence:
             self.do_prior_prediction = True
@@ -314,14 +315,16 @@ class SDTrainer(BaseSDTrainProcess):
             # D-OPSD: the teacher (prior) prediction is the training target
             self.do_prior_prediction = True
         # move vae to device if we did not cache latents
-        if not self.is_latents_cached:
-            self.sd.vae.eval()
-            self.sd.vae.to(self.device_torch)
-        else:
-            # offload it. Already cached
-            self.sd.vae.to('cpu')
-            flush()
-        add_all_snr_to_noise_scheduler(self.sd.noise_scheduler, self.device_torch)
+        if self.sd.vae is not None:
+            if not self.is_latents_cached:
+                self.sd.vae.eval()
+                self.sd.vae.to(self.device_torch)
+            else:
+                # offload it. Already cached
+                self.sd.vae.to('cpu')
+                flush()
+        if self.sd.noise_scheduler is not None:
+            add_all_snr_to_noise_scheduler(self.sd.noise_scheduler, self.device_torch)
         if self.adapter is not None:
             self.adapter.to(self.device_torch)
 
@@ -1418,10 +1421,33 @@ class SDTrainer(BaseSDTrainProcess):
         )
     
 
+    def train_llm_accumulation(self, batch: DataLoaderBatchDTO, accum_scale: float = 1.0):
+        """Text-generating models (BaseModel.is_llm): no noise, scheduler, VAE or prompt
+        encoding. The model computes its own loss from the batch (cached media + captions)
+        and reports per-term logs through additional_loss_logs."""
+        network = self.network if self.network is not None else BlankNetwork()
+        network.multiplier = batch.get_network_weight_list()
+        with torch.no_grad():
+            loss_multiplier = torch.tensor(batch.loss_multiplier_list).to(self.device_torch, dtype=torch.float32)
+        with network:
+            with self.timer('llm_loss'):
+                loss = self.sd.get_llm_loss(batch)
+            self.additional_logs.update(getattr(self.sd, "additional_loss_logs", None) or {})
+            if not torch.isfinite(loss):
+                print_acc("loss is nan")
+                loss = torch.zeros_like(loss).requires_grad_(True)
+            with self.timer('backward'):
+                loss = loss * loss_multiplier.mean()
+                # backward stays inside the network context (see the note in the diffusion path)
+                self.accelerator.backward(loss * accum_scale if accum_scale != 1.0 else loss)
+        return loss.detach()
+
     def train_single_accumulation(self, batch: DataLoaderBatchDTO, accum_scale: float = 1.0):
         # accum_scale: 1 / number of micro-batches accumulated per optimizer step, so the
         # summed gradients equal the mean over the effective batch. Applied to the backward
         # only; the returned loss stays unscaled for logging.
+        if getattr(self.sd, 'is_llm', False):
+            return self.train_llm_accumulation(batch, accum_scale=accum_scale)
         with torch.no_grad():
             self.timer.start('preprocess_batch')
             if isinstance(self.adapter, CustomAdapter):