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,7 +298,8 @@ 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)
+        # 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],
@@ -314,6 +315,7 @@ 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 self.sd.vae is not None:
             if not self.is_latents_cached:
                 self.sd.vae.eval()
                 self.sd.vae.to(self.device_torch)
@@ -321,6 +323,7 @@ class SDTrainer(BaseSDTrainProcess):
                 # 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):

jobs/process/BaseSDTrainProcess.py

@@ -747,6 +747,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
         self.accelerator.even_batches=False
         
         # # prepare all the models stuff for accelerator (hopefully we dont miss any)
+        if self.sd.vae is not None:
             self.sd.vae = self.accelerator.prepare(self.sd.vae)
         if self.sd.unet is not None:
             self.sd.unet = self.accelerator.prepare(self.sd.unet)
@@ -1099,7 +1100,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 is_reg = any(batch.get_is_reg_list())
                 if batch.tensor is not None:
                     imgs = batch.tensor
-                    imgs = imgs.to(self.device_torch, dtype=dtype)
+                    # waveforms stay fp32 into the audio encoder
+                    imgs = imgs.to(self.device_torch, dtype=torch.float32 if getattr(self.sd, 'is_audio_model', False) else dtype)
                     # dont adjust for regs.
                     if self.train_config.img_multiplier is not None and not is_reg:
                         # do it ad contrast
@@ -1792,6 +1794,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
         noise_scheduler = self.sd.noise_scheduler
 
         if self.train_config.xformers:
+            if vae is not None:
                 vae.enable_xformers_memory_efficient_attention()
             unet.enable_xformers_memory_efficient_attention()
             if isinstance(text_encoder, list):
@@ -1868,12 +1871,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
             for te in text_encoder:
                 te.requires_grad_(False)
                 te.eval()
-        else:
+        elif text_encoder is not None:
             text_encoder.requires_grad_(False)
             text_encoder.eval()
         unet.to(self.device_torch, dtype=dtype)
         unet.requires_grad_(False)
         unet.eval()
+        if vae is not None:
             vae = vae.to(torch.device('cpu'), dtype=dtype)
             vae.requires_grad_(False)
             vae.eval()

toolkit/config_modules.py

@@ -1325,7 +1325,11 @@ class GenerateImageConfig:
         # make parent dirs
         os.makedirs(self.output_folder, exist_ok=True)
         self.set_gen_time()
-        if isinstance(image, list):
+        if isinstance(image, str):
+            # text-generating models: the sample is the text itself
+            with open(self.get_prompt_path(count, max_count), 'w', encoding='utf-8') as f:
+                f.write(image)
+        elif isinstance(image, list):
             # video
             if self.num_frames == 1:
                 raise ValueError(f"Expected 1 img but got a list {len(image)}")
@@ -1481,7 +1485,7 @@ class GenerateImageConfig:
         pass
     
     def log_image(self, image, count: int = 0, max_count=0):
-        if self.logger is None:
+        if self.logger is None or isinstance(image, str):
             return
 
         self.logger.log_image(image, count, self.prompt)

toolkit/data_loader.py

@@ -395,6 +395,8 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
         self.dataset_config.bucket_tolerance = sd.get_bucket_divisibility()
         self.is_video = dataset_config.num_frames > 1 or dataset_config.auto_frame_count
         self.is_audio_model = hasattr(sd, 'is_audio_model') and sd.is_audio_model if sd is not None else False
+        # text-generating multimodal models take audio, image and video files in one dataset
+        self.is_multimodal_llm = getattr(sd, 'is_multimodal_llm', False) if sd is not None else False
         super().__init__()
         folder_path = dataset_config.folder_path
         self.dataset_path = dataset_config.dataset_path
@@ -430,6 +432,8 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
             if self.is_audio_model:
                 # only look for audio files
                 extensions = audio_extensions
+            elif self.is_multimodal_llm:
+                extensions = audio_extensions + image_extensions + (video_extensions if self.is_video else [])
             elif self.is_video:
                 # look for videos and images. Video models can train on both;
                 # images are bucketed separately as single-frame items
@@ -532,7 +536,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                 file_item = FileItemDTO(
                     sd=self.sd,
                     path=file,
-                    is_audio_model=self.is_audio_model,
+                    is_audio_model=self.is_audio_model or (self.is_multimodal_llm and file.lower().endswith(tuple(audio_extensions))),
                     dataset_config=dataset_config,
                     dataloader_transforms=self.transform,
                     size_database=self.size_database,
@@ -544,7 +548,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
                     te_padding_side=self.sd.te_padding_side if self.sd else "right",
                     latent_space_version=latent_space_version,
                     temporal_compression=temporal_compression,
-                    sample_rate=self.sd.sample_rate if self.is_audio_model and self.sd is not None else 48000,
+                    sample_rate=self.sd.sample_rate if (self.is_audio_model or self.is_multimodal_llm) and self.sd is not None else 48000,
                 )
                 self.file_list.append(file_item)
             except Exception as e:

toolkit/dataloader_mixins.py

@@ -234,7 +234,7 @@ class BucketsMixin:
         # for file_item in enumerate(file_list):
         for idx, file_item in enumerate(file_list):
             file_item: 'FileItemDTO' = file_item
-            if self.is_audio_model:
+            if file_item.is_audio_model:
                 bucket_key = f"{file_item.width}ms"
                 if bucket_key not in self.buckets:
                     self.buckets[bucket_key] = Bucket(file_item.width, 1)
@@ -423,11 +423,20 @@ class CaptionProcessingDTOMixin:
                 # drop the caption
                 return ''
 
-        # get tokens
+        # tag-style token ops (dropout / shuffle) split on commas; anything else keeps the
-        token_list = raw_caption.split(',')
+        # text exactly as written (splitting and re-joining with ", " doubled the space after
+        # every comma in prose captions)
+        do_token_dropout = (
+            self.dataset_config.token_dropout_rate > 0
+            and not short_caption
+            and not self.dataset_config.cache_text_embeddings
+        )
+        token_list = None
+        if do_token_dropout or self.dataset_config.shuffle_tokens:
+            token_list = [t.strip() for t in raw_caption.split(',')]
 
         # handle token dropout
-        if self.dataset_config.token_dropout_rate > 0 and not short_caption and not self.dataset_config.cache_text_embeddings:
+        if do_token_dropout:
             new_token_list = []
             keep_tokens: int = self.dataset_config.keep_tokens
             for idx, token in enumerate(token_list):
@@ -448,7 +457,7 @@ class CaptionProcessingDTOMixin:
             random.shuffle(token_list)
 
         # join back together
-        caption = ', '.join(token_list)
+        caption = ', '.join(token_list) if token_list is not None else raw_caption
         caption = inject_trigger_into_prompt(caption, trigger, to_replace_list, add_if_not_present)
 
         if self.dataset_config.random_triggers:
@@ -467,7 +476,7 @@ class CaptionProcessingDTOMixin:
 
         if self.dataset_config.shuffle_tokens:
             # shuffle again
-            token_list = caption.split(',')
+            token_list = [t.strip() for t in caption.split(',')]
             random.shuffle(token_list)
             caption = ', '.join(token_list)
         if caption == '':
@@ -2075,7 +2084,7 @@ class LatentCachingMixin:
             # add batch dimension
             cache_uint8 = getattr(self.sd, 'cache_latents_as_uint8', False)
             if self.dataset_config.cache_tensors_to_disk:
-                if not self.is_audio_model:
+                if not file_item.is_audio_model:
                     tensor_uint8 = _latent_to_uint8(file_item.tensor).cpu()
                     if to_disk:
                         state_dict['tensor'] = tensor_uint8
@@ -2092,7 +2101,8 @@ class LatentCachingMixin:
                         file_item._cached_waveform_int16 = waveform_int16
                         file_item._cached_waveform_sample_rate = sample_rate
             try:
-                imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
+                # waveforms stay fp32 into the audio encoder (bf16 mantissa is ~48 dB of noise)
+                imgs = file_item.tensor.unsqueeze(0).to(device, dtype=torch.float32 if file_item.is_audio_model else dtype)
                 latent = self.sd.encode_images(imgs)
                 # a model can return a DTO carrying extra streams alongside the latent
                 latent = latent.map(lambda t: t.squeeze(0)) if isinstance(latent, DTO) else latent.squeeze(0)
@@ -2128,7 +2138,7 @@ class LatentCachingMixin:
                         state_dict['first_frame_latent'] = first_frame_latent.clone().detach().cpu()
 
             # audio (video+audio models only — audio-only models already encoded above via encode_images)
-            if not self.is_audio_model and file_item.audio_data is not None:
+            if not file_item.is_audio_model and file_item.audio_data is not None:
                 audio_latent = self.sd.encode_audio([file_item.audio_data]).squeeze(0)
                 if to_disk:
                     state_dict['audio_latent'] = audio_latent.clone().detach().cpu()

toolkit/models/base_model.py

