Merge branch 'main' of github.com:ostris/ai-toolkit

extensions_built_in/captioner/Qwen3VLCaptioner.py

@@ -1,4 +1,8 @@
-from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
+from transformers import (
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMoeForConditionalGeneration,
+    AutoProcessor,
+)
 from collections import OrderedDict
 
 from optimum.quanto import freeze
@@ -21,7 +25,12 @@ class Qwen3VLCaptioner(BaseCaptioner):
 
     def load_model(self):
         self.print_and_status_update("Loading Qwen3VL model")
-        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
+        ModelClass = (
+            Qwen3VLMoeForConditionalGeneration
+            if "B-A" in self.caption_config.model_name_or_path
+            else Qwen3VLForConditionalGeneration
+        )
+        self.model = ModelClass.from_pretrained(
             self.caption_config.model_name_or_path,
             dtype=self.torch_dtype,
             device_map="cpu",

toolkit/print.py

@@ -21,6 +21,9 @@ class Logger:
     def flush(self):
         self.terminal.flush()
         self.log.flush()
+    
+    def isatty(self):
+        return self.terminal.isatty()
 
 
 def setup_log_to_file(filename):