diff --git a/examples/visual_gen/hf_examples.sh b/examples/visual_gen/hf_examples.sh index 192983015d10..f2bb84dfd4fd 100755 --- a/examples/visual_gen/hf_examples.sh +++ b/examples/visual_gen/hf_examples.sh @@ -78,7 +78,7 @@ echo "" ############################################# echo "============================================" -echo "1/1: WAN Baseline Test" +echo "1/3: WAN Baseline Test" echo "============================================" echo "" @@ -106,6 +106,70 @@ fi echo "" +############################################# +# FLUX.1 Baseline Test +############################################# + +echo "============================================" +echo "2/3: FLUX.1 Baseline Test" +echo "============================================" +echo "" + +FLUX1_MODEL="${MODEL_ROOT}/FLUX.1-dev/" +FLUX1_OUTPUT="${OUTPUT_DIR}/flux1_baseline.png" + +if [ -d "$FLUX1_MODEL" ]; then + echo "Testing FLUX.1 with official diffusers..." + python ${PROJECT_ROOT}/examples/visual_gen/hf_flux.py \ + --model_path "$FLUX1_MODEL" \ + --output_path "$FLUX1_OUTPUT" \ + --prompt "A cat holding a sign that says hello world" \ + --height 1024 \ + --width 1024 \ + --steps 50 \ + --guidance_scale 3.5 \ + --seed 42 + echo "" + echo "✅ FLUX.1 baseline test completed" + echo " Output: $FLUX1_OUTPUT" +else + echo "⚠️ SKIPPED: FLUX.1 model not found at $FLUX1_MODEL" +fi + +echo "" + +############################################# +# FLUX.2 Baseline Test +############################################# + +echo "============================================" +echo "3/3: FLUX.2 Baseline Test" +echo "============================================" +echo "" + +FLUX2_MODEL="${MODEL_ROOT}/FLUX.2-dev/" +FLUX2_OUTPUT="${OUTPUT_DIR}/flux2_baseline.png" + +if [ -d "$FLUX2_MODEL" ]; then + echo "Testing FLUX.2 with official diffusers..." + python ${PROJECT_ROOT}/examples/visual_gen/hf_flux2.py \ + --model_path "$FLUX2_MODEL" \ + --output_path "$FLUX2_OUTPUT" \ + --prompt "A cat holding a sign that says hello world" \ + --height 1024 \ + --width 1024 \ + --steps 50 \ + --guidance_scale 3.5 \ + --seed 42 + echo "" + echo "✅ FLUX.2 baseline test completed" + echo " Output: $FLUX2_OUTPUT" +else + echo "⚠️ SKIPPED: FLUX.2 model not found at $FLUX2_MODEL" +fi + +echo "" + ############################################# # Summary ############################################# diff --git a/examples/visual_gen/hf_flux.py b/examples/visual_gen/hf_flux.py new file mode 100755 index 000000000000..aba1848837d0 --- /dev/null +++ b/examples/visual_gen/hf_flux.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Baseline test for FLUX.1 using official diffusers library.""" + +import sys + +import torch +from output_handler import OutputHandler + +from tensorrt_llm._torch.visual_gen import MediaOutput + + +def test_flux_baseline( + model_path: str, + output_path: str, + prompt: str = "A cat holding a sign that says hello world", + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 3.5, + seed: int = 42, +): + """Test FLUX.1 image generation with official diffusers.""" + from diffusers import FluxPipeline + + print("=" * 80) + print("FLUX.1 Baseline Test (Official Diffusers)") + print("=" * 80) + print() + + # Load pipeline + print(f"Loading FLUX.1 pipeline from {model_path}...") + pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16) + pipe.to("cuda") + print("✅ Pipeline loaded") + print() + + # Check model states + print("Model Training States:") + print(f" text_encoder.training: {pipe.text_encoder.training}") + if hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None: + print(f" text_encoder_2.training: {pipe.text_encoder_2.training}") + print(f" transformer.training: {pipe.transformer.training}") + print(f" vae.training: {pipe.vae.training}") + print() + + # Generate image + print(f"Generating image: '{prompt}'") + print(f"Parameters: {height}x{width}, {num_inference_steps} steps, guidance={guidance_scale}") + print() + + # Set random seed + generator = torch.Generator(device="cuda").manual_seed(seed) + + result = pipe( + prompt=prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + ) + + # Extract PIL image and convert to (H, W, C) uint8 tensor + import numpy as np + + pil_image = result.images[0] + image = torch.from_numpy(np.array(pil_image)) + + print("=" * 80) + print("Generation Complete!") + print("=" * 80) + print(f"Image shape: {image.shape}") + print(f"Image dtype: {image.dtype}") + print() + + # Save output + print(f"Saving output to {output_path}...") + OutputHandler.save(output=MediaOutput(image=image), output_path=output_path) + print(f"✅ Saved to {output_path}") + print() + + print("=" * 80) + print("FLUX.1 BASELINE TEST PASSED ✅") + print("=" * 80) + return image + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="HuggingFace Baseline - FLUX.1 Text-to-Image Generation" + ) + + # Model & Input + parser.add_argument( + "--model_path", + type=str, + default="/llm-models/FLUX.1-dev/", + help="Path to FLUX.1 model", + ) + parser.add_argument( + "--prompt", + type=str, + default="A cat holding a sign that says hello world", + help="Text prompt for generation", + ) + parser.add_argument( + "--output_path", type=str, default="flux1_baseline.png", help="Output file path" + ) + + # Generation parameters + parser.add_argument("--height", type=int, default=1024, help="Image height") + parser.add_argument("--width", type=int, default=1024, help="Image width") + parser.add_argument("--steps", type=int, default=50, help="Number of denoising steps") + parser.add_argument( + "--guidance_scale", type=float, default=3.5, help="Guidance scale (embedded guidance)" + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed") + + args = parser.parse_args() + + try: + test_flux_baseline( + args.model_path, + args.output_path, + prompt=args.prompt, + height=args.height, + width=args.width, + num_inference_steps=args.steps, + guidance_scale=args.guidance_scale, + seed=args.seed, + ) + except Exception as e: + print(f"\n❌ ERROR: {e}") + import traceback + + traceback.print_exc() + sys.exit(1) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index b68dc7f2a20e..52f7b3238055 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -34,6 +34,8 @@ Before running these examples, ensure you have: ```bash trtllm-serve $LLM_MODEL_DIR/Wan2.1-T2V-1.3B-Diffusers --extra_visual_gen_options ./configs/wan.yml + trtllm-serve $LLM_MODEL_DIR/FLUX.1-dev --extra_visual_gen_options ./configs/flux1.yml + trtllm-serve $LLM_MODEL_DIR/FLUX.2-dev --extra_visual_gen_options ./configs/flux2.yml # Run server on background: trtllm-serve $LLM_MODEL_DIR/Wan2.1-T2V-1.3B-Diffusers --extra_visual_gen_options ./configs/wan.yml > /tmp/serve.log 2>&1 & @@ -48,24 +50,29 @@ Before running these examples, ensure you have: Current supported & tested models: 1. WAN T2V/I2V for video generation (t2v, ti2v, delete_video) +2. FLUX.1 for image generation (t2i) +3. FLUX.2 for image generation (t2i) -### 1. Synchronous Image Generation (`sync_t2i.py`) +### 1. Synchronous Image Generation (`sync_image_gen.py`) -Demonstrates synchronous text-to-image generation using the OpenAI SDK. +Demonstrates synchronous text-to-image generation using the OpenAI SDK. Supports FLUX.1 and FLUX.2. **Features:** - Generates images from text prompts -- Supports configurable image size and quality +- Supports configurable model, image size, and quality - Returns base64-encoded images or URLs - Saves generated images to disk **Usage:** ```bash -# Use default localhost server +# FLUX.2 (default) python sync_image_gen.py -# Specify custom server URL -python sync_image_gen.py http://your-server:8000/v1 +# FLUX.1 +python sync_image_gen.py --model flux1 + +# Custom server and prompt +python sync_image_gen.py --base-url http://your-server:8000/v1 --prompt "A sunset" ``` **API Endpoint:** `POST /v1/images/generations` @@ -228,7 +235,7 @@ You can customize these by: ## Common Parameters ### Image Generation -- `model`: Model identifier (e.g., "wan") +- `model`: Model identifier (e.g., "flux1", "flux2") - `prompt`: Text description - `n`: Number of images to generate - `size`: Image dimensions (e.g., "512x512", "1024x1024") diff --git a/examples/visual_gen/serve/configs/flux1.yml b/examples/visual_gen/serve/configs/flux1.yml new file mode 100644 index 000000000000..f97f0016e106 --- /dev/null +++ b/examples/visual_gen/serve/configs/flux1.yml @@ -0,0 +1,10 @@ +linear: + type: default +teacache: + enable_teacache: false + teacache_thresh: 0.2 +attention: + backend: VANILLA +parallel: + dit_cfg_size: 1 + dit_ulysses_size: 1 diff --git a/examples/visual_gen/serve/sync_image_gen.py b/examples/visual_gen/serve/sync_image_gen.py index ca3c33d543c0..9f9c971dc7d9 100755 --- a/examples/visual_gen/serve/sync_image_gen.py +++ b/examples/visual_gen/serve/sync_image_gen.py @@ -3,9 +3,19 @@ Tests: - POST /v1/images/generations - Generate images from text -- POST /v1/images/edits - Edit images with text prompts + +Examples: + # FLUX.2 (default) + python sync_image_gen.py + + # FLUX.1 + python sync_image_gen.py --model flux1 + + # Custom server and prompt + python sync_image_gen.py --base-url http://your-server:8000/v1 --prompt "A sunset" """ +import argparse import base64 import sys @@ -31,6 +41,7 @@ def test_image_generation( client = openai.OpenAI(base_url=base_url, api_key="tensorrt_llm") print("\n1. Generating image...") + print(f" Model: {model}") print(f" Prompt: {prompt}") print(f" Size: {size}") print(f" Quality: {quality}") @@ -78,14 +89,55 @@ def test_image_generation( if __name__ == "__main__": - # Parse command line arguments - base_url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000/v1" + parser = argparse.ArgumentParser( + description="Test image generation API (FLUX.1 / FLUX.2)", + ) + parser.add_argument( + "--base-url", + type=str, + default="http://localhost:8000/v1", + help="Base URL of the API server", + ) + parser.add_argument( + "--model", + type=str, + default="flux2", + help="Model name (e.g., flux1, flux2)", + ) + parser.add_argument( + "--prompt", + type=str, + default="A lovely cat lying on a sofa", + help="Text prompt for image generation", + ) + parser.add_argument( + "--size", + type=str, + default="512x512", + help="Image size in WxH format (e.g., 512x512, 1024x1024)", + ) + parser.add_argument( + "--output", + type=str, + default="output_generation.png", + help="Output image file path", + ) + + args = parser.parse_args() print("\n" + "=" * 80) print("OpenAI SDK - Image Generation Tests") print("=" * 80) - print(f"Base URL: {base_url}") + print(f"Base URL: {args.base_url}") + print(f"Model: {args.model}") print() - # Test image generation - test_image_generation(base_url=base_url) + success = test_image_generation( + base_url=args.base_url, + model=args.model, + prompt=args.prompt, + size=args.size, + output_file=args.output, + ) + + sys.exit(0 if success else 1) diff --git a/examples/visual_gen/visual_gen_examples.sh b/examples/visual_gen/visual_gen_examples.sh index b7697602036e..a55342ad8f24 100755 --- a/examples/visual_gen/visual_gen_examples.sh +++ b/examples/visual_gen/visual_gen_examples.sh @@ -232,6 +232,56 @@ python ${PROJECT_ROOT}/examples/visual_gen/visual_gen_wan_i2v.py \ --guidance_scale_2 5.0 \ --boundary_ratio 0.85 +############################################# +# FLUX.1 Text-to-Image Examples +############################################# + +echo "" +echo "=== FLUX.1 Example 1: Baseline ===" +python ${PROJECT_ROOT}/examples/visual_gen/visual_gen_flux.py \ + --height 1024 \ + --width 1024 \ + --prompt "A cat holding a sign that says hello world" \ + --output_path flux1_cat_sign.png \ + --model_path ${MODEL_ROOT}/FLUX.1-dev/ \ + --guidance_scale 3.5 + +echo "" +echo "=== FLUX.1 Example 2: With FP8 Quantization ===" +python ${PROJECT_ROOT}/examples/visual_gen/visual_gen_flux.py \ + --height 1024 \ + --width 1024 \ + --prompt "A cat holding a sign that says hello world" \ + --output_path flux1_cat_sign_fp8.png \ + --model_path ${MODEL_ROOT}/FLUX.1-dev/ \ + --guidance_scale 3.5 \ + --linear_type trtllm-fp8-per-tensor + +############################################# +# FLUX.2 Text-to-Image Examples +############################################# + +echo "" +echo "=== FLUX.2 Example 1: Baseline ===" +python ${PROJECT_ROOT}/examples/visual_gen/visual_gen_flux.py \ + --height 1024 \ + --width 1024 \ + --prompt "A cat holding a sign that says hello world" \ + --output_path flux2_cat_sign.png \ + --model_path ${MODEL_ROOT}/FLUX.2-dev/ \ + --guidance_scale 4.0 + +echo "" +echo "=== FLUX.2 Example 2: With TeaCache ===" +python ${PROJECT_ROOT}/examples/visual_gen/visual_gen_flux.py \ + --height 1024 \ + --width 1024 \ + --prompt "A cat holding a sign that says hello world" \ + --output_path flux2_cat_sign_teacache.png \ + --model_path ${MODEL_ROOT}/FLUX.2-dev/ \ + --guidance_scale 4.0 \ + --enable_teacache + echo "" echo "============================================" echo "All examples completed successfully!" diff --git a/examples/visual_gen/visual_gen_flux.py b/examples/visual_gen/visual_gen_flux.py new file mode 100755 index 000000000000..235a576dd980 --- /dev/null +++ b/examples/visual_gen/visual_gen_flux.py @@ -0,0 +1,354 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FLUX Text-to-Image generation using TensorRT-LLM Visual Generation. + +Supports both FLUX.1 and FLUX.2 models. The pipeline type is auto-detected +from the model checkpoint (model_index.json). + +Single image mode: + python visual_gen_flux.py --model_path black-forest-labs/FLUX.1-dev \ + --prompt "A cat sitting on a windowsill" --guidance_scale 3.5 + + python visual_gen_flux.py --model_path black-forest-labs/FLUX.2-dev \ + --prompt "A cat sitting on a windowsill" --guidance_scale 4.0 + + # With FP8 quantization + python visual_gen_flux.py --model_path black-forest-labs/FLUX.2-dev \ + --prompt "A cat" --linear_type trtllm-fp8-per-tensor + +Batch mode (generates multiple images from a prompts file): + python visual_gen_flux.py --model_path black-forest-labs/FLUX.1-dev \ + --prompts_file prompts.txt --output_dir results/bf16/ --seed 42 + + # With FP8 quantization + python visual_gen_flux.py --model_path black-forest-labs/FLUX.2-dev \ + --prompts_file prompts.txt --output_dir results/fp8/ \ + --linear_type trtllm-fp8-per-tensor + + # Multi-GPU with CFG + Ulysses parallelism + python visual_gen_flux.py --model_path black-forest-labs/FLUX.1-dev \ + --prompts_file prompts.txt --output_dir results/ \ + --cfg_size 2 --ulysses_size 2 +""" + +import argparse +import json +import os +import time + +from output_handler import OutputHandler + +from tensorrt_llm import logger +from tensorrt_llm.llmapi.visual_gen import VisualGen, VisualGenParams + +# Set logger level to ensure timing logs are printed +logger.set_level("info") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="TRTLLM VisualGen - FLUX Text-to-Image Inference Example (FLUX.1 / FLUX.2)" + ) + + # Model & Input + parser.add_argument( + "--model_path", + type=str, + required=True, + help="Local path or HuggingFace Hub model ID " + "(e.g., black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-dev)", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + help="HuggingFace Hub revision (branch, tag, or commit SHA)", + ) + + # Single image mode + parser.add_argument( + "--prompt", type=str, default=None, help="Text prompt for single image generation" + ) + parser.add_argument( + "--output_path", + type=str, + default="output.png", + help="Path to save the output image (single image mode)", + ) + + # Batch mode + parser.add_argument( + "--prompts_file", + type=str, + default=None, + help="File with prompts (one per line) for batch generation", + ) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="Output directory for batch mode (images named 00.png, 01.png, ...)", + ) + parser.add_argument( + "--num_prompts", + type=int, + default=None, + help="Limit number of prompts from file (batch mode)", + ) + + # Generation Params + parser.add_argument("--height", type=int, default=1024, help="Image height") + parser.add_argument("--width", type=int, default=1024, help="Image width") + parser.add_argument("--steps", type=int, default=50, help="Number of denoising steps") + parser.add_argument( + "--guidance_scale", + type=float, + default=3.5, + help="Embedded guidance scale (3.5 for FLUX.1-dev, 4.0 for FLUX.2-dev)", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed") + + # TeaCache Arguments + parser.add_argument( + "--enable_teacache", action="store_true", help="Enable TeaCache acceleration" + ) + parser.add_argument( + "--teacache_thresh", + type=float, + default=0.2, + help="TeaCache similarity threshold (rel_l1_thresh)", + ) + + # Quantization + parser.add_argument( + "--linear_type", + type=str, + default="default", + choices=["default", "trtllm-fp8-per-tensor", "trtllm-fp8-blockwise", "svd-nvfp4"], + help="Linear layer quantization type", + ) + + # Attention Backend + parser.add_argument( + "--attention_backend", + type=str, + default="VANILLA", + choices=["VANILLA", "TRTLLM"], + help="Attention backend (VANILLA: PyTorch SDPA, TRTLLM: optimized kernels). " + "Note: TRTLLM automatically falls back to VANILLA for cross-attention.", + ) + + # torch.compile + parser.add_argument( + "--disable_torch_compile", action="store_true", help="Disable TorchCompile acceleration" + ) + parser.add_argument( + "--torch_compile_mode", + type=str, + default="default", + help="Torch compile mode", + choices=["default", "max-autotune", "reduce-overhead"], + ) + + # Warmup + parser.add_argument( + "--warmup_steps", + type=int, + default=1, + help="Number of warmup steps (0 to disable)", + ) + + # Parallelism + parser.add_argument( + "--cfg_size", + type=int, + default=1, + choices=[1, 2], + help="CFG parallel size (1 or 2). Set to 2 for CFG Parallelism.", + ) + parser.add_argument( + "--ulysses_size", + type=int, + default=1, + help="Ulysses (sequence) parallel size within each CFG group.", + ) + + args = parser.parse_args() + + # Validate: either --prompt or --prompts_file is required + if args.prompt is None and args.prompts_file is None: + parser.error("Either --prompt or --prompts_file is required") + if args.prompt is not None and args.prompts_file is not None: + parser.error("--prompt and --prompts_file are mutually exclusive") + if args.prompts_file is not None and args.output_dir is None: + parser.error("--output_dir is required when using --prompts_file") + + return args + + +def load_prompts(prompts_file, num_prompts=None): + """Load prompts from file (one per line, skip empty/comments).""" + with open(prompts_file) as f: + prompts = [line.strip() for line in f if line.strip() and not line.startswith("#")] + if num_prompts is not None: + prompts = prompts[:num_prompts] + return prompts + + +def build_diffusion_config(args): + """Build diffusion_config dict from parsed args.""" + # Convert linear_type to quant_config + quant_config = None + if args.linear_type == "trtllm-fp8-per-tensor": + quant_config = {"quant_algo": "FP8", "dynamic": True} + elif args.linear_type == "trtllm-fp8-blockwise": + quant_config = {"quant_algo": "FP8_BLOCK_SCALES", "dynamic": True} + elif args.linear_type == "svd-nvfp4": + quant_config = {"quant_algo": "NVFP4", "dynamic": True} + + # Note: pipeline type (FLUX.1 vs FLUX.2) is auto-detected from model_index.json + diffusion_config = { + "revision": args.revision, + "attention": { + "backend": args.attention_backend, + }, + "teacache": { + "enable_teacache": args.enable_teacache, + "teacache_thresh": args.teacache_thresh, + }, + "parallel": { + "dit_cfg_size": args.cfg_size, + "dit_ulysses_size": args.ulysses_size, + }, + "pipeline": { + "enable_torch_compile": not args.disable_torch_compile, + "torch_compile_mode": args.torch_compile_mode, + "warmup_steps": args.warmup_steps, + }, + } + + if quant_config is not None: + diffusion_config["quant_config"] = quant_config + + return diffusion_config + + +def main(): + args = parse_args() + + # world_size = cfg_size * ulysses_size + n_workers = args.cfg_size * args.ulysses_size + + diffusion_config = build_diffusion_config(args) + + # Initialize VisualGen + logger.info( + f"Initializing VisualGen: world_size={n_workers} " + f"(cfg_size={args.cfg_size}, ulysses_size={args.ulysses_size})" + ) + visual_gen = VisualGen( + model_path=args.model_path, + n_workers=n_workers, + diffusion_config=diffusion_config, + ) + + try: + if args.prompts_file: + # Batch mode + prompts = load_prompts(args.prompts_file, args.num_prompts) + os.makedirs(args.output_dir, exist_ok=True) + + logger.info(f"Batch mode: {len(prompts)} prompts → {args.output_dir}") + logger.info(f"Resolution: {args.height}x{args.width}, Steps: {args.steps}") + + timing_records = [] + total_start = time.time() + + for i, prompt in enumerate(prompts): + logger.info(f"[{i + 1}/{len(prompts)}] {prompt[:60]}...") + start_time = time.time() + + output = visual_gen.generate( + inputs=prompt, + params=VisualGenParams( + height=args.height, + width=args.width, + num_inference_steps=args.steps, + guidance_scale=args.guidance_scale, + seed=args.seed + i, + ), + ) + + elapsed = time.time() - start_time + output_path = os.path.join(args.output_dir, f"{i:02d}.png") + OutputHandler.save(output, output_path) + logger.info(f" Saved {output_path} ({elapsed:.1f}s)") + + timing_records.append( + { + "index": i, + "prompt": prompt, + "time": round(elapsed, 2), + "seed": args.seed + i, + } + ) + + total_elapsed = time.time() - total_start + times = [r["time"] for r in timing_records] + + # Write timing metadata + timing_data = { + "images": timing_records, + "total_time": round(total_elapsed, 2), + "avg_time": round(sum(times) / len(times), 2) if times else 0, + "config": { + "model_path": args.model_path, + "linear_type": args.linear_type, + "attention_backend": args.attention_backend, + "height": args.height, + "width": args.width, + "steps": args.steps, + "guidance_scale": args.guidance_scale, + }, + } + timing_path = os.path.join(args.output_dir, "timing.json") + with open(timing_path, "w") as f: + json.dump(timing_data, f, indent=2) + + logger.info( + f"Batch complete: {len(prompts)} images in {total_elapsed:.1f}s " + f"(avg {timing_data['avg_time']:.1f}s/image)" + ) + logger.info(f"Timing saved to {timing_path}") + + else: + # Single image mode + logger.info(f"Generating image for prompt: '{args.prompt}'") + logger.info(f"Resolution: {args.height}x{args.width}, Steps: {args.steps}") + + start_time = time.time() + + output = visual_gen.generate( + inputs=args.prompt, + params=VisualGenParams( + height=args.height, + width=args.width, + num_inference_steps=args.steps, + guidance_scale=args.guidance_scale, + seed=args.seed, + ), + ) + + end_time = time.time() + logger.info(f"Generation completed in {end_time - start_time:.2f}s") + + OutputHandler.save(output, args.output_path) + + finally: + visual_gen.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index fd2d21642a3f..27409e93246a 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -29,7 +29,9 @@ class PipelineComponent(str, Enum): TRANSFORMER = "transformer" VAE = "vae" TEXT_ENCODER = "text_encoder" + TEXT_ENCODER_2 = "text_encoder_2" TOKENIZER = "tokenizer" + TOKENIZER_2 = "tokenizer_2" SCHEDULER = "scheduler" IMAGE_ENCODER = "image_encoder" IMAGE_PROCESSOR = "image_processor" diff --git a/tensorrt_llm/_torch/visual_gen/models/__init__.py b/tensorrt_llm/_torch/visual_gen/models/__init__.py index 5f726b84ec63..68d59132d94d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/__init__.py @@ -19,11 +19,14 @@ from ..pipeline import BasePipeline from ..pipeline_registry import AutoPipeline, register_pipeline +from .flux import Flux2Pipeline, FluxPipeline from .wan import WanImageToVideoPipeline, WanPipeline __all__ = [ "AutoPipeline", "BasePipeline", + "FluxPipeline", + "Flux2Pipeline", "WanPipeline", "WanImageToVideoPipeline", "register_pipeline", diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/__init__.py b/tensorrt_llm/_torch/visual_gen/models/flux/__init__.py new file mode 100644 index 000000000000..8d1eb08db7ec --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/flux/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .attention import Flux2ParallelSelfAttention, FluxJointAttention +from .pipeline_flux import FluxPipeline +from .pipeline_flux2 import Flux2Pipeline +from .pos_embed_flux import FluxPosEmbed +from .transformer_flux import FluxTransformer2DModel +from .transformer_flux2 import Flux2Transformer2DModel + +__all__ = [ + "FluxJointAttention", + "FluxPipeline", + "FluxTransformer2DModel", + "Flux2Pipeline", + "Flux2Transformer2DModel", + "Flux2ParallelSelfAttention", + "FluxPosEmbed", +] diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/attention.py b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py new file mode 100644 index 000000000000..238ba9a885bf --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/flux/attention.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FLUX attention modules: joint attention and parallel self-attention. + +Key Components: +- FluxJointAttention: Joint attention for dual-stream blocks (FLUX.1 and FLUX.2) +- Flux2ParallelSelfAttention: Fused QKV+MLP for FLUX.2 single-stream blocks +""" + +from typing import TYPE_CHECKING, Optional, Tuple, Union + +import torch + +from tensorrt_llm._torch.modules.linear import Linear, WeightMode, WeightsLoadingConfig +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.swiglu import swiglu +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode, apply_rotary_emb + +if TYPE_CHECKING: + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + + +# ============================================================================= +# Joint Attention (shared by FLUX.1 and FLUX.2 dual-stream blocks) +# ============================================================================= + + +class FluxJointAttention(Attention): + """Joint attention module for FLUX transformer models (FLUX.1 and FLUX.2). + + Extends base Attention with: + - Text-stream QKV projection (add_qkv_proj) for dual-stream blocks + - FLUX-style RoPE on concatenated text+image tokens + - pre_only mode for single-stream blocks (no output projection) + + For dual-stream blocks: Returns (img_attn_output, txt_attn_output) + For single-stream blocks: pre_only=True, returns attention output only + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + head_dim: int = 128, + bias: bool = False, + added_kv_proj_dim: Optional[int] = None, + eps: float = 1e-6, + pre_only: bool = False, + config: Optional["DiffusionModelConfig"] = None, + layer_idx: int = 0, + ): + super().__init__( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + head_dim=head_dim, + qkv_mode=QKVMode.FUSE_QKV, + qk_norm=True, + qk_norm_mode="per_head", + eps=eps, + bias=bias, + config=config, + layer_idx=layer_idx, + ) + + self.pre_only = pre_only + self.added_kv_proj_dim = added_kv_proj_dim + + # Delete output projection for single-stream blocks + if self.pre_only: + del self.to_out + + # Text-stream projections for joint attention (dual-stream blocks only) + if added_kv_proj_dim is not None: + self.add_qkv_proj = Linear( + added_kv_proj_dim, + 3 * self.q_dim, + bias=self.bias, + dtype=self.dtype, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + weights_loading_config=WeightsLoadingConfig( + weight_mode=WeightMode.FUSED_QKV_LINEAR + ), + fused_weight_shard_indices_mapping={ + "q": (0, self.q_dim), + "k": (self.q_dim, self.q_dim), + "v": (2 * self.q_dim, self.q_dim), + }, + ) + + self.norm_added_q = RMSNorm( + hidden_size=head_dim, eps=eps, dtype=self.dtype, has_weights=True + ) + self.norm_added_k = RMSNorm( + hidden_size=head_dim, eps=eps, dtype=self.dtype, has_weights=True + ) + + self.to_add_out = Linear( + self.q_dim, + added_kv_proj_dim, + bias=self.bias, + dtype=self.dtype, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Forward pass of joint attention. + + Args: + hidden_states: Image tokens [batch, img_seq, dim] + encoder_hidden_states: Text tokens [batch, txt_seq, dim] (for dual-stream) + attention_mask: Optional attention mask (unused, for API compat) + image_rotary_emb: Tuple of (cos, sin) for RoPE + + Returns: + For dual-stream: Tuple of (img_attn_output, txt_attn_output) + For single-stream: Attention output tensor + """ + batch_size = hidden_states.shape[0] + + # Image QKV via base (returns 3D), then reshape to 4D for per-head ops + query, key, value = self.get_qkv(hidden_states) + query = query.view(batch_size, -1, self.num_attention_heads, self.head_dim) + key = key.view(batch_size, -1, self.num_attention_heads, self.head_dim) + value = value.view(batch_size, -1, self.num_attention_heads, self.head_dim) + + # Per-head QK normalization via base (per_head mode operates on 4D) + query, key = self.apply_qk_norm(query, key) + + # Text QKV for joint attention (dual-stream blocks) + if encoder_hidden_states is not None and self.added_kv_proj_dim is not None: + txt_seq_len = encoder_hidden_states.shape[1] + + encoder_qkv = self.add_qkv_proj(encoder_hidden_states) + enc_q, enc_k, enc_v = encoder_qkv.chunk(3, dim=-1) + enc_q = enc_q.view(batch_size, -1, self.num_attention_heads, self.head_dim) + enc_k = enc_k.view(batch_size, -1, self.num_attention_heads, self.head_dim) + enc_v = enc_v.view(batch_size, -1, self.num_attention_heads, self.head_dim) + + enc_q = self.norm_added_q(enc_q.reshape(-1, enc_q.shape[-1])).view(enc_q.shape) + enc_k = self.norm_added_k(enc_k.reshape(-1, enc_k.shape[-1])).view(enc_k.shape) + + # Concatenate text + image for joint attention + query = torch.cat([enc_q, query], dim=1) + key = torch.cat([enc_k, key], dim=1) + value = torch.cat([enc_v, value], dim=1) + + # Apply RoPE + if image_rotary_emb is not None: + freqs_cos, freqs_sin = image_rotary_emb + query = apply_rotary_emb(query, freqs_cos, freqs_sin) + key = apply_rotary_emb(key, freqs_cos, freqs_sin) + + # Flatten 4D->3D for base _attn_impl + seq_len = query.shape[1] + query = query.flatten(2) + key = key.flatten(2) + value = value.flatten(2) + + hidden_states = self._attn_impl(query, key, value, batch_size, seq_len) + hidden_states = hidden_states.to(query.dtype) + + # Split and project outputs + if encoder_hidden_states is not None and self.added_kv_proj_dim is not None: + encoder_hidden_states_out, hidden_states = hidden_states.split( + [txt_seq_len, hidden_states.shape[1] - txt_seq_len], dim=1 + ) + + if not self.pre_only: + hidden_states = self.to_out[0](hidden_states) + encoder_hidden_states_out = self.to_add_out(encoder_hidden_states_out) + + return hidden_states, encoder_hidden_states_out + else: + if not self.pre_only: + hidden_states = self.to_out[0](hidden_states) + return hidden_states + + +# ============================================================================= +# Parallel Self-Attention (for FLUX.2 single-stream blocks) +# ============================================================================= + + +class Flux2ParallelSelfAttention(FluxJointAttention): + """FLUX.2 parallel self-attention for single-stream blocks. + + Uses fused QKV+MLP projection: to_qkv_mlp_proj + Output: concatenate attention + MLP outputs, then project with to_out + + This is a key architectural difference from FLUX.1: + - FLUX.1: Separate attention and FFN + - FLUX.2: Fused QKV+MLP projection for efficiency + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + head_dim: int = 128, + mlp_ratio: float = 3.0, + bias: bool = False, + eps: float = 1e-6, + config: Optional["DiffusionModelConfig"] = None, + layer_idx: int = 0, + ): + # Set MLP dims BEFORE super().__init__() — _init_qkv_proj() needs them + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.mlp_mult_factor = 2 # SwiGLU doubles input + + super().__init__( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + head_dim=head_dim, + bias=bias, + added_kv_proj_dim=None, # No text stream + eps=eps, + pre_only=True, # Deletes base to_out + config=config, + layer_idx=layer_idx, + ) + + # Combined output: [q_dim + mlp_hidden_dim] -> [hidden_size] + self.to_out = Linear( + self.q_dim + self.mlp_hidden_dim, + hidden_size, + bias=bias, + dtype=self.dtype, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + ) + + def _init_qkv_proj(self): + """Override: fused QKV+MLP projection instead of standard QKV.""" + qkv_dim = 3 * self.q_dim + mlp_in_dim = self.mlp_hidden_dim * self.mlp_mult_factor + self.to_qkv_mlp_proj = Linear( + self.hidden_size, + qkv_dim + mlp_in_dim, + bias=self.bias, + dtype=self.dtype, + quant_config=self.quant_config, + skip_create_weights_in_init=self.skip_create_weights_in_init, + force_dynamic_quantization=self.force_dynamic_quantization, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + """ + Args: + hidden_states: [batch, seq, dim] + attention_mask: Optional attention mask + image_rotary_emb: Tuple of (freqs_cos, freqs_sin) + + Returns: + hidden_states [batch, seq, dim] + """ + batch_size = hidden_states.shape[0] + + # Fused QKV + MLP projection + proj_out = self.to_qkv_mlp_proj(hidden_states) + qkv, mlp_hidden = torch.split( + proj_out, [3 * self.q_dim, self.mlp_hidden_dim * self.mlp_mult_factor], dim=-1 + ) + + # Split QKV -> 4D + q, k, v = qkv.chunk(3, dim=-1) + q = q.view(batch_size, -1, self.num_attention_heads, self.head_dim) + k = k.view(batch_size, -1, self.num_attention_heads, self.head_dim) + v = v.view(batch_size, -1, self.num_attention_heads, self.head_dim) + + # Per-head QK norm (inherited) + q, k = self.apply_qk_norm(q, k) + + # RoPE + if image_rotary_emb is not None: + freqs_cos, freqs_sin = image_rotary_emb + q = apply_rotary_emb(q, freqs_cos, freqs_sin) + k = apply_rotary_emb(k, freqs_cos, freqs_sin) + + # Flatten 4D->3D for _attn_impl + seq_len = q.shape[1] + q, k, v = q.flatten(2), k.flatten(2), v.flatten(2) + + # Backend dispatch (inherited — handles layout, Ulysses, etc.) + attn_out = self._attn_impl(q, k, v, batch_size, seq_len) + attn_out = attn_out.to(q.dtype) + + # Parallel MLP path (reshape to 2D for Triton kernel, then back) + shape = mlp_hidden.shape + mlp_out = swiglu(mlp_hidden.reshape(-1, shape[-1])).reshape(*shape[:-1], -1) + + # Concatenate + project + return self.to_out(torch.cat([attn_out, mlp_out], dim=-1)) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py new file mode 100644 index 000000000000..3ffb7f6b4635 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -0,0 +1,504 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FLUX.1 Pipeline implementation following WAN pattern.""" + +import time +from typing import Optional, Tuple + +import numpy as np +import torch +from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler +from diffusers.utils.torch_utils import randn_tensor +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast + +from tensorrt_llm._torch.visual_gen.config import PipelineComponent +from tensorrt_llm._torch.visual_gen.output import MediaOutput +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline +from tensorrt_llm._torch.visual_gen.teacache import ExtractorConfig, register_extractor_from_config +from tensorrt_llm.logger import logger + +from .transformer_flux import FluxTransformer2DModel + +# TeaCache coefficients for FLUX.1 variants +FLUX_TEACACHE_COEFFICIENTS = { + "dev": { + "ret_steps": [2.57151496e05, -3.54229917e04, 1.40286849e03, -1.35890334e01, 1.32517977e-01], + "standard": [2.57151496e05, -3.54229917e04, 1.40286849e03, -1.35890334e01, 1.32517977e-01], + }, + "schnell": { + "ret_steps": [1.0, 0.0], # Schnell is already fast, minimal caching + "standard": [1.0, 0.0], + }, +} + + +@register_pipeline("FluxPipeline") +class FluxPipeline(BasePipeline): + """FLUX.1 Text-to-Image Pipeline. + + Supports FLUX.1-dev (50 steps, guidance) and FLUX.1-schnell (4 steps, no guidance). + """ + + @staticmethod + def _compute_flux_timestep_embedding(module, timestep, guidance=None): + """Compute timestep embedding for FLUX transformer. + + FLUX combines timestep and guidance embeddings. + + Args: + module: FluxTransformer2DModel instance + timestep: Timestep tensor [B] + guidance: Guidance scale tensor [B] (optional) + + Returns: + Combined timestep embedding for TeaCache distance calculation + """ + # Cast to embedder's dtype (avoid int8 quantized layers) + te_dtype = next(iter(module.time_text_embed.parameters())).dtype + if timestep.dtype != te_dtype and te_dtype != torch.int8: + timestep = timestep.to(te_dtype) + + temb = module.time_text_embed(timestep) + + if module.guidance_embeds and guidance is not None: + if guidance.dtype != te_dtype and te_dtype != torch.int8: + guidance = guidance.to(te_dtype) + temb = temb + module.guidance_embed(guidance) + + return temb + + @property + def dtype(self): + return self.model_config.torch_dtype + + @property + def device(self): + if self.transformer is not None: + return next(self.transformer.parameters()).device + return torch.device("cuda:0") + + @property + def common_warmup_shapes(self) -> list: + """Return list of common warmup shapes (height, width, num_frames).""" + return [(1024, 1024, 1)] + + def _init_transformer(self) -> None: + """Initialize FLUX transformer with quantization support.""" + logger.info("Creating FLUX transformer with quantization support...") + self.transformer = FluxTransformer2DModel(model_config=self.model_config) + + def _run_warmup(self, warmup_steps: int) -> None: + """Run warmup inference to trigger torch.compile and CUDA init.""" + for height, width, _ in self.common_warmup_shapes: + logger.info(f"Warmup: FLUX.1 {height}x{width}, {warmup_steps} steps") + with torch.no_grad(): + self.forward( + prompt="warmup", + height=height, + width=width, + num_inference_steps=warmup_steps, + guidance_scale=3.5, + seed=0, + max_sequence_length=512, + ) + + def load_standard_components( + self, + checkpoint_dir: str, + device: torch.device, + skip_components: Optional[list] = None, + ) -> None: + """Load VAE, text encoders, tokenizers, and scheduler from checkpoint.""" + skip_components = skip_components or [] + + # CLIP tokenizer and text encoder (for pooled embeddings) + if PipelineComponent.TOKENIZER not in skip_components: + logger.info("Loading CLIP tokenizer...") + self.tokenizer = CLIPTokenizer.from_pretrained( + checkpoint_dir, subfolder=PipelineComponent.TOKENIZER + ) + + if PipelineComponent.TEXT_ENCODER not in skip_components: + logger.info("Loading CLIP text encoder...") + self.text_encoder = CLIPTextModel.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.TEXT_ENCODER, + torch_dtype=self.model_config.torch_dtype, + ).to(device) + + # T5 tokenizer and text encoder (for sequence embeddings) + if PipelineComponent.TOKENIZER_2 not in skip_components: + logger.info("Loading T5 tokenizer...") + self.tokenizer_2 = T5TokenizerFast.from_pretrained( + checkpoint_dir, subfolder=PipelineComponent.TOKENIZER_2 + ) + + if PipelineComponent.TEXT_ENCODER_2 not in skip_components: + logger.info("Loading T5 text encoder...") + self.text_encoder_2 = T5EncoderModel.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.TEXT_ENCODER_2, + torch_dtype=self.model_config.torch_dtype, + ).to(device) + + # VAE + if PipelineComponent.VAE not in skip_components: + logger.info("Loading VAE...") + self.vae = AutoencoderKL.from_pretrained( + checkpoint_dir, subfolder=PipelineComponent.VAE, torch_dtype=torch.bfloat16 + ).to(device) + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + # Scheduler + if PipelineComponent.SCHEDULER not in skip_components: + logger.info("Loading scheduler...") + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + checkpoint_dir, subfolder=PipelineComponent.SCHEDULER + ) + + # Default config values + self.max_sequence_length = 512 + self.default_height = 1024 + self.default_width = 1024 + + def load_weights(self, weights: dict) -> None: + """Load transformer weights.""" + if self.transformer is not None and hasattr(self.transformer, "load_weights"): + logger.info("Loading transformer weights...") + transformer_weights = weights.get("transformer", weights) + self.transformer.load_weights(transformer_weights) + logger.info("Transformer weights loaded successfully.") + + self._target_dtype = self.model_config.torch_dtype + + if self.transformer is not None: + self.transformer.eval() + + def post_load_weights(self) -> None: + """Post-load setup: TeaCache registration.""" + super().post_load_weights() + if self.transformer is not None: + # Register TeaCache extractor for FLUX (must be after device placement) + register_extractor_from_config( + ExtractorConfig( + model_class_name="FluxTransformer2DModel", + timestep_embed_fn=self._compute_flux_timestep_embedding, + guidance_param_name="guidance", + forward_params=[ + "hidden_states", + "encoder_hidden_states", + "pooled_projections", + "timestep", + "img_ids", + "txt_ids", + "guidance", + "return_dict", + ], + return_dict_default=False, + ) + ) + + # Enable TeaCache with FLUX-specific coefficients + self._setup_teacache(self.transformer, coefficients=FLUX_TEACACHE_COEFFICIENTS) + + def infer(self, req): + """Run inference from DiffusionRequest.""" + return self.forward( + prompt=req.prompt, + height=req.height, + width=req.width, + num_inference_steps=req.num_inference_steps, + guidance_scale=req.guidance_scale, + seed=req.seed, + max_sequence_length=req.max_sequence_length, + ) + + @torch.inference_mode() + def forward( + self, + prompt: str, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 3.5, + seed: int = 42, + max_sequence_length: int = 512, + ): + """Generate image from text prompt. + + Args: + prompt: Text prompt for image generation + height: Output image height (default: 1024) + width: Output image width (default: 1024) + num_inference_steps: Number of denoising steps (50 for dev, 4 for schnell) + guidance_scale: Embedded guidance scale (3.5 for dev) + seed: Random seed for reproducibility + max_sequence_length: Maximum text sequence length + + Returns: + MediaOutput with image tensor + """ + pipeline_start = time.time() + generator = torch.Generator(device=self.device).manual_seed(seed) + + # Encode prompt + logger.info("Encoding prompt...") + encode_start = time.time() + prompt_embeds, pooled_prompt_embeds, text_ids = self._encode_prompt( + prompt, max_sequence_length + ) + logger.info(f"Prompt encoding completed in {time.time() - encode_start:.2f}s") + + # Prepare latents + latents, latent_ids = self._prepare_latents(height, width, generator) + logger.info(f"Latents shape: {latents.shape}") + + # Prepare timesteps with dynamic shifting (FLUX uses mu parameter) + image_seq_len = latents.shape[1] + mu = self._compute_mu(image_seq_len) + + # Match HF: create sigmas for non-flow mode, or None for flow mode + use_flow_sigmas = getattr(self.scheduler.config, "use_flow_sigmas", None) + if use_flow_sigmas: + # Flow mode: let scheduler compute sigmas internally + self.scheduler.set_timesteps(num_inference_steps, device=self.device, mu=mu) + else: + # Non-flow mode: provide linear sigmas + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + self.scheduler.set_timesteps(sigmas=sigmas, device=self.device, mu=mu) + + timesteps = self.scheduler.timesteps + + # Prepare guidance (embedded guidance for FLUX) + guidance = None + if self.transformer.guidance_embeds: + guidance = torch.full( + [latents.shape[0]], guidance_scale, device=self.device, dtype=torch.float32 + ) + + # Denoising loop + def forward_fn( + latents, extra_stream_latents, timestep, encoder_hidden_states, extra_tensors + ): + """Forward function for FLUX transformer.""" + return self.transformer( + hidden_states=latents, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_prompt_embeds, + timestep=timestep / 1000, # FLUX expects normalized timesteps + img_ids=latent_ids, + txt_ids=text_ids, + guidance=guidance, + return_dict=False, + )[0] + + latents = self.denoise( + latents=latents, + scheduler=self.scheduler, + prompt_embeds=prompt_embeds, + guidance_scale=1.0, # No CFG: guidance is embedded + forward_fn=forward_fn, + timesteps=timesteps, + ) + + # Decode + logger.info("Decoding image...") + decode_start = time.time() + image = self.decode_latents(latents, lambda lat: self._decode_latents(lat, height, width)) + + if self.rank == 0: + logger.info(f"Image decoded in {time.time() - decode_start:.2f}s") + logger.info(f"Total pipeline time: {time.time() - pipeline_start:.2f}s") + + return MediaOutput(image=image) + + def _encode_prompt( + self, + prompt: str, + max_sequence_length: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode prompt using CLIP and T5. + + Args: + prompt: Text prompt + max_sequence_length: Maximum T5 sequence length + + Returns: + Tuple of (T5 embeddings, CLIP pooled embeddings, text position IDs) + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + + # CLIP encoding (pooled embeddings) + clip_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + clip_input_ids = clip_inputs.input_ids.to(self.device) + + clip_outputs = self.text_encoder(clip_input_ids, output_hidden_states=False) + pooled_prompt_embeds = clip_outputs.pooler_output.to(self.dtype) + + # T5 encoding (sequence embeddings) + t5_inputs = self.tokenizer_2( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_length=False, + return_overflowing_tokens=False, + return_tensors="pt", + ) + t5_input_ids = t5_inputs.input_ids.to(self.device) + + # NOTE: HF diffusers does NOT pass attention_mask to T5 for FLUX.1. + # T5 treats padding tokens as real tokens, and the transformer sees + # their non-zero embeddings. Matching HF behavior for PSNR parity. + t5_outputs = self.text_encoder_2(t5_input_ids, output_hidden_states=False) + prompt_embeds = t5_outputs.last_hidden_state.to(self.dtype) + + # Prepare text position IDs + text_ids = self._prepare_text_ids(prompt_embeds) + text_ids = text_ids.to(self.device) + + return prompt_embeds, pooled_prompt_embeds, text_ids + + def _compute_mu(self, image_seq_len: int) -> float: + """Compute mu parameter for FLUX's dynamic timestep shifting. + + FLUX uses flow matching with dynamic shifting based on image resolution. + Formula matches HuggingFace diffusers implementation. + + Args: + image_seq_len: Number of latent patches (packed format) + + Returns: + mu value for scheduler + """ + # HuggingFace formula: mu = m * image_seq_len + b + m = 0.000169271 + b = 0.456666667 + return float(m * image_seq_len + b) + + def _prepare_text_ids(self, text_embeds: torch.Tensor) -> torch.Tensor: + """Prepare 3D position IDs for text tokens. + + Returns 2D tensor [seq_len, 3] (HF expects unbatched). + """ + batch_size, seq_len, _ = text_embeds.shape + text_ids = torch.zeros(seq_len, 3, device=text_embeds.device) + return text_ids + + def _prepare_latent_ids(self, height: int, width: int) -> torch.Tensor: + """Prepare 3D position IDs for packed latent patches. + + FLUX uses 2x2 spatial packing, so IDs are for (H/2, W/2) grid. + HF convention: index 0 unused, index 1 = row, index 2 = col. + + Returns 2D tensor [seq_len, 3] (HF expects unbatched). + """ + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + + # Packed dimensions (2x2 packing) + packed_h = latent_height // 2 + packed_w = latent_width // 2 + + # Create grid with HF's index convention + latent_ids = torch.zeros(packed_h, packed_w, 3) + latent_ids[..., 1] = torch.arange(packed_h)[:, None] # Row index + latent_ids[..., 2] = torch.arange(packed_w)[None, :] # Col index + + latent_ids = latent_ids.reshape(-1, 3) + return latent_ids # [seq_len, 3] + + def _pack_latents( + self, + latents: torch.Tensor, + batch_size: int, + num_channels: int, + height: int, + width: int, + ) -> torch.Tensor: + """Pack latents from VAE spatial format to FLUX sequence format. + + FLUX uses 2x2 spatial packing: + VAE format: [B, 16, H, W] -> FLUX format: [B, (H/2)*(W/2), 64] + """ + # [B, C, H, W] -> [B, C, H/2, 2, W/2, 2] + latents = latents.view(batch_size, num_channels, height // 2, 2, width // 2, 2) + # -> [B, H/2, W/2, C, 2, 2] + latents = latents.permute(0, 2, 4, 1, 3, 5) + # -> [B, (H/2)*(W/2), C*4] + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels * 4) + return latents + + def _unpack_latents(self, latents: torch.Tensor, height: int, width: int) -> torch.Tensor: + """Unpack latents from FLUX sequence format to VAE spatial format. + + FLUX format: [B, (H/2)*(W/2), 64] -> VAE format: [B, 16, H, W] + """ + batch_size, num_patches, channels = latents.shape + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + + # [B, seq_len, 64] -> [B, H/2, W/2, 16, 2, 2] + latents = latents.view( + batch_size, latent_height // 2, latent_width // 2, channels // 4, 2, 2 + ) + # -> [B, 16, H/2, 2, W/2, 2] + latents = latents.permute(0, 3, 1, 4, 2, 5) + # -> [B, 16, H, W] + latents = latents.reshape(batch_size, channels // 4, latent_height, latent_width) + return latents + + def _prepare_latents( + self, + height: int, + width: int, + generator: torch.Generator, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Prepare random latents in FLUX packed format and position IDs.""" + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + + # Use VAE channels (16), not transformer channels (64) + # The packing will convert 16 -> 64 + vae_channels = self.vae.config.latent_channels # 16 + + # Create random latents in VAE spatial format [B, 16, H, W] + shape = (1, vae_channels, latent_height, latent_width) + latents = randn_tensor(shape, generator=generator, device=self.device, dtype=self.dtype) + + # Prepare position IDs for packed format + latent_ids = self._prepare_latent_ids(height, width) + latent_ids = latent_ids.to(self.device) + + # Pack latents to FLUX sequence format [B, seq_len, 64] + latents = self._pack_latents(latents, 1, vae_channels, latent_height, latent_width) + + return latents, latent_ids + + def _decode_latents(self, latents: torch.Tensor, height: int, width: int) -> torch.Tensor: + """Decode latents to image tensor.""" + # Unpack latents: (batch, seq_len, channels) -> (batch, channels, h, w) + latents = self._unpack_latents(latents, height, width) + + # Scale latents (must include shift_factor, matching HF diffusers) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + # VAE decode + latents = latents.to(self.vae.dtype) + image = self.vae.decode(latents, return_dict=False)[0] + + # Post-process to tensor (H, W, C) uint8 + image = (image / 2 + 0.5).clamp(0, 1) + image = image.permute(0, 2, 3, 1) # (B, C, H, W) -> (B, H, W, C) + image = (image * 255).round().to(torch.uint8) + + return image[0] # Remove batch dimension diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py new file mode 100644 index 000000000000..a6dccd846f3f --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -0,0 +1,644 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FLUX.2 Pipeline implementation following WAN pattern. + +This pipeline uses the TRT-LLM FLUX.2 transformer implementation. + +Supported variants: +- FLUX.2-dev (35B): Mistral3 text encoder, layers (10, 20, 30), embedded guidance +- FLUX.2-klein (4B/9B): Qwen3 text encoder, layers (9, 18, 27), no guidance + +The text encoder type and hidden state layers are auto-detected from model_index.json. + +Key differences from FLUX.1: +- Text encoder: Mistral3 or Qwen3 (decoder-only, used as encoder via hidden state extraction) +- Multi-layer fusion: 3 hidden layers stacked -> 3 x hidden_dim +- No pooled embeddings: Guidance is handled via timestep embedding (dev) or not at all (klein) +- 4-axis RoPE: (32, 32, 32, 32) instead of 3-axis +""" + +import json +import os +import time +from typing import List, Optional, Tuple + +import numpy as np +import torch +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 +from diffusers.utils.torch_utils import randn_tensor +from transformers import ( + AutoModelForCausalLM, + AutoProcessor, + AutoTokenizer, + Mistral3ForConditionalGeneration, +) + +from tensorrt_llm._torch.visual_gen.config import PipelineComponent +from tensorrt_llm._torch.visual_gen.output import MediaOutput +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline +from tensorrt_llm._torch.visual_gen.teacache import ExtractorConfig, register_extractor_from_config +from tensorrt_llm.logger import logger + +from .transformer_flux2 import Flux2Transformer2DModel + +# TeaCache coefficients for FLUX.2 +FLUX2_TEACACHE_COEFFICIENTS = { + "dev": { + "ret_steps": [2.57151496e05, -3.54229917e04, 1.40286849e03, -1.35890334e01, 1.32517977e-01], + "standard": [2.57151496e05, -3.54229917e04, 1.40286849e03, -1.35890334e01, 1.32517977e-01], + }, +} + +# System message for Mistral3 chat template (matches HF diffusers exactly) +SYSTEM_MESSAGE = ( + "You are an AI that reasons about image descriptions. You give structured " + "responses focusing on object relationships, object\nattribution and actions " + "without speculation." +) + + +def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: + """Compute empirical mu for scheduler shift (matches HF diffusers exactly).""" + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + + if image_seq_len > 4300: + return float(a2 * image_seq_len + b2) + + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + return float(a * num_steps + b) + + +def format_input(prompts: List[str], system_message: str) -> List[List[dict]]: + """Format prompts for Mistral3 chat template (PixtralProcessor format).""" + return [ + [ + {"role": "system", "content": [{"type": "text", "text": system_message}]}, + {"role": "user", "content": [{"type": "text", "text": prompt}]}, + ] + for prompt in prompts + ] + + +@register_pipeline("Flux2Pipeline") +class Flux2Pipeline(BasePipeline): + """FLUX.2 Text-to-Image Pipeline. + + Supports FLUX.2 model variants: + - FLUX.2-dev (35B): guidance_embeds=True, embedded guidance + - FLUX.2-klein (4B/9B): guidance_embeds=False, no guidance + + Uses Mistral3 for text encoding and native Flux2Transformer2DModel. + Follows WAN pipeline pattern for DiffusionModelLoader integration. + """ + + # Hidden state layers per text encoder type (auto-detected at load time) + _TEXT_ENCODER_CONFIG = { + "Mistral3ForConditionalGeneration": { + "hidden_layers": (10, 20, 30), + "system_message": SYSTEM_MESSAGE, + }, + "Qwen3ForCausalLM": { + "hidden_layers": (9, 18, 27), + "system_message": None, # Qwen3 uses simple user message + }, + } + # Default for backward compatibility (FLUX.2-dev) + HIDDEN_STATE_LAYERS: Tuple[int, ...] = (10, 20, 30) + + @staticmethod + def _compute_flux2_timestep_embedding(module, timestep, guidance=None): + """Compute timestep embedding for FLUX.2 transformer. + + Always uses time_guidance_embed (handles both guided and unguided variants). + + Args: + module: Flux2Transformer2DModel instance + timestep: Timestep tensor [B] + guidance: Guidance scale tensor [B] (optional, None for klein) + + Returns: + Timestep embedding for TeaCache distance calculation + """ + embed = module.time_guidance_embed + te_dtype = next(embed.timestep_embedder.linear_1.parameters()).dtype + if te_dtype != torch.int8: + t = timestep.to(te_dtype) + g = guidance.to(te_dtype) if guidance is not None else None + else: + t, g = timestep, guidance + return embed(t, g) + + @property + def dtype(self): + return self.model_config.torch_dtype + + @property + def device(self): + if self.transformer is not None: + return next(self.transformer.parameters()).device + return torch.device("cuda:0") + + @property + def common_warmup_shapes(self) -> list: + """Return list of common warmup shapes (height, width, num_frames).""" + return [(1024, 1024, 1)] + + def _init_transformer(self) -> None: + """Initialize FLUX.2 transformer with quantization support.""" + logger.info("Creating FLUX.2 transformer with quantization support...") + self.transformer = Flux2Transformer2DModel(model_config=self.model_config) + + def _run_warmup(self, warmup_steps: int) -> None: + """Run warmup inference to trigger torch.compile and CUDA init.""" + for height, width, _ in self.common_warmup_shapes: + logger.info(f"Warmup: FLUX.2 {height}x{width}, {warmup_steps} steps") + with torch.no_grad(): + self.forward( + prompt="warmup", + height=height, + width=width, + num_inference_steps=warmup_steps, + guidance_scale=3.5, + seed=0, + max_sequence_length=512, + ) + + def _detect_text_encoder_type(self, checkpoint_dir: str) -> str: + """Detect text encoder class from model_index.json.""" + model_index_path = os.path.join(checkpoint_dir, "model_index.json") + if os.path.exists(model_index_path): + with open(model_index_path) as f: + index = json.load(f) + text_encoder_info = index.get("text_encoder", []) + if len(text_encoder_info) >= 2: + return text_encoder_info[1] + return "Mistral3ForConditionalGeneration" # Default for FLUX.2-dev + + def load_standard_components( + self, + checkpoint_dir: str, + device: torch.device, + skip_components: Optional[list] = None, + ) -> None: + """Load VAE, text encoder, tokenizer, and scheduler from checkpoint. + + Auto-detects text encoder type from model_index.json: + - Mistral3ForConditionalGeneration: FLUX.2-dev (uses AutoProcessor) + - Qwen3ForCausalLM: FLUX.2-klein (uses AutoTokenizer) + """ + skip_components = skip_components or [] + + # Auto-detect text encoder type and configure hidden state layers + self._text_encoder_class = self._detect_text_encoder_type(checkpoint_dir) + te_config = self._TEXT_ENCODER_CONFIG.get(self._text_encoder_class, {}) + self.HIDDEN_STATE_LAYERS = te_config.get("hidden_layers", (10, 20, 30)) + logger.info( + f"Detected text encoder: {self._text_encoder_class}, " + f"hidden layers: {self.HIDDEN_STATE_LAYERS}" + ) + + # Tokenizer (type depends on text encoder) + if PipelineComponent.TOKENIZER not in skip_components: + tokenizer_path = os.path.join(checkpoint_dir, PipelineComponent.TOKENIZER) + if self._text_encoder_class == "Qwen3ForCausalLM": + logger.info("Loading tokenizer (Qwen2TokenizerFast)...") + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + else: + logger.info("Loading tokenizer (PixtralProcessor)...") + self.tokenizer = AutoProcessor.from_pretrained(tokenizer_path) + + # Text encoder (loaded based on detected type) + if PipelineComponent.TEXT_ENCODER not in skip_components: + logger.info(f"Loading text encoder ({self._text_encoder_class})...") + text_encoder_path = os.path.join(checkpoint_dir, PipelineComponent.TEXT_ENCODER) + if self._text_encoder_class == "Mistral3ForConditionalGeneration": + # Mistral3 is a multimodal model (not pure CausalLM) + self.text_encoder = Mistral3ForConditionalGeneration.from_pretrained( + text_encoder_path, + torch_dtype=self.model_config.torch_dtype, + ).to(device) + else: + # Qwen3 and other CausalLM text encoders + self.text_encoder = AutoModelForCausalLM.from_pretrained( + text_encoder_path, + torch_dtype=self.model_config.torch_dtype, + ).to(device) + + # VAE (FLUX.2-specific VAE with BatchNorm) + # Use full path to avoid AutoConfig issues + if PipelineComponent.VAE not in skip_components: + logger.info("Loading VAE...") + vae_path = os.path.join(checkpoint_dir, PipelineComponent.VAE) + self.vae = AutoencoderKLFlux2.from_pretrained(vae_path, torch_dtype=torch.bfloat16).to( + device + ) + + self.vae_scale_factor = 8 # FLUX.2 uses scale_factor=8 + + # Scheduler + if PipelineComponent.SCHEDULER not in skip_components: + logger.info("Loading scheduler...") + scheduler_path = os.path.join(checkpoint_dir, PipelineComponent.SCHEDULER) + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(scheduler_path) + + # FLUX.2 config + self.max_sequence_length = 512 + self.default_height = 1024 + self.default_width = 1024 + + def load_weights(self, weights: dict) -> None: + """Load transformer weights.""" + if self.transformer is not None and hasattr(self.transformer, "load_weights"): + logger.info("Loading transformer weights...") + transformer_weights = weights.get("transformer", weights) + self.transformer.load_weights(transformer_weights) + logger.info("Transformer weights loaded successfully.") + + self._target_dtype = self.model_config.torch_dtype + + if self.transformer is not None: + self.transformer.eval() + + def post_load_weights(self) -> None: + """Post-load setup: TeaCache registration.""" + super().post_load_weights() + if self.transformer is not None: + # Register TeaCache extractor for FLUX.2 (must be after device placement) + # Only set guidance_param_name for variants with guidance_embeds + guidance_param = "guidance" if self.transformer.guidance_embeds else None + forward_params = [ + "hidden_states", + "encoder_hidden_states", + "timestep", + "img_ids", + "txt_ids", + "guidance", + "return_dict", + ] + register_extractor_from_config( + ExtractorConfig( + model_class_name="Flux2Transformer2DModel", + timestep_embed_fn=self._compute_flux2_timestep_embedding, + guidance_param_name=guidance_param, + forward_params=forward_params, + return_dict_default=False, + ) + ) + + # Enable TeaCache with FLUX.2-specific coefficients + self._setup_teacache(self.transformer, coefficients=FLUX2_TEACACHE_COEFFICIENTS) + + def infer(self, req): + """Run inference from DiffusionRequest.""" + return self.forward( + prompt=req.prompt, + height=req.height, + width=req.width, + num_inference_steps=req.num_inference_steps, + guidance_scale=req.guidance_scale, + seed=req.seed, + max_sequence_length=req.max_sequence_length, + ) + + @torch.inference_mode() + def forward( + self, + prompt: str, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 3.5, + seed: int = 42, + max_sequence_length: int = 512, + ): + """Generate image from text prompt. + + Args: + prompt: Text prompt for image generation + height: Output image height (default: 1024) + width: Output image width (default: 1024) + num_inference_steps: Number of denoising steps + guidance_scale: Embedded guidance scale + seed: Random seed for reproducibility + max_sequence_length: Maximum text sequence length + + Returns: + Dict with "image" key containing PIL.Image + """ + pipeline_start = time.time() + generator = torch.Generator(device=self.device).manual_seed(seed) + + # Encode prompt using Mistral3 multi-layer extraction + logger.info("Encoding prompt...") + encode_start = time.time() + prompt_embeds, text_ids = self._encode_prompt(prompt, max_sequence_length) + logger.info(f"Prompt encoding completed in {time.time() - encode_start:.2f}s") + + # Prepare latents + latents, latent_ids = self._prepare_latents(height, width, generator) + logger.info(f"Latents shape: {latents.shape}") + + # Prepare timesteps with dynamic shifting + # Use explicit linear sigmas (matches HF diffusers exactly) + # This is critical for step-distilled models like FLUX.2-klein + image_seq_len = latents.shape[1] + mu = compute_empirical_mu(image_seq_len, num_inference_steps) + sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) + + # If scheduler uses flow sigmas, let it compute its own (override linear) + if ( + hasattr(self.scheduler.config, "use_flow_sigmas") + and self.scheduler.config.use_flow_sigmas + ): + self.scheduler.set_timesteps(num_inference_steps, device=self.device, mu=mu) + else: + self.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=self.device, mu=mu) + timesteps = self.scheduler.timesteps + + # Prepare guidance (only for variants with guidance_embeds=True) + guidance = None + if self.transformer.guidance_embeds: + guidance = torch.full( + [latents.shape[0]], guidance_scale, device=self.device, dtype=torch.float32 + ) + + # Denoising loop using forward_fn callback (WAN pattern) + def forward_fn( + latents, extra_stream_latents, timestep, encoder_hidden_states, extra_tensors + ): + """Forward function for FLUX.2 transformer.""" + return self.transformer( + hidden_states=latents, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep / 1000, # FLUX.2 expects normalized timesteps + img_ids=latent_ids, + txt_ids=text_ids, + guidance=guidance, + return_dict=False, + )[0] + + latents = self.denoise( + latents=latents, + scheduler=self.scheduler, + prompt_embeds=prompt_embeds, + guidance_scale=1.0, # No CFG: guidance is embedded + forward_fn=forward_fn, + timesteps=timesteps, + ) + + # Decode + logger.info("Decoding image...") + decode_start = time.time() + image = self.decode_latents(latents, lambda lat: self._decode_latents(lat, latent_ids)) + + if self.rank == 0: + logger.info(f"Image decoded in {time.time() - decode_start:.2f}s") + logger.info(f"Total pipeline time: {time.time() - pipeline_start:.2f}s") + + return MediaOutput(image=image) + + def _encode_prompt( + self, + prompt: str, + max_sequence_length: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode prompt using multi-layer hidden state extraction. + + Supports both text encoder types: + - Mistral3: system message + PixtralProcessor chat template + - Qwen3: simple user message + Qwen2TokenizerFast chat template + + Returns: + Tuple of (prompt_embeds, text_ids) + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + + # Tokenize (format depends on text encoder type) + text_encoder_class = getattr( + self, "_text_encoder_class", "Mistral3ForConditionalGeneration" + ) + if text_encoder_class == "Qwen3ForCausalLM": + input_ids, attention_mask = self._tokenize_qwen3(prompt, max_sequence_length) + else: + input_ids, attention_mask = self._tokenize_mistral3(prompt, max_sequence_length) + + # Forward pass - extract hidden states + outputs = self.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Multi-layer extraction: stack specified layers + stacked = torch.stack([outputs.hidden_states[k] for k in self.HIDDEN_STATE_LAYERS], dim=1) + stacked = stacked.to(dtype=self.dtype, device=self.device) + + # Reshape: [B, num_layers, seq, hidden_dim] -> [B, seq, num_layers * hidden_dim] + batch_size, num_layers, seq_len, hidden_dim = stacked.shape + prompt_embeds = stacked.permute(0, 2, 1, 3).reshape( + batch_size, seq_len, num_layers * hidden_dim + ) + + # NOTE: Do NOT zero out padding tokens for decoder-only models. + # Causal attention means hidden states at padding positions carry + # meaningful context from earlier tokens. (HF diffusers does not zero out either.) + + # Prepare 4-axis text IDs for FLUX.2 RoPE + text_ids = self._prepare_text_ids(prompt_embeds) + text_ids = text_ids.to(self.device) + + return prompt_embeds, text_ids + + def _tokenize_mistral3( + self, prompt: List[str], max_sequence_length: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Tokenize prompt using Mistral3 PixtralProcessor chat template.""" + messages_batch = format_input(prompt, SYSTEM_MESSAGE) + inputs = self.tokenizer.apply_chat_template( + messages_batch, + add_generation_prompt=False, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + return inputs["input_ids"].to(self.device), inputs["attention_mask"].to(self.device) + + def _tokenize_qwen3( + self, prompt: List[str], max_sequence_length: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Tokenize prompt using Qwen3 chat template (matches HF diffusers exactly).""" + all_input_ids = [] + all_attention_masks = [] + + for single_prompt in prompt: + messages = [{"role": "user", "content": single_prompt}] + text = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = self.tokenizer( + text, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + all_input_ids.append(inputs["input_ids"]) + all_attention_masks.append(inputs["attention_mask"]) + + input_ids = torch.cat(all_input_ids, dim=0).to(self.device) + attention_mask = torch.cat(all_attention_masks, dim=0).to(self.device) + return input_ids, attention_mask + + def _prepare_text_ids(self, text_embeds: torch.Tensor) -> torch.Tensor: + """Prepare 4-axis position IDs for text tokens. + + FLUX.2 uses 4-axis: (t, h, w, l) where text has t=h=w=0, l=position. + Returns 2D tensor [seq_len, 4] (unbatched, like FLUX.1). + """ + _batch_size, seq_len, _ = text_embeds.shape + + l_ids = torch.arange(seq_len, device=text_embeds.device) + text_ids = torch.stack( + [ + torch.zeros(seq_len, device=text_embeds.device), # t = 0 + torch.zeros(seq_len, device=text_embeds.device), # h = 0 + torch.zeros(seq_len, device=text_embeds.device), # w = 0 + l_ids.float(), # l = position + ], + dim=-1, + ) + + return text_ids # [seq_len, 4] + + def _prepare_latent_ids(self, height: int, width: int) -> torch.Tensor: + """Prepare 4-axis position IDs for latent patches. + + FLUX.2 uses 4-axis: (t, h, w, l) where image has t=0, l=0, h=row, w=col. + Returns 2D tensor [seq_len, 4] (unbatched, like FLUX.1). + """ + latent_height = height // self.vae_scale_factor // 2 # Account for packing + latent_width = width // self.vae_scale_factor // 2 + + t_dim = torch.arange(1, device=self.device) # [0] + h_dim = torch.arange(latent_height, device=self.device) + w_dim = torch.arange(latent_width, device=self.device) + l_dim = torch.arange(1, device=self.device) # [0] + + latent_ids = torch.cartesian_prod(t_dim, h_dim, w_dim, l_dim).float() + + return latent_ids # [seq_len, 4] + + def _prepare_latents( + self, + height: int, + width: int, + generator: torch.Generator, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Prepare random latents in FLUX.2 packed format and position IDs.""" + # FLUX.2: in_channels=128, VAE scale=8, 2x2 packing + latent_height = 2 * (height // (self.vae_scale_factor * 2)) + latent_width = 2 * (width // (self.vae_scale_factor * 2)) + + in_channels = self.transformer.config.in_channels # 128 + + # Create 4D latents then pack (matches HF for seed reproducibility) + latent_shape = (1, in_channels, latent_height // 2, latent_width // 2) + latents_4d = randn_tensor( + latent_shape, generator=generator, device=self.device, dtype=self.dtype + ) + + # Pack latents: [B, C, H, W] -> [B, H*W, C] + latents = self._pack_latents(latents_4d) + + # Prepare position IDs + latent_ids = self._prepare_latent_ids(height, width) + + return latents, latent_ids + + def _pack_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Pack latents: [B, C, H, W] -> [B, H*W, C]""" + batch_size, num_channels, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels, height * width) + latents = latents.permute(0, 2, 1) # [B, H*W, C] + return latents + + def _unpack_latents_with_ids( + self, + latents: torch.Tensor, + latent_ids: torch.Tensor, + ) -> torch.Tensor: + """Unpack latents using position IDs to scatter tokens into spatial positions.""" + _, ch = latents[0].shape + + h_ids = latent_ids[:, 1].to(torch.int64) + w_ids = latent_ids[:, 2].to(torch.int64) + + h = torch.max(h_ids) + 1 + w = torch.max(w_ids) + 1 + + flat_ids = h_ids * w + w_ids + + x_list = [] + for data in latents: + out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype) + out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data) + out = out.view(h, w, ch).permute(2, 0, 1) + x_list.append(out) + + return torch.stack(x_list, dim=0) + + def _unpatchify_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Unpatchify latents: [B, 128, H, W] -> [B, 32, H*2, W*2]""" + batch_size, num_channels, height, width = latents.shape + + # 128 channels = 32 * 2 * 2 (2x2 patches) + latents = latents.reshape(batch_size, num_channels // 4, 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) # [B, 32, H, 2, W, 2] + latents = latents.reshape(batch_size, num_channels // 4, height * 2, width * 2) + + return latents + + def _decode_latents(self, latents: torch.Tensor, latent_ids: torch.Tensor) -> torch.Tensor: + """Decode latents to image tensor.""" + # Unpack latents using position IDs + latents = self._unpack_latents_with_ids(latents, latent_ids) + + # BatchNorm denormalization (critical for FLUX.2 VAE!) + if hasattr(self.vae, "bn") and hasattr(self.vae.bn, "running_mean"): + bn_eps = getattr(self.vae.config, "batch_norm_eps", 1e-5) + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to( + latents.device, latents.dtype + ) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + bn_eps).to( + latents.device, latents.dtype + ) + latents = latents * latents_bn_std + latents_bn_mean + + # Unpatchify + latents = self._unpatchify_latents(latents) + + # VAE decode + latents = latents.to(self.vae.dtype) + image = self.vae.decode(latents, return_dict=False)[0] + + # Post-process to tensor (H, W, C) uint8 + image = (image / 2 + 0.5).clamp(0, 1) + image = image.permute(0, 2, 3, 1) # (B, C, H, W) -> (B, H, W, C) + image = (image * 255).round().to(torch.uint8) + + return image[0] # Remove batch dimension diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pos_embed_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pos_embed_flux.py new file mode 100644 index 000000000000..db8ed3e04472 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pos_embed_flux.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FLUX position embedding utilities. + +Key Components: +- FluxPosEmbed: Multi-axis rotary position embeddings (FLUX.1: 3-axis, FLUX.2: 4-axis) +- get_1d_rotary_pos_embed: 1D rotary position embedding computation + +RoPE application uses the shared apply_rotary_emb from modules/attention.py. +""" + +from typing import List, Tuple + +import torch +import torch.nn as nn + + +def get_1d_rotary_pos_embed( + dim: int, + pos: torch.Tensor, + theta: "float | torch.Tensor" = 10000.0, + use_real: bool = True, + repeat_interleave_real: bool = True, + freqs_dtype: torch.dtype = torch.float64, +) -> "Tuple[torch.Tensor, torch.Tensor] | torch.Tensor": + """Compute 1D rotary position embeddings. + + Args: + dim: Embedding dimension (must be even) + pos: Position tensor of shape (seq_len,) + theta: RoPE theta parameter + use_real: Return (cos, sin) instead of complex exp + repeat_interleave_real: Interleave or repeat cos/sin + freqs_dtype: Dtype for frequency computation + + Returns: + If use_real=True: Tuple of (cos, sin) tensors each of shape (seq_len, dim) + If use_real=False: Single complex tensor of shape (seq_len, dim/2) + """ + assert dim % 2 == 0, f"dim must be even, got {dim}" + + # Compute frequency bands + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device) / dim)) + + # Outer product: (seq_len, dim/2) + freqs = torch.outer(pos.to(freqs_dtype), freqs) + + if use_real: + if repeat_interleave_real: + # Repeat each frequency: [f0, f0, f1, f1, ...] + freqs = freqs.repeat_interleave(2, dim=-1) + else: + # Repeat pattern: [f0, f1, ..., f0, f1, ...] + freqs = freqs.repeat(1, 2) + + cos = freqs.cos().to(pos.dtype) + sin = freqs.sin().to(pos.dtype) + return cos, sin + else: + # Return complex exponential + return torch.polar(torch.ones_like(freqs), freqs) + + +class FluxPosEmbed(nn.Module): + """Multi-axis Rotary Position Embedding for FLUX models. + + Computes RoPE for each axis independently and concatenates the results. + Parameterized by axes_dim to support different FLUX variants: + - FLUX.1: axes_dim=[16, 56, 56] (3-axis: txt, h, w), theta=10000 + - FLUX.2: axes_dim=[32, 32, 32, 32] (4-axis: t, h, w, l), theta=2000 + + Total dimension = sum(axes_dim) = attention_head_dim (128 for both). + """ + + def __init__(self, theta: int = 10000, axes_dim: List[int] = None): + """Initialize FluxPosEmbed. + + Args: + theta: Base for exponential frequency computation + axes_dim: Dimensions for each axis. Default: [16, 56, 56] (FLUX.1) + """ + super().__init__() + self.register_buffer("theta", torch.tensor(theta, dtype=torch.float64), persistent=False) + self.axes_dim = axes_dim if axes_dim is not None else [16, 56, 56] + + def forward(self, ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute rotary embeddings from position IDs. + + Args: + ids: Position IDs tensor of shape (seq_len, n_axes) + + Returns: + Tuple of (freqs_cos, freqs_sin), each of shape (1, seq_len, 1, head_dim) + """ + n_axes = ids.shape[-1] + assert n_axes == len(self.axes_dim), ( + f"ids has {n_axes} axes but axes_dim has {len(self.axes_dim)} entries" + ) + cos_out = [] + sin_out = [] + pos = ids.float() + + # Determine frequency dtype based on device + is_mps = ids.device.type == "mps" + is_npu = ids.device.type == "npu" + freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 + + # Compute RoPE for each axis + for i in range(n_axes): + cos, sin = get_1d_rotary_pos_embed( + self.axes_dim[i], + pos[:, i], + theta=self.theta, + repeat_interleave_real=True, + use_real=True, + freqs_dtype=freqs_dtype, + ) + cos_out.append(cos) + sin_out.append(sin) + + # Concatenate along dimension axis and reshape to [1, S, 1, D] + # to match WAN's format expected by the shared apply_rotary_emb() + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + freqs_cos = freqs_cos.unsqueeze(0).unsqueeze(2) + freqs_sin = freqs_sin.unsqueeze(0).unsqueeze(2) + + return freqs_cos, freqs_sin diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py new file mode 100644 index 000000000000..a07d44c58c53 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux.py @@ -0,0 +1,937 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FLUX Transformer model implementation. + +FLUX Architecture: +- 19 dual-stream blocks (FluxTransformerBlock): Separate processing for image/text +- 38 single-stream blocks (FluxSingleTransformerBlock): Joint processing +- Joint attention mechanism with QK normalization +- 2D RoPE position embeddings + +Forward Pass Flow: +1. Embed inputs: x_embedder(latents), context_embedder(text), time_text_embed(timestep, pooled) +2. Compute RoPE from position IDs +3. Run 19 dual-stream blocks: (hidden_states, encoder_hidden_states) processed separately +4. Run 38 single-stream blocks: Concatenate, process, split +5. norm_out + proj_out -> noise prediction +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import PixArtAlphaTextProjection as TextProjection +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from tqdm import tqdm + +from tensorrt_llm._torch.modules.layer_norm import LayerNorm +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.utils import maybe_compile +from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention +from tensorrt_llm._torch.visual_gen.models.flux.pos_embed_flux import FluxPosEmbed +from tensorrt_llm._torch.visual_gen.parallelism import setup_sequence_parallelism +from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader +from tensorrt_llm.models.modeling_utils import QuantConfig + +if TYPE_CHECKING: + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + +# HF checkpoint key → our module attribute name +_WEIGHT_KEY_REMAPS = [ + ("net.0.proj.", "up_proj."), # GELU(proj=Linear) → up_proj Linear + ("net.2.", "down_proj."), # net[2] output Linear → down_proj +] + + +def _remap_checkpoint_keys(weights: dict) -> dict: + """Remap HuggingFace checkpoint keys to our module attribute names. + + HF diffusers uses nn.ModuleList wrappers that add numeric indices to + weight key paths. Our simplified module structure uses plain attributes, + so we translate the keys at load time. + """ + remapped = {} + for key, value in weights.items(): + new_key = key + for old, new in _WEIGHT_KEY_REMAPS: + new_key = new_key.replace(old, new) + remapped[new_key] = value + return remapped + + +class _AdaLayerNormBase(nn.Module): + """Base class for adaptive layer normalization variants. + + All variants share the same structure: silu → linear → chunk → norm*scale+shift. + Subclasses only differ in the number of modulation parameters (num_chunks) + and how they parse the chunks in forward(). + """ + + def __init__( + self, + embedding_dim: int, + num_chunks: int, + conditioning_embedding_dim: Optional[int] = None, + eps: float = 1e-5, + elementwise_affine: bool = False, + bias: bool = True, + dtype: torch.dtype = None, + quant_config=None, + skip_create_weights: bool = False, + force_dynamic_quant: bool = False, + ): + super().__init__() + self.silu = nn.SiLU() + in_dim = ( + conditioning_embedding_dim if conditioning_embedding_dim is not None else embedding_dim + ) + self.linear = Linear( + in_dim, + num_chunks * embedding_dim, + bias=bias, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + self.norm = LayerNorm( + hidden_size=embedding_dim, + eps=eps, + has_weights=elementwise_affine, + has_bias=elementwise_affine, + dtype=dtype, + ) + + +class AdaLayerNormZero(_AdaLayerNormBase): + """Adaptive Layer Normalization for dual-stream DiT blocks. + + Returns 5 modulation parameters: (norm_x, gate_msa, shift_mlp, scale_mlp, gate_mlp). + """ + + def __init__(self, embedding_dim: int, num_embeddings: Optional[int] = None, **kwargs): + super().__init__(embedding_dim, num_chunks=6, **kwargs) + + @maybe_compile() + def forward( + self, x: torch.Tensor, emb: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaLayerNormZeroSingle(_AdaLayerNormBase): + """Adaptive Layer Normalization for single-stream blocks. + + Returns 2 modulation parameters: (norm_x, gate). + """ + + def __init__(self, embedding_dim: int, **kwargs): + super().__init__(embedding_dim, num_chunks=3, **kwargs) + + @maybe_compile() + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + emb = self.linear(self.silu(emb)) + shift, scale, gate = emb.chunk(3, dim=1) + x = self.norm(x) * (1 + scale[:, None]) + shift[:, None] + return x, gate + + +class AdaLayerNormContinuous(_AdaLayerNormBase): + """Continuous adaptive layer normalization for output projection. + + Shared by FLUX.1 and FLUX.2. Returns modulated tensor (scale+shift only, no gate). + """ + + def __init__(self, embedding_dim: int, conditioning_embedding_dim: int, **kwargs): + super().__init__( + embedding_dim, + num_chunks=2, + conditioning_embedding_dim=conditioning_embedding_dim, + **kwargs, + ) + + @maybe_compile() + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + emb = self.linear(self.silu(emb)) + scale, shift = emb.unsqueeze(1).chunk(2, dim=-1) + x = self.norm(x) * (1 + scale) + shift + return x + + +class CombinedTimestepTextProjEmbeddings(nn.Module): + """Combined timestep and text projection embeddings (for schnell). + + Uses diffusers embedding classes for compatibility. These small layers + don't need quantization - the bulk of compute is in transformer blocks. + Note: Dtype conversion happens via model.to(dtype) after weight loading. + """ + + def __init__( + self, + embedding_dim: int, + pooled_projection_dim: int, + **kwargs, # Accept but ignore dtype/quant params for API compatibility + ): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.text_embedder = TextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") + + def forward( + self, + timestep: torch.Tensor, + pooled_projection: torch.Tensor, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) + pooled_projections = self.text_embedder(pooled_projection) + return timesteps_emb + pooled_projections + + +class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module): + """Combined timestep, guidance, and text projection embeddings (for dev). + + Uses diffusers embedding classes for compatibility. These small layers + don't need quantization - the bulk of compute is in transformer blocks. + Note: Dtype conversion happens via model.to(dtype) after weight loading. + """ + + def __init__( + self, + embedding_dim: int, + pooled_projection_dim: int, + **kwargs, # Accept but ignore dtype/quant params for API compatibility + ): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + self.text_embedder = TextProjection(pooled_projection_dim, embedding_dim, act_fn="silu") + + def forward( + self, + timestep: torch.Tensor, + guidance: torch.Tensor, + pooled_projection: torch.Tensor, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) + + guidance_proj = self.time_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype)) + + time_guidance_emb = timesteps_emb + guidance_emb + pooled_projections = self.text_embedder(pooled_projection) + + return time_guidance_emb + pooled_projections + + +@torch.compiler.disable +def _gelu_tanh_eager(x: torch.Tensor) -> torch.Tensor: + """GELU(tanh) that always runs in eager mode. + + Inductor's fused GELU kernel introduces per-element numerical drift in BF16 + that compounds through residual transformer blocks, causing ~6.6 dB PSNR loss + over 50 denoising steps. Running GELU in eager avoids this. + See: pytorch/pytorch#145213 + """ + return F.gelu(x, approximate="tanh") + + +class FluxTransformerBlock(nn.Module): + """Dual-stream transformer block for FLUX. + + Processes image and text tokens separately, combines via joint attention, + then applies separate FFNs. + + Architecture: + 1. AdaLN for image (norm1) and text (norm1_context) + 2. Joint attention (FluxJointAttention with added_kv_proj_dim) + 3. Residual + gated attention output + 4. LayerNorm + modulation for FFN + 5. Separate FFNs for image (ff) and text (ff_context) + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + qk_norm: str = "rms_norm", + eps: float = 1e-6, + dtype: torch.dtype = None, + quant_config=None, + skip_create_weights: bool = False, + force_dynamic_quant: bool = False, + config: Optional["DiffusionModelConfig"] = None, + layer_idx: int = 0, + ): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + # AdaLN for image and text + self.norm1 = AdaLayerNormZero( + dim, + eps=eps, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + self.norm1_context = AdaLayerNormZero( + dim, + eps=eps, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + + # Joint attention + self.attn = FluxJointAttention( + hidden_size=dim, + num_attention_heads=num_attention_heads, + head_dim=attention_head_dim, + added_kv_proj_dim=dim, + bias=True, + eps=eps, + config=config, + layer_idx=layer_idx, + ) + + # FFN normalization (TRT-LLM LayerNorm) + self.norm2 = LayerNorm( + hidden_size=dim, eps=1e-6, has_weights=False, has_bias=False, dtype=dtype + ) + self.norm2_context = LayerNorm( + hidden_size=dim, eps=1e-6, has_weights=False, has_bias=False, dtype=dtype + ) + + # FFN layers (shared TRT-LLM MLP module) + # HF key remapping (net.0.proj.* → up_proj.*, net.2.* → down_proj.*) in load_weights() + self.ff = MLP( + hidden_size=dim, + intermediate_size=int(dim * 4.0), + bias=True, + activation=_gelu_tanh_eager, + dtype=dtype, + config=config, + layer_idx=layer_idx, + reduce_output=False, + ) + self.ff_context = MLP( + hidden_size=dim, + intermediate_size=int(dim * 4.0), + bias=True, + activation=_gelu_tanh_eager, + dtype=dtype, + config=config, + layer_idx=layer_idx, + reduce_output=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass. + + Args: + hidden_states: Image tokens (batch, img_seq, dim) + encoder_hidden_states: Text tokens (batch, txt_seq, dim) + temb: Timestep embedding (batch, dim) + image_rotary_emb: RoPE (cos, sin) tuple + joint_attention_kwargs: Additional kwargs for attention + + Returns: + Tuple of (encoder_hidden_states, hidden_states) + """ + # Image: AdaLN modulation + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, emb=temb + ) + + # Text: AdaLN modulation + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = ( + self.norm1_context(encoder_hidden_states, emb=temb) + ) + + # Joint attention + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + # Image: Gated residual for attention + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = hidden_states + attn_output + + # Image: FFN with modulation + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm_hidden_states) + ff_output = gate_mlp.unsqueeze(1) * ff_output + hidden_states = hidden_states + ff_output + + # Text: Gated residual for attention + context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output + encoder_hidden_states = encoder_hidden_states + context_attn_output + + # Text: FFN with modulation + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = ( + norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + ) + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output + + # FP16 overflow protection + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +class FluxSingleTransformerBlock(nn.Module): + """Single-stream transformer block for FLUX. + + Concatenates image and text tokens, processes together, + then splits back. + + Architecture: + 1. Concatenate encoder_hidden_states + hidden_states + 2. AdaLayerNormZeroSingle + 3. Parallel attention + MLP branches + 4. proj_out(concat(attn, mlp)) + 5. Gated residual + 6. Split back into image and text + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + dtype: torch.dtype = None, + quant_config=None, + skip_create_weights: bool = False, + force_dynamic_quant: bool = False, + config: Optional["DiffusionModelConfig"] = None, + layer_idx: int = 0, + ): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.mlp_hidden_dim = int(dim * mlp_ratio) + + # AdaLN + self.norm = AdaLayerNormZeroSingle( + dim, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + + # MLP branch (TRT-LLM Linear for quantization support) + self.proj_mlp = Linear( + dim, + self.mlp_hidden_dim, + bias=True, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + self.act_mlp = _gelu_tanh_eager + + # Output projection (concat of attn + mlp) - TRT-LLM Linear + self.proj_out = Linear( + dim + self.mlp_hidden_dim, + dim, + bias=True, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + + # Attention (no added_kv_proj_dim since tokens are already concatenated) + self.attn = FluxJointAttention( + hidden_size=dim, + num_attention_heads=num_attention_heads, + head_dim=attention_head_dim, + bias=True, + eps=1e-6, + pre_only=True, # No output projection in attention + config=config, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass. + + Args: + hidden_states: Image tokens (batch, img_seq, dim) + encoder_hidden_states: Text tokens (batch, txt_seq, dim) + temb: Timestep embedding (batch, dim) + image_rotary_emb: RoPE (cos, sin) tuple + joint_attention_kwargs: Additional kwargs for attention + + Returns: + Tuple of (encoder_hidden_states, hidden_states) + """ + text_seq_len = encoder_hidden_states.shape[1] + + # Concatenate text + image + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + residual = hidden_states + + # AdaLN + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + + # MLP branch + mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + + # Attention branch + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + # Concat and project + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + gate = gate.unsqueeze(1) + hidden_states = gate * self.proj_out(hidden_states) + + # Residual + hidden_states = residual + hidden_states + + # FP16 overflow protection + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + # Split back into text and image + encoder_hidden_states, hidden_states = ( + hidden_states[:, :text_seq_len], + hidden_states[:, text_seq_len:], + ) + + return encoder_hidden_states, hidden_states + + +class FluxTransformer2DModel(nn.Module): + """FLUX Transformer model for text-to-image generation. + + This is the native TRT-LLM implementation of FLUX transformer. + Supports FP8/NVFP4 quantization for optimized inference. + + Architecture: + - pos_embed: FluxPosEmbed for 2D RoPE + - time_text_embed: Combined timestep + guidance + text projection embeddings + - context_embedder: Linear projection for T5 text embeddings + - x_embedder: Linear projection for latent inputs + - transformer_blocks: 19 dual-stream FluxTransformerBlock + - single_transformer_blocks: 38 single-stream FluxSingleTransformerBlock + - norm_out: AdaLayerNormContinuous + - proj_out: Linear projection to output + """ + + def __init__(self, model_config: "DiffusionModelConfig"): + super().__init__() + self.model_config = model_config + + # Setup sequence parallelism (Ulysses) + num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 24) + self.use_ulysses, self.ulysses_size, self.ulysses_pg, self.ulysses_rank = ( + setup_sequence_parallelism( + model_config=model_config, + num_attention_heads=num_heads, + ) + ) + + # Extract pretrained config from model_config + pretrained_config = model_config.pretrained_config + + # Extract dtype and quantization config for Linear modules (following Wan pattern) + dtype = model_config.torch_dtype + quant_config = model_config.quant_config + skip_create_weights = model_config.skip_create_weights_in_init + force_dynamic_quant = model_config.force_dynamic_quantization + + # Extract FLUX-specific parameters from pretrained config + num_attention_heads = getattr(pretrained_config, "num_attention_heads", 24) + attention_head_dim = getattr(pretrained_config, "attention_head_dim", 128) + in_channels = getattr(pretrained_config, "in_channels", 64) + out_channels = getattr(pretrained_config, "out_channels", in_channels) + num_layers = getattr(pretrained_config, "num_layers", 19) + num_single_layers = getattr(pretrained_config, "num_single_layers", 38) + joint_attention_dim = getattr(pretrained_config, "joint_attention_dim", 4096) + pooled_projection_dim = getattr(pretrained_config, "pooled_projection_dim", 768) + guidance_embeds = getattr(pretrained_config, "guidance_embeds", False) + patch_size = getattr(pretrained_config, "patch_size", 1) + axes_dims_rope = getattr(pretrained_config, "axes_dims_rope", [16, 56, 56]) + theta_rope = getattr(pretrained_config, "theta", 10000) + + # Compute inner dimension + self.inner_dim = num_attention_heads * attention_head_dim + self.out_channels = out_channels + self.in_channels = in_channels + self.guidance_embeds = guidance_embeds + + # Store config for compatibility + self.config = type( + "Config", + (), + { + "num_attention_heads": num_attention_heads, + "attention_head_dim": attention_head_dim, + "in_channels": in_channels, + "out_channels": out_channels, + "num_layers": num_layers, + "num_single_layers": num_single_layers, + "joint_attention_dim": joint_attention_dim, + "pooled_projection_dim": pooled_projection_dim, + "guidance_embeds": guidance_embeds, + "patch_size": patch_size, + "axes_dims_rope": axes_dims_rope, + "theta_rope": theta_rope, + "inner_dim": self.inner_dim, + }, + )() + + # Position embeddings + self.pos_embed = FluxPosEmbed(theta=theta_rope, axes_dim=list(axes_dims_rope)) + + # Time + text embeddings + if guidance_embeds: + self.time_text_embed = CombinedTimestepGuidanceTextProjEmbeddings( + embedding_dim=self.inner_dim, + pooled_projection_dim=pooled_projection_dim, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + else: + self.time_text_embed = CombinedTimestepTextProjEmbeddings( + embedding_dim=self.inner_dim, + pooled_projection_dim=pooled_projection_dim, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + + # Input embedders (TRT-LLM Linear for quantization support) + self.context_embedder = Linear( + joint_attention_dim, + self.inner_dim, + bias=True, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + # NOTE: x_embedder quantization is excluded when in_channels < 128. + # FLUX.1 has in_channels=64, which is below the 128-block size required by + # fp8_block_scaling_gemm (causes NVRTC compilation failure). This layer runs + # once per forward pass (not in the block loop), so the perf impact is negligible. + if in_channels < 128 and quant_config is not None: + if quant_config.exclude_modules is None: + quant_config.exclude_modules = [] + if "*x_embedder*" not in quant_config.exclude_modules: + quant_config.exclude_modules.append("*x_embedder*") + self.x_embedder = Linear( + in_channels, + self.inner_dim, + bias=True, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + + # Dual-stream transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + FluxTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + config=model_config, + layer_idx=i, + ) + for i in range(num_layers) + ] + ) + + # Single-stream transformer blocks + self.single_transformer_blocks = nn.ModuleList( + [ + FluxSingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + config=model_config, + layer_idx=num_layers + i, # Continue numbering from dual-stream blocks + ) + for i in range(num_single_layers) + ] + ) + + # Output layers + self.norm_out = AdaLayerNormContinuous( + self.inner_dim, + self.inner_dim, + elementwise_affine=False, + eps=1e-6, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + # TRT-LLM Linear for quantization support + self.proj_out = Linear( + self.inner_dim, + patch_size * patch_size * self.out_channels, + bias=True, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + + self.__post_init__() + + def __post_init__(self): + self.apply_quant_config_exclude_modules() + + for _, module in self.named_modules(): + if callable(getattr(module, "create_weights", None)): + module.create_weights() + + def apply_quant_config_exclude_modules(self): + quant_config = self.model_config.quant_config + if quant_config is None or quant_config.exclude_modules is None: + return + + kv_cache_quant_algo = quant_config.kv_cache_quant_algo if quant_config else None + no_quant_config = QuantConfig(kv_cache_quant_algo=kv_cache_quant_algo) + + for name, module in self.named_modules(): + if isinstance(module, Linear): + is_excluded = quant_config.is_module_excluded_from_quantization(name) + if is_excluded and getattr(module, "quant_config", None) is not None: + module.quant_config = no_quant_config + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + pooled_projections: torch.Tensor = None, + timestep: torch.Tensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + guidance: torch.Tensor = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + ) -> Union[torch.Tensor, Tuple[torch.Tensor]]: + """Forward pass. + + Args: + hidden_states: Latent image tokens (batch, seq_len, in_channels) + encoder_hidden_states: T5 text embeddings (batch, txt_seq_len, joint_attention_dim) + pooled_projections: CLIP pooled text embeddings (batch, pooled_projection_dim) + timestep: Timestep tensor (batch,) + img_ids: Image position IDs (seq_len, 3) or (batch, seq_len, 3) + txt_ids: Text position IDs (txt_seq_len, 3) or (batch, txt_seq_len, 3) + guidance: Guidance scale tensor (batch,) for FLUX.1-dev + joint_attention_kwargs: Additional kwargs for attention + return_dict: Whether to return dict or tuple + + Returns: + Noise prediction tensor of shape (batch, seq_len, patch_size^2 * out_channels) + """ + # Embed inputs (contiguous needed for FP8 quantize ops) + hidden_states = self.x_embedder(hidden_states.contiguous()) + + # Scale timestep (FLUX convention: multiply by 1000) + timestep = timestep.to(hidden_states.dtype) * 1000 + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + + # Compute timestep + guidance + text embedding + if self.config.guidance_embeds and guidance is not None: + temb = self.time_text_embed(timestep, guidance, pooled_projections) + else: + temb = self.time_text_embed(timestep, pooled_projections) + + # Embed text + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + # Handle 3D IDs (batch dimension) - deprecated format + if txt_ids.ndim == 3: + txt_ids = txt_ids[0] + if img_ids.ndim == 3: + img_ids = img_ids[0] + + # Ulysses: shard sequences and position IDs before RoPE + if self.use_ulysses: + img_seq_len = img_ids.shape[0] + txt_seq_len = txt_ids.shape[0] + + if img_seq_len % self.ulysses_size != 0: + raise ValueError( + f"Image seq len ({img_seq_len}) not divisible by " + f"ulysses_size ({self.ulysses_size})" + ) + if txt_seq_len % self.ulysses_size != 0: + raise ValueError( + f"Text seq len ({txt_seq_len}) not divisible by " + f"ulysses_size ({self.ulysses_size})" + ) + + img_chunk = img_seq_len // self.ulysses_size + txt_chunk = txt_seq_len // self.ulysses_size + r = self.ulysses_rank + + # Shard position IDs (before RoPE computation) + img_ids = img_ids[r * img_chunk : (r + 1) * img_chunk] + txt_ids = txt_ids[r * txt_chunk : (r + 1) * txt_chunk] + + # Shard hidden states + hidden_states = hidden_states[:, r * img_chunk : (r + 1) * img_chunk, :] + encoder_hidden_states = encoder_hidden_states[:, r * txt_chunk : (r + 1) * txt_chunk, :] + + # Compute RoPE embeddings (from potentially sharded IDs) + ids = torch.cat((txt_ids, img_ids), dim=0) + image_rotary_emb = self.pos_embed(ids) + + # Dual-stream blocks + for block in self.transformer_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # Single-stream blocks + for block in self.single_transformer_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # Ulysses: gather output sequence from all ranks + if self.use_ulysses: + hidden_states = hidden_states.contiguous() + gathered = [torch.zeros_like(hidden_states) for _ in range(self.ulysses_size)] + torch.distributed.all_gather(gathered, hidden_states, group=self.ulysses_pg) + hidden_states = torch.cat(gathered, dim=1) + + # Output projection + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if not return_dict: + return (output,) + + return {"sample": output} + + def load_weights(self, weights: dict) -> None: + """Load weights into the transformer. + + Args: + weights: Dictionary of parameter name -> tensor + """ + + # Remap HF checkpoint keys to our module attribute names + weights = _remap_checkpoint_keys(weights) + + # Map fused QKV layer names to original HF checkpoint names + # HF checkpoint has separate to_q, to_k, to_v / add_q_proj, add_k_proj, add_v_proj + # We fuse them into qkv_proj / add_qkv_proj for better performance + params_map = { + "add_qkv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], + "qkv_proj": ["to_q", "to_k", "to_v"], + } + + loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map) + + for name, module in tqdm(self.named_modules(), desc="Loading weights"): + # Create weights for modules with skip_create_weights_in_init=True + # This must be done before loading weights (following Wan pattern) + if callable(getattr(module, "create_weights", None)): + module.create_weights() + + if len(module._parameters) == 0: + continue + + if isinstance(module, Linear): + weight_dicts = loader.get_linear_weights(module, name, weights) + + if weight_dicts: + loader.load_linear_weights(module, name, weight_dicts) + else: + module_weights = loader.filter_weights(name, weights) + for param_name, param in module._parameters.items(): + if param is not None and param_name in module_weights: + param.data.copy_( + module_weights[param_name].to(self.model_config.torch_dtype) + ) + + def post_load_weights(self) -> None: + """Call post_load_weights on all Linear modules and convert embedders to target dtype.""" + # Convert time_text_embed components to target dtype + target_dtype = self.model_config.torch_dtype + if hasattr(self, "time_text_embed"): + if hasattr(self.time_text_embed, "timestep_embedder"): + self.time_text_embed.timestep_embedder.to(target_dtype) + if hasattr(self.time_text_embed, "text_embedder"): + self.time_text_embed.text_embedder.to(target_dtype) + if hasattr(self.time_text_embed, "guidance_embedder"): + self.time_text_embed.guidance_embedder.to(target_dtype) + + # Call post_load_weights on all Linear modules + for _, module in self.named_modules(): + if isinstance(module, Linear): + module.post_load_weights() diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py new file mode 100644 index 000000000000..e08293d13b02 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/flux/transformer_flux2.py @@ -0,0 +1,877 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FLUX.2 Transformer model implementation (Native TRT-LLM). + +FLUX.2 has a DIFFERENT architecture from FLUX.1: +- Different modulation: Flux2Modulation with mod_param_sets +- Different embedding: time_guidance_embed (always, with optional guidance_embedder) +- Different FFN: GatedMLP with swiglu (shared from _torch/modules) +- Different single-stream: Fused QKV+MLP projection (to_qkv_mlp_proj) + +Variants: +- FLUX.2-dev (35B): guidance_embeds=True (default), guidance_embedder active +- FLUX.2-klein (4B/9B): guidance_embeds=False, guidance_embedder=None + +All variants use `self.time_guidance_embed` to match HF checkpoint weight names. +All linear layers use bias=False to match HF weights. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from tqdm import tqdm + +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.layer_norm import LayerNorm +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.visual_gen.models.flux.attention import ( + Flux2ParallelSelfAttention, + FluxJointAttention, +) +from tensorrt_llm._torch.visual_gen.models.flux.pos_embed_flux import FluxPosEmbed +from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux import ( + AdaLayerNormContinuous, + _remap_checkpoint_keys, +) +from tensorrt_llm._torch.visual_gen.parallelism import setup_sequence_parallelism +from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader +from tensorrt_llm.models.modeling_utils import QuantConfig + +if TYPE_CHECKING: + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + +# HF FLUX.2 uses Flux2FeedForward with linear_in/linear_out attribute names. +# We use GatedMLP which uses gate_up_proj/down_proj. Remap at load time. +# NOTE: linear_in is NOT remapped here — it's split into gate/up halves in load_weights() +# because GatedMLP.gate_up_proj uses FUSED_GATE_UP_LINEAR mode (expects 2 separate weights). +_FLUX2_WEIGHT_KEY_REMAPS = [ + ("linear_out.", "down_proj."), +] + +# ============================================================================= +# Time + Guidance Embedding (matches HuggingFace structure exactly) +# ============================================================================= + + +class Flux2TimestepGuidanceEmbeddings(nn.Module): + """Timestep (and optional guidance) embedding for FLUX.2 (matches HuggingFace exactly). + + Used for ALL FLUX.2 variants with the same attribute name `time_guidance_embed`: + - FLUX.2-dev (guidance_embeds=True): timestep_emb + guidance_emb + - FLUX.2-klein (guidance_embeds=False): timestep_emb only (guidance_embedder=None) + + This ensures HF checkpoint weight names (`time_guidance_embed.timestep_embedder.*`) + always match our module attribute name. + + Structure: + - time_proj: Sinusoidal projection (Timesteps) + - timestep_embedder: 2-layer MLP (TimestepEmbedding) + - guidance_embedder: 2-layer MLP (TimestepEmbedding) or None + """ + + def __init__( + self, + in_channels: int = 256, + embedding_dim: int = 6144, + bias: bool = False, + guidance_embeds: bool = True, + dtype: Optional[torch.dtype] = None, + ): + super().__init__() + + # Sinusoidal projection for timesteps (shared for both timestep and guidance) + self.time_proj = Timesteps( + num_channels=in_channels, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + + # Timestep embedder (always present) + self.timestep_embedder = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=embedding_dim, + sample_proj_bias=bias, + ) + + # Guidance embedder (only for variants with guidance_embeds=True) + if guidance_embeds: + self.guidance_embedder = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=embedding_dim, + sample_proj_bias=bias, + ) + else: + self.guidance_embedder = None + + def forward( + self, timestep: torch.Tensor, guidance: Optional[torch.Tensor] = None + ) -> torch.Tensor: + """ + Args: + timestep: [batch] timestep values (already scaled by 1000) + guidance: [batch] guidance scale values (already scaled by 1000), or None + + Returns: + Embedding [batch, embedding_dim] + """ + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(timestep.dtype)) + + if guidance is not None and self.guidance_embedder is not None: + guidance_proj = self.time_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype)) + return timesteps_emb + guidance_emb + + return timesteps_emb + + +# ============================================================================= +# Modulation +# ============================================================================= + + +class Flux2Modulation(nn.Module): + """FLUX.2 modulation layer (matches HuggingFace exactly). + + Projects temb to shift/scale/gate for layer normalization and attention gating. + """ + + def __init__( + self, + dim: int, + mod_param_sets: int = 2, + bias: bool = False, + dtype: Optional[torch.dtype] = None, + quant_config=None, + skip_create_weights: bool = False, + force_dynamic_quant: bool = False, + ): + super().__init__() + self.mod_param_sets = mod_param_sets + self.linear = Linear( + dim, + dim * 3 * mod_param_sets, + bias=bias, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + self.act_fn = nn.SiLU() + + def forward( + self, temb: torch.Tensor + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: + """ + Args: + temb: Time embedding [batch, dim] + + Returns: + Tuple of mod_param_sets 3-tuples, each containing (shift, scale, gate) + Each tensor has shape [batch, 1, dim] + """ + mod = self.act_fn(temb) + mod = self.linear(mod) + + if mod.ndim == 2: + mod = mod.unsqueeze(1) + + # Split into 3*mod_param_sets chunks + mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) + + # Return tuple of 3-tuples (shift, scale, gate) + return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets)) + + +# ============================================================================= +# Feed Forward +# ============================================================================= + + +# ============================================================================= +# Transformer Blocks +# ============================================================================= + + +class Flux2TransformerBlock(nn.Module): + """FLUX.2 dual-stream transformer block (matches HuggingFace). + + Processes image and text tokens with shared attention but separate FFN. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 3.0, + eps: float = 1e-6, + bias: bool = False, + dtype: Optional[torch.dtype] = None, + quant_config=None, + skip_create_weights: bool = False, + force_dynamic_quant: bool = False, + config: Optional["DiffusionModelConfig"] = None, + layer_idx: int = 0, + ): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + + # Layer norms (TRT-LLM - without elementwise affine, modulation provides scale/shift) + self.norm1 = LayerNorm(hidden_size=dim, eps=eps, has_weights=False, has_bias=False) + self.norm1_context = LayerNorm(hidden_size=dim, eps=eps, has_weights=False, has_bias=False) + + # Joint attention + self.attn = FluxJointAttention( + hidden_size=dim, + num_attention_heads=num_attention_heads, + head_dim=attention_head_dim, + bias=bias, + added_kv_proj_dim=dim, + eps=eps, + config=config, + layer_idx=layer_idx, + ) + + # FFN for image stream (shared GatedMLP from _torch/modules) + # HF key remapping (linear_in.* → gate_up_proj.*, linear_out.* → down_proj.*) in load_weights() + self.ff = GatedMLP( + hidden_size=dim, + intermediate_size=int(dim * mlp_ratio), + bias=bias, + dtype=dtype, + config=config, + layer_idx=layer_idx, + reduce_output=False, + ) + # FFN for text stream + self.ff_context = GatedMLP( + hidden_size=dim, + intermediate_size=int(dim * mlp_ratio), + bias=bias, + dtype=dtype, + config=config, + layer_idx=layer_idx, + reduce_output=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + image_rotary_emb: Tuple[torch.Tensor, torch.Tensor], + img_mod: Tuple[Tuple[torch.Tensor, ...], ...], + txt_mod: Tuple[Tuple[torch.Tensor, ...], ...], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + hidden_states: Image features [batch, img_seq, dim] + encoder_hidden_states: Text features [batch, txt_seq, dim] + image_rotary_emb: Tuple of (freqs_cos, freqs_sin) + img_mod: Image modulation ((shift1, scale1, gate1), (shift2, scale2, gate2)) + txt_mod: Text modulation ((shift1, scale1, gate1), (shift2, scale2, gate2)) + + Returns: + Tuple of (encoder_hidden_states, hidden_states) + """ + # Unpack modulation parameters + (img_shift1, img_scale1, img_gate1), (img_shift2, img_scale2, img_gate2) = img_mod + (txt_shift1, txt_scale1, txt_gate1), (txt_shift2, txt_scale2, txt_gate2) = txt_mod + + # Save residuals + img_residual = hidden_states + txt_residual = encoder_hidden_states + + # Pre-norm + modulation + hidden_states = self.norm1(hidden_states) + hidden_states = hidden_states * (1 + img_scale1) + img_shift1 + + encoder_hidden_states = self.norm1_context(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states * (1 + txt_scale1) + txt_shift1 + + # Joint attention + attn_output, encoder_attn_output = self.attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + + # Attention residual with gate + hidden_states = img_residual + attn_output * img_gate1 + encoder_hidden_states = txt_residual + encoder_attn_output * txt_gate1 + + # FFN + img_residual = hidden_states + txt_residual = encoder_hidden_states + + # Modulation for FFN (use scale2/shift2) + hidden_states = self.norm1(hidden_states) + hidden_states = hidden_states * (1 + img_scale2) + img_shift2 + + encoder_hidden_states = self.norm1_context(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states * (1 + txt_scale2) + txt_shift2 + + # FFN with gate + # GatedMLP's swiglu Triton kernel requires 2D input [tokens, features], + # so flatten 3D [batch, seq, dim] before and unflatten after. + b, s, d = hidden_states.shape + hidden_states = ( + img_residual + self.ff(hidden_states.view(b * s, d)).view(b, s, -1) * img_gate2 + ) + b, s, d = encoder_hidden_states.shape + encoder_hidden_states = ( + txt_residual + + self.ff_context(encoder_hidden_states.view(b * s, d)).view(b, s, -1) * txt_gate2 + ) + + return encoder_hidden_states, hidden_states + + +class Flux2SingleTransformerBlock(nn.Module): + """FLUX.2 single-stream transformer block (matches HuggingFace). + + Uses parallel attention and MLP with fused projection. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 3.0, + eps: float = 1e-6, + bias: bool = False, + dtype: Optional[torch.dtype] = None, + quant_config=None, + skip_create_weights: bool = False, + force_dynamic_quant: bool = False, + config: Optional["DiffusionModelConfig"] = None, + layer_idx: int = 0, + ): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.dim = dim + + # Layer norm (TRT-LLM - without elementwise affine) + self.norm = LayerNorm(hidden_size=dim, eps=eps, has_weights=False, has_bias=False) + + # Parallel attention with fused QKV+MLP + self.attn = Flux2ParallelSelfAttention( + hidden_size=dim, + num_attention_heads=num_attention_heads, + head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + bias=bias, + eps=eps, + config=config, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + image_rotary_emb: Tuple[torch.Tensor, torch.Tensor], + mod: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """ + Args: + hidden_states: [batch, seq, dim] + image_rotary_emb: Tuple of (freqs_cos, freqs_sin) + mod: Modulation (shift, scale, gate) + + Returns: + hidden_states [batch, seq, dim] + """ + mod_shift, mod_scale, mod_gate = mod + + # Save residual + residual = hidden_states + + # Pre-norm + modulation + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states * (1 + mod_scale) + mod_shift + + # Parallel attention + MLP + hidden_states = self.attn(hidden_states, image_rotary_emb=image_rotary_emb) + + # Residual with gate + hidden_states = residual + hidden_states * mod_gate + + return hidden_states + + +# ============================================================================= +# Main Model +# ============================================================================= + + +class Flux2Transformer2DModel(nn.Module): + """FLUX.2 Transformer model for image generation (Native TRT-LLM). + + This implements the full FLUX.2 architecture matching HuggingFace diffusers: + - 8 dual-stream transformer blocks with joint attention + - 48 single-stream transformer blocks with fused QKV+MLP + - 4-axis RoPE position embeddings + - Shared modulation layers for all blocks of same type + """ + + def __init__(self, model_config: "DiffusionModelConfig"): + """Initialize FLUX.2 transformer. + + Args: + model_config: DiffusionModelConfig instance (from DiffusionModelLoader) + """ + super().__init__() + self.model_config = model_config + + # Setup sequence parallelism (Ulysses) + num_heads = getattr(model_config.pretrained_config, "num_attention_heads", 48) + self.use_ulysses, self.ulysses_size, self.ulysses_pg, self.ulysses_rank = ( + setup_sequence_parallelism( + model_config=model_config, + num_attention_heads=num_heads, + ) + ) + + # Extract pretrained config from model_config (following WAN/FLUX.1 pattern) + pretrained_config = model_config.pretrained_config + + # Extract dtype and quantization config for Linear modules + dtype = model_config.torch_dtype + quant_config = model_config.quant_config + skip_create_weights = model_config.skip_create_weights_in_init + force_dynamic_quant = model_config.force_dynamic_quantization + + # Extract FLUX.2-specific parameters from pretrained config + patch_size = getattr(pretrained_config, "patch_size", 1) + in_channels = getattr(pretrained_config, "in_channels", 128) + out_channels = getattr(pretrained_config, "out_channels", None) + if out_channels is None: + out_channels = in_channels # Default to in_channels (like Flux2Config.__post_init__) + num_layers = getattr(pretrained_config, "num_layers", 8) + num_single_layers = getattr(pretrained_config, "num_single_layers", 48) + attention_head_dim = getattr(pretrained_config, "attention_head_dim", 128) + num_attention_heads = getattr(pretrained_config, "num_attention_heads", 48) + mlp_ratio = getattr(pretrained_config, "mlp_ratio", 3.0) + joint_attention_dim = getattr(pretrained_config, "joint_attention_dim", 15360) + pooled_projection_dim = getattr(pretrained_config, "pooled_projection_dim", 5120) + timestep_guidance_channels = getattr(pretrained_config, "timestep_guidance_channels", 256) + guidance_embeds = getattr(pretrained_config, "guidance_embeds", True) + axes_dims_rope = tuple(getattr(pretrained_config, "axes_dims_rope", [32, 32, 32, 32])) + theta_rope = getattr(pretrained_config, "rope_theta", 2000.0) + eps = getattr(pretrained_config, "eps", 1e-6) + + # Compute inner dimension + inner_dim = num_attention_heads * attention_head_dim + + # Store key attributes (like FLUX.1) + self.inner_dim = inner_dim + self.in_channels = in_channels + self.out_channels = out_channels + self.guidance_embeds = guidance_embeds + + # Store config for compatibility (like FLUX.1 pattern) + self.config = type( + "Config", + (), + { + "patch_size": patch_size, + "in_channels": in_channels, + "out_channels": out_channels, + "num_layers": num_layers, + "num_single_layers": num_single_layers, + "attention_head_dim": attention_head_dim, + "num_attention_heads": num_attention_heads, + "mlp_ratio": mlp_ratio, + "joint_attention_dim": joint_attention_dim, + "pooled_projection_dim": pooled_projection_dim, + "timestep_guidance_channels": timestep_guidance_channels, + "guidance_embeds": guidance_embeds, + "axes_dims_rope": axes_dims_rope, + "theta_rope": theta_rope, + "eps": eps, + "inner_dim": inner_dim, + }, + )() + + # Position embedding (4-axis RoPE) + self.pos_embed = FluxPosEmbed( + theta=theta_rope, + axes_dim=axes_dims_rope, + ) + + # Time embedding (always stored as time_guidance_embed to match HF weight names) + # When guidance_embeds=False (e.g., klein), guidance_embedder is None + self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( + in_channels=timestep_guidance_channels, + embedding_dim=inner_dim, + bias=False, + guidance_embeds=guidance_embeds, + dtype=dtype, + ) + + # Modulation layers (shared across all blocks of same type) + # mod_param_sets=2 for double stream (attn + ff) + self.double_stream_modulation_img = Flux2Modulation( + inner_dim, + mod_param_sets=2, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + self.double_stream_modulation_txt = Flux2Modulation( + inner_dim, + mod_param_sets=2, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + # mod_param_sets=1 for single stream (parallel attn+ff) + self.single_stream_modulation = Flux2Modulation( + inner_dim, + mod_param_sets=1, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + + # Input embedders + # NOTE: x_embedder quantization is excluded when in_channels < 128. + # FLUX.2 has in_channels=128 (OK), but future variants with smaller latent + # channels would hit the same fp8_block_scaling_gemm NVRTC failure as FLUX.1 + # (in_channels=64). This layer runs once per forward pass, so the perf + # impact is negligible. + if in_channels < 128 and quant_config is not None: + if quant_config.exclude_modules is None: + quant_config.exclude_modules = [] + if "*x_embedder*" not in quant_config.exclude_modules: + quant_config.exclude_modules.append("*x_embedder*") + self.x_embedder = Linear( + in_channels, + inner_dim, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + self.context_embedder = Linear( + joint_attention_dim, + inner_dim, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + + # Dual-stream transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + Flux2TransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + config=model_config, + layer_idx=i, + ) + for i in range(num_layers) + ] + ) + + # Single-stream transformer blocks + self.single_transformer_blocks = nn.ModuleList( + [ + Flux2SingleTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + config=model_config, + layer_idx=num_layers + i, # Continue numbering from dual-stream blocks + ) + for i in range(num_single_layers) + ] + ) + + # Output layers + self.norm_out = AdaLayerNormContinuous( + inner_dim, + inner_dim, + elementwise_affine=False, + eps=eps, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights=skip_create_weights, + force_dynamic_quant=force_dynamic_quant, + ) + self.proj_out = Linear( + inner_dim, + patch_size**2 * out_channels, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights, + force_dynamic_quantization=force_dynamic_quant, + ) + + self.__post_init__() + + def __post_init__(self): + self.apply_quant_config_exclude_modules() + + for _, module in self.named_modules(): + if callable(getattr(module, "create_weights", None)): + module.create_weights() + + def apply_quant_config_exclude_modules(self): + quant_config = self.model_config.quant_config + if quant_config is None or quant_config.exclude_modules is None: + return + + kv_cache_quant_algo = quant_config.kv_cache_quant_algo if quant_config else None + no_quant_config = QuantConfig(kv_cache_quant_algo=kv_cache_quant_algo) + + for name, module in self.named_modules(): + if isinstance(module, Linear): + is_excluded = quant_config.is_module_excluded_from_quantization(name) + if is_excluded and getattr(module, "quant_config", None) is not None: + module.quant_config = no_quant_config + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + img_ids: torch.Tensor, + txt_ids: torch.Tensor, + guidance: Optional[torch.Tensor] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = True, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + """Forward pass. + + Args: + hidden_states: Latent image features [batch, img_seq, in_channels] + encoder_hidden_states: Text features [batch, txt_seq, joint_attention_dim] + timestep: Diffusion timestep [batch] + img_ids: Image position IDs [img_seq, num_axes] or [batch, img_seq, num_axes] + txt_ids: Text position IDs [txt_seq, num_axes] or [batch, txt_seq, num_axes] + guidance: Guidance scale [batch] + joint_attention_kwargs: Additional kwargs for attention (unused) + return_dict: Whether to return a dict + + Returns: + Predicted noise [batch, img_seq, patch_size^2 * out_channels] + """ + txt_seq_len = encoder_hidden_states.shape[1] + + # Embed inputs (contiguous needed for FP8 quantize ops) + hidden_states = self.x_embedder(hidden_states.contiguous()) + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + # Scale timestep and guidance (FLUX convention) + timestep = timestep.to(hidden_states.dtype) * 1000 + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + + # Time embedding (handles both guided and unguided variants) + temb = self.time_guidance_embed(timestep, guidance) + + # Handle batched IDs + if txt_ids.ndim == 3: + txt_ids = txt_ids[0] + if img_ids.ndim == 3: + img_ids = img_ids[0] + + # Ulysses: shard sequences and position IDs before RoPE + if self.use_ulysses: + img_seq_len = img_ids.shape[0] + _txt_seq_len = txt_ids.shape[0] + + if img_seq_len % self.ulysses_size != 0: + raise ValueError( + f"Image seq len ({img_seq_len}) not divisible by " + f"ulysses_size ({self.ulysses_size})" + ) + if _txt_seq_len % self.ulysses_size != 0: + raise ValueError( + f"Text seq len ({_txt_seq_len}) not divisible by " + f"ulysses_size ({self.ulysses_size})" + ) + + img_chunk = img_seq_len // self.ulysses_size + txt_chunk = _txt_seq_len // self.ulysses_size + r = self.ulysses_rank + + # Shard position IDs (before RoPE computation) + img_ids = img_ids[r * img_chunk : (r + 1) * img_chunk] + txt_ids = txt_ids[r * txt_chunk : (r + 1) * txt_chunk] + + # Shard hidden states + hidden_states = hidden_states[:, r * img_chunk : (r + 1) * img_chunk, :] + encoder_hidden_states = encoder_hidden_states[:, r * txt_chunk : (r + 1) * txt_chunk, :] + + # Update txt_seq_len to local (sharded) length for single-stream split + txt_seq_len = txt_chunk + + # Compute RoPE embeddings (4-axis, from potentially sharded IDs) + ids = torch.cat([txt_ids, img_ids], dim=0) + image_rotary_emb = self.pos_embed(ids) + + # Compute modulation parameters (shared across all blocks) + img_mod = self.double_stream_modulation_img(temb) # ((s1,sc1,g1), (s2,sc2,g2)) + txt_mod = self.double_stream_modulation_txt(temb) # ((s1,sc1,g1), (s2,sc2,g2)) + single_mod = self.single_stream_modulation(temb) # ((shift, scale, gate),) + + # Dual-stream blocks + for block in self.transformer_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + img_mod=img_mod, + txt_mod=txt_mod, + ) + + # Concatenate for single-stream blocks + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + # Single-stream blocks + for block in self.single_transformer_blocks: + hidden_states = block( + hidden_states=hidden_states, + image_rotary_emb=image_rotary_emb, + mod=single_mod[0], # Single tuple of (shift, scale, gate) + ) + + # Extract image features (discard text) + hidden_states = hidden_states[:, txt_seq_len:, :] + + # Ulysses: gather output sequence from all ranks + if self.use_ulysses: + hidden_states = hidden_states.contiguous() + gathered = [torch.zeros_like(hidden_states) for _ in range(self.ulysses_size)] + torch.distributed.all_gather(gathered, hidden_states, group=self.ulysses_pg) + hidden_states = torch.cat(gathered, dim=1) + + # Output projection + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if return_dict: + return {"sample": output} + return (output,) + + def load_weights(self, weights: dict) -> None: + """Load weights into the transformer. + + Args: + weights: Dictionary of parameter name -> tensor (from safetensors) + """ + # Remap HF checkpoint keys to our module attribute names + weights = _remap_checkpoint_keys(weights) + # Remap FLUX.2 FFN keys (linear_out -> down_proj) + remapped = {} + for key, value in weights.items(): + new_key = key + for old, new in _FLUX2_WEIGHT_KEY_REMAPS: + new_key = new_key.replace(old, new) + remapped[new_key] = value + weights = remapped + + # Split pre-fused linear_in weights into gate/up halves for GatedMLP. + # HF checkpoint has single linear_in.weight [2*intermediate, hidden], + # but GatedMLP.gate_up_proj uses FUSED_GATE_UP_LINEAR mode which + # expects separate gate and up weights to concatenate during loading. + keys_to_add = {} + keys_to_remove = [] + for key, value in weights.items(): + if ".linear_in.weight" in key: + prefix = key.replace("linear_in.weight", "") + gate, up = value.chunk(2, dim=0) + keys_to_add[f"{prefix}linear_in_gate.weight"] = gate + keys_to_add[f"{prefix}linear_in_up.weight"] = up + keys_to_remove.append(key) + elif ".linear_in.bias" in key: + prefix = key.replace("linear_in.bias", "") + gate, up = value.chunk(2, dim=0) + keys_to_add[f"{prefix}linear_in_gate.bias"] = gate + keys_to_add[f"{prefix}linear_in_up.bias"] = up + keys_to_remove.append(key) + for k in keys_to_remove: + del weights[k] + weights.update(keys_to_add) + + # Map fused layer names to original HF checkpoint names. + # The loader concatenates these checkpoint keys into fused module weights. + params_map = { + "add_qkv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], + "qkv_proj": ["to_q", "to_k", "to_v"], + "gate_up_proj": ["linear_in_gate", "linear_in_up"], + } + + loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map) + + for name, module in tqdm(self.named_modules(), desc="Loading FLUX.2 weights"): + # Create weights for modules with skip_create_weights_in_init=True + if callable(getattr(module, "create_weights", None)): + module.create_weights() + + if len(module._parameters) == 0: + continue + + if isinstance(module, Linear): + weight_dicts = loader.get_linear_weights(module, name, weights) + if weight_dicts: + loader.load_linear_weights(module, name, weight_dicts) + else: + # For non-Linear modules, load weights directly + module_weights = loader.filter_weights(name, weights) + target_dtype = self.model_config.torch_dtype + for param_name, param in module._parameters.items(): + if param is not None and param_name in module_weights: + param.data.copy_(module_weights[param_name].to(target_dtype)) + + def post_load_weights(self) -> None: + """Call post_load_weights on all Linear modules and convert embedders to target dtype.""" + target_dtype = self.model_config.torch_dtype + + # Convert time embedding components to target dtype + if hasattr(self.time_guidance_embed, "timestep_embedder"): + self.time_guidance_embed.timestep_embedder.to(target_dtype) + if self.time_guidance_embed.guidance_embedder is not None: + self.time_guidance_embed.guidance_embedder.to(target_dtype) + + # Call post_load_weights on all Linear modules + for _, module in self.named_modules(): + if isinstance(module, Linear): + module.post_load_weights() diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 0c83bf5e2882..5c5a8013c1d7 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -44,7 +44,9 @@ def __init__( head_dim: Optional[int] = None, qkv_mode: QKVMode = QKVMode.FUSE_QKV, qk_norm: bool = True, - eps: float = 1e-6, # TODO: remove this, we should add this to the config + qk_norm_mode: str = "full", + eps: float = 1e-6, + bias: bool = True, config: Optional["DiffusionModelConfig"] = None, layer_idx: Optional[int] = None, ): @@ -62,6 +64,7 @@ def __init__( self.num_key_value_heads = num_key_value_heads or num_attention_heads self.head_dim = head_dim or (hidden_size // num_attention_heads) self.qkv_mode = QKVMode(qkv_mode) if isinstance(qkv_mode, str) else qkv_mode + self.bias = bias # Select compute backend (orthogonal to parallelism) ulysses_size = config.parallel.dit_ulysses_size @@ -82,11 +85,17 @@ def __init__( self._init_qkv_proj() if self.qk_norm: + if qk_norm_mode == "per_head": + q_norm_dim = self.head_dim + kv_norm_dim = self.head_dim + else: + q_norm_dim = self.q_dim + kv_norm_dim = self.kv_dim self.norm_q = RMSNorm( - hidden_size=self.q_dim, eps=self.eps, dtype=self.dtype, has_weights=True + hidden_size=q_norm_dim, eps=self.eps, dtype=self.dtype, has_weights=True ) self.norm_k = RMSNorm( - hidden_size=self.kv_dim, eps=self.eps, dtype=self.dtype, has_weights=True + hidden_size=kv_norm_dim, eps=self.eps, dtype=self.dtype, has_weights=True ) # TODO: Use weight mapper to create just a Linear module @@ -95,6 +104,7 @@ def __init__( Linear( self.q_dim, self.hidden_size, + bias=self.bias, dtype=self.dtype, mapping=self.mapping, quant_config=self.quant_config, @@ -140,6 +150,7 @@ def _init_qkv_proj(self) -> None: self.qkv_proj = Linear( self.hidden_size, qkv_out_dim, + bias=self.bias, dtype=self.dtype, mapping=self.mapping, quant_config=self.quant_config, @@ -158,6 +169,7 @@ def _init_qkv_proj(self) -> None: self.to_q = Linear( self.hidden_size, self.q_dim, + bias=self.bias, dtype=self.dtype, mapping=self.mapping, quant_config=self.quant_config, @@ -167,6 +179,7 @@ def _init_qkv_proj(self) -> None: self.to_k = Linear( self.hidden_size, self.kv_dim, + bias=self.bias, dtype=self.dtype, mapping=self.mapping, quant_config=self.quant_config, @@ -176,6 +189,7 @@ def _init_qkv_proj(self) -> None: self.to_v = Linear( self.hidden_size, self.kv_dim, + bias=self.bias, dtype=self.dtype, mapping=self.mapping, quant_config=self.quant_config, @@ -202,8 +216,13 @@ def get_qkv( def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: if self.qk_norm: - q = self.norm_q(q) - k = self.norm_k(k) + if q.ndim == 4: + shape = q.shape + q = self.norm_q(q.reshape(-1, shape[-1])).view(shape) + k = self.norm_k(k.reshape(-1, k.shape[-1])).view(k.shape) + else: + q = self.norm_q(q) + k = self.norm_k(k) return q, k def _attn_impl( diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index f4c7fc37da25..25ee07010919 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -2,7 +2,7 @@ Follows: DiffusionArgs → PipelineLoader → DiffusionModelConfig → AutoPipeline → BasePipeline -All pipelines (Wan, Flux2, LTX2) register via @register_pipeline decorator. +All pipelines (Wan, Flux, Flux2, LTX2) register via @register_pipeline decorator. """ import json @@ -83,6 +83,9 @@ def _detect_from_checkpoint(checkpoint_dir: str) -> str: # Generic Wan (T2V) if "Wan" in class_name: return "WanPipeline" + # Check FLUX.2 before FLUX.1 (more specific match first) + if "Flux2" in class_name: + return "Flux2Pipeline" if "Flux" in class_name: return "FluxPipeline" if "LTX" in class_name or "Ltx" in class_name: diff --git a/tensorrt_llm/_torch/visual_gen/quantization/loader.py b/tensorrt_llm/_torch/visual_gen/quantization/loader.py index a4a2a3a11c86..824b5416a511 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/loader.py +++ b/tensorrt_llm/_torch/visual_gen/quantization/loader.py @@ -59,7 +59,7 @@ def get_linear_weights( weights_config = getattr(module, "weights_loading_config", None) if weights_config is not None: weight_mode = getattr(weights_config, "weight_mode", None) - if weight_mode == WeightMode.FUSED_QKV_LINEAR: + if weight_mode in (WeightMode.FUSED_QKV_LINEAR, WeightMode.FUSED_GATE_UP_LINEAR): fused_names = self._get_fused_names(full_name) return self._get_fused_weights(full_name, weights, fused_names) @@ -86,7 +86,7 @@ def filter_weights( def _get_fused_names(self, full_name: str) -> List[str]: """Get checkpoint names for a fused module from params_map.""" for suffix, names in self.params_map.items(): - if full_name.endswith(suffix): + if full_name == suffix or full_name.endswith("." + suffix): return names raise ValueError( f"No params_map entry for fused module '{full_name}'. " diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d355882bf003..e9aecd4ed1a7 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -102,6 +102,9 @@ l0_b200: - unittest/_torch/visual_gen/test_wan.py -k "not TestWanTwoStageTransformer" - unittest/_torch/visual_gen/test_wan_i2v.py - unittest/_torch/visual_gen/test_model_loader.py + - unittest/_torch/visual_gen/test_flux_transformer.py + - unittest/_torch/visual_gen/test_flux_attention.py + - unittest/_torch/visual_gen/test_flux_pipeline.py # - examples/test_visual_gen.py - condition: ranges: diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 9f9bb268bcde..06569f0b0c76 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -36,6 +36,8 @@ l0_dgx_b200: - unittest/_torch/visual_gen/test_wan.py::TestWanCombinedOptimizations::test_all_optimizations_combined - unittest/_torch/visual_gen/test_wan_i2v.py::TestWanI2VParallelism::test_cfg_2gpu_correctness - unittest/_torch/visual_gen/test_wan_i2v.py::TestWanI2VCombinedOptimizations::test_all_optimizations_combined + - unittest/_torch/visual_gen/test_flux_pipeline.py::TestFluxParallelism::test_ulysses_2gpu_correctness + - unittest/_torch/visual_gen/test_flux_pipeline.py::TestFluxCombinedOptimizations::test_all_optimizations_combined - condition: ranges: system_gpu_count: diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py new file mode 100644 index 000000000000..59993bad04f5 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py @@ -0,0 +1,438 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Multi-GPU tests for FLUX Ulysses sequence parallelism. + +Tests that FLUX.1 and FLUX.2 transformers produce correct outputs when using +Ulysses sequence parallelism (sharding sequence across GPUs). + +Run with: + pytest tests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.py -v +""" + +import os + +os.environ["TLLM_DISABLE_MPI"] = "1" + +from types import SimpleNamespace +from typing import Callable + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +try: + from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionModelConfig, + ParallelConfig, + PipelineConfig, + TeaCacheConfig, + ) + from tensorrt_llm._utils import get_free_port + from tensorrt_llm.models.modeling_utils import QuantConfig + + MODULES_AVAILABLE = True +except ImportError: + MODULES_AVAILABLE = False + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + """Clean up TLLM_DISABLE_MPI env var after tests complete.""" + yield + os.environ.pop("TLLM_DISABLE_MPI", None) + + +# ============================================================================= +# Distributed helpers (same pattern as test_ulysses_attention.py) +# ============================================================================= + + +def init_distributed_worker(rank: int, world_size: int, backend: str = "nccl", port: int = 29500): + """Initialize distributed environment for a worker process.""" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + + torch.cuda.set_device(rank % torch.cuda.device_count()) + dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + + +def cleanup_distributed(): + """Clean up distributed environment.""" + if dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_worker(rank, world_size, backend, test_fn, port): + """Worker function that runs in each process. Module-level for pickling.""" + try: + init_distributed_worker(rank, world_size, backend, port) + test_fn(rank, world_size) + except Exception as e: + print(f"Rank {rank} failed with error: {e}") + raise + finally: + cleanup_distributed() + + +def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = True): + """Run a test function in a distributed environment.""" + if not MODULES_AVAILABLE: + pytest.skip("Required modules not available") + + if use_cuda and torch.cuda.device_count() < world_size: + pytest.skip(f"Test requires {world_size} GPUs, only {torch.cuda.device_count()} available") + + backend = "nccl" if use_cuda else "gloo" + port = get_free_port() + + mp.spawn( + _distributed_worker, args=(world_size, backend, test_fn, port), nprocs=world_size, join=True + ) + + +# ============================================================================= +# Model config helpers +# ============================================================================= + +# Small FLUX.1 config for testing (reduced layers, 8 heads, 64 head_dim = 512 inner_dim) +_FLUX1_TEST_CONFIG = dict( + num_attention_heads=8, + attention_head_dim=64, + in_channels=64, + out_channels=64, + num_layers=2, + num_single_layers=4, + joint_attention_dim=256, + pooled_projection_dim=128, + guidance_embeds=False, + patch_size=1, + axes_dims_rope=[16, 24, 24], + theta=10000, +) + +# Small FLUX.2 config for testing +_FLUX2_TEST_CONFIG = dict( + num_attention_heads=8, + attention_head_dim=64, + in_channels=128, + out_channels=128, + num_layers=2, + num_single_layers=4, + joint_attention_dim=256, + pooled_projection_dim=128, + guidance_embeds=False, + patch_size=1, + mlp_ratio=3.0, + axes_dims_rope=[16, 16, 16, 16], + rope_theta=2000.0, + eps=1e-6, + timestep_guidance_channels=256, +) + + +def _make_model_config(pretrained_dict, ulysses_size=1): + """Create DiffusionModelConfig for testing.""" + pretrained_config = SimpleNamespace(**pretrained_dict) + parallel = ParallelConfig(dit_ulysses_size=ulysses_size) + + return DiffusionModelConfig( + pretrained_config=pretrained_config, + quant_config=QuantConfig(), + pipeline=PipelineConfig(enable_torch_compile=False), + attention=AttentionConfig(backend="VANILLA"), + parallel=parallel, + teacache=TeaCacheConfig(), + skip_create_weights_in_init=False, + ) + + +def _stabilize_model_weights(model): + """Reinitialize model weights for stable BF16 forward pass. + + Random default init (std~1.0) causes BF16 overflow through multiple + transformer blocks. Use small uniform init that keeps activations bounded. + """ + with torch.no_grad(): + for name, p in model.named_parameters(): + if p.ndim >= 2: + # Xavier-like: scale by 1/sqrt(fan_in) + fan_in = p.shape[1] if p.ndim >= 2 else p.shape[0] + std = 0.02 / max(1.0, fan_in**0.5) + p.data.uniform_(-std, std) + else: + # Bias/1D params: small values + p.data.uniform_(-0.01, 0.01) + + +# ============================================================================= +# FLUX.1 test logic functions (module-level for pickling) +# ============================================================================= + + +def _logic_flux1_ulysses_forward(rank, world_size): + """FLUX.1 transformer forward with Ulysses — verify shape and no NaN/Inf. + + Uses BF16 (required by flashinfer RMSNorm) with scaled-down weights to + prevent overflow through multiple transformer blocks. + """ + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux import FluxTransformer2DModel + + device = torch.device(f"cuda:{rank}") + + torch.manual_seed(42) # Same seed on all ranks for identical model weights + + model_config = _make_model_config(_FLUX1_TEST_CONFIG, ulysses_size=world_size) + model = FluxTransformer2DModel(model_config).to(device).to(torch.bfloat16) + _stabilize_model_weights(model) + + batch = 1 + img_seq = 16 # Must be divisible by world_size + txt_seq = 8 # Must be divisible by world_size + + # Same inputs on all ranks (required for Ulysses — each rank shards the same input) + torch.manual_seed(100) + hidden_states = torch.randn(batch, img_seq, 64, device=device, dtype=torch.bfloat16) * 0.1 + encoder_hidden_states = ( + torch.randn(batch, txt_seq, 256, device=device, dtype=torch.bfloat16) * 0.1 + ) + pooled_projections = torch.randn(batch, 128, device=device, dtype=torch.bfloat16) * 0.1 + timestep = torch.tensor([0.5], device=device, dtype=torch.bfloat16) + img_ids = torch.randn(img_seq, 3, device=device) + txt_ids = torch.randn(txt_seq, 3, device=device) + + with torch.no_grad(): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + sample = output["sample"] + # Output should have full image sequence (gathered) + assert sample.shape == (batch, img_seq, 64), ( + f"Rank {rank}: Expected shape {(batch, img_seq, 64)}, got {sample.shape}" + ) + assert not torch.isnan(sample).any(), f"Rank {rank}: NaN in output" + assert not torch.isinf(sample).any(), f"Rank {rank}: Inf in output" + + +def _logic_flux1_ulysses_vs_single_gpu(rank, world_size): + """FLUX.1: Ulysses 2-GPU output matches single-GPU reference. + + Uses BF16 (required by flashinfer RMSNorm) with scaled-down weights. + Ulysses all-to-all is a pure data shuffle so results should be nearly + identical — the only drift is from BF16 rounding in attention. + """ + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux import FluxTransformer2DModel + + device = torch.device(f"cuda:{rank}") + compute_dtype = torch.bfloat16 + + batch = 1 + img_seq = 16 + txt_seq = 8 + + # Create single-GPU reference model with shared seed + torch.manual_seed(123) + ref_config = _make_model_config(_FLUX1_TEST_CONFIG, ulysses_size=1) + ref_model = FluxTransformer2DModel(ref_config).to(device).to(compute_dtype) + _stabilize_model_weights(ref_model) + ref_state = ref_model.state_dict() + + # Create Ulysses model with same weights + torch.manual_seed(123) + ulysses_config = _make_model_config(_FLUX1_TEST_CONFIG, ulysses_size=world_size) + ulysses_model = FluxTransformer2DModel(ulysses_config).to(device).to(compute_dtype) + ulysses_model.load_state_dict(ref_state) + + # Same inputs on all ranks + torch.manual_seed(456) + hidden_states = torch.randn(batch, img_seq, 64, device=device, dtype=compute_dtype) * 0.1 + encoder_hidden_states = ( + torch.randn(batch, txt_seq, 256, device=device, dtype=compute_dtype) * 0.1 + ) + pooled_projections = torch.randn(batch, 128, device=device, dtype=compute_dtype) * 0.1 + timestep = torch.tensor([0.5], device=device, dtype=compute_dtype) + img_ids = torch.randn(img_seq, 3, device=device) + txt_ids = torch.randn(txt_seq, 3, device=device) + + with torch.no_grad(): + ref_output = ref_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + ) + ulysses_output = ulysses_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + torch.testing.assert_close( + ulysses_output["sample"], + ref_output["sample"], + rtol=1e-2, + atol=1e-2, + msg=f"Rank {rank}: FLUX.1 Ulysses output differs from single-GPU reference", + ) + + +# ============================================================================= +# FLUX.2 test logic functions (module-level for pickling) +# ============================================================================= + + +def _logic_flux2_ulysses_forward(rank, world_size): + """FLUX.2 transformer forward with Ulysses — verify shape and no NaN/Inf. + + Uses BF16 (required by flashinfer RMSNorm) with scaled-down weights. + """ + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux2 import Flux2Transformer2DModel + + device = torch.device(f"cuda:{rank}") + + torch.manual_seed(42) # Same seed on all ranks for identical model weights + + model_config = _make_model_config(_FLUX2_TEST_CONFIG, ulysses_size=world_size) + model = Flux2Transformer2DModel(model_config).to(device).to(torch.bfloat16) + _stabilize_model_weights(model) + + batch = 1 + img_seq = 16 + txt_seq = 8 + + # Same inputs on all ranks + torch.manual_seed(100) + hidden_states = torch.randn(batch, img_seq, 128, device=device, dtype=torch.bfloat16) * 0.1 + encoder_hidden_states = ( + torch.randn(batch, txt_seq, 256, device=device, dtype=torch.bfloat16) * 0.1 + ) + timestep = torch.tensor([0.5], device=device, dtype=torch.bfloat16) + img_ids = torch.randn(img_seq, 4, device=device) + txt_ids = torch.randn(txt_seq, 4, device=device) + + with torch.no_grad(): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + sample = output["sample"] + # Output should have full image sequence (gathered), out_channels=128 + assert sample.shape == (batch, img_seq, 128), ( + f"Rank {rank}: Expected shape {(batch, img_seq, 128)}, got {sample.shape}" + ) + assert not torch.isnan(sample).any(), f"Rank {rank}: NaN in output" + assert not torch.isinf(sample).any(), f"Rank {rank}: Inf in output" + + +def _logic_flux2_ulysses_vs_single_gpu(rank, world_size): + """FLUX.2: Ulysses 2-GPU output matches single-GPU reference. + + Uses BF16 (required by flashinfer RMSNorm) with scaled-down weights. + """ + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux2 import Flux2Transformer2DModel + + device = torch.device(f"cuda:{rank}") + compute_dtype = torch.bfloat16 + + batch = 1 + img_seq = 16 + txt_seq = 8 + + # Create single-GPU reference model with shared seed + torch.manual_seed(123) + ref_config = _make_model_config(_FLUX2_TEST_CONFIG, ulysses_size=1) + ref_model = Flux2Transformer2DModel(ref_config).to(device).to(compute_dtype) + _stabilize_model_weights(ref_model) + ref_state = ref_model.state_dict() + + # Create Ulysses model with same weights + torch.manual_seed(123) + ulysses_config = _make_model_config(_FLUX2_TEST_CONFIG, ulysses_size=world_size) + ulysses_model = Flux2Transformer2DModel(ulysses_config).to(device).to(compute_dtype) + ulysses_model.load_state_dict(ref_state) + + # Same inputs on all ranks + torch.manual_seed(456) + hidden_states = torch.randn(batch, img_seq, 128, device=device, dtype=compute_dtype) * 0.1 + encoder_hidden_states = ( + torch.randn(batch, txt_seq, 256, device=device, dtype=compute_dtype) * 0.1 + ) + timestep = torch.tensor([0.5], device=device, dtype=compute_dtype) + img_ids = torch.randn(img_seq, 4, device=device) + txt_ids = torch.randn(txt_seq, 4, device=device) + + with torch.no_grad(): + ref_output = ref_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + ) + ulysses_output = ulysses_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + torch.testing.assert_close( + ulysses_output["sample"], + ref_output["sample"], + rtol=1e-2, + atol=1e-2, + msg=f"Rank {rank}: FLUX.2 Ulysses output differs from single-GPU reference", + ) + + +# ============================================================================= +# Test classes +# ============================================================================= + + +class TestFlux1Ulysses: + """Ulysses sequence parallelism tests for FLUX.1 transformer.""" + + def test_flux1_ulysses_forward(self): + """FLUX.1 Ulysses forward: correct output shape and no NaN/Inf.""" + run_test_in_distributed(world_size=2, test_fn=_logic_flux1_ulysses_forward) + + def test_flux1_ulysses_vs_single_gpu(self): + """FLUX.1 Ulysses 2-GPU output matches single-GPU reference.""" + run_test_in_distributed(world_size=2, test_fn=_logic_flux1_ulysses_vs_single_gpu) + + +class TestFlux2Ulysses: + """Ulysses sequence parallelism tests for FLUX.2 transformer.""" + + def test_flux2_ulysses_forward(self): + """FLUX.2 Ulysses forward: correct output shape and no NaN/Inf.""" + run_test_in_distributed(world_size=2, test_fn=_logic_flux2_ulysses_forward) + + def test_flux2_ulysses_vs_single_gpu(self): + """FLUX.2 Ulysses 2-GPU output matches single-GPU reference.""" + run_test_in_distributed(world_size=2, test_fn=_logic_flux2_ulysses_vs_single_gpu) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/visual_gen/test_flux_attention.py b/tests/unittest/_torch/visual_gen/test_flux_attention.py new file mode 100644 index 000000000000..93621497e044 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_flux_attention.py @@ -0,0 +1,313 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for FLUX attention backends. + +Tests cover: +- VANILLA backend (PyTorch SDPA) +- TRTLLM backend (fmha_v2) +- Backend equivalence comparison + +Note: With random weights, attention can produce NaN due to numerical instability. + These tests use scaled inputs and primarily verify correct output shapes. + Full numerical correctness is tested via HuggingFace comparison tests with real weights. +""" + +import unittest +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.visual_gen.config import AttentionConfig, DiffusionModelConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig + + +class TestFluxAttentionBackend(unittest.TestCase): + """Test FLUX attention with different backends.""" + + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + def _create_config(self, backend: str) -> DiffusionModelConfig: + """Create DiffusionModelConfig with specified backend.""" + return DiffusionModelConfig( + pretrained_config=SimpleNamespace(), + quant_config=QuantConfig(), + mapping=Mapping(), + attention=AttentionConfig(backend=backend), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_vanilla_backend_sanity(self): + """Test FLUX attention works with VANILLA backend.""" + from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention + + batch_size = 2 + seq_len = 256 + text_seq_len = 64 + dim = 3072 + heads = 24 + dim_head = 128 + dtype = torch.bfloat16 + + torch.manual_seed(42) + config = self._create_config("VANILLA") + + attn = ( + FluxJointAttention( + hidden_size=dim, + num_attention_heads=heads, + head_dim=dim_head, + added_kv_proj_dim=dim, # Enable dual-stream for text tokens + config=config, + layer_idx=0, + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + # Use scaled inputs to reduce numerical instability + hidden_states = ( + torch.randn(batch_size, seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + encoder_hidden_states = ( + torch.randn(batch_size, text_seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + + with torch.no_grad(): + # Skip RoPE for this sanity test (pass None) + output, text_output = attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=None, + ) + + # With random weights, NaN can occur. For unit tests, we primarily check shapes. + self.assertEqual(output.shape, hidden_states.shape) + self.assertEqual(text_output.shape, encoder_hidden_states.shape) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_trtllm_backend_sanity(self): + """Test FLUX attention works with TRTLLM backend.""" + from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention + + batch_size = 2 + seq_len = 256 + text_seq_len = 64 + dim = 3072 + heads = 24 + dim_head = 128 + dtype = torch.bfloat16 + + torch.manual_seed(42) + config = self._create_config("TRTLLM") + + attn = ( + FluxJointAttention( + hidden_size=dim, + num_attention_heads=heads, + head_dim=dim_head, + added_kv_proj_dim=dim, # Enable dual-stream for text tokens + config=config, + layer_idx=0, + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + # Use scaled inputs to reduce numerical instability + hidden_states = ( + torch.randn(batch_size, seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + encoder_hidden_states = ( + torch.randn(batch_size, text_seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + + with torch.no_grad(): + # Skip RoPE for this sanity test (pass None) + output, text_output = attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=None, + ) + + # With random weights, NaN can occur. For unit tests, we primarily check shapes. + self.assertEqual(output.shape, hidden_states.shape) + self.assertEqual(text_output.shape, encoder_hidden_states.shape) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backend_equivalence(self): + """Test VANILLA and TRTLLM backends produce similar outputs.""" + from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention + + batch_size = 1 + seq_len = 128 + text_seq_len = 32 + dim = 3072 + heads = 24 + dim_head = 128 + dtype = torch.bfloat16 + + torch.manual_seed(42) + + # Create attention modules for both backends + config = self._create_config("VANILLA") + vanilla_attn = ( + FluxJointAttention( + hidden_size=dim, + num_attention_heads=heads, + head_dim=dim_head, + added_kv_proj_dim=dim, # Enable dual-stream for text tokens + config=config, + layer_idx=0, + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + # TRT-LLM Linear uses torch.empty for bias (uninitialized memory). + # After prior tests free GPU memory, recycled memory can contain NaN. + # Initialize all parameters with small random values for numerical stability. + with torch.no_grad(): + for p in vanilla_attn.parameters(): + p.normal_(0, 0.02) + + config = self._create_config("TRTLLM") + trtllm_attn = ( + FluxJointAttention( + hidden_size=dim, + num_attention_heads=heads, + head_dim=dim_head, + added_kv_proj_dim=dim, + config=config, + layer_idx=0, + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + # Copy scaled weights from VANILLA to TRTLLM + trtllm_attn.load_state_dict(vanilla_attn.state_dict()) + attns = {"VANILLA": vanilla_attn, "TRTLLM": trtllm_attn} + + # Create inputs (scaled for numerical stability) + hidden_states = ( + torch.randn(batch_size, seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + encoder_hidden_states = ( + torch.randn(batch_size, text_seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + + # Run both backends (skip RoPE for equivalence test) + outputs = {} + with torch.no_grad(): + for backend in ["VANILLA", "TRTLLM"]: + out, text_out = attns[backend]( + hidden_states=hidden_states.clone(), + encoder_hidden_states=encoder_hidden_states.clone(), + image_rotary_emb=None, + ) + outputs[backend] = (out, text_out) + + # Compare outputs + vanilla_out, vanilla_text = outputs["VANILLA"] + trtllm_out, trtllm_text = outputs["TRTLLM"] + + # Skip comparison if either has NaN or Inf (common with random weights) + has_nan = torch.isnan(vanilla_out).any() or torch.isnan(trtllm_out).any() + has_inf = torch.isinf(vanilla_out).any() or torch.isinf(trtllm_out).any() + if has_nan or has_inf: + self.skipTest("NaN/Inf detected in outputs with random weights - skipping comparison") + + # With random weights, outputs may be all zeros or have numerical issues + # This test is primarily for ensuring both backends can run + # Full equivalence is tested via HuggingFace comparison with real weights + vanilla_norm = vanilla_out.float().norm().item() + trtllm_norm = trtllm_out.float().norm().item() + + print(f"\n[Debug] vanilla_norm={vanilla_norm:.6f}, trtllm_norm={trtllm_norm:.6f}") + + # With random weights, we can only reliably check that both backends produce + # valid outputs (same shapes, non-trivial values). Strict equivalence requires + # real weights from a trained model (tested via HuggingFace comparison). + self.assertEqual(vanilla_out.shape, trtllm_out.shape) + self.assertEqual(vanilla_text.shape, trtllm_text.shape) + + # If both outputs have meaningful norms, compute similarity as informational + if vanilla_norm > 1e-3 and trtllm_norm > 1e-3: + cos_sim = F.cosine_similarity( + vanilla_out.float().flatten().unsqueeze(0), + trtllm_out.float().flatten().unsqueeze(0), + ).item() + print(f" Cosine similarity: {cos_sim:.6f}") + # Note: Not asserting on cos_sim with random weights as it's not meaningful + + +class TestFlux2AttentionBackend(unittest.TestCase): + """Test FLUX.2 attention with different backends.""" + + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + def _create_config(self, backend: str) -> DiffusionModelConfig: + """Create DiffusionModelConfig with specified backend.""" + return DiffusionModelConfig( + pretrained_config=SimpleNamespace(), + quant_config=QuantConfig(), + mapping=Mapping(), + attention=AttentionConfig(backend=backend), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux2_vanilla_backend_sanity(self): + """Test FLUX.2 attention works with VANILLA backend.""" + from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention + + batch_size = 2 + seq_len = 128 + text_seq_len = 64 + dim = 6144 # FLUX.2 has larger dim + heads = 48 + dim_head = 128 + dtype = torch.bfloat16 + + torch.manual_seed(42) + config = self._create_config("VANILLA") + + attn = ( + FluxJointAttention( + hidden_size=dim, + num_attention_heads=heads, + head_dim=dim_head, + added_kv_proj_dim=dim, # Enable dual-stream for text tokens + config=config, + layer_idx=0, + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + # Use scaled inputs to reduce numerical instability + hidden_states = ( + torch.randn(batch_size, seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + encoder_hidden_states = ( + torch.randn(batch_size, text_seq_len, dim, device=self.DEVICE, dtype=dtype) * 0.02 + ) + + with torch.no_grad(): + # Skip RoPE for this sanity test (pass None) + output, text_output = attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=None, + ) + + # With random weights, NaN can occur. For unit tests, we primarily check shapes. + self.assertEqual(output.shape, hidden_states.shape) + self.assertEqual(text_output.shape, encoder_hidden_states.shape) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py new file mode 100644 index 000000000000..608135af92c8 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py @@ -0,0 +1,1167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for FLUX pipelines. + +Tests cover: +- Pipeline loading (FLUX.1 and FLUX.2) +- Quantization (FP8, FP8_BLOCK_SCALES) +- Single-layer numerical correctness (F.linear reference) +- Full transformer E2E numerical correctness +- Memory usage comparison +- Attention backend comparison (VANILLA vs TRTLLM) +- Multi-GPU parallelism (Ulysses sequence parallelism, 2+ GPUs) +""" + +import gc +import os +from pathlib import Path + +import numpy as np +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.visual_gen.config import AttentionConfig, DiffusionArgs, PipelineConfig +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + + +def _llm_models_root() -> str: + """Return LLM_MODELS_ROOT path if it is set in env, assert when it's set but not a valid path.""" + root = Path("/home/scratch.trt_llm_data_ci/llm-models/") + if "LLM_MODELS_ROOT" in os.environ: + root = Path(os.environ["LLM_MODELS_ROOT"]) + if not root.exists(): + root = Path("/scratch.trt_llm_data/llm-models/") + assert root.exists(), ( + "You shall set LLM_MODELS_ROOT env or be able to access scratch.trt_llm_data to run this test" + ) + return str(root) + + +# Checkpoint paths for integration tests +FLUX1_CHECKPOINT_PATH = os.environ.get( + "FLUX1_MODEL_PATH", + os.path.join(_llm_models_root(), "FLUX.1-dev"), +) +FLUX2_CHECKPOINT_PATH = os.environ.get( + "FLUX2_MODEL_PATH", + os.path.join(_llm_models_root(), "FLUX.2-dev"), +) +SKIP_COMPONENTS = ["text_encoder", "text_encoder_2", "vae", "tokenizer", "tokenizer_2", "scheduler"] +# When skip_components includes tokenizer, warmup must be disabled (warmup calls _encode_prompt) +PIPELINE_NO_WARMUP = PipelineConfig(warmup_steps=0) + + +def _get_flux_transformer_inputs(transformer, device="cuda", dtype=torch.bfloat16): + """Create test inputs appropriate for a FLUX.1 transformer. + + Inspects the transformer's config to determine the correct input shapes. + Note: Generates 3-axis position IDs (FLUX.1 only). FLUX.2 uses 4-axis IDs. + """ + torch.manual_seed(42) + batch_size = 1 + seq_len = 64 + text_seq_len = 32 + + config = transformer.config + in_channels = getattr(config, "in_channels", 64) + joint_attention_dim = getattr(config, "joint_attention_dim", 4096) + pooled_projection_dim = getattr(config, "pooled_projection_dim", 768) + + hidden_states = torch.randn(batch_size, seq_len, in_channels, device=device, dtype=dtype) + encoder_hidden_states = torch.randn( + batch_size, text_seq_len, joint_attention_dim, device=device, dtype=dtype + ) + pooled_projections = torch.randn(batch_size, pooled_projection_dim, device=device, dtype=dtype) + timestep = torch.tensor([500.0], device=device, dtype=dtype) + guidance = torch.tensor([3.5], device=device, dtype=dtype) + img_ids = torch.zeros(batch_size, seq_len, 3, device=device) + txt_ids = torch.zeros(batch_size, text_seq_len, 3, device=device) + + return dict( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + guidance=guidance, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + +def _extract_transformer_output(output): + """Extract tensor from transformer output (dict or tuple).""" + if isinstance(output, dict): + return output["sample"] + if isinstance(output, tuple): + return output[0] + return output + + +def _find_first_quantizable_linear(transformer): + """Find the first Linear layer in transformer blocks suitable for testing.""" + # Try attention QKV in first transformer block + if hasattr(transformer, "transformer_blocks") and len(transformer.transformer_blocks) > 0: + block = transformer.transformer_blocks[0] + if hasattr(block, "attn") and hasattr(block.attn, "qkv_proj"): + return block.attn.qkv_proj, "transformer_blocks.0.attn.qkv_proj" + # Try single transformer blocks + if ( + hasattr(transformer, "single_transformer_blocks") + and len(transformer.single_transformer_blocks) > 0 + ): + block = transformer.single_transformer_blocks[0] + if hasattr(block, "attn") and hasattr(block.attn, "qkv_proj"): + return block.attn.qkv_proj, "single_transformer_blocks.0.attn.qkv_proj" + # Fallback: first Linear in any block + for name, module in transformer.named_modules(): + if isinstance(module, Linear) and "blocks" in name: + return module, name + return None, None + + +@pytest.fixture +def flux1_checkpoint_exists(): + """Check if FLUX.1 checkpoint is available locally.""" + if not FLUX1_CHECKPOINT_PATH or not os.path.exists(FLUX1_CHECKPOINT_PATH): + pytest.skip( + f"FLUX.1 checkpoint not found at {FLUX1_CHECKPOINT_PATH}. " + "Set FLUX1_MODEL_PATH or stage checkpoint under LLM_MODELS_ROOT." + ) + return True + + +@pytest.fixture +def flux2_checkpoint_exists(): + """Check if FLUX.2 checkpoint is available locally.""" + if not FLUX2_CHECKPOINT_PATH or not os.path.exists(FLUX2_CHECKPOINT_PATH): + pytest.skip( + f"FLUX.2 checkpoint not found at {FLUX2_CHECKPOINT_PATH}. " + "Set FLUX2_MODEL_PATH or stage checkpoint under LLM_MODELS_ROOT." + ) + return True + + +# ============================================================================= +# Pipeline Loading Tests +# ============================================================================= + + +class TestFluxPipelineLoading: + """Integration tests for FLUX pipeline loading.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_load_flux1_pipeline_basic(self, flux1_checkpoint_exists): + """Test loading FLUX.1 pipeline.""" + args = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + pipeline=PIPELINE_NO_WARMUP, + ) + + pipeline = PipelineLoader(args).load() + + assert pipeline is not None + assert hasattr(pipeline, "transformer") + assert pipeline.transformer is not None + assert pipeline.model_config.attention.backend == "VANILLA" + + del pipeline + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_load_flux2_pipeline_basic(self, flux2_checkpoint_exists): + """Test loading FLUX.2 pipeline.""" + args = DiffusionArgs( + checkpoint_path=FLUX2_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + pipeline=PIPELINE_NO_WARMUP, + ) + + pipeline = PipelineLoader(args).load() + + assert pipeline is not None + assert hasattr(pipeline, "transformer") + assert pipeline.transformer is not None + + del pipeline + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("backend", ["VANILLA", "TRTLLM"]) + def test_load_flux1_with_attention_backend(self, flux1_checkpoint_exists, backend: str): + """Test loading FLUX.1 with different attention backends.""" + args = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + attention=AttentionConfig(backend=backend), + pipeline=PIPELINE_NO_WARMUP, + ) + + pipeline = PipelineLoader(args).load() + + assert pipeline.model_config.attention.backend == backend + + del pipeline + gc.collect() + torch.cuda.empty_cache() + + +# ============================================================================= +# Quantization Tests +# ============================================================================= + + +class TestFluxQuantization: + """Test FLUX quantization loading and FP8 weight verification.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"]) + def test_load_flux1_with_quantization(self, flux1_checkpoint_exists, quant_algo: str): + """Test loading FLUX.1 with FP8 quantization and verify FP8 weights.""" + args = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + quant_config={"quant_algo": quant_algo, "dynamic": True}, + pipeline=PIPELINE_NO_WARMUP, + ) + + pipeline = PipelineLoader(args).load() + + assert pipeline.model_config.quant_config.quant_algo is not None + + # Count quantized Linear layers and verify FP8 weights + quant_count = 0 + found_fp8 = False + for name, module in pipeline.transformer.named_modules(): + if isinstance(module, Linear): + if module.quant_config and module.quant_config.quant_algo: + quant_count += 1 + if "blocks" in name and hasattr(module, "weight") and module.weight is not None: + if not found_fp8: + assert module.weight.dtype == torch.float8_e4m3fn, ( + f"Linear {name} should have FP8 weight, got {module.weight.dtype}" + ) + assert hasattr(module, "weight_scale"), ( + f"Linear {name} missing weight_scale" + ) + found_fp8 = True + print( + f"\n[{quant_algo}] FP8 layer {name}: weight {module.weight.shape}" + ) + + print(f"[{quant_algo}] Quantized {quant_count} Linear layers") + assert quant_count > 0, "No layers were quantized" + assert found_fp8, f"No FP8 Linear modules found in blocks for {quant_algo}" + + del pipeline + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"]) + def test_load_flux2_with_quantization(self, flux2_checkpoint_exists, quant_algo: str): + """Test loading FLUX.2 with FP8 quantization and verify FP8 weights.""" + args = DiffusionArgs( + checkpoint_path=FLUX2_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + quant_config={"quant_algo": quant_algo, "dynamic": True}, + pipeline=PIPELINE_NO_WARMUP, + ) + + pipeline = PipelineLoader(args).load() + + assert pipeline.model_config.quant_config.quant_algo is not None + + quant_count = 0 + found_fp8 = False + for name, module in pipeline.transformer.named_modules(): + if isinstance(module, Linear): + if module.quant_config and module.quant_config.quant_algo: + quant_count += 1 + if "blocks" in name and hasattr(module, "weight") and module.weight is not None: + if not found_fp8: + assert module.weight.dtype == torch.float8_e4m3fn, ( + f"Linear {name} should have FP8 weight, got {module.weight.dtype}" + ) + assert hasattr(module, "weight_scale"), ( + f"Linear {name} missing weight_scale" + ) + found_fp8 = True + print( + f"\n[{quant_algo}] FP8 layer {name}: weight {module.weight.shape}" + ) + + print(f"[{quant_algo}] Quantized {quant_count} Linear layers") + assert quant_count > 0, "No layers were quantized" + assert found_fp8, f"No FP8 Linear modules found in blocks for {quant_algo}" + + del pipeline + gc.collect() + torch.cuda.empty_cache() + + +# ============================================================================= +# FP8 Numerical Correctness Tests +# ============================================================================= + + +class TestFluxFP8NumericalCorrectness: + """Test FP8 vs BF16 numerical accuracy at single-layer and full-transformer levels.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"]) + def test_fp8_vs_bf16_single_layer(self, flux1_checkpoint_exists, quant_algo: str): + """Test FP8 vs BF16 numerical accuracy on a single Linear layer. + + Pattern (matching Wan test_fp8_vs_bf16_numerical_correctness): + 1. Use F.linear() with BF16 weights as ground truth reference + 2. Verify BF16 layer matches F.linear exactly + 3. Compare FP8 layer output against reference + 4. Check max_diff, cosine_similarity, mse_loss + """ + # Load BF16 pipeline (reference) + print(f"\n[Compare {quant_algo}] Loading BF16 pipeline...") + args_bf16 = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_bf16 = PipelineLoader(args_bf16).load() + + # Load FP8 pipeline + print(f"[Compare {quant_algo}] Loading {quant_algo} pipeline...") + args_fp8 = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + quant_config={"quant_algo": quant_algo, "dynamic": True}, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_fp8 = PipelineLoader(args_fp8).load() + + # Get matching Linear layers from both pipelines + linear_bf16, layer_name = _find_first_quantizable_linear(pipeline_bf16.transformer) + linear_fp8, _ = _find_first_quantizable_linear(pipeline_fp8.transformer) + + assert linear_bf16 is not None, "Could not find a Linear layer in BF16 transformer" + assert linear_fp8 is not None, "Could not find a Linear layer in FP8 transformer" + + # Get BF16 weights for F.linear reference + weight_bf16 = linear_bf16.weight.data.clone() + bias_bf16 = linear_bf16.bias.data.clone() if linear_bf16.bias is not None else None + + # Create test input (2D for FP8 kernel compatibility) + torch.manual_seed(42) + hidden_size = linear_bf16.in_features + batch_seq_len = 1024 + + input_tensor = torch.randn(batch_seq_len, hidden_size, dtype=torch.bfloat16, device="cuda") + print(f"[Compare] Layer: {layer_name}, Input shape: {input_tensor.shape}") + + # Compute reference output: F.linear (ground truth) + with torch.no_grad(): + expected = F.linear(input_tensor, weight_bf16, bias_bf16) + + # Compute BF16 layer output + with torch.no_grad(): + result_bf16 = linear_bf16(input_tensor) + + # Compute FP8 output + with torch.no_grad(): + result_fp8 = linear_fp8(input_tensor) + + # Verify BF16 layer matches F.linear reference + assert torch.allclose(result_bf16, expected, rtol=1e-5, atol=1e-6), ( + "BF16 layer should match F.linear reference exactly" + ) + + # Compare FP8 vs reference + max_diff = torch.max(torch.abs(result_fp8 - expected)).item() + cos_sim = F.cosine_similarity( + result_fp8.flatten().float(), expected.flatten().float(), dim=0 + ) + mse = F.mse_loss(result_fp8.flatten().float(), expected.flatten().float()) + + print( + f"\n[{layer_name}] max_diff={max_diff:.6f}, cos_sim={cos_sim.item():.6f}, mse={mse.item():.6f}" + ) + + assert cos_sim > 0.99, f"Cosine similarity too low: {cos_sim.item()}" + assert mse < 1.0, f"MSE too high: {mse.item()}" + + del pipeline_bf16, pipeline_fp8 + torch.cuda.empty_cache() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("quant_algo", ["FP8", "FP8_BLOCK_SCALES"]) + def test_fp8_vs_bf16_full_transformer_e2e(self, flux1_checkpoint_exists, quant_algo: str): + """End-to-end test: Compare full FLUX.1 transformer FP8 vs BF16 output. + + Runs the entire transformer (19 dual + 38 single blocks) and compares outputs. + Errors accumulate across layers, so uses relaxed tolerances vs single-layer test. + """ + # Load BF16 transformer (reference) + print("\n[E2E] Loading BF16 transformer...") + args_bf16 = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_bf16 = PipelineLoader(args_bf16).load() + transformer_bf16 = pipeline_bf16.transformer + + # Load FP8 transformer + print(f"[E2E] Loading {quant_algo} transformer...") + args_fp8 = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + quant_config={"quant_algo": quant_algo, "dynamic": True}, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_fp8 = PipelineLoader(args_fp8).load() + transformer_fp8 = pipeline_fp8.transformer + + # Create test inputs + inputs = _get_flux_transformer_inputs(transformer_bf16) + + # Run both transformers + print("[E2E] Running BF16 transformer forward...") + with torch.no_grad(): + output_bf16 = transformer_bf16(**inputs) + + print(f"[E2E] Running {quant_algo} transformer forward...") + inputs_fp8 = {k: v.clone() for k, v in inputs.items()} + with torch.no_grad(): + output_fp8 = transformer_fp8(**inputs_fp8) + + # Extract outputs + output_bf16 = _extract_transformer_output(output_bf16) + output_fp8 = _extract_transformer_output(output_fp8) + + assert output_bf16.shape == output_fp8.shape, ( + f"Output shape mismatch: BF16={output_bf16.shape}, FP8={output_fp8.shape}" + ) + + # Check for NaN/Inf + assert not torch.isnan(output_bf16).any(), "BF16 output contains NaN" + assert not torch.isinf(output_bf16).any(), "BF16 output contains Inf" + assert not torch.isnan(output_fp8).any(), f"{quant_algo} output contains NaN" + assert not torch.isinf(output_fp8).any(), f"{quant_algo} output contains Inf" + + # Compare numerical accuracy + output_bf16_float = output_bf16.float() + output_fp8_float = output_fp8.float() + + max_diff = torch.max(torch.abs(output_fp8_float - output_bf16_float)).item() + mean_diff = torch.mean(torch.abs(output_fp8_float - output_bf16_float)).item() + + cos_sim = F.cosine_similarity( + output_fp8_float.flatten(), output_bf16_float.flatten(), dim=0 + ).item() + + mse = F.mse_loss(output_fp8_float, output_bf16_float).item() + rel_error = mean_diff / (output_bf16_float.abs().mean().item() + 1e-8) + + num_dual = len(transformer_bf16.transformer_blocks) + num_single = len(transformer_bf16.single_transformer_blocks) + + print(f"\n{'=' * 60}") + print(f"END-TO-END TRANSFORMER COMPARISON ({quant_algo} vs BF16)") + print(f"{'=' * 60}") + print(f"Number of layers: {num_dual} dual + {num_single} single") + print(f"Output shape: {output_bf16.shape}") + print("") + print(f"Max absolute difference: {max_diff:.6f}") + print(f"Mean absolute difference: {mean_diff:.6f}") + print(f"Relative error: {rel_error:.6f}") + print(f"Cosine similarity: {cos_sim:.6f}") + print(f"MSE loss: {mse:.6f}") + print("") + print(f"BF16 output range: [{output_bf16_float.min():.4f}, {output_bf16_float.max():.4f}]") + print( + f"{quant_algo} output range: [{output_fp8_float.min():.4f}, {output_fp8_float.max():.4f}]" + ) + print(f"{'=' * 60}") + + assert cos_sim > 0.95, ( + f"Cosine similarity too low for full transformer: {cos_sim:.6f} (expected >0.95)" + ) + assert rel_error < 0.15, f"Relative error too high: {rel_error:.6f} (expected <0.15)" + + print(f"\n[PASS] {quant_algo} full transformer output matches BF16 within tolerance!") + print(f" Cosine similarity: {cos_sim:.4f} (>0.95)") + print(f" Relative error: {rel_error:.4f} (<0.15)") + + del pipeline_bf16, pipeline_fp8, transformer_bf16, transformer_fp8 + torch.cuda.empty_cache() + + +# ============================================================================= +# FP8 Memory Comparison Tests +# ============================================================================= + + +class TestFluxFP8Memory: + """Test FP8 memory reduction for FLUX models.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_fp8_vs_bf16_memory_comparison(self, flux1_checkpoint_exists): + """Test FP8 uses ~2x less memory than BF16 (matching Wan test).""" + + def get_module_memory_gb(module): + return sum(p.numel() * p.element_size() for p in module.parameters()) / 1024**3 + + # Load BF16 + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + args_bf16 = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_bf16 = PipelineLoader(args_bf16).load() + + bf16_model_mem = get_module_memory_gb(pipeline_bf16.transformer) + bf16_peak_mem = torch.cuda.max_memory_allocated() / 1024**3 + + print(f"\n[BF16] Transformer memory: {bf16_model_mem:.2f} GB") + print(f"[BF16] Peak memory: {bf16_peak_mem:.2f} GB") + + del pipeline_bf16 + torch.cuda.empty_cache() + + # Load FP8 + torch.cuda.reset_peak_memory_stats() + + args_fp8 = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + quant_config={"quant_algo": "FP8", "dynamic": True}, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_fp8 = PipelineLoader(args_fp8).load() + + fp8_model_mem = get_module_memory_gb(pipeline_fp8.transformer) + fp8_peak_mem = torch.cuda.max_memory_allocated() / 1024**3 + + print(f"\n[FP8] Transformer memory: {fp8_model_mem:.2f} GB") + print(f"[FP8] Peak memory: {fp8_peak_mem:.2f} GB") + + # Verify memory savings + model_mem_ratio = bf16_model_mem / fp8_model_mem + peak_mem_ratio = bf16_peak_mem / fp8_peak_mem + + print(f"\n[Comparison] Model memory ratio (BF16/FP8): {model_mem_ratio:.2f}x") + print(f"[Comparison] Peak memory ratio (BF16/FP8): {peak_mem_ratio:.2f}x") + + # FP8 should use ~2x less memory + assert model_mem_ratio > 1.8, f"FP8 should use ~2x less memory, got {model_mem_ratio:.2f}x" + + del pipeline_fp8 + torch.cuda.empty_cache() + + +# ============================================================================= +# Attention Backend Comparison Tests +# ============================================================================= + + +class TestFluxAttentionBackend: + """Test VANILLA vs TRTLLM attention backend numerical correctness.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_attention_backend_comparison(self, flux1_checkpoint_exists): + """Test that VANILLA and TRTLLM backends produce similar outputs. + + FLUX uses joint self-attention (same seq_len for Q and KV), so both + VANILLA and TRTLLM backends should work. This test verifies numerical + consistency between them. + """ + # Run VANILLA first, save output, then free before loading TRTLLM + # (two full transformers don't fit in GPU memory simultaneously) + print("\n[Attention Backend Test] Loading baseline transformer (VANILLA)...") + args_baseline = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + attention=AttentionConfig(backend="VANILLA"), + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_baseline = PipelineLoader(args_baseline).load() + transformer_baseline = pipeline_baseline.transformer + + inputs = _get_flux_transformer_inputs(transformer_baseline) + + print("[Attention Backend Test] Running VANILLA transformer forward...") + with torch.no_grad(): + output_baseline = transformer_baseline(**inputs) + output_baseline = _extract_transformer_output(output_baseline).cpu() + + del pipeline_baseline, transformer_baseline + gc.collect() + torch.cuda.empty_cache() + + # Load and run TRTLLM backend + print("[Attention Backend Test] Loading TRTLLM transformer...") + args_trtllm = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + attention=AttentionConfig(backend="TRTLLM"), + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_trtllm = PipelineLoader(args_trtllm).load() + transformer_trtllm = pipeline_trtllm.transformer + + print("[Attention Backend Test] Running TRTLLM transformer forward...") + with torch.no_grad(): + output_trtllm = transformer_trtllm(**inputs) + output_trtllm = _extract_transformer_output(output_trtllm).cpu() + + assert output_baseline.shape == output_trtllm.shape, ( + f"Output shape mismatch: VANILLA={output_baseline.shape}, TRTLLM={output_trtllm.shape}" + ) + + # Check for NaN/Inf + for name, output in [("VANILLA", output_baseline), ("TRTLLM", output_trtllm)]: + assert not torch.isnan(output).any(), f"{name} output contains NaN" + assert not torch.isinf(output).any(), f"{name} output contains Inf" + + # Compare + output_baseline_float = output_baseline.float() + output_trtllm_float = output_trtllm.float() + + max_diff = torch.max(torch.abs(output_trtllm_float - output_baseline_float)).item() + mean_diff = torch.mean(torch.abs(output_trtllm_float - output_baseline_float)).item() + cos_sim = F.cosine_similarity( + output_trtllm_float.flatten(), output_baseline_float.flatten(), dim=0 + ).item() + mse = F.mse_loss(output_trtllm_float, output_baseline_float).item() + + print(f"\n{'=' * 60}") + print("TRTLLM vs VANILLA Comparison") + print(f"{'=' * 60}") + print(f"Max absolute difference: {max_diff:.6f}") + print(f"Mean absolute difference: {mean_diff:.6f}") + print(f"Cosine similarity: {cos_sim:.6f}") + print(f"MSE loss: {mse:.6f}") + print(f"{'=' * 60}") + + assert cos_sim > 0.99, ( + f"TRTLLM should produce similar results to VANILLA: cos_sim={cos_sim:.6f}" + ) + + print(f"\n[PASS] TRTLLM backend matches VANILLA: cos_sim={cos_sim:.6f} (>0.99)") + + del pipeline_trtllm, transformer_trtllm + gc.collect() + torch.cuda.empty_cache() + + +# ============================================================================= +# End-to-End Pipeline Tests (vs HuggingFace Reference) +# ============================================================================= + + +class TestFluxE2E: + """End-to-end pipeline tests: full generation compared to HuggingFace reference.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux1_e2e_vs_hf(self, flux1_checkpoint_exists): + """Full FLUX.1 pipeline (all components) generates image matching HF reference.""" + from diffusers import FluxPipeline as HFFluxPipeline + + # 1. Generate HF reference image + hf_pipe = HFFluxPipeline.from_pretrained( + FLUX1_CHECKPOINT_PATH, torch_dtype=torch.bfloat16 + ).to("cuda") + hf_result = hf_pipe( + prompt="a tiny astronaut hatching from an egg on the moon", + height=256, + width=256, + num_inference_steps=4, + guidance_scale=3.5, + generator=torch.Generator("cuda").manual_seed(42), + ) + hf_image = np.array(hf_result.images[0]) # PIL -> (H, W, 3) uint8 + del hf_pipe + gc.collect() + torch.cuda.empty_cache() + + # 2. Load TRT-LLM pipeline (full, no skip_components) + args = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + pipeline=PipelineConfig(), + ) + pipeline = PipelineLoader(args).load() + + # 3. Generate native image + result = pipeline.forward( + prompt="a tiny astronaut hatching from an egg on the moon", + height=256, + width=256, + num_inference_steps=4, + guidance_scale=3.5, + seed=42, + ) + native_image = result.image.cpu().numpy() # (H, W, 3) uint8 + + # 4. Compute PSNR + mse = ((hf_image.astype(float) - native_image.astype(float)) ** 2).mean() + psnr = 10 * np.log10(255**2 / mse) if mse > 0 else float("inf") + print(f"\n[E2E FLUX.1] PSNR: {psnr:.2f} dB") + + assert psnr > 20.0, f"PSNR too low: {psnr:.2f} dB (expected >20 dB)" + + del pipeline + gc.collect() + torch.cuda.empty_cache() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux2_e2e_vs_hf(self, flux2_checkpoint_exists): + """Full FLUX.2 pipeline (all components) generates image matching HF reference.""" + from diffusers import Flux2Pipeline as HFFlux2Pipeline + + # 1. Generate HF reference image + hf_pipe = HFFlux2Pipeline.from_pretrained( + FLUX2_CHECKPOINT_PATH, torch_dtype=torch.bfloat16 + ).to("cuda") + hf_result = hf_pipe( + prompt="a tiny astronaut hatching from an egg on the moon", + height=256, + width=256, + num_inference_steps=4, + guidance_scale=3.5, + generator=torch.Generator("cuda").manual_seed(42), + ) + hf_image = np.array(hf_result.images[0]) # PIL -> (H, W, 3) uint8 + del hf_pipe + gc.collect() + torch.cuda.empty_cache() + + # 2. Load TRT-LLM pipeline (full, no skip_components) + args = DiffusionArgs( + checkpoint_path=FLUX2_CHECKPOINT_PATH, + device="cuda", + dtype="bfloat16", + pipeline=PipelineConfig(), + ) + pipeline = PipelineLoader(args).load() + + # 3. Generate native image + result = pipeline.forward( + prompt="a tiny astronaut hatching from an egg on the moon", + height=256, + width=256, + num_inference_steps=4, + guidance_scale=3.5, + seed=42, + ) + native_image = result.image.cpu().numpy() # (H, W, 3) uint8 + + # 4. Compute PSNR + mse = ((hf_image.astype(float) - native_image.astype(float)) ** 2).mean() + psnr = 10 * np.log10(255**2 / mse) if mse > 0 else float("inf") + print(f"\n[E2E FLUX.2] PSNR: {psnr:.2f} dB") + + # from HF is expected (~15 dB) compared to FLUX.1 (~32 dB). + assert psnr > 20.0, f"PSNR too low: {psnr:.2f} dB (expected >20 dB)" + + del pipeline + gc.collect() + torch.cuda.empty_cache() + + +# ============================================================================= +# Multi-GPU Parallelism Tests (Ulysses sequence parallelism) +# ============================================================================= + + +def _setup_distributed(rank, world_size, backend="nccl"): + """Initialize distributed process group for multi-GPU tests.""" + os.environ["TLLM_DISABLE_MPI"] = "1" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "12355" + os.environ["RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + + dist.init_process_group(backend=backend, rank=rank, world_size=world_size) + torch.cuda.set_device(rank) + + +def _cleanup_distributed(): + """Clean up distributed process group.""" + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_ulysses_worker(rank, world_size, checkpoint_path, inputs_cpu, return_dict): + """Worker function for Ulysses multi-GPU test. + + Must be module-level for multiprocessing.spawn() pickling. + """ + try: + _setup_distributed(rank, world_size) + + from tensorrt_llm._torch.visual_gen.config import DiffusionArgs, ParallelConfig + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + + # Load pipeline with Ulysses parallelism + args = DiffusionArgs( + checkpoint_path=checkpoint_path, + device=f"cuda:{rank}", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + parallel=ParallelConfig(dit_ulysses_size=world_size), + pipeline={"warmup_steps": 0}, + ) + pipeline = PipelineLoader(args).load() + + # Load inputs on this GPU + inputs = {k: v.to(f"cuda:{rank}") for k, v in inputs_cpu.items()} + + # Run transformer forward + with torch.no_grad(): + output = pipeline.transformer(**inputs) + + sample = _extract_transformer_output(output) + + # Only rank 0 stores the result + if rank == 0: + return_dict["output"] = sample.cpu() + return_dict["shape"] = list(sample.shape) + + del pipeline + torch.cuda.empty_cache() + + except Exception as e: + return_dict[f"error_{rank}"] = str(e) + raise + finally: + _cleanup_distributed() + + +class TestFluxParallelism: + """Ulysses sequence parallelism tests for FLUX (requires 2+ GPUs).""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.device_count() < 2, + reason="Ulysses parallel test requires at least 2 GPUs", + ) + def test_ulysses_2gpu_correctness(self, flux1_checkpoint_exists): + """Test Ulysses (ulysses_size=2) correctness against single-GPU baseline. + + Similar pattern to WAN's test_cfg_2gpu_correctness: + 1. Load single-GPU reference pipeline, run forward + 2. Spawn 2-GPU Ulysses workers, run same forward + 3. Compare outputs (PSNR > 30 dB) + """ + + torch.manual_seed(42) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(42) + + print("\n" + "=" * 80) + print("ULYSSES SEQUENCE PARALLELISM (ulysses_size=2) CORRECTNESS TEST") + print("=" * 80) + + # Load single-GPU reference + print("\n[1/3] Loading single-GPU reference (ulysses_size=1) on GPU 0...") + args_baseline = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda:0", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_baseline = PipelineLoader(args_baseline).load() + + # Create test inputs (seq_len must be divisible by 2 for Ulysses) + print("\n[2/3] Creating test inputs...") + inputs = _get_flux_transformer_inputs( + pipeline_baseline.transformer, device="cuda:0", dtype=torch.bfloat16 + ) + + # Run single-GPU reference + with torch.no_grad(): + ref_output = pipeline_baseline.transformer(**inputs) + ref_sample = _extract_transformer_output(ref_output) + print(f" Reference output shape: {ref_sample.shape}") + print(f" Reference range: [{ref_sample.min():.4f}, {ref_sample.max():.4f}]") + + # Store inputs on CPU for workers + inputs_cpu = {k: v.cpu() for k, v in inputs.items()} + + # Cleanup baseline + del pipeline_baseline + gc.collect() + torch.cuda.empty_cache() + torch._dynamo.reset() + + # Run Ulysses parallel (2 GPUs) + print("\n[3/3] Running Ulysses (ulysses_size=2) across 2 GPUs...") + manager = mp.Manager() + return_dict = manager.dict() + + mp.spawn( + _run_ulysses_worker, + args=(2, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict), + nprocs=2, + join=True, + ) + + # Check for errors + for i in range(2): + assert f"error_{i}" not in return_dict, ( + f"Rank {i} failed: {return_dict.get(f'error_{i}')}" + ) + + ulysses_sample = return_dict["output"].to("cuda:0") + print(f" Ulysses output shape: {ulysses_sample.shape}") + print(f" Ulysses range: [{ulysses_sample.min():.4f}, {ulysses_sample.max():.4f}]") + + # Compare outputs + assert ref_sample.shape == ulysses_sample.shape, ( + f"Shape mismatch: ref={ref_sample.shape}, ulysses={ulysses_sample.shape}" + ) + + mse = ((ref_sample.float() - ulysses_sample.float()) ** 2).mean().item() + ref_range = (ref_sample.max() - ref_sample.min()).float().item() + psnr = 10 * np.log10(ref_range**2 / mse) if mse > 0 else float("inf") + + print(f"\n MSE: {mse:.6e}") + print(f" PSNR: {psnr:.2f} dB") + + # Ulysses should be nearly identical (only BF16 rounding from all-to-all) + assert psnr > 30.0, f"PSNR too low: {psnr:.2f} dB (expected >30 dB)" + + del ref_sample, ulysses_sample + gc.collect() + torch.cuda.empty_cache() + + +def _run_all_optimizations_worker(rank, world_size, checkpoint_path, inputs_cpu, return_dict): + """Worker for combined optimizations test (FP8 + TeaCache + TRTLLM + Ulysses). + + Must be module-level for multiprocessing.spawn() pickling. + """ + try: + _setup_distributed(rank, world_size) + + from tensorrt_llm._torch.visual_gen.config import ( + AttentionConfig, + DiffusionArgs, + ParallelConfig, + TeaCacheConfig, + ) + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.quantization.mode import QuantAlgo + + # Load pipeline with ALL optimizations + args = DiffusionArgs( + checkpoint_path=checkpoint_path, + device=f"cuda:{rank}", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + quant_config={"quant_algo": "FP8", "dynamic": True}, + teacache=TeaCacheConfig( + enable_teacache=True, + teacache_thresh=0.2, + use_ret_steps=True, + ), + attention=AttentionConfig(backend="TRTLLM"), + parallel=ParallelConfig(dit_ulysses_size=world_size), + pipeline={"warmup_steps": 0}, + ) + pipeline = PipelineLoader(args).load() + transformer = pipeline.transformer.eval() + + # Verify all optimizations are enabled + assert pipeline.model_config.parallel.dit_ulysses_size == world_size, ( + "Ulysses parallel not enabled" + ) + assert transformer.model_config.quant_config.quant_algo == QuantAlgo.FP8, "FP8 not enabled" + assert hasattr(pipeline, "cache_backend"), "TeaCache not enabled" + assert transformer.transformer_blocks[0].attn.attn_backend == "TRTLLM", "TRTLLM not enabled" + + if rank == 0: + print(f" All optimizations verified on rank {rank}:") + print(f" - FP8: {transformer.model_config.quant_config.quant_algo}") + print(" - TeaCache: enabled") + print(f" - TRTLLM attention: {transformer.transformer_blocks[0].attn.attn_backend}") + print(f" - Ulysses: ulysses_size={world_size}") + + # Initialize TeaCache for single-step inference + if hasattr(pipeline, "cache_backend") and pipeline.cache_backend: + pipeline.cache_backend.refresh(num_inference_steps=1) + + # Load inputs on this GPU + inputs = {k: v.to(f"cuda:{rank}") for k, v in inputs_cpu.items()} + + # Run transformer forward (return_dict=False for TeaCache compatibility) + inputs["return_dict"] = False + with torch.no_grad(): + output = transformer(**inputs)[0] + + # Validate output + assert not torch.isnan(output).any(), f"Rank {rank}: Output contains NaN" + assert not torch.isinf(output).any(), f"Rank {rank}: Output contains Inf" + + if rank == 0: + return_dict["output"] = output.cpu() + + del pipeline, transformer + torch.cuda.empty_cache() + + except Exception as e: + return_dict[f"error_{rank}"] = str(e) + raise + finally: + _cleanup_distributed() + + +class TestFluxCombinedOptimizations: + """Test all optimizations combined: FP8 + TeaCache + TRTLLM attention + Ulysses.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.device_count() < 2, + reason="Combined optimization test requires at least 2 GPUs", + ) + def test_all_optimizations_combined(self, flux1_checkpoint_exists): + """Test FP8 + TeaCache + TRTLLM attention + Ulysses=2 combined correctness. + + Validates that all optimizations work together correctly. + Compares against a BF16 single-GPU baseline with relaxed thresholds + since multiple optimizations compound numerical differences. + """ + torch.manual_seed(42) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(42) + + print("\n" + "=" * 80) + print("ALL OPTIMIZATIONS COMBINED TEST") + print("FP8 + TeaCache + TRTLLM Attention + Ulysses (ulysses_size=2)") + print("=" * 80) + + # Load baseline on GPU 0 (no optimizations) + print("\n[1/3] Loading baseline on GPU 0 (BF16, no optimizations)...") + args_baseline = DiffusionArgs( + checkpoint_path=FLUX1_CHECKPOINT_PATH, + device="cuda:0", + dtype="bfloat16", + skip_components=SKIP_COMPONENTS, + pipeline=PIPELINE_NO_WARMUP, + ) + pipeline_baseline = PipelineLoader(args_baseline).load() + + # Reset torch compile state + torch._dynamo.reset() + + # Create test inputs + print("\n[2/3] Creating test inputs...") + inputs = _get_flux_transformer_inputs( + pipeline_baseline.transformer, device="cuda:0", dtype=torch.bfloat16 + ) + + # Run baseline + with torch.no_grad(): + ref_output = pipeline_baseline.transformer(**inputs) + ref_sample = _extract_transformer_output(ref_output) + print(f" Baseline output shape: {ref_sample.shape}") + print(f" Baseline range: [{ref_sample.min():.4f}, {ref_sample.max():.4f}]") + + # Store inputs on CPU for workers + inputs_cpu = {k: v.cpu() for k, v in inputs.items()} + + # Cleanup baseline + del pipeline_baseline + gc.collect() + torch.cuda.empty_cache() + + # Run with ALL optimizations in distributed processes + print("\n[3/3] Running with ALL optimizations (FP8 + TeaCache + TRTLLM + Ulysses=2)...") + manager = mp.Manager() + return_dict = manager.dict() + + mp.spawn( + _run_all_optimizations_worker, + args=(2, FLUX1_CHECKPOINT_PATH, inputs_cpu, return_dict), + nprocs=2, + join=True, + ) + + # Check for errors + for i in range(2): + assert f"error_{i}" not in return_dict, ( + f"Rank {i} failed: {return_dict.get(f'error_{i}')}" + ) + + combined_sample = return_dict["output"].to("cuda:0") + + # Compare outputs with relaxed thresholds (multiple optimizations compound errors) + print("\n[Comparison] Combined Optimizations vs Baseline:") + ref_float = ref_sample.float() + combined_float = combined_sample.float() + + cos_sim = F.cosine_similarity(combined_float.flatten(), ref_float.flatten(), dim=0).item() + + max_diff = torch.max(torch.abs(combined_float - ref_float)).item() + mean_diff = torch.mean(torch.abs(combined_float - ref_float)).item() + + print(f" Cosine similarity: {cos_sim:.6f}") + print(f" Max absolute difference: {max_diff:.6f}") + print(f" Mean absolute difference: {mean_diff:.6f}") + print(f" Combined range: [{combined_float.min():.4f}, {combined_float.max():.4f}]") + print(f" Baseline range: [{ref_float.min():.4f}, {ref_float.max():.4f}]") + + # Relaxed threshold: cos_sim > 0.90 (compounded from FP8 + Ulysses + TeaCache) + assert cos_sim > 0.90, ( + f"Combined optimization cosine similarity {cos_sim:.6f} below threshold 0.90. " + f"This suggests an issue with optimization interactions." + ) + + print(f"\n[PASS] All optimizations validated! cos_sim={cos_sim:.6f}") + print("=" * 80) + + del ref_sample, combined_sample + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/visual_gen/test_flux_transformer.py b/tests/unittest/_torch/visual_gen/test_flux_transformer.py new file mode 100644 index 000000000000..170409a1e3f4 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_flux_transformer.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for FLUX transformer models. + +Tests cover: +- Model structure and instantiation +- Forward pass sanity checks +- Numerical correctness vs HuggingFace +""" + +import unittest +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.visual_gen.config import AttentionConfig, DiffusionModelConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig + +# FLUX.1 dev config (12B params) +FLUX1_CONFIG = { + "attention_head_dim": 128, + "guidance_embeds": True, + "in_channels": 64, + "joint_attention_dim": 4096, + "num_attention_heads": 24, + "num_layers": 19, + "num_single_layers": 38, + "patch_size": 1, + "pooled_projection_dim": 768, + "torch_dtype": "bfloat16", + "axes_dim": [16, 56, 56], +} + +# FLUX.2 dev config (35B params) +FLUX2_CONFIG = { + "attention_head_dim": 128, + "guidance_embeds": False, + "in_channels": 128, + "joint_attention_dim": 4096, + "num_attention_heads": 48, + "num_layers": 8, + "num_single_layers": 48, + "patch_size": 1, + "pooled_projection_dim": 768, + "torch_dtype": "bfloat16", + "axes_dim": [32, 32, 32, 32], +} + + +def reduce_flux_config(mem_for_full_model: int, config_dict: dict): + """Reduce model size if insufficient GPU memory.""" + _, total_mem = torch.cuda.mem_get_info() + if total_mem < mem_for_full_model: + model_fraction = total_mem / mem_for_full_model + num_layers = max(1, int(config_dict["num_layers"] * model_fraction)) + num_single_layers = max(1, int(config_dict["num_single_layers"] * model_fraction)) + config_dict["num_layers"] = min(num_layers, 2) + config_dict["num_single_layers"] = min(num_single_layers, 4) + + +class TestFluxTransformer(unittest.TestCase): + """Unit tests for FLUX transformer models.""" + + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + def _create_model_config( + self, config_dict: dict, backend: str = "VANILLA" + ) -> DiffusionModelConfig: + """Create DiffusionModelConfig from config dict.""" + pretrained_config = SimpleNamespace(**config_dict) + return DiffusionModelConfig( + pretrained_config=pretrained_config, + quant_config=QuantConfig(), + mapping=Mapping(), + attention=AttentionConfig(backend=backend), + skip_create_weights_in_init=False, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux1_model_structure(self): + """Test FLUX.1 model can be instantiated with correct structure.""" + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux import ( + FluxTransformer2DModel, + ) + + config = deepcopy(FLUX1_CONFIG) + config["num_layers"] = 1 + config["num_single_layers"] = 1 + + model_config = self._create_model_config(config) + model = FluxTransformer2DModel(model_config).to(self.DEVICE) + + # Check model structure + self.assertTrue(hasattr(model, "transformer_blocks")) + self.assertTrue(hasattr(model, "single_transformer_blocks")) + self.assertEqual(len(model.transformer_blocks), 1) + self.assertEqual(len(model.single_transformer_blocks), 1) + + # Check key components + self.assertTrue(hasattr(model, "x_embedder")) + self.assertTrue(hasattr(model, "context_embedder")) + self.assertTrue(hasattr(model, "time_text_embed")) + self.assertTrue(hasattr(model, "pos_embed")) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux2_model_structure(self): + """Test FLUX.2 model can be instantiated with correct structure.""" + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux2 import ( + Flux2Transformer2DModel, + ) + + config = deepcopy(FLUX2_CONFIG) + config["num_layers"] = 1 + config["num_single_layers"] = 1 + + model_config = self._create_model_config(config) + model = Flux2Transformer2DModel(model_config).to(self.DEVICE) + + self.assertTrue(hasattr(model, "transformer_blocks")) + self.assertTrue(hasattr(model, "single_transformer_blocks")) + self.assertEqual(len(model.transformer_blocks), 1) + self.assertEqual(len(model.single_transformer_blocks), 1) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux1_forward_sanity(self): + """Test FLUX.1 forward pass produces valid output.""" + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux import ( + FluxTransformer2DModel, + ) + + config = deepcopy(FLUX1_CONFIG) + config["num_layers"] = 1 + config["num_single_layers"] = 1 + + model_config = self._create_model_config(config) + model = FluxTransformer2DModel(model_config).to(self.DEVICE, dtype=torch.bfloat16).eval() + + batch_size = 1 + height, width = 64, 64 + seq_len = (height // 2) * (width // 2) + text_seq_len = 128 + + hidden_states = torch.randn( + batch_size, seq_len, config["in_channels"], device=self.DEVICE, dtype=torch.bfloat16 + ) + encoder_hidden_states = torch.randn( + batch_size, + text_seq_len, + config["joint_attention_dim"], + device=self.DEVICE, + dtype=torch.bfloat16, + ) + pooled_projections = torch.randn( + batch_size, config["pooled_projection_dim"], device=self.DEVICE, dtype=torch.bfloat16 + ) + timestep = torch.tensor([500], device=self.DEVICE, dtype=torch.bfloat16) + guidance = torch.tensor([3.5], device=self.DEVICE, dtype=torch.bfloat16) + img_ids = torch.zeros(batch_size, seq_len, 3, device=self.DEVICE) + txt_ids = torch.zeros(batch_size, text_seq_len, 3, device=self.DEVICE) + + with torch.no_grad(): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + guidance=guidance, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + # Model returns {"sample": tensor} + if isinstance(output, dict): + output = output["sample"] + + # Note: With random weights, NaN can occur. For unit tests, we only check shape. + # Full numerical correctness is tested in TestFluxHuggingFaceComparison. + self.assertEqual(output.shape, hidden_states.shape) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux2_forward_sanity(self): + """Test FLUX.2 forward pass produces valid output.""" + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux2 import ( + Flux2Transformer2DModel, + ) + + config = deepcopy(FLUX2_CONFIG) + config["num_layers"] = 1 + config["num_single_layers"] = 1 + + model_config = self._create_model_config(config) + model = Flux2Transformer2DModel(model_config).to(self.DEVICE, dtype=torch.bfloat16).eval() + + batch_size = 1 + height, width = 32, 32 + seq_len = (height // 2) * (width // 2) + text_seq_len = 128 + + hidden_states = torch.randn( + batch_size, seq_len, config["in_channels"], device=self.DEVICE, dtype=torch.bfloat16 + ) + encoder_hidden_states = torch.randn( + batch_size, + text_seq_len, + config["joint_attention_dim"], + device=self.DEVICE, + dtype=torch.bfloat16, + ) + timestep = torch.tensor([500], device=self.DEVICE, dtype=torch.bfloat16) + guidance = torch.tensor([3.5], device=self.DEVICE, dtype=torch.bfloat16) + # FLUX.2 uses 4-axis RoPE + img_ids = torch.zeros(seq_len, 4, device=self.DEVICE) + txt_ids = torch.zeros(text_seq_len, 4, device=self.DEVICE) + + with torch.no_grad(): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + guidance=guidance, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + # Model returns {"sample": tensor} + if isinstance(output, dict): + output = output["sample"] + + # Note: With random weights, NaN can occur. For unit tests, we only check shape. + # Full numerical correctness is tested in TestFluxHuggingFaceComparison. + self.assertEqual(output.shape, hidden_states.shape) + + +class TestFluxHuggingFaceComparison(unittest.TestCase): + """Test FLUX models match HuggingFace reference implementation.""" + + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + def _create_model_config( + self, config_dict: dict, backend: str = "VANILLA" + ) -> DiffusionModelConfig: + """Create DiffusionModelConfig from config dict.""" + pretrained_config = SimpleNamespace(**config_dict) + return DiffusionModelConfig( + pretrained_config=pretrained_config, + quant_config=QuantConfig(), + mapping=Mapping(), + attention=AttentionConfig(backend=backend), + skip_create_weights_in_init=False, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux1_allclose_to_hf(self): + """Test TRT-LLM FLUX.1 transformer matches HuggingFace output.""" + try: + from diffusers import FluxTransformer2DModel as HFFluxTransformer2DModel + except ImportError: + self.skipTest("diffusers not installed") + + from tensorrt_llm._torch.visual_gen.models.flux.transformer_flux import ( + FluxTransformer2DModel, + ) + + torch.manual_seed(42) + + config = deepcopy(FLUX1_CONFIG) + config["num_layers"] = 1 + config["num_single_layers"] = 2 + + dtype = torch.bfloat16 + + # Create HuggingFace model with random weights + hf_model = ( + HFFluxTransformer2DModel( + patch_size=config["patch_size"], + in_channels=config["in_channels"], + num_layers=config["num_layers"], + num_single_layers=config["num_single_layers"], + attention_head_dim=config["attention_head_dim"], + num_attention_heads=config["num_attention_heads"], + joint_attention_dim=config["joint_attention_dim"], + pooled_projection_dim=config["pooled_projection_dim"], + guidance_embeds=config["guidance_embeds"], + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + # Create TRT-LLM model + model_config = self._create_model_config(config) + trtllm_model = FluxTransformer2DModel(model_config).to(self.DEVICE, dtype=dtype).eval() + + # Copy weights from HF to TRT-LLM + hf_state_dict = hf_model.state_dict() + trtllm_model.load_weights(hf_state_dict) + + # Create test inputs + batch_size = 1 + height, width = 32, 32 + seq_len = (height // 2) * (width // 2) + text_seq_len = 64 + + generator = torch.Generator(device=self.DEVICE).manual_seed(42) + hidden_states = torch.randn( + batch_size, + seq_len, + config["in_channels"], + generator=generator, + device=self.DEVICE, + dtype=dtype, + ) + encoder_hidden_states = torch.randn( + batch_size, + text_seq_len, + config["joint_attention_dim"], + generator=generator, + device=self.DEVICE, + dtype=dtype, + ) + pooled_projections = torch.randn( + batch_size, + config["pooled_projection_dim"], + generator=generator, + device=self.DEVICE, + dtype=dtype, + ) + timestep = torch.tensor([500.0], device=self.DEVICE, dtype=dtype) + guidance = torch.tensor([3.5], device=self.DEVICE, dtype=dtype) + img_ids = torch.zeros(batch_size, seq_len, 3, device=self.DEVICE) + txt_ids = torch.zeros(batch_size, text_seq_len, 3, device=self.DEVICE) + + # Run both models + with ( + torch.no_grad(), + torch.backends.cuda.sdp_kernel( + enable_flash=False, enable_math=True, enable_mem_efficient=False + ), + ): + hf_output = hf_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep / 1000, + guidance=guidance, + img_ids=img_ids, + txt_ids=txt_ids, + return_dict=False, + )[0] + + trtllm_output = trtllm_model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + guidance=guidance, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + # Model returns {"sample": tensor} + if isinstance(trtllm_output, dict): + trtllm_output = trtllm_output["sample"] + + # Compare outputs + hf_output = hf_output.float() + trtllm_output = trtllm_output.float() + + cos_sim = F.cosine_similarity( + hf_output.flatten().unsqueeze(0), trtllm_output.flatten().unsqueeze(0) + ).item() + + max_diff = (hf_output - trtllm_output).abs().max().item() + + print("\n[FLUX.1 HF Comparison]") + print(f" Cosine similarity: {cos_sim:.6f}") + print(f" Max diff: {max_diff:.6f}") + + self.assertGreater(cos_sim, 0.99, f"Cosine similarity too low: {cos_sim}") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])