From 5e013d1a32b7bb36fdfe186eb4f7adc8e5ddc661 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 16:30:38 +0000 Subject: [PATCH 01/19] Add training-free high-res Flux community pipeline (HRDiT) --- benchmarks/benchmarking_flux_hrdit.py | 77 +++ examples/community/README.md | 23 + examples/community/pipeline_flux_hrdit.py | 534 ++++++++++++++++++ .../test_community_pipeline_flux_hrdit.py | 125 ++++ 4 files changed, 759 insertions(+) create mode 100644 benchmarks/benchmarking_flux_hrdit.py create mode 100644 examples/community/pipeline_flux_hrdit.py create mode 100644 tests/others/test_community_pipeline_flux_hrdit.py diff --git a/benchmarks/benchmarking_flux_hrdit.py b/benchmarks/benchmarking_flux_hrdit.py new file mode 100644 index 000000000000..4ae2199abfe6 --- /dev/null +++ b/benchmarks/benchmarking_flux_hrdit.py @@ -0,0 +1,77 @@ +"""Benchmark HRDiT progressive high-resolution generation against naive single-pass generation. + +HRDiT (https://arxiv.org/abs/2608.07003) targets two failure modes of off-the-shelf DiT models at high +resolution: spatial disorder (fixed by SPA, which this benchmark runs with) and long generation time +(addressed by HAP pruning and by the progressive 1024 -> 2048 -> 4096 ladder). This script times the +end-to-end pipeline call for the naive single-pass baseline and for the HRDiT stages. + +Run on a GPU machine with the FLUX.1-dev checkpoint available: + + python benchmarks/benchmarking_flux_hrdit.py --height 2048 --width 2048 +""" + +import argparse +from pathlib import Path + +import torch + +from benchmarking_utils import benchmark_fn, flush + +from diffusers import FluxPipeline +from diffusers.utils.testing_utils import torch_device + + +CKPT_ID = "black-forest-labs/FLUX.1-dev" +CUSTOM_PIPELINE_PATH = str(Path(__file__).resolve().parents[1] / "examples" / "community" / "pipeline_flux_hrdit.py") +RESULT_FILENAME = "flux_hrdit.csv" + + +def load_pipeline(): + return FluxPipeline.from_pretrained( + CKPT_ID, + torch_dtype=torch.bfloat16, + custom_pipeline=CUSTOM_PIPELINE_PATH, + ).to(torch_device) + + +def run_benchmarks(height, width, num_inference_steps, group_num, use_hap): + pipe = load_pipeline() + + settings = { + # Single-pass generation straight at the target resolution: SPA/HAP disabled, one stage. + "naive": dict(resolutions=[max(height, width)], group_num=1, use_hap=False), + # HRDiT: progressive ladder from 1024 up, SPA bundle averaging, optional HAP pruning. + "hrdit": dict(resolutions=None, group_num=group_num, use_hap=use_hap), + } + results = [] + for name, kwargs in settings.items(): + flush() + latency = benchmark_fn( + pipe, + "a photo of a mountain lake at dawn", + height=height, + width=width, + num_inference_steps=num_inference_steps, + **kwargs, + ) + max_memory = torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else float("nan") + results.append((name, latency, max_memory)) + print(f"{name:>6}: {latency:.3f}s, peak memory {max_memory:.2f} GiB") + + with open(RESULT_FILENAME, "w") as f: + f.write("setting,latency_s,peak_memory_gib\n") + for name, latency, max_memory in results: + f.write(f"{name},{latency},{max_memory}\n") + print(f"Results saved to {RESULT_FILENAME}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--height", type=int, default=2048) + parser.add_argument("--width", type=int, default=2048) + parser.add_argument("--num_inference_steps", type=int, default=28) + parser.add_argument("--group_num", type=int, default=4, help="number of SPA bundle variants") + parser.add_argument("--no_hap", action="store_true", help="disable head-adaptive attention pruning") + args = parser.parse_args() + + run_benchmarks(args.height, args.width, args.num_inference_steps, args.group_num, not args.no_hap) diff --git a/examples/community/README.md b/examples/community/README.md index 4ff9c4d77704..c14e875824d0 100644 --- a/examples/community/README.md +++ b/examples/community/README.md @@ -5630,3 +5630,26 @@ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") result.images[0].save(f"flux_fill_controlnet_inpaint_depth{timestamp}.jpg") ``` + +## Flux HRDiT + +Training-free high-resolution (up to 4096x4096) text-to-image generation with off-the-shelf Flux models, adapted from [HRDiT](https://arxiv.org/abs/2608.07003) ([official implementation](https://github.com/zylwithxy/HRDiT)). No fine-tuning and no new weights: the pipeline adds Spatial Position Alignment (position ids wrapped into the trained RoPE window and averaged over sliding bundle variants), optional head-adaptive attention pruning via FlexAttention, and a progressive 1024 -> 2048 -> 4096 generation ladder on top of the stock `FluxPipeline` denoise loop. + +```py +import torch +from diffusers import FluxPipeline + +pipe = FluxPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, custom_pipeline="pipeline_flux_hrdit" +).to("cuda") + +image = pipe( + "a photo of a mountain lake at dawn", + height=4096, + width=4096, + num_inference_steps=28, + group_num=4, # number of SPA bundle variants averaged per step + use_hap=True, # falls back to full attention if FlexAttention (torch >= 2.7) is unavailable +).images[0] +image.save("hrdit_4096.png") +``` diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py new file mode 100644 index 000000000000..64fbf5960a6a --- /dev/null +++ b/examples/community/pipeline_flux_hrdit.py @@ -0,0 +1,534 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on +# an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the +# specific language governing permissions and limitations under the License. + +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +import torch.nn.functional as F + +from diffusers.models.embeddings import apply_rotary_emb +from diffusers.models.transformers.transformer_flux import FluxAttnProcessor, _get_qkv_projections +from diffusers.pipelines.flux.pipeline_flux import FluxPipeline, calculate_shift, retrieve_timesteps +from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput +from diffusers.utils import logging, replace_example_docstring +from diffusers.utils.torch_utils import randn_tensor + + +try: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + FLEX_ATTENTION_AVAILABLE = True +except ImportError: + FLEX_ATTENTION_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import FluxPipeline + + >>> pipe = FluxPipeline.from_pretrained( + ... "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, custom_pipeline="pipeline_flux_hrdit" + ... ).to("cuda") + >>> image = pipe( + ... "a photo of a mountain lake at dawn", height=4096, width=4096, num_inference_steps=28 + ... ).images[0] + >>> image.save("hrdit_4096.png") + ``` +""" + + +# --------------------------------------------------------------------------------------- +# SPA: Spatial Position Alignment +# --------------------------------------------------------------------------------------- + + +def build_bundle_id_variants( + height: int, + width: int, + bundle_size: int = 64, + group_num: int = 4, + device=None, + dtype=None, +) -> List[torch.Tensor]: + """ + Spatial Position Alignment (SPA) position ids for a packed latent grid of `height` x `width` tokens. + + Off-the-shelf Flux models are trained at 1024x1024, i.e. a 64x64 packed latent grid, so the rotary position + ids never exceed 64 during training. Generating at higher resolutions produces out-of-range positions, which + manifests as the "spatial disorder" artifacts described in the HRDiT paper. + + SPA remaps every token's absolute grid coordinate into the trained RoPE window by wrapping it into sliding + `bundle_size` bundles. A single bundle partition would introduce seams at the bundle boundaries, so `group_num` + partitions with shifted origins ("bundle variants") are produced; the pipeline averages the transformer output + over the variants, which removes the seams. + + Adapted from HRDiT (https://arxiv.org/abs/2608.07003), `hrdit/spa.py::build_bundle_id_variants`. At or below + the trained resolution a single variant equal to the stock `FluxPipeline` ids is returned. + + Args: + height (`int`): Packed latent grid height (pixels / 16). + width (`int`): Packed latent grid width (pixels / 16). + bundle_size (`int`, defaults to 64): Side length of one bundle, i.e. the trained packed grid size. + group_num (`int`, defaults to 4): Number of sliding bundle variants. + device, dtype: Placed on / cast to the latent dtype, matching `_prepare_latent_image_ids`. + + Returns: + `List[torch.Tensor]` of shape `(height * width, 3)` each, one per bundle variant. + """ + if height <= bundle_size and width <= bundle_size: + return [FluxPipeline._prepare_latent_image_ids(1, height, width, device, dtype)] + + variants = [] + ys = torch.arange(height, device=device) + xs = torch.arange(width, device=device) + for variant in range(group_num): + # Slide the bundle partition origin; wrap coordinates into the trained RoPE window. + shift = (variant * bundle_size) // group_num + variant_ys = ((ys - shift) % height % bundle_size).to(dtype) + variant_xs = ((xs - shift) % width % bundle_size).to(dtype) + ids = torch.zeros(height, width, 3, device=device, dtype=dtype) + ids[..., 1] = variant_ys[:, None] + ids[..., 2] = variant_xs[None, :] + variants.append(ids.reshape(height * width, 3)) + return variants + + +def upsample_packed_latents(latents: torch.Tensor, old_grid: tuple, new_grid: tuple) -> torch.Tensor: + """ + Bilinearly upsample packed Flux latents (B, old_h * old_w, C * 4) to a `new_grid` (new_h, new_w) packed layout. + + Used between progressive generation stages: the previous stage's denoised latent is the structural prior for + the next, higher-resolution stage. + """ + batch_size, _, channels = latents.shape + num_channels = channels // 4 + old_grid_h, old_grid_w = int(old_grid[0]), int(old_grid[1]) + new_grid_h, new_grid_w = int(new_grid[0]), int(new_grid[1]) + + unpacked = latents.view(batch_size, old_grid_h, old_grid_w, num_channels, 2, 2) + unpacked = unpacked.permute(0, 3, 1, 4, 2, 5).reshape(batch_size, num_channels, old_grid_h * 2, old_grid_w * 2) + upsampled = F.interpolate( + unpacked.float(), size=(new_grid_h * 2, new_grid_w * 2), mode="bilinear", align_corners=False + ).to(latents.dtype) + return upsampled.view(batch_size, num_channels, new_grid_h, 2, new_grid_w, 2).permute(0, 2, 4, 1, 3, 5).reshape( + batch_size, new_grid_h * new_grid_w, channels + ) + + +# --------------------------------------------------------------------------------------- +# HAP: Head-adaptive Attention Pruning +# --------------------------------------------------------------------------------------- + + +def build_head_scope_plan(num_heads: int, window: int = 64, full_period: int = 4) -> torch.Tensor: + """ + Per-head attention scope plan for HAP as an int64 tensor of length `num_heads`. + + An entry of `-1` gives the head full (global) scope; a positive entry is the head's window radius in packed + latent-grid cells. Heads keep text keys inside their scope regardless. + + The paper reads a per-head scope plan from `configs/scope_plan_flux.json`; the checkpoint-specific plan is not + redistributable here, so this deterministic round-robin plan (every `full_period`-th head global, the rest + windowed) substitutes for it. + """ + plan = torch.full((num_heads,), window, dtype=torch.long) + plan[::full_period] = -1 + return plan + + +def build_mask_mod(pos_h: torch.Tensor, pos_w: torch.Tensor, windows: torch.Tensor, num_txt: int) -> Callable: + """ + Build a `flex_attention` mask_mod implementing the per-head scopes of HAP. + + `pos_h` / `pos_w` map an image token index to its packed grid coordinates; `windows` is the output of + [`build_head_scope_plan`]. Text keys and text queries always stay in scope. + """ + + def mask_mod(b, h, q_idx, kv_idx): + window = windows[h] + text_query = q_idx < num_txt + text_key = kv_idx < num_txt + img_kv = (kv_idx - num_txt).clamp(min=0) + img_q = (q_idx - num_txt).clamp(min=0) + row_dist = (pos_h[img_q] - pos_h[img_kv]).abs() + col_dist = (pos_w[img_q] - pos_w[img_kv]).abs() + return (window < 0) | text_query | text_key | ((row_dist <= window) & (col_dist <= window)) + + return mask_mod + + +class _HeadScopeState: + """ + Carries the current HAP scope plan and packed grid between the pipeline and the attention processors. + + The transformer blocks share one processor instance and do not see the grid directly, so the pipeline arms + this module-level state around the denoising loop and the processors read from it. + """ + + def __init__(self): + self.windows = None + self.grid_height = 0 + self.grid_width = 0 + self._block_masks: Dict[tuple, Any] = {} + + @property + def enabled(self): + return self.windows is not None + + def arm(self, windows: torch.Tensor): + self.windows = windows + self._block_masks = {} + + def set_grid(self, grid_height: int, grid_width: int): + self.grid_height = grid_height + self.grid_width = grid_width + self._block_masks = {} + + def disarm(self): + self.windows = None + self._block_masks = {} + + def get_block_mask(self, seq_len: int, num_txt: int, device): + key = (seq_len, num_txt, device) + if key in self._block_masks: + return self._block_masks[key] + + num_img = seq_len - num_txt + pos_h = torch.div(torch.arange(num_img, device=device), self.grid_width, rounding_mode="floor") + pos_w = torch.arange(num_img, device=device) % self.grid_width + mask_mod = build_mask_mod(pos_h, pos_w, self.windows.to(device), num_txt) + block_mask = create_block_mask(mask_mod, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=device) + self._block_masks[key] = block_mask + return block_mask + + +_HAP_STATE = _HeadScopeState() + + +class HRDiTFluxAttnProcessor(FluxAttnProcessor): + """ + Flux attention processor with head-adaptive attention pruning (HAP). + + Only the joint (double) blocks — the calls that pass `encoder_hidden_states` — take the pruned path; the + single blocks fall back to the stock processor, and so does everything when FlexAttention is unavailable. + """ + + def __call__( + self, + attn, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb=None, + ) -> torch.Tensor: + if not (_HAP_STATE.enabled and FLEX_ATTENTION_AVAILABLE) or encoder_hidden_states is None: + return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb) + return self._scoped_attention(attn, hidden_states, encoder_hidden_states, image_rotary_emb) + + def _scoped_attention(self, attn, hidden_states, encoder_hidden_states, image_rotary_emb): + query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( + attn, hidden_states, encoder_hidden_states + ) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + query = attn.norm_q(query) + key = attn.norm_k(key) + + encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) + encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) + encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) + encoder_query = attn.norm_added_q(encoder_query) + encoder_key = attn.norm_added_k(encoder_key) + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + num_txt = encoder_hidden_states.shape[1] + block_mask = _HAP_STATE.get_block_mask(key.shape[1], num_txt, query.device) + hidden_states = flex_attention( + query.transpose(1, 2).contiguous(), + key.transpose(1, 2).contiguous(), + value.transpose(1, 2).contiguous(), + block_mask=block_mask, + ).transpose(1, 2) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + + text_hidden_states, image_hidden_states = hidden_states.split_with_sizes( + [num_txt, hidden_states.shape[1] - num_txt], dim=1 + ) + image_hidden_states = attn.to_out[0](image_hidden_states.contiguous()) + image_hidden_states = attn.to_out[1](image_hidden_states) + text_hidden_states = attn.to_add_out(text_hidden_states.contiguous()) + return image_hidden_states, text_hidden_states + + +# --------------------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------------------- + + +class HRDiTFluxPipeline(FluxPipeline): + r""" + Training-free high-resolution (up to 4096x4096) text-to-image with off-the-shelf Flux models. + + Adapted from HRDiT, "Training-Free High-Resolution Image Generation with Off-the-Shelf Diffusion Transformer + Models" (https://arxiv.org/abs/2608.07003); reference implementation at https://github.com/zylwithxy/HRDiT. + + Three training-free pieces on top of the stock `FluxPipeline` denoise loop: + + - **SPA (Spatial Position Alignment)** — `build_bundle_id_variants` wraps high-resolution rotary position ids + into the trained 64x64 window across `group_num` sliding bundle variants; the transformer output is averaged + over the variants. The paper averages attention inside each attention layer; this pipeline averages at the + transformer output, which keeps the stock `FluxTransformer2DModel` untouched and needs no custom processor. + - **HAP (Head-adaptive Attention Pruning)** — `HRDiTFluxAttnProcessor` prunes attention outside each head's + scope via FlexAttention (torch >= 2.7). The paper's checkpoint-specific scope plan file is replaced by the + deterministic `build_head_scope_plan`; if FlexAttention is unavailable the processor falls back to full + attention and SPA alone remains active. + - **Progressive generation** — denoising climbs a resolution ladder (1024 -> 2048 -> 4096 by default), with the + previous stage's latent bilinearly upsampled and re-noised at the next stage's starting sigma. + + Args: + prompt (`str` or `List[str]`): The prompt to render. + height / width (`int`): Final output resolution. Defaults to 1024. + resolutions (`List[int]`, optional): Progressive resolution ladder (square side lengths). Defaults to + doubling from 1024 up to the target resolution. + group_num (`int`, defaults to 4): Number of SPA bundle variants averaged per step. + bundle_size (`int`, defaults to 64): Trained packed latent grid side (1024px for Flux). + use_hap (`bool`, defaults to True): Enable head-adaptive attention pruning when FlexAttention is available. + hap_window (`int`, defaults to 64): Window radius (packed grid cells) of the windowed heads. + stage_strength (`float`, defaults to 0.6): Fraction of the schedule each upsampled stage re-noises through. + + Example: see `EXAMPLE_DOC_STRING`. + """ + + def _resolution_ladder(self, height: int, width: int, resolutions: Optional[List[int]]) -> List[int]: + if resolutions is None: + target = max(height, width) + ladder = [] + side = min(1024, target) + while side < target: + ladder.append(side) + side = min(side * 2, target) + ladder.append(target) + return ladder + + ladder = [int(res) for res in resolutions] + if not ladder or any(ladder[i] >= ladder[i + 1] for i in range(len(ladder) - 1)): + raise ValueError(f"`resolutions` must be a non-empty, strictly increasing list, got {resolutions}.") + return ladder + + @staticmethod + def _stage_dimensions(height: int, width: int, target: int, side: int, quant: int) -> tuple: + if side >= target: + return height, width + stage_height = max(quant, int(round(height * side / target)) // quant * quant) + stage_width = max(quant, int(round(width * side / target)) // quant * quant) + return stage_height, stage_width + + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + resolutions: Optional[List[int]] = None, + group_num: int = 4, + bundle_size: int = 64, + use_hap: bool = True, + hap_window: int = 64, + stage_strength: float = 0.6, + num_inference_steps: int = 28, + guidance_scale: float = 3.5, + num_images_per_prompt: int = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + pooled_prompt_embeds: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + max_sequence_length: int = 512, + ): + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + quant = self.vae_scale_factor * 2 + if int(height) % quant != 0 or int(width) % quant != 0: + raise ValueError(f"`height` and `width` must be multiples of {quant}, got {height} and {width}.") + if not 0.0 < stage_strength <= 1.0: + raise ValueError(f"`stage_strength` must be in (0, 1], got {stage_strength}.") + + ladder = self._resolution_ladder(height, width, resolutions) + target = max(height, width) + + device = self._execution_device + dtype = prompt_embeds.dtype if prompt_embeds is not None else self.transformer.dtype + + # 1. Encode prompt (guidance-distilled models like FLUX.1-dev need no true CFG pass). + ( + prompt_embeds, + pooled_prompt_embeds, + text_ids, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=None, + ) + batch_size = prompt_embeds.shape[0] + + if self.transformer.config.guidance_embeds: + guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(batch_size) + else: + guidance = None + + if self._joint_attention_kwargs is None: + self._joint_attention_kwargs = {} + + # 2. Optionally arm HAP. + original_attn_processors = None + if use_hap: + if FLEX_ATTENTION_AVAILABLE: + original_attn_processors = dict(self.transformer.attn_processors) + self.transformer.set_attn_processor(HRDiTFluxAttnProcessor()) + num_heads = getattr(self.transformer.config, "num_attention_heads", 24) + _HAP_STATE.arm(build_head_scope_plan(num_heads, window=hap_window)) + else: + logger.warning( + "use_hap=True but FlexAttention is unavailable (needs torch >= 2.7); " + "falling back to full attention. SPA stays active." + ) + + try: + # 3. Progressive denoising over the resolution ladder. + all_sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + num_channels_latents = self.transformer.config.in_channels // 4 + latents = latents.to(device=device, dtype=dtype) if latents is not None else None + old_grid = None + self._num_timesteps = 0 + + for stage, side in enumerate(ladder): + stage_height, stage_width = self._stage_dimensions(height, width, target, side, quant) + latent_height = 2 * (stage_height // quant) + latent_width = 2 * (stage_width // quant) + grid_height, grid_width = latent_height // 2, latent_width // 2 + + if stage == 0: + if latents is not None and latents.shape[-2:] != (latent_height, latent_width): + raise ValueError( + f"Provided `latents` of spatial shape {latents.shape[-2:]} do not match the first " + f"progressive stage ({latent_height}, {latent_width})." + ) + if latents is None: + latents = randn_tensor( + (batch_size, num_channels_latents, latent_height, latent_width), + generator=generator, + device=device, + dtype=dtype, + ) + latents = self._pack_latents( + latents, batch_size, num_channels_latents, latent_height, latent_width + ) + stage_sigmas = all_sigmas + else: + latents = upsample_packed_latents(latents, old_grid, (grid_height, grid_width)) + # Later stages re-noise through the tail of the schedule (`stage_strength` of it). + num_stage_steps = max(1, int(round(num_inference_steps * stage_strength))) + stage_sigmas = all_sigmas[-num_stage_steps:] + old_grid = (grid_height, grid_width) + + image_id_variants = build_bundle_id_variants( + grid_height, grid_width, bundle_size=bundle_size, group_num=group_num, device=device, dtype=dtype + ) + + if _HAP_STATE.enabled: + _HAP_STATE.set_grid(grid_height, grid_width) + + mu = calculate_shift( + grid_height * grid_width, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, _ = retrieve_timesteps( + self.scheduler, len(stage_sigmas), device, sigmas=stage_sigmas, mu=mu + ) + self.scheduler.set_begin_index(0) + self._num_timesteps += len(timesteps) + + if stage > 0: + # Flow-match interpolation at the stage's (shift-adjusted) starting sigma. + start_sigma = float(self.scheduler.sigmas[0]) + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=dtype) + latents = (1.0 - start_sigma) * latents + start_sigma * noise + + with self.progress_bar( + total=len(timesteps), desc=f"HRDiT {stage_width}x{stage_height}" + ) as progress_bar: + for t in timesteps: + self._current_timestep = t + timestep = t.expand(latents.shape[0]).to(latents.dtype) + noise_pred = None + for image_ids in image_id_variants: + with self.transformer.cache_context("cond"): + variant_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + noise_pred = variant_pred if noise_pred is None else noise_pred + variant_pred + noise_pred = noise_pred / len(image_id_variants) + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + progress_bar.update() + + self._current_timestep = None + + if output_type == "latent": + image = latents + else: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + self.maybe_free_model_hooks() + finally: + if original_attn_processors is not None: + self.transformer.set_attn_processor(original_attn_processors) + _HAP_STATE.disarm() + + if not return_dict: + return (image,) + + return FluxPipelineOutput(images=image) diff --git a/tests/others/test_community_pipeline_flux_hrdit.py b/tests/others/test_community_pipeline_flux_hrdit.py new file mode 100644 index 000000000000..6c4d034cbb45 --- /dev/null +++ b/tests/others/test_community_pipeline_flux_hrdit.py @@ -0,0 +1,125 @@ +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on +# an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the +# specific language governing permissions and limitations under the License. + +import importlib.util +import sys +import unittest +from pathlib import Path + +import torch + +from diffusers.models.transformers.transformer_flux import FluxAttnProcessor +from diffusers.pipelines.flux.pipeline_flux import FluxPipeline + + +REPO_ROOT = Path(__file__).parents[2] +PIPELINE_PATH = REPO_ROOT / "examples" / "community" / "pipeline_flux_hrdit.py" +BENCHMARKS_DIR = REPO_ROOT / "benchmarks" + + +def _load_module(name, path, extra_sys_path=None): + if extra_sys_path is not None and str(extra_sys_path) not in sys.path: + sys.path.insert(0, str(extra_sys_path)) + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +hrdit = _load_module("pipeline_flux_hrdit", PIPELINE_PATH) + + +class BuildBundleIdVariantsTests(unittest.TestCase): + def test_single_variant_at_trained_resolution(self): + # Below the trained 64x64 packed grid, SPA must reduce to the stock Flux position ids. + variants = hrdit.build_bundle_id_variants(32, 64, bundle_size=64, group_num=4) + + self.assertEqual(len(variants), 1) + expected = FluxPipeline._prepare_latent_image_ids(1, 32, 64, None, None) + self.assertTrue(torch.equal(variants[0], expected)) + + def test_variants_wrap_into_trained_window(self): + variants = hrdit.build_bundle_id_variants(128, 96, bundle_size=64, group_num=4) + + self.assertEqual(len(variants), 4) + for variant in variants: + self.assertEqual(variant.shape, (128 * 96, 3)) + self.assertLessEqual(variant[:, 1].max().item(), 63) + self.assertLessEqual(variant[:, 2].max().item(), 63) + for i in range(4): + for j in range(i + 1, 4): + self.assertFalse(torch.equal(variants[i], variants[j])) + + def test_first_variant_is_unshifted_partition(self): + variants = hrdit.build_bundle_id_variants(128, 64, bundle_size=64, group_num=4) + + ys = (torch.arange(128) % 64).repeat_interleave(64) + self.assertTrue(torch.equal(variants[0][:, 1], ys.to(variants[0].dtype))) + + +class UpsamplePackedLatentsTests(unittest.TestCase): + def test_upsample_shape(self): + latents = torch.randn(2, 16 * 16, 64) + upsampled = hrdit.upsample_packed_latents(latents, (16, 16), (32, 32)) + + self.assertEqual(upsampled.shape, (2, 32 * 32, 64)) + + def test_upsample_is_identity_at_same_grid(self): + latents = torch.randn(1, 8 * 8, 64) + upsampled = hrdit.upsample_packed_latents(latents, (8, 8), (8, 8)) + + self.assertTrue(torch.allclose(upsampled, latents, atol=1e-5)) + + +class HeadScopeTests(unittest.TestCase): + def test_scope_plan_round_robin(self): + plan = hrdit.build_head_scope_plan(24, window=64, full_period=4) + + self.assertEqual(plan.shape, (24,)) + self.assertEqual((plan == -1).sum().item(), 6) + self.assertEqual((plan == 64).sum().item(), 18) + + def test_mask_mod_respects_scopes(self): + grid_height, grid_width, num_txt = 8, 8, 16 + num_img = grid_height * grid_width + pos_h = torch.div(torch.arange(num_img), grid_width, rounding_mode="floor") + pos_w = torch.arange(num_img) % grid_width + windows = hrdit.build_head_scope_plan(8, window=2, full_period=4) + mask_mod = hrdit.build_mask_mod(pos_h, pos_w, windows, num_txt) + + windowed_head = 1 + full_head = 0 + image_query = torch.tensor(num_txt) # grid position (0, 0) + near_key = torch.tensor(num_txt + 1) # grid position (0, 1) + far_key = torch.tensor(num_txt + 3 * grid_width) # grid position (3, 0) + text_key = torch.tensor(3) + + self.assertTrue(bool(mask_mod(0, windowed_head, image_query, text_key))) + self.assertTrue(bool(mask_mod(0, windowed_head, image_query, near_key))) + self.assertFalse(bool(mask_mod(0, windowed_head, image_query, far_key))) + self.assertTrue(bool(mask_mod(0, full_head, image_query, far_key))) + + +class PipelineIntegrationTests(unittest.TestCase): + def test_pipeline_subclasses_flux_pipeline(self): + self.assertTrue(issubclass(hrdit.HRDiTFluxPipeline, FluxPipeline)) + + def test_attention_processor_subclasses_stock_flux_processor(self): + self.assertTrue(issubclass(hrdit.HRDiTFluxAttnProcessor, FluxAttnProcessor)) + + def test_benchmark_wiring(self): + benchmark = _load_module( + "benchmarking_flux_hrdit", BENCHMARKS_DIR / "benchmarking_flux_hrdit.py", extra_sys_path=BENCHMARKS_DIR + ) + + self.assertEqual(benchmark.RESULT_FILENAME, "flux_hrdit.csv") + self.assertEqual(benchmark.CKPT_ID, "black-forest-labs/FLUX.1-dev") + self.assertTrue(callable(benchmark.run_benchmarks)) From 3343901a5252b96214866c9a4087382d9dc1ca9d Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 11:46:34 -0700 Subject: [PATCH 02/19] fix(hrdit): add __call__ docstring so replace_example_docstring doesn't crash on import --- examples/community/pipeline_flux_hrdit.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 64fbf5960a6a..529f436ddc2c 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -370,6 +370,15 @@ def __call__( return_dict: bool = True, max_sequence_length: int = 512, ): + r"""Generate a high-resolution image, training-free, with HRDiT (SPA + optional HAP). + + Accepts the standard [`FluxPipeline`] arguments plus SPA controls (`resolutions`, + `group_num`, `bundle_size`) and HAP / progressive-ladder controls (`use_hap`, + `hap_window`, `stage_strength`). `height` and `width` set the final resolution; the + pipeline renders progressively up to it. + + Examples: + """ height = height or self.default_sample_size * self.vae_scale_factor width = width or self.default_sample_size * self.vae_scale_factor quant = self.vae_scale_factor * 2 From 075d542470bcb8f3e0eeb6b4d09a5646cc754e5a Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 11:57:00 -0700 Subject: [PATCH 03/19] fix(hrdit): initialize _joint_attention_kwargs before use (read-before-assign at inference) --- examples/community/pipeline_flux_hrdit.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 529f436ddc2c..a81907a76c84 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -415,8 +415,7 @@ def __call__( else: guidance = None - if self._joint_attention_kwargs is None: - self._joint_attention_kwargs = {} + self._joint_attention_kwargs = {} # 2. Optionally arm HAP. original_attn_processors = None From 163b59754c00557f3d5980045d026c2f3d2ba332 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 11:59:17 -0700 Subject: [PATCH 04/19] fix(hrdit): use set_progress_bar_config for stage label (progress_bar takes no desc kwarg) --- examples/community/pipeline_flux_hrdit.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index a81907a76c84..4489faf74ff3 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -495,9 +495,8 @@ def __call__( noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=dtype) latents = (1.0 - start_sigma) * latents + start_sigma * noise - with self.progress_bar( - total=len(timesteps), desc=f"HRDiT {stage_width}x{stage_height}" - ) as progress_bar: + self.set_progress_bar_config(desc=f"HRDiT {stage_width}x{stage_height}") + with self.progress_bar(total=len(timesteps)) as progress_bar: for t in timesteps: self._current_timestep = t timestep = t.expand(latents.shape[0]).to(latents.dtype) From 85435bbef8b4ef95b752a3e28d2ea54ba870a57b Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 12:07:08 -0700 Subject: [PATCH 05/19] fix(hrdit): wrap __call__ in @torch.no_grad() and compile flex_attention (OOM: autograd graph + eager full-score materialization) --- examples/community/pipeline_flux_hrdit.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 4489faf74ff3..89f8f49c7972 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -217,6 +217,17 @@ def get_block_mask(self, seq_len: int, num_txt: int, device): _HAP_STATE = _HeadScopeState() +# flex_attention must be compiled to generate a fused, block-sparse kernel; the eager path +# materializes the full (B, H, S, S) score matrix and OOMs at high resolution. +_FLEX_ATTENTION_COMPILED = None + + +def _compiled_flex_attention(): + global _FLEX_ATTENTION_COMPILED + if _FLEX_ATTENTION_COMPILED is None: + _FLEX_ATTENTION_COMPILED = torch.compile(flex_attention, dynamic=False) + return _FLEX_ATTENTION_COMPILED + class HRDiTFluxAttnProcessor(FluxAttnProcessor): """ @@ -265,7 +276,7 @@ def _scoped_attention(self, attn, hidden_states, encoder_hidden_states, image_ro num_txt = encoder_hidden_states.shape[1] block_mask = _HAP_STATE.get_block_mask(key.shape[1], num_txt, query.device) - hidden_states = flex_attention( + hidden_states = _compiled_flex_attention()( query.transpose(1, 2).contiguous(), key.transpose(1, 2).contiguous(), value.transpose(1, 2).contiguous(), @@ -346,6 +357,7 @@ def _stage_dimensions(height: int, width: int, target: int, side: int, quant: in stage_width = max(quant, int(round(width * side / target)) // quant * quant) return stage_height, stage_width + @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, From 95e9d66c751fd4f67e6e702fcef6b0f6ea444b7f Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 12:41:19 -0700 Subject: [PATCH 06/19] =?UTF-8?q?fix(hrdit):=20faithful=20SPA=20re-port=20?= =?UTF-8?q?=E2=80=94=20monotonic=20bundle=20coarsening=20+=20in-attention?= =?UTF-8?q?=20variant=20averaging=20+=20proportional=20scale?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first draft mis-ported SPA two ways, producing periodic-tiled mush worse than naive high-res generation (caught on GPU validation @2048): - position map used modulo wrapping ((y-shift)%H)%bundle -> periodic tiling; the paper uses a monotonic bundle coarsening phi(x)=ceil((x+1-n1)/size). - averaging happened over 4 full model forwards at the output; the paper averages per-RoPE-variant attention *outputs* inside each layer (identity mean(softmax(A_n))@V == mean(softmax(A_n)@V)), one forward per step. Also adds the reference's proportional attention scale and corrects group_num default to 80. HAP pruning, NTK RoPE scaling and per-step SPA scheduling are documented as follow-ups (not yet ported). Ref: https://github.com/zylwithxy/HRDiT Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/community/pipeline_flux_hrdit.py | 384 +++++++++------------- 1 file changed, 152 insertions(+), 232 deletions(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 89f8f49c7972..32779e266009 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -9,7 +9,8 @@ # an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the # specific language governing permissions and limitations under the License. -from typing import Any, Callable, Dict, List, Optional, Union +import math +from typing import List, Optional, Union import numpy as np import torch @@ -23,16 +24,11 @@ from diffusers.utils.torch_utils import randn_tensor -try: - from torch.nn.attention.flex_attention import create_block_mask, flex_attention - - FLEX_ATTENTION_AVAILABLE = True -except ImportError: - FLEX_ATTENTION_AVAILABLE = False - - logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# FLUX is trained at 1024x1024 -> a 64x64 packed grid (4096 image tokens) plus 512 text tokens. +_TRAIN_SEQ_LEN = 64 ** 2 + 512 + EXAMPLE_DOC_STRING = """ Examples: ```py @@ -55,54 +51,51 @@ # --------------------------------------------------------------------------------------- -def build_bundle_id_variants( - height: int, - width: int, - bundle_size: int = 64, - group_num: int = 4, - device=None, - dtype=None, -) -> List[torch.Tensor]: - """ - Spatial Position Alignment (SPA) position ids for a packed latent grid of `height` x `width` tokens. +def _phi(x: torch.Tensor, n1: int, size: int) -> torch.Tensor: + """Bundle mapping: 0 for x < n1, else ceil((x + 1 - n1) / size). Monotonic non-decreasing.""" + return torch.where(x < n1, torch.zeros_like(x), (x + 1 - n1 + size - 1) // size) - Off-the-shelf Flux models are trained at 1024x1024, i.e. a 64x64 packed latent grid, so the rotary position - ids never exceed 64 during training. Generating at higher resolutions produces out-of-range positions, which - manifests as the "spatial disorder" artifacts described in the HRDiT paper. - SPA remaps every token's absolute grid coordinate into the trained RoPE window by wrapping it into sliding - `bundle_size` bundles. A single bundle partition would introduce seams at the bundle boundaries, so `group_num` - partitions with shifted origins ("bundle variants") are produced; the pipeline averages the transformer output - over the variants, which removes the seams. +def build_bundle_id_variants(img_ids: torch.Tensor, group_num: int) -> List[torch.Tensor]: + """ + Spatial Position Alignment (SPA) bundle-index variants of the packed-latent position ids ``img_ids``. + + Off-the-shelf FLUX is trained on a 64x64 packed grid, so its rotary position ids never exceed ~64. Generating + at higher resolution pushes ids out of the trained range, which is the "spatial disorder" the HRDiT paper + describes. SPA maps each token's grid coordinate into a small number of *bundles* via a monotonic (non-wrapping) + coarsening ``_phi`` -- many neighbouring tokens then share a position id inside the trained range. Because the + mapping is monotonic it introduces no periodic tiling; the residual bundle-boundary seams are averaged out by + sliding the boundary origin across ``group_num`` variants (see [`HRDiTFluxAttnProcessor`], which averages the + per-variant attention outputs -- element-wise identical to averaging the attention maps, at O(T*D) memory). - Adapted from HRDiT (https://arxiv.org/abs/2608.07003), `hrdit/spa.py::build_bundle_id_variants`. At or below - the trained resolution a single variant equal to the stock `FluxPipeline` ids is returned. + Adapted from HRDiT (https://arxiv.org/abs/2608.07003), ``hrdit/spa.py::build_bundle_id_variants``. Args: - height (`int`): Packed latent grid height (pixels / 16). - width (`int`): Packed latent grid width (pixels / 16). - bundle_size (`int`, defaults to 64): Side length of one bundle, i.e. the trained packed grid size. - group_num (`int`, defaults to 4): Number of sliding bundle variants. - device, dtype: Placed on / cast to the latent dtype, matching `_prepare_latent_image_ids`. + img_ids (`torch.Tensor`): Packed-latent position ids of shape `(T, 3)`; column 1 is the row index and + column 2 the column index (as produced by `FluxPipeline._prepare_latent_image_ids`). + group_num (`int`): Controls the bundle size `ceil(max_index / (group_num - 1))`; larger values give finer + bundles (more distinct positions, kept inside the trained window). Must be >= 2. Returns: - `List[torch.Tensor]` of shape `(height * width, 3)` each, one per bundle variant. + `List[torch.Tensor]` of shape `(T, 3)` each, one per sliding bundle-boundary variant. """ - if height <= bundle_size and width <= bundle_size: - return [FluxPipeline._prepare_latent_image_ids(1, height, width, device, dtype)] - - variants = [] - ys = torch.arange(height, device=device) - xs = torch.arange(width, device=device) - for variant in range(group_num): - # Slide the bundle partition origin; wrap coordinates into the trained RoPE window. - shift = (variant * bundle_size) // group_num - variant_ys = ((ys - shift) % height % bundle_size).to(dtype) - variant_xs = ((xs - shift) % width % bundle_size).to(dtype) - ids = torch.zeros(height, width, 3, device=device, dtype=dtype) - ids[..., 1] = variant_ys[:, None] - ids[..., 2] = variant_xs[None, :] - variants.append(ids.reshape(height * width, 3)) + if group_num < 2: + raise ValueError(f"`group_num` must be >= 2 for SPA, got {group_num}.") + + rows = img_ids[:, 1].long() + cols = img_ids[:, 2].long() + s_row = max(1, math.ceil(rows.max().item() / (group_num - 1))) + s_col = max(1, math.ceil(cols.max().item() / (group_num - 1))) + + def variant(n1_row: int, n1_col: int) -> torch.Tensor: + ids = img_ids.clone() + ids[:, 1] = _phi(rows, n1_row, s_row).to(img_ids.dtype) + ids[:, 2] = _phi(cols, n1_col, s_col).to(img_ids.dtype) + return ids + + variants = [variant(s_row, s_col)] + variants += [variant(n, s_col) for n in range(1, s_row)] + variants += [variant(s_row, m) for m in range(1, s_col)] return variants @@ -129,112 +122,46 @@ def upsample_packed_latents(latents: torch.Tensor, old_grid: tuple, new_grid: tu # --------------------------------------------------------------------------------------- -# HAP: Head-adaptive Attention Pruning +# SPA attention processor (averaging happens *inside* attention) # --------------------------------------------------------------------------------------- -def build_head_scope_plan(num_heads: int, window: int = 64, full_period: int = 4) -> torch.Tensor: - """ - Per-head attention scope plan for HAP as an int64 tensor of length `num_heads`. - - An entry of `-1` gives the head full (global) scope; a positive entry is the head's window radius in packed - latent-grid cells. Heads keep text keys inside their scope regardless. - - The paper reads a per-head scope plan from `configs/scope_plan_flux.json`; the checkpoint-specific plan is not - redistributable here, so this deterministic round-robin plan (every `full_period`-th head global, the rest - windowed) substitutes for it. +class _SPAState: """ - plan = torch.full((num_heads,), window, dtype=torch.long) - plan[::full_period] = -1 - return plan - - -def build_mask_mod(pos_h: torch.Tensor, pos_w: torch.Tensor, windows: torch.Tensor, num_txt: int) -> Callable: - """ - Build a `flex_attention` mask_mod implementing the per-head scopes of HAP. - - `pos_h` / `pos_w` map an image token index to its packed grid coordinates; `windows` is the output of - [`build_head_scope_plan`]. Text keys and text queries always stay in scope. - """ - - def mask_mod(b, h, q_idx, kv_idx): - window = windows[h] - text_query = q_idx < num_txt - text_key = kv_idx < num_txt - img_kv = (kv_idx - num_txt).clamp(min=0) - img_q = (q_idx - num_txt).clamp(min=0) - row_dist = (pos_h[img_q] - pos_h[img_kv]).abs() - col_dist = (pos_w[img_q] - pos_w[img_kv]).abs() - return (window < 0) | text_query | text_key | ((row_dist <= window) & (col_dist <= window)) - - return mask_mod - + Module-level carrier for the current SPA rotary-embedding variants. -class _HeadScopeState: - """ - Carries the current HAP scope plan and packed grid between the pipeline and the attention processors. - - The transformer blocks share one processor instance and do not see the grid directly, so the pipeline arms - this module-level state around the denoising loop and the processors read from it. + The transformer blocks share one processor instance and are not aware of SPA, so the pipeline precomputes the + per-variant rotary embeddings once per stage and arms them here; every processor reads from this state. """ def __init__(self): - self.windows = None - self.grid_height = 0 - self.grid_width = 0 - self._block_masks: Dict[tuple, Any] = {} + self.rope_variants = None # List[(cos, sin)] covering the full [text; image] sequence + self.proportional = True @property def enabled(self): - return self.windows is not None - - def arm(self, windows: torch.Tensor): - self.windows = windows - self._block_masks = {} + return self.rope_variants is not None - def set_grid(self, grid_height: int, grid_width: int): - self.grid_height = grid_height - self.grid_width = grid_width - self._block_masks = {} + def arm(self, rope_variants): + self.rope_variants = rope_variants def disarm(self): - self.windows = None - self._block_masks = {} - - def get_block_mask(self, seq_len: int, num_txt: int, device): - key = (seq_len, num_txt, device) - if key in self._block_masks: - return self._block_masks[key] - - num_img = seq_len - num_txt - pos_h = torch.div(torch.arange(num_img, device=device), self.grid_width, rounding_mode="floor") - pos_w = torch.arange(num_img, device=device) % self.grid_width - mask_mod = build_mask_mod(pos_h, pos_w, self.windows.to(device), num_txt) - block_mask = create_block_mask(mask_mod, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=device) - self._block_masks[key] = block_mask - return block_mask + self.rope_variants = None -_HAP_STATE = _HeadScopeState() - -# flex_attention must be compiled to generate a fused, block-sparse kernel; the eager path -# materializes the full (B, H, S, S) score matrix and OOMs at high resolution. -_FLEX_ATTENTION_COMPILED = None - - -def _compiled_flex_attention(): - global _FLEX_ATTENTION_COMPILED - if _FLEX_ATTENTION_COMPILED is None: - _FLEX_ATTENTION_COMPILED = torch.compile(flex_attention, dynamic=False) - return _FLEX_ATTENTION_COMPILED +_SPA_STATE = _SPAState() class HRDiTFluxAttnProcessor(FluxAttnProcessor): """ - Flux attention processor with head-adaptive attention pruning (HAP). - - Only the joint (double) blocks — the calls that pass `encoder_hidden_states` — take the pruned path; the - single blocks fall back to the stock processor, and so does everything when FlexAttention is unavailable. + Flux attention processor implementing HRDiT's Spatial Position Alignment (SPA). + + When SPA is armed (via [`_SPAState`]) the processor ignores the transformer's own rotary embedding and instead + runs attention once per bundle-index variant -- applying that variant's RoPE to the query/key -- then averages + the attention *outputs*. Since `mean_n(softmax(A_n)) @ V == mean_n(softmax(A_n) @ V)`, averaging the outputs is + element-wise identical to the paper's average-over-attention-maps, at O(T*D) memory instead of O(V*T^2). A + proportional attention scale ``sqrt(log_train(seq_len) / head_dim)`` compensates for the longer high-resolution + sequence. When SPA is disarmed the processor is exactly the stock `FluxAttnProcessor`. """ def __call__( @@ -245,11 +172,9 @@ def __call__( attention_mask: torch.Tensor | None = None, image_rotary_emb=None, ) -> torch.Tensor: - if not (_HAP_STATE.enabled and FLEX_ATTENTION_AVAILABLE) or encoder_hidden_states is None: + if not _SPA_STATE.enabled: return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb) - return self._scoped_attention(attn, hidden_states, encoder_hidden_states, image_rotary_emb) - def _scoped_attention(self, attn, hidden_states, encoder_hidden_states, image_rotary_emb): query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( attn, hidden_states, encoder_hidden_states ) @@ -260,38 +185,44 @@ def _scoped_attention(self, attn, hidden_states, encoder_hidden_states, image_ro query = attn.norm_q(query) key = attn.norm_k(key) - encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) - encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) - encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) - encoder_query = attn.norm_added_q(encoder_query) - encoder_key = attn.norm_added_k(encoder_key) - - query = torch.cat([encoder_query, query], dim=1) - key = torch.cat([encoder_key, key], dim=1) - value = torch.cat([encoder_value, value], dim=1) - - if image_rotary_emb is not None: - query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) - key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) - - num_txt = encoder_hidden_states.shape[1] - block_mask = _HAP_STATE.get_block_mask(key.shape[1], num_txt, query.device) - hidden_states = _compiled_flex_attention()( - query.transpose(1, 2).contiguous(), - key.transpose(1, 2).contiguous(), - value.transpose(1, 2).contiguous(), - block_mask=block_mask, - ).transpose(1, 2) - hidden_states = hidden_states.flatten(2, 3) - hidden_states = hidden_states.to(query.dtype) - - text_hidden_states, image_hidden_states = hidden_states.split_with_sizes( - [num_txt, hidden_states.shape[1] - num_txt], dim=1 - ) - image_hidden_states = attn.to_out[0](image_hidden_states.contiguous()) - image_hidden_states = attn.to_out[1](image_hidden_states) - text_hidden_states = attn.to_add_out(text_hidden_states.contiguous()) - return image_hidden_states, text_hidden_states + if encoder_hidden_states is not None: + encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)) + encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)) + encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)) + encoder_query = attn.norm_added_q(encoder_query) + encoder_key = attn.norm_added_k(encoder_key) + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + head_dim = query.shape[-1] + seq_len = query.shape[1] + if _SPA_STATE.proportional and seq_len > 1: + scale = math.sqrt(math.log(seq_len, _TRAIN_SEQ_LEN) / head_dim) + else: + scale = head_dim ** -0.5 + + value_t = value.transpose(1, 2).contiguous() # [B, H, S, D] + acc = None + for rope_variant in _SPA_STATE.rope_variants: + query_v = apply_rotary_emb(query, rope_variant, sequence_dim=1).transpose(1, 2).contiguous() + key_v = apply_rotary_emb(key, rope_variant, sequence_dim=1).transpose(1, 2).contiguous() + out = F.scaled_dot_product_attention(query_v, key_v, value_t, dropout_p=0.0, is_causal=False, scale=scale) + acc = out if acc is None else acc + out + hidden_states = (acc / len(_SPA_STATE.rope_variants)).transpose(1, 2) # [B, S, H, D] + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + if encoder_hidden_states is not None: + num_txt = encoder_hidden_states.shape[1] + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [num_txt, hidden_states.shape[1] - num_txt], dim=1 + ) + hidden_states = attn.to_out[0](hidden_states.contiguous()) + hidden_states = attn.to_out[1](hidden_states) + encoder_hidden_states = attn.to_add_out(encoder_hidden_states.contiguous()) + return hidden_states, encoder_hidden_states + return hidden_states # --------------------------------------------------------------------------------------- @@ -306,28 +237,29 @@ class HRDiTFluxPipeline(FluxPipeline): Adapted from HRDiT, "Training-Free High-Resolution Image Generation with Off-the-Shelf Diffusion Transformer Models" (https://arxiv.org/abs/2608.07003); reference implementation at https://github.com/zylwithxy/HRDiT. - Three training-free pieces on top of the stock `FluxPipeline` denoise loop: + Two training-free pieces on top of the stock `FluxPipeline` denoise loop: - - **SPA (Spatial Position Alignment)** — `build_bundle_id_variants` wraps high-resolution rotary position ids - into the trained 64x64 window across `group_num` sliding bundle variants; the transformer output is averaged - over the variants. The paper averages attention inside each attention layer; this pipeline averages at the - transformer output, which keeps the stock `FluxTransformer2DModel` untouched and needs no custom processor. - - **HAP (Head-adaptive Attention Pruning)** — `HRDiTFluxAttnProcessor` prunes attention outside each head's - scope via FlexAttention (torch >= 2.7). The paper's checkpoint-specific scope plan file is replaced by the - deterministic `build_head_scope_plan`; if FlexAttention is unavailable the processor falls back to full - attention and SPA alone remains active. - - **Progressive generation** — denoising climbs a resolution ladder (1024 -> 2048 -> 4096 by default), with the - previous stage's latent bilinearly upsampled and re-noised at the next stage's starting sigma. + - **SPA (Spatial Position Alignment)** -- `build_bundle_id_variants` maps high-resolution rotary position ids + into the trained ~64x64 window via a monotonic bundle coarsening (no wrapping, so no periodic tiling), across + several sliding bundle-boundary variants. `HRDiTFluxAttnProcessor` runs attention once per variant and averages + the outputs, with a proportional attention scale for the longer sequence. This is the core training-free fix + for high-resolution "spatial disorder". + - **Progressive generation** -- denoising climbs a resolution ladder (1024 -> 2048 -> 4096 by default); each + stage bilinearly upsamples the previous stage's latent and re-noises it through the tail of the schedule. SPA + is active only on the upscale stages (the base stage is in-distribution and uses stock RoPE). + + Not ported from the reference (documented follow-ups): the checkpoint-specific HAP head-scope pruning + (`configs/scope_plan_flux.json`), NTK-aware RoPE scaling, and per-step SPA scheduling. SPA here runs on every + step of every upscale stage and on every attention block. Args: prompt (`str` or `List[str]`): The prompt to render. height / width (`int`): Final output resolution. Defaults to 1024. resolutions (`List[int]`, optional): Progressive resolution ladder (square side lengths). Defaults to doubling from 1024 up to the target resolution. - group_num (`int`, defaults to 4): Number of SPA bundle variants averaged per step. - bundle_size (`int`, defaults to 64): Trained packed latent grid side (1024px for Flux). - use_hap (`bool`, defaults to True): Enable head-adaptive attention pruning when FlexAttention is available. - hap_window (`int`, defaults to 64): Window radius (packed grid cells) of the windowed heads. + group_num (`int`, defaults to 80): SPA bundle granularity; bundle size is `ceil(max_index / (group_num - 1))`. + Larger keeps more distinct positions inside the trained window. `group_num - 1` also bounds the number of + averaged variants. stage_strength (`float`, defaults to 0.6): Fraction of the schedule each upsampled stage re-noises through. Example: see `EXAMPLE_DOC_STRING`. @@ -366,10 +298,7 @@ def __call__( height: Optional[int] = None, width: Optional[int] = None, resolutions: Optional[List[int]] = None, - group_num: int = 4, - bundle_size: int = 64, - use_hap: bool = True, - hap_window: int = 64, + group_num: int = 80, stage_strength: float = 0.6, num_inference_steps: int = 28, guidance_scale: float = 3.5, @@ -382,12 +311,11 @@ def __call__( return_dict: bool = True, max_sequence_length: int = 512, ): - r"""Generate a high-resolution image, training-free, with HRDiT (SPA + optional HAP). + r"""Generate a high-resolution image, training-free, with HRDiT (SPA + progressive generation). - Accepts the standard [`FluxPipeline`] arguments plus SPA controls (`resolutions`, - `group_num`, `bundle_size`) and HAP / progressive-ladder controls (`use_hap`, - `hap_window`, `stage_strength`). `height` and `width` set the final resolution; the - pipeline renders progressively up to it. + Accepts the standard [`FluxPipeline`] arguments plus `resolutions` / `group_num` (SPA) and `stage_strength` + (progressive re-noising). `height` and `width` set the final resolution; the pipeline renders progressively + up to it. Examples: """ @@ -429,19 +357,9 @@ def __call__( self._joint_attention_kwargs = {} - # 2. Optionally arm HAP. - original_attn_processors = None - if use_hap: - if FLEX_ATTENTION_AVAILABLE: - original_attn_processors = dict(self.transformer.attn_processors) - self.transformer.set_attn_processor(HRDiTFluxAttnProcessor()) - num_heads = getattr(self.transformer.config, "num_attention_heads", 24) - _HAP_STATE.arm(build_head_scope_plan(num_heads, window=hap_window)) - else: - logger.warning( - "use_hap=True but FlexAttention is unavailable (needs torch >= 2.7); " - "falling back to full attention. SPA stays active." - ) + # 2. Install the SPA processor (armed per-stage below; disarmed => stock attention). + original_attn_processors = dict(self.transformer.attn_processors) + self.transformer.set_attn_processor(HRDiTFluxAttnProcessor()) try: # 3. Progressive denoising over the resolution ladder. @@ -481,12 +399,19 @@ def __call__( stage_sigmas = all_sigmas[-num_stage_steps:] old_grid = (grid_height, grid_width) - image_id_variants = build_bundle_id_variants( - grid_height, grid_width, bundle_size=bundle_size, group_num=group_num, device=device, dtype=dtype - ) - - if _HAP_STATE.enabled: - _HAP_STATE.set_grid(grid_height, grid_width) + image_ids = self._prepare_latent_image_ids(batch_size, grid_height, grid_width, device, dtype) + + # SPA on the upscale stages only: coarsen the out-of-range ids into the trained window and + # precompute the per-variant rotary embeddings over the full [text; image] sequence. + if grid_height > 64 or grid_width > 64: + variants = build_bundle_id_variants(image_ids, group_num) + rope_variants = [ + self.transformer.pos_embed(torch.cat([text_ids, variant_ids], dim=0)) + for variant_ids in variants + ] + _SPA_STATE.arm(rope_variants) + else: + _SPA_STATE.disarm() mu = calculate_shift( grid_height * grid_width, @@ -512,22 +437,18 @@ def __call__( for t in timesteps: self._current_timestep = t timestep = t.expand(latents.shape[0]).to(latents.dtype) - noise_pred = None - for image_ids in image_id_variants: - with self.transformer.cache_context("cond"): - variant_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = variant_pred if noise_pred is None else noise_pred + variant_pred - noise_pred = noise_pred / len(image_id_variants) + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] progress_bar.update() @@ -543,9 +464,8 @@ def __call__( self.maybe_free_model_hooks() finally: - if original_attn_processors is not None: - self.transformer.set_attn_processor(original_attn_processors) - _HAP_STATE.disarm() + _SPA_STATE.disarm() + self.transformer.set_attn_processor(original_attn_processors) if not return_dict: return (image,) From cfd820e5db822d898bb26e739d4dde5c806b29f3 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 13:38:19 -0700 Subject: [PATCH 07/19] feat(hrdit): NTK-aware RoPE scaling + step-gated SPA (the actual high-res mechanism) GPU validation of the SPA-only port: 2048 coherent but cross-hatched, 4096 washed out. Root cause vs reference inference.py: SPA is a *leading-few-steps* nudge (spa_steps=[3,0] -> 3 steps at 2048, none at 4096), while NTK-aware RoPE scaling (theta*=ntk_factor, [4,10] per stage) applied on *every* step is the primary high-res mechanism. Running SPA every step over-averaged -> washout. - flux_rope(): NTK RoPE (== diffusers FluxPosEmbed at factor 1). - _SPAState: base NTK rope on every step; SPA variants only while spa_active. - per-stage ntk_factor / spa_steps / guidance_scale_highres, gated per step. Remaining (documented): frequency-domain structure guidance + HAP pruning. Ref: https://github.com/zylwithxy/HRDiT (inference.py, hrdit/transformer.py) Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/community/pipeline_flux_hrdit.py | 156 ++++++++++++++++------ 1 file changed, 113 insertions(+), 43 deletions(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 32779e266009..68f432301c19 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -126,27 +126,58 @@ def upsample_packed_latents(latents: torch.Tensor, old_grid: tuple, new_grid: tu # --------------------------------------------------------------------------------------- +def flux_rope(ids: torch.Tensor, axes_dim, theta: float, ntk_factor: float = 1.0): + """ + Flux rotary embeddings for position ids ``ids`` [S, len(axes_dim)], with NTK-aware scaling. + + Identical to diffusers' `FluxPosEmbed` at ``ntk_factor == 1``; NTK scaling multiplies the RoPE base ``theta`` by + ``ntk_factor`` (HRDiT ``hrdit/transformer.py::get_1d_rotary_pos_embed``), which lowers every frequency and thereby + compresses out-of-range high-resolution positions back into the trained band -- the primary training-free + high-resolution mechanism (SPA only augments the leading steps). Returns ``(cos, sin)`` each `[S, sum(axes_dim)]`. + """ + scaled_theta = theta * ntk_factor + cos_out, sin_out = [], [] + for i, dim in enumerate(axes_dim): + pos = ids[:, i].to(torch.float64) + exps = torch.arange(0, dim, 2, dtype=torch.float64, device=ids.device)[: dim // 2] / dim + freqs = torch.outer(pos, 1.0 / (scaled_theta ** exps)) # [S, dim/2] + cos_out.append(freqs.cos().repeat_interleave(2, dim=1).float()) + sin_out.append(freqs.sin().repeat_interleave(2, dim=1).float()) + return torch.cat(cos_out, dim=-1), torch.cat(sin_out, dim=-1) + + class _SPAState: """ - Module-level carrier for the current SPA rotary-embedding variants. + Module-level carrier for the current stage's rotary embeddings. - The transformer blocks share one processor instance and are not aware of SPA, so the pipeline precomputes the - per-variant rotary embeddings once per stage and arms them here; every processor reads from this state. + The transformer blocks share one processor instance and are not aware of SPA/NTK, so the pipeline precomputes the + rotary embeddings once per stage and arms them here; every processor reads from this state. ``base_rope`` is the + single NTK-scaled RoPE used on every step; ``variant_ropes`` are the SPA bundle variants used only while + ``spa_active`` is set (the leading steps of a stage). ``spa_active`` is toggled per denoising step. """ def __init__(self): - self.rope_variants = None # List[(cos, sin)] covering the full [text; image] sequence + self.base_rope = None # (cos, sin) over the full [text; image] sequence + self.variant_ropes = None # List[(cos, sin)] SPA bundle variants + self.spa_active = False self.proportional = True @property def enabled(self): - return self.rope_variants is not None + return self.base_rope is not None - def arm(self, rope_variants): - self.rope_variants = rope_variants + def arm(self, base_rope, variant_ropes): + self.base_rope = base_rope + self.variant_ropes = variant_ropes + self.spa_active = False def disarm(self): - self.rope_variants = None + self.base_rope = None + self.variant_ropes = None + self.spa_active = False + + def current_ropes(self): + return self.variant_ropes if self.spa_active else [self.base_rope] _SPA_STATE = _SPAState() @@ -204,13 +235,14 @@ def __call__( scale = head_dim ** -0.5 value_t = value.transpose(1, 2).contiguous() # [B, H, S, D] + ropes = _SPA_STATE.current_ropes() acc = None - for rope_variant in _SPA_STATE.rope_variants: + for rope_variant in ropes: query_v = apply_rotary_emb(query, rope_variant, sequence_dim=1).transpose(1, 2).contiguous() key_v = apply_rotary_emb(key, rope_variant, sequence_dim=1).transpose(1, 2).contiguous() out = F.scaled_dot_product_attention(query_v, key_v, value_t, dropout_p=0.0, is_causal=False, scale=scale) acc = out if acc is None else acc + out - hidden_states = (acc / len(_SPA_STATE.rope_variants)).transpose(1, 2) # [B, S, H, D] + hidden_states = (acc / len(ropes)).transpose(1, 2) # [B, S, H, D] hidden_states = hidden_states.flatten(2, 3).to(query.dtype) if encoder_hidden_states is not None: @@ -237,20 +269,22 @@ class HRDiTFluxPipeline(FluxPipeline): Adapted from HRDiT, "Training-Free High-Resolution Image Generation with Off-the-Shelf Diffusion Transformer Models" (https://arxiv.org/abs/2608.07003); reference implementation at https://github.com/zylwithxy/HRDiT. - Two training-free pieces on top of the stock `FluxPipeline` denoise loop: + Three training-free pieces on top of the stock `FluxPipeline` denoise loop: - - **SPA (Spatial Position Alignment)** -- `build_bundle_id_variants` maps high-resolution rotary position ids - into the trained ~64x64 window via a monotonic bundle coarsening (no wrapping, so no periodic tiling), across - several sliding bundle-boundary variants. `HRDiTFluxAttnProcessor` runs attention once per variant and averages - the outputs, with a proportional attention scale for the longer sequence. This is the core training-free fix - for high-resolution "spatial disorder". + - **NTK-aware RoPE scaling** -- on every upscale-stage step the rotary base `theta` is multiplied by a per-stage + `ntk_factor`, compressing out-of-range high-resolution positions back into the trained band. This is the + primary high-resolution mechanism (`flux_rope`). + - **SPA (Spatial Position Alignment)** -- for the leading `spa_steps` steps of a stage only, `build_bundle_id_variants` + additionally coarsens position ids into the trained window via a monotonic bundle map (no wrapping, so no + periodic tiling) across sliding boundary variants; `HRDiTFluxAttnProcessor` runs attention once per variant and + averages the outputs, with a proportional attention scale. A light early-step correction for "spatial disorder". - **Progressive generation** -- denoising climbs a resolution ladder (1024 -> 2048 -> 4096 by default); each - stage bilinearly upsamples the previous stage's latent and re-noises it through the tail of the schedule. SPA - is active only on the upscale stages (the base stage is in-distribution and uses stock RoPE). + stage bilinearly upsamples the previous stage's latent and re-noises it through the tail of the schedule. The + base stage is in-distribution and uses stock RoPE. Not ported from the reference (documented follow-ups): the checkpoint-specific HAP head-scope pruning - (`configs/scope_plan_flux.json`), NTK-aware RoPE scaling, and per-step SPA scheduling. SPA here runs on every - step of every upscale stage and on every attention block. + (`configs/scope_plan_flux.json`) and the frequency-domain structure guidance (Butterworth low-pass + alpha/beta + blending in the flow-match step, swin patchify) that further stabilizes the highest stage. Args: prompt (`str` or `List[str]`): The prompt to render. @@ -258,9 +292,13 @@ class HRDiTFluxPipeline(FluxPipeline): resolutions (`List[int]`, optional): Progressive resolution ladder (square side lengths). Defaults to doubling from 1024 up to the target resolution. group_num (`int`, defaults to 80): SPA bundle granularity; bundle size is `ceil(max_index / (group_num - 1))`. - Larger keeps more distinct positions inside the trained window. `group_num - 1` also bounds the number of - averaged variants. + Larger keeps more distinct positions inside the trained window. + ntk_factor (`List[float]`, optional): Per-upscale-stage NTK RoPE-base multiplier. Defaults to `[4.0, 10.0]` + (2048, 4096), extended by its last value for further stages. + spa_steps (`List[int]`, optional): Per-upscale-stage count of leading steps that use SPA. Defaults to + `[3, 0]` -- a light nudge at 2048, none at 4096 (NTK alone). stage_strength (`float`, defaults to 0.6): Fraction of the schedule each upsampled stage re-noises through. + guidance_scale_highres (`List[float]`, optional): Per-upscale-stage guidance. Defaults to `[4.5, 6.0]`. Example: see `EXAMPLE_DOC_STRING`. """ @@ -299,9 +337,12 @@ def __call__( width: Optional[int] = None, resolutions: Optional[List[int]] = None, group_num: int = 80, + ntk_factor: Optional[List[float]] = None, + spa_steps: Optional[List[int]] = None, stage_strength: float = 0.6, num_inference_steps: int = 28, guidance_scale: float = 3.5, + guidance_scale_highres: Optional[List[float]] = None, num_images_per_prompt: int = 1, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.Tensor] = None, @@ -311,11 +352,11 @@ def __call__( return_dict: bool = True, max_sequence_length: int = 512, ): - r"""Generate a high-resolution image, training-free, with HRDiT (SPA + progressive generation). + r"""Generate a high-resolution image, training-free, with HRDiT (NTK RoPE + SPA + progressive generation). - Accepts the standard [`FluxPipeline`] arguments plus `resolutions` / `group_num` (SPA) and `stage_strength` - (progressive re-noising). `height` and `width` set the final resolution; the pipeline renders progressively - up to it. + Accepts the standard [`FluxPipeline`] arguments plus `resolutions` / `group_num` / `ntk_factor` / `spa_steps` + and `stage_strength` / `guidance_scale_highres`. `height` and `width` set the final resolution; the pipeline + renders progressively up to it. Examples: """ @@ -330,6 +371,17 @@ def __call__( ladder = self._resolution_ladder(height, width, resolutions) target = max(height, width) + # Per-upscale-stage schedules (index j = ladder stage - 1). NTK-aware RoPE scaling is applied on every step + # and is the primary high-resolution mechanism; SPA augments only the leading `spa_steps` steps of a stage. + ntk_schedule = ntk_factor if ntk_factor is not None else [4.0, 10.0] + spa_schedule = spa_steps if spa_steps is not None else [3, 0] + guidance_hr = guidance_scale_highres if guidance_scale_highres is not None else [4.5, 6.0] + + def _stage_value(schedule, j, fill): + if not schedule: + return fill + return schedule[j] if j < len(schedule) else schedule[-1] + device = self._execution_device dtype = prompt_embeds.dtype if prompt_embeds is not None else self.transformer.dtype @@ -350,16 +402,20 @@ def __call__( ) batch_size = prompt_embeds.shape[0] - if self.transformer.config.guidance_embeds: - guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(batch_size) - else: - guidance = None + guidance_embeds = self.transformer.config.guidance_embeds + + def _guidance(scale): + if not guidance_embeds: + return None + return torch.full([1], scale, device=device, dtype=torch.float32).expand(batch_size) self._joint_attention_kwargs = {} - # 2. Install the SPA processor (armed per-stage below; disarmed => stock attention). + # 2. Install the SPA/NTK processor (armed per-stage below; disarmed => stock attention). original_attn_processors = dict(self.transformer.attn_processors) self.transformer.set_attn_processor(HRDiTFluxAttnProcessor()) + axes_dim = self.transformer.pos_embed.axes_dim + rope_theta = self.transformer.pos_embed.theta try: # 3. Progressive denoising over the resolution ladder. @@ -401,17 +457,30 @@ def __call__( image_ids = self._prepare_latent_image_ids(batch_size, grid_height, grid_width, device, dtype) - # SPA on the upscale stages only: coarsen the out-of-range ids into the trained window and - # precompute the per-variant rotary embeddings over the full [text; image] sequence. - if grid_height > 64 or grid_width > 64: - variants = build_bundle_id_variants(image_ids, group_num) - rope_variants = [ - self.transformer.pos_embed(torch.cat([text_ids, variant_ids], dim=0)) - for variant_ids in variants - ] - _SPA_STATE.arm(rope_variants) - else: + # Upscale stages: NTK-scaled RoPE on every step (the primary high-res mechanism), with SPA bundle + # variants precomputed for the leading `stage_spa_steps` steps only. The base stage is in-distribution + # and uses stock RoPE (state disarmed). + if stage == 0: _SPA_STATE.disarm() + stage_spa_steps = 0 + stage_guidance = _guidance(guidance_scale) + else: + j = stage - 1 + stage_ntk = float(_stage_value(ntk_schedule, j, 1.0)) + stage_spa_steps = int(_stage_value(spa_schedule, j, 0)) + stage_guidance = _guidance(float(_stage_value(guidance_hr, j, guidance_scale))) + base_rope = flux_rope( + torch.cat([text_ids, image_ids], dim=0), axes_dim, rope_theta, ntk_factor=stage_ntk + ) + if stage_spa_steps > 0: + variants = build_bundle_id_variants(image_ids, group_num) + variant_ropes = [ + flux_rope(torch.cat([text_ids, v], dim=0), axes_dim, rope_theta, ntk_factor=stage_ntk) + for v in variants + ] + else: + variant_ropes = [base_rope] + _SPA_STATE.arm(base_rope, variant_ropes) mu = calculate_shift( grid_height * grid_width, @@ -434,14 +503,15 @@ def __call__( self.set_progress_bar_config(desc=f"HRDiT {stage_width}x{stage_height}") with self.progress_bar(total=len(timesteps)) as progress_bar: - for t in timesteps: + for i, t in enumerate(timesteps): self._current_timestep = t + _SPA_STATE.spa_active = i < stage_spa_steps timestep = t.expand(latents.shape[0]).to(latents.dtype) with self.transformer.cache_context("cond"): noise_pred = self.transformer( hidden_states=latents, timestep=timestep / 1000, - guidance=guidance, + guidance=stage_guidance, pooled_projections=pooled_prompt_embeds, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, From 5291b84608d0da1c00eae647adb59c4de2443488 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 14:18:22 -0700 Subject: [PATCH 08/19] =?UTF-8?q?feat(hrdit):=20high-fidelity=20port=20?= =?UTF-8?q?=E2=80=94=20frequency-domain=20structure=20guidance=20for=20the?= =?UTF-8?q?=20top=20stage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the DemoFusion-style structure guidance the reference relies on at high res, so 4096 stops drifting to a washed-out mean: - shared flow-match schedule across stages (pred_x0 references align per timestep) - per-stage prior: decode -> bicubic upscale -> sharpen -> re-encode - custom flowmatch step: per-step low-frequency injection from the upsampled previous-stage pred_x0 (alpha, Butterworth FFT split) + velocity momentum (beta) Keeps validated NTK RoPE (every step) + gated SPA (leading steps). Reference hyperparameters wired: ntk=[4,10], spa_steps=[3,0], steps=[17,10], guidance=[4.5,6], alphas=[1.0,0.25], betas=[0.5,0.5], filter_ratio=0.2. Component-verified offline vs diffusers: flux_rope==FluxPosEmbed, pack/unpack dims, scheduler tail-index alignment, bundle variants in-range. Deferred: HAP, swin, DWT. Ref: https://github.com/zylwithxy/HRDiT (pipeline.py flowmatch_step, inference.py) Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/community/pipeline_flux_hrdit.py | 484 +++++++++++++--------- 1 file changed, 291 insertions(+), 193 deletions(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 68f432301c19..d40a83f47199 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -24,6 +24,14 @@ from diffusers.utils.torch_utils import randn_tensor +try: + from torchvision.transforms.functional import gaussian_blur as _tv_gaussian_blur + + _TORCHVISION_AVAILABLE = True +except ImportError: + _TORCHVISION_AVAILABLE = False + + logger = logging.get_logger(__name__) # pylint: disable=invalid-name # FLUX is trained at 1024x1024 -> a 64x64 packed grid (4096 image tokens) plus 512 text tokens. @@ -39,7 +47,7 @@ ... "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, custom_pipeline="pipeline_flux_hrdit" ... ).to("cuda") >>> image = pipe( - ... "a photo of a mountain lake at dawn", height=4096, width=4096, num_inference_steps=28 + ... "a photo of a mountain lake at dawn", height=4096, width=4096 ... ).images[0] >>> image.save("hrdit_4096.png") ``` @@ -47,7 +55,7 @@ # --------------------------------------------------------------------------------------- -# SPA: Spatial Position Alignment +# SPA: Spatial Position Alignment + NTK-aware RoPE # --------------------------------------------------------------------------------------- @@ -60,21 +68,14 @@ def build_bundle_id_variants(img_ids: torch.Tensor, group_num: int) -> List[torc """ Spatial Position Alignment (SPA) bundle-index variants of the packed-latent position ids ``img_ids``. - Off-the-shelf FLUX is trained on a 64x64 packed grid, so its rotary position ids never exceed ~64. Generating - at higher resolution pushes ids out of the trained range, which is the "spatial disorder" the HRDiT paper - describes. SPA maps each token's grid coordinate into a small number of *bundles* via a monotonic (non-wrapping) - coarsening ``_phi`` -- many neighbouring tokens then share a position id inside the trained range. Because the - mapping is monotonic it introduces no periodic tiling; the residual bundle-boundary seams are averaged out by - sliding the boundary origin across ``group_num`` variants (see [`HRDiTFluxAttnProcessor`], which averages the - per-variant attention outputs -- element-wise identical to averaging the attention maps, at O(T*D) memory). - - Adapted from HRDiT (https://arxiv.org/abs/2608.07003), ``hrdit/spa.py::build_bundle_id_variants``. + Maps each token's grid coordinate into a small number of bundles via a monotonic (non-wrapping) coarsening + ``_phi`` -- so many neighbouring tokens share a position id inside the trained window, with no periodic tiling. + Residual bundle-boundary seams are averaged out over ``group_num`` sliding-origin variants. Adapted from HRDiT + (https://arxiv.org/abs/2608.07003), ``hrdit/spa.py``. Args: - img_ids (`torch.Tensor`): Packed-latent position ids of shape `(T, 3)`; column 1 is the row index and - column 2 the column index (as produced by `FluxPipeline._prepare_latent_image_ids`). - group_num (`int`): Controls the bundle size `ceil(max_index / (group_num - 1))`; larger values give finer - bundles (more distinct positions, kept inside the trained window). Must be >= 2. + img_ids (`torch.Tensor`): Packed-latent position ids `(T, 3)`; column 1 row index, column 2 column index. + group_num (`int`): Bundle granularity; bundle size is `ceil(max_index / (group_num - 1))`. Must be >= 2. Returns: `List[torch.Tensor]` of shape `(T, 3)` each, one per sliding bundle-boundary variant. @@ -99,41 +100,14 @@ def variant(n1_row: int, n1_col: int) -> torch.Tensor: return variants -def upsample_packed_latents(latents: torch.Tensor, old_grid: tuple, new_grid: tuple) -> torch.Tensor: - """ - Bilinearly upsample packed Flux latents (B, old_h * old_w, C * 4) to a `new_grid` (new_h, new_w) packed layout. - - Used between progressive generation stages: the previous stage's denoised latent is the structural prior for - the next, higher-resolution stage. - """ - batch_size, _, channels = latents.shape - num_channels = channels // 4 - old_grid_h, old_grid_w = int(old_grid[0]), int(old_grid[1]) - new_grid_h, new_grid_w = int(new_grid[0]), int(new_grid[1]) - - unpacked = latents.view(batch_size, old_grid_h, old_grid_w, num_channels, 2, 2) - unpacked = unpacked.permute(0, 3, 1, 4, 2, 5).reshape(batch_size, num_channels, old_grid_h * 2, old_grid_w * 2) - upsampled = F.interpolate( - unpacked.float(), size=(new_grid_h * 2, new_grid_w * 2), mode="bilinear", align_corners=False - ).to(latents.dtype) - return upsampled.view(batch_size, num_channels, new_grid_h, 2, new_grid_w, 2).permute(0, 2, 4, 1, 3, 5).reshape( - batch_size, new_grid_h * new_grid_w, channels - ) - - -# --------------------------------------------------------------------------------------- -# SPA attention processor (averaging happens *inside* attention) -# --------------------------------------------------------------------------------------- - - def flux_rope(ids: torch.Tensor, axes_dim, theta: float, ntk_factor: float = 1.0): """ Flux rotary embeddings for position ids ``ids`` [S, len(axes_dim)], with NTK-aware scaling. Identical to diffusers' `FluxPosEmbed` at ``ntk_factor == 1``; NTK scaling multiplies the RoPE base ``theta`` by - ``ntk_factor`` (HRDiT ``hrdit/transformer.py::get_1d_rotary_pos_embed``), which lowers every frequency and thereby - compresses out-of-range high-resolution positions back into the trained band -- the primary training-free - high-resolution mechanism (SPA only augments the leading steps). Returns ``(cos, sin)`` each `[S, sum(axes_dim)]`. + ``ntk_factor`` (HRDiT ``hrdit/transformer.py::get_1d_rotary_pos_embed``), lowering every frequency and thereby + compressing out-of-range high-resolution positions back into the trained band -- the primary training-free + high-resolution mechanism. Returns ``(cos, sin)`` each `[S, sum(axes_dim)]`. """ scaled_theta = theta * ntk_factor cos_out, sin_out = [], [] @@ -146,19 +120,44 @@ def flux_rope(ids: torch.Tensor, axes_dim, theta: float, ntk_factor: float = 1.0 return torch.cat(cos_out, dim=-1), torch.cat(sin_out, dim=-1) +def butterworth_low_pass_filter_2d(height: int, width: int, ratio: float, device, order: int = 4) -> torch.Tensor: + """Centered 2D Butterworth low-pass mask `[1, 1, H, W]` for frequency-domain structure guidance.""" + if ratio <= 0: + return torch.zeros(1, 1, height, width, device=device) + yy = (2.0 * torch.arange(height, device=device) / height - 1.0).view(height, 1) + xx = (2.0 * torch.arange(width, device=device) / width - 1.0).view(1, width) + d_square = yy ** 2 + xx ** 2 + mask = 1.0 / (1.0 + (d_square / ratio ** 2) ** order) + return mask.view(1, 1, height, width) + + +def split_low_freq(x: torch.Tensor, freq_filter: torch.Tensor) -> torch.Tensor: + """Low-frequency component of ``x`` [B, C, H, W] under a centered ``freq_filter`` (real output).""" + x_freq = torch.fft.fftshift(torch.fft.fft2(x.to(freq_filter.dtype))) + x_low = x_freq * freq_filter + return torch.fft.ifft2(torch.fft.ifftshift(x_low)).real + + +def sharpen(image: torch.Tensor, kernel_size: int = 3, sigma: float = 1.0, alpha: float = 1.0) -> torch.Tensor: + """Unsharp-mask sharpening of an image tensor; no-op if torchvision is unavailable.""" + if not _TORCHVISION_AVAILABLE: + return image + blurred = _tv_gaussian_blur(image, kernel_size=[kernel_size, kernel_size], sigma=[sigma, sigma]) + return (alpha + 1.0) * image - alpha * blurred + + class _SPAState: """ Module-level carrier for the current stage's rotary embeddings. The transformer blocks share one processor instance and are not aware of SPA/NTK, so the pipeline precomputes the - rotary embeddings once per stage and arms them here; every processor reads from this state. ``base_rope`` is the - single NTK-scaled RoPE used on every step; ``variant_ropes`` are the SPA bundle variants used only while - ``spa_active`` is set (the leading steps of a stage). ``spa_active`` is toggled per denoising step. + rotary embeddings once per stage and arms them here. ``base_rope`` is the single NTK-scaled RoPE used on every + step; ``variant_ropes`` are the SPA bundle variants used only while ``spa_active`` is set (the leading steps). """ def __init__(self): - self.base_rope = None # (cos, sin) over the full [text; image] sequence - self.variant_ropes = None # List[(cos, sin)] SPA bundle variants + self.base_rope = None + self.variant_ropes = None self.spa_active = False self.proportional = True @@ -185,14 +184,12 @@ def current_ropes(self): class HRDiTFluxAttnProcessor(FluxAttnProcessor): """ - Flux attention processor implementing HRDiT's Spatial Position Alignment (SPA). - - When SPA is armed (via [`_SPAState`]) the processor ignores the transformer's own rotary embedding and instead - runs attention once per bundle-index variant -- applying that variant's RoPE to the query/key -- then averages - the attention *outputs*. Since `mean_n(softmax(A_n)) @ V == mean_n(softmax(A_n) @ V)`, averaging the outputs is - element-wise identical to the paper's average-over-attention-maps, at O(T*D) memory instead of O(V*T^2). A - proportional attention scale ``sqrt(log_train(seq_len) / head_dim)`` compensates for the longer high-resolution - sequence. When SPA is disarmed the processor is exactly the stock `FluxAttnProcessor`. + Flux attention processor implementing HRDiT's NTK RoPE + Spatial Position Alignment (SPA). + + When armed the processor ignores the transformer's own rotary embedding and uses the NTK-scaled RoPE from + [`_SPAState`] on every step; on the leading SPA steps it instead runs attention once per bundle variant and + averages the outputs (`mean_n(softmax(A_n)) @ V == mean_n(softmax(A_n) @ V)`, so O(T*D) memory), with a + proportional attention scale for the longer sequence. When disarmed it is exactly the stock `FluxAttnProcessor`. """ def __call__( @@ -269,36 +266,37 @@ class HRDiTFluxPipeline(FluxPipeline): Adapted from HRDiT, "Training-Free High-Resolution Image Generation with Off-the-Shelf Diffusion Transformer Models" (https://arxiv.org/abs/2608.07003); reference implementation at https://github.com/zylwithxy/HRDiT. - Three training-free pieces on top of the stock `FluxPipeline` denoise loop: + Training-free pieces on top of the stock `FluxPipeline` denoise loop: - - **NTK-aware RoPE scaling** -- on every upscale-stage step the rotary base `theta` is multiplied by a per-stage - `ntk_factor`, compressing out-of-range high-resolution positions back into the trained band. This is the - primary high-resolution mechanism (`flux_rope`). - - **SPA (Spatial Position Alignment)** -- for the leading `spa_steps` steps of a stage only, `build_bundle_id_variants` - additionally coarsens position ids into the trained window via a monotonic bundle map (no wrapping, so no - periodic tiling) across sliding boundary variants; `HRDiTFluxAttnProcessor` runs attention once per variant and - averages the outputs, with a proportional attention scale. A light early-step correction for "spatial disorder". - - **Progressive generation** -- denoising climbs a resolution ladder (1024 -> 2048 -> 4096 by default); each - stage bilinearly upsamples the previous stage's latent and re-noises it through the tail of the schedule. The - base stage is in-distribution and uses stock RoPE. + - **NTK-aware RoPE scaling** (`flux_rope`) -- the primary high-resolution mechanism. On every upscale-stage step + the rotary base `theta` is multiplied by a per-stage `ntk_factor`, compressing out-of-range positions into the + trained band. + - **SPA (Spatial Position Alignment)** -- for the leading `spa_steps` steps of a stage, `build_bundle_id_variants` + additionally coarsens position ids via a monotonic bundle map (no wrapping) across sliding variants, averaged + inside attention with a proportional scale. A light early-step correction for "spatial disorder". + - **Progressive generation with structure guidance** -- the ladder climbs 1024 -> 2048 -> 4096; each stage decodes + the previous latent, bicubic-upscales + sharpens it, and re-encodes it as a structural prior, then re-noises and + denoises. At every step the low-frequency (coarse-structure) band of the prediction is pulled toward the + upsampled previous-stage `pred_x0` (`alphas`, FFT Butterworth split), with a velocity-momentum term (`betas`). + This is what keeps the highest stage from drifting to a washed-out mean. - Not ported from the reference (documented follow-ups): the checkpoint-specific HAP head-scope pruning - (`configs/scope_plan_flux.json`) and the frequency-domain structure guidance (Butterworth low-pass + alpha/beta - blending in the flow-match step, swin patchify) that further stabilizes the highest stage. + Not ported from the reference (documented follow-ups): HAP head-scope attention pruning + (`configs/scope_plan_flux.json`), the `swin_pachify` shifted-window option, and DWT (as opposed to FFT) guidance. Args: prompt (`str` or `List[str]`): The prompt to render. height / width (`int`): Final output resolution. Defaults to 1024. - resolutions (`List[int]`, optional): Progressive resolution ladder (square side lengths). Defaults to - doubling from 1024 up to the target resolution. - group_num (`int`, defaults to 80): SPA bundle granularity; bundle size is `ceil(max_index / (group_num - 1))`. - Larger keeps more distinct positions inside the trained window. - ntk_factor (`List[float]`, optional): Per-upscale-stage NTK RoPE-base multiplier. Defaults to `[4.0, 10.0]` - (2048, 4096), extended by its last value for further stages. - spa_steps (`List[int]`, optional): Per-upscale-stage count of leading steps that use SPA. Defaults to - `[3, 0]` -- a light nudge at 2048, none at 4096 (NTK alone). - stage_strength (`float`, defaults to 0.6): Fraction of the schedule each upsampled stage re-noises through. + resolutions (`List[int]`, optional): Progressive ladder (square side lengths). Defaults to doubling from 1024. + group_num (`int`, defaults to 80): SPA bundle granularity; bundle size `ceil(max_index / (group_num - 1))`. + ntk_factor (`List[float]`, optional): Per-upscale-stage NTK RoPE-base multiplier. Defaults to `[4.0, 10.0]`. + spa_steps (`List[int]`, optional): Per-upscale-stage count of leading SPA steps. Defaults to `[3, 0]`. + num_inference_steps (`int`, defaults to 30): Base-stage steps (also the shared schedule length). + num_inference_steps_highres (`List[int]`, optional): Steps per upscale stage. Defaults to `[17, 10]`. + guidance_scale (`float`, defaults to 3.5): Base-stage guidance. guidance_scale_highres (`List[float]`, optional): Per-upscale-stage guidance. Defaults to `[4.5, 6.0]`. + alphas / betas (`List[float]`, optional): Per-stage structure-guidance weights (low-freq injection / velocity + momentum). Default `[1.0, 0.25]` and `[0.5, 0.5]`. + filter_ratio (`float`, defaults to 0.2): Butterworth low-pass cutoff for the structure split. Example: see `EXAMPLE_DOC_STRING`. """ @@ -327,6 +325,82 @@ def _stage_dimensions(height: int, width: int, target: int, side: int, quant: in stage_width = max(quant, int(round(width * side / target)) // quant * quant) return stage_height, stage_width + def _encode_image_to_latents(self, image, batch_size, num_channels_latents): + """Encode a pixel image to packed Flux latents (structural prior for a stage).""" + latents = self.vae.encode(image.to(self.vae.dtype).to(self.vae.device)).latent_dist.mode() + latents = (latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + latent_height, latent_width = latents.shape[-2], latents.shape[-1] + latents = self._pack_latents(latents, batch_size, num_channels_latents, latent_height, latent_width) + return latents.to(self.transformer.dtype) + + def _flowmatch_step( + self, + model_output, + timestep, + sample, + *, + structure_on=False, + pred_x0_dict=None, + height_dict=None, + width_dict=None, + batch_size=None, + num_channels_latents=None, + target_height=None, + target_width=None, + filter_ratio=0.0, + alpha=0.0, + beta=0.0, + ): + """ + One flow-match Euler step, optionally with HRDiT structure guidance. + + Structure guidance (``structure_on``) pulls the low-frequency band of the current predicted clean latent + toward the upsampled previous-stage ``pred_x0`` (weight ``alpha``), then applies a cross-step velocity + momentum (weight ``beta``). Returns ``(prev_sample, pred_x0)`` where ``pred_x0`` is the *pre-guidance* + prediction (stored for the next stage's reference, matching the reference implementation). + """ + scheduler = self.scheduler + if scheduler.step_index is None: + scheduler._init_step_index(timestep) + self._mo_high = None + self._mo_ref = None + + sample = sample.to(torch.float32) + sigma = scheduler.sigmas[scheduler.step_index] + sigma_next = scheduler.sigmas[scheduler.step_index + 1] + + pred_x0 = sample - model_output.to(torch.float32) * sigma + original_pred_x0 = pred_x0 + + if structure_on: + x0 = self._unpack_latents(pred_x0, target_height, target_width, self.vae_scale_factor).float() + latent_h, latent_w = x0.shape[-2], x0.shape[-1] + + ref_packed = pred_x0_dict[timestep.item()] + ref = self._unpack_latents( + ref_packed, height_dict[timestep.item()], width_dict[timestep.item()], self.vae_scale_factor + ).float() + ref = F.interpolate(ref, (latent_h, latent_w), mode="bicubic", align_corners=False) + + freq_filter = butterworth_low_pass_filter_2d(latent_h, latent_w, filter_ratio, x0.device) + x0 = x0 + alpha * (split_low_freq(ref, freq_filter) - split_low_freq(x0, freq_filter)) + + x0 = self._pack_latents(x0, batch_size, num_channels_latents, latent_h, latent_w) + ref = self._pack_latents(ref, batch_size, num_channels_latents, latent_h, latent_w) + + model_output = (sample - x0) / (sigma + 1e-6) + model_output_ref = (sample - ref) / (sigma + 1e-6) + if self._mo_high is not None: + model_output = model_output + beta * (self._mo_high + model_output_ref - self._mo_ref - model_output) + self._mo_high = model_output + self._mo_ref = model_output_ref + else: + model_output = model_output.to(torch.float32) + + prev_sample = (sample + (sigma_next - sigma) * model_output).to(self.transformer.dtype) + scheduler._step_index += 1 + return prev_sample, original_pred_x0 + @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( @@ -339,24 +413,22 @@ def __call__( group_num: int = 80, ntk_factor: Optional[List[float]] = None, spa_steps: Optional[List[int]] = None, - stage_strength: float = 0.6, - num_inference_steps: int = 28, + num_inference_steps: int = 30, + num_inference_steps_highres: Optional[List[int]] = None, guidance_scale: float = 3.5, guidance_scale_highres: Optional[List[float]] = None, + alphas: Optional[List[float]] = None, + betas: Optional[List[float]] = None, + filter_ratio: float = 0.2, num_images_per_prompt: int = 1, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.Tensor] = None, prompt_embeds: Optional[torch.Tensor] = None, pooled_prompt_embeds: Optional[torch.Tensor] = None, output_type: Optional[str] = "pil", return_dict: bool = True, max_sequence_length: int = 512, ): - r"""Generate a high-resolution image, training-free, with HRDiT (NTK RoPE + SPA + progressive generation). - - Accepts the standard [`FluxPipeline`] arguments plus `resolutions` / `group_num` / `ntk_factor` / `spa_steps` - and `stage_strength` / `guidance_scale_highres`. `height` and `width` set the final resolution; the pipeline - renders progressively up to it. + r"""Generate a high-resolution image, training-free, with HRDiT (NTK RoPE + SPA + structure-guided progression). Examples: """ @@ -365,17 +437,17 @@ def __call__( quant = self.vae_scale_factor * 2 if int(height) % quant != 0 or int(width) % quant != 0: raise ValueError(f"`height` and `width` must be multiples of {quant}, got {height} and {width}.") - if not 0.0 < stage_strength <= 1.0: - raise ValueError(f"`stage_strength` must be in (0, 1], got {stage_strength}.") ladder = self._resolution_ladder(height, width, resolutions) target = max(height, width) + n_upscale = len(ladder) - 1 - # Per-upscale-stage schedules (index j = ladder stage - 1). NTK-aware RoPE scaling is applied on every step - # and is the primary high-resolution mechanism; SPA augments only the leading `spa_steps` steps of a stage. ntk_schedule = ntk_factor if ntk_factor is not None else [4.0, 10.0] spa_schedule = spa_steps if spa_steps is not None else [3, 0] guidance_hr = guidance_scale_highres if guidance_scale_highres is not None else [4.5, 6.0] + steps_hr = num_inference_steps_highres if num_inference_steps_highres is not None else [17, 10] + alpha_schedule = alphas if alphas is not None else [1.0, 0.25] + beta_schedule = betas if betas is not None else [0.5, 0.5] def _stage_value(schedule, j, fill): if not schedule: @@ -385,12 +457,8 @@ def _stage_value(schedule, j, fill): device = self._execution_device dtype = prompt_embeds.dtype if prompt_embeds is not None else self.transformer.dtype - # 1. Encode prompt (guidance-distilled models like FLUX.1-dev need no true CFG pass). - ( - prompt_embeds, - pooled_prompt_embeds, - text_ids, - ) = self.encode_prompt( + # 1. Encode prompt. + prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt( prompt=prompt, prompt_2=prompt_2, prompt_embeds=prompt_embeds, @@ -401,7 +469,6 @@ def _stage_value(schedule, j, fill): lora_scale=None, ) batch_size = prompt_embeds.shape[0] - guidance_embeds = self.transformer.config.guidance_embeds def _guidance(scale): @@ -410,107 +477,118 @@ def _guidance(scale): return torch.full([1], scale, device=device, dtype=torch.float32).expand(batch_size) self._joint_attention_kwargs = {} + num_channels_latents = self.transformer.config.in_channels // 4 + axes_dim = self.transformer.pos_embed.axes_dim + rope_theta = self.transformer.pos_embed.theta + + # Shared flow-match schedule (same sigmas + mu across stages so per-timestep pred_x0 references align). + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + base_h, base_w = ladder[0], ladder[0] + base_grid = base_h // quant + mu = calculate_shift( + base_grid * base_grid, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) - # 2. Install the SPA/NTK processor (armed per-stage below; disarmed => stock attention). original_attn_processors = dict(self.transformer.attn_processors) self.transformer.set_attn_processor(HRDiTFluxAttnProcessor()) - axes_dim = self.transformer.pos_embed.axes_dim - rope_theta = self.transformer.pos_embed.theta + + pred_x0_dict, height_dict, width_dict = {}, {}, {} try: - # 3. Progressive denoising over the resolution ladder. - all_sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) - num_channels_latents = self.transformer.config.in_channels // 4 - latents = latents.to(device=device, dtype=dtype) if latents is not None else None - old_grid = None - self._num_timesteps = 0 - - for stage, side in enumerate(ladder): - stage_height, stage_width = self._stage_dimensions(height, width, target, side, quant) - latent_height = 2 * (stage_height // quant) - latent_width = 2 * (stage_width // quant) - grid_height, grid_width = latent_height // 2, latent_width // 2 - - if stage == 0: - if latents is not None and latents.shape[-2:] != (latent_height, latent_width): - raise ValueError( - f"Provided `latents` of spatial shape {latents.shape[-2:]} do not match the first " - f"progressive stage ({latent_height}, {latent_width})." - ) - if latents is None: - latents = randn_tensor( - (batch_size, num_channels_latents, latent_height, latent_width), - generator=generator, - device=device, - dtype=dtype, - ) - latents = self._pack_latents( - latents, batch_size, num_channels_latents, latent_height, latent_width - ) - stage_sigmas = all_sigmas - else: - latents = upsample_packed_latents(latents, old_grid, (grid_height, grid_width)) - # Later stages re-noise through the tail of the schedule (`stage_strength` of it). - num_stage_steps = max(1, int(round(num_inference_steps * stage_strength))) - stage_sigmas = all_sigmas[-num_stage_steps:] - old_grid = (grid_height, grid_width) - - image_ids = self._prepare_latent_image_ids(batch_size, grid_height, grid_width, device, dtype) - - # Upscale stages: NTK-scaled RoPE on every step (the primary high-res mechanism), with SPA bundle - # variants precomputed for the leading `stage_spa_steps` steps only. The base stage is in-distribution - # and uses stock RoPE (state disarmed). - if stage == 0: - _SPA_STATE.disarm() - stage_spa_steps = 0 - stage_guidance = _guidance(guidance_scale) - else: - j = stage - 1 - stage_ntk = float(_stage_value(ntk_schedule, j, 1.0)) - stage_spa_steps = int(_stage_value(spa_schedule, j, 0)) - stage_guidance = _guidance(float(_stage_value(guidance_hr, j, guidance_scale))) - base_rope = flux_rope( - torch.cat([text_ids, image_ids], dim=0), axes_dim, rope_theta, ntk_factor=stage_ntk - ) - if stage_spa_steps > 0: - variants = build_bundle_id_variants(image_ids, group_num) - variant_ropes = [ - flux_rope(torch.cat([text_ids, v], dim=0), axes_dim, rope_theta, ntk_factor=stage_ntk) - for v in variants - ] - else: - variant_ropes = [base_rope] - _SPA_STATE.arm(base_rope, variant_ropes) - - mu = calculate_shift( - grid_height * grid_width, - self.scheduler.config.get("base_image_seq_len", 256), - self.scheduler.config.get("max_image_seq_len", 4096), - self.scheduler.config.get("base_shift", 0.5), - self.scheduler.config.get("max_shift", 1.15), - ) - timesteps, _ = retrieve_timesteps( - self.scheduler, len(stage_sigmas), device, sigmas=stage_sigmas, mu=mu + # 2. Base stage (1024): stock RoPE, standard flow-match; record pred_x0 per timestep for guidance. + _SPA_STATE.disarm() + latent_h = 2 * (base_h // quant) + latent_w = 2 * (base_w // quant) + latents = randn_tensor( + (batch_size, num_channels_latents, latent_h, latent_w), generator=generator, device=device, dtype=dtype + ) + latents = self._pack_latents(latents, batch_size, num_channels_latents, latent_h, latent_w) + image_ids = self._prepare_latent_image_ids(batch_size, latent_h // 2, latent_w // 2, device, dtype) + base_guidance = _guidance(guidance_scale) + + timesteps, _ = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu) + self.scheduler._step_index = None + cur_h, cur_w = base_h, base_w + self.set_progress_bar_config(desc=f"HRDiT {base_w}x{base_h}") + with self.progress_bar(total=len(timesteps)) as progress_bar: + for t in timesteps: + self._current_timestep = t + timestep = t.expand(latents.shape[0]).to(latents.dtype) + with self.transformer.cache_context("cond"): + noise_pred = self.transformer( + hidden_states=latents, + timestep=timestep / 1000, + guidance=base_guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=image_ids, + joint_attention_kwargs=self.joint_attention_kwargs, + return_dict=False, + )[0] + latents, pred_x0 = self._flowmatch_step(noise_pred, t, latents) + pred_x0_dict[t.item()] = pred_x0 + height_dict[t.item()] = cur_h + width_dict[t.item()] = cur_w + progress_bar.update() + + # 3. Upscale stages with structure guidance. + for stage in range(1, len(ladder)): + j = stage - 1 + side = ladder[stage] + stage_h, stage_w = self._stage_dimensions(height, width, target, side, quant) + grid_h, grid_w = stage_h // quant, stage_w // quant + stage_ntk = float(_stage_value(ntk_schedule, j, 1.0)) + stage_spa_steps = int(_stage_value(spa_schedule, j, 0)) + stage_guidance = _guidance(float(_stage_value(guidance_hr, j, guidance_scale))) + stage_steps = int(_stage_value(steps_hr, j, max(1, round(num_inference_steps * 0.5)))) + stage_alpha0 = float(_stage_value(alpha_schedule, j, 0.0)) + stage_beta0 = float(_stage_value(beta_schedule, j, 0.0)) + + # Structural prior: decode -> bicubic upscale -> sharpen -> re-encode at the new resolution. + dec = self._unpack_latents(latents, cur_h, cur_w, self.vae_scale_factor) + dec = (dec / self.vae.config.scaling_factor) + self.vae.config.shift_factor + image = self.vae.decode(dec.to(self.vae.dtype), return_dict=False)[0] + image = F.interpolate(image, (stage_h, stage_w), mode="bicubic", align_corners=False) + image = sharpen(image) + latents = self._encode_image_to_latents(image, batch_size, num_channels_latents) + image_ids = self._prepare_latent_image_ids(batch_size, grid_h, grid_w, device, dtype) + + # NTK RoPE (every step) + SPA bundle variants (leading steps only). + base_rope = flux_rope( + torch.cat([text_ids, image_ids], dim=0), axes_dim, rope_theta, ntk_factor=stage_ntk ) - self.scheduler.set_begin_index(0) - self._num_timesteps += len(timesteps) - - if stage > 0: - # Flow-match interpolation at the stage's (shift-adjusted) starting sigma. - start_sigma = float(self.scheduler.sigmas[0]) - noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=dtype) - latents = (1.0 - start_sigma) * latents + start_sigma * noise - - self.set_progress_bar_config(desc=f"HRDiT {stage_width}x{stage_height}") - with self.progress_bar(total=len(timesteps)) as progress_bar: - for i, t in enumerate(timesteps): + if stage_spa_steps > 0: + variants = build_bundle_id_variants(image_ids, group_num) + variant_ropes = [ + flux_rope(torch.cat([text_ids, v], dim=0), axes_dim, rope_theta, ntk_factor=stage_ntk) + for v in variants + ] + else: + variant_ropes = [base_rope] + _SPA_STATE.arm(base_rope, variant_ropes) + + # Re-noise the prior to the tail of the shared schedule, then denoise the last `stage_steps` steps. + retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu) + dlfg_timesteps = self.scheduler.timesteps[-stage_steps:] + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + latents = self.scheduler.scale_noise(latents, dlfg_timesteps[:1], noise).to(self.transformer.dtype) + self.scheduler._step_index = None + + self.set_progress_bar_config(desc=f"HRDiT {stage_w}x{stage_h}") + with self.progress_bar(total=len(dlfg_timesteps)) as progress_bar: + for i, t in enumerate(dlfg_timesteps): self._current_timestep = t _SPA_STATE.spa_active = i < stage_spa_steps - timestep = t.expand(latents.shape[0]).to(latents.dtype) + decay = (stage_steps - i) / stage_steps with self.transformer.cache_context("cond"): noise_pred = self.transformer( hidden_states=latents, - timestep=timestep / 1000, + timestep=t.expand(latents.shape[0]).to(latents.dtype) / 1000, guidance=stage_guidance, pooled_projections=pooled_prompt_embeds, encoder_hidden_states=prompt_embeds, @@ -519,17 +597,37 @@ def _guidance(scale): joint_attention_kwargs=self.joint_attention_kwargs, return_dict=False, )[0] - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + latents, pred_x0 = self._flowmatch_step( + noise_pred, + t, + latents, + structure_on=True, + pred_x0_dict=pred_x0_dict, + height_dict=height_dict, + width_dict=width_dict, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + target_height=stage_h, + target_width=stage_w, + filter_ratio=filter_ratio, + alpha=stage_alpha0 * decay, + beta=stage_beta0 * decay, + ) + pred_x0_dict[t.item()] = pred_x0 + height_dict[t.item()] = stage_h + width_dict[t.item()] = stage_w progress_bar.update() + cur_h, cur_w = stage_h, stage_w + self._current_timestep = None if output_type == "latent": image = latents else: - latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = self._unpack_latents(latents, cur_h, cur_w, self.vae_scale_factor) latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor - image = self.vae.decode(latents, return_dict=False)[0] + image = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0] image = self.image_processor.postprocess(image, output_type=output_type) self.maybe_free_model_hooks() From 702d0c3ce1225aa44c48189ab19b4a29a9637f67 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 14:41:52 -0700 Subject: [PATCH 09/19] test(hrdit): update unit tests to current API (SPA variants, flux_rope NTK, structure-guidance helpers; drop HAP) --- .../test_community_pipeline_flux_hrdit.py | 152 ++++++++++-------- 1 file changed, 83 insertions(+), 69 deletions(-) diff --git a/tests/others/test_community_pipeline_flux_hrdit.py b/tests/others/test_community_pipeline_flux_hrdit.py index 6c4d034cbb45..4b4a2066a2cd 100644 --- a/tests/others/test_community_pipeline_flux_hrdit.py +++ b/tests/others/test_community_pipeline_flux_hrdit.py @@ -10,6 +10,7 @@ # specific language governing permissions and limitations under the License. import importlib.util +import math import sys import unittest from pathlib import Path @@ -38,74 +39,84 @@ def _load_module(name, path, extra_sys_path=None): class BuildBundleIdVariantsTests(unittest.TestCase): - def test_single_variant_at_trained_resolution(self): - # Below the trained 64x64 packed grid, SPA must reduce to the stock Flux position ids. - variants = hrdit.build_bundle_id_variants(32, 64, bundle_size=64, group_num=4) - - self.assertEqual(len(variants), 1) - expected = FluxPipeline._prepare_latent_image_ids(1, 32, 64, None, None) - self.assertTrue(torch.equal(variants[0], expected)) - - def test_variants_wrap_into_trained_window(self): - variants = hrdit.build_bundle_id_variants(128, 96, bundle_size=64, group_num=4) - - self.assertEqual(len(variants), 4) - for variant in variants: - self.assertEqual(variant.shape, (128 * 96, 3)) - self.assertLessEqual(variant[:, 1].max().item(), 63) - self.assertLessEqual(variant[:, 2].max().item(), 63) - for i in range(4): - for j in range(i + 1, 4): - self.assertFalse(torch.equal(variants[i], variants[j])) - - def test_first_variant_is_unshifted_partition(self): - variants = hrdit.build_bundle_id_variants(128, 64, bundle_size=64, group_num=4) - - ys = (torch.arange(128) % 64).repeat_interleave(64) - self.assertTrue(torch.equal(variants[0][:, 1], ys.to(variants[0].dtype))) - - -class UpsamplePackedLatentsTests(unittest.TestCase): - def test_upsample_shape(self): - latents = torch.randn(2, 16 * 16, 64) - upsampled = hrdit.upsample_packed_latents(latents, (16, 16), (32, 32)) - - self.assertEqual(upsampled.shape, (2, 32 * 32, 64)) - - def test_upsample_is_identity_at_same_grid(self): - latents = torch.randn(1, 8 * 8, 64) - upsampled = hrdit.upsample_packed_latents(latents, (8, 8), (8, 8)) - - self.assertTrue(torch.allclose(upsampled, latents, atol=1e-5)) - - -class HeadScopeTests(unittest.TestCase): - def test_scope_plan_round_robin(self): - plan = hrdit.build_head_scope_plan(24, window=64, full_period=4) - - self.assertEqual(plan.shape, (24,)) - self.assertEqual((plan == -1).sum().item(), 6) - self.assertEqual((plan == 64).sum().item(), 18) - - def test_mask_mod_respects_scopes(self): - grid_height, grid_width, num_txt = 8, 8, 16 - num_img = grid_height * grid_width - pos_h = torch.div(torch.arange(num_img), grid_width, rounding_mode="floor") - pos_w = torch.arange(num_img) % grid_width - windows = hrdit.build_head_scope_plan(8, window=2, full_period=4) - mask_mod = hrdit.build_mask_mod(pos_h, pos_w, windows, num_txt) - - windowed_head = 1 - full_head = 0 - image_query = torch.tensor(num_txt) # grid position (0, 0) - near_key = torch.tensor(num_txt + 1) # grid position (0, 1) - far_key = torch.tensor(num_txt + 3 * grid_width) # grid position (3, 0) - text_key = torch.tensor(3) - - self.assertTrue(bool(mask_mod(0, windowed_head, image_query, text_key))) - self.assertTrue(bool(mask_mod(0, windowed_head, image_query, near_key))) - self.assertFalse(bool(mask_mod(0, windowed_head, image_query, far_key))) - self.assertTrue(bool(mask_mod(0, full_head, image_query, far_key))) + def _image_ids(self, grid_h, grid_w): + return FluxPipeline._prepare_latent_image_ids(1, grid_h, grid_w, torch.device("cpu"), torch.float32) + + def test_variant_count_and_shape(self): + # group_num controls the bundle size s = ceil(max_index / (group_num - 1)); there are + # s_row + s_col - 1 sliding-boundary variants. + grid_h, grid_w, group_num = 128, 128, 80 + ids = self._image_ids(grid_h, grid_w) + variants = hrdit.build_bundle_id_variants(ids, group_num) + + s = max(1, math.ceil((grid_h - 1) / (group_num - 1))) + self.assertEqual(len(variants), 2 * s - 1) + for v in variants: + self.assertEqual(v.shape, (grid_h * grid_w, 3)) + + def test_coarsened_ids_stay_in_trained_window(self): + # The whole point of SPA: even at 256x256 (a 4096 image) coarsened positions stay in range. + ids = self._image_ids(256, 256) + for v in hrdit.build_bundle_id_variants(ids, 80): + self.assertLessEqual(v[:, 1].max().item(), 64) + self.assertLessEqual(v[:, 2].max().item(), 64) + + def test_mapping_is_monotonic_non_decreasing(self): + # A monotonic (non-wrapping) coarsening -> no periodic tiling. + grid = 128 + ids = self._image_ids(grid, grid) + rows = hrdit.build_bundle_id_variants(ids, 80)[0][:, 1].reshape(grid, grid) + self.assertTrue(bool((rows[1:] >= rows[:-1]).all())) + + def test_group_num_below_two_raises(self): + ids = self._image_ids(128, 128) + with self.assertRaises(ValueError): + hrdit.build_bundle_id_variants(ids, 1) + + +class FluxRopeTests(unittest.TestCase): + axes_dim = [16, 56, 56] + + def test_shape(self): + ids = torch.zeros(20, 3) + cos, sin = hrdit.flux_rope(ids, self.axes_dim, 10000.0, ntk_factor=1.0) + self.assertEqual(cos.shape, (20, sum(self.axes_dim))) + self.assertEqual(sin.shape, (20, sum(self.axes_dim))) + + def test_zero_position_is_identity_rotation(self): + # Position 0 -> no rotation: cos == 1, sin == 0. + ids = torch.zeros(4, 3) + cos, sin = hrdit.flux_rope(ids, self.axes_dim, 10000.0) + self.assertTrue(torch.allclose(cos, torch.ones_like(cos), atol=1e-5)) + self.assertTrue(torch.allclose(sin, torch.zeros_like(sin), atol=1e-5)) + + def test_ntk_scaling_lowers_rotation(self): + # Larger ntk_factor -> lower frequencies -> less rotation at the same position. + ids = torch.zeros(8, 3) + ids[:, 1] = torch.arange(8) + cos1, _ = hrdit.flux_rope(ids, self.axes_dim, 10000.0, ntk_factor=1.0) + cos10, _ = hrdit.flux_rope(ids, self.axes_dim, 10000.0, ntk_factor=10.0) + self.assertLess((cos10[7] - 1).abs().mean().item(), (cos1[7] - 1).abs().mean().item()) + + +class StructureGuidanceHelperTests(unittest.TestCase): + def test_butterworth_low_pass_shape_and_profile(self): + mask = hrdit.butterworth_low_pass_filter_2d(64, 64, 0.2, torch.device("cpu")) + self.assertEqual(mask.shape, (1, 1, 64, 64)) + self.assertGreater(mask[0, 0, 32, 32].item(), 0.9) # passband at the center + self.assertLess(mask[0, 0, 0, 0].item(), 0.1) # stopband at the corner + + def test_butterworth_zero_ratio_is_all_zeros(self): + mask = hrdit.butterworth_low_pass_filter_2d(16, 16, 0.0, torch.device("cpu")) + self.assertTrue(torch.equal(mask, torch.zeros_like(mask))) + + def test_split_low_freq_reduces_variance(self): + mask = hrdit.butterworth_low_pass_filter_2d(32, 32, 0.2, torch.device("cpu")) + x = torch.randn(1, 4, 32, 32) + low = hrdit.split_low_freq(x, mask) + self.assertEqual(low.shape, x.shape) + self.assertFalse(low.is_complex()) + self.assertLess(low.var().item(), x.var().item()) class PipelineIntegrationTests(unittest.TestCase): @@ -119,7 +130,10 @@ def test_benchmark_wiring(self): benchmark = _load_module( "benchmarking_flux_hrdit", BENCHMARKS_DIR / "benchmarking_flux_hrdit.py", extra_sys_path=BENCHMARKS_DIR ) - self.assertEqual(benchmark.RESULT_FILENAME, "flux_hrdit.csv") self.assertEqual(benchmark.CKPT_ID, "black-forest-labs/FLUX.1-dev") self.assertTrue(callable(benchmark.run_benchmarks)) + + +if __name__ == "__main__": + unittest.main() From d738ed71c879627c9e57f47cfe22e33414e3bd38 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 14:41:54 -0700 Subject: [PATCH 10/19] bench(hrdit): benchmark naive FluxPipeline vs HRDiT at target resolution (drop removed use_hap/group_num=1) --- benchmarks/benchmarking_flux_hrdit.py | 55 +++++++++++++-------------- 1 file changed, 27 insertions(+), 28 deletions(-) diff --git a/benchmarks/benchmarking_flux_hrdit.py b/benchmarks/benchmarking_flux_hrdit.py index 4ae2199abfe6..37946253f4db 100644 --- a/benchmarks/benchmarking_flux_hrdit.py +++ b/benchmarks/benchmarking_flux_hrdit.py @@ -1,13 +1,13 @@ -"""Benchmark HRDiT progressive high-resolution generation against naive single-pass generation. +"""Benchmark HRDiT training-free high-resolution generation against naive single-pass generation. -HRDiT (https://arxiv.org/abs/2608.07003) targets two failure modes of off-the-shelf DiT models at high -resolution: spatial disorder (fixed by SPA, which this benchmark runs with) and long generation time -(addressed by HAP pruning and by the progressive 1024 -> 2048 -> 4096 ladder). This script times the -end-to-end pipeline call for the naive single-pass baseline and for the HRDiT stages. +HRDiT (https://arxiv.org/abs/2608.07003) generates high-resolution images on off-the-shelf FLUX.1-dev +without fine-tuning, via NTK-aware RoPE scaling, Spatial Position Alignment (SPA) and a structure-guided +progressive 1024 -> 2048 -> 4096 ladder. This script times (and measures peak memory of) the end-to-end +call for a naive single-pass FLUX.1-dev baseline and for the HRDiT pipeline at the same target resolution. Run on a GPU machine with the FLUX.1-dev checkpoint available: - python benchmarks/benchmarking_flux_hrdit.py --height 2048 --width 2048 + python benchmarks/benchmarking_flux_hrdit.py --height 4096 --width 4096 """ import argparse @@ -24,6 +24,7 @@ CKPT_ID = "black-forest-labs/FLUX.1-dev" CUSTOM_PIPELINE_PATH = str(Path(__file__).resolve().parents[1] / "examples" / "community" / "pipeline_flux_hrdit.py") RESULT_FILENAME = "flux_hrdit.csv" +PROMPT = "a photo of a mountain lake at dawn" def load_pipeline(): @@ -34,27 +35,27 @@ def load_pipeline(): ).to(torch_device) -def run_benchmarks(height, width, num_inference_steps, group_num, use_hap): - pipe = load_pipeline() +def _peak_memory_gib(): + return torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else float("nan") + + +def run_benchmarks(height, width, num_inference_steps): + hrdit_pipe = load_pipeline() + # Stock single-pass FLUX.1-dev baseline, sharing the loaded components. + naive_pipe = FluxPipeline(**hrdit_pipe.components) settings = { - # Single-pass generation straight at the target resolution: SPA/HAP disabled, one stage. - "naive": dict(resolutions=[max(height, width)], group_num=1, use_hap=False), - # HRDiT: progressive ladder from 1024 up, SPA bundle averaging, optional HAP pruning. - "hrdit": dict(resolutions=None, group_num=group_num, use_hap=use_hap), + # Naive: generate straight at the target resolution in one pass (stock FluxPipeline). + "naive": (naive_pipe, dict(height=height, width=width, num_inference_steps=num_inference_steps)), + # HRDiT: NTK RoPE + SPA + structure-guided progressive ladder up to the target resolution. + "hrdit": (hrdit_pipe, dict(height=height, width=width, num_inference_steps=num_inference_steps)), } + results = [] - for name, kwargs in settings.items(): + for name, (pipe, kwargs) in settings.items(): flush() - latency = benchmark_fn( - pipe, - "a photo of a mountain lake at dawn", - height=height, - width=width, - num_inference_steps=num_inference_steps, - **kwargs, - ) - max_memory = torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else float("nan") + latency = benchmark_fn(pipe, PROMPT, **kwargs) + max_memory = _peak_memory_gib() results.append((name, latency, max_memory)) print(f"{name:>6}: {latency:.3f}s, peak memory {max_memory:.2f} GiB") @@ -67,11 +68,9 @@ def run_benchmarks(height, width, num_inference_steps, group_num, use_hap): if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--height", type=int, default=2048) - parser.add_argument("--width", type=int, default=2048) - parser.add_argument("--num_inference_steps", type=int, default=28) - parser.add_argument("--group_num", type=int, default=4, help="number of SPA bundle variants") - parser.add_argument("--no_hap", action="store_true", help="disable head-adaptive attention pruning") + parser.add_argument("--height", type=int, default=4096) + parser.add_argument("--width", type=int, default=4096) + parser.add_argument("--num_inference_steps", type=int, default=30) args = parser.parse_args() - run_benchmarks(args.height, args.width, args.num_inference_steps, args.group_num, not args.no_hap) + run_benchmarks(args.height, args.width, args.num_inference_steps) From a413b38fbbbf6a70a87d3f07dc16776e7e6cd998 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 14:42:40 -0700 Subject: [PATCH 11/19] docs(hrdit): update community README to the NTK + SPA + structure-guidance method and simplified API --- examples/community/README.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/examples/community/README.md b/examples/community/README.md index c14e875824d0..ee3b3ccc73a3 100644 --- a/examples/community/README.md +++ b/examples/community/README.md @@ -5633,7 +5633,7 @@ result.images[0].save(f"flux_fill_controlnet_inpaint_depth{timestamp}.jpg") ## Flux HRDiT -Training-free high-resolution (up to 4096x4096) text-to-image generation with off-the-shelf Flux models, adapted from [HRDiT](https://arxiv.org/abs/2608.07003) ([official implementation](https://github.com/zylwithxy/HRDiT)). No fine-tuning and no new weights: the pipeline adds Spatial Position Alignment (position ids wrapped into the trained RoPE window and averaged over sliding bundle variants), optional head-adaptive attention pruning via FlexAttention, and a progressive 1024 -> 2048 -> 4096 generation ladder on top of the stock `FluxPipeline` denoise loop. +Training-free high-resolution (up to 4096x4096) text-to-image generation with off-the-shelf Flux models, adapted from [HRDiT](https://arxiv.org/abs/2608.07003) ([official implementation](https://github.com/zylwithxy/HRDiT)). No fine-tuning and no new weights. On top of the stock `FluxPipeline` denoise loop it adds, per upscale stage: **NTK-aware RoPE scaling** (the RoPE base is scaled per stage so out-of-range high-resolution positions fall back into the trained band — the primary high-res mechanism); **Spatial Position Alignment (SPA)** on the leading steps (position ids monotonically coarsened into the trained window and averaged over sliding bundle variants inside attention); and a **structure-guided progressive 1024 -> 2048 -> 4096 ladder** (each stage decodes, upscales and re-encodes the previous latent as a structural prior, then injects its low-frequency band each step to prevent high-resolution drift). The default arguments reproduce the reference configuration. ```py import torch @@ -5647,9 +5647,8 @@ image = pipe( "a photo of a mountain lake at dawn", height=4096, width=4096, - num_inference_steps=28, - group_num=4, # number of SPA bundle variants averaged per step - use_hap=True, # falls back to full attention if FlexAttention (torch >= 2.7) is unavailable ).images[0] image.save("hrdit_4096.png") ``` + +Key arguments (all optional, defaulting to the reference configuration): `ntk_factor` (per-stage RoPE-base multiplier, default `[4.0, 10.0]`), `spa_steps` (leading SPA steps per stage, default `[3, 0]`), `group_num` (SPA bundle granularity, default `80`), `alphas`/`betas` (structure-guidance weights, default `[1.0, 0.25]`/`[0.5, 0.5]`), and `guidance_scale_highres` (default `[4.5, 6.0]`). A 4096x4096 generation runs in ~2 min at ~26 GB peak on an A100-80GB. From fe0a85f845ec26b869c2df33175f988aade8a85e Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 14:57:03 -0700 Subject: [PATCH 12/19] docs(hrdit): add community pipeline overview-table row (with Colab) Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- examples/community/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/community/README.md b/examples/community/README.md index ee3b3ccc73a3..3fc7b57a1992 100644 --- a/examples/community/README.md +++ b/examples/community/README.md @@ -89,6 +89,7 @@ PIXART-α Controlnet pipeline | Implementation of the controlnet model for pixar | Stable Diffusion 3 InstructPix2Pix Pipeline | Implementation of Stable Diffusion 3 InstructPix2Pix Pipeline | [Stable Diffusion 3 InstructPix2Pix Pipeline](#stable-diffusion-3-instructpix2pix-pipeline) | [![Hugging Face Models](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/BleachNick/SD3_UltraEdit_freeform) [![Hugging Face Models](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/CaptainZZZ/sd3-instructpix2pix) | [Jiayu Zhang](https://github.com/xduzhangjiayu) and [Haozhe Zhao](https://github.com/HaozheZhao)| | Flux Kontext multiple images | A modified version of the `FluxKontextPipeline` that supports calling Flux Kontext with multiple reference images.| [Flux Kontext multiple input Pipeline](#flux-kontext-multiple-images) | - | [Net-Mist](https://github.com/Net-Mist) | | Flux Fill ControlNet Pipeline | A modified version of the `FluxFillPipeline` and `FluxControlNetInpaintPipeline` that supports Controlnet with Flux Fill model.| [Flux Fill ControlNet Pipeline](#Flux-Fill-ControlNet-Pipeline) | - | [pratim4dasude](https://github.com/pratim4dasude) | +| Flux HRDiT | Training-free high-resolution (up to 4096×4096) text-to-image on off-the-shelf FLUX.1-dev via NTK-aware RoPE scaling, Spatial Position Alignment, and structure-guided progressive generation. Adapted from [HRDiT](https://arxiv.org/abs/2608.07003). | [Flux HRDiT](#flux-hrdit) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/19IS19j2BvQm4Qnf07v-FDYiiqYZfD7Ja?usp=sharing) | [Terry Rodriguez](https://github.com/smellslikeml) | To load a custom pipeline you just need to pass the `custom_pipeline` argument to `DiffusionPipeline`, as one of the files in `diffusers/examples/community`. Feel free to send a PR with your own pipelines, we will merge them quickly. From 5f5fd012a7293d337ad3a8b7471316cbfd3e251f Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 15:18:39 -0700 Subject: [PATCH 13/19] =?UTF-8?q?style(hrdit):=20satisfy=20ruff=20?= =?UTF-8?q?=E2=80=94=20drop=20unused=20n=5Fupscale,=20format=20exponent=20?= =?UTF-8?q?operators?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- examples/community/pipeline_flux_hrdit.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index d40a83f47199..995f2b4fc9fa 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -35,7 +35,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name # FLUX is trained at 1024x1024 -> a 64x64 packed grid (4096 image tokens) plus 512 text tokens. -_TRAIN_SEQ_LEN = 64 ** 2 + 512 +_TRAIN_SEQ_LEN = 64**2 + 512 EXAMPLE_DOC_STRING = """ Examples: @@ -114,7 +114,7 @@ def flux_rope(ids: torch.Tensor, axes_dim, theta: float, ntk_factor: float = 1.0 for i, dim in enumerate(axes_dim): pos = ids[:, i].to(torch.float64) exps = torch.arange(0, dim, 2, dtype=torch.float64, device=ids.device)[: dim // 2] / dim - freqs = torch.outer(pos, 1.0 / (scaled_theta ** exps)) # [S, dim/2] + freqs = torch.outer(pos, 1.0 / (scaled_theta**exps)) # [S, dim/2] cos_out.append(freqs.cos().repeat_interleave(2, dim=1).float()) sin_out.append(freqs.sin().repeat_interleave(2, dim=1).float()) return torch.cat(cos_out, dim=-1), torch.cat(sin_out, dim=-1) @@ -126,8 +126,8 @@ def butterworth_low_pass_filter_2d(height: int, width: int, ratio: float, device return torch.zeros(1, 1, height, width, device=device) yy = (2.0 * torch.arange(height, device=device) / height - 1.0).view(height, 1) xx = (2.0 * torch.arange(width, device=device) / width - 1.0).view(1, width) - d_square = yy ** 2 + xx ** 2 - mask = 1.0 / (1.0 + (d_square / ratio ** 2) ** order) + d_square = yy**2 + xx**2 + mask = 1.0 / (1.0 + (d_square / ratio**2) ** order) return mask.view(1, 1, height, width) @@ -229,7 +229,7 @@ def __call__( if _SPA_STATE.proportional and seq_len > 1: scale = math.sqrt(math.log(seq_len, _TRAIN_SEQ_LEN) / head_dim) else: - scale = head_dim ** -0.5 + scale = head_dim**-0.5 value_t = value.transpose(1, 2).contiguous() # [B, H, S, D] ropes = _SPA_STATE.current_ropes() @@ -440,7 +440,6 @@ def __call__( ladder = self._resolution_ladder(height, width, resolutions) target = max(height, width) - n_upscale = len(ladder) - 1 ntk_schedule = ntk_factor if ntk_factor is not None else [4.0, 10.0] spa_schedule = spa_steps if spa_steps is not None else [3, 0] From 1166542a360ffef7af907321fde492ff2f370800 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 15:18:40 -0700 Subject: [PATCH 14/19] =?UTF-8?q?style(hrdit):=20satisfy=20ruff=20?= =?UTF-8?q?=E2=80=94=20sort=20imports=20(I001),=20dict()=20->=20literal=20?= =?UTF-8?q?(C408)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- benchmarks/benchmarking_flux_hrdit.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/benchmarks/benchmarking_flux_hrdit.py b/benchmarks/benchmarking_flux_hrdit.py index 37946253f4db..db435c0e2857 100644 --- a/benchmarks/benchmarking_flux_hrdit.py +++ b/benchmarks/benchmarking_flux_hrdit.py @@ -14,7 +14,6 @@ from pathlib import Path import torch - from benchmarking_utils import benchmark_fn, flush from diffusers import FluxPipeline @@ -46,9 +45,9 @@ def run_benchmarks(height, width, num_inference_steps): settings = { # Naive: generate straight at the target resolution in one pass (stock FluxPipeline). - "naive": (naive_pipe, dict(height=height, width=width, num_inference_steps=num_inference_steps)), + "naive": (naive_pipe, {"height": height, "width": width, "num_inference_steps": num_inference_steps}), # HRDiT: NTK RoPE + SPA + structure-guided progressive ladder up to the target resolution. - "hrdit": (hrdit_pipe, dict(height=height, width=width, num_inference_steps=num_inference_steps)), + "hrdit": (hrdit_pipe, {"height": height, "width": width, "num_inference_steps": num_inference_steps}), } results = [] From 928032848ae7b41b7773d7a16bca404909ed8976 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 15:38:30 -0700 Subject: [PATCH 15/19] style(hrdit): align copyright header with Flux community pipelines (Black Forest Labs + HuggingFace Team) Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- examples/community/pipeline_flux_hrdit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 995f2b4fc9fa..0a49ef6cd609 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# Copyright 2025 Black Forest Labs and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at From 6cdcf7855c7d695bb8efb44a1762a05fc159c740 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 15:38:32 -0700 Subject: [PATCH 16/19] style(hrdit): align test copyright header with diffusers convention Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- tests/others/test_community_pipeline_flux_hrdit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/others/test_community_pipeline_flux_hrdit.py b/tests/others/test_community_pipeline_flux_hrdit.py index 4b4a2066a2cd..bd46911ba873 100644 --- a/tests/others/test_community_pipeline_flux_hrdit.py +++ b/tests/others/test_community_pipeline_flux_hrdit.py @@ -1,4 +1,4 @@ -# Copyright 2026 HuggingFace Inc. +# Copyright 2025 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with # the License. You may obtain a copy of the License at From 75cbf7445eb3ee9fcb55ecfa243f50fda1803090 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 16:47:22 -0700 Subject: [PATCH 17/19] docs(hrdit): note MIT license of the reference implementation Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- examples/community/pipeline_flux_hrdit.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index 0a49ef6cd609..aaf4d6742924 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -264,7 +264,8 @@ class HRDiTFluxPipeline(FluxPipeline): Training-free high-resolution (up to 4096x4096) text-to-image with off-the-shelf Flux models. Adapted from HRDiT, "Training-Free High-Resolution Image Generation with Off-the-Shelf Diffusion Transformer - Models" (https://arxiv.org/abs/2608.07003); reference implementation at https://github.com/zylwithxy/HRDiT. + Models" (https://arxiv.org/abs/2608.07003); MIT-licensed reference implementation at + https://github.com/zylwithxy/HRDiT. Training-free pieces on top of the stock `FluxPipeline` denoise loop: From 074dd248519dc54ad76f3d0e7a80473b6d6c4ed4 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 17:03:54 -0700 Subject: [PATCH 18/19] refactor(hrdit): delegate Euler update to scheduler.step; drop private scheduler API + dead code Self-review fixes against the diffusers .ai rubric: - _flowmatch_step now hands the Euler update to self.scheduler.step() instead of inlining prev = sample + (sigma_next-sigma)*model_output and managing scheduler._init_step_index/_step_index by hand (pipelines.md gotcha #3). Reads the current sigma via scheduler.index_for_timestep(); structure guidance still only adjusts the velocity before the step. Behaviour-preserving (same sigma / formula / dtype). - removed the always-true _SPAState.proportional flag + its dead branch, the unused self._current_timestep bookkeeping, and the torchvision try/except fallback (hard-import, matching the reference) per AGENTS.md 'no defensive/just-in-case code'. - momentum state (_mo_high/_mo_ref) now reset per stage by the caller. 13/13 unit tests pass; ruff check + format clean. Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- examples/community/pipeline_flux_hrdit.py | 54 +++++++---------------- 1 file changed, 17 insertions(+), 37 deletions(-) diff --git a/examples/community/pipeline_flux_hrdit.py b/examples/community/pipeline_flux_hrdit.py index aaf4d6742924..2aacecee5dae 100644 --- a/examples/community/pipeline_flux_hrdit.py +++ b/examples/community/pipeline_flux_hrdit.py @@ -15,6 +15,7 @@ import numpy as np import torch import torch.nn.functional as F +from torchvision.transforms.functional import gaussian_blur from diffusers.models.embeddings import apply_rotary_emb from diffusers.models.transformers.transformer_flux import FluxAttnProcessor, _get_qkv_projections @@ -24,14 +25,6 @@ from diffusers.utils.torch_utils import randn_tensor -try: - from torchvision.transforms.functional import gaussian_blur as _tv_gaussian_blur - - _TORCHVISION_AVAILABLE = True -except ImportError: - _TORCHVISION_AVAILABLE = False - - logger = logging.get_logger(__name__) # pylint: disable=invalid-name # FLUX is trained at 1024x1024 -> a 64x64 packed grid (4096 image tokens) plus 512 text tokens. @@ -139,10 +132,8 @@ def split_low_freq(x: torch.Tensor, freq_filter: torch.Tensor) -> torch.Tensor: def sharpen(image: torch.Tensor, kernel_size: int = 3, sigma: float = 1.0, alpha: float = 1.0) -> torch.Tensor: - """Unsharp-mask sharpening of an image tensor; no-op if torchvision is unavailable.""" - if not _TORCHVISION_AVAILABLE: - return image - blurred = _tv_gaussian_blur(image, kernel_size=[kernel_size, kernel_size], sigma=[sigma, sigma]) + """Unsharp-mask sharpening of the upscaled structural prior before it is re-encoded.""" + blurred = gaussian_blur(image, kernel_size=[kernel_size, kernel_size], sigma=[sigma, sigma]) return (alpha + 1.0) * image - alpha * blurred @@ -159,7 +150,6 @@ def __init__(self): self.base_rope = None self.variant_ropes = None self.spa_active = False - self.proportional = True @property def enabled(self): @@ -226,10 +216,8 @@ def __call__( head_dim = query.shape[-1] seq_len = query.shape[1] - if _SPA_STATE.proportional and seq_len > 1: - scale = math.sqrt(math.log(seq_len, _TRAIN_SEQ_LEN) / head_dim) - else: - scale = head_dim**-0.5 + # Proportional attention scale for the longer high-res sequence (equals the stock 1/sqrt(d) at train length). + scale = math.sqrt(math.log(seq_len, _TRAIN_SEQ_LEN) / head_dim) if seq_len > 1 else head_dim**-0.5 value_t = value.transpose(1, 2).contiguous() # [B, H, S, D] ropes = _SPA_STATE.current_ropes() @@ -353,22 +341,17 @@ def _flowmatch_step( beta=0.0, ): """ - One flow-match Euler step, optionally with HRDiT structure guidance. + One flow-match step (the Euler update is delegated to ``self.scheduler.step``), optionally with HRDiT + structure guidance. Structure guidance (``structure_on``) pulls the low-frequency band of the current predicted clean latent toward the upsampled previous-stage ``pred_x0`` (weight ``alpha``), then applies a cross-step velocity - momentum (weight ``beta``). Returns ``(prev_sample, pred_x0)`` where ``pred_x0`` is the *pre-guidance* - prediction (stored for the next stage's reference, matching the reference implementation). + momentum (weight ``beta``); ``self._mo_high`` / ``self._mo_ref`` are reset per stage by the caller. Returns + ``(prev_sample, pred_x0)`` where ``pred_x0`` is the *pre-guidance* prediction (stored as the next stage's + structural reference). """ - scheduler = self.scheduler - if scheduler.step_index is None: - scheduler._init_step_index(timestep) - self._mo_high = None - self._mo_ref = None - sample = sample.to(torch.float32) - sigma = scheduler.sigmas[scheduler.step_index] - sigma_next = scheduler.sigmas[scheduler.step_index + 1] + sigma = self.scheduler.sigmas[self.scheduler.index_for_timestep(timestep)] pred_x0 = sample - model_output.to(torch.float32) * sigma original_pred_x0 = pred_x0 @@ -398,9 +381,9 @@ def _flowmatch_step( else: model_output = model_output.to(torch.float32) - prev_sample = (sample + (sigma_next - sigma) * model_output).to(self.transformer.dtype) - scheduler._step_index += 1 - return prev_sample, original_pred_x0 + # Let the scheduler own the Euler update; structure guidance only adjusts the velocity above. + prev_sample = self.scheduler.step(model_output, timestep, sample, return_dict=False)[0] + return prev_sample.to(self.transformer.dtype), original_pred_x0 @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) @@ -511,12 +494,10 @@ def _guidance(scale): base_guidance = _guidance(guidance_scale) timesteps, _ = retrieve_timesteps(self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu) - self.scheduler._step_index = None cur_h, cur_w = base_h, base_w self.set_progress_bar_config(desc=f"HRDiT {base_w}x{base_h}") with self.progress_bar(total=len(timesteps)) as progress_bar: for t in timesteps: - self._current_timestep = t timestep = t.expand(latents.shape[0]).to(latents.dtype) with self.transformer.cache_context("cond"): noise_pred = self.transformer( @@ -577,12 +558,13 @@ def _guidance(scale): dlfg_timesteps = self.scheduler.timesteps[-stage_steps:] noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) latents = self.scheduler.scale_noise(latents, dlfg_timesteps[:1], noise).to(self.transformer.dtype) - self.scheduler._step_index = None + # Reset the structure-guidance velocity momentum at the start of each stage. + self._mo_high = None + self._mo_ref = None self.set_progress_bar_config(desc=f"HRDiT {stage_w}x{stage_h}") with self.progress_bar(total=len(dlfg_timesteps)) as progress_bar: for i, t in enumerate(dlfg_timesteps): - self._current_timestep = t _SPA_STATE.spa_active = i < stage_spa_steps decay = (stage_steps - i) / stage_steps with self.transformer.cache_context("cond"): @@ -620,8 +602,6 @@ def _guidance(scale): cur_h, cur_w = stage_h, stage_w - self._current_timestep = None - if output_type == "latent": image = latents else: From 2ec80a1eb90ddf3ee8405a6227e58a528642f666 Mon Sep 17 00:00:00 2001 From: smellslikeml Date: Fri, 14 Aug 2026 17:25:44 -0700 Subject: [PATCH 19/19] docs(hrdit): update Colab validation notebook link Co-Authored-By: Claude Opus 4.8 (1M context) Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com> --- examples/community/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/community/README.md b/examples/community/README.md index 3fc7b57a1992..939b62380245 100644 --- a/examples/community/README.md +++ b/examples/community/README.md @@ -89,7 +89,7 @@ PIXART-α Controlnet pipeline | Implementation of the controlnet model for pixar | Stable Diffusion 3 InstructPix2Pix Pipeline | Implementation of Stable Diffusion 3 InstructPix2Pix Pipeline | [Stable Diffusion 3 InstructPix2Pix Pipeline](#stable-diffusion-3-instructpix2pix-pipeline) | [![Hugging Face Models](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/BleachNick/SD3_UltraEdit_freeform) [![Hugging Face Models](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/CaptainZZZ/sd3-instructpix2pix) | [Jiayu Zhang](https://github.com/xduzhangjiayu) and [Haozhe Zhao](https://github.com/HaozheZhao)| | Flux Kontext multiple images | A modified version of the `FluxKontextPipeline` that supports calling Flux Kontext with multiple reference images.| [Flux Kontext multiple input Pipeline](#flux-kontext-multiple-images) | - | [Net-Mist](https://github.com/Net-Mist) | | Flux Fill ControlNet Pipeline | A modified version of the `FluxFillPipeline` and `FluxControlNetInpaintPipeline` that supports Controlnet with Flux Fill model.| [Flux Fill ControlNet Pipeline](#Flux-Fill-ControlNet-Pipeline) | - | [pratim4dasude](https://github.com/pratim4dasude) | -| Flux HRDiT | Training-free high-resolution (up to 4096×4096) text-to-image on off-the-shelf FLUX.1-dev via NTK-aware RoPE scaling, Spatial Position Alignment, and structure-guided progressive generation. Adapted from [HRDiT](https://arxiv.org/abs/2608.07003). | [Flux HRDiT](#flux-hrdit) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/19IS19j2BvQm4Qnf07v-FDYiiqYZfD7Ja?usp=sharing) | [Terry Rodriguez](https://github.com/smellslikeml) | +| Flux HRDiT | Training-free high-resolution (up to 4096×4096) text-to-image on off-the-shelf FLUX.1-dev via NTK-aware RoPE scaling, Spatial Position Alignment, and structure-guided progressive generation. Adapted from [HRDiT](https://arxiv.org/abs/2608.07003). | [Flux HRDiT](#flux-hrdit) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1AU6QNOGDMCPpGdocVyIyKpXXGt3hoWdo?usp=sharing) | [Terry Rodriguez](https://github.com/smellslikeml) | To load a custom pipeline you just need to pass the `custom_pipeline` argument to `DiffusionPipeline`, as one of the files in `diffusers/examples/community`. Feel free to send a PR with your own pipelines, we will merge them quickly.