diff --git a/extensions_built_in/captioner/Qwen3VLCaptioner.py b/extensions_built_in/captioner/Qwen3VLCaptioner.py index 9c76f223..61209b8b 100644 --- a/extensions_built_in/captioner/Qwen3VLCaptioner.py +++ b/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", diff --git a/toolkit/print.py b/toolkit/print.py index e0f6c23b..4d26d243 100644 --- a/toolkit/print.py +++ b/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):