@@ -101,6 +101,9 @@ class BaseModel:
     # rename LoRA keys transformer. <-> diffusion_model. (the ComfyUI-standard
     # prefix) on save/load
     lora_keys_use_comfy_prefix = False
+    # text-generating models: the trainer runs train_llm_accumulation (model-owned
+    # loss via get_llm_loss) instead of the diffusion step; no vae / text encoder
+    is_llm = False
 
     def __init__(
             self,
@@ -568,7 +571,11 @@ class BaseModel:
                             quad_count=4
                         )
 
-                    if self.sample_prompts_cache is not None:
+                    if self.is_llm:
+                        # text-generating models take the prompt and media directly
+                        conditional_embeds = None
+                        unconditional_embeds = None
+                    elif self.sample_prompts_cache is not None:
                         conditional_embeds = self.sample_prompts_cache[i]['conditional'].to(self.device_torch, dtype=self.torch_dtype)
                         unconditional_embeds = self.sample_prompts_cache[i]['unconditional'].to(self.device_torch, dtype=self.torch_dtype)
                     else:
@@ -721,8 +728,10 @@ class BaseModel:
                             raise ValueError(
                                 "Refiner is only supported for XL models")
 
+                    if conditional_embeds is not None:
                         conditional_embeds = conditional_embeds.to(
                             self.device_torch, dtype=self.unet.dtype)
+                    if unconditional_embeds is not None:
                         unconditional_embeds = unconditional_embeds.to(
                             self.device_torch, dtype=self.unet.dtype)
 
@@ -1451,7 +1460,7 @@ class BaseModel:
             'vae': {
                 'training': self.vae.training,
                 'device': self.vae.device,
-            },
+            } if self.vae is not None else None,
             'unet': {
                 'training': self.unet.training,
                 'device': self.unet.device,
@@ -1468,7 +1477,7 @@ class BaseModel:
                     # todo there has to be a better way to do this
                     'requires_grad': te_has_grad
                 })
-        else:
+        elif self.text_encoder is not None:
             te_has_grad = self.get_te_has_grad()
 
             self.device_state['text_encoder'] = {
@@ -1520,6 +1529,7 @@ class BaseModel:
         self.device_state = None
 
     def set_device_state(self, state):
+        if self.vae is not None and state.get('vae') is not None:
             if state['vae']['training']:
                 self.vae.train()
             else:
@@ -1552,7 +1562,7 @@ class BaseModel:
                     encoder.to(state['text_encoder']['device'])
                     encoder.requires_grad_(
                         state['text_encoder']['requires_grad'])
-        else:
+        elif self.text_encoder is not None:
             if state['text_encoder']['training']:
                 self.text_encoder.train()
             else:
@@ -1646,7 +1656,7 @@ class BaseModel:
         if isinstance(self.text_encoder, list):
             for encoder in self.text_encoder:
                 encoder.to(*args, **kwargs)
-        else:
+        elif self.text_encoder is not None:
             self.text_encoder.to(*args, **kwargs)
 
     def component_load_kwargs(self, role: str = "transformer", dtype=None):

toolkit/models/registry.py

@@ -14,7 +14,7 @@ VID = {"width": 832, "height": 480, "num_inference_steps": 20, "guidance_scale":
 
 GENERIC = {"modality": "image", "model": {}, "sample": dict(IMG)}
 
-# modality: image | video | audio
+# modality: image | video | audio | text
 # needs_control_image: the arch requires ctrl_img (edit / i2v models)
 # size_locked: only the default resolution is valid
 ARCH_REGISTRY: Dict[str, dict] = {
@@ -55,6 +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/text_encoders/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}},
 }
 

ui/src/app/api/jobs/[jobID]/samples/route.ts

@@ -28,11 +28,22 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
   // find all img (png, jpg, jpeg) files in the samples folder. Thumbnails
   // live in the hidden .thumbs subfolder (and partial writes in .tmp) — the
   // isFile check keeps those directories out even if their names ever match.
-  const samples = (await fs.promises.readdir(samplesFolder, { withFileTypes: true }))
+  const names = (await fs.promises.readdir(samplesFolder, { withFileTypes: true }))
     .filter(entry => entry.isFile())
