Added flux_shift as timestep type

This commit is contained in:
Jaret Burkett committed 2025-01-27 07:35:00 -07:00
1 parent 2141c6e06c
commit 34a1c6947a
3 files changed
+64 -9

No files matched your search

+6 -6
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,