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

Added flux_shift as timestep type

Browse Source
This commit is contained in:
Jaret Burkett
2025-01-27 07:35:00 -07:00
parent 2141c6e06c
commit 34a1c6947a
3 changed files with 64 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
toolkit/stable_diffusion_model.py
View File

@@ -1027,9 +1027,9 @@ class StableDiffusion:
if self.model_config.use_flux_cfg:
pipeline = FluxWithCFGPipeline(
vae=self.vae,
transformer=self.unet,
text_encoder=self.text_encoder[0],
text_encoder_2=self.text_encoder[1],
transformer=unwrap_model(self.unet),
text_encoder=unwrap_model(self.text_encoder[0]),
text_encoder_2=unwrap_model(self.text_encoder[1]),
tokenizer=self.tokenizer[0],
tokenizer_2=self.tokenizer[1],
scheduler=noise_scheduler,
@@ -1039,9 +1039,9 @@ class StableDiffusion:
else:
pipeline = FluxPipeline(
vae=self.vae,
transformer=self.unet,
text_encoder=self.text_encoder[0],
text_encoder_2=self.text_encoder[1],
transformer=unwrap_model(self.unet),
text_encoder=unwrap_model(self.text_encoder[0]),
text_encoder_2=unwrap_model(self.text_encoder[1]),
tokenizer=self.tokenizer[0],
tokenizer_2=self.tokenizer[1],
scheduler=noise_scheduler,