diff --git a/extensions_built_in/llm_models/__init__.py b/extensions_built_in/llm_models/__init__.py new file mode 100644 index 0000000..393325f --- /dev/null +++ b/extensions_built_in/llm_models/__init__.py @@ -0,0 +1,5 @@ +from .qwen25_omni import Qwen25OmniLLM + +AI_TOOLKIT_MODELS = [ + Qwen25OmniLLM, +] diff --git a/extensions_built_in/llm_models/qwen25_omni.py b/extensions_built_in/llm_models/qwen25_omni.py new file mode 100644 index 0000000..87374b1 --- /dev/null +++ b/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 = "<>" + 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 .txt by GenerateImageConfig.save_image + return self.processor.batch_decode(ids[:, input_len:], skip_special_tokens=True, clean_up_tokenization_spaces=False)[0].strip() diff --git a/extensions_built_in/llm_models/src/__init__.py b/extensions_built_in/llm_models/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/llm_models/src/thinker.py b/extensions_built_in/llm_models/src/thinker.py new file mode 100644 index 0000000..c4ef7bb --- /dev/null +++ b/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 diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 27f990b..69537d4 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -298,15 +298,16 @@ class SDTrainer(BaseSDTrainProcess): # make sure model is on cpu for this part so we don't oom. self.sd.unet.to('cpu') - # cache unconditional embeds (blank prompt) - with torch.no_grad(): - self.unconditional_embeds = self.encode_static_prompt( - [self.train_config.unconditional_prompt], - long_prompts=self.do_long_prompts, - ).to( - self.device_torch, - dtype=self.sd.torch_dtype - ).detach() + # cache unconditional embeds (blank prompt); text-generating models have no text encoder + if not getattr(self.sd, 'is_llm', False): + with torch.no_grad(): + self.unconditional_embeds = self.encode_static_prompt( + [self.train_config.unconditional_prompt], + long_prompts=self.do_long_prompts, + ).to( + self.device_torch, + dtype=self.sd.torch_dtype + ).detach() if self.train_config.do_prior_divergence: self.do_prior_prediction = True @@ -314,14 +315,16 @@ class SDTrainer(BaseSDTrainProcess): # D-OPSD: the teacher (prior) prediction is the training target self.do_prior_prediction = True # move vae to device if we did not cache latents - if not self.is_latents_cached: - self.sd.vae.eval() - self.sd.vae.to(self.device_torch) - else: - # offload it. Already cached - self.sd.vae.to('cpu') - flush() - add_all_snr_to_noise_scheduler(self.sd.noise_scheduler, self.device_torch) + if self.sd.vae is not None: + if not self.is_latents_cached: + self.sd.vae.eval() + self.sd.vae.to(self.device_torch) + else: + # offload it. Already cached + self.sd.vae.to('cpu') + flush() + if self.sd.noise_scheduler is not None: + add_all_snr_to_noise_scheduler(self.sd.noise_scheduler, self.device_torch) if self.adapter is not None: self.adapter.to(self.device_torch) @@ -1418,10 +1421,33 @@ class SDTrainer(BaseSDTrainProcess): ) + def train_llm_accumulation(self, batch: DataLoaderBatchDTO, accum_scale: float = 1.0): + """Text-generating models (BaseModel.is_llm): no noise, scheduler, VAE or prompt + encoding. The model computes its own loss from the batch (cached media + captions) + and reports per-term logs through additional_loss_logs.""" + network = self.network if self.network is not None else BlankNetwork() + network.multiplier = batch.get_network_weight_list() + with torch.no_grad(): + loss_multiplier = torch.tensor(batch.loss_multiplier_list).to(self.device_torch, dtype=torch.float32) + with network: + with self.timer('llm_loss'): + loss = self.sd.get_llm_loss(batch) + self.additional_logs.update(getattr(self.sd, "additional_loss_logs", None) or {}) + if not torch.isfinite(loss): + print_acc("loss is nan") + loss = torch.zeros_like(loss).requires_grad_(True) + with self.timer('backward'): + loss = loss * loss_multiplier.mean() + # backward stays inside the network context (see the note in the diffusion path) + self.accelerator.backward(loss * accum_scale if accum_scale != 1.0 else loss) + return loss.detach() + def train_single_accumulation(self, batch: DataLoaderBatchDTO, accum_scale: float = 1.0): # accum_scale: 1 / number of micro-batches accumulated per optimizer step, so the # summed gradients equal the mean over the effective batch. Applied to the backward # only; the returned loss stays unscaled for logging. + if getattr(self.sd, 'is_llm', False): + return self.train_llm_accumulation(batch, accum_scale=accum_scale) with torch.no_grad(): self.timer.start('preprocess_batch') if isinstance(self.adapter, CustomAdapter): diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 3d3bcaf..0016134 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -747,7 +747,8 @@ class BaseSDTrainProcess(BaseTrainProcess): self.accelerator.even_batches=False # # prepare all the models stuff for accelerator (hopefully we dont miss any) - self.sd.vae = self.accelerator.prepare(self.sd.vae) + 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) # todo always tdo it? @@ -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,7 +1794,8 @@ class BaseSDTrainProcess(BaseTrainProcess): noise_scheduler = self.sd.noise_scheduler if self.train_config.xformers: - vae.enable_xformers_memory_efficient_attention() + if vae is not None: + vae.enable_xformers_memory_efficient_attention() unet.enable_xformers_memory_efficient_attention() if isinstance(text_encoder, list): for te in text_encoder: @@ -1868,15 +1871,16 @@ 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() - vae = vae.to(torch.device('cpu'), dtype=dtype) - vae.requires_grad_(False) - vae.eval() + if vae is not None: + vae = vae.to(torch.device('cpu'), dtype=dtype) + vae.requires_grad_(False) + vae.eval() if self.train_config.learnable_snr_gos: self.snr_gos = LearnableSNRGamma( self.sd.noise_scheduler, device=self.device_torch diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 14f8116..4976e8d 100644 --- a/toolkit/config_modules.py +++ b/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) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 5147870..298cc94 100644 --- a/toolkit/data_loader.py +++ b/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: diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index efad1f4..2f5c0ed 100644 --- a/toolkit/dataloader_mixins.py +++ b/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 - token_list = raw_caption.split(',') + # tag-style token ops (dropout / shuffle) split on commas; anything else keeps the + # 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() diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 4f314f5..b24aefa 100644 --- a/toolkit/models/base_model.py +++ b/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,10 +728,12 @@ class BaseModel: raise ValueError( "Refiner is only supported for XL models") - conditional_embeds = conditional_embeds.to( - self.device_torch, dtype=self.unet.dtype) - unconditional_embeds = unconditional_embeds.to( - self.device_torch, dtype=self.unet.dtype) + 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) img = self.generate_single_image( pipeline, @@ -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,11 +1529,12 @@ class BaseModel: self.device_state = None def set_device_state(self, state): - if state['vae']['training']: - self.vae.train() - else: - self.vae.eval() - self.vae.to(state['vae']['device']) + if self.vae is not None and state.get('vae') is not None: + if state['vae']['training']: + self.vae.train() + else: + self.vae.eval() + self.vae.to(state['vae']['device']) if state['unet']['training']: self.unet.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): diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index 18d90dd..403aca7 100644 --- a/toolkit/models/registry.py +++ b/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}}, } diff --git a/ui/src/app/api/jobs/[jobID]/samples/route.ts b/ui/src/app/api/jobs/[jobID]/samples/route.ts index 04c461b..fc1e911 100644 --- a/ui/src/app/api/jobs/[jobID]/samples/route.ts +++ b/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); diff --git a/ui/src/app/generate/page.tsx b/ui/src/app/generate/page.tsx index d64f07c..e930ace 100644 --- a/ui/src/app/generate/page.tsx +++ b/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 (
- {stoppingForMs > 8000 && ( - )} @@ -680,7 +738,11 @@ function GeneratePageInner() { Engine running {health?.active?.arch && · {health.active.arch}} {vram && · VRAM {vram}} - @@ -702,15 +764,28 @@ function GeneratePageInner() {

{isStarting ? (
- {engineStatus?.engine?.status === 'queued' && ( -
Waiting for the GPU queue to pick the job up. If the queue is stopped, start it from the Queue page or cancel here.
+
+ Waiting for the GPU queue to pick the job up. If the queue is stopped, start it from the Queue + page or cancel here. +
)} {startingSince !== null && now - startingSince > 5000 && ( - )} @@ -718,7 +793,11 @@ function GeneratePageInner() { ) : (
{!isMac() && ( - 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 => (
@@ -756,7 +838,10 @@ function GeneratePageInner() { ) : null}
{running && ( - )} @@ -766,7 +851,10 @@ function GeneratePageInner() { Engine busy: {health.current.arch} {health.current.total_steps ? ` ${health.current.step}/${health.current.total_steps}` : ''} - @@ -774,10 +862,17 @@ function GeneratePageInner() {
{progress?.total ? (
-
+
) : null} - {error &&
{error}
} + {error && ( +
+ {error} +
+ )} {/* live latent preview: shown while generating (or when nothing is selected) */}
@@ -791,7 +886,14 @@ function GeneratePageInner() { {!showPreview && selected && (
{selected.kind === 'video' ? ( -