diffsynth/core/data/custom_units.py

import torch
import torch.nn.functional as F
from diffsynth.pipelines.wan_video import WanVideoPipeline
from diffsynth.diffusion.base_pipeline import PipelineUnit
from PIL import Image


def _frame_to_vae_video_tensor(pipe: WanVideoPipeline, frame):
    if isinstance(frame, Image.Image):
        return pipe.preprocess_video([frame])

    if not torch.is_tensor(frame):
        raise TypeError(f"Unsupported frame type for VAE encoding: {type(frame)}")

    tensor = frame
    if tensor.dim() == 3:
        # (C, H, W) -> (1, C, 1, H, W)
        if tensor.shape[0] not in (1, 3):
            tensor = tensor.permute(2, 0, 1).contiguous()
        tensor = tensor.unsqueeze(0).unsqueeze(2)
    elif tensor.dim() == 4:
        # Accept (1, C, H, W) or (C, T, H, W)
        if tensor.shape[0] == 1 and tensor.shape[1] in (1, 3):
            tensor = tensor.unsqueeze(2)
        elif tensor.shape[0] in (1, 3):
            tensor = tensor.unsqueeze(0)
        else:
            raise ValueError(f"Unexpected 4D frame tensor shape: {tuple(tensor.shape)}")
    elif tensor.dim() == 5:
        pass
    else:
        raise ValueError(f"Unsupported frame tensor rank: {tensor.dim()}")

    return tensor.to(dtype=pipe.torch_dtype, device=pipe.device)
