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

Initial support for finetuning qwen image. Will only work with caching for now, need to add controls everywhere.

Browse Source
This commit is contained in:
Jaret Burkett
2025-08-21 16:41:17 -06:00
parent 38d3814be7
commit bf2700f7be
12 changed files with 399 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
toolkit/dataloader_mixins.py
View File

@@ -30,6 +30,7 @@ import albumentations as A
from toolkit.print import print_acc
from toolkit.accelerator import get_accelerator
from toolkit.prompt_utils import PromptEmbeds
from torchvision.transforms import functional as TF
from toolkit.train_tools import get_torch_dtype
@@ -1802,6 +1803,9 @@ class TextEmbeddingFileItemDTOMixin:
("text_embedding_space_version", self.text_embedding_space_version),
("text_embedding_version", self.text_embedding_version),
])
# if we have a control image, cache the path
if self.encode_control_in_text_embeddings and self.control_path is not None:
item["control_path"] = self.control_path
return item
def get_text_embedding_path(self: 'FileItemDTO', recalculate=False):
@@ -1860,7 +1864,19 @@ class TextEmbeddingCachingMixin:
if not did_move:
self.sd.set_device_state_preset('cache_text_encoder')
did_move = True
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption)
if file_item.encode_control_in_text_embeddings and file_item.control_path is not None:
# load the control image and feed it into the text encoder
ctrl_img = Image.open(file_item.control_path).convert("RGB")
# convert to 0 to 1 tensor
ctrl_img = (
TF.to_tensor(ctrl_img)
.unsqueeze(0)
.to(self.sd.device_torch, dtype=self.sd.torch_dtype)
)
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption, control_images=ctrl_img)
else:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption)
# save it
prompt_embeds.save(text_embedding_path)
del prompt_embeds