Fix Flax LMS scheduler sigma mismatch and derivatives buffer OOB

Remove the erroneous +1e-5 sigma offset in step() so it matches
scale_model_input() and the PyTorch LMSDiscreteScheduler reference.
Cap the multistep fori_loop bound by the pre-allocated derivatives
buffer size to prevent out-of-bounds indexing when order exceeds
max_order from set_timesteps.

Author: cursoragent

src/diffusers/schedulers/scheduling_lms_discrete_flax.py

@@ -271,7 +271,7 @@ def step(
             )
 
         step_index = jnp.where(state.timesteps == timestep, jnp.arange(state.timesteps.shape[0]), 0).sum()
-        sigma = state.sigmas[step_index] + 1e-5
+        sigma = state.sigmas[step_index]
 
         # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
         if self.config.prediction_type == "epsilon":
@@ -291,9 +291,10 @@ def step(
 
         # 3. Compute linear multistep coefficients and the previous sample based on the derivatives path
         effective_order = jnp.minimum(step_index + 1, order)
+        loop_order = jnp.minimum(order, state.derivatives.shape[0])
         prev_sample = jax.lax.fori_loop(
             0,
-            order,
+            loop_order,
             lambda i, val: jnp.where(
                 i < effective_order,
                 val + self.get_lms_coefficient(state, effective_order, step_index, i) * state.derivatives[-(i + 1)],