Skip to content

Commit da1d382

Browse files
authored
Configure fixed sampling sigmas on QwenImage21Pipeline (#14950)
* feat(qwenimage21): configure sampling sigmas on the pipeline * docs(qwenimage21): clarify default and runtime sampling sigmas
1 parent c2798cc commit da1d382

3 files changed

Lines changed: 103 additions & 3 deletions

File tree

‎docs/source/en/api/pipelines/qwenimage21.md‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,26 @@ pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", transformer=tr
7777
)
7878
```
7979

80+
## Sampling sigmas
81+
82+
Model authors can configure a default sampling grid with `sample_sigmas` in the pipeline config. When you load a
83+
released checkpoint, its default grid and scheduler settings are restored automatically.
84+
85+
To experiment with a different grid at runtime, pass `sigmas` to the pipeline call:
86+
87+
```python
88+
# Use the checkpoint's default sampling grid.
89+
image = pipe(prompt).images[0]
90+
91+
# Override the default grid for this call.
92+
image = pipe(prompt, sigmas=[1.0, 0.8, 0.5, 0.2]).images[0]
93+
```
94+
95+
The custom grid above illustrates the API; generation quality depends on the checkpoint and grid. Sigma lists exclude
96+
the terminal sigma, which the scheduler appends. Explicit `sigmas` override the configured `sample_sigmas`, and either
97+
list determines the number of steps instead of `num_inference_steps`. If neither is provided, the pipeline uses
98+
`num_inference_steps` to generate the schedule. The scheduler applies its configured processing to either grid.
99+
80100
## QwenImage21Pipeline
81101

82102
[[autodoc]] QwenImage21Pipeline

‎src/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,10 @@ class QwenImage21Pipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
174174
Processor that builds the chat template and tokenizes prompt and condition images.
175175
transformer ([`QwenImage21Transformer2DModel`]):
176176
The single-stream block-causal transformer that denoises the latents.
177+
sample_sigmas (`list[float]`, *optional*):
178+
Default sampling sigmas configured by the model author, excluding the terminal sigma. Their length
179+
determines the number of denoising steps. To experiment with a different grid at runtime, pass `sigmas` to
180+
`__call__`.
177181
"""
178182

179183
model_cpu_offload_seq = "text_encoder->transformer->vae"
@@ -186,6 +190,7 @@ def __init__(
186190
text_encoder: Qwen3VLForConditionalGeneration,
187191
processor: Qwen3VLProcessor,
188192
transformer: QwenImage21Transformer2DModel,
193+
sample_sigmas: list[float] | None = None,
189194
):
190195
super().__init__()
191196

@@ -196,6 +201,7 @@ def __init__(
196201
transformer=transformer,
197202
scheduler=scheduler,
198203
)
204+
self.register_to_config(sample_sigmas=sample_sigmas)
199205
# The VAE compresses 16x spatially and the transformer consumes latents unpatched, so one token covers a 16x16
200206
# pixel tile.
201207
self.vae_scale_factor = 16
@@ -547,9 +553,12 @@ def __call__(
547553
width (`int`, *optional*):
548554
Width in pixels of the generated image. Derived from the condition image's aspect ratio if omitted.
549555
num_inference_steps (`int`, *optional*, defaults to 40):
550-
Number of denoising steps.
556+
Number of denoising steps. Ignored when `sigmas` or the pipeline's configured `sample_sigmas` is used;
557+
the length of that schedule determines the number of steps.
551558
sigmas (`list[float]`, *optional*):
552-
Custom sigmas for the denoising schedule.
559+
Sampling sigmas to try for this call, excluding the terminal sigma. Overrides the pipeline's configured
560+
`sample_sigmas` and determines the number of denoising steps. The scheduler applies its configured
561+
processing to these values.
553562
num_images_per_prompt (`int`, *optional*, defaults to 1):
554563
Number of images generated per prompt.
555564
generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
@@ -720,7 +729,10 @@ def __call__(
720729
] * batch_size
721730

722731
# 4. Prepare timesteps
723-
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
732+
if sigmas is None:
733+
sigmas = self.config.sample_sigmas
734+
if sigmas is None:
735+
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
724736
mu = calculate_shift(
725737
latents.shape[1],
726738
self.scheduler.config.get("base_image_seq_len", 256),

‎tests/pipelines/qwenimage21/test_qwenimage21.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,74 @@ def get_dummy_inputs(self):
182182

183183

184184
class 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

Comments
 (0)