Division by 0 bug when only one denoising step

[kyle-hoffmeyer]

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:

diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

Line 241 in 9c6a68c

def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:

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:

diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

Line 413 in 9c6a68c

return indices[pos].item()

, 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

  [rank0]: Traceback (most recent call last):
  [rank0]:   File "<frozen runpy>", line 198, in _run_module_as_main
  [rank0]:   File "<frozen runpy>", line 88, in _run_code
  [rank0]:   File "/app/xDiT/xfuser/runner.py", line 123, in <module>
  [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