Update pMF-B-32/pipeline.py

pMF-B-32/pipeline.py

@@ -1,34 +1,16 @@
-# 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: PMFPipeline.
-
-Load with native Hugging Face diffusers and trust_remote_code=True.
-"""
+"""Hub custom pipeline: PMFPipeline."""
 
 from __future__ import annotations
 
+import inspect
 import json
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils.torch_utils import randn_tensor
 
-
 DEFAULT_CFG_BY_MODEL: Dict[str, Dict[str, float]] = {
     "pMF-B/16": {"guidance_scale": 7.5, "guidance_interval_min": 0.1, "guidance_interval_max": 0.8},
     "pMF-B/32": {"guidance_scale": 6.5, "guidance_interval_min": 0.1, "guidance_interval_max": 0.7},
@@ -48,50 +30,37 @@ RECOMMENDED_NOISE_BY_MODEL: Dict[str, float] = {
 }
 
 
-def _set_pmf_timesteps(
-    scheduler: FlowMatchEulerDiscreteScheduler,
-    num_inference_steps: int,
-    device: torch.device,
-) -> torch.Tensor:
-    r"""Set linear flow sigmas from 1.0 to 0.0 for pMF sampling."""
-    flow_sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device, dtype=torch.float32)
-    scheduler.set_timesteps(sigmas=flow_sigmas.tolist(), device=device)
-    return flow_sigmas
-
-
 class PMFPipeline(DiffusionPipeline):
-    r"""
-    Pipeline for ImageNet class-conditional generation with Pixel Mean Flows (pMF).
-
-    Parameters:
-        transformer ([`PMFTransformer2DModel`]):
-            Class-conditioned pMF transformer that predicts mean-flow velocity.
-        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
-            Built-in flow-matching Euler scheduler.
-        id2label (`dict[int, str]`, *optional*):
-            ImageNet class id to English label mapping.
-    """
-
     model_cpu_offload_seq = "transformer"
 
     def __init__(
         self,
-        transformer,
-        scheduler,
+        transformer: Any,
+        scheduler: Any,
         id2label: Optional[Dict[Union[int, str], str]] = None,
-    ):
+    ) -> None:
         super().__init__()
         if scheduler is None:
-            scheduler = FlowMatchEulerDiscreteScheduler(
-                num_train_timesteps=1000,
-                shift=1.0,
-                stochastic_sampling=False,
-            )
+            raise ValueError("PMFPipeline requires a scheduler loaded from the checkpoint.")
         self.register_modules(transformer=transformer, scheduler=scheduler)
         self._id2label = self._normalize_id2label(id2label)
         self.labels = self._build_label2id(self._id2label)
         self._labels_loaded_from_model_index = bool(self._id2label)
 
+    @staticmethod
+    def prepare_extra_step_kwargs(
+        scheduler: Any,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        eta: float | None = None,
+    ) -> Dict[str, Any]:
+        kwargs: Dict[str, Any] = {}
+        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
+
     def _ensure_labels_loaded(self) -> None:
         if self._labels_loaded_from_model_index:
             return
@@ -152,8 +121,8 @@ class PMFPipeline(DiffusionPipeline):
         if isinstance(class_labels, str):
             return self.get_label_ids(class_labels)
         if class_labels and isinstance(class_labels[0], str):
-            return self.get_label_ids(class_labels)
-        return list(class_labels)
+            return self.get_label_ids(class_labels)  # type: ignore[arg-type]
+        return [int(class_id) for class_id in class_labels]
 
     def _recommended_noise_scale(self) -> float:
         model_type = getattr(self.transformer.config, "model_type", None)
@@ -168,6 +137,28 @@ class PMFPipeline(DiffusionPipeline):
             return dict(DEFAULT_CFG_BY_MODEL[model_type])
         return {"guidance_scale": 7.5, "guidance_interval_min": 0.1, "guidance_interval_max": 0.8}
 
+    def predict_u(
+        self,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        class_labels: torch.Tensor,
+        h: torch.Tensor,
+        omega: torch.Tensor,
+        guidance_interval_min: torch.Tensor,
+        guidance_interval_max: torch.Tensor,
+    ) -> torch.Tensor:
+        output = self.transformer(
+            sample=sample,
+            timestep=timestep,
+            class_labels=class_labels,
+            h=h,
+            omega=omega,
+            guidance_interval_min=guidance_interval_min,
+            guidance_interval_max=guidance_interval_max,
+            return_dict=True,
+        )
+        return output.u
+
     @torch.inference_mode()
     def __call__(
         self,
@@ -181,33 +172,6 @@ class PMFPipeline(DiffusionPipeline):
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
     ) -> Union[ImagePipelineOutput, Tuple]:
-        r"""
-        Generate class-conditional images with pMF.
-
-        Args:
-            class_labels (`int`, `str`, or `list`):
-                ImageNet class id(s) or label name(s).
-            num_inference_steps (`int`, *optional*, defaults to 1):
-                Number of flow steps. pMF is typically used with 1 step.
-            guidance_scale (`float`, *optional*):
-                Classifier-free guidance scale. Defaults to model-specific preset.
-            guidance_interval_min (`float`, *optional*):
-                Lower bound of the CFG interval in normalized time.
-            guidance_interval_max (`float`, *optional*):
-                Upper bound of the CFG interval in normalized time.
-            noise_scale (`float`, *optional*):
-                Initial Gaussian noise scale. Defaults to model-specific preset.
-            generator (`torch.Generator`, *optional*):
-                Random generator for reproducibility.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                Output format: `"pil"`, `"np"`, or `"pt"`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether to return an [`~pipelines.ImagePipelineOutput`].
-
-        Returns:
-            [`~pipelines.ImagePipelineOutput`] or `tuple`:
-                Generated images.
-        """
         if num_inference_steps < 1:
             raise ValueError("num_inference_steps must be >= 1.")
         if output_type not in {"pil", "np", "pt"}:
@@ -247,15 +211,17 @@ class PMFPipeline(DiffusionPipeline):
         t_min = torch.full((batch_size,), guidance_interval_min, device=device, dtype=dtype)
         t_max = torch.full((batch_size,), guidance_interval_max, device=device, dtype=dtype)
 
-        flow_sigmas = _set_pmf_timesteps(self.scheduler, num_inference_steps, device)
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator)
+        timesteps = self.scheduler.timesteps
 
         for step_index in self.progress_bar(range(num_inference_steps)):
-            t = flow_sigmas[step_index]
-            t_next = flow_sigmas[step_index + 1]
+            t = timesteps[step_index]
+            t_next = timesteps[step_index + 1]
             h = (t - t_next).expand(batch_size).to(device=device, dtype=dtype)
             t_batch = t.expand(batch_size).to(device=device, dtype=dtype)
 
-            output = self.transformer(
+            u = self.predict_u(
                 sample=latents,
                 timestep=t_batch,
                 class_labels=class_labels_t,
@@ -263,9 +229,8 @@ class PMFPipeline(DiffusionPipeline):
                 omega=omega,
                 guidance_interval_min=t_min,
                 guidance_interval_max=t_max,
-                return_dict=True,
             )
-            latents = self.scheduler.step(output.u, self.scheduler.timesteps[step_index], latents).prev_sample
+            latents = self.scheduler.step(u, t, latents, **extra_step_kwargs).prev_sample
 
         images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
         if output_type == "pt":
@@ -283,3 +248,5 @@ class PMFPipeline(DiffusionPipeline):
 
 
 PMFPipelineOutput = ImagePipelineOutput
+
+__all__ = ["PMFPipeline", "PMFPipelineOutput"]