sefi/runner.py

"""SEFI T2I inference runner with three-phase masked denoising."""

from __future__ import annotations

import json
import math
import os
from typing import Callable, Optional

import torch
from PIL import Image
from torch import Tensor

from .builder import (
    build_components,
    build_lightweight_transformer,
    _derive_semantic_channels,
    _derive_text_output_dim,
    _derive_texture_channels,
    text_encoder_signature,
)
from .config import load_config
from .modeling import Qwen3VLTextEncoder


def _resolve_weight_dtype(config, *, override: Optional[str] = None) -> torch.dtype:
    if override is not None:
        normalized = str(override).strip().lower()
        if normalized == "bf16":
            return torch.bfloat16
        if normalized in {"fp32", "float32"}:
            return torch.float32
        raise ValueError(
            f"Unsupported inference dtype: {override}. Expected one of ['bf16', 'fp32']."
        )

    precision = str(getattr(config.training, "mixed_precision", "bf16")).lower()
    if precision == "fp16":
        return torch.float16
    if precision in {"fp32", "float32", "no"}:
        return torch.float32
    return torch.bfloat16


def _training_sefi_cfg(config):
    cfg = config.training.get("sefi", None)
    if cfg is not None:
        return cfg
    raise ValueError("Config requires training.sefi section.")


def _apply_timestep_shift_unit_interval(u_unit: Tensor, alpha: float) -> Tensor:
    """Apply t' = alpha*t / (1 + (alpha-1)*t) on unit coordinate u in [0, 1]."""
    alpha = float(alpha)
    if alpha <= 0:
        raise ValueError(f"timestep_shift_alpha must be > 0, got {alpha}")
    if alpha == 1.0:
        return u_unit
    denominator = 1.0 + (alpha - 1.0) * u_unit
    return (alpha * u_unit) / denominator


def _combine_guided_velocity(base_pred: Tensor, cond_pred: Tensor, guidance_scale: float) -> Tensor:
    """Shared guidance formula: base + scale * (conditioned - base)."""
    return base_pred + float(guidance_scale) * (cond_pred - base_pred)


def _resolve_guidance_interval_sigma(
    sigma_lo: Optional[float],
    sigma_hi: Optional[float],
) -> tuple[Optional[float], Optional[float]]:
    if sigma_lo is None and sigma_hi is None:
        return None, None
    if sigma_lo is None or sigma_hi is None:
        raise ValueError(
            "Limited interval guidance requires both "
            "guidance_interval_sigma_lo and guidance_interval_sigma_hi, or neither."
        )

    sigma_lo = float(sigma_lo)
    sigma_hi = float(sigma_hi)
    if not math.isfinite(sigma_lo) or not math.isfinite(sigma_hi):
        raise ValueError("guidance interval sigma thresholds must be finite.")
    if sigma_lo < 0.0 or sigma_hi < 0.0:
        raise ValueError("guidance interval sigma thresholds must be >= 0.")
    if sigma_lo >= sigma_hi:
        raise ValueError("guidance_interval_sigma_lo must be < guidance_interval_sigma_hi.")
    return sigma_lo, sigma_hi


def _guidance_interval_is_active(
    sigma: Tensor | float,
    sigma_lo: Optional[float],
    sigma_hi: Optional[float],
) -> bool:
    if sigma_lo is None and sigma_hi is None:
        return True
    if sigma_lo is None or sigma_hi is None:
        raise ValueError("guidance interval sigma bounds must be paired.")

    sigma_value = float(sigma.item()) if isinstance(sigma, Tensor) else float(sigma)
    return float(sigma_lo) < sigma_value <= float(sigma_hi)


def _normalize_optional_path(path: Optional[str]) -> str:
    if path is None:
        return ""
    return str(path).strip()


def _resolve_autoguidance_paths(
    autoguidance_config_path: Optional[str],
    autoguidance_checkpoint_path: Optional[str],
) -> tuple[str, str]:
    config_path = _normalize_optional_path(autoguidance_config_path)
    checkpoint_path = _normalize_optional_path(autoguidance_checkpoint_path)
    if bool(config_path) != bool(checkpoint_path):
        raise ValueError(
            "AutoGuidance requires both --autoguidance_config and "
            "--autoguidance_checkpoint, or neither."
        )
    return config_path, checkpoint_path


