huggingface / huggingface/diffusers

Division by 0 bug when only one denoising step

Aperta Adatta ai principianti
#14,411 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### Describe the bug

When `run_timed_pipe` is called: https://github.com/xdit-project/xDiT/blob/9705ed0d7509537b4d2697f27c089bab5c6e96a3/xfuser/model_executor/models/runner_models/ltx.py#L198, an array `sigmas` is created that represents the denoising steps. 1.0 indicates pure noise and 0.0 indicates a clean output. `sigmas` always starts with 1.0. `sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)`: https://github.com/xdit-project/xDiT/blob/9705ed0d7509537b4d2697f27c089bab5c6e96a3/xfuser/model_executor/pipelines/pipeline_flux.py#L277.

`num_inference_steps` is a parameter in xDiT that controls the number of denoising steps.
- num_inference_steps=2: sigmas = [1.0, 0.5] - two points from full noise to half noise
- num_inference_steps=1: sigmas = [1.0] — a single point at full noise

Dynamic time shifting is applied to `sigmas` via `time_shift` function: https://github.com/huggingface/diffusers/blob/9c6a68c32b3b2a64db91800b624d33cec6e25ab8/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py#L241

After `time_shift`, `sigmas` is compressed and `stretch_shift_to_terminal` linearly rescales the schedule in completion space so the last sigma lands at `shift_terminal` (often 0.1).

**The issue is that `stretch_shift_to_terminal` assumes there to be multiple denoising steps and sets `sigmas` to be [nan] when `num_inference_steps = 1`**:
```
def stretch_shift_to_terminal(self, t):
one_minus_z = 1 - t # 1 - 1 = 0
scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal) -> # 0 / (1 - 0.1) = 0
stretched_t = 1 - (one_minus_z / scale_factor) # 1 - (0 / 0). Division by 0 -> nan
return stretched_t -> nan
```
When `sigmas` is indexed: https://github.com/huggingface/diffusers/blob/9c6a68c32b3b2a64db91800b624d33cec6e25ab8/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py#L413, since it contains `nan`, it errors.

The starting point in a diffusion pipeline is always pure noise and should be 1.0. If there's only one denoising step the denoising should go from 1.0 to 0.0 and `stretch_shift_to_terminal` should be skipped since there's nothing to shift.

```
if self.config.shift_terminal:
sigmas = self.stretch_shift_to_terminal(sigmas)
```

should be changed to

```
if self.config.shift_terminal and len(sigmas) > 1:
sigmas = self.stretch_shift_to_terminal(sigmas)
```

### Reproduction

Modify https://github.com/xdit-project/xDiT/blob/9705ed0d7509537b4d2697f27c089bab5c6e96a3/xfuser/model_executor/models/runner_models/ltx.py#L197 to set `compile_args["num_inference_steps"] = 1`. For context, I am working on a PR to xDiT that involves running `_run_timed_pipe` once with `compile_args["num_inference_steps"] = 1` to profile the CUDA graph capture phase.

Run:
```
xdit \
--model "LTX-2.3" \
--seed "42" \
--prompt "a cat" \
--height "256" \
--width "256" \
--num_inference_steps "4" \
--max_sequence_length "512" \
--warmup_calls "0" \
--ulysses_degree "1" \
--guidance_scale "4" \
--num_iterations "4" \
--attention_backend "aiter" \
--input_images "/app/data/flux_cat.png" \
--output_directory "${XDIT_RUN_DIR}" \
--use_torch_compile \
--profile --profile_wait 1 --profile_warmup 2 --profile_active 1
```

### Logs

```shell
[rank0]: Traceback (most recent call last):
[rank0]: File "", line 198, in _run_module_as_main
[rank0]: File "", line 88, in _run_code
[rank0]: File "/app/xDiT/xfuser/runner.py", line 123, in
[rank0]: runner.initialize(input_args)
[rank0]: File "/app/xDiT/xfuser/runner.py", line 61, in initialize
[rank0]: self.model.initialize(input_args)
[rank0]: File "/app/xDiT/xfuser/model_executor/models/runner_models/base_model.py", line 273, in initialize
[rank0]: self._compile_model(compile_input_args)
[rank0]: File "/app/xDiT/xfuser/model_executor/models/runner_models/ltx.py", line 194, in _compile_model
[rank0]: self._run_compile_warmup(compile_args)
[rank0]: File "/app/xDiT/xfuser/model_executor/models/runner_models/base_model.py", line 645, in _run_compile_warmup
[rank0]: self._run_timed_pipe(compile_args)
[rank0]: File "/app/xDiT/xfuser/model_executor/models/runner_models/base_model.py", line 614, in _run_timed_pipe
[rank0]: out = self._run_pipe(input_args)
[rank0]: File "/app/xDiT/xfuser/model_executor/models/runner_models/ltx.py", line 127, in _run_pipe
[rank0]: video_latent, audio_latent = self.pipe(
[rank0]: File "/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank0]: return func(*args, **kwargs)
[rank0]: File "/app/external/diffusers/src/diffusers/pipelines/ltx2/pipeline_ltx2.py", line 1402, in __call__
[rank0]: latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
[rank0]: File "/app/external/diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 483, in step
[rank0]: self._init_step_index(timestep)
[rank0]: File "/app/external/diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 421, in _init_step_index
[rank0]: self._step_index = self.index_for_timestep(timestep)
[rank0]: File "/app/external/diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 415, in index_for_timestep
[rank0]: return indices[pos].item()
[rank0]: ~~~~~~~^^^^^
[rank0]: IndexError: index 0 is out of bounds for dimension 0 with size 0
```

### System Info

- 🤗 Diffusers version: 0.38.0.dev0
- Platform: Linux-6.8.0-31-generic-x86_64-with-glibc2.39
- Running on Google Colab?: No
- Python version: 3.12.3
- PyTorch version (GPU?): 2.9.1+gitff65f5b (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 1.20.1
- Transformers version: 5.5.4
- Accelerate version: 1.14.0
- PEFT version: 0.19.1
- Bitsandbytes version: not installed
- Safetensors version: 0.8.0
- xFormers version: not installed
- Accelerator: NA
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

### Who can help?

_No response_

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Start with diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py, especially stretch_shift_to_terminal, set_timesteps, and index_for_timestep; use xfuser/model_executor/models/runner_models/ltx.py to reproduce the one-step path. Run the reported LTX-2.3 command with num_inference_steps=1 and verify that the scheduler produces a valid schedule without NaN values or an IndexError, while multi-step behavior remains unchanged.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
2/5
Tempo stimato
1-3 ore
Stato di attività
Tranquilla
Chiarezza
Specificata chiaramente
Idoneità per principianti
72/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.