core/flux_features.py

"""Shared FLUX feature machinery: load the pipeline, extract multi-timestep
activations + concept-attention maps at native resolution, and cache them."""

import numpy as np
import torch
import torch.nn.functional as F
from pathlib import Path
from torchvision.transforms.functional import to_pil_image

from flux_concept_attention import (
    FluxWithConceptAttentionPipeline,
    FluxTransformer2DModelWithConceptAttention,
)


def load_flux_pipeline(config, device="cuda", flux_model=None):
    """Load the FLUX.1-dev concept-attention pipeline. Returns (pipeline, transformer)."""
    flux_model = flux_model or config["model"]["flux_model"]
    transformer = FluxTransformer2DModelWithConceptAttention.from_pretrained(
        flux_model, subfolder="transformer", torch_dtype=torch.float16
    )
    pipeline = FluxWithConceptAttentionPipeline.from_pretrained(
        flux_model, transformer=transformer, torch_dtype=torch.float16
    ).to(device)
    pipeline.set_progress_bar_config(disable=True)
    return pipeline, transformer


class MultiTimestepFeatureCache:
    """On-disk cache of multi-timestep features + concept maps (one dir per image)."""

    def __init__(self, cache_dir: str):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def _image_dir(self, image_name: str) -> Path:
        return self.cache_dir / image_name

    def has_multi_timestep_features(self, image_name: str) -> bool:
        return (self._image_dir(image_name) / "multi_timestep_data.pt").exists()

    def save_multi_timestep_features(self, image_name: str, timestep_data: dict):
        d = self._image_dir(image_name)
        d.mkdir(parents=True, exist_ok=True)
        torch.save(timestep_data, d / "multi_timestep_data.pt")

    def load_multi_timestep_features(self, image_name: str, device="cuda"):
        d = self._image_dir(image_name)
        if not self.has_multi_timestep_features(image_name):
            return None

        # weights_only=False is safe here: we only load our own cached tensors.
        try:
            data = torch.load(d / "multi_timestep_data.pt", map_location=device, weights_only=False)
        except (RuntimeError, EOFError) as e:
            print(f"[CACHE ERROR] Corrupted cache for {image_name}: {e}")
            import shutil
            shutil.rmtree(d, ignore_errors=True)
            return None

        for timestep in data["features"]:
            data["features"][timestep]["single_features"] = [
                f.to(device) for f in data["features"][timestep]["single_features"]
            ]

        if "concept_maps" not in data:
            data["concept_maps"] = {}
        for timestep in data["concept_maps"]:
            for concept in data["concept_maps"][timestep]:
                cmap = data["concept_maps"][timestep][concept]
                if hasattr(cmap, "to"):
                    data["concept_maps"][timestep][concept] = cmap.to(device)
                elif isinstance(cmap, np.ndarray):
                    data["concept_maps"][timestep][concept] = torch.from_numpy(cmap).to(device)

        return data