def _validate_autoguidance_guidance_scale(enabled: bool, guidance_scale: float) -> None:
    if enabled and float(guidance_scale) <= 1.0:
        raise ValueError("AutoGuidance requires guidance_scale > 1.0.")


def _resolve_checkpoint_file(checkpoint_path: str) -> str:
    if os.path.isdir(checkpoint_path):
        transformer_dir = os.path.join(checkpoint_path, "transformer")
        sharded_safetensors = os.path.join(
            transformer_dir,
            "diffusion_pytorch_model.safetensors.index.json",
        )
        if os.path.isfile(sharded_safetensors):
            return sharded_safetensors

        safetensors_state = os.path.join(
            transformer_dir,
            "diffusion_pytorch_model.safetensors",
        )
        if os.path.isfile(safetensors_state):
            return safetensors_state

        torch_state = os.path.join(transformer_dir, "diffusion_pytorch_model.bin")
        if os.path.isfile(torch_state):
            return torch_state

        raise FileNotFoundError(
            f"Unsupported SEFI inference checkpoint directory: {checkpoint_path}. "
            "Expected transformer/diffusion_pytorch_model.safetensors or "
            "transformer/diffusion_pytorch_model.safetensors.index.json."
        )

    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_path}")

    return checkpoint_path


def _extract_state_dict(checkpoint: dict) -> dict:
    if not isinstance(checkpoint, dict):
        raise ValueError("Checkpoint must be a dict-like object.")

    if "model_state_dict" in checkpoint and isinstance(checkpoint["model_state_dict"], dict):
        return checkpoint["model_state_dict"]
    if "module" in checkpoint and isinstance(checkpoint["module"], dict):
        return checkpoint["module"]
    if "state_dict" in checkpoint and isinstance(checkpoint["state_dict"], dict):
        return checkpoint["state_dict"]
    if checkpoint and all(isinstance(v, torch.Tensor) for v in checkpoint.values()):
        return checkpoint

    raise ValueError(
        "Unsupported checkpoint format. Expected one of: "
        "model_state_dict / module / state_dict / plain state_dict."
    )


def _load_checkpoint_payload(checkpoint_file: str):
    if checkpoint_file.endswith(".safetensors.index.json"):
        from safetensors.torch import load_file

        with open(checkpoint_file, "r", encoding="utf-8") as handle:
            index = json.load(handle)
        weight_map = index.get("weight_map", None)
        if not isinstance(weight_map, dict) or not weight_map:
            raise ValueError(f"Invalid safetensors index file: {checkpoint_file}")

        base_dir = os.path.dirname(checkpoint_file)
        state_dict = {}
        for shard_name in sorted(set(weight_map.values())):
            shard_path = os.path.join(base_dir, shard_name)
            if not os.path.isfile(shard_path):
                raise FileNotFoundError(f"Missing safetensors shard: {shard_path}")
            state_dict.update(load_file(shard_path))
        return state_dict

    if checkpoint_file.endswith(".safetensors"):
        from safetensors.torch import load_file

        return load_file(checkpoint_file)

    return torch.load(checkpoint_file, map_location="cpu")


def _strip_prefix_if_needed(state_dict: dict, prefix: str) -> dict:
    if state_dict and all(k.startswith(prefix) for k in state_dict):
        return {k[len(prefix) :]: v for k, v in state_dict.items()}
    return state_dict


def _load_transformer_state_dict_strict_shapes(
    transformer,
    checkpoint_path: str,
    *,
    label: str,
) -> None:
    checkpoint_file = _resolve_checkpoint_file(checkpoint_path)
    print(f"Loading {label} checkpoint from {checkpoint_file}")
    payload = _load_checkpoint_payload(checkpoint_file)
    state_dict = _extract_state_dict(payload)
    state_dict = _strip_prefix_if_needed(state_dict, "module.")

    target_state = transformer.state_dict()
    compatible_state = {}
    mismatched = []
    for key, value in state_dict.items():
        if key not in target_state:
            continue
        if tuple(value.shape) != tuple(target_state[key].shape):
            mismatched.append(
                f"{key}: checkpoint={tuple(value.shape)} vs model={tuple(target_state[key].shape)}"
            )
            continue
        compatible_state[key] = value

    if mismatched:
        raise ValueError(f"{label} checkpoint has shape-mismatched keys: {mismatched[:10]}")
    if not compatible_state:
        raise ValueError(
            f"{label} checkpoint has zero loadable parameters for the constructed model: "
            f"{checkpoint_path}"
        )

    missing, unexpected = transformer.load_state_dict(compatible_state, strict=False)
    if missing:
        print(f"  Warning - {label} missing keys: {missing[:10]}")
    if unexpected:
        print(f"  Warning - {label} unexpected keys: {unexpected[:10]}")


