Upload folder using huggingface_hub

README.md

@@ -1,13 +1,31 @@
 ---
-title: Mmdiff
-emoji: 📉
-colorFrom: purple
-colorTo: gray
+title: MMDiff
+emoji: 🎨
+colorFrom: blue
+colorTo: yellow
 sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.12'
+sdk_version: "5.0.0"
 app_file: app.py
-pinned: false
+short_description: Multi-modal generation with diffusion transformers
+python_version: "3.10"
+startup_duration_timeout: 600
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MMDiff: Extending Diffusion Transformers for Multi-Modal Generation
+
+This Space demonstrates MMDiff, a method that extends frozen diffusion transformers (FLUX.1-dev) to generate images alongside dense predictions (saliency maps, segmentation maps, depth maps) in a single forward pass.
+
+## How it works
+
+1. A text prompt is used to generate an image with FLUX.1-dev
+2. During denoising, intermediate transformer features and concept attention maps are captured
+3. Lightweight trained decoder heads (DPT, DeepLabV3+) decode these features into dense predictions:
+   - **Saliency** (DUTS): Binary foreground/background segmentation
+   - **Segmentation** (Pascal VOC): 21-class semantic segmentation
+   - **Depth** (NYU Depth V2): Monocular depth estimation
+
+## Model
+
+- **Backbone**: [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) (frozen)
+- **Decoder weights**: [yagmurakarken/mmdiff](https://huggingface.co/yagmurakarken/mmdiff)
+- **Paper**: [MMDiff: Extending Diffusion Transformers for Multi-Modal Generation](https://huggingface.co/papers/2606.16673)

app.py

@@ -0,0 +1,842 @@
+"""MMDiff: Multi-Modal Generation with Diffusion Transformers.
+
+This demo generates an image from a text prompt using FLUX.1-dev while simultaneously
+producing dense predictions (saliency, segmentation, depth) from the frozen backbone's
+intermediate features via lightweight trained decoder heads.
+"""
+
+import spaces
+import torch
+import torch.nn.functional as F
+import numpy as np
+import tempfile
+import os
+import sys
+import yaml
+from pathlib import Path
+from PIL import Image
+from torchvision import transforms
+
+# Add local modules to path
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# Download checkpoints before any model loading
+from download_checkpoints import download_all
+download_all()
+
+from core import (
+    load_config, load_flux_pipeline, MultiTimestepFeatureCache,
+    resolve_c_dino, build_dino_extractor, build_hyperfeature_fusion,
+    calculate_distributed_concept_channels, distribute_concepts,
+    distribute_concepts_across_layers,
+)
+from flux_concept_attention import (
+    FluxWithConceptAttentionPipeline,
+    FluxTransformer2DModelWithConceptAttention,
+)
+from models.hyperfeature_fusion import create_hyperfeature_fusion
+from models.dpt_segmentation_decoder import OriginalDPTSegmentationDecoder
+from models.dpt_decoder import DPTHeadSpatial
+from models.dpt_backbone import DPTRefineNetStack
+from models.blocks import FeatureFusionBlock, _make_scratch, ResidualConvUnit
+from models.segmentation_losses import CombinedSegmentationLoss
+import torch.nn as nn
+
+# ---- Model builders (mirroring scripts/inference.py + training scripts) ----
+
+class ASPP(nn.Module):
+    def __init__(self, C_in, C_mid=256, rates=(1, 6, 12, 18), groups=32):
+        super().__init__()
+        def b(conv):
+            return nn.Sequential(conv, nn.GroupNorm(min(groups, C_mid), C_mid), nn.ReLU(True))
+        self.branches = nn.ModuleList([
+            b(nn.Conv2d(C_in, C_mid, 1, bias=False)),
+            b(nn.Conv2d(C_in, C_mid, 3, padding=rates[1], dilation=rates[1], bias=False)),
+            b(nn.Conv2d(C_in, C_mid, 3, padding=rates[2], dilation=rates[2], bias=False)),
+            b(nn.Conv2d(C_in, C_mid, 3, padding=rates[3], dilation=rates[3], bias=False)),
+        ])
+        self.img_pool = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(C_in, C_mid, 1, bias=False),
+            nn.GroupNorm(min(groups, C_mid), C_mid), nn.ReLU(True),
+        )
+        self.project = nn.Sequential(
+            nn.Conv2d(C_mid * 5, C_mid, 1, bias=False),
+            nn.GroupNorm(min(groups, C_mid), C_mid), nn.ReLU(True),
+        )
+
+    def forward(self, x):
+        H, W = x.shape[-2:]
+        feats = [b(x) for b in self.branches]
+        img = F.interpolate(self.img_pool(x), size=(H, W), mode='bilinear', align_corners=False)
+        return self.project(torch.cat(feats + [img], dim=1))
+
+
+class DeepLabV3PlusHead(nn.Module):
+    def __init__(self, C_in=256, C_mid=256, num_classes=21, groups=32):
+        super().__init__()
+        self.aspp = ASPP(C_in, C_mid, groups=groups)
+        self.decode = nn.Sequential(
+            nn.Conv2d(C_mid, C_mid, 3, padding=1, bias=False),
+            nn.GroupNorm(min(groups, C_mid), C_mid), nn.ReLU(True),
+            nn.Conv2d(C_mid, num_classes, 1),
+        )
+        nn.init.constant_(self.decode[-1].bias, -0.5)
+
+    def forward(self, x, target_size):
+        y = self.decode(self.aspp(x))
+        return F.interpolate(y, size=target_size, mode='bilinear', align_corners=True)
+
+
+def build_pascal_decoder(config, c_dino=768, dropout=0.0):
+    concepts = config['concepts'][config['training']['concept_config']]
+    base_channels = 3072
+    num_classes = config['data']['num_classes']
+    per_feature = calculate_distributed_concept_channels(len(concepts), 4)
+    concepts_per_layer = per_feature[0]
+    in_channels = base_channels + c_dino + concepts_per_layer
+
+    def path_block():
+        return nn.Sequential(
+            nn.Conv2d(in_channels, 256, 3, padding=1, bias=False),
+            nn.GroupNorm(32, 256), nn.ReLU(True), nn.Dropout2d(dropout))
+
+    decoder = nn.ModuleDict({f'path{i}': path_block() for i in range(1, 5)})
+    decoder['head'] = DeepLabV3PlusHead(C_in=256, C_mid=256, num_classes=num_classes, groups=32)
+    decoder['aux_head'] = nn.Sequential(nn.Dropout2d(dropout), nn.Conv2d(256, num_classes, 1))
+    decoder['reduce1024to256'] = nn.Sequential(
+        nn.Conv2d(1024, 256, 1, bias=False), nn.GroupNorm(32, 256), nn.ReLU(True), nn.Dropout2d(dropout))
+    return decoder
+
+
+class FluxDinoPascalModel(nn.Module):
+    """Pascal VOC segmentation model (cache_only mode)."""
+    def __init__(self, decoder, config, cache_dir, num_timesteps=4, hidden_dim=768,
+                 num_transformer_layers=3, layer_scale_init=1e-6, dino_model="dinov3_vitb16"):
+        super().__init__()
+        self.decoder = decoder
+        self.config = config
+        self.num_timesteps = num_timesteps
+        self.concepts = config['concepts'][config['training']['concept_config']]
+        self.cache = MultiTimestepFeatureCache(cache_dir)
+
+        self.hyperfeature_fusion = build_hyperfeature_fusion(
+            True, num_timesteps, hidden_dim, num_transformer_layers,
+            layer_scale_init, fusion_type="transformer", return_alpha=True)
+
+        self.dino_extractor = build_dino_extractor(dino_model, "full")
+        self.c_dino = resolve_c_dino("full", dino_model)
+        self.feature_mode = "full"
+        self.use_flux, self.use_dino = True, True
+
+    def forward(self, images, image_name, resolution, timestep_data):
+        device = next(self.parameters()).device
+        images = images.to(device)
+        height, width = resolution
+        patch_h, patch_w = height // 16, width // 16
+
+        dino_features = self.dino_extractor(images)
+
+        multi_timestep_features = {}
+        for timestep in timestep_data['timesteps']:
+            single_features = timestep_data['features'][timestep]['single_features']
+            multi_timestep_features[timestep] = [
+                f.float().permute(0, 2, 1).reshape(1, 3072, patch_h, patch_w) for f in single_features]
+
+        flux_features, alpha_layers = self.hyperfeature_fusion(multi_timestep_features)
+        del multi_timestep_features
+
+        concatenated_features = []
+        for layer_idx in range(4):
+            flux_feat = flux_features[layer_idx]
+            dino_feat = dino_features[layer_idx]
+            if flux_feat.shape[-2:] != dino_feat.shape[-2:]:
+                flux_feat = F.interpolate(flux_feat, size=dino_feat.shape[-2:], mode='bilinear', align_corners=False)
+            concatenated_features.append(torch.cat([flux_feat, dino_feat], dim=1))
+
+        last_timestep = timestep_data['timesteps'][-1]
+        concept_maps = timestep_data['concept_maps'][last_timestep]
+        distributed = distribute_concepts(concept_maps, len(concatenated_features), device)
+        del flux_features
+
+        distributed_resized = []
+        for i, dist in enumerate(distributed):
+            target_size = concatenated_features[i].shape[-2:]
+            if dist.shape[-2:] != target_size:
+                dist = F.interpolate(dist, size=target_size, mode='bilinear', align_corners=False)
+            distributed_resized.append(dist)
+        del distributed
+
+        fused_with_concepts = [
+            torch.cat([feat.to(device), distributed_resized[i].to(device)], dim=1)
+            for i, feat in enumerate(concatenated_features)]
+        del concatenated_features, distributed_resized
+
+        return self._decode(fused_with_concepts, height, width)
+
+    def _decode(self, feats, height, width):
+        path1 = self.decoder['path1'](feats[0])
+        path2 = self.decoder['path2'](feats[1])
+        path3 = self.decoder['path3'](feats[2])
+        path4 = self.decoder['path4'](feats[3])
+
+        target_h, target_w = path1.shape[-2:]
+        path2_up = F.interpolate(path2, size=(target_h, target_w), mode='bilinear', align_corners=False)
+        path3_up = F.interpolate(path3, size=(target_h, target_w), mode='bilinear', align_corners=False)
+        path4_up = F.interpolate(path4, size=(target_h, target_w), mode='bilinear', align_corners=False)
+
+        fused = self.decoder['reduce1024to256'](torch.cat([path1, path2_up, path3_up, path4_up], dim=1))
+        logits = self.decoder['head'](fused, (height, width))
+        return logits
+
+
+def build_duts_decoder(config, c_dino=768):
+    concepts = config['concepts'][config['training']['concept_config']]
+    base_channels = 3072
+    per_feature = calculate_distributed_concept_channels(len(concepts), 4)
+    in_channels = [base_channels + c_dino + c for c in per_feature]
+    return OriginalDPTSegmentationDecoder(
+        in_channels=in_channels,
+        num_classes=config['data']['num_classes'],
+        features=config['model']['decoder']['features'],
+        target_size=None,
+    )
+
+
+class FluxDinoDUTSModel(nn.Module):
+    """DUTS saliency model (cache_only mode)."""
+    def __init__(self, decoder, config, cache_dir, num_timesteps=4, hidden_dim=768,
+                 num_transformer_layers=3, layer_scale_init=1e-6, dino_model="dinov3_vitb16"):
+        super().__init__()
+        self.decoder = decoder
+        self.config = config
+        self.num_timesteps = num_timesteps
+        self.concepts = config['concepts'][config['training']['concept_config']]
+        self.cache = MultiTimestepFeatureCache(cache_dir)
+
+        self.hyperfeature_fusion = build_hyperfeature_fusion(
+            True, num_timesteps, hidden_dim, num_transformer_layers,
+            layer_scale_init, fusion_type="transformer")
+
+        self.dino_extractor = build_dino_extractor(dino_model, "full")
+        self.c_dino = resolve_c_dino("full", dino_model)
+        self.feature_mode = "full"
+        self.use_flux, self.use_dino = True, True
+
+    def forward(self, images, image_name, resolution, timestep_data):
+        device = next(self.parameters()).device
+        images = images.to(device)
+        height, width = resolution
+        patch_h, patch_w = height // 16, width // 16
+
+        for timestep in timestep_data['features']:
+            for key in timestep_data['features'][timestep]:
+                val = timestep_data['features'][timestep][key]
+                if isinstance(val, list):
+                    timestep_data['features'][timestep][key] = [
+                        f.float() if isinstance(f, torch.Tensor) else f for f in val]
+                elif isinstance(val, torch.Tensor):
+                    timestep_data['features'][timestep][key] = val.float()
+
+        dino_features = self.dino_extractor(images)
+
+        multi_timestep_features = {}
+        for timestep in timestep_data['timesteps']:
+            single_features = timestep_data['features'][timestep]['single_features']
+            multi_timestep_features[timestep] = [
+                f.permute(0, 2, 1).reshape(1, 3072, patch_h, patch_w) for f in single_features]
+
+        flux_features = self.hyperfeature_fusion(multi_timestep_features)
+        del multi_timestep_features
+
+        concatenated_features = []
+        for layer_idx in range(4):
+            flux_feat = flux_features[layer_idx]
+            dino_feat = dino_features[layer_idx]
+            if flux_feat.shape[-2:] != dino_feat.shape[-2:]:
+                flux_feat = F.interpolate(flux_feat, size=dino_feat.shape[-2:], mode='bilinear', align_corners=False)
+            concatenated_features.append(torch.cat([flux_feat, dino_feat], dim=1))
+
+        last_timestep = timestep_data['timesteps'][-1]
+        concept_maps = timestep_data['concept_maps'][last_timestep]
+        distributed = distribute_concepts(concept_maps, len(concatenated_features), device)
+        del flux_features
+
+        distributed_resized = []
+        for i, dist in enumerate(distributed):
+            target_size = concatenated_features[i].shape[-2:]
+            if dist.shape[-2:] != target_size:
+                dist = F.interpolate(dist, size=target_size, mode='bilinear', align_corners=False)
+            distributed_resized.append(dist)
+        del distributed
+
+        fused_with_concepts = [
+            torch.cat([feat.to(device), distributed_resized[i].to(device)], dim=1)
+            for i, feat in enumerate(concatenated_features)]
+        del concatenated_features, distributed_resized
+
+        logits = self.decoder(fused_with_concepts)
+        if logits.shape[-2:] != (height, width):
+            logits = F.interpolate(logits, size=(height, width), mode='bilinear', align_corners=False)
+        return logits
+
+
+def build_nyu_decoder(config, c_dino=768, high_res=False):
+    concepts = config['concepts'][config['training']['concept_config']]
+    num_concepts = len(concepts)
+    per_layer = num_concepts // 4
+    remainder = num_concepts % 4
+    in_channels = []
+    for i in range(4):
+        count = per_layer + (1 if i < remainder else 0)
+        in_channels.append(3072 + c_dino + count)
+    return DPTHeadSpatial(in_channels=in_channels, features=256, num_classes=1, use_bn=False)
+
+
+class FluxNYUDepthModel(nn.Module):
+    """NYU Depth model (cache_only mode)."""
+    def __init__(self, decoder, config, cache_dir, num_timesteps=4, hidden_dim=768,
+                 num_transformer_layers=3, layer_scale_init=1e-4, dino_model="dinov3_vitb16",
+                 high_res_decoder=False):
+        super().__init__()
+        self.decoder = decoder
+        self.config = config
+        self.num_timesteps = num_timesteps
+        self.concepts = config['concepts'][config['training']['concept_config']]
+        self.cache = MultiTimestepFeatureCache(cache_dir)
+        self.high_res_decoder = high_res_decoder
+
+        self.hyperfeature_fusion = build_hyperfeature_fusion(
+            True, num_timesteps, hidden_dim, num_transformer_layers,
+            layer_scale_init, fusion_type="transformer", return_alpha=False)
+
+        self.dino_extractor = build_dino_extractor(dino_model, "full")
+        self.c_dino = resolve_c_dino("full", dino_model)
+        self.feature_mode = "full"
+        self.use_flux, self.use_dino = True, True
+
+    def forward(self, images, image_name, resolution, timestep_data):
+        device = next(self.parameters()).device
+        images = images.to(device)
+
+        res = timestep_data.get('resolution')
+        if res is not None:
+            native_h, native_w = res
+        else:
+            native_h, native_w = timestep_data.get('native_h', 896), timestep_data.get('native_w', 1152)
+        patch_h, patch_w = native_h // 16, native_w // 16
+
+        dino_features = self.dino_extractor(images)
+
+        multi_timestep_features = {}
+        for timestep in timestep_data['timesteps']:
+            single_features = timestep_data['features'][timestep]['single_features']
+            multi_timestep_features[timestep] = [
+                f.float().permute(0, 2, 1).reshape(1, 3072, patch_h, patch_w) for f in single_features]
+
+        fused_flux_features = self.hyperfeature_fusion(multi_timestep_features)
+
+        concept_maps_avg = {}
+        for concept in self.concepts:
+            concept_stack = [timestep_data['concept_maps'][t][concept]
+                             for t in timestep_data['timesteps']
+                             if concept in timestep_data['concept_maps'][t]]
+            if concept_stack:
+                concept_maps_avg[concept] = torch.stack([
+                    c.to(device) if hasattr(c, "to") else torch.tensor(c, device=device)
+                    for c in concept_stack]).mean(dim=0).float()
+
+        target_size = dino_features[0].shape[-2:] if self.use_dino else fused_flux_features[0].shape[-2:]
+        distributed_concepts = distribute_concepts_across_layers(
+            concept_maps_avg, num_layers=4, target_size=target_size, device=device)
+
+        final_features = []
+        for layer_idx in range(4):
+            flux_feat = fused_flux_features[layer_idx]
+            concept_feat = distributed_concepts[layer_idx]
+            dino_feat = dino_features[layer_idx]
+            layer_size = dino_feat.shape[-2:]
+            if flux_feat.shape[-2:] != layer_size:
+                flux_feat = F.interpolate(flux_feat, size=layer_size, mode='bilinear', align_corners=False)
+            if concept_feat.shape[-2:] != layer_size:
+                concept_feat = F.interpolate(concept_feat, size=layer_size, mode='bilinear', align_corners=False)
+            final_features.append(torch.cat([flux_feat, dino_feat, concept_feat], dim=1))
+
+        return self.decoder(final_features)
+
+
+# ---- Checkpoint loading ----
+
+def load_checkpoint(model, checkpoint_path):
+    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    state_dict = ckpt.get("state_dict", ckpt)
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    relevant_missing = [k for k in missing if k.startswith(("hyperfeature_fusion", "decoder"))]
+    if relevant_missing:
+        print(f"[WARN] {len(relevant_missing)} fusion/decoder keys NOT found in checkpoint")
+    print(f"[CKPT] Loaded {checkpoint_path} (missing={len(missing)}, unexpected={len(unexpected)})")
+
+
+# ---- Generation with feature capture (from generate.py) ----
+
+def generate_and_capture(pipeline, prompt, concepts, height, width, steps, guidance,
+                         seed, device, concept_attention_kwargs, num_timesteps, group_size=7):
+    transformer = pipeline.transformer
+    target_steps = {steps + (-(i * group_size + 1)) for i in range(num_timesteps)
+                    if i * group_size < steps}
+    captured = {}
+
+    def _capture(pipe, step_idx, t, callback_kwargs):
+        if step_idx in target_steps:
+            tv = int(t.item()) if hasattr(t, "item") else int(t)
+            _, single_features = transformer.get_features()
+            captured[tv] = [f.detach().cpu().clone() for f in single_features]
+        return callback_kwargs
+
+    transformer.stored_features.clear()
+    with torch.no_grad():
+        result = pipeline(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+            generator=torch.Generator(device).manual_seed(seed),
+            concept_attention_kwargs=concept_attention_kwargs,
+            output_type="pil",
+            callback_on_step_end=_capture,
+        )
+
+    raw = result.concept_attention_maps
+    if not raw:
+        raise RuntimeError("Pipeline returned no concept-attention maps")
+    maps_list = raw[0] if (len(raw) == 1 and isinstance(raw[0], list)) else raw
+    if len(maps_list) != len(concepts):
+        raise ValueError(f"{len(concepts)} concepts vs {len(maps_list)} concept maps")
+    concept_maps = {c: maps_list[i] for i, c in enumerate(concepts)}
+
+    return result.images[0], captured, concept_maps
+
+
+def build_timestep_data(captured, concept_maps, concepts, height, width, prompt, stem):
+    timesteps = sorted(captured.keys(), reverse=True)
+    cmaps = {c: (m.cpu() if hasattr(m, "cpu") else m) for c, m in concept_maps.items()}
+    data = {
+        "timesteps": timesteps,
+        "features": {t: {"single_features": captured[t]} for t in timesteps},
+        "concept_maps": {t: dict(cmaps) for t in timesteps},
+        "image_name": stem,
+        "concepts": concepts,
+        "resolution": (height, width),
+        "prompt": prompt,
+        "native_h": height,
+        "native_w": width,
+    }
+    return data
+
+
+# ---- Visualization helpers ----
+
+VOC_PALETTE = None
+
+def voc_color_palette():
+    global VOC_PALETTE
+    if VOC_PALETTE is not None:
+        return VOC_PALETTE
+    palette = [0] * (256 * 3)
+    for i in range(256):
+        r = g = b = 0
+        c = i
+        for j in range(8):
+            r |= ((c >> 0) & 1) << (7 - j)
+            g |= ((c >> 1) & 1) << (7 - j)
+            b |= ((c >> 2) & 1) << (7 - j)
+            c >>= 3
+        palette[i * 3 + 0] = r
+        palette[i * 3 + 1] = g
+        palette[i * 3 + 2] = b
+    VOC_PALETTE = palette
+    return palette
+
+
+def colorize_depth(depth, cmap="magma"):
+    import matplotlib
+    d = depth.astype(np.float32)
+    lo, hi = np.percentile(d, 2), np.percentile(d, 98)
+    d = np.clip((d - lo) / (hi - lo + 1e-8), 0, 1)
+    colormap = matplotlib.colormaps[cmap]
+    rgb = (colormap(d)[:, :, :3] * 255).astype(np.uint8)
+    return rgb
+
+
+def colorize_saliency(prob):
+    prob_norm = (prob - prob.min()) / (prob.max() - prob.min() + 1e-8)
+    rgb = (plt_colormap(prob_norm)[:, :, :3] * 255).astype(np.uint8)
+    return rgb
+
+
+def plt_colormap(x, cmap_name="inferno"):
+    import matplotlib
+    colormap = matplotlib.colormaps[cmap_name]
+    return colormap(x)
+
+
+# ---- Config loading ----
+
+def make_config(task):
+    """Load the task config with env vars expanded."""
+    if task == "pascal":
+        config_path = os.path.join(os.path.dirname(__file__), "configs", "pascal_voc_config.yaml")
+    elif task == "nyu":
+        config_path = os.path.join(os.path.dirname(__file__), "configs", "nyu_depth_config.yaml")
+    else:
+        config_path = os.path.join(os.path.dirname(__file__), "configs", f"{task}_config.yaml")
+    with open(config_path, "r") as f:
+        def _expand_env(value):
+            if isinstance(value, str):
+                return os.path.expanduser(os.path.expandvars(value))
+            if isinstance(value, dict):
+                return {k: _expand_env(v) for k, v in value.items()}
+            if isinstance(value, list):
+                return [_expand_env(v) for v in value]
+            return value
+        config = _expand_env(yaml.safe_load(f))
+    # Set dummy paths since we use a temp dir
+    config['paths'] = config.get('paths', {})
+    config['paths']['permanent_cache_dir'] = '/tmp/mmdiff_cache'
+    # Fix data_root if it's an unexpanded env var path
+    if 'data' in config and 'data_root' in config['data']:
+        if config['data']['data_root'].startswith('$') or '/path/' in config['data']['data_root']:
+            config['data']['data_root'] = '/tmp/dummy_data'
+    return config
+
+
+# ---- Global model loading at module scope ----
+
+print("[SETUP] Loading configs...")
+configs = {
+    'duts': make_config('duts'),
+    'pascal': make_config('pascal_voc'),
+    'nyu': make_config('nyu_depth'),
+}
+
+print("[SETUP] Loading FLUX.1-dev pipeline with concept attention...")
+flux_model = "black-forest-labs/FLUX.1-dev"
+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("cuda")
+pipeline.set_progress_bar_config(disable=True)
+print("[SETUP] FLUX pipeline loaded.")
+
+# Build decoder models
+device = "cuda"
+cache_dir = "/tmp/mmdiff_cache"
+os.makedirs(cache_dir, exist_ok=True)
+
+# Common architecture params (from configs)
+num_timesteps = 4
+hidden_dim = 768
+num_transformer_layers = 3
+layer_scale_init = 1e-6
+dino_model_name = "dinov3_vitb16"
+c_dino = resolve_c_dino("full", dino_model_name)
+
+# Build DUTS saliency model
+print("[SETUP] Building DUTS saliency model...")
+duts_decoder = build_duts_decoder(configs['duts'], c_dino=c_dino)
+duts_model = FluxDinoDUTSModel(
+    duts_decoder, configs['duts'], cache_dir,
+    num_timesteps=num_timesteps, hidden_dim=hidden_dim,
+    num_transformer_layers=num_transformer_layers, layer_scale_init=layer_scale_init,
+    dino_model=dino_model_name)
+duts_ckpt = "/tmp/checkpoints/duts_saliency.ckpt"
+load_checkpoint(duts_model, duts_ckpt)
+duts_model = duts_model.to(device).eval()
+
+# Build Pascal VOC model
+print("[SETUP] Building Pascal VOC segmentation model...")
+pascal_layer_scale_init = 1e-6
+pascal_decoder = build_pascal_decoder(configs['pascal'], c_dino=c_dino, dropout=0.0)
+pascal_model = FluxDinoPascalModel(
+    pascal_decoder, configs['pascal'], cache_dir,
+    num_timesteps=num_timesteps, hidden_dim=hidden_dim,
+    num_transformer_layers=num_transformer_layers, layer_scale_init=pascal_layer_scale_init,
+    dino_model=dino_model_name)
+pascal_ckpt = "/tmp/checkpoints/pascal_segmentation.ckpt"
+load_checkpoint(pascal_model, pascal_ckpt)
+pascal_model = pascal_model.to(device).eval()
+
+# Build NYU Depth model
+print("[SETUP] Building NYU Depth model...")
+nyu_layer_scale_init = 1e-4
+nyu_decoder = build_nyu_decoder(configs['nyu'], c_dino=c_dino, high_res=False)
+nyu_model = FluxNYUDepthModel(
+    nyu_decoder, configs['nyu'], cache_dir,
+    num_timesteps=num_timesteps, hidden_dim=hidden_dim,
+    num_transformer_layers=num_transformer_layers, layer_scale_init=nyu_layer_scale_init,
+    dino_model=dino_model_name, high_res_decoder=False)
+nyu_ckpt = "/tmp/checkpoints/nyu_depth.ckpt"
+load_checkpoint(nyu_model, nyu_ckpt)
+nyu_model = nyu_model.to(device).eval()
+
+print("[SETUP] All models loaded successfully!")
+
+
+# ---- Inference function ----
+
+@spaces.GPU(duration=120)
+def generate(prompt, task_choice, seed, num_steps, guidance_scale):
+    """Generate an image and dense prediction(s) from a text prompt."""
+    height, width = 512, 512
+
+    # Select config and concepts based on task
+    if task_choice == "Saliency (DUTS)":
+        task = "duts"
+        config = configs['duts']
+    elif task_choice == "Segmentation (Pascal VOC)":
+        task = "pascal"
+        config = configs['pascal']
+    elif task_choice == "Depth (NYU)":
+        task = "nyu"
+        config = configs['nyu']
+    else:  # All
+        task = "all"
+        config = configs['duts']  # Use DUTS concepts for generation
+
+    concepts = config['concepts'][config['training']['concept_config']]
+    concept_attention_kwargs = {
+        "concepts": concepts,
+        "timesteps": config["flux"]["concept_timesteps"],
+        "layers": config["flux"]["concept_layers"],
+    }
+
+    # Generate image and capture features
+    with tempfile.TemporaryDirectory(prefix="mmdiff_demo_") as tmp_cache:
+        stem = "demo_image"
+
+        if task == "all":
+            # For "all" mode, we generate once and use each task's own concepts for decoding
+            # First generate with DUTS concepts for the image
+            pil_image, captured, concept_maps = generate_and_capture(
+                pipeline, prompt, concepts, height, width,
+                num_steps, guidance_scale, seed, device,
+                concept_attention_kwargs, num_timesteps)
+
+            # Build timestep data for DUTS
+            timestep_data = build_timestep_data(
+                captured, concept_maps, concepts, height, width, prompt, stem)
+
+            image_tensor = transforms.ToTensor()(pil_image).unsqueeze(0).to(device)
+
+            results = {"image": pil_image}
+
+            # DUTS saliency
+            duts_concepts = configs['duts']['concepts'][configs['duts']['training']['concept_config']]
+            duts_cakw = {
+                "concepts": duts_concepts,
+                "timesteps": configs['duts']["flux"]["concept_timesteps"],
+                "layers": configs['duts']["flux"]["concept_layers"],
+            }
+            # Re-capture with DUTS concepts for proper saliency
+            _, duts_captured, duts_concept_maps = generate_and_capture(
+                pipeline, prompt, duts_concepts, height, width,
+                num_steps, guidance_scale, seed, device,
+                duts_cakw, num_timesteps)
+            duts_td = build_timestep_data(duts_captured, duts_concept_maps, duts_concepts, height, width, prompt, stem)
+            with torch.no_grad():
+                logits = duts_model(image_tensor, stem, (height, width), duts_td)
+                prob = torch.sigmoid(logits.squeeze(1))[0].float().cpu().numpy()
+            sal_vis = colorize_saliency(prob)
+            results["saliancy"] = Image.fromarray(sal_vis)
+
+            # Pascal segmentation
+            pascal_concepts = configs['pascal']['concepts'][configs['pascal']['training']['concept_config']]
+            pascal_cakw = {
+                "concepts": pascal_concepts,
+                "timesteps": configs['pascal']["flux"]["concept_timesteps"],
+                "layers": configs['pascal']["flux"]["concept_layers"],
+            }
+            _, pascal_captured, pascal_concept_maps = generate_and_capture(
+                pipeline, prompt, pascal_concepts, height, width,
+                num_steps, guidance_scale, seed, device,
+                pascal_cakw, num_timesteps)
+            pascal_td = build_timestep_data(pascal_captured, pascal_concept_maps, pascal_concepts, height, width, prompt, stem)
+            with torch.no_grad():
+                logits = pascal_model(image_tensor, stem, (height, width), pascal_td)
+                pred = torch.argmax(logits, dim=1)[0].byte().cpu().numpy()
+            seg_img = Image.fromarray(pred, mode="P")
+            seg_img.putpalette(voc_color_palette())
+            seg_rgb = seg_img.convert("RGB")
+            results["segmentation"] = seg_rgb
+
+            # NYU depth
+            nyu_concepts = configs['nyu']['concepts'][configs['nyu']['training']['concept_config']]
+            nyu_cakw = {
+                "concepts": nyu_concepts,
+                "timesteps": configs['nyu']["flux"]["concept_timesteps"],
+                "layers": configs['nyu']["flux"]["concept_layers"],
+            }
+            _, nyu_captured, nyu_concept_maps = generate_and_capture(
+                pipeline, prompt, nyu_concepts, height, width,
+                num_steps, guidance_scale, seed, device,
+                nyu_cakw, num_timesteps)
+            nyu_td = build_timestep_data(nyu_captured, nyu_concept_maps, nyu_concepts, height, width, prompt, stem)
+            with torch.no_grad():
+                depth = nyu_model(image_tensor, stem, (height, width), nyu_td)
+                depth = F.softplus(depth).squeeze().float().cpu().numpy()
+            depth_vis = colorize_depth(depth)
+            results["depth"] = Image.fromarray(depth_vis)
+
+            return results
+        else:
+            # Single task
+            pil_image, captured, concept_maps = generate_and_capture(
+                pipeline, prompt, concepts, height, width,
+                num_steps, guidance_scale, seed, device,
+                concept_attention_kwargs, num_timesteps)
+
+            timestep_data = build_timestep_data(
+                captured, concept_maps, concepts, height, width, prompt, stem)
+
+            image_tensor = transforms.ToTensor()(pil_image).unsqueeze(0).to(device)
+
+            if task == "duts":
+                with torch.no_grad():
+                    logits = duts_model(image_tensor, stem, (height, width), timestep_data)
+                    prob = torch.sigmoid(logits.squeeze(1))[0].float().cpu().numpy()
+                sal_vis = colorize_saliency(prob)
+                return {"image": pil_image, "saliancy": Image.fromarray(sal_vis)}
+
+            elif task == "pascal":
+                with torch.no_grad():
+                    logits = pascal_model(image_tensor, stem, (height, width), timestep_data)
+                    pred = torch.argmax(logits, dim=1)[0].byte().cpu().numpy()
+                seg_img = Image.fromarray(pred, mode="P")
+                seg_img.putpalette(voc_color_palette())
+                seg_rgb = seg_img.convert("RGB")
+                return {"image": pil_image, "segmentation": seg_rgb}
+
+            elif task == "nyu":
+                with torch.no_grad():
+                    depth = nyu_model(image_tensor, stem, (height, width), timestep_data)
+                    depth = F.softplus(depth).squeeze().float().cpu().numpy()
+                depth_vis = colorize_depth(depth)
+                return {"image": pil_image, "depth": Image.fromarray(depth_vis)}
+
+
+# ---- Gradio UI ----
+
+import gradio as gr
+
+DESCRIPTION = """# MMDiff: Extending Diffusion Transformers for Multi-Modal Generation
+
+Generate an image from a text prompt using FLUX.1-dev while simultaneously producing
+dense predictions (saliency maps, segmentation maps, depth maps) from the frozen
+diffusion transformer's intermediate features via lightweight trained decoder heads.
+
+**Paper**: [MMDiff: Extending Diffusion Transformers for Multi-Modal Generation](https://huggingface.co/papers/2606.16673)
+**Model**: [yagmurakarken/mmdiff](https://huggingface.co/yagmurakarken/mmdiff)
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus()) as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt_input = gr.Textbox(
+                label="Text Prompt",
+                placeholder="A cat sitting on a wooden table...",
+                value="A cat sitting on a wooden table next to a window",
+                lines=2,
+            )
+            task_select = gr.Radio(
+                choices=["Saliency (DUTS)", "Segmentation (Pascal VOC)", "Depth (NYU)", "All (Saliency + Segmentation + Depth)"],
+                label="Task",
+                value="Saliency (DUTS)",
+            )
+            generate_btn = gr.Button("Generate", variant="primary", size="lg")
+
+            with gr.Accordion("Advanced Options", open=False):
+                seed_input = gr.Slider(0, 1000, value=0, step=1, label="Seed")
+                steps_input = gr.Slider(4, 50, value=28, step=1, label="Inference Steps")
+                guidance_input = gr.Slider(1.0, 10.0, value=3.5, step=0.5, label="Guidance Scale")
+
+        with gr.Column(scale=2):
+            # Output gallery - dynamically shown based on task
+            with gr.Row():
+                image_output = gr.Image(label="Generated Image", type="pil", height=300)
+            with gr.Row():
+                saliency_output = gr.Image(label="Saliency Map", type="pil", height=300, visible=False)
+                segmentation_output = gr.Image(label="Segmentation Map", type="pil", height=300, visible=False)
+                depth_output = gr.Image(label="Depth Map", type="pil", height=300, visible=False)
+
+    # Examples
+    gr.Examples(
+        examples=[
+            ["A cat sitting on a wooden table next to a window", "Saliency (DUTS)", 0, 28, 3.5],
+            ["A person riding a bicycle on a city street", "Segmentation (Pascal VOC)", 42, 28, 3.5],
+            ["A modern living room with a sofa and coffee table", "Depth (NYU)", 0, 28, 3.5],
+            ["A dog playing in a grassy park", "All (Saliency + Segmentation + Depth)", 0, 28, 3.5],
+        ],
+        inputs=[prompt_input, task_select, seed_input, steps_input, guidance_input],
+        outputs=[image_output, saliency_output, segmentation_output, depth_output],
+        fn=generate,
+        cache_examples=False,
+        run_on_click=True,
+    )
+
+    def update_visibility(task_choice):
+        """Show/hide output columns based on task."""
+        if "All" in task_choice:
+            return {
+                saliency_output: gr.update(visible=True),
+                segmentation_output: gr.update(visible=True),
+                depth_output: gr.update(visible=True),
+            }
+        elif "Saliency" in task_choice:
+            return {
+                saliency_output: gr.update(visible=True),
+                segmentation_output: gr.update(visible=False),
+                depth_output: gr.update(visible=False),
+            }
+        elif "Segmentation" in task_choice:
+            return {
+                saliency_output: gr.update(visible=False),
+                segmentation_output: gr.update(visible=True),
+                depth_output: gr.update(visible=False),
+            }
+        elif "Depth" in task_choice:
+            return {
+                saliency_output: gr.update(visible=False),
+                segmentation_output: gr.update(visible=False),
+                depth_output: gr.update(visible=True),
+            }
+        return {}
+
+    task_select.change(
+        fn=update_visibility,
+        inputs=[task_select],
+        outputs=[saliency_output, segmentation_output, depth_output],
+    )
+
+    def run_and_route(prompt, task_choice, seed, steps, guidance):
+        """Run generation and route outputs to the right components."""
+        results = generate(prompt, task_choice, seed, steps, guidance)
+        # Return None for hidden outputs
+        img = results.get("image")
+        sal = results.get("saliancy")
+        seg = results.get("segmentation")
+        dep = results.get("depth")
+        return img, sal, seg, dep
+
+    generate_btn.click(
+        fn=run_and_route,
+        inputs=[prompt_input, task_select, seed_input, steps_input, guidance_input],
+        outputs=[image_output, saliency_output, segmentation_output, depth_output],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

configs/duts_config.yaml

@@ -0,0 +1,107 @@
+# DUTS Training Configuration - Native Resolution (Variable Resolution with FluxResizer)
+# Uses FLUX.1-dev and FluxResizer for optimal native resolutions (no padding)
+
+# Experiment settings
+experiment:
+  name: "duts_native_resolution"
+  description: "DUTS with FLUX.1-dev at native resolutions using FluxResizer"
+
+# Concept configurations (same as original)
+concepts:
+  basic: ["object", "background", "detail", "edges"]
+  expanded: ["background", "object", "edges", "salient", "contour"]
+  meta: ["living", "vehicle", "furniture", "object", "background"]
+  combined: ["object", "background", "living", "vehicle", "furniture", "detail", "edges"]
+  combined_expanded: ["object", "background", "living", "vehicle", "furniture", "detail", "edges", "salient", "contour"]
+  with_classes: ["object", "background", "detail", "edges", "airplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "table", "dog", "horse", "motorbike", "person", "plant", "sheep", "sofa", "train", "television"]
+  all_comprehensive: ["living", "vehicle", "furniture", "object", "background", "detail", "edges", "airplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "table", "dog", "horse", "motorbike", "person", "plant", "sheep", "sofa", "train", "television"]
+
+# Training settings
+training:
+  concept_config: "expanded"  # Which concept configuration to use
+  epochs: 100  # Match baseline training duration (was 20)
+  learning_rate: 1e-5  # Reduced from 5e-5 for stability (gradient explosion fix)
+  batch_size: 1  # REQUIRED: Must be 1 for variable resolution
+  accumulate_grad_batches: 4  # Effective batch size = 4
+  gradient_clip_val: 0.1  # Very aggressive clipping for gradient explosion (was 0.5)
+  precision: "32-true"  # FP32 for maximum stability (gradient explosion with FP16)
+  
+# Data settings
+data:
+  dataset: "duts"
+  num_classes: 1
+  # Root of the DUTS dataset. Expected layout:
+  #   <data_root>/DUTS-TR/DUTS-TR-Image, <data_root>/DUTS-TR/DUTS-TR-Mask
+  #   <data_root>/DUTS-TE/DUTS-TE-Image, <data_root>/DUTS-TE/DUTS-TE-Mask
+  data_root: "${DUTS_ROOT}"          # env var; e.g. export DUTS_ROOT=/datasets/.../DUTS
+  # NOTE: No target_size! FluxResizer selects optimal resolution per image
+  
+# Model settings
+model:
+  flux_model: "black-forest-labs/FLUX.1-dev"  # Changed from schnell to dev
+  dtype: "float32"  # Match training precision (was float16)
+  dino_model: "dinov3_vitb16"  # DINOv3 base model
+  feature_locations:
+    transformer_blocks: [4, 9, 13, 18]
+    single_transformer_blocks: [4, 15, 26, 37]
+  decoder:
+    features: 256
+  hyperfeature_fusion:
+    num_timesteps: 4
+    fusion_type: "transformer"
+    hidden_dim: 768
+    num_transformer_layers: 3
+    layer_scale_init: 1e-6
+    
+# FLUX settings
+flux:
+  timesteps: 28
+  guidance_scale: 3.5
+  num_inference_steps: 1
+  concept_timesteps: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]
+  concept_layers: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
+
+# FluxResizer settings
+flux_resizer:
+  # Optimal resolutions (all divisible by 32 for FLUX/DINO compatibility)
+  # FluxResizer will automatically select the closest resolution based on aspect ratio
+  optimal_resolutions:
+    - [1024, 1024]  # 1:1
+    - [896, 1152]   # ~0.78:1
+    - [1152, 896]   # ~1.29:1
+    - [768, 1344]   # ~0.57:1
+    - [1344, 768]   # ~1.75:1
+    - [832, 1216]   # ~0.68:1
+    - [1216, 832]   # ~1.46:1
+    - [704, 1408]   # 0.5:1
+    - [1408, 704]   # 2:1
+    - [960, 1088]   # ~0.88:1
+    - [1088, 960]   # ~1.13:1
+
+# Hardware settings
+hardware:
+  devices: 1  # Use 2 GPUs for faster training
+  #strategy: "ddp_find_unused_parameters_true"  # Multi-GPU strategy
+  num_sanity_val_steps: 1
+  detect_anomaly: true
+  accelerator: "gpu"
+
+# Paths (timestamps will be automatically added)
+paths:
+  # All outputs live under ${MMDIFF_OUTPUT} (set it to a writable dir, e.g. /work/<user>/mmdiff_out).
+  cache_base_dir: "${MMDIFF_OUTPUT}/cache"
+  log_base_dir: "${MMDIFF_OUTPUT}/logs"
+  checkpoint_base_dir: "${MMDIFF_OUTPUT}/checkpoints"
+  use_timestamp: true  # Add timestamp to folder names
+  permanent_cache_dir: "${MMDIFF_OUTPUT}/cache/duts_flux_cache"  # Native resolution feature cache
+  
+# Logging
+logging:
+  log_every_n_steps: 10
+  save_top_k: 3
+  monitor: "val_loss"
+  mode: "min"
+
+# HuggingFace
+huggingface:
+  token: ""  # Leave empty and authenticate via `huggingface-cli login` or the HF_TOKEN env var

configs/nyu_depth_config.yaml

@@ -0,0 +1,92 @@
+# NYU Depth Training Configuration - Native Resolution (Variable Resolution with FluxResizer)
+# Uses FLUX.1-dev and FluxResizer for optimal native resolutions (no padding)
+
+# Experiment settings
+experiment:
+  name: "nyu_depth_native_resolution"
+  description: "NYU Depth with FLUX.1-dev at native resolutions using FluxResizer"
+
+# Concept configurations (same as original)
+concepts:
+  basic: ["depth", "surface", "object", "background"]
+  basic_2: ["depth", "surface", "object", "near", "far"]
+  indoor: ["wall", "floor", "furniture", "object", "depth"]
+  spatial: ["near", "far", "depth", "distance", "surface"]
+  detailed: ["wall", "floor", "ceiling", "furniture", "object", "depth", "surface"]
+  comprehensive: ["wall", "floor", "ceiling", "furniture", "object", "depth", "near", "far", "surface", "indoor", "outdoor", "structure"]
+  semantic: ["bedroom", "kitchen", "bathroom", "living_room", "depth", "furniture", "wall", "floor"]
+  geometric: ["depth", "surface", "edge", "plane", "corner", "boundary", "gradient", "distance"]
+
+# Training settings
+training:
+  concept_config: "basic_2"  # Which concept configuration to use
+  epochs: 20
+  learning_rate: 1e-4
+  batch_size: 1  # REQUIRED: Must be 1 for variable resolution
+  accumulate_grad_batches: 4  # Effective batch size = 4
+  gradient_clip_val: 1.0
+  precision: "16-mixed"
+  
+# Data settings
+data:
+  dataset: "nyu_depth_v2"
+  # NYU Depth V2 standard train/test split. Point the trainer at it via CLI args
+  # (--data_path / --filenames_path); see shell_scripts/train_nyu.sh.
+  num_classes: 1  # Depth is regression, but num_classes used for decoder
+  min_depth: 0.1  # Minimum valid depth (meters)
+  max_depth: 10.0  # Maximum valid depth (meters)
+  # NOTE: No target_size! FluxResizer selects optimal resolution per image
+  
+# Model settings
+model:
+  flux_model: "black-forest-labs/FLUX.1-dev"  # Changed from schnell to dev
+  dtype: "float16"
+  dino_model: "dinov3_vitb16"  # DINOv3 base model
+  feature_locations:
+    transformer_blocks: [4, 9, 13, 18]
+    single_transformer_blocks: [4, 15, 26, 37]
+  decoder:
+    features: 256
+  hyperfeature_fusion:
+    num_timesteps: 4
+    fusion_type: "transformer"
+    hidden_dim: 768
+    num_transformer_layers: 3
+    layer_scale_init: 1e-4
+    
+# FLUX settings
+flux:
+  timesteps: 28
+  guidance_scale: 3.5
+  num_inference_steps: 1
+  concept_timesteps: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]
+  concept_layers: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
+
+# Hardware settings
+hardware:
+  devices: 1  # Number of GPUs (DDP for multi-GPU)
+  #strategy: "ddp_find_unused_parameters_true"  # Required for DDP with conditional logic
+  num_sanity_val_steps: 1
+  detect_anomaly: true
+  accelerator: "gpu"
+  num_workers: 8
+
+# Paths (timestamps will be automatically added)
+paths:
+  # All outputs live under ${MMDIFF_OUTPUT} (set it to a writable dir, e.g. /work/<user>/mmdiff_out).
+  cache_base_dir: "${MMDIFF_OUTPUT}/cache"
+  log_base_dir: "${MMDIFF_OUTPUT}/logs"
+  checkpoint_base_dir: "${MMDIFF_OUTPUT}/checkpoints"
+  use_timestamp: true
+  permanent_cache_dir: "${MMDIFF_OUTPUT}/cache/nyu_native_feature_cache"  # Native resolution feature cache
+  
+# Logging
+logging:
+  log_every_n_steps: 10
+  save_top_k: 3
+  monitor: "val_loss"
+  mode: "min"
+
+# HuggingFace
+huggingface:
+  token: ""  # Leave empty and authenticate via `huggingface-cli login` or the HF_TOKEN env var

configs/pascal_voc_config.yaml

@@ -0,0 +1,89 @@
+# Pascal VOC Training Configuration - Native Resolution (Variable Resolution with FluxResizer)
+# Uses FLUX.1-dev and FluxResizer for optimal native resolutions (no padding)
+
+# Experiment settings
+experiment:
+  name: "pascal_voc_native_resolution"
+  description: "Pascal VOC with FLUX.1-dev at native resolutions using FluxResizer"
+
+# Concept configurations (same as original)
+concepts:
+  basic: ["object", "background", "detail", "edges"]
+  meta: ["living", "vehicle", "furniture", "object", "background"]
+  meta_expanded: ["living", "vehicle", "furniture", "object", "background", "person", "car", "airplane", "table", "plant", "bird", "bicycle"]
+  combined: ["object", "background", "living", "vehicle", "furniture", "detail", "edges"]
+  with_classes: ["object", "background", "detail", "edges", "airplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "table", "dog", "horse", "motorbike", "person", "plant", "sheep", "sofa", "train", "television"]
+  all_comprehensive: ["living", "vehicle", "furniture", "object", "background", "detail", "edges", "airplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "table", "dog", "horse", "motorbike", "person", "plant", "sheep", "sofa", "train", "television"]
+
+# Training settings
+training:
+  concept_config: "basic"  # Which concept configuration to use
+  epochs: 100  # Match baseline (was 20)
+  learning_rate: 1e-4  # Match baseline (capped at 3e-5 in optimizer code)
+  batch_size: 1  # REQUIRED: Must be 1 for variable resolution
+  accumulate_grad_batches: 4  # Effective batch size = 4
+  gradient_clip_val: 0.5
+  precision: "16-mixed"
+  
+# Data settings
+data:
+  dataset: "pascal_voc"
+  data_root: "${VOC_ROOT}"                       # dir holding JPEGImages/
+  train_split_file: "${VOC_SPLITS}/train_aug.txt"
+  val_split_file: "${VOC_SPLITS}/val.txt"
+  mask_root: "${VOC_MASKS}"                      # dir holding SegmentationClassAug/ (masks)
+  num_classes: 21
+  # NOTE: No target_size! FluxResizer selects optimal resolution per image
+  
+# Model settings
+model:
+  flux_model: "black-forest-labs/FLUX.1-dev"  # Changed from schnell to dev
+  dtype: "float16"
+  dino_model: "dinov3_vitb16"  # DINOv3 base model
+  feature_locations:
+    transformer_blocks: [4, 9, 13, 18]
+    single_transformer_blocks: [4, 15, 26, 37]
+  decoder:
+    features: 256
+  hyperfeature_fusion:
+    num_timesteps: 4
+    fusion_type: "transformer"
+    hidden_dim: 768
+    num_transformer_layers: 3
+    layer_scale_init: 1e-6
+    
+# FLUX settings
+flux:
+  timesteps: 28
+  guidance_scale: 3.5
+  num_inference_steps: 1
+  concept_timesteps: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]
+  concept_layers: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
+
+# Hardware settings
+hardware:
+  devices: 1  # Use 2 GPUs for faster training
+  #strategy: "ddp_find_unused_parameters_true"  # Multi-GPU strategy
+  num_sanity_val_steps: 1
+  detect_anomaly: true
+  accelerator: "gpu"
+
+# Paths (timestamps will be automatically added)
+paths:
+  # All outputs live under ${MMDIFF_OUTPUT} (set it to a writable dir, e.g. /work/<user>/mmdiff_out).
+  cache_base_dir: "${MMDIFF_OUTPUT}/cache"
+  log_base_dir: "${MMDIFF_OUTPUT}/logs"
+  checkpoint_base_dir: "${MMDIFF_OUTPUT}/checkpoints"
+  use_timestamp: true
+  permanent_cache_dir: "${MMDIFF_OUTPUT}/cache/pascal_feature_cache"  # Native resolution feature cache
+  
+# Logging
+logging:
+  log_every_n_steps: 10
+  save_top_k: 3
+  monitor: "val_loss"
+  mode: "min"
+
+# HuggingFace
+huggingface:
+  token: ""  # Leave empty and authenticate via `huggingface-cli login` or the HF_TOKEN env var

core/__init__.py

@@ -0,0 +1,39 @@
+"""Shared building blocks for the mmdiff training / extraction / inference scripts."""
+
+from .flux_features import (
+    load_flux_pipeline,
+    MultiTimestepFeatureCache,
+    MultiTimestepFeatureExtractor,
+    distribute_concepts,
+    distribute_concepts_across_layers,
+)
+from .training_utils import (
+    EMA,
+    load_config,
+    setup_paths,
+    calculate_distributed_concept_channels,
+    feature_mode_flags,
+    resolve_c_dino,
+    build_dino_extractor,
+    build_hyperfeature_fusion,
+    add_shared_training_args,
+    add_model_args,
+)
+
+__all__ = [
+    "load_flux_pipeline",
+    "MultiTimestepFeatureCache",
+    "MultiTimestepFeatureExtractor",
+    "distribute_concepts",
+    "distribute_concepts_across_layers",
+    "EMA",
+    "load_config",
+    "setup_paths",
+    "calculate_distributed_concept_channels",
+    "feature_mode_flags",
+    "resolve_c_dino",
+    "build_dino_extractor",
+    "build_hyperfeature_fusion",
+    "add_shared_training_args",
+    "add_model_args",
+]

core/flux_features.py

@@ -0,0 +1,287 @@
+"""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

core/training_utils.py

@@ -0,0 +1,181 @@
+"""Shared trainer utilities: EMA, config/paths, DINO/fusion builders and argparse helpers."""
+
+import os
+from datetime import datetime
+
+import torch
+import yaml
+
+from models.hyperfeature_fusion import create_hyperfeature_fusion
+from models.dinov3_hf_extractor import create_dinov3_hf_extractor
+
+
+DINO_REPO_MAP = {
+    "dinov3_vits16": "facebook/dinov3-vits16-pretrain-lvd1689m",
+    "dinov3_vitb16": "facebook/dinov3-vitb16-pretrain-lvd1689m",
+    "dinov3_vitl16": "facebook/dinov3-vitl16-pretrain-lvd1689m",
+    "dinov3_vitg16": "facebook/dinov3-vitg16-pretrain-lvd1689m",
+}
+
+
+class EMA:
+    """Exponential Moving Average of trainable, floating-point model parameters."""
+
+    def __init__(self, model, beta=0.999):
+        self.beta = beta
+        self.shadow = {}
+        self.backup = {}
+        for name, param in model.named_parameters():
+            if param.requires_grad and param.dtype.is_floating_point:
+                self.shadow[name] = param.data.detach().clone()
+
+    def update(self, model):
+        with torch.no_grad():
+            for name, param in model.named_parameters():
+                if name in self.shadow:
+                    if self.shadow[name].device != param.device:
+                        self.shadow[name] = self.shadow[name].to(param.device)
+                    self.shadow[name].mul_(self.beta).add_(param.data.detach(), alpha=1.0 - self.beta)
+
+    def apply_shadow(self, model):
+        self.backup = {}
+        with torch.no_grad():
+            for name, param in model.named_parameters():
+                if name in self.shadow:
+                    if self.shadow[name].device != param.device:
+                        self.shadow[name] = self.shadow[name].to(param.device)
+                    self.backup[name] = param.data.detach().clone()
+                    param.data.copy_(self.shadow[name])
+        return self.backup
+
+    def restore_backup(self, model):
+        with torch.no_grad():
+            for name, param in model.named_parameters():
+                if name in self.backup:
+                    param.data.copy_(self.backup[name])
+        self.backup = {}
+
+    def restore(self, model, backup=None):
+        b = backup if backup is not None else self.backup
+        with torch.no_grad():
+            for name, param in model.named_parameters():
+                if name in b:
+                    if b[name].device != param.device:
+                        b[name] = b[name].to(param.device)
+                    param.data.copy_(b[name])
+        self.backup = {}
+
+
+def _expand_env(value):
+    """Recursively expand ${VAR}/$VAR and ~ in all string values of a config."""
+    if isinstance(value, str):
+        return os.path.expanduser(os.path.expandvars(value))
+    if isinstance(value, dict):
+        return {k: _expand_env(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_expand_env(v) for v in value]
+    return value
+
+
+def load_config(config_path: str):
+    with open(config_path, "r") as f:
+        return _expand_env(yaml.safe_load(f))
+
+
+def setup_paths(config, task: str):
+    """Timestamped log/checkpoint dirs plus the (permanent) feature cache dir."""
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    cache_dir = config["paths"].get("permanent_cache_dir") or f"{config['paths']['cache_base_dir']}/{task}_cache"
+    return {
+        "cache_dir": cache_dir,
+        "log_dir": f"{config['paths']['log_base_dir']}/{task}_logs_{timestamp}",
+        "checkpoint_dir": f"{config['paths']['checkpoint_base_dir']}/{task}_checkpoints_{timestamp}",
+    }
+
+
+def calculate_distributed_concept_channels(num_concepts, num_features=4):
+    """Per-layer concept channel counts when distributing concepts across layers."""
+    per = num_concepts // num_features
+    rem = num_concepts % num_features
+    return [per + (1 if i < rem else 0) for i in range(num_features)]
+
+
+def feature_mode_flags(feature_mode):
+    """Return ``(use_flux, use_dino)`` for an ablation mode."""
+    assert feature_mode in ("full", "flux_only", "dino_only"), feature_mode
+    return feature_mode in ("full", "flux_only"), feature_mode in ("full", "dino_only")
+
+
+def resolve_c_dino(feature_mode, dino_model):
+    """DINO feature width for the given backbone (0 when DINO is disabled)."""
+    if feature_mode == "flux_only":
+        return 0
+    if "vits" in dino_model:
+        return 384
+    if "vitl" in dino_model:
+        return 1024
+    if "vitg" in dino_model:
+        return 1536
+    return 768  # vitb
+
+
+def build_dino_extractor(dino_model, feature_mode, layer_indices=(2, 5, 8, 11)):
+    """Build the DINOv3 extractor (HF when known, torch.hub fallback otherwise)."""
+    trainable = feature_mode == "dino_only"
+    repo_id = DINO_REPO_MAP.get(dino_model)
+    if repo_id:
+        return create_dinov3_hf_extractor(
+            repo_id=repo_id, take_indices=list(layer_indices), trainable=trainable
+        )
+    from models.dino_fusion import create_dino_extractor
+    return create_dino_extractor(model_name=dino_model, take_indices=list(layer_indices))
+
+
+def build_hyperfeature_fusion(use_flux, num_timesteps, hidden_dim, num_transformer_layers,
+                              layer_scale_init, fusion_type="transformer", return_alpha=False):
+    """Build the multi-timestep hyperfeature fusion module, or ``None`` without FLUX."""
+    if not use_flux:
+        return None
+    return create_hyperfeature_fusion(
+        num_timesteps=num_timesteps,
+        num_layers=4,
+        fusion_type=fusion_type,
+        hidden_dim=hidden_dim,
+        num_transformer_layers=num_transformer_layers,
+        layer_scale_init=layer_scale_init,
+        return_alpha=return_alpha,
+    )
+
+
+def add_shared_training_args(parser, *, config_default, hidden_dim, num_transformer_layers,
+                             layer_scale_init, include_cache_args=True):
+    """Register the backbone/ablation arguments common to all task trainers."""
+    parser.add_argument("--config", type=str, default=config_default)
+    parser.add_argument("--num_timesteps", type=int, default=4)
+    parser.add_argument("--hidden_dim", type=int, default=hidden_dim)
+    parser.add_argument("--num_transformer_layers", type=int, default=num_transformer_layers)
+    parser.add_argument("--layer_scale_init", type=float, default=layer_scale_init)
+    parser.add_argument("--dino_model", type=str, default="dinov3_vitb16")
+    parser.add_argument("--feature_mode", type=str, default="full",
+                        choices=["full", "flux_only", "dino_only"],
+                        help="Ablation: 'full' (FLUX+concepts+DINO), 'flux_only' (FLUX+concepts), "
+                             "'dino_only' (trainable DINO only)")
+    parser.add_argument("--fast_dev_run", action="store_true",
+                        help="Smoke test: run a single train+val batch then exit (verifies the "
+                             "training pipeline starts end-to-end)")
+    if include_cache_args:
+        parser.add_argument("--limit_images", type=int, default=None, help="Limit dataset size for testing")
+        parser.add_argument("--cache_only", action="store_true",
+                            help="Load features from cache only (no FLUX pipeline)")
+
+
+def add_model_args(parser):
+    """Architecture flags for inference/generation; must match the checkpoint."""
+    parser.add_argument("--dino_model", default="dinov3_vitb16")
+    parser.add_argument("--num_timesteps", type=int, default=4)
+    parser.add_argument("--hidden_dim", type=int, default=768)
+    parser.add_argument("--num_transformer_layers", type=int, default=3)
+    parser.add_argument("--layer_scale_init", type=float, default=1e-6)
+    parser.add_argument("--fusion_type", default="transformer", choices=["attention", "transformer"])
+    parser.add_argument("--high_res_decoder", action="store_true", help="(nyu) high-res depth decoder")
+    parser.add_argument("--device", default="cuda")

download_checkpoints.py

@@ -0,0 +1,28 @@
+"""Download model checkpoints from yagmurakarken/mmdiff at Space startup."""
+import os
+from huggingface_hub import hf_hub_download
+
+CHECKPOINT_DIR = "/tmp/checkpoints"
+MODEL_REPO = "yagmurakarken/mmdiff"
+CHECKPOINT_FILES = ["duts_saliency.ckpt", "nyu_depth.ckpt", "pascal_segmentation.ckpt"]
+
+
+def download_all():
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+    for fname in CHECKPOINT_FILES:
+        path = os.path.join(CHECKPOINT_DIR, fname)
+        if not os.path.exists(path):
+            print(f"[DOWNLOAD] {fname} ...")
+            downloaded = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename=fname,
+                repo_type="model",
+                local_dir=CHECKPOINT_DIR,
+            )
+            print(f"[DOWNLOAD] {fname} -> {downloaded}")
+        else:
+            print(f"[DOWNLOAD] {fname} already exists at {path}")
+
+
+if __name__ == "__main__":
+    download_all()

flux_concept_attention/__init__.py

@@ -0,0 +1,4 @@
+from .flux_with_concept_attention_pipeline import FluxWithConceptAttentionPipeline
+from .flux_dit_with_concept_attention import FluxTransformer2DModelWithConceptAttention
+
+__all__ = ["FluxWithConceptAttentionPipeline", "FluxTransformer2DModelWithConceptAttention"]

flux_concept_attention/flux_dit_block_with_concept_attention.py

@@ -0,0 +1,342 @@
+import torch
+import torch.nn.functional as F
+from typing import Any, Dict, Optional, Tuple
+from torch import nn
+import einops
+
+from diffusers.models.transformers.transformer_flux import FluxTransformerBlock
+from diffusers.models.attention import Attention
+from diffusers.models.embeddings import apply_rotary_emb
+
+
+def _trim_rope(rope, target_len: int):
+    """
+    Ensure rotary embedding length matches the attention sequence length.
+    Works for Diffusers' rope tuples (cos, sin) or a single tensor.
+    
+    Rope format: [seq_len, embed_dim], so dimension 0 is the sequence.
+    
+    Args:
+        rope: Either a tuple (cos, sin) or a single tensor, or None
+        target_len: Target sequence length
+    
+    Returns:
+        Trimmed rope in the same format as input (takes last target_len positions)
+    """
+    if rope is None:
+        return None
+    
+    # Rope is a (cos, sin) tuple (Diffusers format)
+    # Each tensor is [seq_len, embed_dim]
+    if isinstance(rope, tuple):
+        cos, sin = rope
+        if cos.shape[0] != target_len:
+            # Trim dimension 0 (sequence length)
+            cos = cos[-target_len:, ...]
+            sin = sin[-target_len:, ...]
+        return (cos, sin)
+    
+    # Rope is a single tensor
+    if rope.shape[0] != target_len:
+        rope = rope[-target_len:, ...]
+    return rope
+
+
+class FluxConceptAttentionProcessor:
+    """
+    Custom attention processor for FLUX that implements concept attention.
+    Exactly matches original FLUX attention pattern while adding concept observation stream.
+    """
+    
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("FluxConceptAttentionProcessor requires PyTorch 2.0")
+    
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        concept_hidden_states: Optional[torch.Tensor] = None,
+        concept_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        q_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        kv_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        # Concept-specific arguments are now explicit parameters (required for inspect.signature)
+        
+        batch_size, _, _ = encoder_hidden_states.shape
+        
+        # *** MAIN GENERATION STREAM *** - Exact FLUX Pattern
+        
+        # 1. `sample` projections (image hidden states)
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        image_query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        image_key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        image_value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            image_query = attn.norm_q(image_query)
+        if attn.norm_k is not None:
+            image_key = attn.norm_k(image_key)
+
+        # 2. `context` projections (text encoder hidden states) - FLUX specific
+        encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+        encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+        encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+        encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+            batch_size, -1, attn.heads, head_dim
+        ).transpose(1, 2)
+        encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+            batch_size, -1, attn.heads, head_dim
+        ).transpose(1, 2)
+        encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+            batch_size, -1, attn.heads, head_dim
+        ).transpose(1, 2)
+
+        if attn.norm_added_q is not None:
+            encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+        if attn.norm_added_k is not None:
+            encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+        # 3. FLUX concatenation order: [encoder, image]
+        query = torch.cat([encoder_hidden_states_query_proj, image_query], dim=2)
+        key = torch.cat([encoder_hidden_states_key_proj, image_key], dim=2)
+        value = torch.cat([encoder_hidden_states_value_proj, image_value], dim=2)
+
+        # 4. Apply rotary embeddings to FULL concatenated tensors (FLUX way)
+        # Use explicit q/kv ropes; fallback to image_rotary_emb
+        rope_q = q_rotary_emb if q_rotary_emb is not None else image_rotary_emb
+        rope_kv = kv_rotary_emb if kv_rotary_emb is not None else image_rotary_emb
+        
+        if rope_q is not None:
+            # query shape after transpose: [B, heads, seq, head_dim]
+            q_len = query.shape[2]  # sequence dimension is at index 2
+            rope_q = _trim_rope(rope_q, q_len)
+            query = apply_rotary_emb(query, rope_q, sequence_dim=2)
+        
+        if rope_kv is not None:
+            k_len = key.shape[2]  # sequence dimension is at index 2
+            rope_kv = _trim_rope(rope_kv, k_len)
+            key = apply_rotary_emb(key, rope_kv, sequence_dim=2)
+
+        # 5. Main text+image attention (FLUX standard)
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # *** CONCEPT STREAM *** - Exact same logic as text-image (following reference)
+        concept_hidden_states_output = None
+        concept_attention_maps = None
+        
+        if concept_hidden_states is not None:
+            # Use TEXT projections for concepts (like reference: txt_attn.qkv)
+            concept_query = attn.add_q_proj(concept_hidden_states)   # Same as text!
+            concept_key = attn.add_k_proj(concept_hidden_states)     # Same as text!
+            concept_value = attn.add_v_proj(concept_hidden_states)   # Same as text!
+            
+            concept_query = concept_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            concept_key = concept_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            concept_value = concept_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            
+            if attn.norm_added_q is not None:
+                concept_query = attn.norm_added_q(concept_query)    # Text normalization
+            if attn.norm_added_k is not None:
+                concept_key = attn.norm_added_k(concept_key)        # Text normalization
+
+            concept_image_q = torch.cat([concept_query, image_query], dim=2)
+            concept_image_k = torch.cat([concept_key, image_key], dim=2)
+            concept_image_v = torch.cat([concept_value, image_value], dim=2)
+            
+            # Apply concept rotary embeddings to FULL concatenated tensor (like reference)
+            if concept_rotary_emb is not None:
+                # concept_image_q shape after transpose: [B, heads, seq, head_dim]
+                cq_len = concept_image_q.shape[2]  # sequence dimension is at index 2
+                ck_len = concept_image_k.shape[2]
+                
+                rope_cq = _trim_rope(concept_rotary_emb, cq_len)
+                rope_ck = _trim_rope(concept_rotary_emb, ck_len)
+                
+                concept_image_q = apply_rotary_emb(concept_image_q, rope_cq, sequence_dim=2)
+                concept_image_k = apply_rotary_emb(concept_image_k, rope_ck, sequence_dim=2)
+            
+            # Do the joint attention operation (like reference)
+            concept_image_attn = F.scaled_dot_product_attention(
+                concept_image_q,
+                concept_image_k, 
+                concept_image_v,
+                dropout_p=0.0,
+                is_causal=False
+            )
+            
+            # Separate the concept attention (like reference: concept_attn = concept_image_attn[:, :, :concepts.shape[1]])
+            concept_attn = concept_image_attn[:, :, :concept_hidden_states.size(1)]
+            concept_hidden_states_output = concept_attn.transpose(1, 2).reshape(
+                batch_size, -1, attn.heads * head_dim
+            )
+            
+            # Compute attention maps from concept and image queries (before attention)
+            # Save vectors for postprocessing (like reference implementation)
+            concept_attention_maps = {
+                'concept_vectors': concept_query,  # (batch, heads, concepts, dim)
+                'image_vectors': image_query       # (batch, heads, patches, dim)
+            }
+
+        # 6. FLUX output processing
+        encoder_hidden_states_out, hidden_states_out = (
+            hidden_states[:, : encoder_hidden_states.shape[1]],
+            hidden_states[:, encoder_hidden_states.shape[1] :],
+        )
+
+        # linear proj
+        hidden_states_out = attn.to_out[0](hidden_states_out)
+        # dropout
+        hidden_states_out = attn.to_out[1](hidden_states_out)
+
+        # FLUX specific: separate output projection for encoder
+        encoder_hidden_states_out = attn.to_add_out(encoder_hidden_states_out)
+        
+        # Process concept outputs with same projections
+        if concept_hidden_states_output is not None:
+            concept_hidden_states_output = attn.to_out[0](concept_hidden_states_output)
+            concept_hidden_states_output = attn.to_out[1](concept_hidden_states_output)
+            
+            # Return vectors for postprocessing (like reference implementation)
+            concept_attention_maps = {
+                'concept_vectors': concept_hidden_states_output,  # Final processed concept features
+                'image_vectors': hidden_states_out               # Final processed image features
+            }
+        
+        return hidden_states_out, encoder_hidden_states_out, concept_hidden_states_output, concept_attention_maps
+
+
+class FluxTransformerBlockWithConceptAttention(FluxTransformerBlock):
+    """
+    Simplified FLUX transformer block with concept attention.
+    Uses the elegant CogVideoX approach with custom attention processor.
+    """
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.attn.processor = FluxConceptAttentionProcessor()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        concept_hidden_states: Optional[torch.Tensor],
+        temb: torch.Tensor,
+        concept_temb: Optional[torch.Tensor],
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        concept_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        concept_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+            hidden_states, 
+            emb=temb
+        )
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            encoder_hidden_states, 
+            emb=temb
+        )
+        
+        # Concept normalization (use same temb as main stream if concept_temb is None)
+        norm_concept_hidden_states = None
+        concept_gate_msa = concept_shift_mlp = concept_scale_mlp = concept_gate_mlp = None
+        concept_gate_ff = None
+        
+        if concept_hidden_states is not None:
+            effective_concept_temb = concept_temb if concept_temb is not None else temb
+            norm_concept_hidden_states, concept_gate_msa, concept_shift_mlp, concept_scale_mlp, concept_gate_mlp = self.norm1_context(
+                concept_hidden_states, 
+                emb=effective_concept_temb
+            )
+        
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        
+        # Pass concept-specific args through joint_attention_kwargs
+        # (they're not accepted as direct args by FluxAttention.forward)
+        if norm_concept_hidden_states is not None:
+            joint_attention_kwargs['concept_hidden_states'] = norm_concept_hidden_states
+            joint_attention_kwargs['concept_rotary_emb'] = concept_rotary_emb
+        
+        # Attention with concept attention (using our custom processor)
+        attention_outputs = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        if len(attention_outputs) == 4:
+            attn_output, context_attn_output, concept_attn_output, concept_attention_maps = attention_outputs
+            ip_attn_output = None
+        elif len(attention_outputs) == 5:
+            attn_output, context_attn_output, concept_attn_output, concept_attention_maps, ip_attn_output = attention_outputs
+        else:
+            # Fallback for when no concept attention
+            attn_output, context_attn_output = attention_outputs[:2]
+            concept_attn_output = None
+            concept_attention_maps = None
+            ip_attn_output = attention_outputs[2] if len(attention_outputs) > 2 else None
+
+        ################## Process Concept Features FIRST (like CogVideoX) ##################
+        if concept_attn_output is not None and concept_hidden_states is not None:
+            # Apply concept attention gate and residual
+            concept_attn_output = concept_gate_msa.unsqueeze(1) * concept_attn_output
+            concept_hidden_states = concept_hidden_states + concept_attn_output
+            
+            # Concept feedforward processing (norm2_context is regular LayerNorm, not adaptive)
+            norm_concept_hidden_states = self.norm2_context(concept_hidden_states)
+            norm_concept_hidden_states = norm_concept_hidden_states * (1 + concept_scale_mlp[:, None]) + concept_shift_mlp[:, None]
+            concept_ff_output = self.ff_context(norm_concept_hidden_states)
+            concept_hidden_states = concept_hidden_states + concept_gate_mlp.unsqueeze(1) * concept_ff_output
+            
+            if concept_hidden_states.dtype == torch.float16:
+                concept_hidden_states = concept_hidden_states.clip(-65504, 65504)
+
+        ################## Now Process Main Generation Stream ##################
+        # Standard FLUX processing for image features
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+        hidden_states = hidden_states + ff_output
+        
+        if ip_attn_output is not None:
+            hidden_states = hidden_states + ip_attn_output
+
+        # Standard FLUX processing for text features
+        if context_attn_output is not None:
+            context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+            encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+            norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+            norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+            context_ff_output = self.ff_context(norm_encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+            
+            if encoder_hidden_states.dtype == torch.float16:
+                encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+        
+        return encoder_hidden_states, hidden_states, concept_hidden_states, concept_attention_maps
+    

flux_concept_attention/flux_dit_with_concept_attention.py

@@ -0,0 +1,400 @@
+from typing import List, Dict
+
+import torch
+import numpy as np
+from typing import Any, Dict, Optional, Tuple, Union
+from torch import nn
+from torch import Tensor
+
+from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
+from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock
+from diffusers.models.normalization import AdaLayerNormContinuous
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers, BaseOutput
+from diffusers.utils.import_utils import is_torch_npu_available
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
+
+from .flux_dit_block_with_concept_attention import FluxTransformerBlockWithConceptAttention
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class FluxTransformer2DOutputWithConceptAttention(BaseOutput):
+    sample: torch.Tensor
+    concept_attention_maps: torch.Tensor
+
+
+class FluxTransformer2DModelWithConceptAttention(FluxTransformer2DModel):
+    """
+        The Transformer model introduced in Flux with Concept Attention.
+    """
+
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        out_channels: Optional[int] = None,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        guidance_embeds: bool = True,
+        axes_dims_rope: Tuple[int] = (16, 56, 56),
+        feature_locations: Optional[Dict[str, List[int]]] = None,
+    ):
+        super().__init__(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            num_layers=num_layers,
+            num_single_layers=num_single_layers,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+            joint_attention_dim=joint_attention_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            guidance_embeds=guidance_embeds,
+            axes_dims_rope=axes_dims_rope,
+        )
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if self.config.guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
+        self.x_embedder = nn.Linear(self.config.in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlockWithConceptAttention(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+        self.stored_features: Dict[str, Tensor] = {}
+        self.feature_locations = feature_locations or {
+            "transformer_blocks": [4, 9, 13, 18],
+            "single_transformer_blocks": [4, 16, 27, 36],
+        }
+        self._register_feature_hooks()
+    
+    def get_features(self) -> Tuple[List[Tensor], List[Tensor]]:
+        """
+        Get the stored feature maps as raw tokens for downstream reshaping.
+
+        For dual stream transformer blocks: Returns the second item in the tuple (image tokens)
+          Shape: [B, H*W, C] where H*W is the actual spatial size
+
+        For single stream transformer blocks: Extracts image tokens from full sequence
+          Shape: [B, H*W, C] where H*W is the actual spatial size (excluding text tokens)
+
+        Returns:
+            Tuple containing:
+                - List of transformer block features as tokens [B, H*W, C]
+                - List of single transformer block features as tokens [B, H*W, C]
+        """
+        transformer_features = []
+        single_transformer_features = []
+
+        # Debug flag to print shapes on first call
+        for name, feature_output in self.stored_features.items():
+            if "single_transformer_blocks" in name:
+                # Single blocks return (encoder_hidden_states, hidden_states) tuple
+                if isinstance(feature_output, tuple) and len(feature_output) >= 2:
+                    image_feature = feature_output[1]  # [B, H*W, C] - image tokens only
+                    single_transformer_features.append(image_feature)
+                else:
+                    # Fallback if not a tuple (shouldn't happen)
+                    single_transformer_features.append(feature_output)
+            elif "transformer_blocks" in name:
+                if isinstance(feature_output, tuple) and len(feature_output) >= 2:
+                    image_feature = feature_output[1]  # [B, H*W, C]
+                    transformer_features.append(image_feature)
+        
+        return (transformer_features, single_transformer_features)
+    
+    def _get_hook(self, name: str):
+        """
+        Create a forward hook function for feature extraction.
+
+        Args:
+            name: Identifier for the layer where the hook will be attached
+
+        Returns:
+            Callable hook function that stores the layer's output tensor
+        """
+
+        def hook(
+                module: nn.Module, input: Union[Tensor, Tuple[Tensor, ...]], output: Tensor
+        ) -> None:
+            self.stored_features[name] = output
+
+        return hook
+
+    def _register_feature_hooks(self) -> None:
+        """
+        Register forward hooks on the specified layers to capture their outputs.
+
+        Attaches hooks based on the feature_locations configuration:
+        - transformer_blocks: Main transformer blocks (indexed from 0 to 18)
+        - single_transformer_blocks: Single transformer blocks (indexed from 0 to 37)
+        """
+        for block_type, indices in self.feature_locations.items():
+            if block_type == "transformer_blocks":
+                for idx in indices:
+                    if 0 <= idx < len(self.transformer_blocks):
+                        self.transformer_blocks[idx].register_forward_hook(
+                            self._get_hook(f"{block_type}_{idx}")
+                        )
+            elif block_type == "single_transformer_blocks":
+                for idx in indices:
+                    if 0 <= idx < len(self.single_transformer_blocks):
+                        self.single_transformer_blocks[idx].register_forward_hook(
+                            self._get_hook(f"{block_type}_{idx}")
+                        )
+
+    @torch.no_grad()
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        concept_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        pooled_concept_embeds: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        concept_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        concept_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+    ) -> Union[torch.Tensor, FluxTransformer2DOutputWithConceptAttention]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            concept_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary with parameters for Concept Attention.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        concept_temb = None
+        if pooled_concept_embeds is not None:
+            if guidance is None:
+                concept_temb = self.time_text_embed(timestep, pooled_concept_embeds)
+            else:
+                concept_temb = self.time_text_embed(timestep, guidance, pooled_concept_embeds)
+
+        # Apply the context embedder to the concept_hidden_states
+        if concept_hidden_states is not None:
+            concept_hidden_states = self.context_embedder(concept_hidden_states)
+
+        if txt_ids.ndim == 3:
+            logger.warning(
+                "Passing `txt_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            logger.warning(
+                "Passing `img_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            img_ids = img_ids[0]
+
+        # Build rotary embeddings for different attention patterns:
+        
+        # 1. Image-only rotary (for single blocks)
+        image_rotary_emb = self.pos_embed(img_ids)
+        
+        # 2. Joint rotary for text+image (for dual blocks' vanilla attention)
+        #    Dual blocks concatenate encoder+image queries → need joint rope
+        ids_joint = torch.cat((txt_ids, img_ids), dim=0)  # [512 + 1024, 2]
+        rope_joint = self.pos_embed(ids_joint)             # (cos, sin) with len=1536
+        
+        # 3. Concept rotary for concept attention (concept + image sequence)
+        concept_image_ids = torch.cat((concept_ids, img_ids), dim=0)
+        concept_rotary_emb = self.pos_embed(concept_image_ids)
+
+        if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+            ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+            ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+            joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+        # Initialize concept attention processing (collect raw maps only - no processing!)
+        all_concept_attention_maps = []
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                raise NotImplementedError("Gradient checkpointing is not implemented for concept attention.")
+            else:
+                # Prepare kwargs for dual block: pass joint rope via kwargs
+                # Remove concept_* from joint_attention_kwargs to avoid warnings
+                block_joint_kwargs = dict(joint_attention_kwargs or {})
+                block_joint_kwargs.pop("concept_hidden_states", None)
+                block_joint_kwargs.pop("concept_rotary_emb", None)
+                
+                # Pass joint rope for vanilla attention (encoder+image queries)
+                block_joint_kwargs["q_rotary_emb"] = rope_joint
+                block_joint_kwargs["kv_rotary_emb"] = rope_joint
+                
+                block_output = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    concept_hidden_states=concept_hidden_states,
+                    temb=temb,
+                    concept_temb=concept_temb,
+                    image_rotary_emb=None,  # Use q/kv_rotary_emb from kwargs instead
+                    concept_rotary_emb=concept_rotary_emb,
+                    joint_attention_kwargs=block_joint_kwargs,
+                    concept_attention_kwargs=concept_attention_kwargs,
+                )
+                
+                encoder_hidden_states, hidden_states, concept_hidden_states, current_concept_attention_maps = block_output
+                
+                # Collect raw attention maps only (no processing here!)
+                if (current_concept_attention_maps is not None and 
+                    concept_attention_kwargs is not None and 
+                    index_block in concept_attention_kwargs["layers"]):
+                    all_concept_attention_maps.append(current_concept_attention_maps)
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        if concept_hidden_states is not None:
+            concept_hidden_states = concept_hidden_states.cpu()
+        
+        # Single-stream blocks: pass encoder separately
+        # IMPORTANT: Single blocks also concatenate encoder+hidden internally,
+        # so they need the joint rope (1536 = 512 text + 1024 image), not image-only!
+        
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    rope_joint,  # Use joint rope, not image_rotary_emb!
+                )
+            else:
+                # Single blocks return (encoder_hidden_states, hidden_states) tuple
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=rope_joint,  # Use joint rope, not image_rotary_emb!
+                )
+            # controlnet residual (no slicing needed - output is image-only)
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states = hidden_states + controlnet_single_block_samples[index_block // interval_control]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        # Process collected attention maps (pass through dictionaries to pipeline)
+        concept_attention_maps = None
+        if all_concept_attention_maps:
+            # Return the collected dictionaries as-is for pipeline postprocessing
+            concept_attention_maps = all_concept_attention_maps
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output, concept_attention_maps)
+        
+        return FluxTransformer2DOutputWithConceptAttention(sample=output, concept_attention_maps=concept_attention_maps)

