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

Performance fixes and bug fixes

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-28 22:21:09 -06:00
parent 5ddc5f8ca7
commit 7195abc32b
21 changed files with 511 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -206,11 +206,13 @@ class ChromaModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: text encoders stay on cpu; get_prompt_embeds moves them
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[1].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
text_encoder[1].to(self.device_torch)
text_encoder[1].requires_grad_(False)
text_encoder[1].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

8
extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py
View File

@@ -192,11 +192,13 @@ class ChromaRadianceModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: text encoders stay on cpu; get_prompt_embeds moves them
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[1].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
text_encoder[1].to(self.device_torch)
text_encoder[1].requires_grad_(False)
text_encoder[1].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

6
extensions_built_in/diffusion_models/ernie_image/ernie_image.py
View File

@@ -133,8 +133,10 @@ class ErnieImageModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

6
extensions_built_in/diffusion_models/f_light/f_light.py
View File

@@ -115,8 +115,10 @@ class FLiteModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

13
extensions_built_in/diffusion_models/flux2/flux2_model.py
View File

@@ -98,12 +98,21 @@ class Flux2Model(BaseModel):
self.print_and_status_update("Loading Mistral")
# load + quantize + offload + placement, all driven by model_config
# tie_word_embeddings=False: the checkpoint carries both embed_tokens
# and lm_head with different values; the config's tie claim is wrong
text_encoder = Mistral3TextEncoder.load(
MISTRAL_PATH, subfolder="", **self.component_load_kwargs("te")
MISTRAL_PATH,
subfolder="",
tie_word_embeddings=False,
**self.component_load_kwargs("te"),
)
flush()
tokenizer = AutoProcessor.from_pretrained(MISTRAL_PATH)
# fix_mistral_regex=False: keep the exact tokenization flux2 has always
# used (True would change the pre-tokenizer and shift conditioning)
tokenizer = AutoProcessor.from_pretrained(
MISTRAL_PATH, fix_mistral_regex=False
)
return text_encoder, tokenizer
def load_model(self):

10
extensions_built_in/diffusion_models/flux2/src/pipeline.py
View File

@@ -109,15 +109,19 @@ class Flux2Pipeline(DiffusionPipeline):
# Process all messages at once
# with image processing a too short max length can throw an error in here.
try:
# tokenization kwargs ride in processor_kwargs (same values end up
# in the same place; loose **kwargs just warn on new transformers)
inputs = self.tokenizer.apply_chat_template(
messages_batch,
add_generation_prompt=False,
tokenize=True,
return_dict=True,
return_tensors="pt",
padding="max_length",
truncation=True,
max_length=max_sequence_length,
processor_kwargs={
"padding": "max_length",
"truncation": True,
"max_length": max_sequence_length,
},
)
except ValueError as e:
print(

8
extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
View File

@@ -137,11 +137,13 @@ class FluxKontextModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: text encoders stay on cpu; get_prompt_embeds moves them
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[1].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
text_encoder[1].to(self.device_torch)
text_encoder[1].requires_grad_(False)
text_encoder[1].eval()
pipe.transformer = pipe.transformer.to(self.device_torch)

6
extensions_built_in/diffusion_models/ltx2/ltx2.py
View File

@@ -490,8 +490,10 @@ class LTX2Model(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

6
extensions_built_in/diffusion_models/nucleus_image/nucleus_image_model.py
View File

@@ -144,8 +144,10 @@ class NucleusImageModel(QwenImageVAEHolderMixin, BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

6
extensions_built_in/diffusion_models/qwen_image/qwen_image.py
View File

@@ -161,8 +161,10 @@ class QwenImageModel(QwenImageVAEHolderMixin, BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()

20
extensions_built_in/diffusion_models/wan22/wan22_14b_model.py
View File

@@ -95,6 +95,13 @@ class DualWanTransformer3DModel(torch.nn.Module):
def dtype(self) -> torch.dtype:
return self.torch_dtype
def get_offload_ignore_modules(self):
# both DiTs' fp32 modulation tables stay resident on the compute device
return (
self.transformer_1.get_offload_ignore_modules()
+ self.transformer_2.get_offload_ignore_modules()
)
@property
def config(self):
return self.transformer_1.config
@@ -128,17 +135,26 @@ class DualWanTransformer3DModel(torch.nn.Module):
else:
t_name = "transformer_2"
# memory-managed transformers own their placement (layers bounce
# from cpu per forward); whole-model swaps would haul the full
# 14B up and defeat the offloading
managed = (
hasattr(self, "_memory_manager")
or hasattr(self.transformer_1, "_memory_manager")
or hasattr(self.transformer_2, "_memory_manager")
)
# check if we are changing the active transformer, if so, we need to swap the one in
# vram if low_vram is enabled
# todo swap the loras as well
if t_name != self._active_transformer_name:
if self.low_vram:
if self.low_vram and not managed:
getattr(self, self._active_transformer_name).to("cpu")
getattr(self, t_name).to(self.device_torch)
torch.cuda.empty_cache()
self._active_transformer_name = t_name
if self.transformer.device != hidden_states.device:
if not managed and self.transformer.device != hidden_states.device:
if self.low_vram:
# move other transformer to cpu
other_tname = (

6
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py
View File

@@ -149,8 +149,10 @@ class ZetaChromaModel(BaseModel):
pipe.transformer = pipe.transformer.to(self.device_torch)
flush()
# just to make sure everything is on the right device and dtype
text_encoder[0].to(self.device_torch)
# low_vram: the text encoder stays on cpu; get_prompt_embeds moves it
# to the gpu on demand
if not self.low_vram:
text_encoder[0].to(self.device_torch)
text_encoder[0].requires_grad_(False)
text_encoder[0].eval()
flush()