github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2026-04-17 06:05:40 -06:00
parent be54094704 a513a1583e
commit 67048df9f9
2 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
extensions_built_in/captioner/Qwen3VLCaptioner.py
View File

@@ -1,4 +1,8 @@
from transformers import Qwen3VLForConditionalGeneration, AutoProcessor from transformers import (
Qwen3VLForConditionalGeneration,
Qwen3VLMoeForConditionalGeneration,
AutoProcessor,
)
from collections import OrderedDict from collections import OrderedDict
from optimum.quanto import freeze from optimum.quanto import freeze
@@ -21,7 +25,12 @@ class Qwen3VLCaptioner(BaseCaptioner):
def load_model(self): def load_model(self):
self.print_and_status_update("Loading Qwen3VL model") 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, self.caption_config.model_name_or_path,
dtype=self.torch_dtype, dtype=self.torch_dtype,
device_map="cpu", device_map="cpu",

3
toolkit/print.py
View File

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