Update JiT-B-16-SIM/pipeline.py

JiT-B-16-SIM/pipeline.py

@@ -0,0 +1,280 @@
+"""Hub custom pipeline: JiTPipeline for FD-Loss post-trained JiT checkpoints.
+
+Uses FD-Loss sampling (legacy time convention, velocity Euler/Heun, t: 1→0).
+See libs/FD-Loss-diffusers and scripts/evaluate_released_ckpt.sh (JiT_B preset).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+from scheduling_flow_match_fd import FDLossFlowMatchScheduler
+
+RECOMMENDED_NOISE_BY_SIZE = {
+    256: 1.0,
+    512: 2.0,
+}
+
+RECOMMENDED_CFG_BY_VARIANT = {
+    "JiT-B": 3.0,
+    "JiT-L": 2.4,
+    "JiT-H": 2.2,
+}
+
+
+class JiTPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for FD-Loss post-trained JiT (flow matching, legacy time convention).
+
+    Parameters:
+        transformer ([`JiTTransformer2DModel`]):
+            Class-conditioned JiT backbone.
+        scheduler ([`FDLossFlowMatchScheduler`]):
+            Flow timesteps from 1 (noise) to 0 (data).
+        legacy_time_convention (`bool`, *optional*, defaults to `True`):
+            Flip flow time when passing to the backbone (`t_bb = 1 - t`), as in FD-Loss training.
+        id2label (`dict[int, str]`, *optional*):
+            ImageNet class id to English label mapping.
+    """
+
+    model_cpu_offload_seq = "transformer"
+
+    def __init__(
+        self,
+        transformer,
+        scheduler=None,
+        id2label: Optional[Dict[Union[int, str], str]] = None,
+        legacy_time_convention: bool = True,
+    ):
+        super().__init__()
+        if scheduler is None:
+            scheduler = FDLossFlowMatchScheduler()
+        self.register_modules(transformer=transformer, scheduler=scheduler)
+        self.legacy_time_convention = legacy_time_convention
+        self._id2label = self._normalize_id2label(id2label)
+        self.labels = self._build_label2id(self._id2label)
+        self._labels_loaded_from_model_index = bool(self._id2label)
+
+    def _backbone_t(self, t: torch.Tensor) -> torch.Tensor:
+        if self.legacy_time_convention:
+            return 1.0 - t
+        return t
+
+    def _normalize_class_labels(
+        self,
+        class_labels: Union[int, str, List[Union[int, str]]],
+    ) -> List[int]:
+        if isinstance(class_labels, int):
+            return [class_labels]
+        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)
+
+    def _forward_velocity(
+        self,
+        z: torch.Tensor,
+        t: torch.Tensor,
+        labels: torch.Tensor,
+        guidance_scale: float,
+        cfg_interval: Optional[Tuple[float, float]],
+        t_eps: float,
+    ) -> torch.Tensor:
+        t_view = t.reshape(-1, *([1] * (z.ndim - 1)))
+        t_bb = self._backbone_t(t).flatten().expand(z.shape[0])
+
+        x_cond = self.transformer(
+            z,
+            timestep=t_bb,
+            class_labels=labels,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+        ).sample
+        v_cond = (z - x_cond) / t_view.clamp_min(t_eps)
+
+        if guidance_scale == 1.0:
+            return v_cond
+
+        null_class = int(
+            getattr(self.transformer.config, "num_classes", getattr(self.transformer.config, "num_class_embeds", 1000))
+        )
+        class_null = torch.full_like(labels, null_class)
+        x_uncond = self.transformer(
+            z,
+            timestep=t_bb,
+            class_labels=class_null,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+        ).sample
+        v_uncond = (z - x_uncond) / t_view.clamp_min(t_eps)
+
+        if cfg_interval is None:
+            return v_uncond + guidance_scale * (v_cond - v_uncond)
+
+        low, high = cfg_interval
+        mask = (t < high) & ((low == 0) | (t > low))
+        scale = torch.where(
+            mask,
+            torch.tensor(guidance_scale, device=z.device, dtype=z.dtype),
+            torch.tensor(1.0, device=z.device, dtype=z.dtype),
+        )
+        while scale.ndim < v_cond.ndim:
+            scale = scale.unsqueeze(-1)
+        return v_uncond + scale * (v_cond - v_uncond)
+
+    @staticmethod
+    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
+        if not id2label:
+            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 {}
+        variant_dir = Path(variant_path).resolve()
+        model_index_path = variant_dir / "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] = {}
+        for class_id, value in id2label.items():
+            for synonym in value.split(","):
+                synonym = synonym.strip()
+                if synonym:
+                    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]:
+        self._ensure_labels_loaded()
+        label2id = self.labels
+        if not label2id:
+            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")
+        if isinstance(label, str):
+            label = [label]
+        missing = [item for item in label if item not in label2id]
+        if missing:
+            preview = ", ".join(list(label2id.keys())[:8])
+            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
+        return [label2id[item] for item in label]
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        class_labels: Union[int, str, List[Union[int, str]]],
+        num_inference_steps: int = 1,
+        guidance_scale: float = 3.0,
+        guidance_interval_min: float = 0.1,
+        guidance_interval_max: float = 1.0,
+        sampling_method: str = "euler",
+        noise_scale: Optional[float] = None,
+        t_eps: float = 5e-2,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        interpolate_pos_encoding: bool = True,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        if num_inference_steps < 1:
+            raise ValueError("num_inference_steps must be >= 1.")
+        if output_type not in {"pil", "np", "pt"}:
+            raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")
+        if sampling_method not in {"euler", "heun"}:
+            raise ValueError("sampling_method must be 'euler' or 'heun'.")
+
+        class_label_ids = self._normalize_class_labels(class_labels)
+        batch_size = len(class_label_ids)
+        image_size = int(self.transformer.config.sample_size)
+        patch_size = int(self.transformer.config.patch_size)
+        height = int(height or image_size)
+        width = int(width or image_size)
+        if height % patch_size != 0 or width % patch_size != 0:
+            raise ValueError(
+                f"height and width must be divisible by patch_size={patch_size}. Got {(height, width)}."
+            )
+        channels = int(self.transformer.config.in_channels)
+
+        if noise_scale is None:
+            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)
+
+        z = randn_tensor(
+            shape=(batch_size, channels, height, width),
+            generator=generator,
+            device=self._execution_device,
+            dtype=self.transformer.dtype,
+        ) * noise_scale
+
+        labels = torch.tensor(class_label_ids, device=self._execution_device, dtype=torch.long).reshape(-1)
+        null_class_val = int(
+            getattr(self.transformer.config, "num_classes", getattr(self.transformer.config, "num_class_embeds", 1000))
+        )
+        labels = labels.clamp(0, null_class_val - 1)
+
+        cfg_interval = [
+            float(self._backbone_t(torch.tensor(guidance_interval_min, device=z.device))),
+            float(self._backbone_t(torch.tensor(guidance_interval_max, device=z.device))),
+        ]
+        cfg_interval = (min(cfg_interval), max(cfg_interval))
+
+        timesteps = self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
+        ts = timesteps.view(-1, *([1] * z.ndim)).expand(-1, batch_size, -1, -1, -1)
+
+        for i in self.progress_bar(range(num_inference_steps - 1)):
+            t_cur = ts[i]
+            t_next = ts[i + 1]
+            if sampling_method == "heun":
+                dt = t_next - t_cur
+                v1 = self._forward_velocity(z, t_cur, labels, guidance_scale, cfg_interval, t_eps)
+                z_mid = z + dt * v1
+                v2 = self._forward_velocity(z_mid, t_next, labels, guidance_scale, cfg_interval, t_eps)
+                z = z + dt * 0.5 * (v1 + v2)
+            else:
+                v = self._forward_velocity(z, t_cur, labels, guidance_scale, cfg_interval, t_eps)
+                z = z + (t_next - t_cur) * v
+
+        if num_inference_steps >= 1:
+            t_cur = ts[-2]
+            t_next = ts[-1]
+            v = self._forward_velocity(z, t_cur, labels, guidance_scale, cfg_interval, t_eps)
+            z = z + (t_next - t_cur) * v
+
+        images_pt = ((z.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
+        if output_type == "pt":
+            images = images_pt
+        elif output_type == "np":
+            images = images_pt.permute(0, 2, 3, 1).numpy()
+        else:
+            images = self.numpy_to_pil(images_pt.permute(0, 2, 3, 1).numpy())
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return ImagePipelineOutput(images=images)