class MultiTimestepFeatureExtractor:
    """Extract FLUX features at native resolution.

    ``extract_or_load`` does a resolution-aware cache check then extracts (training/
    inference); ``extract_features`` returns a CPU dict for the extraction scripts to save.
    """

    def __init__(self, pipeline, config, cache_dir: str, num_timesteps: int = 4, captions: dict = None):
        self.pipeline = pipeline
        self.config = config
        self.cache = MultiTimestepFeatureCache(cache_dir)
        self.num_timesteps = num_timesteps
        self.captions = captions or {}

    def _calculate_shift(self, image_seq_len, base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15):
        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
        b = base_shift - m * base_seq_len
        return max(base_shift, min(max_shift, image_seq_len * m + b))

    def _setup_scheduler_for_resolution(self, height: int, width: int):
        # FLUX uses 16x VAE downsampling; the scheduler shift depends on sequence length.
        image_seq_len = (height // 16) * (width // 16)
        mu = self._calculate_shift(image_seq_len)
        num_inference_steps = self.config.get("flux", {}).get("timesteps", 28)
        self.pipeline.scheduler.set_timesteps(num_inference_steps, mu=mu)
        return mu

    def _get_evenly_spaced_timesteps(self):
        all_timesteps = self.pipeline.scheduler.timesteps
        group_size = 7
        indices = [-(i * group_size + 1) for i in range(self.num_timesteps) if i * group_size < len(all_timesteps)]
        selected = [int(all_timesteps[i]) for i in indices]
        selected = [t for t in selected if t != 1000][:self.num_timesteps]
        return sorted(selected, reverse=True)

    def _run_extraction(self, pil_image, image_name, concepts, resolution, prompt):
        """Core extraction loop. Returns a CPU-tensor timestep_data dict."""
        height, width = resolution
        mu = self._setup_scheduler_for_resolution(height, width)
        selected_timesteps = self._get_evenly_spaced_timesteps()

        timestep_data = {
            "timesteps": selected_timesteps,
            "features": {},
            "concept_maps": {},
            "image_name": image_name,
            "concepts": concepts,
            "resolution": (height, width),
            "prompt": prompt,
            "mu": mu,
        }
        concept_attention_kwargs = {
            "concepts": concepts,
            "timesteps": self.config["flux"]["concept_timesteps"],
            "layers": self.config["flux"]["concept_layers"],
        }

        for timestep in selected_timesteps:
            with torch.no_grad():
                output = self.pipeline(
                    prompt=prompt,
                    image=pil_image,
                    height=height,
                    width=width,
                    timesteps=[timestep],
                    num_inference_steps=1,
                    guidance_scale=self.config["flux"]["guidance_scale"],
                    concept_attention_kwargs=concept_attention_kwargs,
                    generator=torch.Generator("cuda").manual_seed(42),
                    output_type="latent",
                )

            _, single_features = self.pipeline.transformer.get_features()
            concept_maps_raw = output.concept_attention_maps
            maps_list = (
                concept_maps_raw[0]
                if (len(concept_maps_raw) == 1 and isinstance(concept_maps_raw[0], list))
                else concept_maps_raw
            )
            if len(maps_list) != len(concepts):
                raise ValueError(f"Mismatch: {len(concepts)} concepts vs {len(maps_list)} maps")
            concept_maps = {c: maps_list[i] for i, c in enumerate(concepts)}

            timestep_data["features"][timestep] = {
                "single_features": [f.cpu() for f in single_features]
            }
            timestep_data["concept_maps"][timestep] = {
                concept: (cmap.cpu() if hasattr(cmap, "cpu") else cmap)
                for concept, cmap in concept_maps.items()
            }

        return timestep_data

    def _to_device(self, timestep_data, device):
        for timestep in timestep_data["features"]:
            timestep_data["features"][timestep]["single_features"] = [
                f.to(device) for f in timestep_data["features"][timestep]["single_features"]
            ]
        for timestep in timestep_data["concept_maps"]:
            for concept in timestep_data["concept_maps"][timestep]:
                cmap = timestep_data["concept_maps"][timestep][concept]
                if hasattr(cmap, "to"):
                    timestep_data["concept_maps"][timestep][concept] = cmap.to(device)
        return timestep_data

    def extract_or_load(self, image, image_name, concepts, resolution, prompt=None, force=False):
        """Extract or load resolution-aware cached features; returns tensors on device."""
        device = next(self.pipeline.transformer.parameters()).device
        height, width = resolution

        if not force and self.cache.has_multi_timestep_features(image_name):
            cached = self.cache.load_multi_timestep_features(image_name, device)
            if cached is not None and cached.get("resolution") == (height, width):
                return cached

        if prompt is None:
            prompt = self.captions.get(image_name, "")
        pil_image = to_pil_image(image[0].cpu())
        timestep_data = self._run_extraction(pil_image, image_name, concepts, (height, width), prompt)

        self.cache.save_multi_timestep_features(image_name, timestep_data)
        return self._to_device(timestep_data, device)

    def extract_features(self, images, prompts, image_names, concepts, height, width):
        """Batch-style extraction (single image per call); returns a CPU dict to save."""
        pil_image = to_pil_image(images[0].cpu())
        return self._run_extraction(pil_image, image_names[0], concepts, (height, width), prompts[0])


def distribute_concepts(concept_maps, num_features, device):
    """Split concept maps roughly evenly across ``num_features`` feature layers."""
    processed = []
    for _, t in concept_maps.items():
        if not hasattr(t, "to"):
            t = torch.from_numpy(t).float().to(device)
        elif t.device != device:
            t = t.to(device)

        if t.dim() == 2:
            t = t.view(1, 1, t.shape[0], t.shape[1])
        elif t.dim() == 3:
            if t.shape[0] == 1:
                t = t.unsqueeze(1)
            elif t.shape[2] == 1:
                t = t.permute(2, 0, 1).unsqueeze(0)
            else:
                t = t.unsqueeze(1)
        processed.append(t)

    n = len(processed)
    per = n // num_features
    rem = n % num_features
    out = []
    start = 0
    for i in range(num_features):
        cnt = per + (1 if i < rem else 0)
        end = start + cnt
        if cnt > 0:
            out.append(torch.cat(processed[start:end], dim=1))
        else:
            spatial_h, spatial_w = processed[0].shape[-2:]
            out.append(torch.zeros(1, 1, spatial_h, spatial_w, device=device))
        start = end
    return out


def distribute_concepts_across_layers(concept_maps, num_layers, target_size, device):
    """Depth variant of :func:`distribute_concepts`: resizes every map to
    ``target_size`` first and emits zero-channel tensors for empty layers."""
    processed = []
    for _, v in concept_maps.items():
        if isinstance(v, np.ndarray):
            t = torch.from_numpy(v).float().to(device)
        elif hasattr(v, "to"):
            t = v.float().to(device)
        else:
            t = torch.tensor(v).float().to(device)

        if t.dim() == 2:
            t = t.view(1, 1, t.shape[0], t.shape[1])
        elif t.dim() == 3:
            if t.shape[0] == 1:
                t = t.unsqueeze(1)
            elif t.shape[2] == 1:
                t = t.permute(2, 0, 1).unsqueeze(0)
            else:
                t = t.unsqueeze(1)

        if t.shape[-2:] != target_size:
            t = F.interpolate(t, size=target_size, mode="bilinear", align_corners=False)
        processed.append(t)

    num_concepts = len(processed)
    if num_concepts == 0:
        return [torch.zeros(1, 0, target_size[0], target_size[1], device=device) for _ in range(num_layers)]

    per_layer = num_concepts // num_layers
    remainder = num_concepts % num_layers
    distributed = []
    start_idx = 0
    for i in range(num_layers):
        count = per_layer + (1 if i < remainder else 0)
        end_idx = start_idx + count
        if count > 0:
            distributed.append(torch.cat(processed[start_idx:end_idx], dim=1))
        else:
            distributed.append(torch.zeros(1, 0, target_size[0], target_size[1], device=device))
        start_idx = end_idx
    return distributed