Add support for training VLLMs, omni models, and LLMs

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
-        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()

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):

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}},
 }