[TRTLLM-10616][feat] Add FLUX.1 and FLUX.2 text-to-image pipeline support

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
 #############################################

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)

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")

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

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)