@@ -182,6 +182,74 @@ def get_dummy_inputs(self):
182182
183183
184184class TestQwenImage21Pipeline (QwenImage21PipelineTesterConfig , PipelineTesterMixin ):
185+ @pytest .mark .parametrize ("num_inference_steps" , [None , 2 ])
186+ def test_sample_sigmas_from_pipeline_config (self , num_inference_steps ):
187+ pipe = self .get_pipeline ()
188+ sample_sigmas = [1.0 , 0.978453 , 0.954180 , 0.926626 , 0.895080 , 0.845148 , 0.704534 , 0.414568 ]
189+ pipe .scheduler = FlowMatchEulerDiscreteScheduler .from_config (
190+ pipe .scheduler .config ,
191+ use_dynamic_shifting = False ,
192+ shift = 1.0 ,
193+ shift_terminal = None ,
194+ )
195+ pipe = self .pipeline_class (** pipe .components , sample_sigmas = sample_sigmas )
196+ inputs = self .get_dummy_inputs ()
197+ inputs ["output_type" ] = "latent"
198+ if num_inference_steps is None :
199+ inputs .pop ("num_inference_steps" )
200+ else :
201+ inputs ["num_inference_steps" ] = num_inference_steps
202+ seen_timesteps = []
203+
204+ def callback (pipe , step , timestep , callback_kwargs ):
205+ seen_timesteps .append (timestep .clone ())
206+ return callback_kwargs
207+
208+ pipe (** inputs , callback_on_step_end = callback )
209+ expected_sigmas = torch .tensor (sample_sigmas + [0.0 ], device = pipe .scheduler .sigmas .device )
210+ expected_timesteps = expected_sigmas [:- 1 ] * pipe .scheduler .config .num_train_timesteps
211+ assert torch .equal (pipe .scheduler .sigmas , expected_sigmas )
212+ assert torch .equal (pipe .scheduler .timesteps , expected_timesteps )
213+ assert torch .equal (torch .stack (seen_timesteps ), expected_timesteps )
214+ assert pipe .scheduler .num_inference_steps == len (sample_sigmas )
215+
216+ def test_explicit_sigmas_override_pipeline_config (self ):
217+ pipe = self .get_pipeline ()
218+ pipe .scheduler = FlowMatchEulerDiscreteScheduler .from_config (
219+ pipe .scheduler .config , use_dynamic_shifting = True , shift_terminal = 0.02
220+ )
221+ inputs = self .get_dummy_inputs ()
222+ inputs .update (sigmas = [1.0 , 0.6 , 0.2 ], output_type = "latent" )
223+ expected = pipe (** inputs ).images
224+ expected_sigmas = pipe .scheduler .sigmas .clone ()
225+ pipe .register_to_config (sample_sigmas = [1.0 , 0.8 ])
226+ inputs ["generator" ] = self .get_generator (0 )
227+ actual = pipe (** inputs ).images
228+ assert torch .equal (pipe .scheduler .sigmas , expected_sigmas )
229+ assert torch .equal (actual , expected )
230+
231+ def test_sample_sigmas_save_load (self , tmp_path ):
232+ pipe = self .get_pipeline ()
233+ sample_sigmas = [1.0 , 0.8 , 0.4 ]
234+ pipe .scheduler = FlowMatchEulerDiscreteScheduler .from_config (
235+ pipe .scheduler .config , use_dynamic_shifting = False , shift = 1.0 , shift_terminal = None
236+ )
237+ pipe = self .pipeline_class (** pipe .components , sample_sigmas = sample_sigmas )
238+ inputs = self .get_dummy_inputs ()
239+ inputs ["output_type" ] = "latent"
240+ expected = pipe (** inputs ).images
241+ pipe .save_pretrained (tmp_path )
242+ restored = self .pipeline_class .from_pretrained (tmp_path )
243+ assert restored .config .sample_sigmas == sample_sigmas
244+ assert "sample_sigmas" not in restored .scheduler .config
245+ inputs ["generator" ] = self .get_generator (0 )
246+ actual = restored (** inputs ).images
247+ assert torch .equal (actual , expected )
248+ overridden = self .pipeline_class .from_pretrained (tmp_path , sample_sigmas = [1.0 , 0.5 ])
249+ inputs ["generator" ] = self .get_generator (0 )
250+ overridden (** inputs )
251+ assert torch .equal (overridden .scheduler .sigmas .cpu (), torch .tensor ([1.0 , 0.5 , 0.0 ]))
252+
185253 @PROCESSOR_REQUIRED_AT_INIT
186254 def test_encode_prompt_works_in_isolation (self ):
187255 super ().test_encode_prompt_works_in_isolation ()
0 commit comments