base_schedule

Author: hlky

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -820,7 +820,7 @@ def __call__(
         )
 
         # 5. Prepare timesteps
-        if self.scheduler.schedule.__class__.__name__ != "FlowMatchFlux":
+        if self.scheduler._schedule.base_schedule.__class__.__name__ != "FlowMatchFlux":
             self.scheduler._schedule.set_base_schedule("FlowMatchFlux")
         image_seq_len = latents.shape[1]
         mu = calculate_shift(

src/diffusers/pipelines/mochi/pipeline_mochi.py

@@ -652,7 +652,7 @@ def __call__(
             prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
 
         # 5. Prepare timestep
-        if self.scheduler.schedule.__class__.__name__ != "FlowMatchLinearQuadratic":
+        if self.scheduler._schedule.base_schedule.__class__.__name__ != "FlowMatchLinearQuadratic":
             self.scheduler._schedule.set_base_schedule("FlowMatchLinearQuadratic")
 
         timesteps, num_inference_steps = retrieve_timesteps(