"""Shared FLUX feature machinery: load the pipeline, extract multi-timestep activations + concept-attention maps at native resolution, and cache them.""" import numpy as np import torch import torch.nn.functional as F from pathlib import Path from torchvision.transforms.functional import to_pil_image from flux_concept_attention import ( FluxWithConceptAttentionPipeline, FluxTransformer2DModelWithConceptAttention, ) def load_flux_pipeline(config, device="cuda", flux_model=None): """Load the FLUX.1-dev concept-attention pipeline. Returns (pipeline, transformer).""" flux_model = flux_model or config["model"]["flux_model"] transformer = FluxTransformer2DModelWithConceptAttention.from_pretrained( flux_model, subfolder="transformer", torch_dtype=torch.float16 ) pipeline = FluxWithConceptAttentionPipeline.from_pretrained( flux_model, transformer=transformer, torch_dtype=torch.float16 ).to(device) pipeline.set_progress_bar_config(disable=True) return pipeline, transformer class MultiTimestepFeatureCache: """On-disk cache of multi-timestep features + concept maps (one dir per image).""" def __init__(self, cache_dir: str): self.cache_dir = Path(cache_dir) self.cache_dir.mkdir(parents=True, exist_ok=True) def _image_dir(self, image_name: str) -> Path: return self.cache_dir / image_name def has_multi_timestep_features(self, image_name: str) -> bool: return (self._image_dir(image_name) / "multi_timestep_data.pt").exists() def save_multi_timestep_features(self, image_name: str, timestep_data: dict): d = self._image_dir(image_name) d.mkdir(parents=True, exist_ok=True) torch.save(timestep_data, d / "multi_timestep_data.pt") def load_multi_timestep_features(self, image_name: str, device="cuda"): d = self._image_dir(image_name) if not self.has_multi_timestep_features(image_name): return None # weights_only=False is safe here: we only load our own cached tensors. try: data = torch.load(d / "multi_timestep_data.pt", map_location=device, weights_only=False) except (RuntimeError, EOFError) as e: print(f"[CACHE ERROR] Corrupted cache for {image_name}: {e}") import shutil shutil.rmtree(d, ignore_errors=True) return None for timestep in data["features"]: data["features"][timestep]["single_features"] = [ f.to(device) for f in data["features"][timestep]["single_features"] ] if "concept_maps" not in data: data["concept_maps"] = {} for timestep in data["concept_maps"]: for concept in data["concept_maps"][timestep]: cmap = data["concept_maps"][timestep][concept] if hasattr(cmap, "to"): data["concept_maps"][timestep][concept] = cmap.to(device) elif isinstance(cmap, np.ndarray): data["concept_maps"][timestep][concept] = torch.from_numpy(cmap).to(device) return data class MultiTimestepFeatureExtractor: """Extract FLUX features at native resolution. ``extract_or_load`` does a resolution-aware cache check then extracts (training/ inference); ``extract_features`` returns a CPU dict for the extraction scripts to save. """ def __init__(self, pipeline, config, cache_dir: str, num_timesteps: int = 4, captions: dict = None): self.pipeline = pipeline self.config = config self.cache = MultiTimestepFeatureCache(cache_dir) self.num_timesteps = num_timesteps self.captions = captions or {} def _calculate_shift(self, image_seq_len, base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len return max(base_shift, min(max_shift, image_seq_len * m + b)) def _setup_scheduler_for_resolution(self, height: int, width: int): # FLUX uses 16x VAE downsampling; the scheduler shift depends on sequence length. image_seq_len = (height // 16) * (width // 16) mu = self._calculate_shift(image_seq_len) num_inference_steps = self.config.get("flux", {}).get("timesteps", 28) self.pipeline.scheduler.set_timesteps(num_inference_steps, mu=mu) return mu def _get_evenly_spaced_timesteps(self): all_timesteps = self.pipeline.scheduler.timesteps group_size = 7 indices = [-(i * group_size + 1) for i in range(self.num_timesteps) if i * group_size < len(all_timesteps)] selected = [int(all_timesteps[i]) for i in indices] selected = [t for t in selected if t != 1000][:self.num_timesteps] return sorted(selected, reverse=True) def _run_extraction(self, pil_image, image_name, concepts, resolution, prompt): """Core extraction loop. Returns a CPU-tensor timestep_data dict.""" height, width = resolution mu = self._setup_scheduler_for_resolution(height, width) selected_timesteps = self._get_evenly_spaced_timesteps() timestep_data = { "timesteps": selected_timesteps, "features": {}, "concept_maps": {}, "image_name": image_name, "concepts": concepts, "resolution": (height, width), "prompt": prompt, "mu": mu, } concept_attention_kwargs = { "concepts": concepts, "timesteps": self.config["flux"]["concept_timesteps"], "layers": self.config["flux"]["concept_layers"], } for timestep in selected_timesteps: with torch.no_grad(): output = self.pipeline( prompt=prompt, image=pil_image, height=height, width=width, timesteps=[timestep], num_inference_steps=1, guidance_scale=self.config["flux"]["guidance_scale"], concept_attention_kwargs=concept_attention_kwargs, generator=torch.Generator("cuda").manual_seed(42), output_type="latent", ) _, single_features = self.pipeline.transformer.get_features() concept_maps_raw = output.concept_attention_maps maps_list = ( concept_maps_raw[0] if (len(concept_maps_raw) == 1 and isinstance(concept_maps_raw[0], list)) else concept_maps_raw ) if len(maps_list) != len(concepts): raise ValueError(f"Mismatch: {len(concepts)} concepts vs {len(maps_list)} maps") concept_maps = {c: maps_list[i] for i, c in enumerate(concepts)} timestep_data["features"][timestep] = { "single_features": [f.cpu() for f in single_features] } timestep_data["concept_maps"][timestep] = { concept: (cmap.cpu() if hasattr(cmap, "cpu") else cmap) for concept, cmap in concept_maps.items() } return timestep_data def _to_device(self, timestep_data, device): for timestep in timestep_data["features"]: timestep_data["features"][timestep]["single_features"] = [ f.to(device) for f in timestep_data["features"][timestep]["single_features"] ] for timestep in timestep_data["concept_maps"]: for concept in timestep_data["concept_maps"][timestep]: cmap = timestep_data["concept_maps"][timestep][concept] if hasattr(cmap, "to"): timestep_data["concept_maps"][timestep][concept] = cmap.to(device) return timestep_data def extract_or_load(self, image, image_name, concepts, resolution, prompt=None, force=False): """Extract or load resolution-aware cached features; returns tensors on device.""" device = next(self.pipeline.transformer.parameters()).device height, width = resolution if not force and self.cache.has_multi_timestep_features(image_name): cached = self.cache.load_multi_timestep_features(image_name, device) if cached is not None and cached.get("resolution") == (height, width): return cached if prompt is None: prompt = self.captions.get(image_name, "") pil_image = to_pil_image(image[0].cpu()) timestep_data = self._run_extraction(pil_image, image_name, concepts, (height, width), prompt) self.cache.save_multi_timestep_features(image_name, timestep_data) return self._to_device(timestep_data, device) def extract_features(self, images, prompts, image_names, concepts, height, width): """Batch-style extraction (single image per call); returns a CPU dict to save.""" pil_image = to_pil_image(images[0].cpu()) return self._run_extraction(pil_image, image_names[0], concepts, (height, width), prompts[0]) def distribute_concepts(concept_maps, num_features, device): """Split concept maps roughly evenly across ``num_features`` feature layers.""" processed = [] for _, t in concept_maps.items(): if not hasattr(t, "to"): t = torch.from_numpy(t).float().to(device) elif t.device != device: t = t.to(device) if t.dim() == 2: t = t.view(1, 1, t.shape[0], t.shape[1]) elif t.dim() == 3: if t.shape[0] == 1: t = t.unsqueeze(1) elif t.shape[2] == 1: t = t.permute(2, 0, 1).unsqueeze(0) else: t = t.unsqueeze(1) processed.append(t) n = len(processed) per = n // num_features rem = n % num_features out = [] start = 0 for i in range(num_features): cnt = per + (1 if i < rem else 0) end = start + cnt if cnt > 0: out.append(torch.cat(processed[start:end], dim=1)) else: spatial_h, spatial_w = processed[0].shape[-2:] out.append(torch.zeros(1, 1, spatial_h, spatial_w, device=device)) start = end return out def distribute_concepts_across_layers(concept_maps, num_layers, target_size, device): """Depth variant of :func:`distribute_concepts`: resizes every map to ``target_size`` first and emits zero-channel tensors for empty layers.""" processed = [] for _, v in concept_maps.items(): if isinstance(v, np.ndarray): t = torch.from_numpy(v).float().to(device) elif hasattr(v, "to"): t = v.float().to(device) else: t = torch.tensor(v).float().to(device) if t.dim() == 2: t = t.view(1, 1, t.shape[0], t.shape[1]) elif t.dim() == 3: if t.shape[0] == 1: t = t.unsqueeze(1) elif t.shape[2] == 1: t = t.permute(2, 0, 1).unsqueeze(0) else: t = t.unsqueeze(1) if t.shape[-2:] != target_size: t = F.interpolate(t, size=target_size, mode="bilinear", align_corners=False) processed.append(t) num_concepts = len(processed) if num_concepts == 0: return [torch.zeros(1, 0, target_size[0], target_size[1], device=device) for _ in range(num_layers)] per_layer = num_concepts // num_layers remainder = num_concepts % num_layers distributed = [] start_idx = 0 for i in range(num_layers): count = per_layer + (1 if i < remainder else 0) end_idx = start_idx + count if count > 0: distributed.append(torch.cat(processed[start_idx:end_idx], dim=1)) else: distributed.append(torch.zeros(1, 0, target_size[0], target_size[1], device=device)) start_idx = end_idx return distributed