Community Pipeline: Add HRDiT (training-free high-resolution FLUX.1-dev)

benchmarks/benchmarking_flux_hrdit.py

@@ -0,0 +1,75 @@
+"""Benchmark HRDiT training-free high-resolution generation against naive single-pass generation.
+
+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 4096 --width 4096
+"""
+
+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"
+PROMPT = "a photo of a mountain lake at dawn"
+
+
+def load_pipeline():
+    return FluxPipeline.from_pretrained(
+        CKPT_ID,
+        torch_dtype=torch.bfloat16,
+        custom_pipeline=CUSTOM_PIPELINE_PATH,
+    ).to(torch_device)
+
+
+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 = {
+        # Naive: generate straight at the target resolution in one pass (stock FluxPipeline).
+        "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, {"height": height, "width": width, "num_inference_steps": num_inference_steps}),
+    }
+
+    results = []
+    for name, (pipe, kwargs) in settings.items():
+        flush()
+        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")
+
+    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=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)

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/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.
 
@@ -5630,3 +5631,25 @@ 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. 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
+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,
+).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.