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

42
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -37,6 +37,8 @@ from toolkit.models.diffusion_feature_extraction import DiffusionFeatureExtracto
from toolkit.util.wavelet_loss import wavelet_loss
import torch.nn.functional as F
from toolkit.unloader import unload_text_encoder
from PIL import Image
from torchvision.transforms import functional as TF
def flush():
@@ -127,9 +129,28 @@ class SDTrainer(BaseSDTrainProcess):
prompt=prompt, # it will autoparse the prompt
negative_prompt=sample_item.neg,
output_path=output_path,
ctrl_img=sample_item.ctrl_img
)
positive = self.sd.encode_prompt(gen_img_config.prompt).to('cpu')
negative = self.sd.encode_prompt(gen_img_config.negative_prompt).to('cpu')
# see if we need to encode the control images
if self.sd.encode_control_in_text_embeddings and gen_img_config.ctrl_img is not None:
ctrl_img = Image.open(gen_img_config.ctrl_img).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)
)
positive = self.sd.encode_prompt(
gen_img_config.prompt,
control_images=ctrl_img
).to('cpu')
negative = self.sd.encode_prompt(
gen_img_config.negative_prompt,
control_images=ctrl_img
).to('cpu')
else:
positive = self.sd.encode_prompt(gen_img_config.prompt).to('cpu')
negative = self.sd.encode_prompt(gen_img_config.negative_prompt).to('cpu')
self.sd.sample_prompts_cache.append({
'conditional': positive,
@@ -177,9 +198,15 @@ class SDTrainer(BaseSDTrainProcess):
# cache unconditional embeds (blank prompt)
with torch.no_grad():
kwargs = {}
if self.sd.encode_control_in_text_embeddings:
# just do a blank image for unconditionals
control_image = torch.zeros((1, 3, 224, 224), device=self.sd.device_torch, dtype=self.sd.torch_dtype)
kwargs['control_images'] = control_image
self.unconditional_embeds = self.sd.encode_prompt(
[self.train_config.unconditional_prompt],
long_prompts=self.do_long_prompts
long_prompts=self.do_long_prompts,
**kwargs
).to(
self.device_torch,
dtype=self.sd.torch_dtype
@@ -241,9 +268,14 @@ class SDTrainer(BaseSDTrainProcess):
print_acc("***********************************")
print_acc("")
self.sd.text_encoder_to(self.device_torch)
self.cached_blank_embeds = self.sd.encode_prompt("")
encode_kwargs = {}
if self.sd.encode_control_in_text_embeddings:
# just do a blank image for unconditionals
control_image = torch.zeros((1, 3, 224, 224), device=self.sd.device_torch, dtype=self.sd.torch_dtype)
encode_kwargs['control_images'] = control_image
self.cached_blank_embeds = self.sd.encode_prompt("", **encode_kwargs)
if self.trigger_word is not None:
self.cached_trigger_embeds = self.sd.encode_prompt(self.trigger_word)
self.cached_trigger_embeds = self.sd.encode_prompt(self.trigger_word, **encode_kwargs)
if self.train_config.diff_output_preservation:
self.diff_output_preservation_embeds = self.sd.encode_prompt(self.train_config.diff_output_preservation_class)