diffusers_patch/src/diffusers/pipelines/pixeldit/pipeline_pixeldit_styled.py

"""
PixelDiT Styled Pipeline — ControlNet + IP-Adapter style transfer.

Combines:
  • ControlNet scribble conditioning (HED edge map from the reference image)
  • IP-Adapter SigLIP style conditioning
  • Flow-matching img2img variation

Reference image drives both structure (ControlNet) and style (IP-Adapter).
A text prompt is optional; leave it empty for pure reference-driven generation.
"""

from typing import Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

try:
    from .pipeline_pixeldit import PixelDiTPipeline
    from .pipeline_output import PixelDiTPipelineOutput
    from .modeling_pixeldit_controlnet import (
        PixelDiTControlNet,
        load_checkpoint,
        load_ip_adapter_checkpoint,
        unwrap_transformer,
    )
    from .image_processor_hed import HEDExtractor, control_to_tensor, hed_to_scribble
except ImportError:
    import importlib, os as _os, sys as _sys
    _pkg = _os.path.dirname(_os.path.abspath(__file__))
    if _pkg not in _sys.path:
        _sys.path.insert(0, _pkg)
    from pipeline_pixeldit import PixelDiTPipeline                          # noqa: E402
    from pipeline_output import PixelDiTPipelineOutput                      # noqa: E402
    from modeling_pixeldit_controlnet import (                               # noqa: E402
        PixelDiTControlNet,
        load_checkpoint,
        load_ip_adapter_checkpoint,
        unwrap_transformer,
    )
    from image_processor_hed import HEDExtractor, control_to_tensor, hed_to_scribble  # noqa: E402


def _pil_to_u8(image: Image.Image, width: int, height: int) -> np.ndarray:
    return np.asarray(image.convert("RGB").resize((width, height)), dtype=np.uint8).copy()


def _sigma_schedule(steps: int, flow_shift: float, device, dtype) -> torch.Tensor:
    t = torch.linspace(1.0, 0.0, steps + 1)
    return (flow_shift * t / (1.0 + (flow_shift - 1.0) * t)).to(device=device, dtype=dtype)


