app.py

"""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().to(device).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().to(device) if isinstance(f, torch.Tensor) else f for f in val]
                elif isinstance(val, torch.Tensor):
                    timestep_data['features'][timestep][key] = val.float().to(device)

        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().to(device).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"
hf_token = os.environ.get("HF_TOKEN")
transformer = FluxTransformer2DModelWithConceptAttention.from_pretrained(
    flux_model, subfolder="transformer", torch_dtype=torch.float16, token=hf_token
)
pipeline = FluxWithConceptAttentionPipeline.from_pretrained(
    flux_model, transformer=transformer, torch_dtype=torch.float16, token=hf_token
).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()