-    .map(entry => entry.name)
+    .map(entry => entry.name);
+  const isMedia = (file: string) =>
+    file.endsWith('.png') ||
+    file.endsWith('.jpg') ||
+    file.endsWith('.jpeg') ||
+    file.endsWith('.webp') ||
+    file.endsWith('.mp4') ||
+    file.endsWith('mp3');
+  const mediaStems = new Set(names.filter(isMedia).map(file => file.slice(0, file.lastIndexOf('.'))));
+  const samples = names
     .filter(file => {
-      return file.endsWith('.png') || file.endsWith('.jpg') || file.endsWith('.jpeg') || file.endsWith('.webp') || file.endsWith('.mp4') || file.endsWith('mp3') || file.endsWith('wav') || file.endsWith('flac') || file.endsWith('ogg');
+      if (isMedia(file)) return true;
+      // text samples (LLM models); a .txt next to a media sample is its prompt sidecar, not a sample
+      return file.endsWith('.txt') && !mediaStems.has(file.slice(0, -4));
     })
     .map(file => {
       return path.join(samplesFolder, file);

ui/src/app/generate/page.tsx

@@ -3,10 +3,30 @@
 import React, { Suspense, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useSearchParams } from 'next/navigation';
 import { Button } from '@headlessui/react';
-import { ChevronDown, ChevronLeft, ChevronRight, Loader2, OctagonX, Play, Plus, Square, Sparkles, Trash2, X } from 'lucide-react';
+import {
+  ChevronDown,
+  ChevronLeft,
+  ChevronRight,
+  Loader2,
+  OctagonX,
+  Play,
+  Plus,
+  Square,
+  Sparkles,
+  Trash2,
+  X,
+} from 'lucide-react';
 import { openConfirm } from '@/components/ConfirmModal';
 import { TopBar, MainContent } from '@/components/layout';
-import { Checkbox, CreatableSelectInput, NumberInput, SelectInput, SliderInput, TextAreaInput, TextInput } from '@/components/formInputs';
+import {
+  Checkbox,
+  CreatableSelectInput,
+  NumberInput,
+  SelectInput,
+  SliderInput,
+  TextAreaInput,
+  TextInput,
+} from '@/components/formInputs';
 import { apiClient } from '@/utils/api';
 import { startJob, stopJob } from '@/utils/jobs';
 import { startQueue } from '@/utils/queue';
@@ -20,7 +40,6 @@ import { modelArchs, getGenerateDefaults, GenerateDefaults } from '@/app/jobs/ne
 import GenerateFooter from '@/components/generate/GenerateFooter';
 import LoraBrowserModal, { LoraPick } from '@/components/generate/LoraBrowserModal';
 
-
 interface EngineStatus {
   running: boolean;
   engine: { jobId: string; jobName: string; gpuIds: string; status: string; ready: boolean } | null;
@@ -90,7 +109,11 @@ function Card({
 }) {
   return (
     <div className="bg-gray-900 rounded-xl border border-gray-800">
-      <button type="button" onClick={onToggle} className="w-full flex items-center justify-between gap-2 px-3 py-1.5 text-xs text-gray-200 font-semibold">
+      <button
+        type="button"
+        onClick={onToggle}
+        className="w-full flex items-center justify-between gap-2 px-3 py-1.5 text-xs text-gray-200 font-semibold"
+      >
         <span className="flex items-baseline gap-2 min-w-0">
           <span>{title}</span>
           {subtitle && <span className="text-[11px] font-normal text-gray-500 truncate">{subtitle}</span>}
@@ -124,7 +147,11 @@ function GeneratePageInner() {
     return () => clearInterval(t);
   }, []);
   const engineJobId = preferredJob || engineStatus?.engine?.jobId || null;
-  const proxy = useCallback((path: string) => `/api/inference/${path}${engineJobId ? `${path.includes('?') ? '&' : '?'}job=${engineJobId}` : ''}`, [engineJobId]);
+  const proxy = useCallback(
+    (path: string) =>
+      `/api/inference/${path}${engineJobId ? `${path.includes('?') ? '&' : '?'}job=${engineJobId}` : ''}`,
+    [engineJobId],
+  );
 
   usePollLoop(
     () =>
@@ -139,7 +166,10 @@ function GeneratePageInner() {
               setStartingJobId(null);
             } else if (!mine) {
               // no longer queued/running: the launch failed (or was stopped)
-              const job = await apiClient.get('/api/jobs', { params: { id: startingJobId } }).then(r => r.data).catch(() => null);
+              const job = await apiClient
+                .get('/api/jobs', { params: { id: startingJobId } })
+                .then(r => r.data)
+                .catch(() => null);
               if (!job || ['error', 'stopped', 'completed'].includes(job.status)) {
                 setStartingJobId(null);
                 alert(`The engine did not start${job?.info ? `: ${job.info}` : ''}. Check the job log in the Queue.`);
@@ -172,7 +202,9 @@ function GeneratePageInner() {
     try {
       const gpu = isMac() ? 'mps' : gpuId;
       const name = `inference_engine_gpu${gpu}`;
-      const existing = await apiClient.get('/api/jobs', { params: { job_type: 'inference' } }).then(r => r.data.jobs || []);
+      const existing = await apiClient
+        .get('/api/jobs', { params: { job_type: 'inference' } })
+        .then(r => r.data.jobs || []);
       const match = existing.find((j: any) => j.name === name);
       const res = await apiClient.post('/api/jobs', {
         id: match ? match.id : null,
@@ -206,7 +238,10 @@ function GeneratePageInner() {
       let alive = true;
       while (Date.now() < deadline) {
         await new Promise(r => setTimeout(r, 750));
-        const st = await apiClient.get(`/api/inference/status?job=${id}`).then(r => r.data).catch(() => null);
+        const st = await apiClient
+          .get(`/api/inference/status?job=${id}`)
+          .then(r => r.data)
+          .catch(() => null);
         alive = !!st?.engines?.some((e: any) => e.jobId === id && (e.ready || e.status === 'stopping'));
         if (!alive) break;
       }
@@ -222,7 +257,8 @@ function GeneratePageInner() {
 
   // a row that is stopped but whose process is still alive shows as 'stopping' from the status route
   const isStopping = !!stoppingJobId || engineStatus?.engine?.status === 'stopping';
-  const isStarting = !isStopping && (engineBusy || !!startingJobId || (!!engineStatus?.engine && !engineStatus.running));
+  const isStarting =
+    !isStopping && (engineBusy || !!startingJobId || (!!engineStatus?.engine && !engineStatus.running));
   const stoppingForMs = stoppingSince ? now - stoppingSince : isStopping ? 999_999 : 0;
   // cancel a start that is not getting anywhere: a queued job (no process)
   // is just marked stopped; a launched one gets the normal stop
@@ -259,12 +295,18 @@ function GeneratePageInner() {
   };
 
   // ---- models: the same arch list the training UI uses (jobs/new/options.tsx) ----
-  const archs: GenerateDefaults[] = useMemo(() => modelArchs.map(getGenerateDefaults), []);
+  // text-generating archs have no Generate page path yet
+  const archs: GenerateDefaults[] = useMemo(
+    () => modelArchs.filter(a => a.group !== 'llm').map(getGenerateDefaults),
+    [],
+  );
   const ready = !!engineStatus?.running;
 
   const [arch, setArch] = useState<string>(persisted.arch || 'zimage');
   const [model, setModel] = useState<{ [key: string]: any }>(persisted.model || {});
-  const [sample, setSample] = useState<{ [key: string]: any }>(persisted.sample || { prompt: '', negative_prompt: '', seed: -1 });
+  const [sample, setSample] = useState<{ [key: string]: any }>(
+    persisted.sample || { prompt: '', negative_prompt: '', seed: -1 },
+  );
   const entry = useMemo(() => archs.find(a => a.arch === arch) || null, [archs, arch]);
 
   const applyArch = (name: string) => {
@@ -272,7 +314,14 @@ function GeneratePageInner() {
     const e = archs.find(a => a.arch === name);
     if (!e) return;
     setModel({ ...e.model });
-    setSample(s => ({ ...s, ...e.sample, prompt: s.prompt, negative_prompt: s.negative_prompt, seed: s.seed ?? -1, ctrl_img: undefined }));
+    setSample(s => ({
+      ...s,
+      ...e.sample,
+      prompt: s.prompt,
+      negative_prompt: s.negative_prompt,
+      seed: s.seed ?? -1,
+      ctrl_img: undefined,
+    }));
   };
   useEffect(() => {
     if (archs.length && !Object.keys(model).length) applyArch(archs.some(a => a.arch === arch) ? arch : archs[0].arch);
@@ -652,7 +701,9 @@ function GeneratePageInner() {
         : engineStatus?.engine
           ? `Engine ${engineStatus.engine.status}`
           : 'Engine off';
-  const vram = health?.vram ? `${(health.vram.used / 1e9).toFixed(1)} / ${(health.vram.total / 1e9).toFixed(0)} GB` : null;
+  const vram = health?.vram
+    ? `${(health.vram.used / 1e9).toFixed(1)} / ${(health.vram.total / 1e9).toFixed(0)} GB`
+    : null;
   const pool = health?.pool;
 
   return (
@@ -666,11 +717,18 @@ function GeneratePageInner() {
         <div className="flex items-center gap-2 text-xs sm:text-sm">
           {isStopping ? (
             <>
-              <Button disabled className="px-2 py-1 rounded-md bg-red-900 text-white/70 flex items-center gap-1 cursor-wait">
+              <Button
+                disabled
+                className="px-2 py-1 rounded-md bg-red-900 text-white/70 flex items-center gap-1 cursor-wait"
+              >
                 <Loader2 className="w-4 h-4 animate-spin" /> Stopping engine…
               </Button>
               {stoppingForMs > 8000 && (
-                <Button onClick={forceStop} className="px-2 py-1 rounded-md bg-red-700 hover:bg-red-600 text-white flex items-center gap-1" title="Kill the engine process">
+                <Button
+                  onClick={forceStop}
+                  className="px-2 py-1 rounded-md bg-red-700 hover:bg-red-600 text-white flex items-center gap-1"
+                  title="Kill the engine process"
+                >
                   <OctagonX className="w-4 h-4" /> Force stop
                 </Button>
               )}
@@ -680,7 +738,11 @@ function GeneratePageInner() {
               <span className="text-green-400">Engine running</span>
               {health?.active?.arch && <span className="text-gray-400 hidden sm:inline">· {health.active.arch}</span>}
               {vram && <span className="text-gray-400 hidden md:inline">· VRAM {vram}</span>}
-              <Button onClick={stopEngine} disabled={engineBusy} className="ml-2 px-2 py-1 rounded-md bg-red-700 hover:bg-red-600 disabled:opacity-60 text-white flex items-center gap-1">
+              <Button
+                onClick={stopEngine}
+                disabled={engineBusy}
+                className="ml-2 px-2 py-1 rounded-md bg-red-700 hover:bg-red-600 disabled:opacity-60 text-white flex items-center gap-1"
+              >
                 {engineBusy ? <Loader2 className="w-4 h-4 animate-spin" /> : <Square className="w-4 h-4" />} Stop
               </Button>
             </>
@@ -702,15 +764,28 @@ function GeneratePageInner() {
               </p>
               {isStarting ? (
                 <div className="flex flex-col gap-2">
-                  <Button disabled className="w-full px-3 py-2 rounded-md bg-blue-900 text-white/70 flex items-center justify-center gap-2 cursor-wait">
+                  <Button
+                    disabled
+                    className="w-full px-3 py-2 rounded-md bg-blue-900 text-white/70 flex items-center justify-center gap-2 cursor-wait"
+                  >
                     <Loader2 className="w-4 h-4 animate-spin" />
-                    {engineStatus?.engine?.status === 'running' ? 'Engine loading…' : engineStatus?.engine ? `Engine ${engineStatus.engine.status}…` : 'Starting engine…'}
+                    {engineStatus?.engine?.status === 'running'
+                      ? 'Engine loading…'
+                      : engineStatus?.engine
+                        ? `Engine ${engineStatus.engine.status}…`
+                        : 'Starting engine…'}
                   </Button>
                   {engineStatus?.engine?.status === 'queued' && (
-                    <div className="text-[11px] text-gray-500">Waiting for the GPU queue to pick the job up. If the queue is stopped, start it from the Queue page or cancel here.</div>
+                    <div className="text-[11px] text-gray-500">
+                      Waiting for the GPU queue to pick the job up. If the queue is stopped, start it from the Queue
+                      page or cancel here.
+                    </div>
                   )}
                   {startingSince !== null && now - startingSince > 5000 && (
-                    <Button onClick={cancelStart} className="w-full px-3 py-1.5 rounded-md bg-gray-800 hover:bg-gray-700 text-gray-200 text-sm">
+                    <Button
+                      onClick={cancelStart}
+                      className="w-full px-3 py-1.5 rounded-md bg-gray-800 hover:bg-gray-700 text-gray-200 text-sm"
+                    >
                       Cancel
                     </Button>
                   )}
@@ -718,7 +793,11 @@ function GeneratePageInner() {
               ) : (
                 <div className="flex flex-col gap-3">
                   {!isMac() && (
-                    <select value={gpuId} onChange={e => setGpuId(e.target.value)} className="w-full bg-gray-800 text-gray-200 rounded-md px-3 py-2 border border-gray-700">
+                    <select
+                      value={gpuId}
+                      onChange={e => setGpuId(e.target.value)}
+                      className="w-full bg-gray-800 text-gray-200 rounded-md px-3 py-2 border border-gray-700"
+                    >
                       {gpuList.map(g => (
                         <option key={g.index} value={`${g.index}`}>
                           GPU #{g.index} {g.name}
@@ -726,7 +805,10 @@ function GeneratePageInner() {
                       ))}
                     </select>
                   )}
-                  <Button onClick={startEngine} className="w-full px-3 py-2 rounded-md bg-blue-700 hover:bg-blue-600 text-white flex items-center justify-center gap-2">
+                  <Button
+                    onClick={startEngine}
+                    className="w-full px-3 py-2 rounded-md bg-blue-700 hover:bg-blue-600 text-white flex items-center justify-center gap-2"
+                  >
                     <Play className="w-4 h-4" /> Start engine
                   </Button>
                 </div>
@@ -756,7 +838,10 @@ function GeneratePageInner() {
                 ) : null}
                 <div className="flex-1" />
                 {running && (
-                  <Button onClick={cancel} className="px-2 py-0.5 rounded bg-gray-800/80 hover:bg-gray-700 text-gray-200 text-xs">
+                  <Button
+                    onClick={cancel}
+                    className="px-2 py-0.5 rounded bg-gray-800/80 hover:bg-gray-700 text-gray-200 text-xs"
+                  >
                     Cancel
                   </Button>
                 )}
@@ -766,7 +851,10 @@ function GeneratePageInner() {
                       Engine busy: {health.current.arch}
                       {health.current.total_steps ? ` ${health.current.step}/${health.current.total_steps}` : ''}
                     </span>
-                    <Button onClick={cancelBusy} className="px-2 py-0.5 rounded bg-gray-800/80 hover:bg-gray-700 text-gray-200 text-xs">
+                    <Button
+                      onClick={cancelBusy}
+                      className="px-2 py-0.5 rounded bg-gray-800/80 hover:bg-gray-700 text-gray-200 text-xs"
+                    >
                       Cancel
                     </Button>
                   </>
@@ -774,10 +862,17 @@ function GeneratePageInner() {
               </div>
               {progress?.total ? (
                 <div className="absolute top-0 left-0 right-0 h-0.5 bg-gray-800 z-10">
-                  <div className="h-0.5 bg-blue-500 transition-all" style={{ width: `${Math.min(100, (progress.step / progress.total) * 100)}%` }} />
+                  <div
+                    className="h-0.5 bg-blue-500 transition-all"
+                    style={{ width: `${Math.min(100, (progress.step / progress.total) * 100)}%` }}
+                  />
                 </div>
               ) : null}
-              {error && <div className="absolute bottom-3 left-3 right-3 z-10 text-rose-300 text-sm bg-rose-950/70 rounded px-3 py-2 whitespace-pre-wrap">{error}</div>}
+              {error && (
+                <div className="absolute bottom-3 left-3 right-3 z-10 text-rose-300 text-sm bg-rose-950/70 rounded px-3 py-2 whitespace-pre-wrap">
+                  {error}
+                </div>
+              )}
 
               {/* live latent preview: shown while generating (or when nothing is selected) */}
               <div className={`absolute inset-3 flex items-center justify-center ${showPreview ? '' : 'hidden'}`}>
@@ -791,7 +886,14 @@ function GeneratePageInner() {
               {!showPreview && selected && (
                 <div className="absolute inset-3 flex items-center justify-center">
                   {selected.kind === 'video' ? (
-                    <video key={selected.path} src={resultUrl(selected)} controls autoPlay loop className="max-w-full max-h-full rounded" />
+                    <video
+                      key={selected.path}
+                      src={resultUrl(selected)}
+                      controls
+                      autoPlay
+                      loop
+                      className="max-w-full max-h-full rounded"
+                    />
                   ) : selected.kind === 'audio' ? (
                     <div className="w-full max-w-xl bg-gray-900 rounded-xl p-6 text-center">
                       <div className="text-gray-300 text-sm mb-4 truncate" title={selected.prompt}>
@@ -800,9 +902,19 @@ function GeneratePageInner() {
                       <audio key={selected.path} src={resultUrl(selected)} controls autoPlay className="w-full" />
                     </div>
                   ) : (
-                    <a href={resultUrl(selected)} target="_blank" rel="noreferrer" className="w-full h-full flex items-center justify-center">
+                    <a
+                      href={resultUrl(selected)}
+                      target="_blank"
+                      rel="noreferrer"
+                      className="w-full h-full flex items-center justify-center"
+                    >
                       {/* eslint-disable-next-line @next/next/no-img-element */}
-                      <img key={selected.path} src={resultUrl(selected)} alt={selected.prompt} className="max-w-full max-h-full rounded object-contain" />
+                      <img
+                        key={selected.path}
+                        src={resultUrl(selected)}
+                        alt={selected.prompt}
+                        className="max-w-full max-h-full rounded object-contain"
+                      />
                     </a>
                   )}
                 </div>
@@ -901,20 +1013,40 @@ function GeneratePageInner() {
             >
               {sidebarOpen ? <ChevronRight className="w-3.5 h-3.5" /> : <ChevronLeft className="w-3.5 h-3.5" />}
             </button>
-            <div className={`h-full flex flex-col border-l border-gray-800 bg-gray-900/60 ${sidebarOpen ? '' : 'hidden'}`}>
+            <div
+              className={`h-full flex flex-col border-l border-gray-800 bg-gray-900/60 ${sidebarOpen ? '' : 'hidden'}`}
+            >
               <div className="compact-form flex-1 min-h-0 overflow-y-auto p-2 space-y-2">
-
+                <Card
-                <Card title="Model" subtitle={entry?.label || arch} open={cardOpen('model')} onToggle={() => toggleCard('model')}>
+                  title="Model"
-                  <SelectInput label="Architecture" value={arch} onChange={v => applyArch(v as string)} options={archOptions} disabled={!ready} />
+                  subtitle={entry?.label || arch}
+                  open={cardOpen('model')}
+                  onToggle={() => toggleCard('model')}
+                >
+                  <SelectInput
+                    label="Architecture"
+                    value={arch}
+                    onChange={v => applyArch(v as string)}
+                    options={archOptions}
+                    disabled={!ready}
+                  />
                   <CreatableSelectInput
                     label="Name or Path"
                     value={model.name_or_path || ''}
                     onChange={(v: string | null) => setModel(m => ({ ...m, name_or_path: v || '' }))}
-                    options={entry?.model?.name_or_path ? [{ value: entry.model.name_or_path, label: entry.model.name_or_path }] : []}
+                    options={
+                      entry?.model?.name_or_path
+                        ? [{ value: entry.model.name_or_path, label: entry.model.name_or_path }]
+                        : []
+                    }
                     placeholder="hub repo, local folder, or .safetensors"
                   />
                   {'extras_name_or_path' in (entry?.model || {}) && (
-                    <TextInput label="Extras Name or Path" value={model.extras_name_or_path || ''} onChange={v => setModel(m => ({ ...m, extras_name_or_path: v }))} />
+                    <TextInput
+                      label="Extras Name or Path"
+                      value={model.extras_name_or_path || ''}
+                      onChange={v => setModel(m => ({ ...m, extras_name_or_path: v }))}
+                    />
                   )}
                   <div className="grid grid-cols-2 gap-2">
                     <SelectInput
@@ -931,8 +1063,16 @@ function GeneratePageInner() {
                     />
                   </div>
                   <div className="grid grid-cols-2 gap-2 mt-2">
-                    <Checkbox label="Low VRAM" checked={!!model.low_vram} onChange={v => setModel(m => ({ ...m, low_vram: v }))} />
+                    <Checkbox
-                    <Checkbox label="Layer offloading" checked={!!model.layer_offloading} onChange={v => setModel(m => ({ ...m, layer_offloading: v }))} />
+                      label="Low VRAM"
+                      checked={!!model.low_vram}
+                      onChange={v => setModel(m => ({ ...m, low_vram: v }))}
+                    />
+                    <Checkbox
+                      label="Layer offloading"
+                      checked={!!model.layer_offloading}
+                      onChange={v => setModel(m => ({ ...m, layer_offloading: v }))}
+                    />
                   </div>
                   {model.layer_offloading && (
                     <div className="mt-2 space-y-1">
@@ -956,10 +1096,18 @@ function GeneratePageInner() {
                   )}
                 </Card>
 
-                <Card title="LoRAs" subtitle={loras.length ? `${activeLoras.length}/${loras.length} active` : 'none'} open={cardOpen('loras')} onToggle={() => toggleCard('loras')}>
+                <Card
+                  title="LoRAs"
+                  subtitle={loras.length ? `${activeLoras.length}/${loras.length} active` : 'none'}
+                  open={cardOpen('loras')}
+                  onToggle={() => toggleCard('loras')}
+                >
                   <div className="space-y-2">
                     {loras.map((l, i) => (
-                      <div key={l.path} className={`bg-gray-950/60 border border-gray-800 rounded-md px-2 py-1.5 ${l.disabled ? 'opacity-50' : ''}`}>
+                      <div
+                        key={l.path}
+                        className={`bg-gray-950/60 border border-gray-800 rounded-md px-2 py-1.5 ${l.disabled ? 'opacity-50' : ''}`}
+                      >
                         <div className="flex items-center gap-2">
                           <button
                             type="button"
@@ -977,7 +1125,10 @@ function GeneratePageInner() {
                               className={`absolute top-0.5 left-0.5 w-3 h-3 rounded-full bg-white transition-transform ${l.disabled ? '' : 'translate-x-3'}`}
                             />
                           </button>
-                          <span className={`text-xs truncate flex-1 ${l.disabled ? 'text-gray-500 line-through' : 'text-gray-200'}`} title={l.path}>
+                          <span
+                            className={`text-xs truncate flex-1 ${l.disabled ? 'text-gray-500 line-through' : 'text-gray-200'}`}
+                            title={l.path}
+                          >
                             {l.name}
                           </span>
                           <input
@@ -1022,13 +1173,15 @@ function GeneratePageInner() {
                         description: (
                           <>
                             <p>
-                              <strong>Off</strong> (default): each LoRA is applied dynamically as an extra term on the layers it targets. Strength
+                              <strong>Off</strong> (default): each LoRA is applied dynamically as an extra term on the
-                              changes take effect on the next generation and adding or removing a LoRA never reloads the model.
+                              layers it targets. Strength changes take effect on the next generation and adding or
+                              removing a LoRA never reloads the model.
                             </p>
                             <p className="mt-2">
-                              <strong>On</strong>: the LoRA deltas are added into the model weights. Quantized weights are dequantized, merged, and
+                              <strong>On</strong>: the LoRA deltas are added into the model weights. Quantized weights
-                              re-quantized with stochastic rounding so small deltas are not rounded away. Merged weights cannot be un-merged, so
+                              are dequantized, merged, and re-quantized with stochastic rounding so small deltas are not
-                              changing the LoRA set or strengths reloads the affected components. Slightly faster per step than the dynamic path.
+                              rounded away. Merged weights cannot be un-merged, so changing the LoRA set or strengths
+                              reloads the affected components. Slightly faster per step than the dynamic path.
                             </p>
                           </>
                         ),
@@ -1041,40 +1194,125 @@ function GeneratePageInner() {
                     // tagged prompt archs (audio): one field per tag, stored as the tagged prompt string
                     Object.entries(entry.sampleTags).map(([tagKey, tag]) => {
                       const tags = tagsToObj(sample.prompt || '');
-                      const setTag = (v: any) => setSample(s => ({ ...s, prompt: objToTags({ ...tagsToObj(s.prompt || ''), [tagKey]: v }) }));
+                      const setTag = (v: any) =>
+                        setSample(s => ({ ...s, prompt: objToTags({ ...tagsToObj(s.prompt || ''), [tagKey]: v }) }));
                       const value = tags[tagKey] ?? '';
-                      if (tag.type === 'multiline') return <TextAreaInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                      if (tag.type === 'multiline')
-                      if (tag.type === 'number') return <NumberInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                        return (
-                      return <TextInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                          <TextAreaInput
+                            key={tagKey}
+                            label={tag.title}
+                            value={value}
+                            onChange={setTag}
+                            placeholder={`Enter ${tag.title.toLowerCase()}`}
+                          />
+                        );
+                      if (tag.type === 'number')
+                        return (
+                          <NumberInput
+                            key={tagKey}
+                            label={tag.title}
+                            value={value}
+                            onChange={setTag}
+                            placeholder={`Enter ${tag.title.toLowerCase()}`}
+                          />
+                        );
+                      return (
+                        <TextInput
+                          key={tagKey}
+                          label={tag.title}
+                          value={value}
+                          onChange={setTag}
+                          placeholder={`Enter ${tag.title.toLowerCase()}`}
+                        />
+                      );
                     })
                   ) : (
-                    <TextAreaInput label="Prompt" value={sample.prompt || ''} onChange={v => setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} />
+                    <TextAreaInput
+                      label="Prompt"
+                      value={sample.prompt || ''}
+                      onChange={v => setSample(s => ({ ...s, prompt: v }))}
+                      placeholder={
+                        modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'
+                      }
+                    />
                   )}
                   {(sample.guidance_scale ?? 4) > 1 && modality !== 'audio' && (
-                    <TextAreaInput label="Negative prompt" value={sample.negative_prompt || ''} onChange={v => setSample(s => ({ ...s, negative_prompt: v }))} />
+                    <TextAreaInput
+                      label="Negative prompt"
+                      value={sample.negative_prompt || ''}
+                      onChange={v => setSample(s => ({ ...s, negative_prompt: v }))}
+                    />
                   )}
                   <div className="grid grid-cols-3 gap-2">
                     {modality !== 'audio' && (
                       <>
-                        <NumberInput label="Width" value={sample.width ?? 1024} onChange={v => setSample(s => ({ ...s, width: v }))} min={64} max={4096} />
+                        <NumberInput
-                        <NumberInput label="Height" value={sample.height ?? 1024} onChange={v => setSample(s => ({ ...s, height: v }))} min={64} max={4096} />
+                          label="Width"
+                          value={sample.width ?? 1024}
+                          onChange={v => setSample(s => ({ ...s, width: v }))}
+                          min={64}
+                          max={4096}
+                        />
+                        <NumberInput
+                          label="Height"
+                          value={sample.height ?? 1024}
+                          onChange={v => setSample(s => ({ ...s, height: v }))}
+                          min={64}
+                          max={4096}
+                        />
                       </>
                     )}
-                    <NumberInput label="Steps" value={sample.num_inference_steps ?? 25} onChange={v => setSample(s => ({ ...s, num_inference_steps: v }))} min={1} max={200} />
+                    <NumberInput
-                    <NumberInput label="Guidance" value={sample.guidance_scale ?? 4} onChange={v => setSample(s => ({ ...s, guidance_scale: v }))} min={1} max={30} />
+                      label="Steps"
-                    <NumberInput label="Seed (-1 random)" value={sample.seed ?? -1} onChange={v => setSample(s => ({ ...s, seed: v }))} min={-1} max={4294967295} />
+                      value={sample.num_inference_steps ?? 25}
+                      onChange={v => setSample(s => ({ ...s, num_inference_steps: v }))}
+                      min={1}
+                      max={200}
+                    />
+                    <NumberInput
+                      label="Guidance"
+                      value={sample.guidance_scale ?? 4}
+                      onChange={v => setSample(s => ({ ...s, guidance_scale: v }))}
+                      min={1}
+                      max={30}
+                    />
+                    <NumberInput
+                      label="Seed (-1 random)"
+                      value={sample.seed ?? -1}
+                      onChange={v => setSample(s => ({ ...s, seed: v }))}
+                      min={-1}
+                      max={4294967295}
+                    />
                     {modality === 'audio' && sample.duration !== undefined && (
-                      <NumberInput label="Duration (s)" value={sample.duration} onChange={v => setSample(s => ({ ...s, duration: v }))} min={1} max={600} />
+                      <NumberInput
+                        label="Duration (s)"
+                        value={sample.duration}
+                        onChange={v => setSample(s => ({ ...s, duration: v }))}
+                        min={1}
+                        max={600}
+                      />
                     )}
                     {modality === 'video' && (
                       <>
-                        <NumberInput label="Frames" value={sample.num_frames ?? 33} onChange={v => setSample(s => ({ ...s, num_frames: v }))} min={1} max={1000} />
+                        <NumberInput
-                        <NumberInput label="FPS" value={sample.fps ?? 16} onChange={v => setSample(s => ({ ...s, fps: v }))} min={1} max={60} />
+                          label="Frames"
+                          value={sample.num_frames ?? 33}
+                          onChange={v => setSample(s => ({ ...s, num_frames: v }))}
+                          min={1}
+                          max={1000}
+                        />
+                        <NumberInput
+                          label="FPS"
+                          value={sample.fps ?? 16}
+                          onChange={v => setSample(s => ({ ...s, fps: v }))}
+                          min={1}
+                          max={60}
+                        />
                       </>
                     )}
                   </div>
                 </Card>
-    
               </div>
               {/* always-visible action bar */}
               <div className="shrink-0 p-2 border-t border-gray-800 bg-gray-900 flex gap-2">
@@ -1096,7 +1334,12 @@ function GeneratePageInner() {
           </div>
         </div>
       </MainContent>
-      <GenerateFooter jobId={engineJobId} status={footerStatus} busy={running || isStarting} progress={running ? progress : null} />
+      <GenerateFooter
+        jobId={engineJobId}
+        status={footerStatus}
+        busy={running || isStarting}
+        progress={running ? progress : null}
+      />
       <LoraBrowserModal isOpen={loraModalOpen} onClose={() => setLoraModalOpen(false)} onPick={addLora} />
     </>
   );

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

@@ -82,6 +82,8 @@ export default function SimpleJob({
 
   const isVideoModel = !!(modelArch?.group === 'video');
   const isAudioModel = !!(modelArch?.group === 'audio');
+  // text-generating models: samples are media in, text out (no size)
+  const isLlmModel = !!(modelArch?.group === 'llm');
 
   const taggedSampleArr: Record<string, any>[] | null = useMemo(() => {
     if (!modelArch) return null;
@@ -213,6 +215,9 @@ export default function SimpleJob({
     numDatasetCols -= 1;
     numSampleTopCols -= 1;
   }
+  if (isLlmModel) {
+    numSampleTopCols -= 1;
+  }
   if (numDatasetCols == 3) {
     datasetStyleClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-6';
   }
@@ -296,6 +301,18 @@ export default function SimpleJob({
               placeholder=""
               required
             />
+            {modelArch?.additionalSections?.includes('model.model_kwargs.instruction') && (
+              <TextAreaInput
+                label="LLM Prompt"
+                className="pt-2"
+                rows={2}
+                docKey="model.model_kwargs.instruction"
+                value={jobConfig.config.process[0].model.model_kwargs?.instruction ?? ''}
+                onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.instruction')}
+                placeholder="Describe this in detail."
+                required
+              />
+            )}
             {modelArch?.additionalSections?.includes('model.assistant_lora_path') && (
               <TextInput
                 label="Training Adapter Path"
@@ -489,7 +506,7 @@ export default function SimpleJob({
                 }}
                 options={transformerQuantizationOptions}
               />
-              {!disableSections.includes('model.quantize_te') && (
+              {!disableSections.includes('model.quantize_te') && !isLlmModel && (
                 <SelectInput
                   label="Text Encoder"
                   value={
@@ -507,6 +524,8 @@ export default function SimpleJob({
                   options={quantizationOptions}
                 />
               )}
+              {!isLlmModel && (
+                <>
                   <FormGroup label="Compile Options">
                     <></>
                   </FormGroup>
@@ -526,6 +545,8 @@ export default function SimpleJob({
                       }
                     }}
                   />
+                </>
+              )}
             </Card>
           )}
           {modelArch?.additionalSections?.includes('model.multistage') && (
@@ -737,6 +758,7 @@ export default function SimpleJob({
                   required
                 />
               </div>
+              {!isLlmModel && (
                 <div>
                   {disableSections.includes('train.timestep_type') ? null : (
                     <SelectInput
@@ -787,6 +809,7 @@ export default function SimpleJob({
                     />
                   )}
                 </div>
+              )}
               <div>
                 <FormGroup label="EMA (Exponential Moving Average)">
                   <Checkbox
@@ -807,6 +830,7 @@ export default function SimpleJob({
                   />
                 )}
 
+                {!isLlmModel && (
                   <FormGroup label="Text Encoder Optimizations" className="pt-2">
                     {!disableSections.includes('train.unload_text_encoder') && (
                       <Checkbox
@@ -833,6 +857,7 @@ export default function SimpleJob({
                       }}
                     />
                   </FormGroup>
+                )}
               </div>
               <div>
                 {disableSections.includes('train.diff_output_preservation') ||
@@ -914,6 +939,7 @@ export default function SimpleJob({
                     )}
                   </>
                 )}
+                {!isLlmModel && (
                   <FormGroup label="Other" className="pt-2">
                     <>
                       <Checkbox
@@ -947,11 +973,13 @@ export default function SimpleJob({
                       )}
                     </>
                   </FormGroup>
+                )}
               </div>
             </div>
           </Card>
         </div>
         <div>
+          {!isLlmModel && (
             <Card
               title="Validation"
               toggled={!!validationConfig}
@@ -978,8 +1006,8 @@ export default function SimpleJob({
                     and predicted at the selected sigmas with fixed seeds, so the result is always deterministic and
                     comparable across the run. The average loss is logged as val/loss every time validation runs. The
                     images need to match the concept of your dataset, but{' '}
-                  <span className="font-bold text-gray-300">do not include the validation images in the dataset</span>.
+                    <span className="font-bold text-gray-300">do not include the validation images in the dataset</span>
-                  They must be images containing the concept you want to train, but not an image trained on.
+                    . They must be images containing the concept you want to train, but not an image trained on.
                   </p>
                   <div className="grid grid-cols-1 md:grid-cols-3 gap-6">
                     <NumberInput
@@ -1080,8 +1108,10 @@ export default function SimpleJob({
                 </>
               )}
             </Card>
+          )}
         </div>
         <div>
+          {!isLlmModel && (
             <Card title="Advanced" collapsible>
               <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 gap-6">
                 <div>
@@ -1119,6 +1149,7 @@ export default function SimpleJob({
                 </div>
               </div>
             </Card>
+          )}
         </div>
         <div>
           <Card title="Datasets">
@@ -1359,7 +1390,7 @@ export default function SimpleJob({
                           />
                         )}
                       </FormGroup>
-                      {!isAudioModel && (
+                      {!isAudioModel && !isLlmModel && (
                         <FormGroup label="Flipping" docKey={'datasets.flip'} className="mt-2">
                           <Checkbox
                             label={
@@ -1458,6 +1489,8 @@ export default function SimpleJob({
                   min={0}
                   required
                 />
+                {!isLlmModel && (
+                  <>
                     <SelectInput
                       label="Sampler"
                       className="pt-2"
@@ -1486,9 +1519,11 @@ export default function SimpleJob({
                       min={1}
                       required
                     />
+                  </>
+                )}
               </div>
 
-              {!isAudioModel && (
+              {!isAudioModel && !isLlmModel && (
                 <div>
                   <NumberInput
                     label="Width"
@@ -1532,6 +1567,7 @@ export default function SimpleJob({
                 </div>
               )}
 
+              {!isLlmModel && (
                 <div>
                   <NumberInput
                     label="Seed"
@@ -1559,6 +1595,7 @@ export default function SimpleJob({
                     onChange={value => setJobConfig(value, 'config.process[0].sample.walk_seed')}
                   />
                 </div>
+              )}
               <div>
                 <FormGroup label="Advanced Sampling" className="pt-2">
                   <div>
@@ -1713,8 +1750,8 @@ export default function SimpleJob({
                                 label={`Prompt`}
                                 value={sample.prompt}
                                 onChange={value => setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)}
-                                placeholder="Enter prompt"
+                                placeholder={isLlmModel ? 'blank = LLM Prompt' : 'Enter prompt'}
-                                required
+                                required={!isLlmModel}
                               />
                             )}
                           </>
@@ -1746,7 +1783,7 @@ export default function SimpleJob({
                         )}
 
                         <div className="grid w-full lg:grid-flow-col lg:auto-cols-fr gap-4 mt-2">
-                          {!isAudioModel && (
+                          {!isAudioModel && !isLlmModel && (
                             <TextInput
                               label={`Width`}
                               value={sample.width ? `${sample.width}` : ''}
@@ -1775,7 +1812,7 @@ export default function SimpleJob({
                               placeholder={`${jobConfig.config.process[0].sample.width} (default)`}
                             />
                           )}
-                          {!isAudioModel && (
+                          {!isAudioModel && !isLlmModel && (
                             <TextInput
                               label={`Height`}
                               value={sample.height ? `${sample.height}` : ''}
@@ -1804,6 +1841,7 @@ export default function SimpleJob({
                               placeholder={`${jobConfig.config.process[0].sample.height} (default)`}
                             />
                           )}
+                          {!isLlmModel && (
                             <TextInput
                               label={`Seed`}
                               value={sample.seed ? `${sample.seed}` : ''}
@@ -1831,6 +1869,7 @@ export default function SimpleJob({
                               }}
                               placeholder={`${jobConfig.config.process[0].sample.walk_seed ? jobConfig.config.process[0].sample.seed + i : jobConfig.config.process[0].sample.seed} (default)`}
                             />
+                          )}
                           <TextInput
                             label={`LoRA Scale`}
                             value={sample.network_multiplier ? `${sample.network_multiplier}` : ''}
@@ -1883,6 +1922,8 @@ export default function SimpleJob({
                       {modelArch?.additionalSections?.includes('sample.ctrl_img') && (
                         <SampleControlImage
                           className="mt-6 ml-4"
+                          instruction={isLlmModel ? 'Add Media' : undefined}
+                          allowAudio={isLlmModel}
                           src={sample.ctrl_img}
                           onNewImageSelected={imagePath => {
                             if (!imagePath) {

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

@@ -7,6 +7,7 @@ import {
   defaultSampleConfig,
   defaultIdeogramSamplesConfig,
   defaultYue2SampleConfig,
+  defaultQwen25OmniSampleConfig,
 } from '@/helpers/defaultSamples';
 
 type Control = 'depth' | 'line' | 'pose' | 'inpaint';
@@ -42,9 +43,10 @@ type AdditionalSections =
   | 'model.assistant_lora_path'
   | 'model.unconditional_lora_path'
   | 'model.model_kwargs.kv_cache'
+  | 'model.model_kwargs.instruction'
   | 'ideogram_4_prompt';
 
-type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio';
+type ModelGroup = 'image' | 'instruction' | 'video' | 'experimental' | 'audio' | 'llm';
 
 export interface CustomModelSelectOption {
   label: string;
@@ -1378,6 +1380,62 @@ export const modelArchs: ModelArch[] = [
     disableSections: ['network.conv'],
     additionalSections: ['sample.multi_ctrl_imgs', 'model.low_vram', 'model.layer_offloading'],
   },
+  {
+    name: 'qwen25_omni',
+    label: 'Qwen2.5-Omni',
+    group: 'llm',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': [
+        'ostris/qwen2_5_omni/text_encoders/qwen2_5_omni_7b_convrot8.safetensors',
+        defaultNameOrPath,
+      ],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [false, false],
+      'config.process[0].model.low_vram': [false, false],
+      // the single-file thinker ships convrot8 layers; requesting convrot8 keeps them as-is
+      'config.process[0].model.qtype': ['convrot8', 'qfloat8'],
+      'config.process[0].train.unload_text_encoder': [false, false],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.batch_size': [1, 1],
+      'config.process[0].sample': [defaultQwen25OmniSampleConfig, defaultSampleConfig],
+      // media is encoded on the GPU per step; the cache is optional and large (~54 MB per 300 s of audio)
+      'config.process[0].datasets[x].cache_latents_to_disk': [false, true],
+      'config.process[0].datasets[x].resolution': [[512], [512, 768, 1024]],
+      // the caption is the training target; a blank one trains nothing
+      'config.process[0].datasets[x].caption_dropout_rate': [0, 0.05],
+      'config.process[0].model.model_kwargs': [{ instruction: 'Describe this in detail.' }, {}],
+    },
+    disableSections: [
+      'network.conv',
+      'trigger_word',
+      'train.diff_output_preservation',
+      'train.blank_prompt_preservation',
+      'train.unload_text_encoder',
+      'slider',
+    ],
+    additionalSections: ['model.model_kwargs.instruction', 'sample.ctrl_img', 'datasets.num_frames'],
+    modelNotes: (
+      <div className="space-y-2">
+        <p>
+          Qwen2.5-Omni 7B thinker as a text-generating model: audio, image or video in, text out. Each dataset item is a
+          media file (mp3, wav, flac, ogg, jpg, png, webp, mp4, ...) with a caption file next to it; the caption is the
+          text the model learns to produce for that media. One dataset can mix all three kinds. Video files are used
+          when Num Frames is above 1 and are seen as frames only (no audio track).
+        </p>
+        <p>
+          The instruction in <code>model_kwargs.instruction</code> is the user turn for every training item; use the
+          same wording when captioning with the trained LoRA. Samples take a media file per prompt and write the
+          generated text as a .txt file.
+        </p>
+        <p>
+          Media is encoded by the frozen audio and vision towers on the GPU each step, so Cache Latents to Disk can stay
+          off (the cache is about 54 MB per 300 seconds of audio). Batch size is 1. The LoRA trains the text stack only.
+          Watch <code>loss/ce</code>: next-token cross-entropy on the caption.
+        </p>
+      </div>
+    ),
+  },
   {
     name: 'yue2',
     label: 'YuE2',
@@ -1402,7 +1460,10 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].datasets[x].resolution': [[512], [512, 768, 1024]],
       // blank captions break lyric following; the AR must always see the prefix
       'config.process[0].datasets[x].caption_dropout_rate': [0, 0.05],
-      'config.process[0].model.model_kwargs': [{ cot: 'full', abc_dropout: 0.5, sample_ar_repetition_penalty: 1.2, ar_kl_weight: 0.2 }, {}],
+      'config.process[0].model.model_kwargs': [
+        { cot: 'full', abc_dropout: 0.5, sample_ar_repetition_penalty: 1.2, ar_kl_weight: 0.2 },
+        {},
+      ],
     },
     // native YuE2 prompt: style text, a [Lyrics] line, the lyrics
     hasMultiLinePrompts: true,

ui/src/components/SampleControlImage.tsx

@@ -9,13 +9,17 @@ import { encodeFilePathForUrl } from '@/utils/basic';
 import type { AxiosProgressEvent } from 'axios';
 
 const VIDEO_EXTS = ['.mp4', '.mov', '.webm', '.mkv', '.avi', '.m4v', '.wmv', '.flv'];
+const AUDIO_EXTS = ['.mp3', '.wav', '.flac', '.ogg'];
 const isVideoPath = (p: string) => VIDEO_EXTS.some(ext => p.toLowerCase().endsWith(ext));
+const isAudioPath = (p: string) => AUDIO_EXTS.some(ext => p.toLowerCase().endsWith(ext));
 
 interface Props {
   src: string | null | undefined;
   className?: string;
   instruction?: string;
   onNewImageSelected: (imagePath: string | null) => void;
+  /** also accept audio files (text-generating models take any media) */
+  allowAudio?: boolean;
 }
 
 export default function SampleControlImage({
@@ -23,6 +27,7 @@ export default function SampleControlImage({
   className,
   instruction = 'Add Control Image',
   onNewImageSelected,
+  allowAudio = false,
 }: Props) {
   const [isUploading, setIsUploading] = useState(false);
   const [uploadProgress, setUploadProgress] = useState(0);
@@ -31,7 +36,8 @@ export default function SampleControlImage({
 
   const backgroundUrl = useMemo(() => {
     if (localPreview) return localPreview;
-    // videos preview as a server-generated thumbnail
+    // videos preview as a server-generated thumbnail, audio as its waveform art
+    if (src && isAudioPath(src)) return `/api/audio/art/${encodeURIComponent(src)}`;
     if (src) return `/api/img/${encodeFilePathForUrl(src)}${isVideoPath(src) ? '?thumb=1' : ''}`;
     return null;
   }, [src, localPreview]);
@@ -104,6 +110,7 @@ export default function SampleControlImage({
     accept: {
       'image/*': ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'],
       'video/*': VIDEO_EXTS,
+      ...(allowAudio ? { 'audio/*': AUDIO_EXTS } : {}),
     },
     multiple: false,
     noClick: true,
@@ -138,7 +145,7 @@ export default function SampleControlImage({
       <input
         ref={fileInputRef}
         type="file"
-        accept="image/*,video/*"
+        accept={allowAudio ? 'image/*,video/*,audio/*' : 'image/*,video/*'}
         className="hidden"
         onChange={e => {
           const file = e.currentTarget.files?.[0];

ui/src/components/SampleImageCard.tsx

@@ -1,5 +1,5 @@
 import React, { useRef, useEffect, useState, ReactNode } from 'react';
-import { isVideo, isAudio, encodeFilePathForUrl } from '@/utils/basic';
+import { isVideo, isAudio, isText, encodeFilePathForUrl } from '@/utils/basic';
 
 interface SampleImageCardProps {
   imageUrl: string;
@@ -39,6 +39,28 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
 
   const isItAudio = isAudio(imageUrl);
   const isItVideo = isVideo(imageUrl);
+  const isItText = isText(imageUrl);
+  const [text, setText] = useState<string | null>(null);
+
+  // text samples: fetch the file body and render it in the card
+  useEffect(() => {
+    if (!isItText || !isVisible) return;
+    const controller = new AbortController();
+    fetch(`/api/img/${encodeFilePathForUrl(imageUrl)}`, { signal: controller.signal })
+      .then(r => (r.ok ? r.text() : Promise.reject(new Error(`HTTP ${r.status}`))))
+      .then(t => {
+        setText(t);
+        setLoaded(true);
+      })
+      .catch(err => {
+        if (err?.name !== 'AbortError') console.error('Sample text fetch failed:', err);
+      });
+    return () => {
+      controller.abort();
+      setText(null);
+      setLoaded(false);
+    };
+  }, [isItText, isVisible, imageUrl]);
 
   // Observe both enter and exit
   useEffect(() => {
@@ -69,7 +91,7 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
   // the element unmounts). A short debounce skips requests entirely during fast
   // scrolls where the card is only briefly visible.
   useEffect(() => {
-    if (isItAudio) return;
+    if (isItAudio || isItText) return;
     if (!isVisible) return;
 
     const controller = new AbortController();
@@ -114,7 +136,7 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
       setLoaded(false);
       setVideoFallback(false);
     };
-  }, [isVisible, isItAudio, isItVideo, imageUrl]);
+  }, [isVisible, isItAudio, isItText, isItVideo, imageUrl]);
 
   return (
     <div className={`flex flex-col ${className}`}>
@@ -127,14 +149,23 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
         onClick={onClick}
       >
         <div
-          className={`absolute rounded-t-lg shadow-md bg-gray-900 overflow-hidden transition-all duration-200 ${
+          className={`absolute rounded-t-lg shadow-md bg-gray-900 overflow-hidden transition-all duration-200 [container-type:inline-size] ${
             selected ? 'inset-2' : 'inset-0'
-          } ${
+          } ${isVisible && !isItAudio && !loaded ? 'animate-pulse' : ''}`}
-            isVisible && !isItAudio && !loaded ? 'animate-pulse' : ''
-          }`}
         >
           {isVisible ? (
-            isItAudio ? (
+            isItText ? (
+              // font size scales with the card and the text length: a square of side W holds
+              // ~W^2 / (0.6 f^2) characters at size f, so f ~ 120/sqrt(n) cqw fills the box
+              <div
+                className="w-full h-full overflow-hidden p-[3cqw] leading-snug text-gray-200 whitespace-pre-wrap break-words text-left bg-gray-900"
+                style={{
+                  fontSize: `clamp(2.5cqw, ${(120 / Math.sqrt(Math.max((text ?? '').length, 1))).toFixed(2)}cqw, 12cqw)`,
+                }}
+              >
+                {text ?? ''}
+              </div>
+            ) : isItAudio ? (
               <div className="w-full h-full flex items-center justify-center bg-gray-900">
                 <img
                   src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}

ui/src/components/SampleImageViewer.tsx

@@ -8,7 +8,7 @@ import classNames from 'classnames';
 import { Menu, MenuButton, MenuItem, MenuItems } from '@headlessui/react';
 import { openConfirm } from './ConfirmModal';
 import { apiClient } from '@/utils/api';
-import { isVideo, isAudio, encodeFilePathForUrl } from '@/utils/basic';
+import { isVideo, isAudio, isText, encodeFilePathForUrl } from '@/utils/basic';
 import AudioPlayer from './AudioPlayer';
 import { TransformWrapper, TransformComponent } from 'react-zoom-pan-pinch';
 import BoundingBoxOverlay, { parseBoundingBoxes } from './BoundingBoxOverlay';
@@ -199,6 +199,27 @@ export default function SampleImageViewer({
     return imgPath;
   }, [showingControlIdx, controlImages, imgPath]);
 
+  // text samples (LLM models): the file body is the sample
+  const [sampleText, setSampleText] = useState<string | null>(null);
+  useEffect(() => {
+    if (!displayedImgPath || !isText(displayedImgPath)) {
+      setSampleText(null);
+      return;
+    }
+    const controller = new AbortController();
+    fetch(`/api/img/${encodeFilePathForUrl(displayedImgPath)}`, { signal: controller.signal })
+      .then(r => (r.ok ? r.text() : Promise.reject(new Error(`HTTP ${r.status}`))))
+      .then(setSampleText)
+      .catch(err => {
+        if (err?.name !== 'AbortError') console.error('Sample text fetch failed:', err);
+      });
+    return () => controller.abort();
+  }, [displayedImgPath]);
+
+  // control thumbnails: audio has no image thumb, it gets its waveform art
+  const thumbSrc = (p: string) =>
+    isAudio(p) ? `/api/audio/art/${encodeURIComponent(p)}` : `/api/img/${encodeFilePathForUrl(p)}?thumb=1`;
+
   // The sample's prompt is what generated it; if it's an Ideogram bbox-JSON we can
   // overlay the boxes on the generated image. Only on the main image (not controls).
   const boundingBoxes = useMemo(
@@ -311,7 +332,11 @@ export default function SampleImageViewer({
           >
             <div className="overflow-hidden flex items-center justify-center">
               {displayedImgPath &&
-                (isAudio(displayedImgPath) ? (
+                (isText(displayedImgPath) ? (
+                  <div className="w-[640px] max-w-full sm:max-w-[95vw] max-h-[82vh] overflow-y-auto p-6 text-sm text-gray-100 whitespace-pre-wrap break-words text-left">
+                    {sampleText ?? 'Loading…'}
+                  </div>
+                ) : isAudio(displayedImgPath) ? (
                   <div className="w-[500px] h-[500px] max-w-full sm:max-w-[95vw] max-h-[82vh]">
                     <AudioPlayer
                       src={`/api/img/${encodeFilePathForUrl(displayedImgPath)}`}
@@ -370,19 +395,28 @@ export default function SampleImageViewer({
               </div>
               {controlImages.length > 0 && (
                 <div key={imgPath} className="flex space-x-2 mr-4">
-                  {showingControlIdx !== null && (
+                  {showingControlIdx !== null &&
+                    (isText(imgPath!) ? (
+                      <div
+                        className="h-12 w-12 flex items-center justify-center text-xs bg-black border-2 border-gray-700 hover:border-gray-500 rounded cursor-pointer"
+                        onClick={() => setShowingControlIdx(null)}
+                        title="Generated text"
+                      >
+                        Txt
+                      </div>
+                    ) : (
                       <img
-                      src={`/api/img/${encodeFilePathForUrl(imgPath!)}?thumb=1`}
+                        src={thumbSrc(imgPath!)}
                         alt="Main"
                         className="max-h-12 max-w-12 object-contain bg-black border-2 border-gray-700 hover:border-gray-500 rounded cursor-pointer"
                         onClick={() => setShowingControlIdx(null)}
                         title="Main image"
                       />
-                  )}
+                    ))}
                   {controlImages.map((ci, idx) => (
                     <img
                       key={idx}
-                      src={`/api/img/${encodeFilePathForUrl(ci)}?thumb=1`}
+                      src={thumbSrc(ci)}
                       alt={`Control ${idx + 1}`}
                       className={`max-h-12 max-w-12 object-contain bg-black border-2 rounded cursor-pointer ${
                         showingControlIdx === idx ? 'border-blue-500' : 'border-gray-700 hover:border-gray-500'
@@ -429,7 +463,7 @@ export default function SampleImageViewer({
                     anchor="bottom end"
                     className="bg-gray-900 border border-gray-700 rounded shadow-lg w-48 px-2 py-2 mt-1 z-50"
                   >
-                    {imgPath && isAudio(imgPath) && (
+                    {imgPath && (isAudio(imgPath) || isText(imgPath)) && (
                       <MenuItem>
                         <a
                           className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded block"

ui/src/docs.tsx

@@ -340,24 +340,30 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'model.model_kwargs.instruction': {
+    title: 'LLM Prompt',
+    description: (
+      <>
+        The instruction the model is asked for every training item, as the user turn of the chat. The item's caption
+        file is the answer it learns to give. Use the same wording when running the trained LoRA, and it is also what
+        samples use when their own prompt is left blank.
+      </>
+    ),
+  },
   'model.model_kwargs.kv_cache': {
     title: 'KV Cache',
     description: (
       <>
-        This will enable KV Cache for control images in a model that supports it. LoRAs trained with this on
+        This will enable KV Cache for control images in a model that supports it. LoRAs trained with this on need to
-        need to also be inferenced with it, and vice versa. This does not speed up or slow down training, but on inference,
+        also be inferenced with it, and vice versa. This does not speed up or slow down training, but on inference, the
-        the control images only need to be processed once for the entire generation, vs being processed for every step.
+        control images only need to be processed once for the entire generation, vs being processed for every step.
         Which leads to a significant speedup on inference.
       </>
     ),
   },
   'train.guidance_loss_target': {
     title: 'Guidance Loss Target',
-    description: (
+    description: <>For contrastive guidance loss, this is the target CGF to amplify predictions to.</>,
-      <>
-        For contrastive guidance loss, this is the target CGF to amplify predictions to. 
-      </>
-    ),
   },
   'datasets.caption_dropout_rate': {
     title: 'Caption Dropout Rate',
@@ -365,15 +371,15 @@ const docs: { [key: string]: ConfigDoc } = {
       <>
         Caption dropout rate is the probability that the caption for an image will be dropped (replaced with a blank
         caption) for any given training step. For example, a value of 0.05 will drop the caption around 5% of the time.
-        Dropping captions helps the model learn the concept being trained without relying entirely on the caption,
+        Dropping captions helps the model learn the concept being trained without relying entirely on the caption, and
-        and helps preserve the model&apos;s ability to generate without a prompt. If a trigger word is set, the trigger
+        helps preserve the model&apos;s ability to generate without a prompt. If a trigger word is set, the trigger word
-        word is still used when the caption is dropped, so the model still associates the dropped samples with your
+        is still used when the caption is dropped, so the model still associates the dropped samples with your trigger
-        trigger word. Regularization images, or images without a trigger word, drop to a fully blank caption.
+        word. Regularization images, or images without a trigger word, drop to a fully blank caption.
         <br />
         <br />
-        Caption dropout also works when caching text embeddings. An additional embedding for the dropout caption
+        Caption dropout also works when caching text embeddings. An additional embedding for the dropout caption (blank,
-        (blank, or the trigger word alone) is cached to disk alongside the normal one, and it is randomly swapped in
+        or the trigger word alone) is cached to disk alongside the normal one, and it is randomly swapped in at train
-        at train time at this rate.
+        time at this rate.
       </>
     ),
   },

ui/src/helpers/defaultSamples.ts

@@ -985,6 +985,16 @@ const tagValue = (prompt: string, tag: string) => {
   return m ? m[1].trim() : '';
 };
 
+// text-generating models: the prompt is the instruction, ctrl_img the media; output is a .txt
+export const defaultQwen25OmniSampleConfig: SampleConfig = {
+  ...defaultSampleConfig,
+  width: 512,
+  height: 512,
+  guidance_scale: 1,
+  sample_steps: 1,
+  samples: [{ prompt: 'Describe this in detail.' }, { prompt: 'Describe this in detail.' }],
+};
+
 export const defaultYue2SampleConfig: SampleConfig = {
   ...defaultAudioSampleConfig,
   samples: defaultAudioSampleConfig.samples.map(s => ({

version.py

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