# custom_units.py
class WanVideoUnit_GeoNeXtFused(PipelineUnit):
    """
    1. Receives noise from NoiseInitializer.
    2. Receives input_video (multi-frame).
    3. Outputs input_latents (clean GT latents) for loss computation.
    4. Outputs latents (noisy latents) for model_fn input.
    5. Outputs clean_reference_latents used by model_fn for frame replacement.
    """
    def __init__(self, num_condition_frames=1, rgb_condition_mode="first_frame", rgb_condition_scale=1.0, loss_mask_mode="depth_only"):
        super().__init__(
            
            input_params=("input_video", "noise", "tiled", "tile_size", "tile_stride", "valid_mask", "sky_mask", "target_frame_loss_mask"),
            output_params=("latents", "input_latents", "clean_reference_latents", "rgb_condition_latents", "rgb_condition_mode", "rgb_condition_scale"), 
            onload_model_names=("vae",)
        )
        self.num_condition_frames = num_condition_frames
        self.rgb_condition_mode = rgb_condition_mode
        self.rgb_condition_scale = float(rgb_condition_scale)
        self.loss_mask_mode = loss_mask_mode

    def _downsample_mask(self, mask, latent_h, latent_w):
        if mask is None:
            return None
        if mask.dim() == 2:
            mask = mask.unsqueeze(0).unsqueeze(0)
        elif mask.dim() == 3:
            mask = mask.unsqueeze(0)
        mask = mask.float()
        invalid_mask = ~mask.bool()
        src_h, src_w = invalid_mask.shape[-2:]
        scale_h = max(src_h // latent_h, 1)
        scale_w = max(src_w // latent_w, 1)
        if src_h % latent_h == 0 and src_w % latent_w == 0:
            invalid_down = F.max_pool2d(invalid_mask.float(), kernel_size=(scale_h, scale_w), stride=(scale_h, scale_w))
        else:
            invalid_down = F.interpolate(invalid_mask.float(), size=(latent_h, latent_w), mode="nearest")
        return (~invalid_down.bool()).to(mask.device)

    def _build_loss_mask(self, input_latents, valid_mask=None, sky_mask=None, target_frame_loss_mask=None):
        bsz, channels, frames, latent_h, latent_w = input_latents.shape
        loss_mask = torch.zeros((bsz, channels, frames, latent_h, latent_w), device=input_latents.device, dtype=input_latents.dtype)

        default_mask = torch.ones((1, 1, latent_h, latent_w), device=input_latents.device, dtype=torch.bool)
        valid_down = self._downsample_mask(valid_mask, latent_h, latent_w) if valid_mask is not None else default_mask
        sky_down = self._downsample_mask(sky_mask, latent_h, latent_w) if sky_mask is not None else None

        for frame_idx in range(frames):
            frame_mask = default_mask
            if frame_idx >= self.num_condition_frames:
                if self.loss_mask_mode == "all_targets":
                    frame_mask = valid_down
                    if sky_down is not None:
                        frame_mask = frame_mask | sky_down
                elif self.loss_mask_mode == "depth_only":
                    # Backward-compatible behavior: only apply geometric masks on
                    # the first predicted frame (typically depth).
                    if frame_idx == self.num_condition_frames:
                        frame_mask = valid_down
                        if sky_down is not None:
                            frame_mask = frame_mask | sky_down
                else:
                    raise ValueError(f"Unsupported loss_mask_mode: {self.loss_mask_mode}")
            loss_mask[:, :, frame_idx, :, :] = frame_mask.to(dtype=input_latents.dtype)

        # Optional source-specific supervision gating for predicted target frames.
        # target_frame_loss_mask is expected to align with frames after condition frames.
        if target_frame_loss_mask is not None:
            tfm = target_frame_loss_mask
            if isinstance(tfm, torch.Tensor):
                tfm = tfm.to(device=input_latents.device, dtype=input_latents.dtype).flatten()
            else:
                tfm = torch.tensor(tfm, device=input_latents.device, dtype=input_latents.dtype).flatten()
            n_pred = max(frames - self.num_condition_frames, 0)
            if tfm.numel() > 0 and n_pred > 0:
                n = min(n_pred, int(tfm.numel()))
                gate = tfm[:n].view(1, 1, n, 1, 1)
                loss_mask[:, :, self.num_condition_frames:self.num_condition_frames + n, :, :] *= gate
        return loss_mask

    def process(self, pipe: WanVideoPipeline, input_video, noise, tiled, tile_size, tile_stride, valid_mask=None, sky_mask=None, target_frame_loss_mask=None):
        if input_video is None or len(input_video) == 0:
            return {}
        if pipe.scheduler.training and len(input_video) < 2:
            raise ValueError(f"Training expects at least 2-frame input_video, got {len(input_video)}")
        if not pipe.scheduler.training and len(input_video) < self.num_condition_frames:
            raise ValueError(f"Inference expects at least {self.num_condition_frames} input frames, got {len(input_video)}")
        if pipe.scheduler.training and not (1 <= self.num_condition_frames < len(input_video)):
            raise ValueError(f"num_condition_frames must be in [1, {len(input_video) - 1}], got {self.num_condition_frames}")
            
        pipe.load_models_to_device(self.onload_model_names)

        if not pipe.scheduler.training:
            cond_latents = []
            for img in input_video[:self.num_condition_frames]:
                video_tensor = _frame_to_vae_video_tensor(pipe, img)
                z = pipe.vae.encode(
                    video_tensor,
                    device=pipe.device,
                    tiled=tiled,
                    tile_size=tile_size,
                    tile_stride=tile_stride
                ).to(dtype=pipe.torch_dtype, device=pipe.device)
                cond_latents.append(z)
            clean_ref = torch.cat(cond_latents, dim=2)
            return {
                "latents": noise,
                "clean_reference_latents": clean_ref,
                "rgb_condition_latents": clean_ref[:, :, :1],
                "rgb_condition_mode": self.rgb_condition_mode,
                "rgb_condition_scale": self.rgb_condition_scale,
            }
        
        encoded_latents_list = []

        for img in input_video:
            video_tensor = _frame_to_vae_video_tensor(pipe, img)

            z = pipe.vae.encode(
                video_tensor, 
                device=pipe.device, 
                tiled=tiled, 
                tile_size=tile_size, 
                tile_stride=tile_stride
            ).to(dtype=pipe.torch_dtype, device=pipe.device)
            
            encoded_latents_list.append(z)

        full_clean_latents = torch.cat(encoded_latents_list, dim=2)
        
        clean_ref = torch.cat(encoded_latents_list[:self.num_condition_frames], dim=2)
        rgb_condition = encoded_latents_list[0]

        mask = self._build_loss_mask(
            full_clean_latents,
            valid_mask=valid_mask,
            sky_mask=sky_mask,
            target_frame_loss_mask=target_frame_loss_mask,
        )

        if pipe.scheduler.training:
            return {
                "latents": noise,                 
                "input_latents": full_clean_latents, 
                "clean_reference_latents": clean_ref,
                "rgb_condition_latents": rgb_condition,
                "rgb_condition_mode": self.rgb_condition_mode,
                "rgb_condition_scale": self.rgb_condition_scale,
                "loss_mask": mask 
            }
        else:
            return {
                "latents": noise,
                "clean_reference_latents": clean_ref,
                "rgb_condition_latents": clean_ref[:, :, :1],
                "rgb_condition_mode": self.rgb_condition_mode,
                "rgb_condition_scale": self.rgb_condition_scale,
            }


# Backward-compatible alias.
WanVideoUnit_ThreeFrameFused = WanVideoUnit_GeoNeXtFused