flux_concept_attention/flux_with_concept_attention_pipeline.py

@@ -0,0 +1,1461 @@
+"""
+    Here we make various wrapper classes for the FluxPipeline from diffusers
+    to add the concept attention functionality. 
+
+    We opt for a wrapper functionality
+"""
+import inspect
+
+import torch
+import numpy as np
+from typing import List, Union, Optional, Dict, Any, Callable
+import PIL.Image
+import einops
+import matplotlib.pyplot as plt
+
+from diffusers import DiffusionPipeline
+from diffusers.image_processor import PipelineImageInput
+from diffusers.pipelines.flux.pipeline_flux import retrieve_timesteps, calculate_shift
+from diffusers.utils import is_torch_xla_available, BaseOutput, logging, USE_PEFT_BACKEND, \
+     scale_lora_layers, unscale_lora_layers
+
+from diffusers.utils.torch_utils import randn_tensor
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.models.transformers import FluxTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def retrieve_timesteps(
+        scheduler,
+        num_inference_steps: Optional[int] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        timesteps: Optional[List[int]] = None,
+        sigmas: Optional[List[float]] = None,
+        **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class FluxConceptAttentionOutput(BaseOutput):
+    """
+    Output class for the FluxPipeline with concept attention functionality.
+
+    Args:
+        images (`List[PIL.Image.Image]` or `np.ndarray`)
+            The generated images.
+        concept_attention_maps (`List[PIL.Image.Image]` or `np.ndarray`)
+            The concept attention maps.
+    """
+    images: Union[List[PIL.Image.Image], np.ndarray]
+    concept_attention_maps: Union[List[PIL.Image.Image], np.ndarray]
+
+class FluxWithConceptAttentionPipeline(
+    DiffusionPipeline,
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+    FluxIPAdapterMixin,
+):
+    r"""
+    The Flux pipeline for text-to-image generation with added Concept Attention. 
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae"
+    _optional_components = ["image_encoder", "feature_extractor"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_2: T5TokenizerFast,
+        transformer: FluxTransformer2DModel,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels)) if hasattr(self, "vae") and self.vae is not None else 16
+        )
+        self.latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16
+        self.image_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor * 2, vae_latent_channels=self.latent_channels
+        )
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = 64
+        self.latent_width = self.default_sample_size
+        self.latent_height = self.default_sample_size
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 512,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer_2)
+
+        text_inputs = self.tokenizer_2(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+
+        dtype = self.text_encoder_2.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+    ):
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds.pooler_output
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        max_sequence_length: int = 512,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            # We only use the pooled prompt output from the CLIPTextModel
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype
+        text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+
+        return prompt_embeds, pooled_prompt_embeds, text_ids
+    
+    def encode_concepts(self, concepts: List[str], device: Optional[torch.device] = None):
+        """
+            Encodes our concept vectors using the T5 Encoder.  
+        """
+        """
+            # Utils for concept encoding
+            def embed_concepts(
+                clip,
+                t5,
+                concepts: list[str],
+                batch_size=1
+            ):
+                # Code pulled from concept_attention.flux/sampling.py: prepare()
+                # Embed each concept separately
+                concept_embeddings = []
+                for concept in concepts:
+                    concept_embedding = t5(concept)
+                    # Pull out the first token
+                    token_embedding = concept_embedding[0, 0, :] # First token of first prompt
+                    concept_embeddings.append(token_embedding)
+                concept_embeddings = torch.stack(concept_embeddings).unsqueeze(0)
+                # Add filler tokens of zeros
+                concept_ids = torch.zeros(batch_size, concept_embeddings.shape[1], 3)
+
+                # Embed the concepts to a clip vector
+                prompt = " ".join(concepts)
+                vec = clip(prompt)
+                vec = torch.zeros_like(vec).to(vec.device)
+
+                return concept_embeddings, concept_ids, vec
+        """
+
+        concept_embeds = self._get_t5_prompt_embeds(
+            prompt=concepts,
+            num_images_per_prompt=1,
+            max_sequence_length=64,
+            device=device,
+        )
+        # Pull out the first token of each embedded concept to get the concept embeddings
+        concept_embeds = concept_embeds[:, 0, :]
+        concept_embeds = concept_embeds.unsqueeze(0)
+        # Make the CLIP vector for the concepts 
+        clip_vec = self._get_clip_prompt_embeds(
+            prompt=" ".join(concepts),
+            num_images_per_prompt=1,
+            device=device,
+        )   
+        # # Set the vec to zero
+        # clip_vec = torch.zeros_like(clip_vec).to(clip_vec.device)
+        # # Add filler tokens of zeros
+        concept_ids = torch.zeros(concept_embeds.shape[1], 3).to(device=device, dtype=concept_embeds.dtype)
+
+        return concept_embeds, clip_vec, concept_ids
+
+    def encode_image(self, image, device, num_images_per_prompt):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if not isinstance(image, torch.Tensor):
+            image = self.feature_extractor(image, return_tensors="pt").pixel_values
+
+        image = image.to(device=device, dtype=dtype)
+        image_embeds = self.image_encoder(image).image_embeds
+        image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+        return image_embeds
+
+    def prepare_ip_adapter_image_embeds(
+        self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt
+    ):
+        image_embeds = []
+        if ip_adapter_image_embeds is None:
+            if not isinstance(ip_adapter_image, list):
+                ip_adapter_image = [ip_adapter_image]
+
+            if len(ip_adapter_image) != len(self.transformer.encoder_hid_proj.image_projection_layers):
+                raise ValueError(
+                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.transformer.encoder_hid_proj.image_projection_layers)} IP Adapters."
+                )
+
+            for single_ip_adapter_image, image_proj_layer in zip(
+                ip_adapter_image, self.transformer.encoder_hid_proj.image_projection_layers
+            ):
+                single_image_embeds = self.encode_image(single_ip_adapter_image, device, 1)
+
+                image_embeds.append(single_image_embeds[None, :])
+        else:
+            for single_image_embeds in ip_adapter_image_embeds:
+                image_embeds.append(single_image_embeds)
+
+        ip_adapter_image_embeds = []
+        for i, single_image_embeds in enumerate(image_embeds):
+            single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
+            single_image_embeds = single_image_embeds.to(device=device)
+            ip_adapter_image_embeds.append(single_image_embeds)
+
+        return ip_adapter_image_embeds
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        height,
+        width,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+    @staticmethod
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        height = height // vae_scale_factor
+        width = width // vae_scale_factor
+
+        latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height * 2, width * 2)
+
+        return latents
+    
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
+        if isinstance(generator, list):
+            image_latents = [
+                retrieve_latents(self.vae.encode(image[i: i + 1]), generator=generator[i])
+                for i in range(image.shape[0])
+            ]
+            image_latents = torch.cat(image_latents, dim=0)
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), generator=generator)
+
+        image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+        return image_latents
+
+    def prepare_latents(
+        self,
+        image,
+        timestep,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        height = 2 * (int(height) // self.vae_scale_factor)
+        width = 2 * (int(width) // self.vae_scale_factor)
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is not None:
+            latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype)
+            return latents.to(device=device, dtype=dtype), latent_image_ids
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+        
+        if image is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+            latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype)
+
+            return latents, latent_image_ids
+
+        shape = (batch_size, num_channels_latents, height, width)
+        latent_image_ids = self._prepare_latent_image_ids(batch_size, height, width, device, dtype)
+        image = image.to(device=device, dtype=dtype)
+        if image.shape[1] != self.latent_channels:
+            image_latents = self._encode_vae_image(image=image, generator=generator)
+        else:
+            image_latents = image
+        if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+            # expand init_latents for batch_size
+            additional_image_per_prompt = batch_size // image_latents.shape[0]
+            image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+        elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            image_latents = torch.cat([image_latents], dim=0)
+
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self.scheduler.scale_noise(image_latents, timestep, noise)
+        latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+        return latents, latent_image_ids
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        true_cfg_scale: float = 1.0,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        image: PipelineImageInput = None,
+        timesteps: List[int] = None,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_ip_adapter_image: Optional[PipelineImageInput] = None,
+        negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        concept_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
+                not greater than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            true_cfg_scale (`float`, *optional*, defaults to 1.0):
+                When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`, *optional*):
+                Input image for img2img generation. If provided, the pipeline will perform image-to-image generation
+                instead of text-to-image generation.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_ip_adapter_image:
+                (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+        # Verify the concept kwargs inputs
+        if concept_attention_kwargs is not None:
+            assert "concepts" in concept_attention_kwargs, "Concepts must be passed in the concept_attention_kwargs"
+            assert isinstance(concept_attention_kwargs["concepts"], list), "Concepts must be a list of strings"
+            assert len(concept_attention_kwargs["concepts"]) > 0, "Concepts must not be an empty list"
+            assert "timesteps" in concept_attention_kwargs, "Timesteps must be passed in the concept_attention_kwargs"
+            assert isinstance(concept_attention_kwargs["timesteps"], list), "Timesteps must be a list of integers"
+            assert len(concept_attention_kwargs["timesteps"]) > 0, "Timesteps must not be an empty list"
+            assert "layers" in concept_attention_kwargs, "Layers must be passed in the concept_attention_kwargs"
+            assert isinstance(concept_attention_kwargs["layers"], list), "Layers must be a list of integers"
+            assert len(concept_attention_kwargs["layers"]) > 0, "Layers must not be an empty list"
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        init_image = None
+        if image is not None:
+            width, height = image.size
+            init_image = self.image_processor.preprocess(image, height=height, width=width)
+            init_image = init_image.to(dtype=torch.float32)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        if do_true_cfg:
+            (
+                negative_prompt_embeds,
+                negative_pooled_prompt_embeds,
+                _,
+            ) = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_2=negative_prompt_2,
+                prompt_embeds=negative_prompt_embeds,
+                pooled_prompt_embeds=negative_pooled_prompt_embeds,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                lora_scale=lora_scale,
+            )
+
+        # Embed concepts
+        concept_embeddings, pooled_concept_embeds, concept_ids = self.encode_concepts(
+            concept_attention_kwargs["concepts"], 
+            device=device
+        )
+        # Add the concept embeddings to the concept_attention_kwargs
+        # if concept_attention_kwargs is not None:
+        #     concept_attention_kwargs["concept_embeddings"] = concept_embeddings
+        #     concept_attention_kwargs["concept_vec"] = concept_vec
+
+        # 4. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if timesteps is None else None
+        image_seq_len = (int(height) // self.vae_scale_factor // 2) * (int(width) // self.vae_scale_factor // 2)
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+        latents, latent_image_ids = self.prepare_latents(
+            init_image,
+            latent_timestep,
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
+            negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
+        ):
+            negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+        elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
+            negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
+        ):
+            ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        image_embeds = None
+        negative_image_embeds = None
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+        if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
+            negative_image_embeds = self.prepare_ip_adapter_image_embeds(
+                negative_ip_adapter_image,
+                negative_ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+
+        # Make concept attention maps
+        all_concept_attention_maps = []
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+                if image_embeds is not None:
+                    self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+                # Don't do concept attention if the timestep is not in the concept_attention_kwargs
+                if concept_attention_kwargs is not None and i not in concept_attention_kwargs["timesteps"]:
+                    current_concept_embeddings = None
+                elif concept_attention_kwargs is None:
+                    # Use concept embeddings for all timesteps if no specific config
+                    current_concept_embeddings = concept_embeddings
+                else:
+                    # Use concept embeddings for specified timesteps
+                    current_concept_embeddings = concept_embeddings
+
+
+                transformer_output = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    pooled_concept_embeds=pooled_concept_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    concept_hidden_states=current_concept_embeddings,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    concept_ids=concept_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    concept_attention_kwargs=concept_attention_kwargs,
+                    return_dict=False,
+                )
+                noise_pred, current_concept_attention_maps = transformer_output
+                
+                # Process attention maps immediately with softmax (critical!)
+                if concept_attention_kwargs is not None and i in concept_attention_kwargs["timesteps"] and current_concept_attention_maps is not None:
+                    if isinstance(current_concept_attention_maps, list):
+                        # Transformer now returns list of dictionaries (one per layer)
+                        for layer_dict in current_concept_attention_maps:
+                            if isinstance(layer_dict, dict):
+                                all_concept_attention_maps.append(layer_dict)
+                    elif isinstance(current_concept_attention_maps, dict):
+                        # Collect vector dictionaries for proper postprocessing (like reference)
+                        all_concept_attention_maps.append(current_concept_attention_maps)
+
+                if do_true_cfg:
+                    if negative_image_embeds is not None:
+                        self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latents,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        pooled_projections=negative_pooled_prompt_embeds,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_image_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        elif output_type == "visualization":
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = image.detach()
+            # Force PIL for visualization
+            image = self.image_processor.postprocess(image, output_type="pil")
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = image.detach()
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+
+        # Process final attention maps
+        if all_concept_attention_maps:
+            # Extract and stack vectors across timesteps (like reference)
+            concept_vectors = []
+            image_vectors = []
+            
+            for timestep_data in all_concept_attention_maps:
+                concept_vectors.append(timestep_data['concept_vectors'].detach().cpu())
+                image_vectors.append(timestep_data['image_vectors'].detach().cpu())
+            
+            # Stack across timesteps: (timesteps, batch, heads, tokens, dim)
+            concept_vectors = torch.stack(concept_vectors, dim=0)
+            image_vectors = torch.stack(image_vectors, dim=0)
+            
+            concept_vectors = concept_vectors / (concept_vectors.norm(dim=-1, keepdim=True) + 1e-8)
+            
+            concept_attention_maps = einops.einsum(
+                image_vectors,
+                concept_vectors,
+                "time batch patches dim, time batch concepts dim -> time batch concepts patches"
+            )
+            
+            concept_attention_maps = torch.softmax(concept_attention_maps, dim=-2)
+            
+            concept_attention_maps = concept_attention_maps.mean(dim=0)
+            
+            # Convert to numpy for further processing
+            concept_attention_maps = concept_attention_maps.float().numpy()
+            
+            # Reshape to spatial format (64x64 for FLUX)
+            concept_attention_maps = einops.rearrange(
+                concept_attention_maps,
+                "batch concepts (h w) -> batch concepts h w",
+                h=height // 16,
+                w=width // 16,
+            )
+
+            # Normalize concept maps - always return raw normalized maps (not colored PIL images)
+            # Global min-max normalization per batch for consistency
+            processed = []
+            for b in range(concept_attention_maps.shape[0]):
+                maps = concept_attention_maps[b]  # (concepts, H, W)
+                vmin, vmax = maps.min(), maps.max()
+                
+                if vmax > vmin:
+                    maps = (maps - vmin) / (vmax - vmin)
+                else:
+                    maps = np.zeros_like(maps)
+                
+                # If visualization PIL images are needed, create them separately
+                if output_type == "visualization":  # New output type for colored visualizations
+                    batch_imgs = []
+                    for i, m in enumerate(maps):
+                        colored = plt.get_cmap("plasma")(m)
+                        rgb = (colored[:, :, :3] * 255).astype(np.uint8)
+                        batch_imgs.append(PIL.Image.fromarray(rgb))
+                    processed.append(batch_imgs)
+                else:
+                    # Return raw normalized arrays as list of numpy arrays per concept
+                    processed.append([maps[i] for i in range(maps.shape[0])])
+            
+            concept_attention_maps = processed
+        else:
+            concept_attention_maps = []
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image, concept_attention_maps)
+
+        return FluxConceptAttentionOutput(
+            images=image,
+            concept_attention_maps=concept_attention_maps,
+        )
+        
+
+    def encode_image(
+        self,
+        image: Union[PIL.Image.Image, torch.Tensor],
+        concepts: List[str],
+        prompt: str = "",  # Optional prompt context
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_samples: int = 1,
+        num_inference_steps: int = 28,
+        noise_timestep: int = 15,  # How much noise to add (higher = more noise)
+        guidance_scale: float = 3.5,
+        generator: Optional[torch.Generator] = None,
+        layers: List[int] = [15, 16, 17, 18],
+        timesteps_to_analyze: List[int] = [24, 25, 26, 27],
+        device: Optional[torch.device] = None,
+        return_dict: bool = True,
+        max_sequence_length: int = 512,
+    ) -> Union[torch.Tensor, Dict]:
+        """
+        Encode an existing image to analyze concept attention patterns.
+        
+        This method adds controlled noise to an input image and runs it through
+        the denoising process to capture how the model attends to different concepts.
+        
+        Args:
+            image: Input image to analyze (PIL Image or tensor)
+            concepts: List of concept strings to analyze
+            prompt: Optional text prompt for context
+            height, width: Output dimensions (inferred from image if not provided)
+            num_samples: Number of noise samples to average over
+            num_inference_steps: Total denoising steps
+            noise_timestep: Which timestep to add noise at (higher = more noise)
+            guidance_scale: Guidance scale for generation
+            generator: Random generator for reproducibility
+            layers: Which transformer layers to analyze
+            timesteps_to_analyze: Which denoising timesteps to capture attention from
+            device: Compute device
+            return_dict: Whether to return dict or tuple
+            max_sequence_length: Max sequence length for text encoding
+            
+        Returns:
+            Dictionary containing original image and concept attention maps
+        """
+        device = device or self._execution_device
+        
+        # Process input image
+        if isinstance(image, PIL.Image.Image):
+            if height is None or width is None:
+                width, height = image.size
+            # Ensure dimensions are compatible
+            height = height - (height % (self.vae_scale_factor * 2))
+            width = width - (width % (self.vae_scale_factor * 2))
+            
+            # Preprocess image
+            processed_image = self.image_processor.preprocess(
+                image, height=height, width=width
+            ).to(device=device, dtype=self.vae.dtype)
+        else:
+            processed_image = image.to(device=device, dtype=self.vae.dtype)
+            if height is None or width is None:
+                _, _, height, width = processed_image.shape
+        
+        # Encode image to latent space
+        with torch.no_grad():
+            if processed_image.shape[1] != self.latent_channels:
+                image_latents = self._encode_vae_image(processed_image, generator)
+            else:
+                image_latents = processed_image
+        
+        # Setup text embeddings
+        prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
+            prompt=prompt or "",
+            prompt_2=None,
+            device=device,
+            num_images_per_prompt=1,
+            max_sequence_length=max_sequence_length,
+        )
+        
+        # Setup concept embeddings
+        concept_embeddings, pooled_concept_embeds, concept_ids = self.encode_concepts(
+            concepts, device=device
+        )
+        
+        # Prepare timesteps
+        image_seq_len = (height // self.vae_scale_factor // 2) * (width // self.vae_scale_factor // 2)
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        
+        # Get full timestep schedule
+        timesteps, _ = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps=None,
+            sigmas=None,
+            mu=mu,
+        )
+        
+        # Prepare latent image IDs
+        latent_image_ids = self._prepare_latent_image_ids(
+            1, 
+            2 * (height // self.vae_scale_factor),
+            2 * (width // self.vae_scale_factor), 
+            device, 
+            image_latents.dtype
+        )
+        
+        # Handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+        else:
+            guidance = None
+        
+        # Collect attention maps across samples
+        all_concept_attention_maps = []
+
+        for sample_idx in range(num_samples):
+            # Add noise at specified timestep
+            if generator is not None:
+                # Use different seed for each sample
+                sample_generator = torch.Generator(device=device).manual_seed(
+                    generator.initial_seed() + sample_idx
+                )
+                noise = torch.randn(
+                    image_latents.shape, 
+                    generator=sample_generator, 
+                    device=device, 
+                    dtype=image_latents.dtype
+                )
+            else:
+                noise = torch.randn_like(image_latents)
+            
+            # Get the timestep tensor for noise addition
+            t_noise = timesteps[noise_timestep]
+            noisy_latents = self.scheduler.scale_noise(
+                image_latents, t_noise.unsqueeze(0), noise
+            )
+            
+            # Pack latents for transformer
+            packed_latents = self._pack_latents(
+                noisy_latents, 
+                1, 
+                self.transformer.config.in_channels // 4,
+                2 * (height // self.vae_scale_factor),
+                2 * (width // self.vae_scale_factor)
+            )
+            
+            # Run single denoising step with concept attention
+            timestep_tensor = t_noise.expand(1).to(packed_latents.dtype) / 1000
+            
+            concept_attention_kwargs = {
+                "concepts": concepts,
+                "timesteps": [0],  # We're only doing one step
+                "layers": layers
+            }
+            
+            with torch.no_grad():
+                transformer_output = self.transformer(
+                    hidden_states=packed_latents,
+                    timestep=timestep_tensor,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    pooled_concept_embeds=pooled_concept_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    concept_hidden_states=concept_embeddings,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    concept_ids=concept_ids,
+                    joint_attention_kwargs=None,
+                    concept_attention_kwargs=concept_attention_kwargs,
+                    return_dict=False,
+                )
+                
+                _, sample_concept_attention_maps = transformer_output
+                
+                if sample_concept_attention_maps:
+                    all_concept_attention_maps.extend(sample_concept_attention_maps)
+        
+        # Process collected attention maps
+        concept_attention_maps = []
+        if all_concept_attention_maps:
+            # Extract and stack vectors across samples
+            concept_vectors = []
+            image_vectors = []
+            
+            for sample_data in all_concept_attention_maps:
+                concept_vectors.append(sample_data['concept_vectors'].detach().cpu())
+                image_vectors.append(sample_data['image_vectors'].detach().cpu())
+            
+            if concept_vectors and image_vectors:
+                # Stack and average across samples
+                concept_vectors = torch.stack(concept_vectors, dim=0).mean(dim=0)
+                image_vectors = torch.stack(image_vectors, dim=0).mean(dim=0)
+                
+                # Normalize concept vectors
+                concept_vectors = concept_vectors / (concept_vectors.norm(dim=-1, keepdim=True) + 1e-8)
+                
+                # Compute attention maps
+                attention_maps = einops.einsum(
+                    image_vectors,
+                    concept_vectors,
+                    "batch patches dim, batch concepts dim -> batch concepts patches"
+                )
+                
+                # Apply softmax normalization
+                attention_maps = torch.softmax(attention_maps, dim=-2)
+                
+                # Convert to numpy and reshape to spatial format
+                attention_maps = attention_maps.float().numpy()
+                spatial_maps = einops.rearrange(
+                    attention_maps,
+                    "batch concepts (h w) -> batch concepts h w",
+                    h=height // 16,
+                    w=width // 16,
+                )
+                
+                # Process maps per batch (usually just 1 batch for encode_image)
+                for b in range(spatial_maps.shape[0]):
+                    maps = spatial_maps[b]  # (concepts, H, W)
+                    
+                    # Normalize maps
+                    vmin, vmax = maps.min(), maps.max()
+                    if vmax > vmin:
+                        maps = (maps - vmin) / (vmax - vmin)
+                    else:
+                        maps = np.zeros_like(maps)
+                    
+                    # Convert to list of arrays per concept
+                    concept_maps = [maps[i] for i in range(maps.shape[0])]
+                    concept_attention_maps.append(concept_maps)
+        
+        # Prepare output
+        result = {
+            'image': image,  # Original input image
+            'concept_attention_maps': concept_attention_maps,
+            'concepts': concepts,
+            'height': height,
+            'width': width
+        }
+        
+        if not return_dict:
+            return (image, concept_attention_maps)
+        
+        return result
+
+    

models/blocks.py

@@ -0,0 +1,127 @@
+import torch.nn as nn
+
+
+def _make_scratch(in_shape, out_shape, groups=1, expand=False):
+    scratch = nn.Module()
+
+    out_shape1 = out_shape
+    out_shape2 = out_shape
+    out_shape3 = out_shape
+    if len(in_shape) >= 4:
+        out_shape4 = out_shape
+
+    if expand:
+        out_shape1 = out_shape
+        out_shape2 = out_shape*2
+        out_shape3 = out_shape*4
+        if len(in_shape) >= 4:
+            out_shape4 = out_shape*8
+
+    scratch.layer1_rn = nn.Conv2d(
+        in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    scratch.layer2_rn = nn.Conv2d(
+        in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    scratch.layer3_rn = nn.Conv2d(
+        in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+    )
+    if len(in_shape) >= 4:
+        scratch.layer4_rn = nn.Conv2d(
+            in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+        )
+
+    return scratch
+
+
+class ResidualConvUnit(nn.Module):
+    """Residual convolution module."""
+
+    def __init__(self, features, activation, bn):
+        super().__init__()
+
+        self.bn = bn
+
+        self.groups=1
+
+        self.conv1 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
+        )
+        
+        self.conv2 = nn.Conv2d(
+            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
+        )
+
+        if self.bn==True:
+            self.bn1 = nn.BatchNorm2d(features)
+            self.bn2 = nn.BatchNorm2d(features)
+
+        self.activation = activation
+
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x):
+        out = self.activation(x)
+        out = self.conv1(out)
+        if self.bn==True:
+            out = self.bn1(out)
+       
+        out = self.activation(out)
+        out = self.conv2(out)
+        if self.bn==True:
+            out = self.bn2(out)
+
+        if self.groups > 1:
+            out = self.conv_merge(out)
+
+        return self.skip_add.add(out, x)
+
+
+class FeatureFusionBlock(nn.Module):
+    """Feature fusion block."""
+
+    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
+        super(FeatureFusionBlock, self).__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+
+        self.groups=1
+
+        self.expand = expand
+        out_features = features
+        if self.expand==True:
+            out_features = features//2
+        
+        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
+
+        self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
+        self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
+        
+        self.skip_add = nn.quantized.FloatFunctional()
+
+        self.size=size
+
+    def forward(self, *xs, size=None):
+        output = xs[0]
+
+        if len(xs) == 2:
+            res = self.resConfUnit1(xs[1])
+            output = self.skip_add.add(output, res)
+
+        output = self.resConfUnit2(output)
+
+        if (size is None) and (self.size is None):
+            modifier = {"scale_factor": 2}
+        elif size is None:
+            modifier = {"size": self.size}
+        else:
+            modifier = {"size": size}
+
+        output = nn.functional.interpolate(
+            output, **modifier, mode="bilinear", align_corners=self.align_corners
+        )
+
+        output = self.out_conv(output)
+
+        return output

models/dino_fusion.py

@@ -0,0 +1,83 @@
+"""torch.hub DINOv3/DINOv2 feature extractor (fallback for the HF extractor).
+
+Used by `core.build_dino_extractor` only when a backbone name is not in the
+HuggingFace repo map. Extracts intermediate transformer layers as spatial maps.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DINOv3FeatureExtractor(nn.Module):
+    """Extract intermediate ViT layers from a DINOv3/DINOv2 model loaded via torch.hub.
+
+    ``take_indices`` selects the layers to return (a list of block indices, matching
+    the HuggingFace extractor); the backbone is frozen.
+    """
+    def __init__(self, model_name="dinov3_vitb16", take_indices=(2, 5, 8, 11)):
+        super().__init__()
+        self.model_name = model_name
+        self.take_indices = list(take_indices)
+
+        if model_name.startswith('facebook/'):
+            hf_to_hub = {
+                'facebook/dinov2-base': 'dinov2_vitb14',
+                'facebook/dinov2-small': 'dinov2_vits14',
+                'facebook/dinov2-large': 'dinov2_vitl14',
+                'facebook/dinov2-giant': 'dinov2_vitg14',
+                'facebook/dinov3-vitb16-pretrain-lvd1689m': 'dinov3_vitb16',
+                'facebook/dinov3-vits16-pretrain-lvd1689m': 'dinov3_vits16',
+                'facebook/dinov3-vitl16-pretrain-lvd1689m': 'dinov3_vitl16',
+            }
+            model_name = hf_to_hub.get(model_name, 'dinov2_vitb14')
+            self.model_name = model_name
+
+        if 'dinov2' in model_name:
+            self.dino = torch.hub.load('facebookresearch/dinov2', model_name)
+            self.patch_size = 14
+        elif 'dinov3' in model_name:
+            try:
+                self.dino = torch.hub.load('facebookresearch/dinov3', model_name)
+            except Exception as e:
+                # torch.hub gated/unavailable -> fall back to HuggingFace weights.
+                print(f"[DINO] torch.hub failed ({e}); loading from HuggingFace")
+                from transformers import AutoModel
+                hf_map = {
+                    'dinov3_vits16': 'facebook/dinov3-vits16-pretrain-lvd1689m',
+                    'dinov3_vitb16': 'facebook/dinov3-vitb16-pretrain-lvd1689m',
+                    'dinov3_vitl16': 'facebook/dinov3-vitl16-pretrain-lvd1689m',
+                }
+                self.dino = AutoModel.from_pretrained(
+                    hf_map.get(model_name, 'facebook/dinov3-vitb16-pretrain-lvd1689m'),
+                    trust_remote_code=True)
+            self.patch_size = 16
+        else:
+            raise ValueError(f"Unsupported model name: {model_name}. Use dinov2_* or dinov3_*")
+
+        self.dino.eval()
+        for p in self.dino.parameters():
+            p.requires_grad = False
+
+    @torch.no_grad()
+    def forward(self, images):
+        """images: [B, 3, H, W] in [0, 1] -> list of feature maps [B, C, H//p, W//p]."""
+        h, w = images.shape[-2:]
+        if h % self.patch_size != 0 or w % self.patch_size != 0:
+            new_h = ((h + self.patch_size - 1) // self.patch_size) * self.patch_size
+            new_w = ((w + self.patch_size - 1) // self.patch_size) * self.patch_size
+            images = F.interpolate(images, size=(new_h, new_w), mode='bilinear', align_corners=False)
+            h, w = new_h, new_w
+
+        features = self.dino.get_intermediate_layers(images, self.take_indices, return_class_token=False)
+
+        patch_h, patch_w = h // self.patch_size, w // self.patch_size
+        outs = []
+        for feat in features:
+            B, N, C = feat.shape
+            outs.append(feat.permute(0, 2, 1).reshape(B, C, patch_h, patch_w).contiguous())
+        return outs
+
+
+def create_dino_extractor(model_name="dinov3_vitb16", take_indices=(2, 5, 8, 11)):
+    return DINOv3FeatureExtractor(model_name=model_name, take_indices=take_indices)

models/dinov3_hf_extractor.py

@@ -0,0 +1,135 @@
+"""DINOv3 feature extraction via HuggingFace, handling CLS and register tokens."""
+
+import torch
+import torch.nn as nn
+from transformers import AutoImageProcessor, AutoModel
+
+
+class DINOv3HFExtractor(nn.Module):
+    """
+    Extracts intermediate features from DINOv3 via HuggingFace transformers.
+    
+    Properly handles DINOv3 register tokens:
+    - Output shape: [B, 1 + P + R, C] where:
+      - 1 = CLS token
+      - P = spatial patches (1024 for 512x512 with 16x16 patches)
+      - R = register tokens (typically 4 for DINOv3)
+    
+    Returns 4 feature maps of shape [B, C_dino, 32, 32] from last 4 layers.
+    Input images must be [B, 3, 512, 512] in [0, 1] range.
+    """
+    
+    def __init__(self, repo_id="facebook/dinov3-vitb16-pretrain-lvd1689m", take_last=None, take_indices=None, trainable=False):
+        super().__init__()
+
+        self.proc = AutoImageProcessor.from_pretrained(repo_id)
+        
+        # Disable resizing/cropping so native 512x512 maps to 32x32 patches
+        for k in ("do_resize", "do_center_crop"):
+            if hasattr(self.proc, k):
+                setattr(self.proc, k, False)
+        
+        self.model = AutoModel.from_pretrained(repo_id)
+        self.model.config.output_hidden_states = True
+        
+        # trainable=True is used by the `dino_only` ablation (fine-tune the backbone);
+        # otherwise the backbone is frozen and kept in eval mode.
+        self._frozen = not trainable
+        if self._frozen:
+            self.model.eval()
+            for p in self.model.parameters():
+                p.requires_grad = False
+        else:
+            self.model.train()
+            for p in self.model.parameters():
+                p.requires_grad = True
+        
+        # ImageNet normalization stats: mean [0.485, 0.456, 0.406], std [0.229, 0.224, 0.225]
+        mean = torch.tensor(self.proc.image_mean).view(1, 3, 1, 1)
+        std = torch.tensor(self.proc.image_std).view(1, 3, 1, 1)
+        self.register_buffer("mean", mean, persistent=False)
+        self.register_buffer("std", std, persistent=False)
+        
+        if take_indices is not None:
+            self.take_indices = take_indices
+            self.take_last = None
+        else:
+            self.take_last = take_last if take_last is not None else 4
+            self.take_indices = None
+        
+        self.patch_size = getattr(self.model.config, "patch_size", 16)
+        self.num_register_tokens = getattr(self.model.config, "num_register_tokens", 0)
+        
+        hidden_size = getattr(self.model.config, "hidden_size", 768)
+        
+        layers = self.take_indices if self.take_indices is not None else f"last {self.take_last}"
+        trainable_str = "trainable" if not self._frozen else "frozen"
+        print(f"[DINOv3] {repo_id.split('/')[-1]}: dim={hidden_size}, patch={self.patch_size}, "
+              f"layers={layers} ({trainable_str})")
+    
+    def train(self, mode: bool = True):
+        """Keep a frozen backbone in eval mode; otherwise follow `mode`."""
+        self.training = mode
+        if self._frozen:
+            self.model.eval()
+        else:
+            self.model.train(mode)
+        return self
+    
+    def forward(self, images_512: torch.Tensor):
+        """Extract DINOv3 features from [B, 3, H, W] images in [0, 1]; returns maps [B, C, H//16, W//16]."""
+        # Disable grad only when the backbone is frozen; otherwise allow fine-tuning
+        with torch.set_grad_enabled(not self._frozen):
+            return self._forward(images_512)
+
+    def _forward(self, images_512: torch.Tensor):
+        x = (images_512 - self.mean) / self.std
+        
+        out = self.model(pixel_values=x, output_hidden_states=True)
+        hidden_states = out.hidden_states  # Tuple: [emb, blk1, ..., blkL]
+        
+        B, _, H, W = images_512.shape
+        H_patches = H // self.patch_size
+        W_patches = W // self.patch_size
+        P = H_patches * W_patches  # total spatial patches
+        R = self.num_register_tokens
+        
+        # Each hidden state is [B, 1 + P + R, C], laid out as [CLS][P spatial patches][R register tokens]
+        maps = []
+        if self.take_indices is not None:
+            for idx in self.take_indices:
+                hidden = hidden_states[idx]
+                spatial = hidden[:, 1:1+P, :]  # [B, P, C], drop CLS and register tokens
+                C = spatial.shape[-1]
+                spatial_map = spatial.transpose(1, 2).reshape(B, C, H_patches, W_patches).contiguous()
+                maps.append(spatial_map)
+        else:
+            for hidden in hidden_states[-self.take_last:]:
+                spatial = hidden[:, 1:1+P, :]  # [B, P, C], drop CLS and register tokens
+                C = spatial.shape[-1]
+                spatial_map = spatial.transpose(1, 2).reshape(B, C, H_patches, W_patches).contiguous()
+                maps.append(spatial_map)
+        
+        return maps
+
+
+def create_dinov3_hf_extractor(repo_id="facebook/dinov3-vitb16-pretrain-lvd1689m", take_last=None, take_indices=None, trainable=False):
+    """
+    Factory for DINOv3HFExtractor (frozen in eval mode unless trainable=True).
+    repo_id options: vits16 (384-dim), vitb16 (768-dim, recommended), vitl16 (1024-dim).
+    """
+    return DINOv3HFExtractor(repo_id=repo_id, take_last=take_last, take_indices=take_indices, trainable=trainable)
+
+
+if __name__ == "__main__":
+    extractor = create_dinov3_hf_extractor().cuda()
+    
+    dummy_input = torch.randn(2, 3, 512, 512).cuda()
+    print(f"\nInput: {dummy_input.shape}")
+    
+    features = extractor(dummy_input)
+    
+    print(f"\nExtracted {len(features)} feature maps:")
+    for i, feat in enumerate(features):
+        print(f"  Layer {i}: {feat.shape}")
+

models/dpt_backbone.py

@@ -0,0 +1,44 @@
+"""Shared DPT reassemble + RefineNet cascade.
+
+`DPTRefineNetStack` owns the `scratch` reassemble layers and the four
+`FeatureFusionBlock` refinenets used by every DPT decoder in this repo (the depth
+heads in ``dpt_decoder.py`` and the saliency decoder in
+``dpt_segmentation_decoder.py``). Decoders subclass it so the parameter names stay
+flat (``scratch.*`` / ``refinenet{1..4}.*``) and existing checkpoints keep loading;
+each subclass provides its own input projection and output head.
+"""
+
+import torch.nn as nn
+
+from .blocks import FeatureFusionBlock, _make_scratch
+
+
+class DPTRefineNetStack(nn.Module):
+    def __init__(self, features=256, use_bn=False, out_channels=(256, 512, 1024, 1024)):
+        super().__init__()
+        self.features = features
+        self.scratch = _make_scratch(list(out_channels), features)
+        self.refinenet4 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
+        self.refinenet3 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
+        self.refinenet2 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
+        self.refinenet1 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
+
+    def fuse(self, layers, keep_layer1_size=False):
+        """Run the coarse-to-fine RefineNet cascade; returns the layer-1 feature map.
+
+        ``keep_layer1_size=True`` stops the final block from doing its default 2x
+        upsample (used by the high-res depth decoder, which upsamples in its head).
+        """
+        l1, l2, l3, l4 = layers
+        l1 = self.scratch.layer1_rn(l1)
+        l2 = self.scratch.layer2_rn(l2)
+        l3 = self.scratch.layer3_rn(l3)
+        l4 = self.scratch.layer4_rn(l4)
+        path = self.refinenet4(l4, size=l3.shape[2:])
+        path = self.refinenet3(path, l3, size=l2.shape[2:])
+        path = self.refinenet2(path, l2, size=l1.shape[2:])
+        if keep_layer1_size:
+            path = self.refinenet1(path, l1, size=l1.shape[2:])
+        else:
+            path = self.refinenet1(path, l1)
+        return path

models/dpt_decoder.py

@@ -0,0 +1,107 @@
+"""DPT-style depth decoders that consume spatial feature maps.
+
+Both heads take a list of 4 spatial feature maps (FLUX + DINO + concepts) with
+possibly different channel counts, run the shared DPT RefineNet cascade, and
+upsample to a depth map:
+  - DPTHeadSpatial : single 2x upsample in the head (bilinear resize to GT after).
+  - DPTHeadHighRes : 4x learned 2x upsampling (16x) with optional image skip connections.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .dpt_backbone import DPTRefineNetStack
+
+
+class DPTHeadSpatial(DPTRefineNetStack):
+    def __init__(self, in_channels=[3840, 3840, 3840, 3840], features=256,
+                 num_classes=1, use_bn=False, head_features=32):
+        super().__init__(features=features, use_bn=use_bn)
+        self.head_features = head_features
+        self.num_classes = num_classes
+
+        out_channels = [256, 512, 1024, 1024]
+        self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)])
+        self.pre_fuse = nn.ModuleList([
+            nn.Sequential(nn.GroupNorm(1, oc), nn.Conv2d(oc, oc, 1), nn.ReLU(inplace=True))
+            for oc in out_channels])
+
+        self.head = nn.Sequential(
+            nn.Conv2d(features, features // 2, 3, padding=1),
+            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+            nn.Conv2d(features // 2, head_features, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(head_features, num_classes, 1),
+        )
+
+    def forward(self, features):
+        resized = [pf(proj(f)) for f, proj, pf in zip(features, self.projects, self.pre_fuse)]
+        return self.head(self.fuse(resized))
+
+
+class DPTHeadHighRes(DPTRefineNetStack):
+    """High-resolution variant: progressive 4x (2x) learned upsampling = 16x total,
+    with optional skip connections from the original image."""
+    def __init__(self, in_channels=[3840, 3840, 3840, 3840], features=256,
+                 num_classes=1, use_bn=False, head_features=64, use_skip_connections=True):
+        super().__init__(features=features, use_bn=use_bn)
+        self.head_features = head_features
+        self.num_classes = num_classes
+        self.use_skip_connections = use_skip_connections
+
+        out_channels = [256, 512, 1024, 1024]
+        self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)])
+        self.pre_fuse = nn.ModuleList([
+            nn.Sequential(nn.GroupNorm(1, oc), nn.Conv2d(oc, oc, 1), nn.ReLU(inplace=True))
+            for oc in out_channels])
+
+        if use_skip_connections:
+            def skip_enc(stride, out_ch):
+                return nn.Sequential(
+                    nn.Conv2d(3, out_ch, 3, stride=stride, padding=1), nn.ReLU(inplace=True),
+                    nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.ReLU(inplace=True))
+            self.skip_enc_2x = skip_enc(2, 16)
+            self.skip_enc_4x = skip_enc(4, 16)
+            self.skip_enc_8x = skip_enc(8, 16)
+            self.skip_enc_full = skip_enc(1, 8)
+
+        def up_block(c_in, c_out):
+            return nn.Sequential(
+                nn.Conv2d(c_in, c_out, 3, padding=1), nn.ReLU(inplace=True),
+                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
+                nn.Conv2d(c_out, c_out, 3, padding=1), nn.ReLU(inplace=True))
+
+        skip = 16 if use_skip_connections else 0
+        self.up1 = up_block(features, features // 2)
+        self.fuse1 = nn.Conv2d(features // 2 + skip, features // 2, 1) if use_skip_connections else None
+        self.up2 = up_block(features // 2, features // 4)
+        self.fuse2 = nn.Conv2d(features // 4 + skip, features // 4, 1) if use_skip_connections else None
+        self.up3 = up_block(features // 4, head_features)
+        self.fuse3 = nn.Conv2d(head_features + skip, head_features, 1) if use_skip_connections else None
+        self.up4 = up_block(head_features, head_features // 2)
+        self.fuse4 = nn.Conv2d(head_features // 2 + (8 if use_skip_connections else 0), head_features // 2, 1) if use_skip_connections else None
+        self.output = nn.Conv2d(head_features // 2, num_classes, 3, padding=1)
+
+    def _apply_skip(self, x, fuse, skip):
+        if not self.use_skip_connections or skip is None:
+            return x
+        if skip.shape[-2:] != x.shape[-2:]:
+            skip = F.interpolate(skip, size=x.shape[-2:], mode='bilinear', align_corners=True)
+        return fuse(torch.cat([x, skip], dim=1))
+
+    def forward(self, features, image=None):
+        if self.use_skip_connections and image is not None:
+            skip_8x, skip_4x = self.skip_enc_8x(image), self.skip_enc_4x(image)
+            skip_2x, skip_full = self.skip_enc_2x(image), self.skip_enc_full(image)
+        else:
+            skip_8x = skip_4x = skip_2x = skip_full = None
+
+        resized = [pf(proj(f)) for f, proj, pf in zip(features, self.projects, self.pre_fuse)]
+        path_1 = self.fuse(resized, keep_layer1_size=True)  # keep patch resolution
+
+        x = self._apply_skip(self.up1(path_1), self.fuse1, skip_8x)
+        x = self._apply_skip(self.up2(x), self.fuse2, skip_4x)
+        x = self._apply_skip(self.up3(x), self.fuse3, skip_2x)
+        x = self._apply_skip(self.up4(x), self.fuse4, skip_full)
+        return self.output(x)

models/dpt_segmentation_decoder.py

@@ -0,0 +1,61 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .dpt_backbone import DPTRefineNetStack
+
+
+class ProgressiveOutputHead(nn.Module):
+    """Progressive channel-reduction head (GroupNorm for AMP/small-batch stability)."""
+    def __init__(self, in_features=256, num_classes=21, intermediate_features=32, groups=8):
+        super().__init__()
+        self.num_classes = num_classes
+        self.output_conv1 = nn.Sequential(
+            nn.Conv2d(in_features, in_features // 2, 3, padding=1, bias=False),
+            nn.GroupNorm(num_groups=groups, num_channels=in_features // 2),
+            nn.ReLU(inplace=True),
+        )
+        self.output_conv2 = nn.Sequential(
+            nn.Conv2d(in_features // 2, intermediate_features, 3, padding=1, bias=False),
+            nn.GroupNorm(num_groups=max(1, min(groups, intermediate_features)), num_channels=intermediate_features),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(intermediate_features, num_classes, 1),
+        )
+        self.final_activation = nn.Identity()  # sigmoid/softmax applied in the loss
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        # Background-prior bias on the final logits to avoid early sigmoid overflow.
+        final_conv = self.output_conv2[-1]
+        if isinstance(final_conv, nn.Conv2d) and final_conv.bias is not None:
+            nn.init.constant_(final_conv.bias, -2.2 if self.num_classes == 1 else -0.5)
+
+    def forward(self, x):
+        return self.final_activation(self.output_conv2(self.output_conv1(x)))
+
+
+class OriginalDPTSegmentationDecoder(DPTRefineNetStack):
+    def __init__(self, in_channels, num_classes=21, features=256, target_size=(512, 512)):
+        super().__init__(features=features, use_bn=False)
+        self.target_size = target_size
+
+        out_channels = [256, 512, 1024, 1024]
+        self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)])
+        self.resize_layers = nn.ModuleList([
+            nn.ConvTranspose2d(out_channels[0], out_channels[0], kernel_size=4, stride=4),
+            nn.ConvTranspose2d(out_channels[1], out_channels[1], kernel_size=2, stride=2),
+            nn.Identity(),
+            nn.Conv2d(out_channels[3], out_channels[3], kernel_size=3, stride=2, padding=1),
+        ])
+        self.output_head = ProgressiveOutputHead(
+            in_features=features, num_classes=num_classes, intermediate_features=32, groups=8)
+
+    def forward(self, features):
+        proj = [p(f) for p, f in zip(self.projects, features)]
+        resized = [r(f) for r, f in zip(self.resize_layers, proj)]
+        logits = self.output_head(self.fuse(resized))
+        if self.target_size is not None:
+            logits = F.interpolate(logits, size=self.target_size, mode='bilinear', align_corners=True)
+        return logits

models/flux_resizer.py

@@ -0,0 +1,90 @@
+"""Image resizing to FLUX-compatible resolutions (dimensions divisible by 32)."""
+
+import cv2
+import numpy as np
+from typing import Tuple
+from PIL import Image
+
+
+class FluxResizer:
+    """
+    Resizer that ensures images are compatible with FLUX requirements.
+    - FLUX: Dimensions divisible by 32 (due to 2x2 packing on top of 16-stride VAE)
+    """
+    
+    # Predefined optimal resolutions (all divisible by 32)
+    OPTIMAL_RESOLUTIONS = [
+        # Square and near-square
+        (1024, 1024),  # 1:1 (64×64, 64×64)
+        (896, 1152),   # ~0.78:1 (56×72)
+        (1152, 896),   # ~1.29:1 (72×56)
+        (768, 1344),   # ~0.57:1 (48×84)
+        (1344, 768),   # ~1.75:1 (84×48)
+        
+        # Additional common ratios
+        (832, 1216),   # ~0.68:1 (52×76)
+        (1216, 832),   # ~1.46:1 (76×52)
+        (704, 1408),   # 0.5:1 (44×88)
+        (1408, 704),   # 2:1 (88×44)
+        (960, 1088),   # ~0.88:1 (60×68)
+        (1088, 960),   # ~1.13:1 (68×60)
+    ]
+    
+    def __init__(self):
+        self.resolution_aspects = [
+            (h, w, w / h) for h, w in self.OPTIMAL_RESOLUTIONS
+        ]
+    
+    def select_best_resolution(self, original_h: int, original_w: int) -> Tuple[int, int]:
+        """Pick the optimal resolution whose aspect ratio best matches the input."""
+        original_aspect = original_w / original_h
+        
+        best_resolution = None
+        min_aspect_diff = float('inf')
+        
+        for h, w, aspect in self.resolution_aspects:
+            aspect_diff = abs(original_aspect - aspect)
+            
+            if aspect_diff < min_aspect_diff:
+                min_aspect_diff = aspect_diff
+                best_resolution = (h, w)
+        
+        return best_resolution
+    
+    def resize_image(self, image: np.ndarray) -> Tuple[np.ndarray, Tuple[int, int]]:
+        """Resize an (H, W, C) numpy image to the optimal FLUX-compatible resolution."""
+        original_h, original_w = image.shape[:2]
+        target_h, target_w = self.select_best_resolution(original_h, original_w)
+        
+        resized_image = cv2.resize(image, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
+        
+        return resized_image, (target_h, target_w)
+    
+    def resize_pil_image(self, image: Image.Image) -> Tuple[Image.Image, Tuple[int, int]]:
+        """Resize a PIL image to the optimal FLUX-compatible resolution."""
+        original_w, original_h = image.size  # PIL uses (W, H)
+        target_h, target_w = self.select_best_resolution(original_h, original_w)
+        
+        resized_image = image.resize((target_w, target_h), Image.LANCZOS)
+        
+        return resized_image, (target_h, target_w)
+    
+    def resize_mask(self, mask: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray:
+        """Resize a mask to target_size=(H, W) using nearest-neighbor interpolation."""
+        target_h, target_w = target_size
+        
+        if len(mask.shape) == 3 and mask.shape[2] == 1:
+            mask = mask.squeeze(2)
+        
+        resized_mask = cv2.resize(mask, (target_w, target_h), interpolation=cv2.INTER_NEAREST)
+        
+        return resized_mask
+    
+    def get_compatible_resolutions(self) -> list:
+        """Return list of all compatible resolutions."""
+        return self.OPTIMAL_RESOLUTIONS.copy()
+    
+    @staticmethod
+    def verify_compatibility(height: int, width: int) -> bool:
+        """True if both dimensions are divisible by 32 (FLUX requirement)."""
+        return (height % 32 == 0) and (width % 32 == 0)

models/hyperfeature_fusion.py

@@ -0,0 +1,441 @@
+"""
+Aggregates features across multiple timesteps and layers using learned attention,
+allowing the model to adaptively combine information from different denoising stages.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Dict, Tuple
+
+
+class FeatureNormalizer(nn.Module):
+    """
+    Normalizes features from different timesteps/layers to comparable scales.
+    Uses learnable normalization parameters per timestep.
+    """
+    def __init__(self, feature_dim: int = 3072, num_timesteps: int = 4):
+        super().__init__()
+        self.feature_dim = feature_dim
+        self.num_timesteps = num_timesteps
+        
+        # Learnable normalization parameters (per timestep)
+        self.layer_norms = nn.ModuleList([
+            nn.LayerNorm(feature_dim) for _ in range(num_timesteps)
+        ])
+    
+    def forward(self, features_per_timestep: List[torch.Tensor]) -> List[torch.Tensor]:
+        normalized = []
+        for i, feat in enumerate(features_per_timestep):
+            # LayerNorm expects channels last: [B, C, H, W] -> [B, H, W, C]
+            B, C, H, W = feat.shape
+            feat_reshaped = feat.permute(0, 2, 3, 1).contiguous()  # [B, H, W, C]
+            
+            normalized_feat = self.layer_norms[i](feat_reshaped)
+            
+            normalized_feat = normalized_feat.permute(0, 3, 1, 2).contiguous()  # back to [B, C, H, W]
+            normalized.append(normalized_feat)
+        
+        return normalized
+
+
+class FeatureProjector(nn.Module):
+    """
+    Projects features from different layers/timesteps to a common dimension.
+    Uses 1x1 convolutions for efficient channel reduction.
+    """
+    def __init__(self, in_channels: int = 3072, out_channels: int = 512, num_timesteps: int = 4):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        
+        # Separate linear projection for each timestep (allows timestep-specific transformations)
+        # No activation here: the transformer and output_proj provide nonlinearities,
+        # and keeping this linear ensures the skip connection can correct in both directions.
+        self.projectors = nn.ModuleList([
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
+            for _ in range(num_timesteps)
+        ])
+    
+    def forward(self, features_per_timestep: List[torch.Tensor]) -> List[torch.Tensor]:
+        return [self.projectors[i](feat) for i, feat in enumerate(features_per_timestep)]
+
+
+class CrossTimestepAttention(nn.Module):
+    """
+    Pixel-wise cross-timestep attention: for each spatial location, attends across
+    features from different timesteps to learn per-pixel timestep importance.
+    """
+    def __init__(self, feature_dim: int = 512, num_timesteps: int = 4, num_heads: int = 4):
+        super().__init__()
+        self.feature_dim = feature_dim
+        self.num_timesteps = num_timesteps
+        self.num_heads = num_heads
+        self.head_dim = feature_dim // num_heads
+        
+        assert feature_dim % num_heads == 0, "feature_dim must be divisible by num_heads"
+        
+        self.q_proj = nn.Linear(feature_dim, feature_dim)
+        self.k_proj = nn.Linear(feature_dim, feature_dim)
+        self.v_proj = nn.Linear(feature_dim, feature_dim)
+        
+        self.out_proj = nn.Linear(feature_dim, feature_dim)
+        
+        self.scale = self.head_dim ** -0.5
+    
+    def forward(self, features_per_timestep: List[torch.Tensor]) -> torch.Tensor:
+        B, C, H, W = features_per_timestep[0].shape
+        T = len(features_per_timestep)
+        
+        stacked = torch.stack(features_per_timestep, dim=1)  # [B, T, C, H, W]
+        
+        stacked = stacked.permute(0, 3, 4, 1, 2).contiguous()  # [B, H, W, T, C]
+        stacked = stacked.view(B * H * W, T, C)
+        
+        # Query from last (most refined) timestep; keys/values from all timesteps
+        q = self.q_proj(stacked[:, -1:, :])  # [B*H*W, 1, C]
+        k = self.k_proj(stacked)  # [B*H*W, T, C]
+        v = self.v_proj(stacked)  # [B*H*W, T, C]
+        
+        q = q.view(B * H * W, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B * H * W, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B * H * W, T, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale  # [B*H*W, num_heads, 1, T]
+        attn = F.softmax(attn, dim=-1)
+        
+        out = torch.matmul(attn, v)  # [B*H*W, num_heads, 1, head_dim]
+        
+        out = out.transpose(1, 2).contiguous().view(B * H * W, 1, C)
+        
+        out = self.out_proj(out)
+        
+        out = out.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
+        
+        return out
+
+
+class LayerScaleTransformerLayer(nn.Module):
+    """
+    Transformer layer with LayerScale (from CaiT: "Going deeper with Image Transformers").
+    LayerScale helps stabilize training of deeper transformers by adding learnable
+    diagonal scaling matrices initialized to small values (e.g., 1e-4).
+    """
+    def __init__(self, feature_dim: int, nhead: int = 8, dropout: float = 0.1, 
+                 layer_scale_init: float = 1e-4):
+        super().__init__()
+        self.feature_dim = feature_dim
+        
+        self.self_attn = nn.MultiheadAttention(
+            feature_dim, nhead, dropout=dropout, batch_first=True
+        )
+        
+        self.ffn = nn.Sequential(
+            nn.Linear(feature_dim, feature_dim * 2),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(feature_dim * 2, feature_dim),
+            nn.Dropout(dropout)
+        )
+        
+        self.norm1 = nn.LayerNorm(feature_dim)
+        self.norm2 = nn.LayerNorm(feature_dim)
+        
+        # LayerScale diagonal scaling, initialized small to stabilize deep transformers
+        self.layer_scale_1 = nn.Parameter(
+            torch.ones(feature_dim) * layer_scale_init
+        )
+        self.layer_scale_2 = nn.Parameter(
+            torch.ones(feature_dim) * layer_scale_init
+        )
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Pre-norm self-attention with LayerScale (CaiT style)
+        x_normed = self.norm1(x)
+        attn_out, _ = self.self_attn(x_normed, x_normed, x_normed)
+        x = x + self.layer_scale_1 * attn_out
+        
+        ffn_out = self.ffn(self.norm2(x))
+        x = x + self.layer_scale_2 * ffn_out
+        
+        return x
+
+
+class TemporalTransformerAggregator(nn.Module):
+    """
+    Transformer that aggregates features across timesteps with LayerScale.
+    Predicts content-adaptive per-pixel fusion weights rather than global weights.
+    """
+    def __init__(self, feature_dim: int = 512, num_timesteps: int = 4, num_layers: int = 2,
+                 layer_scale_init: float = 1e-4):
+        super().__init__()
+        self.feature_dim = feature_dim
+        self.num_timesteps = num_timesteps
+        self.num_layers = num_layers
+        
+        # Learnable temporal positional encoding, small init to avoid early dominance
+        self.temporal_pos_embed = nn.Parameter(torch.randn(1, num_timesteps, feature_dim))
+        nn.init.trunc_normal_(self.temporal_pos_embed, std=0.02)
+        
+        self.layers = nn.ModuleList([
+            LayerScaleTransformerLayer(
+                feature_dim=feature_dim,
+                nhead=8,
+                dropout=0.1,
+                layer_scale_init=layer_scale_init
+            ) for _ in range(num_layers)
+        ])
+        
+        # Per-pixel timestep weights: [B*H*W, C, T] -> [B*H*W, 1, T]
+        self.alpha_head = nn.Conv1d(self.feature_dim, 1, kernel_size=1)
+        # Small init so weighting starts nearly uniform across timesteps
+        nn.init.normal_(self.alpha_head.weight, mean=0.0, std=0.01)
+        nn.init.zeros_(self.alpha_head.bias)
+    
+    def forward(self, features_per_timestep: List[torch.Tensor], return_alpha: bool = False):
+        """If return_alpha, also returns the per-pixel timestep weight map [B, T, H, W]."""
+        B, C, H, W = features_per_timestep[0].shape
+        T = len(features_per_timestep)
+        
+        stacked = torch.stack(features_per_timestep, dim=1)  # [B, T, C, H, W]
+        stacked = stacked.permute(0, 3, 4, 1, 2).contiguous()  # [B, H, W, T, C]
+        stacked = stacked.view(B * H * W, T, C)
+        
+        stacked = stacked + self.temporal_pos_embed
+        
+        transformed = stacked
+        for layer in self.layers:
+            transformed = layer(transformed)  # [B*H*W, T, C]
+        
+        transformed_t = transformed.transpose(1, 2)  # [B*H*W, C, T]
+        
+        logits_alpha = self.alpha_head(transformed_t)  # [B*H*W, 1, T]
+        logits_alpha = logits_alpha.squeeze(1)  # [B*H*W, T]
+        
+        # Softmax over timesteps per pixel
+        alpha = torch.softmax(logits_alpha, dim=-1)  # [B*H*W, T]
+        
+        alpha_expanded = alpha.unsqueeze(1)  # [B*H*W, 1, T]
+        pooled = (alpha_expanded * transformed_t).sum(dim=-1)  # [B*H*W, C]
+        
+        out = pooled.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
+        
+        if return_alpha:
+            alpha_map = alpha.view(B, H, W, T).permute(0, 3, 1, 2).contiguous()  # [B, T, H, W]
+            return out, alpha_map
+        
+        return out
+
+
+class CBAM(nn.Module):
+    """
+    Convolutional Block Attention Module (CBAM, ECCV 2018).
+    Sequentially applies channel then spatial attention to refine features.
+    """
+    def __init__(self, channels: int, reduction: int = 16, kernel_size: int = 7):
+        super().__init__()
+        self.channels = channels
+        
+        self.channel_mlp = nn.Sequential(
+            nn.Linear(channels, channels // reduction),
+            nn.ReLU(inplace=True),
+            nn.Linear(channels // reduction, channels)
+        )
+        
+        self.spatial_conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = x.shape
+        
+        avg_pool = x.mean(dim=(2, 3), keepdim=False)  # [B, C]
+        max_pool = x.amax(dim=(2, 3), keepdim=False)  # [B, C]
+        
+        # Shared MLP applied to both pooled descriptors
+        channel_weight = torch.sigmoid(
+            self.channel_mlp(avg_pool) + self.channel_mlp(max_pool)
+        ).view(B, C, 1, 1)
+        
+        x = x * channel_weight
+        
+        avg_spatial = x.mean(dim=1, keepdim=True)  # [B, 1, H, W]
+        max_spatial = x.amax(dim=1, keepdim=True)  # [B, 1, H, W]
+        
+        spatial_weight = torch.sigmoid(
+            self.spatial_conv(torch.cat([avg_spatial, max_spatial], dim=1))
+        )
+        
+        return x * spatial_weight
+
+
+class HyperfeatureFusion(nn.Module):
+    """
+    Complete Hyperfeature Fusion module.
+    Implements the full pipeline: normalize -> project -> attend -> fuse -> refine (CBAM).
+    """
+    def __init__(
+        self,
+        in_channels: int = 3072,
+        out_channels: int = 3072,
+        hidden_dim: int = 512,
+        num_timesteps: int = 4,
+        fusion_type: str = 'attention',  # 'attention' or 'transformer'
+        num_attention_heads: int = 4,
+        num_transformer_layers: int = 2,
+        layer_scale_init: float = 1e-4,
+        return_alpha: bool = False
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_dim = hidden_dim
+        self.num_timesteps = num_timesteps
+        self.fusion_type = fusion_type
+        self.return_alpha = return_alpha
+        
+        self.normalizer = FeatureNormalizer(in_channels, num_timesteps)
+        
+        self.projector = FeatureProjector(in_channels, hidden_dim, num_timesteps)
+        
+        if fusion_type == 'attention':
+            self.aggregator = CrossTimestepAttention(hidden_dim, num_timesteps, num_attention_heads)
+        elif fusion_type == 'transformer':
+            self.aggregator = TemporalTransformerAggregator(
+                hidden_dim, num_timesteps, num_transformer_layers, layer_scale_init
+            )
+        else:
+            raise ValueError(f"Unknown fusion_type: {fusion_type}")
+        
+        self.output_proj = nn.Sequential(
+            nn.Conv2d(hidden_dim, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=1)
+        )
+        
+        # Align hidden_dim to out_channels for the skip from the latest timestep
+        if hidden_dim != out_channels:
+            self.skip_align = nn.Conv2d(hidden_dim, out_channels, kernel_size=1, bias=False)
+        else:
+            self.skip_align = nn.Identity()
+        
+        # Learnable skip gate initialized near zero so the aggregated path dominates
+        # early and the skip is gradually introduced during training.
+        self.skip_gate = nn.Parameter(torch.tensor(0.1))
+        
+        self.cbam = CBAM(out_channels, reduction=16, kernel_size=7)
+    
+    def forward(self, features_per_timestep: List[torch.Tensor]):
+        """features_per_timestep: list of [B, C_in, H, W] tensors ordered early to late."""
+        normalized = self.normalizer(features_per_timestep)
+        
+        projected = self.projector(normalized)
+        
+        if self.return_alpha and self.fusion_type == 'transformer':
+            aggregated, alpha_map = self.aggregator(projected, return_alpha=True)
+        else:
+            aggregated = self.aggregator(projected)
+            alpha_map = None
+        
+        output = self.output_proj(aggregated)
+        
+        # Gated skip connection from the latest timestep preserves refined features
+        latest_skip = self.skip_align(projected[-1])
+        output = output + self.skip_gate * latest_skip
+        
+        output = self.cbam(output)
+        
+        if self.return_alpha and alpha_map is not None:
+            return output, alpha_map
+        
+        return output
+
+
+class MultiLayerHyperfeatureFusion(nn.Module):
+    """
+    Applies Hyperfeature fusion independently for each FLUX feature layer.
+    This allows different layers to learn different temporal aggregation strategies.
+    """
+    def __init__(
+        self,
+        in_channels: int = 3072,
+        out_channels: int = 3072,
+        hidden_dim: int = 512,
+        num_layers: int = 4,
+        num_timesteps: int = 4,
+        fusion_type: str = 'attention',
+        num_attention_heads: int = 4,
+        num_transformer_layers: int = 2,
+        layer_scale_init: float = 1e-4,
+        return_alpha: bool = False
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.num_timesteps = num_timesteps
+        self.return_alpha = return_alpha
+        
+        self.layer_fusions = nn.ModuleList([
+            HyperfeatureFusion(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                hidden_dim=hidden_dim,
+                num_timesteps=num_timesteps,
+                fusion_type=fusion_type,
+                num_attention_heads=num_attention_heads,
+                num_transformer_layers=num_transformer_layers,
+                layer_scale_init=layer_scale_init,
+                return_alpha=return_alpha
+            ) for _ in range(num_layers)
+        ])
+    
+    def forward(self, multi_timestep_features: Dict[int, List[torch.Tensor]]):
+        """multi_timestep_features: dict mapping timestep -> list of per-layer maps [B, C, H, W]."""
+        # Regroup features by layer instead of by timestep
+        features_per_layer = []
+        timesteps = sorted(multi_timestep_features.keys())
+        
+        for layer_idx in range(self.num_layers):
+            layer_features_across_timesteps = [
+                multi_timestep_features[t][layer_idx] for t in timesteps
+            ]
+            features_per_layer.append(layer_features_across_timesteps)
+        
+        fused_layers = []
+        alpha_layers = []
+        for layer_idx, layer_features in enumerate(features_per_layer):
+            result = self.layer_fusions[layer_idx](layer_features)
+            if self.return_alpha:
+                fused, alpha = result
+                fused_layers.append(fused)
+                alpha_layers.append(alpha)
+            else:
+                fused_layers.append(result)
+        
+        if self.return_alpha:
+            return fused_layers, alpha_layers
+        
+        return fused_layers
+
+
+def create_hyperfeature_fusion(
+    num_timesteps: int = 4,
+    num_layers: int = 4,
+    fusion_type: str = 'attention',
+    hidden_dim: int = 512,
+    num_transformer_layers: int = 2,
+    layer_scale_init: float = 1e-4,
+    return_alpha: bool = False,
+    feature_dim: int = 3072
+) -> MultiLayerHyperfeatureFusion:
+    """Factory for MultiLayerHyperfeatureFusion (feature_dim 3072 for FLUX, 1536 for SD3.5)."""
+    return MultiLayerHyperfeatureFusion(
+        in_channels=feature_dim,
+        out_channels=feature_dim,
+        hidden_dim=hidden_dim,
+        num_layers=num_layers,
+        num_timesteps=num_timesteps,
+        fusion_type=fusion_type,
+        num_attention_heads=8,
+        num_transformer_layers=num_transformer_layers,
+        layer_scale_init=layer_scale_init,
+        return_alpha=return_alpha
+    )
+

models/segmentation_losses.py

@@ -0,0 +1,257 @@
+"""Segmentation losses: Lovasz-Softmax, Focal, Dice, and a combined loss."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, List
+import numpy as np
+
+
+def lovasz_grad(gt_sorted):
+    """
+    Compute gradient of the Lovasz extension w.r.t sorted errors.
+    See Algorithm 1 in paper.
+    """
+    p = len(gt_sorted)
+    gts = gt_sorted.sum()
+    intersection = gts - gt_sorted.float().cumsum(0)
+    union = gts + (1 - gt_sorted).float().cumsum(0)
+    jaccard = 1.0 - intersection / union
+    if p > 1:
+        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
+    return jaccard
+
+
+def lovasz_softmax_flat(probas, labels, classes='present', ignore_index=255):
+    """Multi-class Lovasz-Softmax loss over flat [P,C] probabilities and [P] labels."""
+    if probas.numel() == 0:
+        return probas * 0.0
+    
+    C = probas.size(1)
+    losses = []
+    
+    class_to_sum = list(range(C)) if classes == 'all' else []
+    
+    for c in range(C):
+        fg = (labels == c).float()  # foreground for class c
+        if classes == 'present' and fg.sum() == 0:
+            continue
+        if c == ignore_index:
+            continue
+        
+        class_to_sum.append(c) if classes != 'all' else None
+        errors = (fg - probas[:, c]).abs()
+        errors_sorted, perm = torch.sort(errors, 0, descending=True)
+        perm = perm.data
+        fg_sorted = fg[perm]
+        losses.append(torch.dot(errors_sorted, lovasz_grad(fg_sorted)))
+    
+    if len(losses) == 0:
+        return probas.sum() * 0.0
+    
+    return torch.stack(losses).mean()
+
+
+def flatten_probas(probas, labels, ignore_index=255):
+    """Flatten [B,C,H,W] preds and [B,H,W] labels to [P,C]/[P], dropping ignored pixels."""
+    B, C, H, W = probas.size()
+    probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # [B*H*W, C]
+    labels = labels.view(-1)  # [B*H*W]
+    
+    if ignore_index is not None:
+        valid = (labels != ignore_index)
+        probas = probas[valid]
+        labels = labels[valid]
+    
+    return probas, labels
+
+
+class LovaszSoftmaxLoss(nn.Module):
+    """
+    Lovasz-Softmax loss for multi-class semantic segmentation.
+    Directly optimizes the mean IoU (Jaccard index).
+    """
+    
+    def __init__(self, classes='present', ignore_index=255):
+        super().__init__()
+        self.classes = classes
+        self.ignore_index = ignore_index
+    
+    def forward(self, logits, labels):
+        probas = F.softmax(logits, dim=1)
+        probas, labels = flatten_probas(probas, labels, self.ignore_index)
+        return lovasz_softmax_flat(probas, labels, self.classes, self.ignore_index)
+
+
+class FocalLoss(nn.Module):
+    """
+    Focal Loss for multi-class classification.
+    FL(p_t) = -alpha_t * (1 - p_t)^gamma * log(p_t)
+    """
+    
+    def __init__(
+        self,
+        gamma: float = 2.0,
+        alpha: Optional[torch.Tensor] = None,
+        ignore_index: int = 255,
+        reduction: str = 'mean'
+    ):
+        super().__init__()
+        self.gamma = gamma
+        self.alpha = alpha
+        self.ignore_index = ignore_index
+        self.reduction = reduction
+    
+    def forward(self, logits, labels):
+        B, C, H, W = logits.shape
+        
+        ce_loss = F.cross_entropy(
+            logits, labels,
+            weight=self.alpha,
+            ignore_index=self.ignore_index,
+            reduction='none'
+        )  # [B, H, W]
+        
+        logits_flat = logits.permute(0, 2, 3, 1).contiguous().view(-1, C)  # [B*H*W, C]
+        labels_flat = labels.view(-1)  # [B*H*W]
+        
+        valid_mask = (labels_flat != self.ignore_index)
+        
+        probs = F.softmax(logits_flat, dim=1)  # [B*H*W, C]
+        labels_clamped = labels_flat.clamp(0, C-1)  # clamp for safe indexing of ignored pixels
+        p_t = probs.gather(1, labels_clamped.unsqueeze(1)).squeeze(1)  # [B*H*W]
+        
+        focal_weight = (1 - p_t) ** self.gamma  # [B*H*W]
+        focal_weight = focal_weight.view(B, H, W)  # [B, H, W]
+        
+        focal_loss = focal_weight * ce_loss
+        
+        valid_mask_2d = (labels != self.ignore_index)
+        
+        if self.reduction == 'mean':
+            return focal_loss[valid_mask_2d].mean() if valid_mask_2d.sum() > 0 else focal_loss.sum() * 0.0
+        elif self.reduction == 'sum':
+            return focal_loss[valid_mask_2d].sum()
+        else:
+            return focal_loss
+
+
+class DiceLoss(nn.Module):
+    """
+    Multi-class Dice loss with optional class weighting.
+    """
+    
+    def __init__(
+        self,
+        num_classes: int = 21,
+        ignore_index: int = 255,
+        smooth: float = 1e-6,
+        class_weights: Optional[torch.Tensor] = None,
+        reduction: str = 'mean'
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.ignore_index = ignore_index
+        self.smooth = smooth
+        self.class_weights = class_weights
+        self.reduction = reduction
+    
+    def forward(self, logits, labels):
+        probs = F.softmax(logits, dim=1)  # [B, C, H, W]
+        valid_mask = (labels != self.ignore_index).float()  # [B, H, W]
+        
+        dice_loss = 0.0
+        valid_classes = 0
+        class_losses = []
+        
+        for c in range(self.num_classes):
+            target_c = (labels == c).float()  # [B, H, W]
+            pred_c = probs[:, c, :, :]  # [B, H, W]
+            
+            target_c = target_c * valid_mask
+            pred_c = pred_c * valid_mask
+            
+            intersection = (pred_c * target_c).sum()
+            union = pred_c.sum() + target_c.sum()
+            
+            if union > 0:
+                dice = (2.0 * intersection + self.smooth) / (union + self.smooth)
+                class_loss = 1.0 - dice
+                
+                if self.class_weights is not None:
+                    class_loss = class_loss * self.class_weights[c]
+                
+                class_losses.append(class_loss)
+                valid_classes += 1
+        
+        if valid_classes > 0:
+            if self.reduction == 'mean':
+                return torch.stack(class_losses).mean()
+            else:
+                return torch.stack(class_losses).sum()
+        
+        return logits.sum() * 0.0
+
+
+class CombinedSegmentationLoss(nn.Module):
+    """Weighted combination of CE/Focal, Dice, and Lovasz-Softmax losses."""
+    
+    def __init__(
+        self,
+        num_classes: int = 21,
+        ignore_index: int = 255,
+        weight_ce: float = 1.0,
+        weight_dice: float = 1.0,
+        weight_lovasz: float = 1.0,
+        use_focal: bool = True,
+        focal_gamma: float = 2.0,
+        class_weights: Optional[torch.Tensor] = None,
+    ):
+        super().__init__()
+        
+        self.weight_ce = weight_ce
+        self.weight_dice = weight_dice
+        self.weight_lovasz = weight_lovasz
+        self.use_focal = use_focal
+        
+        if use_focal:
+            self.ce_loss = FocalLoss(
+                gamma=focal_gamma,
+                alpha=class_weights,
+                ignore_index=ignore_index,
+            )
+        else:
+            self.ce_loss = nn.CrossEntropyLoss(
+                weight=class_weights,
+                ignore_index=ignore_index,
+            )
+        
+        self.dice_loss = DiceLoss(
+            num_classes=num_classes,
+            ignore_index=ignore_index,
+            class_weights=class_weights,
+        )
+        
+        self.lovasz_loss = LovaszSoftmaxLoss(
+            classes='present',
+            ignore_index=ignore_index,
+        )
+    
+    def forward(self, logits, labels):
+        loss_ce = self.ce_loss(logits, labels)
+        loss_dice = self.dice_loss(logits, labels)
+        loss_lovasz = self.lovasz_loss(logits, labels)
+        
+        loss_total = (
+            self.weight_ce * loss_ce +
+            self.weight_dice * loss_dice +
+            self.weight_lovasz * loss_lovasz
+        )
+        
+        return {
+            'loss_ce': loss_ce,
+            'loss_dice': loss_dice,
+            'loss_lovasz': loss_lovasz,
+            'loss_total': loss_total,
+        }
+

requirements.txt

@@ -0,0 +1,15 @@
+torch
+torchvision
+diffusers
+transformers
+accelerate
+einops
+pytorch-lightning
+torchmetrics
+segmentation-models-pytorch
+opencv-python-headless
+matplotlib
+Pillow
+PyYAML
+huggingface-hub
+safetensors