"""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()