Fix generator determinism: forward generator through scheduler steps and seeded noise

LightningDit-XL-1-256/model_index.json

@@ -5,8 +5,8 @@
   ],
   "_diffusers_version": "0.36.0",
   "scheduler": [
-    "diffusers",
-    "FlowMatchHeunDiscreteScheduler"
+    "scheduling_flow_match_lightningdit",
+    "LightningDiTFlowMatchScheduler"
   ],
   "vae": [
     "diffusers",

LightningDit-XL-1-256/pipeline.py

@@ -1,131 +1,89 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Hub custom pipeline for class-conditional LightningDiT image generation.
-
-Load with native Hugging Face diffusers via ``DiffusionPipeline.from_pretrained`` and
-``trust_remote_code=True``.
+"""Hub custom pipeline: LightningDiTPipeline.
+Load with native Hugging Face diffusers and trust_remote_code=True.
 """
 
 from __future__ import annotations
 
 import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
 
-import torch
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
 
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
-from diffusers.schedulers import KarrasDiffusionSchedulers
-from diffusers.utils import replace_example_docstring
-from diffusers.utils.torch_utils import randn_tensor
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> from pathlib import Path
-        >>> import torch
-        >>> from diffusers import DiffusionPipeline
-
-        >>> model_dir = Path("BiliSakura/LightningDiT-diffusers/LightningDit-XL-1-256")
-        >>> pipe = DiffusionPipeline.from_pretrained(
-        ...     str(model_dir),
-        ...     local_files_only=True,
-        ...     custom_pipeline=str(model_dir / "pipeline.py"),
-        ...     trust_remote_code=True,
-        ...     torch_dtype=torch.bfloat16,
-        ... ).to("cuda")
-
-        >>> class_id = pipe.get_label_ids("golden retriever")[0]
-        >>> image = pipe(
-        ...     class_labels=class_id,
-        ...     num_inference_steps=250,
-        ...     guidance_scale=6.7,
-        ...     cfg_interval_start=0.125,
-        ...     generator=torch.Generator(device="cuda").manual_seed(0),
-        ... ).images[0]
-        ```
-"""
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union, Any
+
+import json
+import torch
 
+try:
+    from diffusers.image_processor import VaeImageProcessor
+    from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+    from diffusers.utils import BaseOutput
+    from diffusers.utils.torch_utils import randn_tensor
+except Exception:  # pragma: no cover
+    class BaseOutput(dict):
+        def __post_init__(self):
+            self.update(self.__dict__)
+
+    class DiffusionPipeline:
+        def register_modules(self, **kwargs):
+            for name, module in kwargs.items():
+                setattr(self, name, module)
+
+        @property
+        def _execution_device(self):
+            return torch.device("cpu")
+
+        def maybe_free_model_hooks(self):
+            pass
+
+    class VaeImageProcessor:
+        def postprocess(self, image, output_type="pil"):
+            return image
 
-def _uses_explicit_next_timestep_scheduler(scheduler: KarrasDiffusionSchedulers) -> bool:
-    """True for LightningDiTFlowMatchScheduler (explicit t, t_next); False for built-in FlowMatch schedulers."""
-    try:
-        return "next_timestep" in inspect.signature(scheduler.step).parameters
-    except (TypeError, ValueError):
-        return False
+    def randn_tensor(shape, generator=None, device=None, dtype=None):
+        return torch.randn(shape, generator=generator, device=device, dtype=dtype)
 
+@dataclass
+class LightningDiTPipelineOutput(BaseOutput):
+    images: Union[torch.FloatTensor, List]
 
 class LightningDiTPipeline(DiffusionPipeline):
     r"""
-    Pipeline for class-conditional image generation with [LightningDiT](https://github.com/hustvl/LightningDiT).
-
-    Uses VA-VAE latents and flow-matching velocity prediction. The bundled checkpoint defaults to
-    [`FlowMatchHeunDiscreteScheduler`] with `shift=0.3` (2nd-order Heun). Flow time passed to the
-    transformer is `1 - sigma` (`t=0` noise, `t=1` data). Latents are denormalized from VAE
-    `latents_mean` / `latents_std` before decode.
-
-    Recommended settings for `LightningDiT-XL/1` ImageNet-256 (800 epochs), matching official inference:
-
-    - `num_inference_steps=250`
-    - `guidance_scale=6.7`
-    - `cfg_interval_start=0.125`
-    - `cfg_channels=3`
-    - `timestep_shift=0.3` (only when the scheduler supports `set_shift`; otherwise set `shift` in
-      `scheduler/scheduler_config.json`)
-
-    Parameters:
-        transformer ([`LightningDiTTransformer2DModel`]):
-            LightningDiT transformer predicting flow-matching velocity in latent space.
-        scheduler ([`FlowMatchHeunDiscreteScheduler`]):
-            Flow-matching scheduler. Other [`KarrasDiffusionSchedulers`] may be swapped at load time.
-        vae ([`AutoencoderKL`]):
-            VA-VAE used to decode latents to pixels.
-        id2label (`dict[int, str]`, *optional*):
-            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
-    """
+    Class-conditional image generation with LightningDiT and a flow-matching scheduler.
 
-    model_cpu_offload_seq = "transformer->vae"
+    Components are stored in separate subfolders (`transformer`, `scheduler`, optional `vae`) for
+    `DiffusionPipeline.from_pretrained` compatibility.
+    """
 
-    def __init__(
-        self,
-        transformer,
-        vae,
+    @staticmethod
+    def prepare_extra_step_kwargs(
         scheduler,
-        id2label=None,
-        null_class_id=None,
+        generator=None,
+        eta: float | None = None,
     ):
-        super().__init__()
-        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        kwargs = {}
+        step_params = set(inspect.signature(scheduler.step).parameters.keys())
+        if "generator" in step_params:
+            kwargs["generator"] = generator
+        if eta is not None and "eta" in step_params:
+            kwargs["eta"] = eta
+        return kwargs
+
 
-        if null_class_id is None:
-            null_class_id = int(getattr(self.transformer.config, "num_classes", 1000))
-        self.register_to_config(null_class_id=int(null_class_id))
+    model_cpu_offload_seq = "transformer->vae"
+    _optional_components = ["vae"]
 
+    def __init__(self, transformer, scheduler, vae=None, id2label: Optional[Dict[Union[int, str], str]] = None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
+        self.image_processor = VaeImageProcessor()
         self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
-
-    @property
-    def vae_scale_factor(self) -> int:
-        block_out_channels = getattr(self.vae.config, "block_out_channels", None)
-        if block_out_channels:
-            return int(2 ** (len(block_out_channels) - 1))
-        downsample_ratio = getattr(self.vae.config, "downsample_ratio", None)
-        if downsample_ratio is not None:
-            return int(downsample_ratio)
-        return 16
+        self._labels_loaded_from_model_index = bool(self._id2label)
 
     @staticmethod
     def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
@@ -133,6 +91,19 @@ class LightningDiTPipeline(DiffusionPipeline):
             return {}
         return {int(key): value for key, value in id2label.items()}
 
+    @staticmethod
+    def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]:
+        if not variant_path:
+            return {}
+        model_index_path = Path(variant_path).resolve() / "model_index.json"
+        if not model_index_path.exists():
+            return {}
+        raw = json.loads(model_index_path.read_text(encoding="utf-8"))
+        id2label = raw.get("id2label")
+        if not isinstance(id2label, dict):
+            return {}
+        return {int(key): value for key, value in id2label.items()}
+
     @staticmethod
     def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]:
         label2id: Dict[str, int] = {}
@@ -143,81 +114,77 @@ class LightningDiTPipeline(DiffusionPipeline):
                     label2id[synonym] = int(class_id)
         return dict(sorted(label2id.items()))
 
+    def _ensure_labels_loaded(self) -> None:
+        if self._labels_loaded_from_model_index:
+            return
+        loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None))
+        if loaded:
+            self._id2label = loaded
+            self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = True
+
     @property
     def id2label(self) -> Dict[int, str]:
+        self._ensure_labels_loaded()
         return self._id2label
 
     def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
-        r"""Map English ImageNet labels to class ids."""
-        labels = [label] if isinstance(label, str) else label
+        self._ensure_labels_loaded()
         if not self.labels:
-            raise ValueError("No id2label mapping is available in this checkpoint.")
+            raise ValueError("No labels loaded. Ensure `id2label` exists in model_index.json.")
+        labels = [label] if isinstance(label, str) else label
         missing = [item for item in labels if item not in self.labels]
         if missing:
             preview = ", ".join(list(self.labels.keys())[:8])
-            raise ValueError(f"Unknown labels: {missing}. Example valid labels: {preview}, ...")
+            raise ValueError(f"Unknown label(s): {missing}. Example valid labels: {preview}, ...")
         return [self.labels[item] for item in labels]
 
     def _normalize_class_labels(
         self,
-        class_labels: Union[int, str, List[Union[int, str]], torch.Tensor],
+        class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor],
     ) -> torch.LongTensor:
         if isinstance(class_labels, torch.Tensor):
-            return class_labels.to(device=self._execution_device, dtype=torch.long).reshape(-1)
+            return class_labels.to(dtype=torch.long).reshape(-1)
         if isinstance(class_labels, int):
-            class_label_ids = [class_labels]
+            class_labels = [class_labels]
         elif isinstance(class_labels, str):
-            class_label_ids = self.get_label_ids(class_labels)
+            class_labels = self.get_label_ids(class_labels)
         elif class_labels and isinstance(class_labels[0], str):
-            class_label_ids = self.get_label_ids(class_labels)  # type: ignore[arg-type]
-        else:
-            class_label_ids = [int(class_id) for class_id in class_labels]  # type: ignore[union-attr]
-        return torch.tensor(class_label_ids, device=self._execution_device, dtype=torch.long).reshape(-1)
+            class_labels = self.get_label_ids(class_labels)  # type: ignore[arg-type]
+        return torch.tensor(class_labels, dtype=torch.long).reshape(-1)
 
-    def _default_image_size(self) -> int:
-        return int(self.transformer.config.input_size) * self.vae_scale_factor
-
-    def check_inputs(
+    def _prepare_latents(
         self,
+        batch_size: int,
         height: int,
         width: int,
-        num_inference_steps: int,
-        output_type: str,
-    ) -> None:
-        if num_inference_steps < 1:
-            raise ValueError("num_inference_steps must be >= 1.")
-        if output_type not in {"pil", "np", "pt", "latent"}:
-            raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.")
-        if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
-            raise ValueError(
-                f"height and width must be divisible by the VAE downsample factor {self.vae_scale_factor}."
-            )
-        latent_height = height // self.vae_scale_factor
-        latent_width = width // self.vae_scale_factor
-        expected_size = int(self.transformer.config.input_size)
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
+    ) -> torch.Tensor:
+        downsample = 16
+        if self.vae is not None:
+            block_out = getattr(self.vae.config, "block_out_channels", None)
+            if block_out is not None:
+                downsample = 2 ** (len(block_out) - 1)
+            elif hasattr(self.vae.config, "downsample_ratio"):
+                downsample = int(self.vae.config.downsample_ratio)
+
+        if height % downsample != 0 or width % downsample != 0:
+            raise ValueError(f"height and width must be divisible by the VAE downsample factor {downsample}.")
+
+        latent_height = height // downsample
+        latent_width = width // downsample
         patch_size = int(self.transformer.config.patch_size)
-        if latent_height != expected_size or latent_width != expected_size:
-            raise ValueError(
-                f"Requested latent size {(latent_height, latent_width)} does not match transformer "
-                f"input_size={expected_size}. Use height=width={self._default_image_size()}."
-            )
         if latent_height % patch_size != 0 or latent_width % patch_size != 0:
-            raise ValueError("Latent height and width must be divisible by transformer patch_size.")
-
-    @staticmethod
-    def prepare_extra_step_kwargs(
-        scheduler: KarrasDiffusionSchedulers,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
-    ) -> Dict[str, Any]:
-        extra_step_kwargs: Dict[str, Any] = {}
-        if "generator" in inspect.signature(scheduler.step).parameters:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
+            raise ValueError("Latent height and width must be divisible by the transformer patch_size.")
 
-    @staticmethod
-    def _flow_time_from_sigma_timestep(timestep: torch.Tensor, num_train_timesteps: int) -> torch.Tensor:
-        """Map FlowMatch scheduler timestep (sigma * N) to LightningDiT flow time in [0, 1]."""
-        return 1.0 - timestep.to(dtype=torch.float32) / float(num_train_timesteps)
+        return randn_tensor(
+            (batch_size, self.transformer.config.in_channels, latent_height, latent_width),
+            generator=generator,
+            device=device,
+            dtype=dtype,
+        )
 
     @staticmethod
     def _apply_cfg(
@@ -230,10 +197,11 @@ class LightningDiTPipeline(DiffusionPipeline):
             return model_output
         eps, rest = model_output[:, :cfg_channels], model_output[:, cfg_channels:]
         cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0)
-        guided_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
+        half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
         if rest.numel() == 0:
-            return guided_eps
-        return torch.cat([guided_eps, rest[: cond_eps.shape[0]]], dim=1)
+            return half_eps
+        cond_rest, _ = torch.chunk(rest, 2, dim=0)
+        return torch.cat([half_eps, cond_rest], dim=1)
 
     def _resolve_latent_stats(
         self,
@@ -263,58 +231,30 @@ class LightningDiTPipeline(DiffusionPipeline):
     ) -> torch.Tensor:
         return (latents * latent_std) / latent_multiplier + latent_mean
 
-    def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"):
-        if output_type == "latent":
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        if self.vae is None:
             return latents
+
         vae_dtype = next(self.vae.parameters()).dtype
         latents = latents.to(dtype=vae_dtype)
         scaling_factor = getattr(self.vae.config, "scaling_factor", None)
         if scaling_factor not in (None, 0):
             latents = latents / scaling_factor
-        image = self.vae.decode(latents).sample
-        if output_type == "pt":
-            return image
-        return self.image_processor.postprocess(image, output_type=output_type)
-
-    def _configure_scheduler(self, num_inference_steps: int, device: torch.device, timestep_shift: float):
-        if hasattr(self.scheduler, "set_shift"):
-            self.scheduler.set_shift(float(timestep_shift))
-        if _uses_explicit_next_timestep_scheduler(self.scheduler):
-            return self.scheduler.set_timesteps(
-                num_inference_steps,
-                device=device,
-                timestep_shift=float(timestep_shift),
-            )
-        if getattr(self.scheduler.config, "stochastic_sampling", False):
-            raise ValueError(
-                "LightningDiT expects deterministic FlowMatch scheduler stepping "
-                "(scheduler.config.stochastic_sampling=False)."
-            )
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-        return self.scheduler.timesteps
-
-    def _guidance_active(
-        self,
-        flow_time: float,
-        guidance_interval: Tuple[float, float],
-        cfg_interval_start: float,
-    ) -> bool:
-        if flow_time < float(cfg_interval_start):
-            return False
-        return guidance_interval[0] <= flow_time <= guidance_interval[1]
+        image = self.vae.decode(latents)
+        return image.sample if hasattr(image, "sample") else image
 
     @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
     def __call__(
         self,
-        class_labels: Union[int, str, List[Union[int, str]], torch.Tensor],
-        height: Optional[int] = None,
-        width: Optional[int] = None,
+        class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor],
+        height: int = 256,
+        width: int = 256,
         num_inference_steps: int = 250,
-        guidance_scale: float = 6.7,
+        guidance_scale: float = 1.0,
         guidance_interval: Tuple[float, float] = (0.0, 1.0),
         cfg_interval_start: float = 0.125,
-        timestep_shift: Optional[float] = None,
+        timestep_shift: float = 0.3,
+        heun: bool = False,
         cfg_channels: int = 3,
         latent_mean: Optional[torch.Tensor] = None,
         latent_std: Optional[torch.Tensor] = None,
@@ -322,138 +262,69 @@ class LightningDiTPipeline(DiffusionPipeline):
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         output_type: str = "pil",
         return_dict: bool = True,
-    ) -> Union[ImagePipelineOutput, Tuple]:
-        r"""
-        Generate class-conditional images at the transformer's native latent resolution.
-
-        Args:
-            class_labels (`int`, `str`, `list[int]`, `list[str]`, or `torch.Tensor`):
-                ImageNet class indices or human-readable English label strings (comma-separated synonyms
-                in `id2label` are supported).
-            height (`int`, *optional*):
-                Output image height in pixels. Defaults to `input_size * vae_scale_factor` (256 for XL/1-256).
-            width (`int`, *optional*):
-                Output image width in pixels. Defaults to the same value as `height`.
-            num_inference_steps (`int`, defaults to `250`):
-                Number of flow-matching steps. With [`FlowMatchHeunDiscreteScheduler`], each step may use
-                two model evaluations (2nd-order Heun).
-            guidance_scale (`float`, defaults to `6.7`):
-                Classifier-free guidance scale on the first `cfg_channels` latent channels. CFG is active when
-                `guidance_scale > 1.0` and flow time is at least `cfg_interval_start`.
-            guidance_interval (`tuple[float, float]`, defaults to `(0.0, 1.0)`):
-                Flow-time interval `[low, high]` where CFG is allowed (in addition to `cfg_interval_start`).
-            cfg_interval_start (`float`, defaults to `0.125`):
-                Minimum flow time before CFG is applied (official LightningDiT XL/1 setting).
-            timestep_shift (`float`, *optional*):
-                Timestep schedule shift. Defaults to `scheduler.config.shift`. Only applied at runtime if the
-                scheduler implements `set_shift` (e.g. [`FlowMatchEulerDiscreteScheduler`]); for
-                [`FlowMatchHeunDiscreteScheduler`], set `shift` in `scheduler_config.json` when loading.
-            cfg_channels (`int`, defaults to `3`):
-                Number of latent channels to apply CFG on.
-            latent_mean (`torch.Tensor`, *optional*):
-                Per-channel latent mean for denormalization before VAE decode. Read from the VAE config when omitted.
-            latent_std (`torch.Tensor`, *optional*):
-                Per-channel latent std for denormalization before VAE decode. Read from the VAE config when omitted.
-            latent_multiplier (`float`, defaults to `1.0`):
-                Divisor applied with `latent_std` during denormalization (`latents * std / multiplier + mean`).
-            generator (`torch.Generator`, *optional*):
-                RNG for reproducible noise initialization (and scheduler stochastic paths if enabled).
-            output_type (`str`, defaults to `"pil"`):
-                `"pil"`, `"np"`, `"pt"`, or `"latent"`.
-            return_dict (`bool`, defaults to `True`):
-                Return [`~pipelines.pipeline_utils.ImagePipelineOutput`] if `True`, else a `(images,)` tuple.
-
-        Examples:
-            <!-- this section is replaced by replace_example_docstring -->
-
-        Returns:
-            [`~pipelines.pipeline_utils.ImagePipelineOutput`] or `tuple`:
-                Generated images.
-        """
-        default_size = self._default_image_size()
-        height = int(height or default_size)
-        width = int(width or default_size)
-        self.check_inputs(height, width, num_inference_steps, output_type)
-
+    ) -> Union[LightningDiTPipelineOutput, Tuple]:
         device = self._execution_device
         model_dtype = next(self.transformer.parameters()).dtype
-        class_labels_tensor = self._normalize_class_labels(class_labels)
-        batch_size = class_labels_tensor.numel()
-        null_labels = torch.full_like(class_labels_tensor, int(self.config.null_class_id))
-
-        if timestep_shift is None:
-            timestep_shift = float(getattr(self.scheduler.config, "shift", 0.3))
-
-        schedule = self._configure_scheduler(num_inference_steps, device, timestep_shift)
-        num_train_timesteps = int(self.scheduler.config.num_train_timesteps)
-        use_builtin_flow_match = not _uses_explicit_next_timestep_scheduler(self.scheduler)
-        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator) if use_builtin_flow_match else {}
-
-        latents = randn_tensor(
-            (
-                batch_size,
-                int(self.transformer.config.in_channels),
-                height // self.vae_scale_factor,
-                width // self.vae_scale_factor,
-            ),
-            generator=generator,
-            device=device,
-            dtype=model_dtype,
-        )
-
-        if use_builtin_flow_match:
-            for timestep in self.progress_bar(schedule):
-                flow_time = float(self._flow_time_from_sigma_timestep(timestep, num_train_timesteps))
-                guidance_active = self._guidance_active(flow_time, guidance_interval, cfg_interval_start)
-                do_cfg = guidance_scale > 1.0 and guidance_active
-
-                if do_cfg:
-                    model_input = torch.cat([latents, latents], dim=0)
-                    labels = torch.cat([class_labels_tensor, null_labels], dim=0)
-                else:
-                    model_input = latents
-                    labels = class_labels_tensor
-
-                flow_time_batch = torch.full((labels.shape[0],), flow_time, device=device, dtype=model_dtype)
-                velocity = self.transformer(
-                    hidden_states=model_input,
-                    timestep=flow_time_batch,
-                    class_labels=labels,
-                    return_dict=True,
-                ).sample
-                velocity = self._apply_cfg(velocity, guidance_scale, guidance_active, cfg_channels)
 
-                # FlowMatchEuler/Heun: integrate in sigma space; model expects -velocity
-                latents = self.scheduler.step(
-                    -velocity,
-                    timestep,
-                    latents,
-                    **extra_step_kwargs,
+        class_labels = self._normalize_class_labels(class_labels).to(device=device)
+        batch_size = class_labels.numel()
+
+        latents = self._prepare_latents(batch_size, height, width, model_dtype, device, generator)
+        timesteps = self.scheduler.set_timesteps(num_inference_steps, device=device, timestep_shift=timestep_shift)
+
+        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator)
+
+        null_labels = torch.full_like(class_labels, self.transformer.config.num_classes)
+        for index, timestep in enumerate(timesteps[:-1]):
+            next_timestep = timesteps[index + 1]
+            guidance_active = guidance_interval[0] <= float(timestep) <= guidance_interval[1]
+            if cfg_interval_start is not None and float(timestep) < cfg_interval_start:
+                guidance_active = False
+
+            if guidance_scale > 1.0 and guidance_active:
+                model_input = torch.cat([latents, latents], dim=0)
+                labels = torch.cat([class_labels, null_labels], dim=0)
+            else:
+                model_input = latents
+                labels = class_labels
+
+            timestep_batch = torch.full((labels.shape[0],), float(timestep), device=device, dtype=model_dtype)
+            model_output = self.transformer(
+                model_input.to(dtype=model_dtype),
+                timestep_batch,
+                labels,
+                return_dict=True,
+            ).sample
+            model_output = self._apply_cfg(model_output, guidance_scale, guidance_active, cfg_channels)
+
+            if heun and index < len(timesteps) - 2:
+                provisional = self.scheduler.step(
+                    model_output, timestep[None], latents, next_timestep[None], **extra_step_kwargs
                 ).prev_sample
-        else:
-            for index, timestep in enumerate(self.progress_bar(schedule[:-1])):
-                next_timestep = schedule[index + 1]
-                flow_time = float(timestep)
-                guidance_active = self._guidance_active(flow_time, guidance_interval, cfg_interval_start)
-
                 if guidance_scale > 1.0 and guidance_active:
-                    model_input = torch.cat([latents, latents], dim=0)
-                    labels = torch.cat([class_labels_tensor, null_labels], dim=0)
+                    prime_input = torch.cat([provisional, provisional], dim=0)
+                    prime_labels = torch.cat([class_labels, null_labels], dim=0)
                 else:
-                    model_input = latents
-                    labels = class_labels_tensor
-
-                flow_time_batch = torch.full((labels.shape[0],), flow_time, device=device, dtype=model_dtype)
-                velocity = self.transformer(
-                    hidden_states=model_input,
-                    timestep=flow_time_batch,
-                    class_labels=labels,
+                    prime_input = provisional
+                    prime_labels = class_labels
+                next_timestep_batch = torch.full(
+                    (prime_labels.shape[0],), float(next_timestep), device=device, dtype=model_dtype
+                )
+                next_model_output = self.transformer(
+                    prime_input.to(dtype=model_dtype),
+                    next_timestep_batch,
+                    prime_labels,
                     return_dict=True,
                 ).sample
-                velocity = self._apply_cfg(velocity, guidance_scale, guidance_active, cfg_channels)
-
+                next_model_output = self._apply_cfg(
+                    next_model_output, guidance_scale, guidance_active, cfg_channels
+                )
+                latents = self.scheduler.step_heun(
+                    model_output, next_model_output, timestep[None], latents, next_timestep[None]
+                ).prev_sample
+            else:
                 latents = self.scheduler.step(
-                    velocity, timestep[None], latents, next_timestep[None]
+                    model_output, timestep[None], latents, next_timestep[None], **extra_step_kwargs
                 ).prev_sample
 
         latent_mean, latent_std = self._resolve_latent_stats(
@@ -466,12 +337,12 @@ class LightningDiTPipeline(DiffusionPipeline):
             )
         latents = self._denormalize_latents(latents, latent_mean, latent_std, latent_multiplier)
 
-        image = self.decode_latents(latents, output_type=output_type)
-        self.maybe_free_model_hooks()
+        image = self._decode_latents(latents)
+        if self.vae is not None:
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image = self.image_processor.postprocess(image, output_type=output_type)
 
+        self.maybe_free_model_hooks()
         if not return_dict:
             return (image,)
-        return ImagePipelineOutput(images=image)
-
-
-__all__ = ["LightningDiTPipeline"]
+        return LightningDiTPipelineOutput(images=image)

LightningDit-XL-1-256/scheduler/scheduler_config.json

@@ -1,5 +1,5 @@
 {
-  "_class_name": "FlowMatchHeunDiscreteScheduler",
+  "_class_name": "LightningDiTFlowMatchScheduler",
   "_diffusers_version": "0.36.0",
   "num_train_timesteps": 1000,
   "shift": 0.3

LightningDit-XL-1-256/scheduler/scheduling_flow_match_lightningdit.py

@@ -0,0 +1,108 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+
+try:
+    from diffusers.configuration_utils import ConfigMixin, register_to_config
+    from diffusers.schedulers.scheduling_utils import SchedulerMixin
+    from diffusers.utils import BaseOutput
+except Exception:  # pragma: no cover
+    class BaseOutput(dict):
+        def __post_init__(self):
+            self.update(self.__dict__)
+
+    class ConfigMixin:
+        config_name = "scheduler_config.json"
+
+    class SchedulerMixin:
+        pass
+
+    def register_to_config(init):
+        return init
+
+
+@dataclass
+class LightningDiTFlowMatchSchedulerOutput(BaseOutput):
+    prev_sample: torch.FloatTensor
+
+
+class LightningDiTFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Flow-matching ODE scheduler for LightningDiT (linear path, velocity prediction).
+
+    Integrates from t=0 (noise) to t=1 (data) with optional timestep shifting used in LightningDiT sampling.
+    """
+
+    config_name = "scheduler_config.json"
+    order = 1
+
+    @register_to_config
+    def __init__(self, path_type: str = "linear", num_train_timesteps: int = 1000):
+        if path_type not in {"linear", "cosine"}:
+            raise ValueError("path_type must be either 'linear' or 'cosine'.")
+        self.path_type = path_type
+        self.num_train_timesteps = num_train_timesteps
+        self.timesteps = torch.linspace(0.0, 1.0, num_train_timesteps + 1, dtype=torch.float64)
+
+    @staticmethod
+    def _apply_timestep_shift(timesteps: torch.Tensor, timestep_shift: float) -> torch.Tensor:
+        if timestep_shift <= 0:
+            return timesteps
+        return timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps)
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Optional[torch.device] = None,
+        timestep_shift: float = 0.0,
+    ):
+        timesteps = torch.linspace(0.0, 1.0, num_inference_steps + 1, dtype=torch.float64)
+        timesteps = self._apply_timestep_shift(timesteps, timestep_shift)
+        self.timesteps = timesteps.to(device=device)
+        return self.timesteps
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: torch.Tensor,
+        sample: torch.Tensor,
+        next_timestep: torch.Tensor,
+        return_dict: bool = True,
+    ) -> LightningDiTFlowMatchSchedulerOutput:
+        sample_dtype = sample.dtype
+        sample = sample.to(dtype=torch.float64)
+        model_output = model_output.to(dtype=torch.float64)
+        timestep = timestep.to(device=sample.device, dtype=torch.float64).flatten()
+        next_timestep = next_timestep.to(device=sample.device, dtype=torch.float64).flatten()
+        prev_sample = sample + (next_timestep[0] - timestep[0]) * model_output
+        prev_sample = prev_sample.to(sample_dtype)
+        if not return_dict:
+            return (prev_sample,)
+        return LightningDiTFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def step_heun(
+        self,
+        model_output: torch.Tensor,
+        next_model_output: torch.Tensor,
+        timestep: torch.Tensor,
+        sample: torch.Tensor,
+        next_timestep: torch.Tensor,
+        return_dict: bool = True,
+    ) -> LightningDiTFlowMatchSchedulerOutput:
+        sample_dtype = sample.dtype
+        sample = sample.to(dtype=torch.float64)
+        model_output = model_output.to(dtype=torch.float64)
+        next_model_output = next_model_output.to(dtype=torch.float64)
+        timestep = timestep.to(device=sample.device, dtype=torch.float64).flatten()
+        next_timestep = next_timestep.to(device=sample.device, dtype=torch.float64).flatten()
+        prev_sample = sample + (next_timestep[0] - timestep[0]) * (0.5 * model_output + 0.5 * next_model_output)
+        prev_sample = prev_sample.to(sample_dtype)
+        if not return_dict:
+            return (prev_sample,)
+        return LightningDiTFlowMatchSchedulerOutput(prev_sample=prev_sample)