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)