diffsynth/core/data/data_profiles.py

import json, os, random
import numpy as np
from PIL import Image
from diffsynth.core.data.operators import ImageCropAndResize,SequencialProcess,LoadImage
from diffsynth.core.data.custom_operators import (
    LoadVideoRange,
    ComputeCutMask,
    LoadRGBDepthNormalItem,
    LoadDepthMaskItem,
    ResizeShortestEdge,
    VKITTICropAndResize,
)
from diffsynth.core.data.custom_units import WanVideoUnit_GeoNeXtFused
from diffsynth.core.data.custom_model_fn import model_fn_wan_video_geonext
class BaseDataProfile:
    def __init__(self, args):
        self.args = args
    def load_and_transform(self, path): raise NotImplementedError
    def get_operator_map(self): return {}
    def get_data_keys(self): return []
    def get_extra_inputs(self): return []
    def configure_pipeline(self, pipe):pass
    def postprocess_sample(self, data): return data


import torch
import torch.nn as nn

class DepthNormalProfile(BaseDataProfile):
    def __init__(self, args):
        super().__init__(args)
        self.fixed_sample_rgb_path = str(getattr(args, "fixed_sample_rgb_path", "") or "").strip()
        self.random_flip = bool(getattr(args, "random_flip", False))
        self.interiorverse_use_mask = bool(getattr(args, "interiorverse_use_mask", True))
        self.interiorverse_cache_index = bool(getattr(args, "interiorverse_cache_index", True))
        self.refresh_interiorverse_index = bool(getattr(args, "refresh_interiorverse_index", False))
        self.interiorverse_index_cache_path = str(getattr(args, "interiorverse_index_cache_path", "") or "").strip()
        self.loss_mask_mode = str(getattr(args, "loss_mask_mode", "depth_only"))
        target_modalities = str(getattr(args, "target_modalities", "depth,normal"))
        self.target_modalities = [m.strip() for m in target_modalities.split(",") if m.strip()]
        if len(self.target_modalities) == 0:
            self.target_modalities = ["depth", "normal"]
        valid_modalities = {"depth", "normal"}
        invalid = [m for m in self.target_modalities if m not in valid_modalities]
        if invalid:
            raise ValueError(
                f"Unsupported target_modalities={invalid}. Allowed: {sorted(valid_modalities)}"
            )
        # Keep depth as first predicted slot when present so legacy depth-only mask
        # semantics still make sense.
        if "depth" in self.target_modalities:
            self.target_modalities = ["depth"] + [m for m in self.target_modalities if m != "depth"]
        self.num_condition_frames = int(getattr(args, "num_condition_frames", 1))
        self.modalities = ["rgb"] + list(self.target_modalities)
        self.supervise_modalities_by_source = {
            "hypersim": self._parse_source_supervise_modalities(getattr(args, "supervise_modalities_hypersim", "")),
            "vkitti": self._parse_source_supervise_modalities(getattr(args, "supervise_modalities_vkitti", "")),
            "interiorverse": self._parse_source_supervise_modalities(getattr(args, "supervise_modalities_interiorverse", "")),
        }

    def _parse_source_supervise_modalities(self, text):
        text = str(text or "").strip()
        if len(text) == 0:
            return set(self.target_modalities)
        items = {m.strip() for m in text.split(",") if m.strip()}
        return {m for m in items if m in set(self.target_modalities)}

    def _default_interiorverse_cache_path(self, root_dir, fov_mode):
        safe_modalities = "_".join(self.target_modalities)
        norm_root = os.path.normpath(root_dir)
        # Preferred layout for current cluster dataset path:
        #   /.../datasets/interiorVerse/interiorverse/interverse
        # Save cache at:
        #   /.../datasets/interiorVerse/.index_cache
        if norm_root.endswith(os.path.join("interiorverse", "interverse")):
            cache_base = os.path.dirname(os.path.dirname(norm_root))
            cache_dir = os.path.join(cache_base, ".index_cache")
        else:
            cache_dir = os.path.join(root_dir, ".index_cache")
        return os.path.join(cache_dir, f"interiorverse_index_{fov_mode}_{safe_modalities}.json")

    def _mix_sources(self, source_to_data, mix_seed):
        available = {k: v for k, v in source_to_data.items() if len(v) > 0}
        if len(available) == 0:
            return []
        if len(available) == 1:
            only_source = next(iter(available.keys()))
            print(f"[DepthNormalProfile] Only {only_source} is available; sampling ratio options are ignored.")
            return list(next(iter(available.values())))

        p_h = float(getattr(self.args, "p_hypersim", 0.5))
        p_i = float(getattr(self.args, "p_interiorverse", 0.0))
        p_h = max(0.0, p_h)
        p_i = max(0.0, p_i)
        p_v = max(0.0, 1.0 - p_h - p_i)
        raw_weights = {
            "hypersim": p_h,
            "vkitti": p_v,
            "interiorverse": p_i,
        }
        weights = {k: raw_weights.get(k, 0.0) for k in available.keys()}
        total_w = sum(weights.values())
        if total_w <= 0:
            uniform = 1.0 / len(available)
            weights = {k: uniform for k in available.keys()}
        else:
            weights = {k: v / total_w for k, v in weights.items()}

        total_len = sum(len(v) for v in available.values())
        counts = {k: int(round(total_len * w)) for k, w in weights.items()}
        # Ensure every available source contributes at least one sample.
        for k in list(counts.keys()):
            counts[k] = max(1, counts[k])
        # Fix total count drift.
        drift = sum(counts.values()) - total_len
        keys_sorted = sorted(counts.keys(), key=lambda k: counts[k], reverse=(drift > 0))
        idx = 0
        while drift != 0 and len(keys_sorted) > 0:
            key = keys_sorted[idx % len(keys_sorted)]
            if drift > 0 and counts[key] > 1:
                counts[key] -= 1
                drift -= 1
            elif drift < 0:
                counts[key] += 1
                drift += 1
            idx += 1
            if idx > 100000:
                break

        rng = random.Random(mix_seed)
        mixed = []
        for key, data in available.items():
            n = counts[key]
            if n <= len(data):
                sampled = rng.sample(data, n)
            else:
                sampled = [rng.choice(data) for _ in range(n)]
            mixed.extend(sampled)
        rng.shuffle(mixed)
        print(
            "[DepthNormalProfile] Mixed sources with ratios="
            + ", ".join([f"{k}:{weights[k]:.3f}" for k in sorted(weights.keys())])
            + " -> counts="
            + ", ".join([f"{k}:{counts[k]}" for k in sorted(counts.keys())])
            + f", total={len(mixed)}"
        )
        return mixed

    def _expand_patch_embedding_for_concat(self, dit, extra_in_channels):
        old_conv = dit.patch_embedding
        new_in_channels = old_conv.in_channels + int(extra_in_channels)
        if old_conv.in_channels == new_in_channels:
            return
        if old_conv.in_channels > new_in_channels:
            raise ValueError(
                f"Cannot shrink patch_embedding from {old_conv.in_channels} to {new_in_channels} channels."
            )

        new_conv = nn.Conv3d(
            new_in_channels,
            old_conv.out_channels,
            kernel_size=old_conv.kernel_size,
            stride=old_conv.stride,
            padding=old_conv.padding,
            dilation=old_conv.dilation,
            groups=old_conv.groups,
            bias=old_conv.bias is not None,
            padding_mode=old_conv.padding_mode,
        ).to(device=old_conv.weight.device, dtype=old_conv.weight.dtype)
        with torch.no_grad():
            new_conv.weight.zero_()
            new_conv.weight[:, :old_conv.in_channels].copy_(old_conv.weight)
            if old_conv.bias is not None:
                new_conv.bias.copy_(old_conv.bias)

        dit.patch_embedding = new_conv
        dit.in_dim = new_in_channels
        print(
            f"[Config] Expanded DiT patch_embedding input channels "
            f"{old_conv.in_channels} -> {new_in_channels} for rgb_condition_mode=concat."
        )

    def _build_entry(self, source, modality_paths, valid_mask_item=None, sky_mask_item=None):
        rgb_kind = "rgb_exr" if source == "interiorverse" else "rgb"
        video = [{"path": modality_paths["rgb"], "kind": rgb_kind, "source": source}]
        target_kinds_in_video = []
        for kind in self.target_modalities:
            path = modality_paths.get(kind)
            if path is None:
                continue
            video.append({"path": path, "kind": kind, "source": source})
            target_kinds_in_video.append(kind)
        if len(video) < 2:
            return None
        first_target = video[min(1, len(video) - 1)].copy()
        if valid_mask_item is None and source in ("vkitti", "hypersim"):
            valid_mask_item = first_target.copy()
        if sky_mask_item is None and source in ("vkitti", "hypersim"):
            sky_mask_item = first_target.copy()
        entry = {
            "video": video,
            "prompt": "",
            "first_frame": video[0].copy(),
            "second_frame": first_target,
            "sample_rgb_path": modality_paths["rgb"],
            "sample_source": source,
        }
        supervise_set = self.supervise_modalities_by_source.get(source, set(self.target_modalities))
        target_frame_loss_mask = [1.0 if kind in supervise_set else 0.0 for kind in target_kinds_in_video]
        entry["target_frame_loss_mask"] = torch.tensor(target_frame_loss_mask, dtype=torch.float32)
        if valid_mask_item is not None:
            entry["valid_mask"] = valid_mask_item
        if sky_mask_item is not None:
            entry["sky_mask"] = sky_mask_item
        return entry

    def _load_vkitti(self, root_dir):
        flattened_data = []
        scenes = ["02", "06", "18", "20"]
        conditions = [
            "15-deg-left", "15-deg-right", "30-deg-left", "30-deg-right", "clone",
            "fog", "morning", "overcast", "rain", "sunset"
        ]
        cameras = ["0", "1"]

        for scene in scenes:
            for condition in conditions:
                for camera in cameras:
                    image_dir = os.path.join(root_dir, f"Scene{scene}/{condition}/frames/rgb/Camera_{camera}")
                    if not os.path.isdir(image_dir):
                        continue
                    for image_name in sorted(os.listdir(image_dir)):
                        rgb_path = os.path.join(image_dir, image_name)
                        depth_path = rgb_path.replace("/rgb/", "/depth/").replace("rgb_", "depth_").replace(".jpg", ".png")
                        normal_path = rgb_path.replace("/rgb/", "/normal/").replace("rgb_", "normal_").replace(".jpg", ".png")
                        modality_paths = {"rgb": rgb_path}
                        if "depth" in self.target_modalities:
                            modality_paths["depth"] = depth_path
                        if "normal" in self.target_modalities:
                            modality_paths["normal"] = normal_path
                        missing = [k for k in self.target_modalities if k not in modality_paths or not os.path.exists(modality_paths[k])]
                        if len(missing) > 0 and any(m in ("depth", "normal") for m in missing):
                            continue
                        entry = self._build_entry("vkitti", modality_paths)
                        if entry is not None:
                            flattened_data.append(entry)
        print(f"[DepthNormalProfile] Loaded {len(flattened_data)} VKITTI samples.")
        return flattened_data

    def _load_hypersim(self, root_dir):
        flattened_data = []
        split = getattr(self.args, "depthnormal_hypersim_split", "train")
        split_dir = os.path.join(root_dir, split)
        if not os.path.isdir(split_dir):
            print(f"[DepthNormalProfile] Hypersim split directory not found: {split_dir}")
            return flattened_data

        for current_root, _, files in os.walk(split_dir):
            for file_name in sorted(files):
                if not file_name.endswith("tonemap.jpg"):
                    continue
                rgb_path = os.path.join(current_root, file_name)
                modality_paths = {"rgb": rgb_path}
                if "depth" in self.target_modalities:
                    modality_paths["depth"] = rgb_path.replace("final_preview", "geometry_hdf5").replace("tonemap.jpg", "depth_meters.hdf5")
                if "normal" in self.target_modalities:
                    modality_paths["normal"] = rgb_path.replace("final_preview", "geometry_hdf5").replace("tonemap.jpg", "normal_cam.hdf5")
                missing = [k for k in self.target_modalities if k not in modality_paths or not os.path.exists(modality_paths[k])]
                if len(missing) > 0 and any(m in ("depth", "normal") for m in missing):
                    continue
                entry = self._build_entry("hypersim", modality_paths)
                if entry is not None:
                    flattened_data.append(entry)
        print(f"[DepthNormalProfile] Loaded {len(flattened_data)} Hypersim samples.")
        return flattened_data

    def _load_interiorverse(self, root_dir):
        def _is_nonempty_file(path):
            return os.path.isfile(path) and (os.path.getsize(path) > 0)

        fov_mode = str(getattr(self.args, "interiorverse_fov", "85")).lower()
        cache_path = self.interiorverse_index_cache_path or self._default_interiorverse_cache_path(root_dir, fov_mode)
        use_cache = self.interiorverse_cache_index and not self.refresh_interiorverse_index
        if use_cache and os.path.isfile(cache_path):
            try:
                with open(cache_path, "r", encoding="utf-8") as handle:
                    cached = json.load(handle)
                if isinstance(cached, list) and len(cached) > 0:
                    print(f"[DepthNormalProfile] Loaded InteriorVerse index cache: {cache_path} ({len(cached)} samples)")
                    return cached
            except Exception as exc:
                print(f"[DepthNormalProfile] Failed to read InteriorVerse index cache ({cache_path}): {exc}")

        flattened_data = []
        subset_dirs = []
        if fov_mode in ("85", "both", "all"):
            subset_dirs.append(os.path.join(root_dir, "85"))
        if fov_mode in ("120", "both", "all"):
            subset_dirs.append(os.path.join(root_dir, "120_part"))

        for subset in subset_dirs:
            if not os.path.isdir(subset):
                print(f"[DepthNormalProfile] InteriorVerse subset not found: {subset}")
                continue
            for scene_name in sorted(os.listdir(subset)):
                scene_dir = os.path.join(subset, scene_name)
                if not os.path.isdir(scene_dir):
                    continue
                rgb_names = [n for n in os.listdir(scene_dir) if n.endswith("_im.exr")]
                for rgb_name in sorted(rgb_names):
                    stem = rgb_name[:-7]
                    modality_paths = {"rgb": os.path.join(scene_dir, rgb_name)}
                    if "depth" in self.target_modalities:
                        modality_paths["depth"] = os.path.join(scene_dir, f"{stem}_depth.exr")
                    if "normal" in self.target_modalities:
                        modality_paths["normal"] = os.path.join(scene_dir, f"{stem}_normal.exr")
                    required_keys = ["rgb"] + list(self.target_modalities)
                    missing = [
                        k
                        for k in required_keys
                        if (k not in modality_paths) or (not _is_nonempty_file(modality_paths[k]))
                    ]
                    if len(missing) > 0:
                        continue

                    valid_mask_item = None
                    sky_mask_item = None
                    if self.interiorverse_use_mask:
                        mask_path = os.path.join(scene_dir, f"{stem}_mask.exr")
                        if _is_nonempty_file(mask_path):
                            valid_mask_item = {"path": mask_path, "kind": "mask", "source": "interiorverse"}
                            sky_mask_item = {"path": mask_path, "kind": "mask", "source": "interiorverse"}
                    entry = self._build_entry(
                        "interiorverse",
                        modality_paths,
                        valid_mask_item=valid_mask_item,
                        sky_mask_item=sky_mask_item,
                    )
                    if entry is not None:
                        flattened_data.append(entry)

        if self.interiorverse_cache_index and len(flattened_data) > 0:
            try:
                os.makedirs(os.path.dirname(cache_path), exist_ok=True)
                with open(cache_path, "w", encoding="utf-8") as handle:
                    json.dump(flattened_data, handle)
                print(f"[DepthNormalProfile] Saved InteriorVerse index cache: {cache_path} ({len(flattened_data)} samples)")
            except Exception as exc:
                print(f"[DepthNormalProfile] Failed to write InteriorVerse index cache ({cache_path}): {exc}")

        print(f"[DepthNormalProfile] Loaded {len(flattened_data)} InteriorVerse samples.")
        return flattened_data

    def load_and_transform(self, metadata_path):
        flattened_data = []
        hypersim_root = getattr(self.args, "hypersim_root", None)
        vkitti_root = getattr(self.args, "vkitti_root", None)
        interiorverse_root = getattr(self.args, "interiorverse_root", None)
        mix_seed = int(getattr(self.args, "mix_seed", 42))

        hypersim_data = []
        vkitti_data = []
        interiorverse_data = []
        if hypersim_root:
            hypersim_data = self._load_hypersim(hypersim_root)
        if vkitti_root:
            vkitti_data = self._load_vkitti(vkitti_root)
        if interiorverse_root:
            interiorverse_data = self._load_interiorverse(interiorverse_root)

        if len(hypersim_data) == 0 and len(vkitti_data) == 0 and len(interiorverse_data) == 0:
            raise ValueError(
                "DepthNormalProfile found no samples. Set --hypersim_root and/or --vkitti_root and/or --interiorverse_root "
                "to valid dataset directories."
            )

        # Hard lock to one exact source sample if requested.
        if len(self.fixed_sample_rgb_path) > 0:
            all_entries = hypersim_data + vkitti_data + interiorverse_data
            target = os.path.normpath(self.fixed_sample_rgb_path)
            matched = None
            for entry in all_entries:
                try:
                    first_item = entry["video"][0]
                    rgb_path = first_item["path"] if isinstance(first_item, dict) else first_item
                except Exception:
                    continue
                if os.path.normpath(str(rgb_path)) == target:
                    matched = entry
                    break
            if matched is None:
                raise ValueError(
                    f"fixed_sample_rgb_path not found in loaded data: {self.fixed_sample_rgb_path}"
                )
            print(f"[DepthNormalProfile] Using fixed_sample_rgb_path: {self.fixed_sample_rgb_path}")
            return [matched]

        flattened_data = self._mix_sources(
            {
                "hypersim": hypersim_data,
                "vkitti": vkitti_data,
                "interiorverse": interiorverse_data,
            },
            mix_seed=mix_seed,
        )

        print(f"[DepthNormalProfile] Total mixed samples: {len(flattened_data)}")
        max_samples = int(getattr(self.args, "depthnormal_max_samples", 0) or 0)
        if max_samples > 0 and len(flattened_data) > max_samples:
            flattened_data = flattened_data[:max_samples]
            print(f"[DepthNormalProfile] Truncated to depthnormal_max_samples={max_samples}.")
        return flattened_data

    def get_operator_map(self):
        hypersim_resolution = getattr(self.args, "resolution_hypersim", 576)
        interiorverse_resolution = getattr(self.args, "resolution_interiorverse", 576)
        processor_map = {
            ("hypersim", "rgb"): ResizeShortestEdge(hypersim_resolution, interpolation=Image.BILINEAR),
            ("hypersim", "depth"): ResizeShortestEdge(hypersim_resolution, interpolation=Image.NEAREST),
            ("hypersim", "normal"): ResizeShortestEdge(hypersim_resolution, interpolation=Image.NEAREST),
            ("vkitti", "rgb"): VKITTICropAndResize(interpolation=Image.BILINEAR),
            ("vkitti", "depth"): VKITTICropAndResize(interpolation=Image.NEAREST),
            ("vkitti", "normal"): VKITTICropAndResize(interpolation=Image.NEAREST),
            ("interiorverse", "rgb_exr"): ResizeShortestEdge(interiorverse_resolution, interpolation=Image.BILINEAR),
            ("interiorverse", "depth"): ResizeShortestEdge(interiorverse_resolution, interpolation=Image.NEAREST),
            ("interiorverse", "normal"): ResizeShortestEdge(interiorverse_resolution, interpolation=Image.NEAREST),
        }
        load_item = LoadRGBDepthNormalItem(
            frame_processor_map=processor_map,
            norm_type=getattr(self.args, "norm_type", "trunc_disparity"),
            truncnorm_min=getattr(self.args, "truncnorm_min", 0.02),
            align_cam_normal=bool(getattr(self.args, "align_cam_normal", False)),
            output_tensor=True,
        )
        valid_mask_loader = LoadDepthMaskItem(mask_kind="valid")
        sky_mask_loader = LoadDepthMaskItem(mask_kind="sky")
        return {
            "video": SequencialProcess(load_item),
            "first_frame": load_item,
            "second_frame": load_item,
            "valid_mask": valid_mask_loader,
            "sky_mask": sky_mask_loader,
        }

    def get_data_keys(self):
        return ["video", "first_frame", "second_frame", "valid_mask", "sky_mask"]

    def get_extra_inputs(self):
        return ["first_frame", "second_frame", "valid_mask", "sky_mask", "target_frame_loss_mask"]

    def postprocess_sample(self, data):
        if not self.random_flip or random.random() <= 0.5:
            return data

        if "video" in data and len(data["video"]) >= 1:
            for frame_idx, frame in enumerate(data["video"]):
                if isinstance(frame, Image.Image):
                    flipped = frame.transpose(Image.FLIP_LEFT_RIGHT)
                    data["video"][frame_idx] = flipped
                elif isinstance(frame, torch.Tensor):
                    flipped = torch.flip(frame, dims=[-1]).clone()
                    # For target normal frame, invert x after horizontal flip.
                    if frame_idx > 0:
                        tgt_idx = frame_idx - 1
                        if tgt_idx < len(self.target_modalities) and self.target_modalities[tgt_idx] == "normal":
                            flipped[0] = -flipped[0]
                    data["video"][frame_idx] = flipped

        if "first_frame" in data and isinstance(data["first_frame"], Image.Image):
            data["first_frame"] = data["first_frame"].transpose(Image.FLIP_LEFT_RIGHT)
        elif "first_frame" in data and isinstance(data["first_frame"], torch.Tensor):
            data["first_frame"] = torch.flip(data["first_frame"], dims=[-1])

        if "second_frame" in data and isinstance(data["second_frame"], Image.Image):
            data["second_frame"] = data["second_frame"].transpose(Image.FLIP_LEFT_RIGHT)
        elif "second_frame" in data and isinstance(data["second_frame"], torch.Tensor):
            data["second_frame"] = torch.flip(data["second_frame"], dims=[-1])

        if "valid_mask" in data and isinstance(data["valid_mask"], torch.Tensor):
            data["valid_mask"] = torch.flip(data["valid_mask"], dims=[-1])
        if "sky_mask" in data and isinstance(data["sky_mask"], torch.Tensor):
            data["sky_mask"] = torch.flip(data["sky_mask"], dims=[-1])

        return data

    def configure_pipeline(self, pipe):
        print(
            "[Config] Configuring GeoNeXt RGB->depth/normal fused inference "
            f"(targets={self.target_modalities}, num_condition_frames={self.num_condition_frames})."
        )

        current_in_dim = pipe.dit.in_dim
        if current_in_dim not in [16, 48]:
            raise ValueError(
                f"[Config] Unsupported GeoNeXt base-model input dimension: {current_in_dim}. "
                "Expected a 16-channel T2V or 48-channel Ti2V model."
            )

        pipe.dit.seperated_timestep = True
        units_to_remove = [
            "WanVideoUnit_ShapeChecker",
            "WanVideoUnit_InputVideoEmbedder",
            "WanVideoUnit_ImageEmbedderVAE",
            "WanVideoUnit_ImageEmbedderFused",
        ]
        pipe.units = [u for u in pipe.units if u.__class__.__name__ not in units_to_remove]

        insert_index = 0
        for i, unit in enumerate(pipe.units):
            if unit.__class__.__name__ == "WanVideoUnit_NoiseInitializer":
                insert_index = i + 1
                break

        rgb_condition_mode = getattr(self.args, "rgb_condition_mode", "first_frame")
        rgb_condition_scale = float(getattr(self.args, "rgb_condition_scale", 1.0))
        if rgb_condition_mode == "concat":
            self._expand_patch_embedding_for_concat(pipe.dit, extra_in_channels=16)
        pipe.units.insert(
            insert_index,
            WanVideoUnit_GeoNeXtFused(
                num_condition_frames=self.num_condition_frames,
                rgb_condition_mode=rgb_condition_mode,
                rgb_condition_scale=rgb_condition_scale,
                loss_mask_mode=self.loss_mask_mode,
            ),
        )
        pipe.model_fn = model_fn_wan_video_geonext
        print(
            "[Config] Switched to model_fn_wan_video_geonext "
            f"(rgb_condition_mode={rgb_condition_mode}, rgb_condition_scale={rgb_condition_scale}, "
            f"loss_mask_mode={self.loss_mask_mode})."
        )