class PixelDiTStyledPipeline(PixelDiTPipeline):
    """
    Style-transfer pipeline for PixelDiT using ControlNet scribble conditioning
    and IP-Adapter SigLIP image conditioning.

    Load via :meth:`from_pretrained_styled`:

    .. code-block:: python

        from diffusers.pipelines.pixeldit import PixelDiTStyledPipeline

        pipe = PixelDiTStyledPipeline.from_pretrained_styled(
            "madtune/pixeldit-diffusers",
            controlnet_path="/path/to/controlnet_scribble_ip_768.pt",
            ip_adapter_path="/path/to/ip_adapter_v2.pt",   # optional
            hed_ckpt_path="/path/to/ControlNetHED.pth",    # optional
            torch_dtype=torch.bfloat16,
        )
        pipe.enable_model_cpu_offload(gpu_id=1)

        out = pipe(
            image=Image.open("style_ref.jpg"),
            prompt="gothic pale woman, dramatic rim lighting",
            variation_strength=0.85,
            ctrl_strength=0.25,
            ip_strength=0.85,
        ).images[0]

    Or from HuggingFace Hub:

    .. code-block:: python

        from huggingface_hub import hf_hub_download

        pipe = PixelDiTStyledPipeline.from_pretrained_styled(
            "madtune/pixeldit-diffusers",
            controlnet_path=hf_hub_download("madtune/pixeldit-controlnet-ip", "controlnet_scribble_ip_768.pt"),
            ip_adapter_path=hf_hub_download("madtune/pixeldit-controlnet-ip", "ip_adapter_v2.pt"),
            hed_ckpt_path=hf_hub_download("madtune/pixeldit-controlnet-ip", "ControlNetHED.pth"),
            torch_dtype=torch.bfloat16,
        )
    """

    # siglip_model and siglip_processor are optional — set as instance attrs
    # after from_pretrained_styled, not registered modules, because they use
    # a different loading path (transformers, not diffusers).
    _optional_components = ["siglip_model", "siglip_processor"]

    def __init__(self, transformer, scheduler, text_encoder, tokenizer, controlnet):
        super().__init__(
            transformer=transformer,
            scheduler=scheduler,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
        )
        self.register_modules(controlnet=controlnet)
        self.siglip_model     = None
        self.siglip_processor = None
        self._hed_extractor   = None

    # ------------------------------------------------------------------
    # Factory
    # ------------------------------------------------------------------

    @classmethod
    def from_pretrained_styled(
        cls,
        pretrained_model_name_or_path: str,
        controlnet_path: str,
        ip_adapter_path: Optional[str] = None,
        hed_ckpt_path:   Optional[str] = None,
        copy_blocks_num: int = 7,
        siglip_model_id: str = "google/siglip-so400m-patch14-384",
        **kwargs,
    ) -> "PixelDiTStyledPipeline":
        """
        Load the base PixelDiT model then attach ControlNet + IP-Adapter weights.

        Args:
            pretrained_model_name_or_path: HF repo or local path for the base model
                (e.g. ``"madtune/pixeldit-diffusers"``).
            controlnet_path: Local path to ``controlnet_scribble_ip_768.pt``.
                This checkpoint may also contain the IP-Adapter weights — if so,
                ``ip_adapter_path`` is optional.
            ip_adapter_path: Optional separate ``ip_adapter_v2.pt``.  Loaded on top
                of ``controlnet_path`` when provided.
            hed_ckpt_path: Optional path to ``ControlNetHED.pth``.  Required if you
                want automatic edge extraction from the reference image.  Omit when
                you will always pass an explicit ``control_image`` to ``__call__``.
            copy_blocks_num: Number of transformer blocks copied into the ControlNet
                branch.  Must match the training config (default 7).
            siglip_model_id: HF model id for the SigLIP encoder (default:
                ``google/siglip-so400m-patch14-384``).
            **kwargs: Forwarded to ``PixelDiTPipeline.from_pretrained``
                (e.g. ``torch_dtype``, ``device_map``).
        """
        import diffusers
        try:
            from .modeling_pixeldit_hf import PixelDiTModel
        except ImportError:
            from modeling_pixeldit_hf import PixelDiTModel

        if not hasattr(diffusers, "PixelDiTModel"):
            diffusers.PixelDiTModel = PixelDiTModel

        dtype = kwargs.get("torch_dtype", torch.float32)

        print("[PixelDiTStyledPipeline] Loading base model…")
        t2i = PixelDiTPipeline.from_pretrained(pretrained_model_name_or_path, **kwargs)

        print("[PixelDiTStyledPipeline] Building ControlNet…")
        inner = unwrap_transformer(t2i.transformer)
        controlnet = PixelDiTControlNet(inner, copy_blocks_num=copy_blocks_num)

        print(f"[PixelDiTStyledPipeline] Loading ControlNet checkpoint: {controlnet_path}")
        step = load_checkpoint(controlnet, controlnet_path)
        print(f"  step={step}")

        if ip_adapter_path is not None:
            print(f"[PixelDiTStyledPipeline] Loading IP-Adapter checkpoint: {ip_adapter_path}")
            ip_step = load_ip_adapter_checkpoint(controlnet, ip_adapter_path)
            print(f"  ip_step={ip_step}")

        controlnet = controlnet.to(dtype=dtype)

        pipe = cls(
            transformer=t2i.transformer,
            scheduler=t2i.scheduler,
            text_encoder=t2i.text_encoder,
            tokenizer=t2i.tokenizer,
            controlnet=controlnet,
        )

        print(f"[PixelDiTStyledPipeline] Loading SigLIP: {siglip_model_id}")
        from transformers import AutoImageProcessor, SiglipVisionModel
        pipe.siglip_processor = AutoImageProcessor.from_pretrained(siglip_model_id)
        pipe.siglip_model     = SiglipVisionModel.from_pretrained(
            siglip_model_id, torch_dtype=dtype
        ).eval()

        if hed_ckpt_path is not None:
            print(f"[PixelDiTStyledPipeline] Loading HED extractor: {hed_ckpt_path}")
            pipe._hed_extractor = HEDExtractor(hed_ckpt_path, device="cpu")

        return pipe

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @torch.no_grad()
    def _encode_siglip(self, image: Image.Image, device, dtype) -> torch.Tensor:
        """Return IP-Adapter features for ``image`` via SigLIP + controlnet projection."""
        inputs  = self.siglip_processor(images=image, return_tensors="pt").to(device)
        patches = self.siglip_model(
            pixel_values=inputs["pixel_values"].to(dtype)
        ).last_hidden_state                                 # [1, N, 1152]
        return self.controlnet.encode_siglip(patches)       # [1, 256, 1536]

    def _extract_control(
        self,
        image_u8: np.ndarray,
        control_image: Optional[Image.Image],
        width: int,
        height: int,
        hed_thickness: int,
    ) -> np.ndarray:
        """Return HW uint8 scribble map from either a provided image or auto-HED."""
        if control_image is not None:
            ctrl = np.asarray(
                control_image.convert("L").resize((width, height), Image.NEAREST),
                dtype=np.uint8,
            ).copy()
            return np.where(ctrl > 127, 255, 0).astype(np.uint8)

        if self._hed_extractor is None:
            raise ValueError(
                "No control_image provided and no HED extractor loaded.  "
                "Pass hed_ckpt_path to from_pretrained_styled, or supply control_image."
            )
        return self._hed_extractor(image_u8, thickness=hed_thickness)

    # ------------------------------------------------------------------
    # __call__
    # ------------------------------------------------------------------

    @torch.no_grad()
    def __call__(
        self,
        image: Image.Image,
        prompt: Union[str, List[str]] = "",
        negative_prompt: Optional[Union[str, List[str]]] = None,
        control_image: Optional[Image.Image] = None,
        width:  Optional[int] = None,
        height: Optional[int] = None,
        variation_strength: float = 0.85,
        ctrl_strength:      float = 0.25,
        ip_strength:        float = 0.85,
        flow_shift:         float = 8.0,
        guidance_scale:     float = 4.5,
        num_inference_steps: int  = 50,
        hed_thickness:      int   = 2,
        generator: Optional[torch.Generator] = None,
        output_type: str = "pil",
        return_dict: bool = True,
        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
        callback_on_step_end_tensor_inputs: List[str] = [],
        **kwargs,
    ) -> Union[PixelDiTPipelineOutput, Tuple]:
        """
        Run styled image generation.

        Args:
            image: Reference image — drives both ControlNet structure and IP-Adapter style.
            prompt: Optional text prompt.  Leave empty for pure reference-driven output.
            negative_prompt: Optional negative text.
            control_image: Optional pre-computed scribble map (white edges on black, PIL L or RGB).
                If ``None``, HED edges are extracted automatically from ``image``.
            width / height: Output resolution.  Defaults to reference image size (snapped to ×16).
            variation_strength: How much to vary from the reference (0 = copy, 1 = full noise).
                Recommended: 0.65–0.95.
            ctrl_strength: ControlNet skip scale.  0.25 is a good starting point.
            ip_strength: IP-Adapter style scale.  0.35 = subtle, 0.85 = strong.
            flow_shift: Flow-matching shift parameter.  Higher = more detail (7–8 for 768+ px).
            guidance_scale: CFG scale.  3.5–5.0 works well.
            num_inference_steps: Total denoising steps (≥ 50 recommended).
            hed_thickness: Scribble line thickness (0 = thin/erode, 2 = default, 4+ = thick).
            generator: Torch RNG for reproducibility.
            output_type: ``"pil"`` or ``"np"``.
            return_dict: Return :class:`PixelDiTPipelineOutput` if ``True``, else tuple.
            callback_on_step_end: Called at the end of each denoising step with
                ``(step_index, sigma, kwargs_dict)``.
        """
        device = self._execution_device
        dtype  = next(self.controlnet.parameters()).dtype

        # ── image size ────────────────────────────────────────────
        orig_w, orig_h = image.size
        if width is None:
            width  = (orig_w // 16) * 16
        if height is None:
            height = (orig_h // 16) * 16
        width  = (width  // 16) * 16
        height = (height // 16) * 16

        image_rgb = _pil_to_u8(image, width, height)

        # ── text encoding ─────────────────────────────────────────
        if isinstance(prompt, str):
            prompt = [prompt]
        batch_size = len(prompt)

        has_prompt = any(p.strip() for p in prompt)
        if has_prompt:
            y_text, _ = self.encode_prompt(
                prompt,
                device=device,
                dtype=dtype,
                do_classifier_free_guidance=False,
            )
        else:
            # null text — match original flow-matching behaviour (pure zeros)
            y_text = torch.zeros(batch_size, 300, 2304, dtype=dtype, device=device)
        y_null = torch.zeros_like(y_text)

        # ── control map ───────────────────────────────────────────
        if ctrl_strength > 0.0:
            scribble_u8 = self._extract_control(image_rgb, control_image, width, height, hed_thickness)
        else:
            scribble_u8 = np.zeros((height, width), dtype=np.uint8)
        ref_x = control_to_tensor(scribble_u8).unsqueeze(0).to(device, dtype=dtype)
        if batch_size > 1:
            ref_x = ref_x.expand(batch_size, -1, -1, -1).contiguous()

        # ── SigLIP / IP-Adapter features ──────────────────────────
        if ip_strength > 0.0 and self.siglip_model is not None:
            siglip_dev  = next(self.siglip_model.parameters()).device
            ip_features = self._encode_siglip(image, siglip_dev, dtype).to(device)
            if batch_size > 1:
                ip_features = ip_features.expand(batch_size, -1, -1).contiguous()
        else:
            ip_features = None

        # ── sigma schedule ────────────────────────────────────────
        sigmas    = _sigma_schedule(num_inference_steps, flow_shift, device, dtype)
        variation = float(np.clip(variation_strength, 0.0, 1.0))
        start     = max(0, min(num_inference_steps, round((1.0 - variation) * num_inference_steps)))

        # ── reference image tensor ────────────────────────────────
        ref_img = torch.from_numpy(image_rgb).to(device, dtype=dtype).permute(2, 0, 1).unsqueeze(0)
        ref_img = ref_img / 127.5 - 1.0
        if batch_size > 1:
            ref_img = ref_img.expand(batch_size, -1, -1, -1).contiguous()

        noise = torch.randn(ref_img.shape, generator=generator, dtype=dtype).to(device)

        if start >= num_inference_steps:
            x = ref_img.clone()
        else:
            sigma_s = sigmas[start]
            x = (1.0 - sigma_s) * ref_img + sigma_s * noise

        # ── CFG scale tensors ─────────────────────────────────────
        # [0] = uncond branch, [1] = cond branch
        ctrl_scales = torch.tensor([0.0, ctrl_strength], dtype=dtype, device=device)
        ip_scales   = torch.tensor([0.0, ip_strength],   dtype=dtype, device=device)

        # ── denoising loop ────────────────────────────────────────
        total = num_inference_steps - start
        self._num_timesteps = total

        ctx = (
            torch.amp.autocast("cuda", dtype=dtype)
            if device != "cpu" and torch.cuda.is_available()
            else torch.no_grad()
        )

        with torch.inference_mode(), ctx:
            for step_i, i in enumerate(self.progress_bar(range(start, num_inference_steps))):
                sigma      = sigmas[i].item()
                sigma_next = sigmas[i + 1].item()
                t_val      = torch.full((2 * batch_size,), sigma * 1000, dtype=dtype, device=device)

                x_in   = x.repeat(2, 1, 1, 1)
                y_in   = torch.cat([y_null, y_text])
                ref_in = ref_x.repeat(2, 1, 1, 1)
                ip_in  = ip_features.repeat(2, 1, 1) if ip_features is not None else None

                v_batch = self.controlnet(
                    x_in, t_val, y_in, ref_in,
                    ctrl_scale=ctrl_scales,
                    ip_features=ip_in,
                    ip_strength=ip_scales,
                )
                v_u, v_c = v_batch.chunk(2)
                v = v_u + guidance_scale * (v_c - v_u)
                x = x + (sigma_next - sigma) * v

                if callback_on_step_end is not None:
                    cb_kwargs = {k: locals().get(k) for k in callback_on_step_end_tensor_inputs}
                    callback_on_step_end(step_i, sigma, cb_kwargs)

        # ── decode ────────────────────────────────────────────────
        image_out = ((x.clamp(-1, 1) + 1) * 127.5).byte().permute(0, 2, 3, 1).cpu().numpy()

        if output_type == "pil":
            image_out = [Image.fromarray(img) for img in image_out]

        self.maybe_free_model_hooks()

        if not return_dict:
            return (image_out,)
        return PixelDiTPipelineOutput(images=image_out)