Spaces:
Running on Zero
Running on Zero
| """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 ---- | |
| 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() |