Add support for training VLLMs, omni models, and LLMs

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