class SEFIInferenceRunner:
    """Inference runner for SEFI-T2I with three-phase masked denoising."""

    def __init__(
        self,
        config,
        *,
        checkpoint_path: str = "",
        device: str = "cuda",
        debug_assert_schedule: bool = False,
        delta_t_override: Optional[float] = None,
        inference_dtype: Optional[str] = None,
        timestep_shift_alpha: float = 1.0,
        autoguidance_config_path: Optional[str] = None,
        autoguidance_checkpoint_path: Optional[str] = None,
        guidance_interval_sigma_lo: Optional[float] = None,
        guidance_interval_sigma_hi: Optional[float] = None,
    ):
        from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor

        self.config = config
        self.device = torch.device(device)
        self.component_dtype = _resolve_weight_dtype(config)
        self.weight_dtype = _resolve_weight_dtype(config, override=inference_dtype)
        (
            self.autoguidance_config_path,
            self.autoguidance_checkpoint_path,
        ) = _resolve_autoguidance_paths(
            autoguidance_config_path,
            autoguidance_checkpoint_path,
        )
        self.autoguidance_enabled = bool(self.autoguidance_config_path)
        self.autoguidance_transformer = None
        self.autoguidance_text_encoder = None
        self.autoguidance_reuse_main_text_encoder = True
        (
            self.guidance_interval_sigma_lo,
            self.guidance_interval_sigma_hi,
        ) = _resolve_guidance_interval_sigma(
            guidance_interval_sigma_lo,
            guidance_interval_sigma_hi,
        )
        self.guidance_interval_enabled = self.guidance_interval_sigma_lo is not None

        components = build_components(config, component_dtype=self.component_dtype)
        self.transformer = components.transformer.to(
            device=self.device,
            dtype=self.weight_dtype,
        ).eval()
        self.text_encoder = components.text_encoder.to(
            device=self.device,
            dtype=self.component_dtype,
        ).eval()
        self.texture_codec = components.texture_codec.to(
            device=self.device,
            dtype=self.component_dtype,
        ).eval()
        self.noise_scheduler = components.noise_scheduler
        self.pipeline_cls = components.pipeline_cls
        self.semantic_channels = int(components.semantic_channels)
        self.texture_channels = int(components.texture_channels)
        self.total_channels = int(components.total_channels)

        self.debug_assert_schedule = bool(debug_assert_schedule)
        self.timestep_shift_alpha = float(timestep_shift_alpha)
        if self.timestep_shift_alpha <= 0:
            raise ValueError(
                "timestep_shift_alpha must be > 0. "
                f"Got {self.timestep_shift_alpha}."
            )
        self._configure_delta_t(delta_t_override)
        shift_enabled = self.timestep_shift_alpha != 1.0
        print(
            "Inference timestep schedule: "
            f"timestep_shift_alpha={self.timestep_shift_alpha:.6f}, "
            f"delta_t={self.delta_t:.6f}, shift_enabled={shift_enabled}"
        )
        if self.guidance_interval_enabled:
            print(
                "Limited interval guidance enabled on base sigma: "
                f"({self.guidance_interval_sigma_lo:.6f}, "
                f"{self.guidance_interval_sigma_hi:.6f}]"
            )

        texture_vae_cfg = self.texture_codec.texture_vae.config
        self.vae_scale_factor = 2 ** (len(texture_vae_cfg.block_out_channels) - 1)
        self.image_processor = Flux2ImageProcessor(
            vae_scale_factor=self.vae_scale_factor * 2
        )

        for module in (self.transformer, self.text_encoder, self.texture_codec):
            for param in module.parameters():
                param.requires_grad = False

        if checkpoint_path:
            self.load_checkpoint(checkpoint_path)

        if self.autoguidance_enabled:
            self._load_autoguidance_model()

    def _configure_delta_t(self, delta_t_override: Optional[float]) -> None:
        sefi_cfg = _training_sefi_cfg(self.config)

        delta_t_min_raw = sefi_cfg.get("delta_t_min", None)
        delta_t_max_raw = sefi_cfg.get("delta_t_max", None)
        if delta_t_min_raw is None or delta_t_max_raw is None:
            raise ValueError("training.sefi.delta_t_min and delta_t_max are required.")

        self.delta_t_min = float(delta_t_min_raw)
        self.delta_t_max = float(delta_t_max_raw)
        if self.delta_t_min < 0 or self.delta_t_min > 1:
            raise ValueError("training.sefi.delta_t_min must be in [0, 1].")
        if self.delta_t_max < 0 or self.delta_t_max > 1:
            raise ValueError("training.sefi.delta_t_max must be in [0, 1].")
        if self.delta_t_min > self.delta_t_max:
            raise ValueError("training.sefi.delta_t_min must be <= delta_t_max.")

        if delta_t_override is None:
            self.delta_t = self.delta_t_max
            print(
                "Warning: --delta-t not provided. "
                f"Using training.sefi.delta_t_max={self.delta_t_max:.6f} for inference."
            )
            return

        self.delta_t = float(delta_t_override)
        if self.delta_t < 0 or self.delta_t > 1:
            raise ValueError("inference delta_t must be in [0, 1].")
        if self.delta_t < self.delta_t_min or self.delta_t > self.delta_t_max:
            print(
                "Warning: inference delta_t is outside training range "
                f"[{self.delta_t_min:.6f}, {self.delta_t_max:.6f}]. "
                f"Got delta_t={self.delta_t:.6f}."
            )

    def load_checkpoint(self, checkpoint_path: str):
        ckpt_file = _resolve_checkpoint_file(checkpoint_path)
        print(f"Loading checkpoint from {ckpt_file}")
        ckpt = _load_checkpoint_payload(ckpt_file)
        state_dict = _extract_state_dict(ckpt)
        state_dict = _strip_prefix_if_needed(state_dict, "module.")

        missing, unexpected = self.transformer.load_state_dict(state_dict, strict=False)
        if missing:
            raise ValueError(f"Checkpoint is missing transformer keys: {missing[:10]}")
        if unexpected:
            raise ValueError(f"Checkpoint has unexpected transformer keys: {unexpected[:10]}")

    def _load_autoguidance_model(self) -> None:
        autoguidance_config = load_config(self.autoguidance_config_path)
        self.autoguidance_config = autoguidance_config

        ag_semantic_channels = _derive_semantic_channels(autoguidance_config)
        ag_texture_channels = _derive_texture_channels(autoguidance_config)
        if ag_semantic_channels != self.semantic_channels:
            raise ValueError(
                "AutoGuidance semantic channel mismatch: "
                f"main={self.semantic_channels}, small={ag_semantic_channels}."
            )
        if ag_texture_channels != self.texture_channels:
            raise ValueError(
                "AutoGuidance texture channel mismatch: "
                f"main={self.texture_channels}, small={ag_texture_channels}."
            )

        ag_text_output_dim = _derive_text_output_dim(autoguidance_config)
        autoguidance_transformer = build_lightweight_transformer(
            autoguidance_config,
            total_channels=self.total_channels,
            text_output_dim=ag_text_output_dim,
        )
        _load_transformer_state_dict_strict_shapes(
            autoguidance_transformer,
            self.autoguidance_checkpoint_path,
            label="AutoGuidance",
        )
        self.autoguidance_transformer = autoguidance_transformer.to(
            device=self.device,
            dtype=self.weight_dtype,
        ).eval()
        for param in self.autoguidance_transformer.parameters():
            param.requires_grad = False

        self.autoguidance_reuse_main_text_encoder = (
            text_encoder_signature(self.config) == text_encoder_signature(autoguidance_config)
        )
        if self.autoguidance_reuse_main_text_encoder:
            print("AutoGuidance reuses main prompt embeddings.")
        else:
            text_cfg = autoguidance_config.model.text_encoder
            self.autoguidance_text_encoder = Qwen3VLTextEncoder(
                model_name=str(text_cfg.model_name),
                weights_root=str(text_cfg.get("weights_root", "outputs/model_weights")),
                max_length=int(text_cfg.max_length),
                hidden_layers=[int(x) for x in text_cfg.hidden_layers],
                torch_dtype=self.component_dtype,
            ).to(device=self.device, dtype=self.component_dtype).eval()
            if int(self.autoguidance_text_encoder.output_dim) != int(ag_text_output_dim):
                raise ValueError(
                    "AutoGuidance text encoder output dim mismatch: "
                    f"loaded={self.autoguidance_text_encoder.output_dim}, "
                    f"expected={ag_text_output_dim}."
                )
            for param in self.autoguidance_text_encoder.parameters():
                param.requires_grad = False
            print("AutoGuidance uses a separate small-model text encoder.")

        print(
            "Loaded AutoGuidance model: "
            f"config={self.autoguidance_config_path}, "
            f"checkpoint={self.autoguidance_checkpoint_path}"
        )

    def _timesteps_and_sigmas(
        self,
        u_continuous: Tensor,
        *,
        n_dim: int,
        dtype: torch.dtype,
    ) -> tuple[Tensor, Tensor]:
        num_steps = int(self.noise_scheduler.config.num_train_timesteps)
        indices = (u_continuous * (num_steps - 1)).long().clamp(0, num_steps - 1)

        timesteps = self.noise_scheduler.timesteps[indices.cpu()].to(self.device)
        sigmas = self.noise_scheduler.sigmas[indices.cpu()].to(
            device=self.device,
            dtype=dtype,
        )
        while sigmas.ndim < n_dim:
            sigmas = sigmas.unsqueeze(-1)
        return timesteps, sigmas

    def _assert_shifted_schedule(
        self,
        u_base_unit: Tensor,
        u_sem_raw_schedule: Tensor,
        eps: float = 1e-6,
    ) -> None:
        if u_base_unit.ndim != 1 or u_sem_raw_schedule.ndim != 1:
            raise ValueError("u_base_unit and u_sem_raw_schedule must be 1D tensors.")
        if u_base_unit.shape != u_sem_raw_schedule.shape:
            raise ValueError("u_base_unit and u_sem_raw_schedule must have the same shape.")

        expected_u_max = 1.0 + self.delta_t
        if abs(float(u_base_unit[0].item()) - 0.0) > eps:
            raise ValueError(
                f"Invalid u_base_unit[0], expected 0, got {float(u_base_unit[0].item()):.6f}"
            )
        if abs(float(u_base_unit[-1].item()) - 1.0) > eps:
            raise ValueError(
                f"Invalid u_base_unit[-1], expected 1, got {float(u_base_unit[-1].item()):.6f}"
            )
        if abs(float(u_sem_raw_schedule[0].item()) - 0.0) > eps:
            raise ValueError(
                "Invalid shifted schedule start, expected 0, "
                f"got {float(u_sem_raw_schedule[0].item()):.6f}"
            )
        if abs(float(u_sem_raw_schedule[-1].item()) - expected_u_max) > eps:
            raise ValueError(
                "Invalid shifted schedule end, expected 1+delta_t, "
                f"got {float(u_sem_raw_schedule[-1].item()):.6f}, "
                f"expected={expected_u_max:.6f}"
            )

        diffs = u_sem_raw_schedule[1:] - u_sem_raw_schedule[:-1]
        if torch.any(diffs < -eps):
            index = int(torch.nonzero(diffs < -eps, as_tuple=False)[0, 0].item())
            raise ValueError(
                "Shifted u_sem_raw schedule must be monotonic non-decreasing, "
                f"but got decrease at step {index}: "
                f"{float(u_sem_raw_schedule[index].item()):.6f} -> "
                f"{float(u_sem_raw_schedule[index + 1].item()):.6f}"
            )

    def _assert_dual_time_invariants(
        self,
        u_sem: Tensor,
        u_tex: Tensor,
        sigmas_sem: Tensor,
        sigmas_tex: Tensor,
        eps: float = 1e-6,
    ) -> None:
        u_violation = u_sem < u_tex
        if torch.any(u_violation):
            index = int(torch.nonzero(u_violation, as_tuple=False)[0, 0].item())
            raise ValueError(
                "Dual-time invariant violated: expected u_sem >= u_tex, got "
                f"u_sem[{index}]={float(u_sem[index].item()):.6f}, "
                f"u_tex[{index}]={float(u_tex[index].item()):.6f}."
            )

        sigma_violation = sigmas_sem > (sigmas_tex + eps)
        if torch.any(sigma_violation):
            index = int(torch.nonzero(sigma_violation, as_tuple=False)[0, 0].item())
            sigma_sem_flat = sigmas_sem.reshape(sigmas_sem.shape[0], -1)
            sigma_tex_flat = sigmas_tex.reshape(sigmas_tex.shape[0], -1)
            raise ValueError(
                "Dual-time invariant violated: expected sigmas_sem <= sigmas_tex, got "
                f"sigmas_sem[{index}]={float(sigma_sem_flat[index, 0].item()):.6f}, "
                f"sigmas_tex[{index}]={float(sigma_tex_flat[index, 0].item()):.6f}."
            )

    def _prepare_latents(
        self,
        *,
        batch_size: int,
        height: int,
        width: int,
        generator: Optional[torch.Generator],
    ) -> tuple[Tensor, Tensor, int, int]:
        height = 2 * (int(height) // (self.vae_scale_factor * 2))
        width = 2 * (int(width) // (self.vae_scale_factor * 2))

        latents = torch.randn(
            (batch_size, self.total_channels, height // 2, width // 2),
            generator=generator,
            device=self.device,
            dtype=self.weight_dtype,
        )
        latent_ids = self.pipeline_cls._prepare_latent_ids(latents).to(self.device)
        return latents, latent_ids, height, width

    def _predict_velocity(
        self,
        transformer,
        *,
        packed_latents: Tensor,
        timesteps_sem: Tensor,
        timesteps_tex: Tensor,
        encoder_hidden_states: Tensor,
        txt_ids: Tensor,
        img_ids: Tensor,
    ) -> Tensor:
        pred = transformer(
            hidden_states=packed_latents,
            timestep_sem=timesteps_sem / 1000,
            timestep_tex=timesteps_tex / 1000,
            encoder_hidden_states=encoder_hidden_states,
            txt_ids=txt_ids,
            img_ids=img_ids,
        )
        pred = pred[:, : packed_latents.size(1)]
        return self.pipeline_cls._unpack_latents_with_ids(pred, img_ids)

    @torch.no_grad()
    def generate_batch(
        self,
        *,
        prompts: list[str],
        num_inference_steps: int,
        guidance_scale: float,
        height: int,
        width: int,
        generator: Optional[torch.Generator] = None,
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> list[Image.Image]:
        if num_inference_steps <= 0:
            raise ValueError("num_inference_steps must be > 0")

        batch_size = len(prompts)
        if batch_size == 0:
            return []

        prompt_embeds, text_ids = self.text_encoder.encode(prompts, dtype=self.weight_dtype)
        _validate_autoguidance_guidance_scale(
            self.autoguidance_enabled,
            guidance_scale,
        )

        if self.autoguidance_enabled:
            if self.autoguidance_reuse_main_text_encoder:
                autoguidance_prompt_embeds = prompt_embeds
                autoguidance_text_ids = text_ids
            else:
                autoguidance_prompt_embeds, autoguidance_text_ids = (
                    self.autoguidance_text_encoder.encode(
                        prompts,
                        dtype=self.weight_dtype,
                    )
                )
            neg_prompt_embeds = None
            neg_text_ids = None
        elif guidance_scale > 1.0:
            neg_prompts = [""] * batch_size
            neg_prompt_embeds, neg_text_ids = self.text_encoder.encode(
                neg_prompts,
                dtype=self.weight_dtype,
            )
            autoguidance_prompt_embeds = None
            autoguidance_text_ids = None
        else:
            autoguidance_prompt_embeds = None
            autoguidance_text_ids = None
            neg_prompt_embeds = None
            neg_text_ids = None

        latents, latent_ids, _, _ = self._prepare_latents(
            batch_size=batch_size,
            height=height,
            width=width,
            generator=generator,
        )

        u_base_unit = torch.linspace(
            0.0,
            1.0,
            steps=num_inference_steps + 1,
            device=self.device,
            dtype=torch.float32,
        )
        u_shifted_unit = _apply_timestep_shift_unit_interval(
            u_base_unit,
            self.timestep_shift_alpha,
        )
        _, base_sigmas_schedule = self._timesteps_and_sigmas(
            u_shifted_unit,
            n_dim=1,
            dtype=torch.float32,
        )
        u_sem_raw_schedule = u_shifted_unit * (1.0 + self.delta_t)
        if self.debug_assert_schedule:
            self._assert_shifted_schedule(
                u_base_unit=u_base_unit,
                u_sem_raw_schedule=u_sem_raw_schedule,
            )

        if progress_callback is not None:
            progress_callback(0, num_inference_steps)

        for step in range(num_inference_steps):
            u_sem_raw_cur = torch.full(
                (batch_size,),
                float(u_sem_raw_schedule[step].item()),
                device=self.device,
            )
            u_sem_raw_next = torch.full(
                (batch_size,),
                float(u_sem_raw_schedule[step + 1].item()),
                device=self.device,
            )

            u_tex_cur = torch.clamp(u_sem_raw_cur - self.delta_t, min=0.0, max=1.0)
            u_sem_cur = torch.clamp(u_sem_raw_cur, max=1.0)
            u_tex_next = torch.clamp(u_sem_raw_next - self.delta_t, min=0.0, max=1.0)
            u_sem_next = torch.clamp(u_sem_raw_next, max=1.0)

            timesteps_sem_cur, sigmas_sem_cur = self._timesteps_and_sigmas(
                u_sem_cur,
                n_dim=latents.ndim,
                dtype=latents.dtype,
            )
            timesteps_tex_cur, sigmas_tex_cur = self._timesteps_and_sigmas(
                u_tex_cur,
                n_dim=latents.ndim,
                dtype=latents.dtype,
            )
            _, sigmas_sem_next = self._timesteps_and_sigmas(
                u_sem_next,
                n_dim=latents.ndim,
                dtype=latents.dtype,
            )
            _, sigmas_tex_next = self._timesteps_and_sigmas(
                u_tex_next,
                n_dim=latents.ndim,
                dtype=latents.dtype,
            )
            if self.debug_assert_schedule:
                self._assert_dual_time_invariants(
                    u_sem_cur,
                    u_tex_cur,
                    sigmas_sem_cur,
                    sigmas_tex_cur,
                )

            guidance_active = _guidance_interval_is_active(
                base_sigmas_schedule[step],
                self.guidance_interval_sigma_lo,
                self.guidance_interval_sigma_hi,
            )
            packed_latents = self.pipeline_cls._pack_latents(latents)
            pred_cond = self._predict_velocity(
                self.transformer,
                packed_latents=packed_latents,
                timesteps_sem=timesteps_sem_cur,
                timesteps_tex=timesteps_tex_cur,
                encoder_hidden_states=prompt_embeds,
                txt_ids=text_ids,
                img_ids=latent_ids,
            )

            if not guidance_active:
                velocity = pred_cond
            elif self.autoguidance_enabled:
                pred_base = self._predict_velocity(
                    self.autoguidance_transformer,
                    packed_latents=packed_latents,
                    timesteps_sem=timesteps_sem_cur,
                    timesteps_tex=timesteps_tex_cur,
                    encoder_hidden_states=autoguidance_prompt_embeds,
                    txt_ids=autoguidance_text_ids,
                    img_ids=latent_ids,
                )
                velocity = _combine_guided_velocity(
                    pred_base,
                    pred_cond,
                    guidance_scale,
                )
            elif guidance_scale > 1.0:
                pred_uncond = self._predict_velocity(
                    self.transformer,
                    packed_latents=packed_latents,
                    timesteps_sem=timesteps_sem_cur,
                    timesteps_tex=timesteps_tex_cur,
                    encoder_hidden_states=neg_prompt_embeds,
                    txt_ids=neg_text_ids,
                    img_ids=latent_ids,
                )
                velocity = _combine_guided_velocity(
                    pred_uncond,
                    pred_cond,
                    guidance_scale,
                )
            else:
                velocity = pred_cond

            vel_sem = velocity[:, : self.semantic_channels]
            vel_tex = velocity[:, self.semantic_channels :]
            lat_sem = latents[:, : self.semantic_channels]
            lat_tex = latents[:, self.semantic_channels :]

            dt_sem = sigmas_sem_next - sigmas_sem_cur
            dt_tex = sigmas_tex_next - sigmas_tex_cur

            lat_sem = lat_sem + dt_sem * vel_sem
            lat_tex = lat_tex + dt_tex * vel_tex
            latents = torch.cat([lat_sem, lat_tex], dim=1)
            if progress_callback is not None:
                progress_callback(step + 1, num_inference_steps)

        texture_latents = latents[:, self.semantic_channels :]
        decoded = self.texture_codec.decode_texture(
            texture_latents.to(dtype=self.component_dtype),
            pipeline_cls=self.pipeline_cls,
        )
        return self.image_processor.postprocess(decoded, output_type="pil")