From 07cbe8c7f2e3d6387ce3a959e7fb4b0343bd1334 Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Sun, 20 Sep 2026 18:01:46 +0200 Subject: [PATCH 01/14] add translator and FP16 --- image_gen/flux2/model_export/export_flux2.py | 28 ++++- image_gen/flux2/model_export/translator.py | 85 ++++++++++++++++ image_gen/flux2/model_export/verify_flux2.py | 102 +++++++++++++++++-- 3 files changed, 203 insertions(+), 12 deletions(-) create mode 100644 image_gen/flux2/model_export/translator.py diff --git a/image_gen/flux2/model_export/export_flux2.py b/image_gen/flux2/model_export/export_flux2.py index 3ec28af..3163192 100644 --- a/image_gen/flux2/model_export/export_flux2.py +++ b/image_gen/flux2/model_export/export_flux2.py @@ -33,13 +33,14 @@ from pathlib import Path import torch -from diffusers import Flux2KleinPipeline as FluxPipeline from huggingface_hub import snapshot_download +from translator import add_translator_arguments, load_translator_encoder EXTERNAL_DATA_NAME = "model.onnx_data" DEFAULT_MODEL_NAME = "black-forest-labs/FLUX.2-klein-4b" DEFAULT_MODEL_OPSETS = { "text_encoder": 25, + "text_encoder_translator": 25, "transformer": 25, "vae_encoder": 22, "vae_decoder": 22, @@ -315,6 +316,8 @@ def copy_scheduler_and_tokenizer(model_snapshot: Path, output_path: Path): for name in ("scheduler", "tokenizer"): src = model_snapshot / name dst = output_path / name + if src.resolve() == dst.resolve(): + continue if src.exists(): if dst.exists(): shutil.rmtree(dst) @@ -428,7 +431,7 @@ def main(): nargs="+", default=None, metavar="NAME", - help="Model(s) to export: one or more of all, text_encoder, transformer, vae_encoder, vae_decoder (default: all)", + help="Models to export: all, text_encoder, text_encoder_translator, transformer, vae_encoder, vae_decoder. Translator is opt-in.", ) ap.add_argument( "--opset", @@ -486,6 +489,7 @@ def main(): default=[], help="Additional argument to pass to tensorrt_rtx. May be specified more than once.", ) + add_translator_arguments(ap) args = ap.parse_args() model_snapshot = resolve_model_snapshot(args.model_name, args.local_files_only) @@ -495,7 +499,7 @@ def main(): seq_len = 512 # static text sequence length for FLUX.2-klein # Parse --model: default all; otherwise one or more of all, text_encoder, transformer, … - model_choices = ["all"] + list(MODELS) + model_choices = ["all", "text_encoder_translator"] + list(MODELS) raw = args.model if args.model is not None else ["all"] to_export: list[str] = [] for m in raw: @@ -518,12 +522,26 @@ def main(): f"Loading pipeline from {args.model_name} ({model_snapshot}) " f"(dtype={dtype}, image_size={image_size}, io_precision={args.io_precision})..." ) - pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype).to(device) + pipe = None + if any(name != "text_encoder_translator" for name in to_export): + from diffusers import Flux2KleinPipeline as FluxPipeline + pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype).to(device) exported_model_paths: list[Path] = [] for name in to_export: model_opset = args.opset if args.opset is not None else DEFAULT_MODEL_OPSETS[name] print(f"Exporting {name}...") + if name == "text_encoder_translator": + encoder = load_translator_encoder(args, device) + ids = torch.zeros((1, seq_len), dtype=torch.int64, device=device) + mask = torch.ones_like(ids) + path = model_output_path(output_path, name, args.transformer_precision) + _onnx_export(encoder, (ids, mask), path, + ["input_ids", "attention_mask"], ["prompt_embeds"], model_opset) + exported_model_paths.append(path) + del encoder + torch.cuda.empty_cache() + continue fn = EXPORTERS[name] if name == "text_encoder": fn(pipe, output_path, device, model_opset, @@ -539,7 +557,7 @@ def main(): fn(pipe, output_path, device, model_opset, io_dtype=io_dtype) exported_model_paths.append(model_output_path(output_path, name, args.transformer_precision)) - if export_all: + if export_all or "text_encoder_translator" in to_export: copy_scheduler_and_tokenizer(model_snapshot, output_path) write_model_index(output_path) diff --git a/image_gen/flux2/model_export/translator.py b/image_gen/flux2/model_export/translator.py new file mode 100644 index 0000000..e228a1d --- /dev/null +++ b/image_gen/flux2/model_export/translator.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deployment architecture for gafr1/translator; shared by export and verification.""" +from pathlib import Path + +import torch +from torch import nn +from torch.nn import functional as F + +DEFAULT_TRANSLATOR_REPO = "gafr1/translator" +DEFAULT_TRANSLATOR_FILE = "mlp_d3072_b4_e2.pt" + + +class ResMLP(nn.Module): + def __init__(self, width, expansion): + super().__init__() + self.norm = nn.LayerNorm(width) + self.fc1 = nn.Linear(width, width * expansion) + self.fc2 = nn.Linear(width * expansion, width) + + def forward(self, x): + return x + self.fc2(F.gelu(self.fc1(self.norm(x)))) + + +class MLPTranslator(nn.Module): + def __init__(self, in_dim, width, blocks, expansion): + super().__init__() + self.in_proj = nn.Linear(in_dim, width) + self.blocks = nn.ModuleList([ResMLP(width, expansion) for _ in range(blocks)]) + self.norm = nn.LayerNorm(width) + self.out = nn.Linear(width, 7680) + + def forward(self, x): + x = self.in_proj(x) + for block in self.blocks: + x = block(x) + return self.out(self.norm(x)) + + +class TranslatorEncoder(nn.Module): + def __init__(self, qwen, translator, layers): + super().__init__() + self.qwen = qwen + self.translator = translator + self.layers = layers + + def forward(self, input_ids, attention_mask): + output = self.qwen(input_ids=input_ids, attention_mask=attention_mask, + output_hidden_states=True, use_cache=False) + features = torch.cat([output.hidden_states[i] for i in self.layers], dim=-1) + return self.translator(features).float() + + +def add_translator_arguments(parser): + parser.add_argument("--qwen-model", default="Qwen/Qwen3-0.6B", help="HF repository or local Qwen3-0.6B snapshot") + parser.add_argument("--translator-checkpoint", type=Path, + help="Local checkpoint; takes precedence over HF download") + parser.add_argument("--translator-repo", default=DEFAULT_TRANSLATOR_REPO) + parser.add_argument("--translator-file", default=DEFAULT_TRANSLATOR_FILE) + parser.add_argument("--translator-revision", default="main", + help="HF revision; pin a commit for reproducible exports") + + +def load_translator_encoder(args, device): + from huggingface_hub import hf_hub_download + from transformers import AutoModel + + checkpoint = args.translator_checkpoint + if checkpoint is None: + checkpoint = Path(hf_hub_download( + args.translator_repo, args.translator_file, revision=args.translator_revision, + local_files_only=args.local_files_only)) + # Checkpoints contain only state tensors plus plain config/scalar dictionaries. + data = torch.load(checkpoint, map_location="cpu", weights_only=True) + cfg = data["cfg"] + if cfg["kind"] != "mlp" or cfg["layers"] != "q4" or data["in_dim"] != 4096: + raise ValueError("Expected a Qwen3-0.6B q4 MLP translator with input dimension 4096") + net = MLPTranslator(data["in_dim"], cfg["d"], cfg["nblocks"], cfg["expansion"]) + net.load_state_dict(data["model"], strict=True) + qwen = AutoModel.from_pretrained( + args.qwen_model, torch_dtype=torch.bfloat16, local_files_only=args.local_files_only, + attn_implementation="eager") + if qwen.config.hidden_size != 1024 or qwen.config.num_hidden_layers != 28: + raise ValueError("Translator requires Qwen3-0.6B (28 layers, width 1024)") + return TranslatorEncoder(qwen, net, (7, 14, 21, 28)).eval().to(device=device, dtype=torch.bfloat16) diff --git a/image_gen/flux2/model_export/verify_flux2.py b/image_gen/flux2/model_export/verify_flux2.py index 9dd004d..0607910 100644 --- a/image_gen/flux2/model_export/verify_flux2.py +++ b/image_gen/flux2/model_export/verify_flux2.py @@ -26,9 +26,9 @@ import onnx import torch import torch.nn as nn -from diffusers import Flux2KleinPipeline as FluxPipeline from huggingface_hub import snapshot_download from PIL import Image as PILImage +from translator import add_translator_arguments, load_translator_encoder EP_NAME = "nv_tensorrt_rtx" _REGISTERED_EP_NAME: str | None = None @@ -281,7 +281,7 @@ def transformer_model_path(onnx_dir: Path, precision: str) -> Path: def validate_onnx_dir(onnx_dir: Path, precision: str, provider: str, args) -> None: print(f"\n[Validate] {onnx_dir}") model_paths = [ - onnx_dir / "text_encoder" / "model.onnx", + onnx_dir / ("text_encoder_translator" if args.encoder == "translator" else "text_encoder") / "model.onnx", transformer_model_path(onnx_dir, precision), onnx_dir / "vae_encoder" / "model.onnx", onnx_dir / "vae_decoder" / "model.onnx", @@ -310,11 +310,13 @@ def validate_onnx_dir(onnx_dir: Path, precision: str, provider: str, args) -> No raise RuntimeError(f"No model.onnx files found under {onnx_dir}") -def run_pytorch_pipeline(pipe, prompt: str, seed: int, image_size: int, num_steps: int) -> PILImage.Image: +def run_pytorch_pipeline(pipe, prompt: str, seed: int, image_size: int, num_steps: int, + prompt_embeds=None) -> PILImage.Image: print("\n[PyTorch] Running pipeline...") generator = torch.Generator(device="cuda").manual_seed(seed) result = pipe( - prompt=prompt, + prompt=prompt if prompt_embeds is None else None, + prompt_embeds=prompt_embeds, height=image_size, width=image_size, num_inference_steps=num_steps, @@ -339,7 +341,7 @@ def run_ort_pipeline( prompt_embeds = None original_text_encoder = None - te_path = onnx_dir / "text_encoder" / "model.onnx" + te_path = onnx_dir / ("text_encoder_translator" if args.encoder == "translator" else "text_encoder") / "model.onnx" if te_path.exists(): original_text_encoder = pipe.text_encoder original_text_dtype = original_text_encoder.dtype @@ -439,12 +441,46 @@ def compare_images(ref: PILImage.Image, ort: PILImage.Image, output_dir: Path) - print(f"[Comparison] Side-by-side saved -> {cmp_path}") +def verify_translator_encoder(onnx_dir, args): + from transformers import AutoTokenizer + + if args.encoder != "translator": + raise ValueError("--encoder-only requires --encoder translator") + tokenizer = AutoTokenizer.from_pretrained(str(onnx_dir / "tokenizer")) + reference = load_translator_encoder(args, "cuda") + runner = OrtSessionRunner.from_path(onnx_dir / "text_encoder_translator/model.onnx", args.provider, args) + _assert_fp32_float_io(runner, onnx_dir / "text_encoder_translator/model.onnx") + if set(runner.input_names) != {"input_ids", "attention_mask"} or runner.output_names != ["prompt_embeds"]: + raise RuntimeError("Translator ONNX graph IO names do not match the C++ contract") + prompts = [args.prompt, "", "A café in Zürich, 日本の庭", "a detailed landscape " * 600] + for index, prompt in enumerate(prompts): + text = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], tokenize=False, + add_generation_prompt=True, enable_thinking=False) + expected_text = "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + if text != expected_text: + raise RuntimeError("Tokenizer chat template differs from the C++ runtime") + tokens = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=512).to("cuda") + with torch.no_grad(): + expected = reference(tokens["input_ids"], tokens["attention_mask"]).float().cpu().numpy() + actual = runner.run({key: _to_numpy(tokens[key], np.int64) for key in ("input_ids", "attention_mask")})[0] + if actual.shape != (1, 512, 7680) or actual.dtype != np.float32 or not np.isfinite(actual).all(): + raise RuntimeError("Translator output must be finite fp32 [1,512,7680]") + cosine = float(np.mean(np.sum(expected * actual, axis=-1) / np.maximum( + np.linalg.norm(expected, axis=-1) * np.linalg.norm(actual, axis=-1), 1e-12))) + nrmse = float(np.sqrt(np.mean((expected - actual) ** 2)) / max(float(np.sqrt(np.mean(expected ** 2))), 1e-12)) + print(f"[Translator prompt {index}] cosine={cosine:.6f}, NRMSE={nrmse:.6f}", flush=True) + if not np.isfinite(cosine + nrmse) or cosine < args.embedding_cosine_min or nrmse > args.embedding_nrmse_max: + raise RuntimeError("Translator embedding parity failed") + + def main(): p = argparse.ArgumentParser(description="Verify FLUX.2-klein ONNX graphs with PyTorch and ONNX Runtime.") p.add_argument("--model_name", type=str, default=DEFAULT_MODEL_NAME, help="Hugging Face model repo ID, or a local snapshot directory") p.add_argument("--local_files_only", action="store_true", help="Resolve --model_name from the local Hugging Face cache only") p.add_argument("--onnx_dir", type=str, default="./flux2_klein_onnx") - p.add_argument("--precision", choices=["bf16", "fp8", "nvfp4"], default="bf16", help="Transformer precision to verify from transformer_.") + p.add_argument("--precision", choices=["bf16", "fp16", "fp8", "nvfp4"], default="bf16", + help="Transformer precision to verify from transformer_.") p.add_argument("--output_dir", type=str, default=None) p.add_argument("--prompt", type=str, default=DEFAULT_PROMPT) p.add_argument("--seed", type=int, default=42) @@ -456,6 +492,12 @@ def main(): p.add_argument("--ep_dll_dir", default=None, help="Dir with the plugin EP's bundled runtime DLLs") p.add_argument("--trt_bin", default=None, help="TensorRT-RTX bin dir, used as a DLL search path") p.add_argument("--validate_only", action="store_true", help="Only run ONNX Runtime validation; do not generate images") + p.add_argument("--encoder", choices=["4b", "translator"], default="4b") + p.add_argument("--encoder-only", action="store_true", + help="Compare fused translator PyTorch and ONNX embeddings without loading the diffusion pipeline") + p.add_argument("--embedding-cosine-min", type=float, default=0.99) + p.add_argument("--embedding-nrmse-max", type=float, default=0.10) + add_translator_arguments(p) args = p.parse_args() configure_stdio() @@ -463,6 +505,12 @@ def main(): output_dir = Path(args.output_dir).resolve() if args.output_dir else onnx_dir output_dir.mkdir(parents=True, exist_ok=True) + if args.encoder_only: + verify_translator_encoder(onnx_dir, args) + return + + if args.encoder == "translator" and not (onnx_dir / "text_encoder_translator/model.onnx").is_file(): + raise FileNotFoundError("Selected translator ONNX is missing") if args.validate_only: validate_onnx_dir(onnx_dir, args.precision, args.provider, args) return @@ -472,15 +520,55 @@ def main(): device = "cuda" model_snapshot = resolve_model_snapshot(args.model_name, args.local_files_only) + from diffusers import Flux2KleinPipeline as FluxPipeline print(f"Loading pipeline from {args.model_name} ({model_snapshot})...") - pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype).to(device) + pipe = FluxPipeline.from_pretrained(str(model_snapshot), torch_dtype=dtype) + if args.encoder == "translator": + # Avoid loading the unused 4B encoder on the GPU. + pipe.transformer.to(device) + pipe.vae.to(device) + else: + pipe.to(device) + reference_embeds = None + if args.encoder == "translator": + # Use the same fused student/translator for numerical parity, not the 4B teacher. + pipe.text_encoder.to("cpu") + torch.cuda.empty_cache() + reference = load_translator_encoder(args, device) + text = pipe.tokenizer.apply_chat_template( + [{"role": "user", "content": args.prompt}], tokenize=False, + add_generation_prompt=True, enable_thinking=False) + tokens = pipe.tokenizer(text, return_tensors="pt", padding="max_length", + truncation=True, max_length=512).to(device) + with torch.no_grad(): + reference_embeds = reference(tokens["input_ids"], tokens["attention_mask"]).detach() + del reference + torch.cuda.empty_cache() + path = onnx_dir / "text_encoder_translator/model.onnx" + runner = OrtSessionRunner.from_path(path, args.provider, args) + expected = reference_embeds.float().cpu().numpy() + actual = runner.run({k: _to_numpy(tokens[k], np.int64) + for k in ("input_ids", "attention_mask")})[0] + if actual.shape != (1, 512, 7680) or actual.dtype != np.float32 or not np.isfinite(actual).all(): + raise RuntimeError("Translator output must be finite fp32 [1,512,7680]") + cosine = float(np.mean(np.sum(expected * actual, axis=-1) / + np.maximum(np.linalg.norm(expected, axis=-1) * np.linalg.norm(actual, axis=-1), 1e-12))) + nrmse = float(np.sqrt(np.mean((expected - actual) ** 2)) / max(float(np.sqrt(np.mean(expected ** 2))), 1e-12)) + print(f"[Translator embeddings] cosine={cosine:.6f}, NRMSE={nrmse:.6f}") + if cosine < args.embedding_cosine_min or nrmse > args.embedding_nrmse_max: + raise RuntimeError("Translator embedding parity failed") + del runner + reference_embeds = reference_embeds.to(dtype=dtype) + # Precomputed embeddings bypass the teacher; keep pipeline device discovery on CUDA. + pipe.text_encoder = _DeviceAnchor(device, dtype) pytorch_image = run_pytorch_pipeline( pipe, prompt=args.prompt, seed=args.seed, image_size=args.image_size, num_steps=args.num_steps, + prompt_embeds=reference_embeds, ) pt_path = output_dir / "pytorch_output.png" pytorch_image.save(str(pt_path)) From eb74b700fdac3170c9bfdbd70879a2e3318283bd Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Wed, 23 Sep 2026 00:32:12 +0200 Subject: [PATCH 02/14] flux2 UI --- .github/workflows/ci.yml | 4 +- cmake/imgui.cmake | 41 ++ common/ort_session.cpp | 7 +- common/ort_session.h | 1 + image_gen/flux2/CMakeLists.txt | 6 + image_gen/flux2/README.md | 75 ++-- image_gen/flux2/flux2_cli.cpp | 31 +- image_gen/flux2/flux2_cli.h | 40 +- image_gen/flux2/flux2_cuda_backend.cpp | 86 +++-- image_gen/flux2/flux2_dx_backend.cpp | 77 ++-- image_gen/flux2/flux2_pipeline.cpp | 56 ++- image_gen/flux2/flux2_telemetry.h | 75 ++++ image_gen/flux2/flux2_ui.cpp | 500 +++++++++++++++++++++++++ image_gen/flux2/flux2_vk_backend.cpp | 88 +++-- image_gen/flux2/helpers/flux2.h | 40 +- 15 files changed, 1004 insertions(+), 123 deletions(-) create mode 100644 cmake/imgui.cmake create mode 100644 image_gen/flux2/flux2_telemetry.h create mode 100644 image_gen/flux2/flux2_ui.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4351b31..12222c1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,7 +54,7 @@ jobs: run: | set -euxo pipefail apt-get update - apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev + apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev - name: Check out source uses: actions/checkout@v7 @@ -97,7 +97,7 @@ jobs: run: | set -euxo pipefail apt-get update - apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev + apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev - name: Check out source uses: actions/checkout@v7 diff --git a/cmake/imgui.cmake b/cmake/imgui.cmake new file mode 100644 index 0000000..599e65a --- /dev/null +++ b/cmake/imgui.cmake @@ -0,0 +1,41 @@ +include_guard(GLOBAL) +include(FetchContent) + +# Attach the shared GLFW/OpenGL ImGui implementation to an existing UI target. +# Dependencies are fetched only when a UI target is requested. +function(din_link_imgui TARGET) + if (NOT TARGET din_imgui) + find_package(OpenGL REQUIRED GLOBAL) + find_package(Threads REQUIRED GLOBAL) + set(GLFW_BUILD_DOCS OFF CACHE BOOL "" FORCE) + set(GLFW_BUILD_TESTS OFF CACHE BOOL "" FORCE) + set(GLFW_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) + set(GLFW_INSTALL OFF CACHE BOOL "" FORCE) + if (UNIX AND NOT APPLE) + set(GLFW_BUILD_WAYLAND OFF CACHE BOOL "" FORCE) + set(GLFW_BUILD_X11 ON CACHE BOOL "" FORCE) + endif () + FetchContent_Declare(glfw + GIT_REPOSITORY https://github.com/glfw/glfw.git GIT_TAG 3.4 GIT_SHALLOW TRUE) + FetchContent_Declare(imgui + GIT_REPOSITORY https://github.com/ocornut/imgui.git GIT_TAG v1.91.5 GIT_SHALLOW TRUE) + FetchContent_MakeAvailable(glfw imgui) + add_library(din_imgui STATIC + ${imgui_SOURCE_DIR}/imgui.cpp + ${imgui_SOURCE_DIR}/imgui_draw.cpp + ${imgui_SOURCE_DIR}/imgui_tables.cpp + ${imgui_SOURCE_DIR}/imgui_widgets.cpp + ${imgui_SOURCE_DIR}/backends/imgui_impl_glfw.cpp + ${imgui_SOURCE_DIR}/backends/imgui_impl_opengl3.cpp) + target_include_directories(din_imgui PUBLIC ${imgui_SOURCE_DIR} ${imgui_SOURCE_DIR}/backends) + target_link_libraries(din_imgui PUBLIC glfw OpenGL::GL Threads::Threads) + endif () + target_link_libraries(${TARGET} PRIVATE din_imgui) +endfunction() + +# Usage: add_din_ui_executable(my_app main.cpp [other sources...]) +# App-specific libraries and settings belong in the caller's CMakeLists.txt. +function(add_din_ui_executable TARGET) + add_din_executable(${TARGET} ${ARGN}) + din_link_imgui(${TARGET}) +endfunction() diff --git a/common/ort_session.cpp b/common/ort_session.cpp index 81e6e66..149e797 100644 --- a/common/ort_session.cpp +++ b/common/ort_session.cpp @@ -487,7 +487,10 @@ void AppendTensorRTRTXEP(Ort::Env& env, Ort::SessionOptions& options, const std: const auto stem = fs::path(model_path).stem().string(); const auto cache_path = (fs::path(cache_dir) / (profile.cache_subpath.empty() ? stem : profile.cache_subpath)).string(); - ep_options.Add("nv_runtime_cache_path", cache_path.c_str()); + if (ep_context.enable_cache) + { + ep_options.Add("nv_runtime_cache_path", cache_path.c_str()); + } if (!profile.min_shapes.empty()) { ep_options.Add("nv_profile_min_shapes", profile.min_shapes.c_str()); @@ -833,7 +836,7 @@ OrtRunner::OrtRunner(Ort::Env& env_in, const std::string& model_path, const std: compute_stream = &owned_compute_stream_; } - if (!profile.skip_compile) + if (!profile.skip_compile && ep_context.enable_cache) { [[maybe_unused]] din::common::nvtx_scoped_range range{"load_or_compile_ep_context"}; session_model_path = CompileEpContextModel(env, model_path, cache_dir, ep_context, profile); diff --git a/common/ort_session.h b/common/ort_session.h index 3e0a09f..5491255 100644 --- a/common/ort_session.h +++ b/common/ort_session.h @@ -77,6 +77,7 @@ QueryCudaGraphicsInteropSharedMemoryInfo(int cuda_device_ordinal, bool high_prio struct EpContextOptions { std::string output_dir; + bool enable_cache = true; }; struct ModelProfile diff --git a/image_gen/flux2/CMakeLists.txt b/image_gen/flux2/CMakeLists.txt index 2308a15..0641675 100644 --- a/image_gen/flux2/CMakeLists.txt +++ b/image_gen/flux2/CMakeLists.txt @@ -4,6 +4,7 @@ enable_language(CUDA) include(slang) option(FLUX2_ENABLE_VULKAN_VALIDATION "Enable Vulkan validation layers for the Flux2 Vulkan target" ON) +option(DIN_BUILD_FLUX2_UI "Build the GLFW/OpenGL Flux2 desktop sample" ON) add_library(din_flux2_common STATIC utils.cpp @@ -135,6 +136,11 @@ add_din_executable(din_flux2_cli flux2_cli.cpp ) target_link_libraries(din_flux2_cli PRIVATE argparse din_flux2_runtime) +if (DIN_BUILD_FLUX2_UI) + include(imgui) + add_din_ui_executable(din_flux2_ui flux2_ui.cpp) + target_link_libraries(din_flux2_ui PRIVATE argparse din_flux2_runtime CUDA::cudart) +endif () if(TARGET din_flux2_dx) target_link_libraries(din_flux2_cli PRIVATE din_flux2_dx) endif() diff --git a/image_gen/flux2/README.md b/image_gen/flux2/README.md index b0f66a0..1f5a01d 100644 --- a/image_gen/flux2/README.md +++ b/image_gen/flux2/README.md @@ -1,58 +1,71 @@ # FLUX.2-klein -CLI image generation with ONNX Runtime. +Interactive image generation and a CLI for FLUX.2-klein-4B using ONNX Runtime and TensorRT RTX. + +The interactive studio was used to produce the results in +[The Weight Is Over – Interactive Diffusion on Consumer GPUs](https://arxiv.org/abs/2609.21849) +(Frieder Ganz and Maximilian Müller). See the paper for the small text-encoder translator, +interactive workflow, and speed/quality/memory trade-offs implemented here. ## Supported models -| Model | Hugging Face ID | Checkpoint | ONNX export | -| --- | --- | --- | --- | -| FLUX.2-klein-4B | `black-forest-labs/FLUX.2-klein-4B` | BF16 | FP32 I/O | -| FLUX.2-klein-4B FP8 | `black-forest-labs/FLUX.2-klein-4b-fp8` | FP8 | FP32 I/O | -| FLUX.2-klein-4B NVFP4 | `black-forest-labs/FLUX.2-klein-4b-nvfp4` | NVFP4 | FP32 I/O | +| Model | Hugging Face ID | Checkpoint | ONNX export | +|-----------------------|-------------------------------------------|------------|-------------| +| FLUX.2-klein-4B | `black-forest-labs/FLUX.2-klein-4B` | BF16 | FP32 I/O | +| FLUX.2-klein-4B FP8 | `black-forest-labs/FLUX.2-klein-4b-fp8` | FP8 | FP32 I/O | +| FLUX.2-klein-4B NVFP4 | `black-forest-labs/FLUX.2-klein-4b-nvfp4` | NVFP4 | FP32 I/O | -The CLI supports CPU or TensorRT RTX execution with CPU, CUDA, DirectX, DirectX CIG, and Vulkan processing where available. +The backend supports CPU or TensorRT RTX execution with CPU, CUDA, DirectX, DirectX CIG, and Vulkan processing where +available. -## Backends +## Export -`CiG` means compute-in-graphics: inference runs in the graphics GPU context so it can coexist with rendering workloads. +Run from the repository root: -| Mode | CLI | Inference and sampling | -| --- | --- | --- | -| Direct TensorRT RTX | `--provider trt-rtx --processing cuda` | TensorRT RTX inference with CUDA sampling. | -| Direct TensorRT RTX | `--provider trt-rtx --processing cpu` | TensorRT RTX inference with CPU sampling. | -| Vulkan interop | `--provider trt-rtx --processing vk` | TensorRT RTX inference with Vulkan sampling; shared buffers and timeline semaphores keep intermediate data on the GPU. | -| DirectX CiG | `--provider trt-rtx --processing dx-cig` | TensorRT RTX inference in a DirectX CiG context with DirectX sampling. | +```powershell +python ./model_export/export_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --output models/flux_full --model all --transformer-precision bf16 +python ./model_export/quantize_onnx.py --input models/flux_full/transformer_bf16 --quant fp8 --output models/flux_full/transformer_fp8 +``` -CiG is useful when inference must share the graphics GPU with rendering. It can reduce standalone inference performance because of graphics-concurrency limits. TensorRT RTX currently documents DirectX CiG support; Vulkan CiG support is planned. See [Simultaneous Compute and Graphics](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/inference-library/compute-graphics.html). +Use `--quant nvfp4` with an output directory of `transformer_nvfp4` for NVFP4. +All precisions share the text encoder, VAE, and tokenizer under one model root. -## Export +Optionally export Qwen3-0.6B + `gafr1/translator` as the small text encoder: ```powershell -cd model_export -python export_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --output D:\models\flux_full --model all --transformer-precision bf16 -python quantize_onnx.py --input D:\models\flux_full\transformer_bf16 --quant fp8 --output D:\models\flux_full\transformer_fp8 -python quantize_onnx.py --input D:\models\flux_full\transformer_bf16 --quant nvfp4 --output D:\models\flux_full\transformer_nvfp4 +python image_gen/flux2/model_export/export_flux2.py --model_name models/flux_full --output models/flux_full --model text_encoder_translator +python image_gen/flux2/model_export/verify_flux2.py --onnx_dir models/flux_full --encoder translator --encoder-only ``` -All precisions share one root directory. The export contains `text_encoder/model.onnx`, `vae_decoder/model.onnx`, and `tokenizer/tokenizer.json` once, plus one transformer per precision: `transformer_bf16/model.onnx`, `transformer_fp8/model.onnx`, and `transformer_nvfp4/model.onnx`. +Use `verify_flux2.py --help` for full-pipeline verification options. -## Verify +## Build and run ```powershell -python verify_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --onnx_dir D:\models\flux_full --precision fp8 --output_dir flux2-fp8 --provider trt-rtx +cmake --preset windows-x64 +cmake --build --preset windows-x64-release --target din_flux2_ui din_flux2_cli ``` -## Build +From the repo root you can then run: -```powershell -cmake --preset windows-x64 -cmake --build --preset windows-x64-debug --target din_flux2_cli ``` +out/build/windows-x64/bin/Release/din_flux2_ui.exe --model-dir /flux_full +``` + +The studio provides prompt editing, encoder/precision/backend selection, seed sweeps, +weight-streaming controls, and an image gallery with optional PNG output. It selects the +translator when exported; override with `--encoder 4b|translator`. +Runtime and compiled-model caching are always enabled; clear affected caches after re-exporting models in place. + +The UI uses GLFW, Dear ImGui, and OpenGL. For Debian/Ubuntu, install +`libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev`. +Use `-DDIN_BUILD_FLUX2_UI=OFF` for headless builds. -## Run +For pure CLI you can use: ```powershell -out\build\windows-x64\bin\Debug\din_flux2_cli.exe --model-dir D:\models\flux_full --precision nvfp4 --provider trt-rtx --processing cuda --prompt "a red fox" --output out +out/build/windows-x64/bin/Debug/din_flux2_cli.exe --model-dir models/flux_full --precision fp8 --provider trt-rtx --processing cuda --prompt "a red fox" --output out ``` -`--precision` selects `transformer_` and defaults to `bf16`. Only the transformer TensorRT RTX runtime-cache and EP-context paths are precision-qualified, so switching precisions never reuses an incompatible compiled transformer engine while the shared text encoder and VAE reuse their existing caches. +Use `--help` for available processing backends and options, including `--encoder`, +`--steps`, and startup-only `--ws`. CPU, CUDA, DirectX, and Vulkan processing are available where built. diff --git a/image_gen/flux2/flux2_cli.cpp b/image_gen/flux2/flux2_cli.cpp index c8d7e08..c40571c 100644 --- a/image_gen/flux2/flux2_cli.cpp +++ b/image_gen/flux2/flux2_cli.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -23,7 +24,8 @@ unsigned int parse_uint(const std::string& value, const char* flag_name) { size_t parsed_chars = 0; const unsigned long parsed = std::stoul(value, &parsed_chars, 10); - if (parsed_chars != value.size()) + if (value.empty() || value.front() == '-' || parsed_chars != value.size() || + parsed > std::numeric_limits::max()) { throw std::invalid_argument(std::string(flag_name) + " must be an unsigned integer"); } @@ -118,11 +120,11 @@ std::string parse_precision(std::string value) { c = static_cast(std::tolower(static_cast(c))); } - if (value == "bf16" || value == "fp8" || value == "nvfp4") + if (value == "bf16" || value == "fp16" || value == "fp8" || value == "nvfp4") { return value; } - throw std::invalid_argument("Precision must be one of: bf16, fp8, nvfp4"); + throw std::invalid_argument("Precision must be one of: bf16, fp16, fp8, nvfp4"); } void validate_config(const Flux2Config& config) @@ -159,7 +161,7 @@ Flux2Config parse_args(int argc, char* argv[]) parser.add_argument("--precision") .default_value(config.precision) .nargs(1) - .metavar("bf16|fp8|nvfp4") + .metavar("bf16|fp16|fp8|nvfp4") .help("Transformer precision to load from transformer_."); parser.add_argument("--ep-cache") .default_value(config.ep_cache_dir.string()) @@ -182,6 +184,11 @@ Flux2Config parse_args(int argc, char* argv[]) .metavar("TEXT") .help("Prompt text to encode with the Flux2 tokenizer."); parser.add_argument("--seed").default_value(std::to_string(config.seed)).nargs(1).metavar("N").help("Random seed."); + parser.add_argument("--encoder").default_value(std::string("4b")).help("4b or translator."); + parser.add_argument("--steps").default_value(std::string("4")).help("Denoise steps, 1 through 50."); + parser.add_argument("--ws") + .default_value(std::string("off")) + .help("Transformer resident budget: off or 0% through 100%."); parser.add_argument("--num-images") .default_value(std::to_string(config.num_images)) .nargs(1) @@ -206,12 +213,24 @@ Flux2Config parse_args(int argc, char* argv[]) config.ep_context_dir = parser.get("--ep-context-dir"); config.prompt = parser.get("--prompt"); config.seed = parse_uint(parser.get("--seed"), "--seed"); + const auto encoder = parser.get("--encoder"); + if (encoder != "4b" && encoder != "translator") + throw std::invalid_argument("--encoder must be 4b or translator"); + config.text_encoder = encoder == "4b" ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + const auto steps = parse_uint(parser.get("--steps"), "--steps"); + if (steps < 1 || steps > 50) + throw std::invalid_argument("--steps must be 1 through 50"); + config.steps = static_cast(steps); + config.weight_streaming_budget = parser.get("--ws"); + if (config.weight_streaming_budget == "off") + config.weight_streaming_budget.clear(); config.num_images = parse_uint(parser.get("--num-images"), "--num-images"); const std::filesystem::path output_dir = parser.get("--output"); std::filesystem::create_directories(output_dir); config.output_path = output_dir / "flux2.png"; validate_config(config); + ValidateFlux2Config(config); return config; } @@ -268,6 +287,10 @@ int main(int argc, char* argv[]) << " (seed=" << current_seed << ") ==========" << std::endl; const auto image_start = std::chrono::steady_clock::now(); Flux2Image image = pipeline->GenerateImage(current_seed); + const auto& times = image.timings; + std::cout << "Timings (ms): encode=" << times.encode_ms << " rng=" << times.rng_ms + << " denoise=" << times.denoise_ms << " decode=" << times.decode_ms << " total=" << times.total_ms + << std::endl; const auto image_end = std::chrono::steady_clock::now(); save_image(make_output_path(config, image_index), image); const std::chrono::duration image_duration = image_end - image_start; diff --git a/image_gen/flux2/flux2_cli.h b/image_gen/flux2/flux2_cli.h index 75ceeeb..9561370 100644 --- a/image_gen/flux2/flux2_cli.h +++ b/image_gen/flux2/flux2_cli.h @@ -4,7 +4,9 @@ #pragma once #include +#include #include +#include #include #include @@ -24,12 +26,35 @@ enum class Flux2ExecutionProvider TrtRtx, }; +enum class Flux2TextEncoder +{ + Qwen3_4B, + Qwen3_06BTranslator +}; + +// Called on the generation thread. Completed denoising steps are in [0, total]. +using Flux2Progress = std::function; + +struct Flux2Timings +{ + double encode_ms = 0; + double rng_ms = 0; + double denoise_ms = 0; + double decode_ms = 0; + double total_ms = 0; + double per_step_ms = 0; +}; + struct Flux2Config { Flux2ProcessingBackend processing = Flux2ProcessingBackend::Cuda; Flux2ExecutionProvider provider = Flux2ExecutionProvider::TrtRtx; std::filesystem::path model_dir; std::string precision = "bf16"; + Flux2TextEncoder text_encoder = Flux2TextEncoder::Qwen3_4B; + int steps = 4; + // Startup budget: empty keeps a non-streaming engine; -1 enables automatic streaming. + std::string weight_streaming_budget; std::filesystem::path ep_cache_dir = "artifacts/flux2/trt_rtx_cache"; std::filesystem::path ep_context_dir = "artifacts/flux2/ep_context"; std::filesystem::path output_path; @@ -38,11 +63,18 @@ struct Flux2Config unsigned int num_images = 5; }; +struct Flux2GenerationOptions +{ + // Omit to restore the startup budget. Applied via transformer session dynamic options only. + std::optional weight_streaming_budget; +}; + struct Flux2Image { std::vector data; int height = 0; int width = 0; + Flux2Timings timings; }; class Flux2ProcessingPipeline @@ -51,7 +83,13 @@ class Flux2ProcessingPipeline virtual ~Flux2ProcessingPipeline() = default; virtual void Initialize() = 0; virtual void SetPrompt(std::string prompt) = 0; - virtual Flux2Image GenerateImage(unsigned int seed) = 0; + // Opaque identity for diagnostics; never dereference or retain after pipeline destruction. + virtual const void* TransformerSessionIdentity() const = 0; + virtual Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress = {}, + const Flux2GenerationOptions& options = {}) = 0; }; std::unique_ptr CreateFlux2Pipeline(const Flux2Config& config); + +void ValidateFlux2Config(const Flux2Config& config); +bool IsFlux2BackendAvailable(Flux2ProcessingBackend backend); diff --git a/image_gen/flux2/flux2_cuda_backend.cpp b/image_gen/flux2/flux2_cuda_backend.cpp index 420c629..65ad090 100644 --- a/image_gen/flux2/flux2_cuda_backend.cpp +++ b/image_gen/flux2/flux2_cuda_backend.cpp @@ -27,6 +27,7 @@ #include "cuda_kernels.h" #include "flux2.h" #include "flux2_runtime_context.h" +#include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" #include @@ -155,19 +156,6 @@ void fill_position_ids(std::vector& img_ids, std::vector& txt_ } } -std::vector create_time_schedule() -{ - constexpr double MU = 2.291179894115571; - std::vector time_schedule(FLOW_STEPS + 1); - for (int i = 0; i <= FLOW_STEPS; ++i) - { - const double t = T_START - (T_START - T_END) * static_cast(i) / FLOW_STEPS; - time_schedule[i] = static_cast(std::exp(MU) / (std::exp(MU) + std::pow(1.0 / t - 1.0, 1.0))); - std::cout << " sigma[" << i << "] = " << time_schedule[i] << std::endl; - } - return time_schedule; -} - class EulerStage { public: @@ -394,9 +382,16 @@ struct CudaPipelineState }; void run_pipeline(FluxPipeline& pipeline, const std::vector& time_schedule, SamplingBackend sampling_backend, - bool device_is_cuda, Ort::SyncStream* compute_stream, bool encode_prompt) + bool device_is_cuda, Ort::SyncStream* compute_stream, bool encode_prompt, Flux2Telemetry& telemetry) { nvtx3::scoped_range nvtx_pipeline("run_pipeline"); + Ort::RunOptions transformer_options; + if (compute_stream) + { + transformer_options.SetSyncStream(*compute_stream); + transformer_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); + } + Ort::RunOptions run_options; if (compute_stream != nullptr) { @@ -404,13 +399,16 @@ void run_pipeline(FluxPipeline& pipeline, const std::vector& time_schedul run_options.AddConfigEntry(kOrtRunOptionsConfigDisableSynchronizeExecutionProviders, "1"); } + const int steps = static_cast(time_schedule.size()) - 1; if (encode_prompt) { + telemetry.Begin("encode", &Flux2Timings::encode_ms); nvtx3::scoped_range nvtx("text_encoder"); pipeline.text_encoder_session.Run(run_options, pipeline.text_encoder_io); } - for (int step = 0; step < FLOW_STEPS; ++step) + telemetry.Begin("denoise", &Flux2Timings::denoise_ms); + for (int step = 0; step < steps; ++step) { nvtx3::scoped_range nvtx_step("diffusion_step_" + std::to_string(step)); const float t_curr = time_schedule[step]; @@ -425,14 +423,16 @@ void run_pipeline(FluxPipeline& pipeline, const std::vector& time_schedul { nvtx3::scoped_range nvtx_transformer("transformer"); - pipeline.transformer_session.Run(run_options, pipeline.transformer_io); + pipeline.transformer_session.Run(transformer_options, pipeline.transformer_io); } { nvtx3::scoped_range nvtx_euler("euler_step"); pipeline.euler.Run(compute_stream, sampling_backend, device_is_cuda, t_curr, t_next); } + telemetry.Step(step + 1); } + telemetry.Begin("decode", &Flux2Timings::decode_ms); { nvtx3::scoped_range nvtx_post("postprocess"); @@ -477,8 +477,9 @@ Ort::ConstEpDevice find_cpu_device(Ort::Env& env) void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const Flux2Config& config, ExecutionProviderMode provider_mode, SamplingBackend sampling_backend) { - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision); - const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision); + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); + const Flux2ModelCachePaths cache_paths = + MakeFlux2ModelCachePaths(config.precision, "", config.text_encoder, !config.weight_streaming_budget.empty()); const bool trt_rtx_device = is_trt_rtx_device(ep_device); const bool cpu_device = is_cpu_device(ep_device); const bool has_separate_binding = !cpu_device; @@ -522,15 +523,18 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const } din::common::EpContextOptions ep_context; ep_context.output_dir = config.ep_context_dir.string(); + ep_context.enable_cache = true; const std::string cache_dir = config.ep_cache_dir.string(); - auto make_profile = [](std::string cache_subpath) + auto make_profile = [&](std::string cache_subpath, bool transformer = false) { din::common::ModelProfile profile; profile.cache_subpath = std::move(cache_subpath); profile.enable_cuda_graph = false; profile.embed_ep_context = false; profile.extra_ep_options.emplace_back("nv_length_aux_stream_array", "0"); + if (transformer && !config.weight_streaming_budget.empty()) + profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.weight_streaming_budget); return profile; }; @@ -538,9 +542,9 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const state.env, model_paths.text_encoder_model.string(), provider, cache_dir, ep_context, make_profile(cache_paths.text_encoder), compute_stream_ptr); std::cout << " Text encoder loaded" << std::endl; - state.transformer_runner = - std::make_unique(state.env, model_paths.transformer_model.string(), provider, cache_dir, - ep_context, make_profile(cache_paths.transformer), compute_stream_ptr); + state.transformer_runner = std::make_unique( + state.env, model_paths.transformer_model.string(), provider, cache_dir, ep_context, + make_profile(cache_paths.transformer, true), compute_stream_ptr); std::cout << " Transformer loaded" << std::endl; state.vae_decoder_runner = std::make_unique(state.env, model_paths.vae_decoder_model.string(), provider, cache_dir, @@ -638,7 +642,7 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const bn_std_ready.Sync(); } - state.time_schedule = create_time_schedule(); + state.time_schedule = MakeFlux2Schedule(config.steps); state.euler.BindInput("hidden_states", *state.hidden_states); state.euler.BindInput("sample", *state.transformer_output); @@ -651,10 +655,16 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const state.postprocess.BindOutput("latent_sample", *state.decoder_latent); } -Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed) +Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed, const Flux2Progress& progress) { Ort::SyncStream* compute_stream_ptr = state.compute_stream ? &*state.compute_stream : nullptr; + Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, + [&] + { + if (compute_stream_ptr) + signal_and_wait_on_host(*compute_stream_ptr); + }); FluxPipeline pipeline{state.text_encoder_runner->session, state.transformer_runner->session, state.vae_decoder_runner->session, @@ -678,7 +688,7 @@ Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed) } run_pipeline(pipeline, state.time_schedule, state.sampling_backend, state.binding_is_cuda, compute_stream_ptr, - !state.prompt_embeds_valid); + !state.prompt_embeds_valid, telemetry); state.prompt_embeds_valid = true; Flux2Image image; @@ -692,6 +702,7 @@ Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed) } const float* image_data = state.decoded_image->HostData(); std::copy_n(image_data, image.data.size(), image.data.data()); + image.timings = telemetry.Finish(); return image; } @@ -746,6 +757,8 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline void SetPrompt(std::string prompt) override { + if (state_ && state_->prompt_embeds_valid && config_.prompt == prompt) + return; config_.prompt = std::move(prompt); if (!state_) { @@ -753,7 +766,8 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir, config_.precision); + const Flux2ModelPaths model_paths = + MakeFlux2ModelPaths(config_.model_dir, config_.precision, config_.text_encoder); const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state_->token->HostData(), state_->attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); @@ -767,17 +781,33 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } } - Flux2Image GenerateImage(unsigned int seed) override + const void* TransformerSessionIdentity() const override + { + return state_ ? static_cast(state_->transformer_runner->session) : nullptr; + } + + Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, + const Flux2GenerationOptions& options) override { if (!state_) { throw std::runtime_error("Flux2 CUDA pipeline was not initialized"); } - return run_initialized_pipeline(*state_, seed); + const std::string budget = options.weight_streaming_budget.value_or(config_.weight_streaming_budget); + const std::string& previous = applied_budget_ ? *applied_budget_ : config_.weight_streaming_budget; + if (budget != previous) + { + const char* keys[] = {"nv_weight_streaming_budget"}; + const char* values[] = {budget.c_str()}; + state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); + applied_budget_ = budget; + } + return run_initialized_pipeline(*state_, seed, progress); } private: Flux2Config config_; + std::optional applied_budget_; Flux2RuntimeContext& runtime_; std::unique_ptr state_; }; diff --git a/image_gen/flux2/flux2_dx_backend.cpp b/image_gen/flux2/flux2_dx_backend.cpp index 8cd2acb..cb19bcc 100644 --- a/image_gen/flux2/flux2_dx_backend.cpp +++ b/image_gen/flux2/flux2_dx_backend.cpp @@ -19,6 +19,7 @@ #include "dx_helper.h" #include "flux2.h" #include "flux2_runtime_context.h" +#include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" #include @@ -289,13 +290,18 @@ static void run_pipeline_dx(Ort::Session& text_encoder_session, Ort::Session& tr DxPostprocessPipeline& postprocess_pipeline, OrtD3D12Fence& sync_fence, Ort::SyncStream& ort_stream, DxBuffer& timestep_buf, DxBuffer& hidden_states, DxBuffer& transformer_output, DxBuffer& decoder_latent, DxBuffer& bn_mean, DxBuffer& bn_std, - bool encode_prompt) + bool encode_prompt, Flux2Telemetry& telemetry) { + Ort::RunOptions transformer_options; + transformer_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); + Ort::RunOptions run_options; run_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); + const int steps = static_cast(time_schedule.size()) - 1; if (encode_prompt) { + telemetry.Begin("encode", &Flux2Timings::encode_ms); text_encoder_session.Run(run_options, text_encoder_io); } @@ -305,12 +311,13 @@ static void run_pipeline_dx(Ort::Session& text_encoder_session, Ort::Session& tr const uint64_t initial_timestep_uploaded = sync_fence.signal_d3d12(dx.queue.Get()); sync_fence.wait_ort(ort_stream, initial_timestep_uploaded); - for (int step = 0; step < FLOW_STEPS; ++step) + telemetry.Begin("denoise", &Flux2Timings::denoise_ms); + for (int step = 0; step < steps; ++step) { const float t_curr = time_schedule[step]; const float t_next = time_schedule[step + 1]; - transformer_session.Run(run_options, transformer_io); + transformer_session.Run(transformer_options, transformer_io); const uint64_t transformer_done = sync_fence.signal_ort(ort_stream); sync_fence.wait_d3d12(dx.queue.Get(), transformer_done); @@ -322,19 +329,21 @@ static void run_pipeline_dx(Ort::Session& text_encoder_session, Ort::Session& tr dx.begin(); euler_pipeline.record(dx, hidden_states, transformer_output, euler); - if (step + 1 < FLOW_STEPS) + if (step + 1 < steps) { dx.record_upload(timestep_buf, &t_next, sizeof(float)); } dx.execute(); const uint64_t d3d_done = sync_fence.signal_d3d12(dx.queue.Get()); - if (step + 1 < FLOW_STEPS) + if (step + 1 < steps) { sync_fence.wait_ort(ort_stream, d3d_done); } + telemetry.Step(step + 1); } + telemetry.Begin("decode", &Flux2Timings::decode_ms); DxPostprocessConstants constants{}; constants.C = static_cast(LATENT_CHANNELS); constants.I = static_cast(LATENT_HEIGHT); @@ -449,9 +458,10 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi { throw std::runtime_error("DirectX processing requires --provider trt-rtx"); } - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision); + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); const Flux2ModelCachePaths cache_paths = - MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "dx_cig" : "dx"); + MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "dx_cig" : "dx", config.text_encoder, + !config.weight_streaming_budget.empty()); std::cout << "Model dir: " << model_paths.base_dir.string() << "\n" << "CIG: " << (state.use_cig ? "enabled" : "disabled") << "\n" @@ -491,14 +501,17 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi std::cout << "=== Loading ONNX Models ===" << std::endl; din::common::EpContextOptions ep_context; ep_context.output_dir = config.ep_context_dir.string(); + ep_context.enable_cache = true; const std::string cache_dir = config.ep_cache_dir.string(); - auto make_profile = [&](std::string cache_subpath) + auto make_profile = [&](std::string cache_subpath, bool transformer = false) { din::common::ModelProfile profile; profile.cache_subpath = std::move(cache_subpath); profile.enable_cuda_graph = state.use_cig ? false : true; profile.embed_ep_context = false; profile.extra_ep_options = graphics_ep_options; + if (transformer && !config.weight_streaming_budget.empty()) + profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.weight_streaming_budget); return profile; }; @@ -507,7 +520,7 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi make_profile(cache_paths.text_encoder), &*state.sync_stream); state.transformer_runner = std::make_unique( state.env, model_paths.transformer_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.transformer), &*state.sync_stream); + make_profile(cache_paths.transformer, true), &*state.sync_stream); state.vae_decoder_runner = std::make_unique( state.env, model_paths.vae_decoder_model.string(), "trt-rtx", cache_dir, ep_context, make_profile(cache_paths.vae_decoder), &*state.sync_stream); @@ -627,14 +640,7 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.dx.record_upload(state.txt_ids_buf, txt_ids_cpu.data(), state.txt_ids_buf.size); } - constexpr double MU = 2.291179894115571; - state.time_schedule.resize(FLOW_STEPS + 1); - for (int i = 0; i <= FLOW_STEPS; ++i) - { - double t = T_START - (T_START - T_END) * static_cast(i) / FLOW_STEPS; - state.time_schedule[i] = static_cast(std::exp(MU) / (std::exp(MU) + std::pow(1.0 / t - 1.0, 1.0))); - std::cout << " sigma[" << i << "] = " << state.time_schedule[i] << std::endl; - } + state.time_schedule = MakeFlux2Schedule(config.steps); std::cout << "\n=== Tokenizing Prompt ===" << std::endl; { @@ -653,8 +659,15 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.sync_fence->wait_ort(*state.sync_stream, static_uploads_done); } -static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed) +static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, const Flux2Progress& progress) { + Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, + [&] + { + const auto completed = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, completed); + signal_and_wait_on_host(*state.sync_stream); + }); std::vector latent_cpu(state.hidden_count); initialize_latent(latent_cpu.data(), latent_cpu.size(), seed); state.dx.begin(); @@ -668,7 +681,7 @@ static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed) *state.vae_decoder_io, state.time_schedule, state.dx, *state.euler_pipeline, *state.postprocess_pipeline, *state.sync_fence, *state.sync_stream, state.timestep_buf, state.hidden_states, state.transformer_output, state.decoder_latent, state.bn_mean_buf, - state.bn_std_buf, !state.prompt_embeds_valid); + state.bn_std_buf, !state.prompt_embeds_valid, telemetry); state.prompt_embeds_valid = true; state.dx.readback(state.decoded_image_buf, state.readback_buffer, state.decoded_image_bytes); @@ -681,6 +694,7 @@ static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed) std::copy_n(image_data, image.data.size(), image.data.data()); state.dx.unmap(state.readback_buffer); + image.timings = telemetry.Finish(); return image; } @@ -707,6 +721,8 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline void SetPrompt(std::string prompt) override { + if (state_ && state_->prompt_embeds_valid && config_.prompt == prompt) + return; config_.prompt = std::move(prompt); if (!state_) { @@ -714,7 +730,8 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir, config_.precision); + const Flux2ModelPaths model_paths = + MakeFlux2ModelPaths(config_.model_dir, config_.precision, config_.text_encoder); const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); @@ -729,17 +746,33 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline state_->sync_fence->wait_ort(*state_->sync_stream, prompt_uploaded); } - Flux2Image GenerateImage(unsigned int seed) override + const void* TransformerSessionIdentity() const override + { + return state_ ? static_cast(state_->transformer_runner->session) : nullptr; + } + + Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, + const Flux2GenerationOptions& options) override { if (!state_) { throw std::runtime_error("Flux2 DirectX pipeline was not initialized"); } - return run_dx_image(*state_, seed); + const std::string budget = options.weight_streaming_budget.value_or(config_.weight_streaming_budget); + const std::string& previous = applied_budget_ ? *applied_budget_ : config_.weight_streaming_budget; + if (budget != previous) + { + const char* keys[] = {"nv_weight_streaming_budget"}; + const char* values[] = {budget.c_str()}; + state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); + applied_budget_ = budget; + } + return run_dx_image(*state_, seed, progress); } private: Flux2Config config_; + std::optional applied_budget_; Flux2RuntimeContext& runtime_; std::unique_ptr state_; }; diff --git a/image_gen/flux2/flux2_pipeline.cpp b/image_gen/flux2/flux2_pipeline.cpp index 3b69570..6f2b247 100644 --- a/image_gen/flux2/flux2_pipeline.cpp +++ b/image_gen/flux2/flux2_pipeline.cpp @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +#include #include #include @@ -10,7 +11,10 @@ namespace { Flux2RuntimeContext& GetProcessRuntime(Flux2ExecutionProvider provider) { - static Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "Flux2"}; + static Ort::Env env{ + std::getenv("DIN_FLUX2_VERBOSE") ? ORT_LOGGING_LEVEL_VERBOSE : ORT_LOGGING_LEVEL_WARNING, + "Flux2" + }; static Flux2RuntimeContext runtime{env}; static std::once_flag register_trt_once; @@ -28,6 +32,7 @@ Flux2RuntimeContext& GetProcessRuntime(Flux2ExecutionProvider provider) std::unique_ptr CreateFlux2Pipeline(const Flux2Config& config) { + ValidateFlux2Config(config); Flux2RuntimeContext& runtime = GetProcessRuntime(config.provider); if (config.processing == Flux2ProcessingBackend::Cpu || config.processing == Flux2ProcessingBackend::Cuda) @@ -55,3 +60,52 @@ std::unique_ptr CreateFlux2Pipeline(const Flux2Config& throw std::invalid_argument("Unknown Flux2 processing backend"); } + +bool IsFlux2BackendAvailable(Flux2ProcessingBackend backend) +{ + switch (backend) + { + case Flux2ProcessingBackend::Cpu: + case Flux2ProcessingBackend::Cuda: + return true; +#if defined(DIN_FLUX2_BUILD_DX) + case Flux2ProcessingBackend::Dx: + case Flux2ProcessingBackend::DxCig: + return true; +#endif +#if defined(DIN_FLUX2_BUILD_VK) + case Flux2ProcessingBackend::Vk: + case Flux2ProcessingBackend::VkCig: + return true; +#endif + default: + return false; + } +} + +void ValidateFlux2Config(const Flux2Config& config) +{ + if (!IsFlux2BackendAvailable(config.processing)) + throw std::invalid_argument("Selected processing backend was not built"); + if (config.provider == Flux2ExecutionProvider::Cpu && config.processing != Flux2ProcessingBackend::Cpu) + throw std::invalid_argument("CPU execution requires CPU processing"); + if (config.steps < 1 || config.steps > 50) + throw std::invalid_argument("Denoise steps must be in [1, 50]"); + if (config.num_images == 0) + throw std::invalid_argument("Image count must be positive"); + if (config.precision != "bf16" && config.precision != "fp16" && config.precision != "fp8" && + config.precision != "nvfp4") + throw std::invalid_argument("Precision must be bf16, fp16, fp8, or nvfp4"); + if (!config.weight_streaming_budget.empty()) + { + if (config.provider != Flux2ExecutionProvider::TrtRtx) + throw std::invalid_argument("Weight streaming requires TensorRT RTX"); + const auto& budget = config.weight_streaming_budget; + if (budget == "-1") + return; + size_t parsed = 0; + const int value = std::stoi(budget, &parsed); + if (parsed + 1 != budget.size() || budget.back() != '%' || value < 0 || value > 100) + throw std::invalid_argument("Weight streaming budget must be 0% through 100%"); + } +} diff --git a/image_gen/flux2/flux2_telemetry.h b/image_gen/flux2/flux2_telemetry.h new file mode 100644 index 0000000..7c623fb --- /dev/null +++ b/image_gen/flux2/flux2_telemetry.h @@ -0,0 +1,75 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include + +#include "flux2_cli.h" + +// Stage timings are synchronized wall times, including stage transfers. +// Synchronization is supplied by each backend to respect its ORT/graphics ordering. +class Flux2Telemetry +{ +public: + Flux2Telemetry(int steps, Flux2Progress progress, std::function synchronize) + : steps_(steps) + , progress_(std::move(progress)) + , synchronize_(std::move(synchronize)) + , start_(Clock::now()) + , stage_start_(start_) + { + Report("rng", 0); + } + + void Begin(const char* stage, double Flux2Timings::* field) + { + synchronize_(); + const auto now = Clock::now(); + timings_.*field_ += Milliseconds(now - stage_start_); + stage_start_ = now; + field_ = field; + Report(stage, 0); + } + + void Step(int completed) + { + if (progress_) + { + synchronize_(); + Report("denoise", completed); + } + } + + Flux2Timings Finish() + { + synchronize_(); + const auto now = Clock::now(); + timings_.*field_ += Milliseconds(now - stage_start_); + timings_.total_ms = Milliseconds(now - start_); + timings_.per_step_ms = timings_.denoise_ms / steps_; + Report("done", steps_); + return timings_; + } + +private: + using Clock = std::chrono::steady_clock; + + static double Milliseconds(Clock::duration duration) + { + return std::chrono::duration(duration).count(); + } + + void Report(const char* stage, int completed) + { + if (progress_) + progress_(stage, completed, steps_); + } + + int steps_; + Flux2Progress progress_; + std::function synchronize_; + Clock::time_point start_, stage_start_; + Flux2Timings timings_; + double Flux2Timings::* field_ = &Flux2Timings::rng_ms; +}; diff --git a/image_gen/flux2/flux2_ui.cpp b/image_gen/flux2/flux2_ui.cpp new file mode 100644 index 0000000..ef31a9a --- /dev/null +++ b/image_gen/flux2/flux2_ui.cpp @@ -0,0 +1,500 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "flux2.h" +#include "flux2_cli.h" +#include "io/image.h" +#include +#include +#include +#include +#include + +namespace +{ + // Stable per-root namespace prevents a different model directory from reusing stale engines. + std::string CacheNamespace(const std::filesystem::path& root) + { + const auto path = std::filesystem::weakly_canonical(root).generic_u8string(); + uint64_t hash = 14695981039346656037ull; + for (const auto c : path) + { + hash ^= static_cast(c); + hash *= 1099511628211ull; + } + return "root_" + std::to_string(hash); + } + + struct Request + { + Flux2Config config; + Flux2GenerationOptions generation; + bool random_seeds = false; + }; + + struct Result + { + Flux2Image image; + unsigned int seed = 0; + GLuint texture = 0; + }; + + // Only the worker touches the pipeline; all graphics resources belong to the UI thread. + class Generator + { + public: + Generator() + : worker_( + [this] + { + Run(); + }) + { + } + + ~Generator() + { + { + std::lock_guard lock(mutex); + closing_ = true; + } + cv_.notify_one(); + worker_.join(); + } + + void Submit(Request request) + { + std::lock_guard lock(mutex); + if (busy) + return; + busy = true; + failed = false; + status = "Loading models (first compilation can take several minutes)"; + completed = 0; + total = request.config.steps * static_cast(request.config.num_images); + request_ = std::move(request); + cv_.notify_one(); + } + + std::mutex mutex; + bool busy = false, failed = false; + std::string status = "Ready"; + int completed = 0, total = 1; + size_t used_mb = 0, total_mb = 0; + std::vector ready; + + private: + static bool SameModels(const Flux2Config& a, const Flux2Config& b) + { + return a.model_dir == b.model_dir && a.processing == b.processing && a.provider == b.provider && + a.precision == b.precision && a.text_encoder == b.text_encoder && a.steps == b.steps && + a.weight_streaming_budget == b.weight_streaming_budget && + a.ep_cache_dir == b.ep_cache_dir && a.ep_context_dir == b.ep_context_dir; + } + + void UpdateMemory() + { + size_t free = 0, total_bytes = 0; + if (cudaMemGetInfo(&free, &total_bytes) == cudaSuccess) + { + std::lock_guard lock(mutex); + used_mb = (total_bytes - free) / (1024 * 1024); + total_mb = total_bytes / (1024 * 1024); + } + else + cudaGetLastError(); + } + + void Run() + { + std::unique_ptr pipeline; + std::optional loaded; + std::mt19937 random(std::random_device{}()); + for (;;) + { + std::unique_lock lock(mutex); + cv_.wait_for(lock, std::chrono::milliseconds(500), + [&] + { + return closing_ || request_.has_value(); + }); + if (closing_) + break; + if (!request_) + { + lock.unlock(); + if (loaded && loaded->provider == Flux2ExecutionProvider::TrtRtx) + UpdateMemory(); + continue; + } + Request request = std::move(*request_); + request_.reset(); + lock.unlock(); + try + { + const auto& config = request.config; + ValidateFlux2Config(config); + if (!loaded || !SameModels(*loaded, config)) + { + pipeline.reset(); // Release old engines before allocating the replacement. + loaded.reset(); + pipeline = CreateFlux2Pipeline(config); + pipeline->Initialize(); + loaded = config; + } + pipeline->SetPrompt(config.prompt); + for (unsigned int i = 0; i < config.num_images; ++i) + { + const unsigned int seed = request.random_seeds ? random() : config.seed + i; + auto image = pipeline->GenerateImage( + seed, + [&](const char* stage, int step, int) + { + { + std::lock_guard guard(mutex); + completed = static_cast(i) * config.steps + step; + status = "Image " + std::to_string(i + 1) + "/" + std::to_string(config.num_images) + + + ": " + stage; + } + if (config.provider == Flux2ExecutionProvider::TrtRtx) + UpdateMemory(); + }, + request.generation); + if (!config.output_path.empty()) + { + std::filesystem::create_directories(config.output_path); + const auto file = + config.output_path / ("flux2_" + std::to_string(i) + "_" + std::to_string(seed) + + ".png"); + if (!din::io::SaveRgbFloatImage(file.string(), image.data.data(), image.height, image.width, + din::io::ImageValueRange::MinusOneToOne)) + throw std::runtime_error("Could not save " + file.string()); + } + std::lock_guard guard(mutex); + ready.push_back({std::move(image), seed, 0}); + } + std::lock_guard guard(mutex); + completed = total; + status = "Done"; + } + catch (const std::exception& error) + { + pipeline.reset(); + loaded.reset(); + std::lock_guard guard(mutex); + status = std::string("Error: ") + error.what(); + failed = true; + } + std::lock_guard guard(mutex); + busy = false; + } + // Destroy sessions, streams and GPU allocations on their owning thread. + } + + std::condition_variable cv_; + bool closing_ = false; + std::optional request_; + std::thread worker_; + }; + + void Upload(Result& result) + { + const auto& img = result.image; + const size_t pixels = static_cast(img.width) * img.height; + if (!pixels || img.data.size() != pixels * 3) + throw std::runtime_error("Invalid image returned by pipeline"); + std::vector rgb(pixels * 3); + for (size_t p = 0; p < pixels; ++p) + for (size_t c = 0; c < 3; ++c) + { + const float v = img.data[c * pixels + p]; + rgb[p * 3 + c] = + std::isfinite(v) ? static_cast(std::clamp(v * 0.5f + 0.5f, 0.0f, 1.0f) * 255.0f) : 0; + } + glGenTextures(1, &result.texture); + glBindTexture(GL_TEXTURE_2D, result.texture); + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR); + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR); + glPixelStorei(GL_UNPACK_ALIGNMENT, 1); + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, img.width, img.height, 0, GL_RGB, GL_UNSIGNED_BYTE, rgb.data()); + } + + int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit_after) + { + Generator generator; + Flux2Config config = initial; + std::array model{}, output{}; + std::array prompt{}; + std::snprintf(model.data(), model.size(), "%s", config.model_dir.string().c_str()); + std::snprintf(output.data(), output.size(), "%s", config.output_path.string().c_str()); + std::snprintf(prompt.data(), prompt.size(), "%s", config.prompt.c_str()); + int streaming_mode = 0; // Auto / Manual / Off; independent of session configuration. + int budget = 50; + int count = static_cast(config.num_images); + int encoder = config.text_encoder == Flux2TextEncoder::Qwen3_4B ? 0 : 1; + int precision = config.precision == "bf16" + ? 0 + : config.precision == "fp16" + ? 1 + : config.precision == "fp8" + ? 2 + : 3; + const char* precisions[] = {"bf16", "fp16", "fp8", "nvfp4"}; + const char* backends[] = {"CPU", "CUDA", "DirectX", "DirectX CIG", "Vulkan", "Vulkan CIG"}; + int backend = static_cast(config.processing); + int provider = config.provider == Flux2ExecutionProvider::Cpu ? 0 : 1; + std::vector results; + int selected = -1, frames = 0; + bool launched = false; + int exit_code = 0; + GLFWwindow* window = glfwGetCurrentContext(); + while (!glfwWindowShouldClose(window)) + { + glfwPollEvents(); + ImGui_ImplOpenGL3_NewFrame(); + ImGui_ImplGlfw_NewFrame(); + ImGui::NewFrame(); + ImGui::SetNextWindowPos({0, 0}); + ImGui::SetNextWindowSize(ImGui::GetIO().DisplaySize); + ImGui::Begin("Flux2 Studio", nullptr, + ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize); + ImGui::TextUnformatted("DIN Deploy - FLUX.2-klein Interactive Studio"); + ImGui::Separator(); + bool busy, failed; + int completed, total; + std::string status; + size_t used_mb, total_mb; + { + std::lock_guard lock(generator.mutex); + busy = generator.busy; + failed = generator.failed; + status = generator.status; + completed = generator.completed; + total = generator.total; + used_mb = generator.used_mb; + total_mb = generator.total_mb; + for (auto& result : generator.ready) + { + results.push_back(std::move(result)); + selected = static_cast(results.size()) - 1; + } + generator.ready.clear(); + } + if (launched && !busy && exit_after) + { + exit_code = failed ? 1 : 0; + glfwSetWindowShouldClose(window, true); + } + ImGui::BeginChild("controls", ImVec2(430, 0), true); + ImGui::BeginDisabled(busy); + ImGui::InputText("Model root", model.data(), model.size()); + ImGui::InputText("Save directory", output.data(), output.size()); + ImGui::TextUnformatted("Prompt"); + ImGui::InputTextMultiline("##prompt", prompt.data(), prompt.size(), ImVec2(-1, 100)); + ImGui::RadioButton("Qwen3-4B", &encoder, 0); + const bool translator_exists = std::filesystem::is_regular_file(std::filesystem::path(model.data()) / + "text_encoder_translator/model.onnx"); + ImGui::BeginDisabled(!translator_exists); + ImGui::RadioButton("Qwen3-0.6B + translator", &encoder, 1); + ImGui::EndDisabled(); + if (!translator_exists) + ImGui::TextWrapped("Translator ONNX is missing from this model root."); + ImGui::Combo("Precision", &precision, precisions, IM_ARRAYSIZE(precisions)); + ImGui::Combo("Provider", &provider, "CPU\0TensorRT RTX\0"); + if (provider == 0) + { + backend = 0; + streaming_mode = 0; + } + if (ImGui::BeginCombo("Processing", backends[backend])) + { + for (int i = 0; i < 6; ++i) + { + if (!IsFlux2BackendAvailable(static_cast(i)) || (provider == 0 && i != 0)) + continue; + if (ImGui::Selectable(backends[i], backend == i)) + backend = i; + } + ImGui::EndCombo(); + } + ImGui::BeginDisabled(provider == 0); + ImGui::Combo("Weight streaming", &streaming_mode, "Auto\0Manual\0Off\0"); + if (streaming_mode == 1) + ImGui::SliderInt("Resident weights %", &budget, 10, 100); + ImGui::EndDisabled(); + ImGui::InputScalar("Seed", ImGuiDataType_U32, &config.seed); + ImGui::SliderInt("Denoise steps", &config.steps, 1, 50); + ImGui::SliderInt("Images", &count, 1, 16); + bool generate = ImGui::Button("Generate", ImVec2(-1, 35)); + bool random = ImGui::Button("Random sweep", ImVec2(-1, 30)); + ImGui::EndDisabled(); + if (!busy && (generate || random || (auto_generate && !launched))) + { + for (auto& result : results) + if (result.texture) + glDeleteTextures(1, &result.texture); + results.clear(); + selected = -1; + config.model_dir = model.data(); + config.output_path = output.data(); + config.prompt = prompt.data(); + const auto cache_namespace = CacheNamespace(config.model_dir); + config.ep_cache_dir = initial.ep_cache_dir / cache_namespace; + config.ep_context_dir = initial.ep_context_dir / cache_namespace; + config.text_encoder = encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + config.precision = precisions[precision]; + config.num_images = static_cast(count); + config.processing = static_cast(backend); + config.provider = provider == 0 ? Flux2ExecutionProvider::Cpu : Flux2ExecutionProvider::TrtRtx; + config.weight_streaming_budget = provider == 0 ? "" : "-1"; + Flux2GenerationOptions generation; + if (provider != 0 && streaming_mode != 0) + generation.weight_streaming_budget = streaming_mode == 1 ? std::to_string(budget) + "%" : "0"; + generator.Submit({config, generation, random}); + launched = true; + } + ImGui::ProgressBar(static_cast(completed) / std::max(total, 1), ImVec2(-1, 0)); + ImGui::TextWrapped("%s", status.c_str()); + if (provider != 0 && total_mb) + ImGui::Text("Device memory: %zu / %zu MiB", used_mb, total_mb); + if (selected >= 0) + { + const auto& t = results[selected].image.timings; + ImGui::Separator(); + ImGui::Text("Encode %.0f ms | RNG %.0f ms", t.encode_ms, t.rng_ms); + ImGui::Text("Denoise %.0f ms (%.0f ms/step)", t.denoise_ms, t.per_step_ms); + ImGui::Text("Decode %.0f ms | Total %.2f s", t.decode_ms, t.total_ms / 1000); + } + ImGui::EndChild(); + ImGui::SameLine(); + ImGui::BeginChild("images", ImVec2(0, 0), true); + for (auto& result : results) + if (!result.texture) + Upload(result); + ImGui::BeginChild("preview", ImVec2(0, -115), false); + if (selected >= 0) + { + const auto available = ImGui::GetContentRegionAvail(); + const float side = std::max(1.0f, std::min(available.x, available.y)); + ImGui::Image(static_cast(results[selected].texture), ImVec2(side, side)); + } + else + ImGui::TextDisabled("Generated images will appear here."); + ImGui::EndChild(); + ImGui::BeginChild("thumbnails", ImVec2(0, 110), false, ImGuiWindowFlags_HorizontalScrollbar); + for (int i = 0; i < static_cast(results.size()); ++i) + { + if (i) + ImGui::SameLine(); + ImGui::PushID(i); + ImGui::BeginGroup(); + if (ImGui::ImageButton("image", static_cast(results[i].texture), ImVec2(72, 72))) + selected = i; + ImGui::Text("%u", results[i].seed); + ImGui::EndGroup(); + ImGui::PopID(); + } + ImGui::EndChild(); + ImGui::EndChild(); + ImGui::End(); + ImGui::Render(); + int width, height; + glfwGetFramebufferSize(window, &width, &height); + glViewport(0, 0, width, height); + glClearColor(0.09f, 0.09f, 0.10f, 1); + glClear(GL_COLOR_BUFFER_BIT); + ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData()); + glfwSwapBuffers(window); + if (smoke && ++frames >= 3) + break; + } + for (auto& result : results) + if (result.texture) + glDeleteTextures(1, &result.texture); + return exit_code; + } +} // namespace + +int main(int argc, char** argv) +{ + GLFWwindow* window = nullptr; + bool imgui = false, glfw_backend = false, gl_backend = false; + int result = 1; + try + { + argparse::ArgumentParser args("din_flux2_ui"); + args.add_argument("--model-dir").default_value(DEFAULT_MODEL_BASE_PATH.string()); + args.add_argument("--output").default_value(std::string{}); + args.add_argument("--ep-cache").default_value(Flux2Config{}.ep_cache_dir.string()); + args.add_argument("--ep-context-dir").default_value(Flux2Config{}.ep_context_dir.string()); + args.add_argument("--prompt").default_value(std::string(DEFAULT_PROMPT)); + args.add_argument("--encoder").default_value(std::string("auto")).choices("auto", "4b", "translator"); + args.add_argument("--precision").default_value(std::string("bf16")).choices("bf16", "fp16", "fp8", "nvfp4"); + args.add_argument("--smoke-test").default_value(false).implicit_value(true); + args.add_argument("--auto-generate").default_value(false).implicit_value(true); + args.add_argument("--exit-after-generation").default_value(false).implicit_value(true); + args.parse_args(argc, argv); + Flux2Config config; + config.model_dir = args.get("--model-dir"); + config.output_path = args.get("--output"); + config.ep_cache_dir = args.get("--ep-cache"); + config.ep_context_dir = args.get("--ep-context-dir"); + config.prompt = args.get("--prompt"); + config.precision = args.get("--precision"); + config.num_images = 1; + const auto encoder = args.get("--encoder"); + const bool translator = encoder == "translator" || + (encoder == "auto" && std::filesystem::is_regular_file( + config.model_dir / "text_encoder_translator/model.onnx")); + config.text_encoder = translator ? Flux2TextEncoder::Qwen3_06BTranslator : Flux2TextEncoder::Qwen3_4B; + config.weight_streaming_budget = "-1"; + if (!glfwInit()) + throw std::runtime_error("GLFW initialization failed"); + glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3); + glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 2); + window = glfwCreateWindow(1280, 850, "DIN Flux2 Studio", nullptr, nullptr); + if (!window) + throw std::runtime_error("OpenGL 3.2 window creation failed"); + glfwMakeContextCurrent(window); + glfwSwapInterval(1); + IMGUI_CHECKVERSION(); + ImGui::CreateContext(); + imgui = true; + ImGui::GetIO().IniFilename = nullptr; + ImGui::StyleColorsDark(); + glfw_backend = ImGui_ImplGlfw_InitForOpenGL(window, true); + gl_backend = ImGui_ImplOpenGL3_Init("#version 130"); + if (!glfw_backend || !gl_backend) + throw std::runtime_error("ImGui graphics initialization failed"); + result = Studio(config, args.get("--smoke-test"), args.get("--auto-generate"), + args.get("--exit-after-generation")); + } + catch (const std::exception& error) + { + std::cerr << error.what() << '\n'; + } + if (gl_backend) + ImGui_ImplOpenGL3_Shutdown(); + if (glfw_backend) + ImGui_ImplGlfw_Shutdown(); + if (imgui) + ImGui::DestroyContext(); + if (window) + glfwDestroyWindow(window); + glfwTerminate(); + return result; +} diff --git a/image_gen/flux2/flux2_vk_backend.cpp b/image_gen/flux2/flux2_vk_backend.cpp index 865f564..1298c89 100644 --- a/image_gen/flux2/flux2_vk_backend.cpp +++ b/image_gen/flux2/flux2_vk_backend.cpp @@ -27,6 +27,7 @@ #include "flux2.h" #include "flux2_runtime_context.h" +#include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" #include "vk_helper.h" @@ -438,17 +439,23 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor VkCommandBuffer cmd_postprocess, ComputePipelineResources& euler_shader, ComputePipelineResources& postprocess_shader, VulkanBuffer& hidden_states_buf, VulkanBuffer& transformer_output_buf, VulkanBuffer& decoder_latent_buf, StagingBuffer& staging, - bool encode_prompt) + bool encode_prompt, Flux2Telemetry& telemetry) { nvtx3::scoped_range nvtx_pipeline("run_pipeline"); + Ort::RunOptions transformer_options; + transformer_options.SetSyncStream(ort_stream); + transformer_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); + Ort::RunOptions run_options; run_options.SetSyncStream(ort_stream); run_options.AddConfigEntry("disable_synchronize_execution_providers", "1"); // --- Text Encoder --- + const int steps = static_cast(time_schedule.size()) - 1; if (encode_prompt) { + telemetry.Begin("encode", &Flux2Timings::encode_ms); nvtx3::scoped_range nvtx("text_encoder"); text_encoder_session.Run(run_options, text_encoder_io); } @@ -482,7 +489,8 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor sync_semaphore.wait_ort(ort_stream, timestep_done); } - for (int step = 0; step < FLOW_STEPS; ++step) + telemetry.Begin("denoise", &Flux2Timings::denoise_ms); + for (int step = 0; step < steps; ++step) { nvtx3::scoped_range nvtx_step("diffusion_step_" + std::to_string(step)); @@ -491,7 +499,7 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor { nvtx3::scoped_range nvtx_transformer("transformer"); - transformer_session.Run(run_options, transformer_io); + transformer_session.Run(transformer_options, transformer_io); } { @@ -527,7 +535,7 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT); - if (step + 1 < FLOW_STEPS) + if (step + 1 < steps) { vkCmdUpdateBuffer(cmd_euler[step], timestep_buf.buffer, 0, sizeof(float), &t_next); } @@ -538,8 +546,10 @@ uint64_t run_pipeline(Ort::Session& text_encoder_session, Ort::Session& transfor const uint64_t euler_done = sync_semaphore.submit_vulkan(vk, cmd_euler[step], transformer_done); sync_semaphore.wait_ort(ort_stream, euler_done); } + telemetry.Step(step + 1); } + telemetry.Begin("decode", &Flux2Timings::decode_ms); // --- Post-process Latents (Vulkan compute) --- { nvtx3::scoped_range nvtx_post("postprocess"); @@ -685,8 +695,8 @@ struct VkPipelineState Ort::Value decoded_image_tensor{nullptr}; VkCommandPool command_pool = VK_NULL_HANDLE; - std::array cmd_bufs_euler{}; - std::array cmd_bufs_time{}; + std::vector cmd_bufs_euler{}; + std::vector cmd_bufs_time{}; VkCommandBuffer cmd_upload_lat = VK_NULL_HANDLE; VkCommandBuffer cmd_download = VK_NULL_HANDLE; VkCommandBuffer cmd_postprocess = VK_NULL_HANDLE; @@ -705,10 +715,11 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi { throw std::runtime_error("Vulkan processing requires --provider trt-rtx"); } - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision); + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); state.use_cig = config.processing == Flux2ProcessingBackend::VkCig; const Flux2ModelCachePaths cache_paths = - MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "vk_cig" : "vk"); + MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "vk_cig" : "vk", config.text_encoder, + !config.weight_streaming_budget.empty()); std::cout << "Model dir: " << model_paths.base_dir.string() << "\n" << "CIG: " << (state.use_cig ? "enabled" : "disabled") << "\n" @@ -736,7 +747,7 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi { state.cig_external_compute_queue_data = state.vk->createCudaGraphicsInteropData(); state.graphics_interop = std::make_unique(*state.interop_api, state.trt_device, - state.cig_external_compute_queue_data); + state.cig_external_compute_queue_data); } // ----------------------------------------------------------------- @@ -746,6 +757,7 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi auto nvtx_scope_load = nvtx3::start_range("load_onnx_models"); din::common::EpContextOptions ep_context; ep_context.output_dir = config.ep_context_dir.string(); + ep_context.enable_cache = true; const std::string cache_dir = config.ep_cache_dir.string(); std::vector> graphics_ep_options; if (state.use_cig) @@ -767,13 +779,15 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // graphics_ep_options.emplace_back("nv_use_sync_gpu_allocator", "1"); } - auto make_profile = [&](std::string cache_subpath) + auto make_profile = [&](std::string cache_subpath, bool transformer = false) { din::common::ModelProfile profile; profile.cache_subpath = std::move(cache_subpath); profile.embed_ep_context = false; profile.enable_cuda_graph = !state.use_cig; profile.extra_ep_options = graphics_ep_options; + if (transformer && !config.weight_streaming_budget.empty()) + profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.weight_streaming_budget); return profile; }; @@ -908,7 +922,7 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi state.transformer_runner = std::make_unique( state.env, model_paths.transformer_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.transformer), &*state.sync_stream); + make_profile(cache_paths.transformer, true), &*state.sync_stream); std::cout << " Transformer loaded" << std::endl; state.vae_decoder_runner = std::make_unique( @@ -922,10 +936,12 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // ----------------------------------------------------------------- // Euler sampler - state.vk->device.allocateCommandBuffers(state.command_pool, FLOW_STEPS, state.cmd_bufs_euler.data()); + state.cmd_bufs_euler.resize(config.steps); + state.cmd_bufs_time.resize(config.steps); + state.vk->device.allocateCommandBuffers(state.command_pool, config.steps, state.cmd_bufs_euler.data()); // Time scheduler - state.vk->device.allocateCommandBuffers(state.command_pool, FLOW_STEPS, state.cmd_bufs_time.data()); + state.vk->device.allocateCommandBuffers(state.command_pool, config.steps, state.cmd_bufs_time.data()); VkCommandBuffer cmd_bufs[3]; state.vk->device.allocateCommandBuffers(state.command_pool, 3, cmd_bufs); @@ -1003,14 +1019,7 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // ----------------------------------------------------------------- // Time schedule (computed once, constant across images) // ----------------------------------------------------------------- - constexpr double MU = 2.291179894115571; - state.time_schedule.resize(FLOW_STEPS + 1); - for (int i = 0; i <= FLOW_STEPS; ++i) - { - double t = T_START - (T_START - T_END) * static_cast(i) / FLOW_STEPS; - state.time_schedule[i] = static_cast(std::exp(MU) / (std::exp(MU) + std::pow(1.0 / t - 1.0, 1.0))); - std::cout << " sigma[" << i << "] = " << state.time_schedule[i] << std::endl; - } + state.time_schedule = MakeFlux2Schedule(config.steps); // ----------------------------------------------------------------- // Load prompt tokens & upload to GPU (once for all images) @@ -1040,8 +1049,14 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi } } -static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed) +static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, const Flux2Progress& progress) { + Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, + [&] + { + signal_and_wait_on_host(*state.sync_stream); + VK_CHECK(vkQueueWaitIdle(state.vk->queue)); + }); initialize_latent(static_cast(state.staging.mapped), state.hidden_count, seed); const uint64_t vae_done = run_pipeline( @@ -1049,7 +1064,8 @@ static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed) *state.text_encoder_io, *state.transformer_io, *state.vae_decoder_io, state.time_schedule, state.timestep_buf, *state.vk, *state.sync_semaphore, *state.sync_stream, state.cmd_upload_lat, state.cmd_bufs_time.data(), state.cmd_bufs_euler.data(), state.cmd_postprocess, state.euler_shader, state.postprocess_shader, - state.hidden_states, state.transformer_output, state.decoder_latent, state.staging, !state.prompt_embeds_valid); + state.hidden_states, state.transformer_output, state.decoder_latent, state.staging, !state.prompt_embeds_valid, + telemetry); state.prompt_embeds_valid = true; vkResetCommandBuffer(state.cmd_download, 0); @@ -1069,6 +1085,7 @@ static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed) image.width = static_cast(IMAGE_WIDTH); image.data.resize(state.decoded_image_bytes / sizeof(float)); std::copy_n(static_cast(state.staging.mapped), image.data.size(), image.data.data()); + image.timings = telemetry.Finish(); return image; } @@ -1094,6 +1111,8 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline void SetPrompt(std::string prompt) override { + if (state_ && state_->prompt_embeds_valid && config_.prompt == prompt) + return; config_.prompt = std::move(prompt); if (!state_) { @@ -1101,7 +1120,8 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config_.model_dir, config_.precision); + const Flux2ModelPaths model_paths = + MakeFlux2ModelPaths(config_.model_dir, config_.precision, config_.text_encoder); const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); @@ -1121,17 +1141,33 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline state_->vk->endAndSubmitCommandBuffer(cmd, state_->command_pool); } - Flux2Image GenerateImage(unsigned int seed) override + const void* TransformerSessionIdentity() const override + { + return state_ ? static_cast(state_->transformer_runner->session) : nullptr; + } + + Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, + const Flux2GenerationOptions& options) override { if (!state_) { throw std::runtime_error("Flux2 Vulkan pipeline was not initialized"); } - return run_vk_image(*state_, seed); + const std::string budget = options.weight_streaming_budget.value_or(config_.weight_streaming_budget); + const std::string& previous = applied_budget_ ? *applied_budget_ : config_.weight_streaming_budget; + if (budget != previous) + { + const char* keys[] = {"nv_weight_streaming_budget"}; + const char* values[] = {budget.c_str()}; + state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); + applied_budget_ = budget; + } + return run_vk_image(*state_, seed, progress); } private: Flux2Config config_; + std::optional applied_budget_; Flux2RuntimeContext& runtime_; std::unique_ptr state_; }; diff --git a/image_gen/flux2/helpers/flux2.h b/image_gen/flux2/helpers/flux2.h index 9ef75e9..24b4be5 100644 --- a/image_gen/flux2/helpers/flux2.h +++ b/image_gen/flux2/helpers/flux2.h @@ -15,6 +15,7 @@ #include #include +#include "flux2_cli.h" #include "tokenizer.h" // ============================================================================ @@ -43,7 +44,8 @@ struct Flux2ModelCachePaths std::string vae_decoder; }; -inline Flux2ModelPaths MakeFlux2ModelPaths(std::filesystem::path model_dir, const std::string& precision) +inline Flux2ModelPaths MakeFlux2ModelPaths(std::filesystem::path model_dir, const std::string& precision, + Flux2TextEncoder encoder = Flux2TextEncoder::Qwen3_4B) { if (model_dir.empty()) { @@ -51,19 +53,21 @@ inline Flux2ModelPaths MakeFlux2ModelPaths(std::filesystem::path model_dir, cons } return { model_dir, - model_dir / "text_encoder/model.onnx", + model_dir / (encoder == Flux2TextEncoder::Qwen3_4B ? "text_encoder" : "text_encoder_translator") / "model.onnx", model_dir / ("transformer_" + precision) / "model.onnx", model_dir / "vae_decoder/model.onnx", model_dir / "tokenizer", }; } -inline Flux2ModelCachePaths MakeFlux2ModelCachePaths(const std::string& precision, const std::string& prefix = {}) +inline Flux2ModelCachePaths MakeFlux2ModelCachePaths(const std::string& precision, const std::string& prefix = {}, + Flux2TextEncoder encoder = Flux2TextEncoder::Qwen3_4B, + bool weight_streaming = false) { const std::string prefix_separator = prefix.empty() ? "" : prefix + "_"; return { - prefix_separator + "text_encoder", - prefix_separator + "transformer_" + precision, + prefix_separator + (encoder == Flux2TextEncoder::Qwen3_4B ? "text_encoder" : "text_encoder_translator"), + prefix_separator + "transformer_" + precision + (weight_streaming ? "_ws" : ""), prefix_separator + "vae_decoder", }; } @@ -114,12 +118,36 @@ struct Flux2TextEncoderInputs int64_t pad_token_id = 0; }; +inline std::string FormatFlux2Prompt(const std::string& prompt) +{ + return "<|im_start|>user\n" + prompt + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; +} + +inline std::vector MakeFlux2Schedule(int steps) +{ + if (steps < 1 || steps > 50) + throw std::invalid_argument("Denoise steps must be in [1, 50]"); + constexpr double m200 = 0.00016927 * IMAGE_SEQUENCE + 0.45666666; + constexpr double m10 = 8.73809524e-05 * IMAGE_SEQUENCE + 1.89833333; + constexpr double a = (m200 - m10) / 190.0; + const double mu = a * steps + m200 - 200.0 * a; + const double e = std::exp(mu); + std::vector schedule(steps + 1); + for (int i = 0; i < steps; ++i) + { + const double t = 1.0 - static_cast(i) / steps; + schedule[i] = static_cast(e / (e + 1.0 / t - 1.0)); + } + schedule.back() = 0.0f; + return schedule; +} + inline Flux2TextEncoderInputs TokenizeFlux2Prompt(const Flux2ModelPaths& model_paths, const std::string& prompt) { din::io::Tokenizer tokenizer((model_paths.tokenizer_dir / "tokenizer.json").string(), din::io::TokenizerFormat::Json); - const std::string chat_text = "<|im_start|>user\n" + prompt + "<|im_end|>"; + const std::string chat_text = FormatFlux2Prompt(prompt); Flux2TextEncoderInputs inputs; inputs.token_ids = tokenizer.Encode(chat_text, false); inputs.pad_token_id = tokenizer.TokenId("<|endoftext|>"); From 4c3658c14f5bdf5f3eb2320798834e66508e1e5a Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Wed, 23 Sep 2026 01:28:13 +0200 Subject: [PATCH 03/14] fix version --- cmake/onnxruntime_trt_rtx_ep.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/onnxruntime_trt_rtx_ep.cmake b/cmake/onnxruntime_trt_rtx_ep.cmake index ccd5109..d6668e4 100644 --- a/cmake/onnxruntime_trt_rtx_ep.cmake +++ b/cmake/onnxruntime_trt_rtx_ep.cmake @@ -9,7 +9,7 @@ set(ONNXRUNTIME_TRT_RTX_EP_GIT_REPOSITORY "https://github.com/NVIDIA/TensorRT-RTX-EP-ABI.git" CACHE STRING "Git repository for the ONNX Runtime TensorRT RTX Execution Provider") set(ONNXRUNTIME_TRT_RTX_EP_GIT_TAG - "v0.4.0" + "v0.5.0" CACHE STRING "Git tag, branch, or commit for the ONNX Runtime TensorRT RTX Execution Provider") set(ONNXRUNTIME_TRT_RTX_EP_GIT_HTTP_TOKEN "" From e79d51cac34a8555457c8a5daf1db468eafebc7c Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Wed, 23 Sep 2026 16:05:05 +0200 Subject: [PATCH 04/14] normalization based on exported json --- image_gen/flux2/CMakeLists.txt | 2 + image_gen/flux2/flux2_cuda_backend.cpp | 9 ++-- image_gen/flux2/flux2_dx_backend.cpp | 7 +-- image_gen/flux2/flux2_ui.cpp | 54 +++++++++++++++----- image_gen/flux2/flux2_vk_backend.cpp | 7 +-- image_gen/flux2/helpers/flux2.h | 49 +----------------- image_gen/flux2/model_export/export_flux2.py | 15 ++++++ image_gen/flux2/utils.cpp | 39 ++++++++++++++ image_gen/flux2/utils.h | 8 +++ 9 files changed, 120 insertions(+), 70 deletions(-) diff --git a/image_gen/flux2/CMakeLists.txt b/image_gen/flux2/CMakeLists.txt index 0641675..8689f14 100644 --- a/image_gen/flux2/CMakeLists.txt +++ b/image_gen/flux2/CMakeLists.txt @@ -19,6 +19,8 @@ target_include_directories(din_flux2_common target_link_libraries(din_flux2_common PUBLIC din_common_ort + PRIVATE + nlohmann_json::nlohmann_json ) set_target_properties(din_flux2_common PROPERTIES diff --git a/image_gen/flux2/flux2_cuda_backend.cpp b/image_gen/flux2/flux2_cuda_backend.cpp index 65ad090..4affbbf 100644 --- a/image_gen/flux2/flux2_cuda_backend.cpp +++ b/image_gen/flux2/flux2_cuda_backend.cpp @@ -308,7 +308,7 @@ class PostprocessStage constexpr int HW = static_cast(LATENT_HEIGHT * LATENT_WIDTH); unpack_latents_with_ids(hidden, scratch, LATENT_HEIGHT, LATENT_WIDTH, LATENT_CHANNELS); - denormalize_latents(scratch, BN_MEAN, BN_STD, LATENT_CHANNELS, HW); + denormalize_latents(scratch, bn_mean_->HostData(), bn_std_->HostData(), LATENT_CHANNELS, HW); unpatchify_latents(scratch, decoder_input, 1, LATENT_CHANNELS / PATCH_SIZE / PATCH_SIZE, LATENT_HEIGHT, LATENT_WIDTH, PATCH_SIZE, PATCH_SIZE); if (device_is_cuda) @@ -478,6 +478,8 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const ExecutionProviderMode provider_mode, SamplingBackend sampling_backend) { const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); + const auto latent_stats = LoadFlux2LatentStats( + model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision, "", config.text_encoder, !config.weight_streaming_budget.empty()); const bool trt_rtx_device = is_trt_rtx_device(ep_device); @@ -523,7 +525,6 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const } din::common::EpContextOptions ep_context; ep_context.output_dir = config.ep_context_dir.string(); - ep_context.enable_cache = true; const std::string cache_dir = config.ep_cache_dir.string(); auto make_profile = [&](std::string cache_subpath, bool transformer = false) @@ -623,8 +624,8 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state.token->HostData(), state.attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); - std::copy_n(BN_MEAN, LATENT_CHANNELS, state.bn_mean->HostData()); - std::copy_n(BN_STD, LATENT_CHANNELS, state.bn_std->HostData()); + std::copy_n(latent_stats.mean.data(), LATENT_CHANNELS, state.bn_mean->HostData()); + std::copy_n(latent_stats.std.data(), LATENT_CHANNELS, state.bn_std->HostData()); if (has_separate_binding) { diff --git a/image_gen/flux2/flux2_dx_backend.cpp b/image_gen/flux2/flux2_dx_backend.cpp index cb19bcc..b737619 100644 --- a/image_gen/flux2/flux2_dx_backend.cpp +++ b/image_gen/flux2/flux2_dx_backend.cpp @@ -459,6 +459,8 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi throw std::runtime_error("DirectX processing requires --provider trt-rtx"); } const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); + const auto latent_stats = LoadFlux2LatentStats( + model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "dx_cig" : "dx", config.text_encoder, !config.weight_streaming_budget.empty()); @@ -501,7 +503,6 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi std::cout << "=== Loading ONNX Models ===" << std::endl; din::common::EpContextOptions ep_context; ep_context.output_dir = config.ep_context_dir.string(); - ep_context.enable_cache = true; const std::string cache_dir = config.ep_cache_dir.string(); auto make_profile = [&](std::string cache_subpath, bool transformer = false) { @@ -601,8 +602,8 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi std::cout << "\n=== Uploading Static Inputs (D3D12) ===" << std::endl; state.dx.begin(); - state.dx.record_upload(state.bn_mean_buf, BN_MEAN, state.bn_mean_buf.size); - state.dx.record_upload(state.bn_std_buf, BN_STD, state.bn_std_buf.size); + state.dx.record_upload(state.bn_mean_buf, latent_stats.mean.data(), state.bn_mean_buf.size); + state.dx.record_upload(state.bn_std_buf, latent_stats.std.data(), state.bn_std_buf.size); { std::vector img_ids_cpu(shape_numel(img_ids_shape)); diff --git a/image_gen/flux2/flux2_ui.cpp b/image_gen/flux2/flux2_ui.cpp index ef31a9a..1d2b889 100644 --- a/image_gen/flux2/flux2_ui.cpp +++ b/image_gen/flux2/flux2_ui.cpp @@ -16,6 +16,7 @@ #include "flux2.h" #include "flux2_cli.h" +#include "utils.h" #include "io/image.h" #include #include @@ -223,7 +224,9 @@ namespace { const float v = img.data[c * pixels + p]; rgb[p * 3 + c] = - std::isfinite(v) ? static_cast(std::clamp(v * 0.5f + 0.5f, 0.0f, 1.0f) * 255.0f) : 0; + std::isfinite(v) + ? static_cast(std::lround(std::clamp(v * 127.5f + 127.5f, 0.0f, 255.0f))) + : 0; } glGenTextures(1, &result.texture); glBindTexture(GL_TEXTURE_2D, result.texture); @@ -233,13 +236,18 @@ namespace glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, img.width, img.height, 0, GL_RGB, GL_UNSIGNED_BYTE, rgb.data()); } + int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit_after) { Generator generator; Flux2Config config = initial; - std::array model{}, output{}; + std::array model{}, output{}, cache{}; + bool cache_changed = false; + const auto cache_utf8 = initial.ep_cache_dir.parent_path().u8string(); + std::snprintf(cache.data(), cache.size(), "%s", reinterpret_cast(cache_utf8.c_str())); std::array prompt{}; - std::snprintf(model.data(), model.size(), "%s", config.model_dir.string().c_str()); + const auto model_utf8 = config.model_dir.u8string(); + std::snprintf(model.data(), model.size(), "%s", reinterpret_cast(model_utf8.c_str())); std::snprintf(output.data(), output.size(), "%s", config.output_path.string().c_str()); std::snprintf(prompt.data(), prompt.size(), "%s", config.prompt.c_str()); int streaming_mode = 0; // Auto / Manual / Off; independent of session configuration. @@ -302,11 +310,12 @@ namespace ImGui::BeginChild("controls", ImVec2(430, 0), true); ImGui::BeginDisabled(busy); ImGui::InputText("Model root", model.data(), model.size()); + cache_changed |= ImGui::InputText("Cache root", cache.data(), cache.size()); ImGui::InputText("Save directory", output.data(), output.size()); ImGui::TextUnformatted("Prompt"); ImGui::InputTextMultiline("##prompt", prompt.data(), prompt.size(), ImVec2(-1, 100)); ImGui::RadioButton("Qwen3-4B", &encoder, 0); - const bool translator_exists = std::filesystem::is_regular_file(std::filesystem::path(model.data()) / + const bool translator_exists = std::filesystem::is_regular_file(std::filesystem::u8path(model.data()) / "text_encoder_translator/model.onnx"); ImGui::BeginDisabled(!translator_exists); ImGui::RadioButton("Qwen3-0.6B + translator", &encoder, 1); @@ -339,22 +348,29 @@ namespace ImGui::InputScalar("Seed", ImGuiDataType_U32, &config.seed); ImGui::SliderInt("Denoise steps", &config.steps, 1, 50); ImGui::SliderInt("Images", &count, 1, 16); + const bool paths_valid = model[0] && (!cache_changed || cache[0]); + ImGui::BeginDisabled(!paths_valid); bool generate = ImGui::Button("Generate", ImVec2(-1, 35)); bool random = ImGui::Button("Random sweep", ImVec2(-1, 30)); ImGui::EndDisabled(); - if (!busy && (generate || random || (auto_generate && !launched))) + ImGui::EndDisabled(); + if (!busy && paths_valid && (generate || random || (auto_generate && !launched))) { for (auto& result : results) if (result.texture) glDeleteTextures(1, &result.texture); results.clear(); selected = -1; - config.model_dir = model.data(); + config.model_dir = std::filesystem::u8path(model.data()); config.output_path = output.data(); config.prompt = prompt.data(); const auto cache_namespace = CacheNamespace(config.model_dir); - config.ep_cache_dir = initial.ep_cache_dir / cache_namespace; - config.ep_context_dir = initial.ep_context_dir / cache_namespace; + config.ep_cache_dir = (cache_changed + ? std::filesystem::u8path(cache.data()) / "runtime" + : initial.ep_cache_dir) / cache_namespace; + config.ep_context_dir = (cache_changed + ? std::filesystem::u8path(cache.data()) / "ep_context" + : initial.ep_context_dir) / cache_namespace; config.text_encoder = encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; config.precision = precisions[precision]; config.num_images = static_cast(count); @@ -437,10 +453,22 @@ int main(int argc, char** argv) try { argparse::ArgumentParser args("din_flux2_ui"); +#ifdef _WIN32 + const wchar_t* local_app_data = _wgetenv(L"LOCALAPPDATA"); + if (!local_app_data || !*local_app_data) + throw std::runtime_error("LOCALAPPDATA is not set; cannot determine the default cache directory"); + const auto cache_root = std::filesystem::path(local_app_data) / "din_deploy"; +#else + const auto cache_root = std::filesystem::path("/tmp/din_deploy/cache"); +#endif args.add_argument("--model-dir").default_value(DEFAULT_MODEL_BASE_PATH.string()); args.add_argument("--output").default_value(std::string{}); - args.add_argument("--ep-cache").default_value(Flux2Config{}.ep_cache_dir.string()); - args.add_argument("--ep-context-dir").default_value(Flux2Config{}.ep_context_dir.string()); + const auto runtime_cache_utf8 = (cache_root / "runtime").u8string(); + const auto context_cache_utf8 = (cache_root / "ep_context").u8string(); + args.add_argument("--ep-cache").default_value(std::string( + reinterpret_cast(runtime_cache_utf8.c_str()))); + args.add_argument("--ep-context-dir").default_value(std::string( + reinterpret_cast(context_cache_utf8.c_str()))); args.add_argument("--prompt").default_value(std::string(DEFAULT_PROMPT)); args.add_argument("--encoder").default_value(std::string("auto")).choices("auto", "4b", "translator"); args.add_argument("--precision").default_value(std::string("bf16")).choices("bf16", "fp16", "fp8", "nvfp4"); @@ -451,8 +479,10 @@ int main(int argc, char** argv) Flux2Config config; config.model_dir = args.get("--model-dir"); config.output_path = args.get("--output"); - config.ep_cache_dir = args.get("--ep-cache"); - config.ep_context_dir = args.get("--ep-context-dir"); + config.ep_cache_dir = std::filesystem::u8path(args.get("--ep-cache")); + config.ep_context_dir = std::filesystem::u8path(args.get("--ep-context-dir")); + std::filesystem::create_directories(config.ep_cache_dir); + std::filesystem::create_directories(config.ep_context_dir); config.prompt = args.get("--prompt"); config.precision = args.get("--precision"); config.num_images = 1; diff --git a/image_gen/flux2/flux2_vk_backend.cpp b/image_gen/flux2/flux2_vk_backend.cpp index 1298c89..0b2eb66 100644 --- a/image_gen/flux2/flux2_vk_backend.cpp +++ b/image_gen/flux2/flux2_vk_backend.cpp @@ -716,6 +716,8 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi throw std::runtime_error("Vulkan processing requires --provider trt-rtx"); } const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); + const auto latent_stats = LoadFlux2LatentStats( + model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); state.use_cig = config.processing == Flux2ProcessingBackend::VkCig; const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "vk_cig" : "vk", config.text_encoder, @@ -757,7 +759,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi auto nvtx_scope_load = nvtx3::start_range("load_onnx_models"); din::common::EpContextOptions ep_context; ep_context.output_dir = config.ep_context_dir.string(); - ep_context.enable_cache = true; const std::string cache_dir = config.ep_cache_dir.string(); std::vector> graphics_ep_options; if (state.use_cig) @@ -884,14 +885,14 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi { VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - memcpy(state.staging.mapped, BN_MEAN, LATENT_CHANNELS * sizeof(float)); + memcpy(state.staging.mapped, latent_stats.mean.data(), LATENT_CHANNELS * sizeof(float)); VkBufferCopy region = {0, 0, state.bn_mean_buf.size}; vkCmdCopyBuffer(cmd, state.staging.buffer, state.bn_mean_buf.buffer, 1, ®ion); state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); } { VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - memcpy(state.staging.mapped, BN_STD, LATENT_CHANNELS * sizeof(float)); + memcpy(state.staging.mapped, latent_stats.std.data(), LATENT_CHANNELS * sizeof(float)); VkBufferCopy region = {0, 0, state.bn_std_buf.size}; vkCmdCopyBuffer(cmd, state.staging.buffer, state.bn_std_buf.buffer, 1, ®ion); state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); diff --git a/image_gen/flux2/helpers/flux2.h b/image_gen/flux2/helpers/flux2.h index 24b4be5..252ebd1 100644 --- a/image_gen/flux2/helpers/flux2.h +++ b/image_gen/flux2/helpers/flux2.h @@ -130,7 +130,7 @@ inline std::vector MakeFlux2Schedule(int steps) constexpr double m200 = 0.00016927 * IMAGE_SEQUENCE + 0.45666666; constexpr double m10 = 8.73809524e-05 * IMAGE_SEQUENCE + 1.89833333; constexpr double a = (m200 - m10) / 190.0; - const double mu = a * steps + m200 - 200.0 * a; + const double mu = IMAGE_SEQUENCE > 4300 ? m200 : a * steps + m200 - 200.0 * a; const double e = std::exp(mu); std::vector schedule(steps + 1); for (int i = 0; i < steps; ++i) @@ -268,50 +268,3 @@ inline void stats(const std::vector& vec) { stats(vec.data(), vec.size()); } - -// ============================================================================ -// Batch-norm parameters (exported from vae.bn) -// ============================================================================ - -static const float BN_STD[128] = { - 1.8046875, 1.7734375, 1.7890625, 1.7890625, 1.7734375, 1.7578125, 1.7578125, 1.75, 1.734375, 1.734375, - 1.734375, 1.734375, 1.859375, 1.8515625, 1.859375, 1.859375, 1.7578125, 1.7578125, 1.7578125, 1.7578125, - 1.734375, 1.7421875, 1.734375, 1.7421875, 1.734375, 1.7265625, 1.7421875, 1.734375, 1.7578125, 1.75, - 1.7578125, 1.75, 1.84375, 1.84375, 1.8515625, 1.8515625, 1.828125, 1.78125, 1.7890625, 1.796875, - 1.8046875, 1.7890625, 1.765625, 1.765625, 1.7734375, 1.765625, 1.796875, 1.7890625, 1.7578125, 1.7421875, - 1.75, 1.734375, 1.796875, 1.796875, 1.796875, 1.796875, 1.7734375, 1.7734375, 1.7734375, 1.7734375, - 1.7578125, 1.7421875, 1.765625, 1.75, 1.7734375, 1.7734375, 1.78125, 1.78125, 1.8125, 1.8046875, - 1.8046875, 1.8046875, 1.7578125, 1.734375, 1.7578125, 1.75, 1.7578125, 1.75, 1.765625, 1.765625, - 1.765625, 1.7578125, 1.78125, 1.7734375, 1.7890625, 1.796875, 1.7890625, 1.7734375, 1.765625, 1.765625, - 1.7578125, 1.7578125, 1.7734375, 1.765625, 1.78125, 1.78125, 1.71875, 1.703125, 1.7265625, 1.7109375, - 1.78125, 1.7578125, 1.75, 1.7421875, 1.7578125, 1.75, 1.7578125, 1.734375, 1.765625, 1.7578125, - 1.7421875, 1.7421875, 1.75, 1.7578125, 1.765625, 1.765625, 1.7265625, 1.7109375, 1.734375, 1.7265625, - 1.765625, 1.765625, 1.7734375, 1.765625, 1.7421875, 1.7421875, 1.7578125, 1.7578125}; - -static const float BN_MEAN[128] = { - -0.0673828125f, -0.0712890625f, -0.0751953125f, -0.07470703125f, 0.0223388671875f, - 0.0179443359375f, 0.01422119140625f, 0.018310546875f, -6.29425048828125e-05f, -0.006256103515625f, - -0.000209808349609375f, -0.003143310546875f, -0.0272216796875f, -0.028076171875f, -0.027587890625f, - -0.029052734375f, -0.07666015625f, -0.0673828125f, -0.09033203125f, -0.08935546875f, - 0.016845703125f, 0.01519775390625f, 0.00787353515625f, 0.00860595703125f, 0.00836181640625f, - 0.0015411376953125f, 0.0002574920654296875f, -0.0042724609375f, -0.0439453125f, -0.0419921875f, - -0.043701171875f, -0.043212890625f, -0.01025390625f, -0.01318359375f, -0.006622314453125f, - -0.0047607421875f, -0.031005859375f, -0.030517578125f, -0.0279541015625f, -0.0179443359375f, - 0.003021240234375f, 0.00150299072265625f, 0.0125732421875f, 0.01446533203125f, 0.03466796875f, - 0.03369140625f, 0.03369140625f, 0.0283203125f, 0.001983642578125f, 0.004730224609375f, - 0.004669189453125f, 0.004974365234375f, 0.01226806640625f, 0.00811767578125f, 0.008056640625f, - 0.01458740234375f, 0.06787109375f, 0.06787109375f, 0.07666015625f, 0.0732421875f, - -0.046142578125f, -0.04736328125f, -0.039306640625f, -0.051025390625f, -0.052734375f, - -0.0478515625f, -0.047119140625f, -0.0517578125f, -0.03173828125f, -0.03173828125f, - -0.034423828125f, -0.0281982421875f, 0.051025390625f, 0.04443359375f, 0.057861328125f, - 0.0458984375f, -0.041259765625f, -0.0458984375f, -0.048828125f, -0.046630859375f, - -0.00885009765625f, -0.0106201171875f, -0.0087890625f, -0.004608154296875f, -0.03759765625f, - -0.043212890625f, -0.04345703125f, -0.0498046875f, 0.0118408203125f, 0.0166015625f, - 0.020263671875f, 0.0279541015625f, 0.01129150390625f, 0.01287841796875f, 0.001556396484375f, - 0.00714111328125f, -0.01177978515625f, -0.00183868408203125f, -0.01416015625f, -0.00537109375f, - -0.00909423828125f, -0.0137939453125f, -0.01446533203125f, -0.0186767578125f, 0.0322265625f, - 0.030517578125f, 0.02587890625f, 0.0299072265625f, 0.053955078125f, 0.0615234375f, - 0.049560546875f, 0.05908203125f, -0.051025390625f, -0.060302734375f, -0.0478515625f, - -0.052490234375f, -0.022705078125f, -0.0274658203125f, -0.015380859375f, -0.0255126953125f, - -0.05712890625f, -0.056396484375f, -0.0517578125f, -0.049560546875f, 0.0115966796875f, - 0.00543212890625f, 0.016357421875f, 0.0103759765625f}; diff --git a/image_gen/flux2/model_export/export_flux2.py b/image_gen/flux2/model_export/export_flux2.py index 3163192..c6a7e5a 100644 --- a/image_gen/flux2/model_export/export_flux2.py +++ b/image_gen/flux2/model_export/export_flux2.py @@ -303,6 +303,21 @@ def forward(self, latent_sample, return_dict=False): opset=opset, ) print(" vae_decoder exported.") + # Denormalization happens outside the decoder ONNX graph, in patchified channel order. + # Use learned BN buffers, not optional config.latents_mean/config.latents_std. + mean = vae.bn.running_mean + std = torch.sqrt(vae.bn.running_var + 1e-4) + expected_channels = vae_latent_channels * patch_size[0] * patch_size[1] + if mean.numel() != expected_channels or std.numel() != expected_channels: + raise ValueError("VAE latent statistics have an unexpected channel count") + if not torch.isfinite(mean).all() or not torch.isfinite(std).all() or not (std > 0).all(): + raise ValueError("VAE latent statistics must be finite, with positive standard deviations") + stats_path = output_path / "vae_decoder" / "latent_stats.json" + stats_path.write_text( + json.dumps({"bn_mean": mean.float().cpu().tolist(), "bn_std": std.float().cpu().tolist()}, allow_nan=False), + encoding="utf-8", + ) + print(f" VAE latent statistics exported to {stats_path}.") def resolve_model_snapshot(model_name: str, local_files_only: bool) -> Path: diff --git a/image_gen/flux2/utils.cpp b/image_gen/flux2/utils.cpp index 86b4343..156a2ec 100644 --- a/image_gen/flux2/utils.cpp +++ b/image_gen/flux2/utils.cpp @@ -19,6 +19,45 @@ #include #include #include +#include +#include +#include + +Flux2LatentStats LoadFlux2LatentStats(const std::filesystem::path& path, size_t channels) +{ + std::ifstream file(path); + if (!file) + throw std::runtime_error("Missing VAE latent statistics: " + path.string() + + ". Re-export the VAE decoder."); + try + { + const auto document = nlohmann::json::parse(file); + const auto read = [&](const char* key, bool positive) + { + const auto& values = document.at(key); + if (!values.is_array() || values.size() != channels) + throw std::runtime_error(std::string(key) + " must contain " + + std::to_string(channels) + " values"); + std::vector result; + result.reserve(channels); + for (const auto& value : values) + { + if (!value.is_number()) + throw std::runtime_error(std::string(key) + " must contain only numbers"); + const float number = value.get(); + if (!std::isfinite(number) || (positive && number <= 0.0f)) + throw std::runtime_error(std::string(key) + " contains an invalid value"); + result.push_back(number); + } + return result; + }; + return {read("bn_mean", false), read("bn_std", true)}; + } + catch (const std::exception& error) + { + throw std::runtime_error("Invalid VAE latent statistics in " + path.string() + ": " + error.what()); + } +} #ifdef _WIN32 #include diff --git a/image_gen/flux2/utils.h b/image_gen/flux2/utils.h index 2d3a174..d4aff19 100644 --- a/image_gen/flux2/utils.h +++ b/image_gen/flux2/utils.h @@ -42,6 +42,14 @@ OrtFileString toOrtFileString(const std::filesystem::path& path); std::filesystem::path get_executable_path(); std::filesystem::path get_executable_parent_path(); +struct Flux2LatentStats +{ + std::vector mean; + std::vector std; +}; + +Flux2LatentStats LoadFlux2LatentStats(const std::filesystem::path& path, size_t channels); + void flush_ort_stream(Ort::SyncStream& stream); void copy_tensors(Ort::Env& env, std::initializer_list src, std::initializer_list dst, Ort::SyncStream& stream); From 65e376f965ea0b73b9af7f0d720391207c24ecfd Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Wed, 23 Sep 2026 16:12:31 +0200 Subject: [PATCH 05/14] remove default model paths --- image_gen/flux2/flux2_cli.cpp | 3 +-- image_gen/flux2/flux2_ui.cpp | 9 ++++++--- image_gen/flux2/helpers/flux2.h | 8 +------- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/image_gen/flux2/flux2_cli.cpp b/image_gen/flux2/flux2_cli.cpp index c40571c..6bec33c 100644 --- a/image_gen/flux2/flux2_cli.cpp +++ b/image_gen/flux2/flux2_cli.cpp @@ -138,7 +138,6 @@ void validate_config(const Flux2Config& config) Flux2Config parse_args(int argc, char* argv[]) { Flux2Config config; - config.model_dir = DEFAULT_MODEL_BASE_PATH; config.prompt = DEFAULT_PROMPT; argparse::ArgumentParser parser("din_flux2"); @@ -154,7 +153,7 @@ Flux2Config parse_args(int argc, char* argv[]) .metavar("cpu|trt-rtx") .help("Select the ONNX Runtime execution provider."); parser.add_argument("--model-dir") - .default_value(config.model_dir.string()) + .required() .nargs(1) .metavar("PATH") .help("Root directory with shared Flux2 ONNX artifacts and transformer_ directories."); diff --git a/image_gen/flux2/flux2_ui.cpp b/image_gen/flux2/flux2_ui.cpp index 1d2b889..cb5a819 100644 --- a/image_gen/flux2/flux2_ui.cpp +++ b/image_gen/flux2/flux2_ui.cpp @@ -315,7 +315,8 @@ namespace ImGui::TextUnformatted("Prompt"); ImGui::InputTextMultiline("##prompt", prompt.data(), prompt.size(), ImVec2(-1, 100)); ImGui::RadioButton("Qwen3-4B", &encoder, 0); - const bool translator_exists = std::filesystem::is_regular_file(std::filesystem::u8path(model.data()) / + const bool translator_exists = model[0] && std::filesystem::is_regular_file( + std::filesystem::u8path(model.data()) / "text_encoder_translator/model.onnx"); ImGui::BeginDisabled(!translator_exists); ImGui::RadioButton("Qwen3-0.6B + translator", &encoder, 1); @@ -461,7 +462,7 @@ int main(int argc, char** argv) #else const auto cache_root = std::filesystem::path("/tmp/din_deploy/cache"); #endif - args.add_argument("--model-dir").default_value(DEFAULT_MODEL_BASE_PATH.string()); + args.add_argument("--model-dir").default_value(std::string{}); args.add_argument("--output").default_value(std::string{}); const auto runtime_cache_utf8 = (cache_root / "runtime").u8string(); const auto context_cache_utf8 = (cache_root / "ep_context").u8string(); @@ -478,6 +479,8 @@ int main(int argc, char** argv) args.parse_args(argc, argv); Flux2Config config; config.model_dir = args.get("--model-dir"); + if (config.model_dir.empty() && args.get("--auto-generate")) + throw std::runtime_error("--auto-generate requires --model-dir"); config.output_path = args.get("--output"); config.ep_cache_dir = std::filesystem::u8path(args.get("--ep-cache")); config.ep_context_dir = std::filesystem::u8path(args.get("--ep-context-dir")); @@ -488,7 +491,7 @@ int main(int argc, char** argv) config.num_images = 1; const auto encoder = args.get("--encoder"); const bool translator = encoder == "translator" || - (encoder == "auto" && std::filesystem::is_regular_file( + (encoder == "auto" && !config.model_dir.empty() && std::filesystem::is_regular_file( config.model_dir / "text_encoder_translator/model.onnx")); config.text_encoder = translator ? Flux2TextEncoder::Qwen3_06BTranslator : Flux2TextEncoder::Qwen3_4B; config.weight_streaming_budget = "-1"; diff --git a/image_gen/flux2/helpers/flux2.h b/image_gen/flux2/helpers/flux2.h index 252ebd1..68c5fc0 100644 --- a/image_gen/flux2/helpers/flux2.h +++ b/image_gen/flux2/helpers/flux2.h @@ -22,12 +22,6 @@ // Configuration // ============================================================================ -#ifdef _WIN32 -static const std::filesystem::path DEFAULT_MODEL_BASE_PATH = "S:/din_deploy_artifacts/FLUX.2-klein-4B-onnx"; -#else -static const std::filesystem::path DEFAULT_MODEL_BASE_PATH = "/mnt/share/onnx/FLUX.2-klein-4B-onnx/"; -#endif - struct Flux2ModelPaths { std::filesystem::path base_dir; @@ -49,7 +43,7 @@ inline Flux2ModelPaths MakeFlux2ModelPaths(std::filesystem::path model_dir, cons { if (model_dir.empty()) { - model_dir = DEFAULT_MODEL_BASE_PATH; + throw std::invalid_argument("Model directory must be specified"); } return { model_dir, From cc56996c9c1bd2dce5d6cdce9cc0efbaf525f5bc Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Wed, 23 Sep 2026 16:18:53 +0200 Subject: [PATCH 06/14] remove unnecessary changes --- common/ort_session.cpp | 7 ++----- common/ort_session.h | 1 - 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/common/ort_session.cpp b/common/ort_session.cpp index 149e797..81e6e66 100644 --- a/common/ort_session.cpp +++ b/common/ort_session.cpp @@ -487,10 +487,7 @@ void AppendTensorRTRTXEP(Ort::Env& env, Ort::SessionOptions& options, const std: const auto stem = fs::path(model_path).stem().string(); const auto cache_path = (fs::path(cache_dir) / (profile.cache_subpath.empty() ? stem : profile.cache_subpath)).string(); - if (ep_context.enable_cache) - { - ep_options.Add("nv_runtime_cache_path", cache_path.c_str()); - } + ep_options.Add("nv_runtime_cache_path", cache_path.c_str()); if (!profile.min_shapes.empty()) { ep_options.Add("nv_profile_min_shapes", profile.min_shapes.c_str()); @@ -836,7 +833,7 @@ OrtRunner::OrtRunner(Ort::Env& env_in, const std::string& model_path, const std: compute_stream = &owned_compute_stream_; } - if (!profile.skip_compile && ep_context.enable_cache) + if (!profile.skip_compile) { [[maybe_unused]] din::common::nvtx_scoped_range range{"load_or_compile_ep_context"}; session_model_path = CompileEpContextModel(env, model_path, cache_dir, ep_context, profile); diff --git a/common/ort_session.h b/common/ort_session.h index 5491255..3e0a09f 100644 --- a/common/ort_session.h +++ b/common/ort_session.h @@ -77,7 +77,6 @@ QueryCudaGraphicsInteropSharedMemoryInfo(int cuda_device_ordinal, bool high_prio struct EpContextOptions { std::string output_dir; - bool enable_cache = true; }; struct ModelProfile From 448b6905424950591b08920960b6d7aab432fd00 Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Tue, 6 Oct 2026 14:02:15 +0200 Subject: [PATCH 07/14] miniaudio bump version to work on windows on arm --- cmake/miniaudio.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/miniaudio.cmake b/cmake/miniaudio.cmake index 398643a..e0844dc 100644 --- a/cmake/miniaudio.cmake +++ b/cmake/miniaudio.cmake @@ -3,7 +3,7 @@ include_guard(GLOBAL) include(FetchContent) FetchContent_Declare(din_miniaudio GIT_REPOSITORY https://github.com/mackron/miniaudio.git - GIT_TAG f40cf03f80cdb7e741d43e53b7e706e8c1394bcf + GIT_TAG "0.11.25" SOURCE_SUBDIR header-only) FetchContent_MakeAvailable(din_miniaudio) From 0797fc5991629d3af406af4134e7d2a4d0404fe3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Tue, 6 Oct 2026 15:19:27 +0200 Subject: [PATCH 08/14] formatting --- cmake/miniaudio.cmake | 2 +- image_gen/flux2/flux2_cli.cpp | 8 +- image_gen/flux2/flux2_cuda_backend.cpp | 4 +- image_gen/flux2/flux2_dx_backend.cpp | 4 +- image_gen/flux2/flux2_pipeline.cpp | 6 +- image_gen/flux2/flux2_telemetry.h | 8 +- image_gen/flux2/flux2_ui.cpp | 758 ++++++++++++------------- image_gen/flux2/flux2_vk_backend.cpp | 6 +- image_gen/flux2/utils.cpp | 9 +- 9 files changed, 397 insertions(+), 408 deletions(-) diff --git a/cmake/miniaudio.cmake b/cmake/miniaudio.cmake index e0844dc..0c2cf2e 100644 --- a/cmake/miniaudio.cmake +++ b/cmake/miniaudio.cmake @@ -3,7 +3,7 @@ include_guard(GLOBAL) include(FetchContent) FetchContent_Declare(din_miniaudio GIT_REPOSITORY https://github.com/mackron/miniaudio.git - GIT_TAG "0.11.25" + GIT_TAG "0.11.25" SOURCE_SUBDIR header-only) FetchContent_MakeAvailable(din_miniaudio) diff --git a/image_gen/flux2/flux2_cli.cpp b/image_gen/flux2/flux2_cli.cpp index 6bec33c..83bd212 100644 --- a/image_gen/flux2/flux2_cli.cpp +++ b/image_gen/flux2/flux2_cli.cpp @@ -186,8 +186,8 @@ Flux2Config parse_args(int argc, char* argv[]) parser.add_argument("--encoder").default_value(std::string("4b")).help("4b or translator."); parser.add_argument("--steps").default_value(std::string("4")).help("Denoise steps, 1 through 50."); parser.add_argument("--ws") - .default_value(std::string("off")) - .help("Transformer resident budget: off or 0% through 100%."); + .default_value(std::string("off")) + .help("Transformer resident budget: off or 0% through 100%."); parser.add_argument("--num-images") .default_value(std::to_string(config.num_images)) .nargs(1) @@ -288,8 +288,8 @@ int main(int argc, char* argv[]) Flux2Image image = pipeline->GenerateImage(current_seed); const auto& times = image.timings; std::cout << "Timings (ms): encode=" << times.encode_ms << " rng=" << times.rng_ms - << " denoise=" << times.denoise_ms << " decode=" << times.decode_ms << " total=" << times.total_ms - << std::endl; + << " denoise=" << times.denoise_ms << " decode=" << times.decode_ms << " total=" << times.total_ms + << std::endl; const auto image_end = std::chrono::steady_clock::now(); save_image(make_output_path(config, image_index), image); const std::chrono::duration image_duration = image_end - image_start; diff --git a/image_gen/flux2/flux2_cuda_backend.cpp b/image_gen/flux2/flux2_cuda_backend.cpp index 4affbbf..b1c1f74 100644 --- a/image_gen/flux2/flux2_cuda_backend.cpp +++ b/image_gen/flux2/flux2_cuda_backend.cpp @@ -478,8 +478,8 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const ExecutionProviderMode provider_mode, SamplingBackend sampling_backend) { const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); - const auto latent_stats = LoadFlux2LatentStats( - model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); + const auto latent_stats = + LoadFlux2LatentStats(model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision, "", config.text_encoder, !config.weight_streaming_budget.empty()); const bool trt_rtx_device = is_trt_rtx_device(ep_device); diff --git a/image_gen/flux2/flux2_dx_backend.cpp b/image_gen/flux2/flux2_dx_backend.cpp index b737619..ae06dbe 100644 --- a/image_gen/flux2/flux2_dx_backend.cpp +++ b/image_gen/flux2/flux2_dx_backend.cpp @@ -459,8 +459,8 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi throw std::runtime_error("DirectX processing requires --provider trt-rtx"); } const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); - const auto latent_stats = LoadFlux2LatentStats( - model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); + const auto latent_stats = + LoadFlux2LatentStats(model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "dx_cig" : "dx", config.text_encoder, !config.weight_streaming_budget.empty()); diff --git a/image_gen/flux2/flux2_pipeline.cpp b/image_gen/flux2/flux2_pipeline.cpp index 6f2b247..9c4efe2 100644 --- a/image_gen/flux2/flux2_pipeline.cpp +++ b/image_gen/flux2/flux2_pipeline.cpp @@ -11,10 +11,8 @@ namespace { Flux2RuntimeContext& GetProcessRuntime(Flux2ExecutionProvider provider) { - static Ort::Env env{ - std::getenv("DIN_FLUX2_VERBOSE") ? ORT_LOGGING_LEVEL_VERBOSE : ORT_LOGGING_LEVEL_WARNING, - "Flux2" - }; + static Ort::Env env{std::getenv("DIN_FLUX2_VERBOSE") ? ORT_LOGGING_LEVEL_VERBOSE : ORT_LOGGING_LEVEL_WARNING, + "Flux2"}; static Flux2RuntimeContext runtime{env}; static std::once_flag register_trt_once; diff --git a/image_gen/flux2/flux2_telemetry.h b/image_gen/flux2/flux2_telemetry.h index 7c623fb..c985f4b 100644 --- a/image_gen/flux2/flux2_telemetry.h +++ b/image_gen/flux2/flux2_telemetry.h @@ -14,10 +14,10 @@ class Flux2Telemetry public: Flux2Telemetry(int steps, Flux2Progress progress, std::function synchronize) : steps_(steps) - , progress_(std::move(progress)) - , synchronize_(std::move(synchronize)) - , start_(Clock::now()) - , stage_start_(start_) + , progress_(std::move(progress)) + , synchronize_(std::move(synchronize)) + , start_(Clock::now()) + , stage_start_(start_) { Report("rng", 0); } diff --git a/image_gen/flux2/flux2_ui.cpp b/image_gen/flux2/flux2_ui.cpp index cb5a819..1ba925c 100644 --- a/image_gen/flux2/flux2_ui.cpp +++ b/image_gen/flux2/flux2_ui.cpp @@ -16,8 +16,8 @@ #include "flux2.h" #include "flux2_cli.h" -#include "utils.h" #include "io/image.h" +#include "utils.h" #include #include #include @@ -26,425 +26,416 @@ namespace { - // Stable per-root namespace prevents a different model directory from reusing stale engines. - std::string CacheNamespace(const std::filesystem::path& root) +// Stable per-root namespace prevents a different model directory from reusing stale engines. +std::string CacheNamespace(const std::filesystem::path& root) +{ + const auto path = std::filesystem::weakly_canonical(root).generic_u8string(); + uint64_t hash = 14695981039346656037ull; + for (const auto c : path) { - const auto path = std::filesystem::weakly_canonical(root).generic_u8string(); - uint64_t hash = 14695981039346656037ull; - for (const auto c : path) - { - hash ^= static_cast(c); - hash *= 1099511628211ull; - } - return "root_" + std::to_string(hash); + hash ^= static_cast(c); + hash *= 1099511628211ull; } + return "root_" + std::to_string(hash); +} - struct Request - { - Flux2Config config; - Flux2GenerationOptions generation; - bool random_seeds = false; - }; +struct Request +{ + Flux2Config config; + Flux2GenerationOptions generation; + bool random_seeds = false; +}; - struct Result +struct Result +{ + Flux2Image image; + unsigned int seed = 0; + GLuint texture = 0; +}; + +// Only the worker touches the pipeline; all graphics resources belong to the UI thread. +class Generator +{ +public: + Generator() + : worker_( + [this] + { + Run(); + }) { - Flux2Image image; - unsigned int seed = 0; - GLuint texture = 0; - }; + } - // Only the worker touches the pipeline; all graphics resources belong to the UI thread. - class Generator + ~Generator() { - public: - Generator() - : worker_( - [this] - { - Run(); - }) { + std::lock_guard lock(mutex); + closing_ = true; } + cv_.notify_one(); + worker_.join(); + } - ~Generator() - { - { - std::lock_guard lock(mutex); - closing_ = true; - } - cv_.notify_one(); - worker_.join(); - } + void Submit(Request request) + { + std::lock_guard lock(mutex); + if (busy) + return; + busy = true; + failed = false; + status = "Loading models (first compilation can take several minutes)"; + completed = 0; + total = request.config.steps * static_cast(request.config.num_images); + request_ = std::move(request); + cv_.notify_one(); + } - void Submit(Request request) - { - std::lock_guard lock(mutex); - if (busy) - return; - busy = true; - failed = false; - status = "Loading models (first compilation can take several minutes)"; - completed = 0; - total = request.config.steps * static_cast(request.config.num_images); - request_ = std::move(request); - cv_.notify_one(); - } + std::mutex mutex; + bool busy = false, failed = false; + std::string status = "Ready"; + int completed = 0, total = 1; + size_t used_mb = 0, total_mb = 0; + std::vector ready; - std::mutex mutex; - bool busy = false, failed = false; - std::string status = "Ready"; - int completed = 0, total = 1; - size_t used_mb = 0, total_mb = 0; - std::vector ready; +private: + static bool SameModels(const Flux2Config& a, const Flux2Config& b) + { + return a.model_dir == b.model_dir && a.processing == b.processing && a.provider == b.provider && + a.precision == b.precision && a.text_encoder == b.text_encoder && a.steps == b.steps && + a.weight_streaming_budget == b.weight_streaming_budget && a.ep_cache_dir == b.ep_cache_dir && + a.ep_context_dir == b.ep_context_dir; + } - private: - static bool SameModels(const Flux2Config& a, const Flux2Config& b) + void UpdateMemory() + { + size_t free = 0, total_bytes = 0; + if (cudaMemGetInfo(&free, &total_bytes) == cudaSuccess) { - return a.model_dir == b.model_dir && a.processing == b.processing && a.provider == b.provider && - a.precision == b.precision && a.text_encoder == b.text_encoder && a.steps == b.steps && - a.weight_streaming_budget == b.weight_streaming_budget && - a.ep_cache_dir == b.ep_cache_dir && a.ep_context_dir == b.ep_context_dir; + std::lock_guard lock(mutex); + used_mb = (total_bytes - free) / (1024 * 1024); + total_mb = total_bytes / (1024 * 1024); } + else + cudaGetLastError(); + } - void UpdateMemory() + void Run() + { + std::unique_ptr pipeline; + std::optional loaded; + std::mt19937 random(std::random_device{}()); + for (;;) { - size_t free = 0, total_bytes = 0; - if (cudaMemGetInfo(&free, &total_bytes) == cudaSuccess) + std::unique_lock lock(mutex); + cv_.wait_for(lock, std::chrono::milliseconds(500), + [&] + { + return closing_ || request_.has_value(); + }); + if (closing_) + break; + if (!request_) { - std::lock_guard lock(mutex); - used_mb = (total_bytes - free) / (1024 * 1024); - total_mb = total_bytes / (1024 * 1024); + lock.unlock(); + if (loaded && loaded->provider == Flux2ExecutionProvider::TrtRtx) + UpdateMemory(); + continue; } - else - cudaGetLastError(); - } - - void Run() - { - std::unique_ptr pipeline; - std::optional loaded; - std::mt19937 random(std::random_device{}()); - for (;;) + Request request = std::move(*request_); + request_.reset(); + lock.unlock(); + try { - std::unique_lock lock(mutex); - cv_.wait_for(lock, std::chrono::milliseconds(500), - [&] - { - return closing_ || request_.has_value(); - }); - if (closing_) - break; - if (!request_) + const auto& config = request.config; + ValidateFlux2Config(config); + if (!loaded || !SameModels(*loaded, config)) { - lock.unlock(); - if (loaded && loaded->provider == Flux2ExecutionProvider::TrtRtx) - UpdateMemory(); - continue; + pipeline.reset(); // Release old engines before allocating the replacement. + loaded.reset(); + pipeline = CreateFlux2Pipeline(config); + pipeline->Initialize(); + loaded = config; } - Request request = std::move(*request_); - request_.reset(); - lock.unlock(); - try + pipeline->SetPrompt(config.prompt); + for (unsigned int i = 0; i < config.num_images; ++i) { - const auto& config = request.config; - ValidateFlux2Config(config); - if (!loaded || !SameModels(*loaded, config)) - { - pipeline.reset(); // Release old engines before allocating the replacement. - loaded.reset(); - pipeline = CreateFlux2Pipeline(config); - pipeline->Initialize(); - loaded = config; - } - pipeline->SetPrompt(config.prompt); - for (unsigned int i = 0; i < config.num_images; ++i) - { - const unsigned int seed = request.random_seeds ? random() : config.seed + i; - auto image = pipeline->GenerateImage( - seed, - [&](const char* stage, int step, int) - { - { - std::lock_guard guard(mutex); - completed = static_cast(i) * config.steps + step; - status = "Image " + std::to_string(i + 1) + "/" + std::to_string(config.num_images) - + - ": " + stage; - } - if (config.provider == Flux2ExecutionProvider::TrtRtx) - UpdateMemory(); - }, - request.generation); - if (!config.output_path.empty()) + const unsigned int seed = request.random_seeds ? random() : config.seed + i; + auto image = pipeline->GenerateImage( + seed, + [&](const char* stage, int step, int) { - std::filesystem::create_directories(config.output_path); - const auto file = - config.output_path / ("flux2_" + std::to_string(i) + "_" + std::to_string(seed) + - ".png"); - if (!din::io::SaveRgbFloatImage(file.string(), image.data.data(), image.height, image.width, - din::io::ImageValueRange::MinusOneToOne)) - throw std::runtime_error("Could not save " + file.string()); - } - std::lock_guard guard(mutex); - ready.push_back({std::move(image), seed, 0}); + { + std::lock_guard guard(mutex); + completed = static_cast(i) * config.steps + step; + status = "Image " + std::to_string(i + 1) + "/" + std::to_string(config.num_images) + + ": " + stage; + } + if (config.provider == Flux2ExecutionProvider::TrtRtx) + UpdateMemory(); + }, + request.generation); + if (!config.output_path.empty()) + { + std::filesystem::create_directories(config.output_path); + const auto file = + config.output_path / ("flux2_" + std::to_string(i) + "_" + std::to_string(seed) + ".png"); + if (!din::io::SaveRgbFloatImage(file.string(), image.data.data(), image.height, image.width, + din::io::ImageValueRange::MinusOneToOne)) + throw std::runtime_error("Could not save " + file.string()); } std::lock_guard guard(mutex); - completed = total; - status = "Done"; - } - catch (const std::exception& error) - { - pipeline.reset(); - loaded.reset(); - std::lock_guard guard(mutex); - status = std::string("Error: ") + error.what(); - failed = true; + ready.push_back({std::move(image), seed, 0}); } std::lock_guard guard(mutex); - busy = false; + completed = total; + status = "Done"; } - // Destroy sessions, streams and GPU allocations on their owning thread. - } - - std::condition_variable cv_; - bool closing_ = false; - std::optional request_; - std::thread worker_; - }; - - void Upload(Result& result) - { - const auto& img = result.image; - const size_t pixels = static_cast(img.width) * img.height; - if (!pixels || img.data.size() != pixels * 3) - throw std::runtime_error("Invalid image returned by pipeline"); - std::vector rgb(pixels * 3); - for (size_t p = 0; p < pixels; ++p) - for (size_t c = 0; c < 3; ++c) + catch (const std::exception& error) { - const float v = img.data[c * pixels + p]; - rgb[p * 3 + c] = - std::isfinite(v) - ? static_cast(std::lround(std::clamp(v * 127.5f + 127.5f, 0.0f, 255.0f))) - : 0; + pipeline.reset(); + loaded.reset(); + std::lock_guard guard(mutex); + status = std::string("Error: ") + error.what(); + failed = true; } - glGenTextures(1, &result.texture); - glBindTexture(GL_TEXTURE_2D, result.texture); - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR); - glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR); - glPixelStorei(GL_UNPACK_ALIGNMENT, 1); - glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, img.width, img.height, 0, GL_RGB, GL_UNSIGNED_BYTE, rgb.data()); + std::lock_guard guard(mutex); + busy = false; + } + // Destroy sessions, streams and GPU allocations on their owning thread. } + std::condition_variable cv_; + bool closing_ = false; + std::optional request_; + std::thread worker_; +}; + +void Upload(Result& result) +{ + const auto& img = result.image; + const size_t pixels = static_cast(img.width) * img.height; + if (!pixels || img.data.size() != pixels * 3) + throw std::runtime_error("Invalid image returned by pipeline"); + std::vector rgb(pixels * 3); + for (size_t p = 0; p < pixels; ++p) + for (size_t c = 0; c < 3; ++c) + { + const float v = img.data[c * pixels + p]; + rgb[p * 3 + c] = + std::isfinite(v) + ? static_cast(std::lround(std::clamp(v * 127.5f + 127.5f, 0.0f, 255.0f))) + : 0; + } + glGenTextures(1, &result.texture); + glBindTexture(GL_TEXTURE_2D, result.texture); + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR); + glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR); + glPixelStorei(GL_UNPACK_ALIGNMENT, 1); + glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, img.width, img.height, 0, GL_RGB, GL_UNSIGNED_BYTE, rgb.data()); +} - int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit_after) +int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit_after) +{ + Generator generator; + Flux2Config config = initial; + std::array model{}, output{}, cache{}; + bool cache_changed = false; + const auto cache_utf8 = initial.ep_cache_dir.parent_path().u8string(); + std::snprintf(cache.data(), cache.size(), "%s", reinterpret_cast(cache_utf8.c_str())); + std::array prompt{}; + const auto model_utf8 = config.model_dir.u8string(); + std::snprintf(model.data(), model.size(), "%s", reinterpret_cast(model_utf8.c_str())); + std::snprintf(output.data(), output.size(), "%s", config.output_path.string().c_str()); + std::snprintf(prompt.data(), prompt.size(), "%s", config.prompt.c_str()); + int streaming_mode = 0; // Auto / Manual / Off; independent of session configuration. + int budget = 50; + int count = static_cast(config.num_images); + int encoder = config.text_encoder == Flux2TextEncoder::Qwen3_4B ? 0 : 1; + int precision = config.precision == "bf16" ? 0 : config.precision == "fp16" ? 1 : config.precision == "fp8" ? 2 : 3; + const char* precisions[] = {"bf16", "fp16", "fp8", "nvfp4"}; + const char* backends[] = {"CPU", "CUDA", "DirectX", "DirectX CIG", "Vulkan", "Vulkan CIG"}; + int backend = static_cast(config.processing); + int provider = config.provider == Flux2ExecutionProvider::Cpu ? 0 : 1; + std::vector results; + int selected = -1, frames = 0; + bool launched = false; + int exit_code = 0; + GLFWwindow* window = glfwGetCurrentContext(); + while (!glfwWindowShouldClose(window)) { - Generator generator; - Flux2Config config = initial; - std::array model{}, output{}, cache{}; - bool cache_changed = false; - const auto cache_utf8 = initial.ep_cache_dir.parent_path().u8string(); - std::snprintf(cache.data(), cache.size(), "%s", reinterpret_cast(cache_utf8.c_str())); - std::array prompt{}; - const auto model_utf8 = config.model_dir.u8string(); - std::snprintf(model.data(), model.size(), "%s", reinterpret_cast(model_utf8.c_str())); - std::snprintf(output.data(), output.size(), "%s", config.output_path.string().c_str()); - std::snprintf(prompt.data(), prompt.size(), "%s", config.prompt.c_str()); - int streaming_mode = 0; // Auto / Manual / Off; independent of session configuration. - int budget = 50; - int count = static_cast(config.num_images); - int encoder = config.text_encoder == Flux2TextEncoder::Qwen3_4B ? 0 : 1; - int precision = config.precision == "bf16" - ? 0 - : config.precision == "fp16" - ? 1 - : config.precision == "fp8" - ? 2 - : 3; - const char* precisions[] = {"bf16", "fp16", "fp8", "nvfp4"}; - const char* backends[] = {"CPU", "CUDA", "DirectX", "DirectX CIG", "Vulkan", "Vulkan CIG"}; - int backend = static_cast(config.processing); - int provider = config.provider == Flux2ExecutionProvider::Cpu ? 0 : 1; - std::vector results; - int selected = -1, frames = 0; - bool launched = false; - int exit_code = 0; - GLFWwindow* window = glfwGetCurrentContext(); - while (!glfwWindowShouldClose(window)) + glfwPollEvents(); + ImGui_ImplOpenGL3_NewFrame(); + ImGui_ImplGlfw_NewFrame(); + ImGui::NewFrame(); + ImGui::SetNextWindowPos({0, 0}); + ImGui::SetNextWindowSize(ImGui::GetIO().DisplaySize); + ImGui::Begin("Flux2 Studio", nullptr, + ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize); + ImGui::TextUnformatted("DIN Deploy - FLUX.2-klein Interactive Studio"); + ImGui::Separator(); + bool busy, failed; + int completed, total; + std::string status; + size_t used_mb, total_mb; { - glfwPollEvents(); - ImGui_ImplOpenGL3_NewFrame(); - ImGui_ImplGlfw_NewFrame(); - ImGui::NewFrame(); - ImGui::SetNextWindowPos({0, 0}); - ImGui::SetNextWindowSize(ImGui::GetIO().DisplaySize); - ImGui::Begin("Flux2 Studio", nullptr, - ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize); - ImGui::TextUnformatted("DIN Deploy - FLUX.2-klein Interactive Studio"); - ImGui::Separator(); - bool busy, failed; - int completed, total; - std::string status; - size_t used_mb, total_mb; - { - std::lock_guard lock(generator.mutex); - busy = generator.busy; - failed = generator.failed; - status = generator.status; - completed = generator.completed; - total = generator.total; - used_mb = generator.used_mb; - total_mb = generator.total_mb; - for (auto& result : generator.ready) - { - results.push_back(std::move(result)); - selected = static_cast(results.size()) - 1; - } - generator.ready.clear(); - } - if (launched && !busy && exit_after) - { - exit_code = failed ? 1 : 0; - glfwSetWindowShouldClose(window, true); - } - ImGui::BeginChild("controls", ImVec2(430, 0), true); - ImGui::BeginDisabled(busy); - ImGui::InputText("Model root", model.data(), model.size()); - cache_changed |= ImGui::InputText("Cache root", cache.data(), cache.size()); - ImGui::InputText("Save directory", output.data(), output.size()); - ImGui::TextUnformatted("Prompt"); - ImGui::InputTextMultiline("##prompt", prompt.data(), prompt.size(), ImVec2(-1, 100)); - ImGui::RadioButton("Qwen3-4B", &encoder, 0); - const bool translator_exists = model[0] && std::filesystem::is_regular_file( - std::filesystem::u8path(model.data()) / - "text_encoder_translator/model.onnx"); - ImGui::BeginDisabled(!translator_exists); - ImGui::RadioButton("Qwen3-0.6B + translator", &encoder, 1); - ImGui::EndDisabled(); - if (!translator_exists) - ImGui::TextWrapped("Translator ONNX is missing from this model root."); - ImGui::Combo("Precision", &precision, precisions, IM_ARRAYSIZE(precisions)); - ImGui::Combo("Provider", &provider, "CPU\0TensorRT RTX\0"); - if (provider == 0) - { - backend = 0; - streaming_mode = 0; - } - if (ImGui::BeginCombo("Processing", backends[backend])) - { - for (int i = 0; i < 6; ++i) - { - if (!IsFlux2BackendAvailable(static_cast(i)) || (provider == 0 && i != 0)) - continue; - if (ImGui::Selectable(backends[i], backend == i)) - backend = i; - } - ImGui::EndCombo(); - } - ImGui::BeginDisabled(provider == 0); - ImGui::Combo("Weight streaming", &streaming_mode, "Auto\0Manual\0Off\0"); - if (streaming_mode == 1) - ImGui::SliderInt("Resident weights %", &budget, 10, 100); - ImGui::EndDisabled(); - ImGui::InputScalar("Seed", ImGuiDataType_U32, &config.seed); - ImGui::SliderInt("Denoise steps", &config.steps, 1, 50); - ImGui::SliderInt("Images", &count, 1, 16); - const bool paths_valid = model[0] && (!cache_changed || cache[0]); - ImGui::BeginDisabled(!paths_valid); - bool generate = ImGui::Button("Generate", ImVec2(-1, 35)); - bool random = ImGui::Button("Random sweep", ImVec2(-1, 30)); - ImGui::EndDisabled(); - ImGui::EndDisabled(); - if (!busy && paths_valid && (generate || random || (auto_generate && !launched))) + std::lock_guard lock(generator.mutex); + busy = generator.busy; + failed = generator.failed; + status = generator.status; + completed = generator.completed; + total = generator.total; + used_mb = generator.used_mb; + total_mb = generator.total_mb; + for (auto& result : generator.ready) { - for (auto& result : results) - if (result.texture) - glDeleteTextures(1, &result.texture); - results.clear(); - selected = -1; - config.model_dir = std::filesystem::u8path(model.data()); - config.output_path = output.data(); - config.prompt = prompt.data(); - const auto cache_namespace = CacheNamespace(config.model_dir); - config.ep_cache_dir = (cache_changed - ? std::filesystem::u8path(cache.data()) / "runtime" - : initial.ep_cache_dir) / cache_namespace; - config.ep_context_dir = (cache_changed - ? std::filesystem::u8path(cache.data()) / "ep_context" - : initial.ep_context_dir) / cache_namespace; - config.text_encoder = encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; - config.precision = precisions[precision]; - config.num_images = static_cast(count); - config.processing = static_cast(backend); - config.provider = provider == 0 ? Flux2ExecutionProvider::Cpu : Flux2ExecutionProvider::TrtRtx; - config.weight_streaming_budget = provider == 0 ? "" : "-1"; - Flux2GenerationOptions generation; - if (provider != 0 && streaming_mode != 0) - generation.weight_streaming_budget = streaming_mode == 1 ? std::to_string(budget) + "%" : "0"; - generator.Submit({config, generation, random}); - launched = true; + results.push_back(std::move(result)); + selected = static_cast(results.size()) - 1; } - ImGui::ProgressBar(static_cast(completed) / std::max(total, 1), ImVec2(-1, 0)); - ImGui::TextWrapped("%s", status.c_str()); - if (provider != 0 && total_mb) - ImGui::Text("Device memory: %zu / %zu MiB", used_mb, total_mb); - if (selected >= 0) + generator.ready.clear(); + } + if (launched && !busy && exit_after) + { + exit_code = failed ? 1 : 0; + glfwSetWindowShouldClose(window, true); + } + ImGui::BeginChild("controls", ImVec2(430, 0), true); + ImGui::BeginDisabled(busy); + ImGui::InputText("Model root", model.data(), model.size()); + cache_changed |= ImGui::InputText("Cache root", cache.data(), cache.size()); + ImGui::InputText("Save directory", output.data(), output.size()); + ImGui::TextUnformatted("Prompt"); + ImGui::InputTextMultiline("##prompt", prompt.data(), prompt.size(), ImVec2(-1, 100)); + ImGui::RadioButton("Qwen3-4B", &encoder, 0); + const bool translator_exists = + model[0] && std::filesystem::is_regular_file(std::filesystem::u8path(model.data()) / + "text_encoder_translator/model.onnx"); + ImGui::BeginDisabled(!translator_exists); + ImGui::RadioButton("Qwen3-0.6B + translator", &encoder, 1); + ImGui::EndDisabled(); + if (!translator_exists) + ImGui::TextWrapped("Translator ONNX is missing from this model root."); + ImGui::Combo("Precision", &precision, precisions, IM_ARRAYSIZE(precisions)); + ImGui::Combo("Provider", &provider, "CPU\0TensorRT RTX\0"); + if (provider == 0) + { + backend = 0; + streaming_mode = 0; + } + if (ImGui::BeginCombo("Processing", backends[backend])) + { + for (int i = 0; i < 6; ++i) { - const auto& t = results[selected].image.timings; - ImGui::Separator(); - ImGui::Text("Encode %.0f ms | RNG %.0f ms", t.encode_ms, t.rng_ms); - ImGui::Text("Denoise %.0f ms (%.0f ms/step)", t.denoise_ms, t.per_step_ms); - ImGui::Text("Decode %.0f ms | Total %.2f s", t.decode_ms, t.total_ms / 1000); + if (!IsFlux2BackendAvailable(static_cast(i)) || (provider == 0 && i != 0)) + continue; + if (ImGui::Selectable(backends[i], backend == i)) + backend = i; } - ImGui::EndChild(); - ImGui::SameLine(); - ImGui::BeginChild("images", ImVec2(0, 0), true); + ImGui::EndCombo(); + } + ImGui::BeginDisabled(provider == 0); + ImGui::Combo("Weight streaming", &streaming_mode, "Auto\0Manual\0Off\0"); + if (streaming_mode == 1) + ImGui::SliderInt("Resident weights %", &budget, 10, 100); + ImGui::EndDisabled(); + ImGui::InputScalar("Seed", ImGuiDataType_U32, &config.seed); + ImGui::SliderInt("Denoise steps", &config.steps, 1, 50); + ImGui::SliderInt("Images", &count, 1, 16); + const bool paths_valid = model[0] && (!cache_changed || cache[0]); + ImGui::BeginDisabled(!paths_valid); + bool generate = ImGui::Button("Generate", ImVec2(-1, 35)); + bool random = ImGui::Button("Random sweep", ImVec2(-1, 30)); + ImGui::EndDisabled(); + ImGui::EndDisabled(); + if (!busy && paths_valid && (generate || random || (auto_generate && !launched))) + { for (auto& result : results) - if (!result.texture) - Upload(result); - ImGui::BeginChild("preview", ImVec2(0, -115), false); - if (selected >= 0) - { - const auto available = ImGui::GetContentRegionAvail(); - const float side = std::max(1.0f, std::min(available.x, available.y)); - ImGui::Image(static_cast(results[selected].texture), ImVec2(side, side)); - } - else - ImGui::TextDisabled("Generated images will appear here."); - ImGui::EndChild(); - ImGui::BeginChild("thumbnails", ImVec2(0, 110), false, ImGuiWindowFlags_HorizontalScrollbar); - for (int i = 0; i < static_cast(results.size()); ++i) - { - if (i) - ImGui::SameLine(); - ImGui::PushID(i); - ImGui::BeginGroup(); - if (ImGui::ImageButton("image", static_cast(results[i].texture), ImVec2(72, 72))) - selected = i; - ImGui::Text("%u", results[i].seed); - ImGui::EndGroup(); - ImGui::PopID(); - } - ImGui::EndChild(); - ImGui::EndChild(); - ImGui::End(); - ImGui::Render(); - int width, height; - glfwGetFramebufferSize(window, &width, &height); - glViewport(0, 0, width, height); - glClearColor(0.09f, 0.09f, 0.10f, 1); - glClear(GL_COLOR_BUFFER_BIT); - ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData()); - glfwSwapBuffers(window); - if (smoke && ++frames >= 3) - break; + if (result.texture) + glDeleteTextures(1, &result.texture); + results.clear(); + selected = -1; + config.model_dir = std::filesystem::u8path(model.data()); + config.output_path = output.data(); + config.prompt = prompt.data(); + const auto cache_namespace = CacheNamespace(config.model_dir); + config.ep_cache_dir = + (cache_changed ? std::filesystem::u8path(cache.data()) / "runtime" : initial.ep_cache_dir) / + cache_namespace; + config.ep_context_dir = + (cache_changed ? std::filesystem::u8path(cache.data()) / "ep_context" : initial.ep_context_dir) / + cache_namespace; + config.text_encoder = encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + config.precision = precisions[precision]; + config.num_images = static_cast(count); + config.processing = static_cast(backend); + config.provider = provider == 0 ? Flux2ExecutionProvider::Cpu : Flux2ExecutionProvider::TrtRtx; + config.weight_streaming_budget = provider == 0 ? "" : "-1"; + Flux2GenerationOptions generation; + if (provider != 0 && streaming_mode != 0) + generation.weight_streaming_budget = streaming_mode == 1 ? std::to_string(budget) + "%" : "0"; + generator.Submit({config, generation, random}); + launched = true; + } + ImGui::ProgressBar(static_cast(completed) / std::max(total, 1), ImVec2(-1, 0)); + ImGui::TextWrapped("%s", status.c_str()); + if (provider != 0 && total_mb) + ImGui::Text("Device memory: %zu / %zu MiB", used_mb, total_mb); + if (selected >= 0) + { + const auto& t = results[selected].image.timings; + ImGui::Separator(); + ImGui::Text("Encode %.0f ms | RNG %.0f ms", t.encode_ms, t.rng_ms); + ImGui::Text("Denoise %.0f ms (%.0f ms/step)", t.denoise_ms, t.per_step_ms); + ImGui::Text("Decode %.0f ms | Total %.2f s", t.decode_ms, t.total_ms / 1000); } + ImGui::EndChild(); + ImGui::SameLine(); + ImGui::BeginChild("images", ImVec2(0, 0), true); for (auto& result : results) - if (result.texture) - glDeleteTextures(1, &result.texture); - return exit_code; + if (!result.texture) + Upload(result); + ImGui::BeginChild("preview", ImVec2(0, -115), false); + if (selected >= 0) + { + const auto available = ImGui::GetContentRegionAvail(); + const float side = std::max(1.0f, std::min(available.x, available.y)); + ImGui::Image(static_cast(results[selected].texture), ImVec2(side, side)); + } + else + ImGui::TextDisabled("Generated images will appear here."); + ImGui::EndChild(); + ImGui::BeginChild("thumbnails", ImVec2(0, 110), false, ImGuiWindowFlags_HorizontalScrollbar); + for (int i = 0; i < static_cast(results.size()); ++i) + { + if (i) + ImGui::SameLine(); + ImGui::PushID(i); + ImGui::BeginGroup(); + if (ImGui::ImageButton("image", static_cast(results[i].texture), ImVec2(72, 72))) + selected = i; + ImGui::Text("%u", results[i].seed); + ImGui::EndGroup(); + ImGui::PopID(); + } + ImGui::EndChild(); + ImGui::EndChild(); + ImGui::End(); + ImGui::Render(); + int width, height; + glfwGetFramebufferSize(window, &width, &height); + glViewport(0, 0, width, height); + glClearColor(0.09f, 0.09f, 0.10f, 1); + glClear(GL_COLOR_BUFFER_BIT); + ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData()); + glfwSwapBuffers(window); + if (smoke && ++frames >= 3) + break; } -} // namespace + for (auto& result : results) + if (result.texture) + glDeleteTextures(1, &result.texture); + return exit_code; +} +} // namespace int main(int argc, char** argv) { @@ -466,10 +457,10 @@ int main(int argc, char** argv) args.add_argument("--output").default_value(std::string{}); const auto runtime_cache_utf8 = (cache_root / "runtime").u8string(); const auto context_cache_utf8 = (cache_root / "ep_context").u8string(); - args.add_argument("--ep-cache").default_value(std::string( - reinterpret_cast(runtime_cache_utf8.c_str()))); - args.add_argument("--ep-context-dir").default_value(std::string( - reinterpret_cast(context_cache_utf8.c_str()))); + args.add_argument("--ep-cache") + .default_value(std::string(reinterpret_cast(runtime_cache_utf8.c_str()))); + args.add_argument("--ep-context-dir") + .default_value(std::string(reinterpret_cast(context_cache_utf8.c_str()))); args.add_argument("--prompt").default_value(std::string(DEFAULT_PROMPT)); args.add_argument("--encoder").default_value(std::string("auto")).choices("auto", "4b", "translator"); args.add_argument("--precision").default_value(std::string("bf16")).choices("bf16", "fp16", "fp8", "nvfp4"); @@ -490,9 +481,10 @@ int main(int argc, char** argv) config.precision = args.get("--precision"); config.num_images = 1; const auto encoder = args.get("--encoder"); - const bool translator = encoder == "translator" || - (encoder == "auto" && !config.model_dir.empty() && std::filesystem::is_regular_file( - config.model_dir / "text_encoder_translator/model.onnx")); + const bool translator = + encoder == "translator" || + (encoder == "auto" && !config.model_dir.empty() && + std::filesystem::is_regular_file(config.model_dir / "text_encoder_translator/model.onnx")); config.text_encoder = translator ? Flux2TextEncoder::Qwen3_06BTranslator : Flux2TextEncoder::Qwen3_4B; config.weight_streaming_budget = "-1"; if (!glfwInit()) diff --git a/image_gen/flux2/flux2_vk_backend.cpp b/image_gen/flux2/flux2_vk_backend.cpp index 0b2eb66..b1b9559 100644 --- a/image_gen/flux2/flux2_vk_backend.cpp +++ b/image_gen/flux2/flux2_vk_backend.cpp @@ -716,8 +716,8 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi throw std::runtime_error("Vulkan processing requires --provider trt-rtx"); } const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); - const auto latent_stats = LoadFlux2LatentStats( - model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); + const auto latent_stats = + LoadFlux2LatentStats(model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); state.use_cig = config.processing == Flux2ProcessingBackend::VkCig; const Flux2ModelCachePaths cache_paths = MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "vk_cig" : "vk", config.text_encoder, @@ -749,7 +749,7 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi { state.cig_external_compute_queue_data = state.vk->createCudaGraphicsInteropData(); state.graphics_interop = std::make_unique(*state.interop_api, state.trt_device, - state.cig_external_compute_queue_data); + state.cig_external_compute_queue_data); } // ----------------------------------------------------------------- diff --git a/image_gen/flux2/utils.cpp b/image_gen/flux2/utils.cpp index 156a2ec..a645da4 100644 --- a/image_gen/flux2/utils.cpp +++ b/image_gen/flux2/utils.cpp @@ -16,19 +16,19 @@ #include "utils.h" #include +#include #include #include #include -#include #include + #include Flux2LatentStats LoadFlux2LatentStats(const std::filesystem::path& path, size_t channels) { std::ifstream file(path); if (!file) - throw std::runtime_error("Missing VAE latent statistics: " + path.string() + - ". Re-export the VAE decoder."); + throw std::runtime_error("Missing VAE latent statistics: " + path.string() + ". Re-export the VAE decoder."); try { const auto document = nlohmann::json::parse(file); @@ -36,8 +36,7 @@ Flux2LatentStats LoadFlux2LatentStats(const std::filesystem::path& path, size_t { const auto& values = document.at(key); if (!values.is_array() || values.size() != channels) - throw std::runtime_error(std::string(key) + " must contain " + - std::to_string(channels) + " values"); + throw std::runtime_error(std::string(key) + " must contain " + std::to_string(channels) + " values"); std::vector result; result.reserve(channels); for (const auto& value : values) From ef54af6175932b712c871a675b7af51568c710d3 Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:35:42 +0000 Subject: [PATCH 09/14] Fix FLUX.2 README export script paths and Release CLI path --- image_gen/flux2/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/image_gen/flux2/README.md b/image_gen/flux2/README.md index 1f5a01d..0647569 100644 --- a/image_gen/flux2/README.md +++ b/image_gen/flux2/README.md @@ -23,8 +23,8 @@ available. Run from the repository root: ```powershell -python ./model_export/export_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --output models/flux_full --model all --transformer-precision bf16 -python ./model_export/quantize_onnx.py --input models/flux_full/transformer_bf16 --quant fp8 --output models/flux_full/transformer_fp8 +python image_gen/flux2/model_export/export_flux2.py --model_name black-forest-labs/FLUX.2-klein-4b --output models/flux_full --model all --transformer-precision bf16 +python image_gen/flux2/model_export/quantize_onnx.py --input models/flux_full/transformer_bf16 --quant fp8 --output models/flux_full/transformer_fp8 ``` Use `--quant nvfp4` with an output directory of `transformer_nvfp4` for NVFP4. @@ -64,7 +64,7 @@ Use `-DDIN_BUILD_FLUX2_UI=OFF` for headless builds. For pure CLI you can use: ```powershell -out/build/windows-x64/bin/Debug/din_flux2_cli.exe --model-dir models/flux_full --precision fp8 --provider trt-rtx --processing cuda --prompt "a red fox" --output out +out/build/windows-x64/bin/Release/din_flux2_cli.exe --model-dir models/flux_full --precision fp8 --provider trt-rtx --processing cuda --prompt "a red fox" --output out ``` Use `--help` for available processing backends and options, including `--encoder`, From b35cb6050c5a68be6d095c608f11e242ed91e4b3 Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Fri, 9 Oct 2026 15:08:24 +0200 Subject: [PATCH 10/14] support binary lookup for 0.5.0 or earlier release --- cmake/onnxruntime_trt_rtx_ep.cmake | 45 ++++++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/cmake/onnxruntime_trt_rtx_ep.cmake b/cmake/onnxruntime_trt_rtx_ep.cmake index d6668e4..aecef0e 100644 --- a/cmake/onnxruntime_trt_rtx_ep.cmake +++ b/cmake/onnxruntime_trt_rtx_ep.cmake @@ -155,14 +155,53 @@ if(_onnxruntime_trt_rtx_ep_runtime_dependencies) endif() +# Resolve at build time, after ExternalProject has produced the library. +# Keep the compatibility logic inline; generate its executable script in the build tree. +set(_onnxruntime_trt_rtx_ep_copy_script "${CMAKE_CURRENT_BINARY_DIR}/copy_onnxruntime_trt_rtx_ep.cmake") +file(GENERATE OUTPUT "${_onnxruntime_trt_rtx_ep_copy_script}" CONTENT [=[ +if(EP_MULTI_CONFIG) + set(_ep_candidates + "${EP_BINARY_DIR}/bin/${EP_CONFIGURATION}/${EP_LIBRARY_NAME}" + "${EP_BINARY_DIR}/${EP_CONFIGURATION}/${EP_LIBRARY_NAME}" + ) +else() + set(_ep_candidates + "${EP_BINARY_DIR}/bin/${EP_LIBRARY_NAME}" + "${EP_BINARY_DIR}/${EP_LIBRARY_NAME}" + ) +endif() + +foreach(_ep_candidate IN LISTS _ep_candidates) + if(EXISTS "${_ep_candidate}") + file(MAKE_DIRECTORY "${EP_RUNTIME_DIR}") + execute_process( + COMMAND "${CMAKE_COMMAND}" -E copy_if_different + "${_ep_candidate}" "${EP_RUNTIME_DIR}" + RESULT_VARIABLE _ep_copy_result + ERROR_VARIABLE _ep_copy_error + ) + if(NOT _ep_copy_result STREQUAL "0") + message(FATAL_ERROR "Failed to copy TensorRT RTX EP: ${_ep_copy_error}") + endif() + return() + endif() +endforeach() + +message(FATAL_ERROR "TensorRT RTX EP library was not found. Checked: ${_ep_candidates}") +]=]) + set(_onnxruntime_trt_rtx_ep_runtime_commands COMMAND "${CMAKE_COMMAND}" -E make_directory "${ONNXRUNTIME_TRT_RTX_EP_RUNTIME_DIR}" ) list(APPEND _onnxruntime_trt_rtx_ep_runtime_commands - COMMAND "${CMAKE_COMMAND}" -E copy_if_different - "${_onnxruntime_trt_rtx_ep_binary_dir}/$/${_onnxruntime_trt_rtx_ep_runtime_library_name}" - "${ONNXRUNTIME_TRT_RTX_EP_RUNTIME_DIR}" + COMMAND "${CMAKE_COMMAND}" + "-DEP_BINARY_DIR=${_onnxruntime_trt_rtx_ep_binary_dir}" + "-DEP_CONFIGURATION=$" + "-DEP_MULTI_CONFIG=$" + "-DEP_LIBRARY_NAME=${_onnxruntime_trt_rtx_ep_runtime_library_name}" + "-DEP_RUNTIME_DIR=${ONNXRUNTIME_TRT_RTX_EP_RUNTIME_DIR}" + -P "${_onnxruntime_trt_rtx_ep_copy_script}" ) foreach(_onnxruntime_trt_rtx_ep_runtime_dependency IN LISTS _onnxruntime_trt_rtx_ep_runtime_dependencies) From 234bed5de7297b17e64ef4cfeb15f54fc9115172 Mon Sep 17 00:00:00 2001 From: Maximilian Mueller Date: Wed, 14 Oct 2026 13:05:04 +0200 Subject: [PATCH 11/14] add configurable RNG and relax model reload conditions to make model switching from the UI faster --- CMakeLists.txt | 4 + cmake/random123.cmake | 49 +++++ common/CMakeLists.txt | 17 ++ common/ort_session.cpp | 32 +++- common/ort_session.h | 7 +- common/rng.cpp | 88 +++++++++ common/rng.h | 31 +++ common/rng_cuda.cu | 86 +++++++++ common/rng_cuda.h | 21 +++ image_gen/flux2/CMakeLists.txt | 4 +- image_gen/flux2/README.md | 3 + image_gen/flux2/cuda_kernels.cu | 1 + image_gen/flux2/cuda_kernels.h | 1 + image_gen/flux2/flux2_cli.cpp | 47 +++-- image_gen/flux2/flux2_cli.h | 70 ++++++- image_gen/flux2/flux2_cuda_backend.cpp | 247 +++++++++++++----------- image_gen/flux2/flux2_dx_backend.cpp | 194 ++++++++++--------- image_gen/flux2/flux2_models.cpp | 77 ++++++++ image_gen/flux2/flux2_models.h | 66 +++++++ image_gen/flux2/flux2_pipeline.cpp | 10 +- image_gen/flux2/flux2_ui.cpp | 86 ++++++--- image_gen/flux2/flux2_vk_backend.cpp | 250 ++++++++++++------------- image_gen/flux2/helpers/flux2.h | 70 ++++--- 23 files changed, 1023 insertions(+), 438 deletions(-) create mode 100644 cmake/random123.cmake create mode 100644 common/rng.cpp create mode 100644 common/rng.h create mode 100644 common/rng_cuda.cu create mode 100644 common/rng_cuda.h create mode 100644 image_gen/flux2/flux2_models.cpp create mode 100644 image_gen/flux2/flux2_models.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 05e7abf..0a6e04c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -44,6 +44,10 @@ set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY "${DIN_RUNTIME_DIR}") list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") +# Load optional local URL overrides when present. +if (EXISTS "${CMAKE_CURRENT_LIST_DIR}/cmake/internal_urls.cmake") + include("${CMAKE_CURRENT_LIST_DIR}/cmake/internal_urls.cmake") +endif () include(onnxruntime) include(nlohmann_json) include(argparse) diff --git a/cmake/random123.cmake b/cmake/random123.cmake new file mode 100644 index 0000000..31d4966 --- /dev/null +++ b/cmake/random123.cmake @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +include_guard(GLOBAL) +include(FetchContent) +set(FETCHCONTENT_UPDATES_DISCONNECTED ON) + +FetchContent_Declare(random123 + GIT_REPOSITORY https://github.com/DEShawResearch/random123.git + GIT_TAG v1.14.0 + GIT_SHALLOW TRUE +) +FetchContent_MakeAvailable(random123) + +# Random123 v1.14.0's MSVC feature header rejects ARM64 and otherwise selects +# x86 intrinsics for newer MSVC versions. Philox4x32 uses the scalar path. +if(MSVC AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(ARM64|arm64)$") + set(_din_r123_msvc_features "${random123_SOURCE_DIR}/include/Random123/features/msvcfeatures.h") + file(READ "${_din_r123_msvc_features}" _din_r123_features) + set(_din_r123_old_guard "#if !defined(_M_IX86) && !defined(_M_X64)") + set(_din_r123_new_guard "#if !defined(_M_IX86) && !defined(_M_X64) && !defined(_M_ARM64)") + if(NOT _din_r123_features MATCHES "_M_ARM64") + if(NOT _din_r123_features MATCHES "#if !defined\\(_M_IX86\\) && !defined\\(_M_X64\\)") + message(FATAL_ERROR "Random123 MSVC architecture guard changed; review the ARM64 patch") + endif() + string(REPLACE "${_din_r123_old_guard}" "${_din_r123_new_guard}" + _din_r123_features "${_din_r123_features}") + file(WRITE "${_din_r123_msvc_features}" "${_din_r123_features}") + endif() + unset(_din_r123_msvc_features) + unset(_din_r123_features) + unset(_din_r123_old_guard) + unset(_din_r123_new_guard) +endif() + +add_library(din_random123 INTERFACE) +add_library(Random123::Random123 ALIAS din_random123) +target_include_directories(din_random123 SYSTEM INTERFACE "${random123_SOURCE_DIR}/include") +if(MSVC AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(ARM64|arm64)$") + target_compile_definitions(din_random123 INTERFACE + R123_USE_SSE=0 + R123_USE_SSE4_1=0 + R123_USE_SSE4_2=0 + R123_USE_EMMINTRIN_H=0 + R123_USE_SMMINTRIN_H=0 + R123_USE_WMMINTRIN_H=0 + R123_USE_CPUID_MSVC=0 + ) +endif() diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index ebe3a1c..e7e2c3b 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -33,3 +33,20 @@ set_target_properties(din_common_ort PROPERTIES ) add_subdirectory(io) + +# CPU callers do not depend on CUDA headers or libraries. +include(random123) +add_library(din_common_rng STATIC rng.cpp) +target_include_directories(din_common_rng PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}") +target_compile_features(din_common_rng PUBLIC cxx_std_20) +target_link_libraries(din_common_rng PRIVATE Random123::Random123) +set_target_properties(din_common_rng PROPERTIES POSITION_INDEPENDENT_CODE ON) + +enable_language(CUDA) +add_library(din_common_rng_cuda STATIC rng_cuda.cu) +target_include_directories(din_common_rng_cuda PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}") +target_link_libraries(din_common_rng_cuda PUBLIC CUDA::cudart) +set_target_properties(din_common_rng_cuda PROPERTIES + CUDA_ARCHITECTURES "${CUDA_ARCHITECTURES}" + POSITION_INDEPENDENT_CODE ON +) diff --git a/common/ort_session.cpp b/common/ort_session.cpp index 81e6e66..c636e32 100644 --- a/common/ort_session.cpp +++ b/common/ort_session.cpp @@ -916,17 +916,31 @@ Ort::ConstMemoryInfo OrtRunner::DeviceMemory() template TensorBuffer::TensorBuffer(OrtRunner& runner, const std::vector& shape, bool use_device_io, bool disable_uma) - : runner_(&runner) + : TensorBuffer(runner.env, runner.ep_device, runner.compute_stream, shape, use_device_io, disable_uma) +{ +} + +template +TensorBuffer::TensorBuffer(Ort::Env& env, Ort::ConstEpDevice device, Ort::SyncStream* stream, + const std::vector& shape, bool use_device_io, bool disable_uma) + : env_(&env) + , compute_stream_(stream) , use_device_io_(use_device_io) - , use_unified_memory_(use_device_io && runner.UsesUnifiedMemory() && !disable_uma) + , use_unified_memory_(use_device_io && IsCudaUnifiedMemoryDevice(device) && !disable_uma) , shape_(shape) { if (use_device_io_) { - host_tensor_ = Ort::Value::CreateTensor(runner.PinnedAllocator(), shape_.data(), shape_.size()); + if (!device || stream == nullptr) + throw std::invalid_argument("Device tensors require a device and a persistent stream"); + const auto pinned = env.GetSharedAllocator(device.GetMemoryInfo(OrtDeviceMemoryType_HOST_ACCESSIBLE)); + if (!pinned) throw std::runtime_error("Missing shared pinned allocator"); + host_tensor_ = Ort::Value::CreateTensor(pinned, shape_.data(), shape_.size()); if (!use_unified_memory_) { - device_tensor_ = Ort::Value::CreateTensor(runner.DeviceAllocator(), shape_.data(), shape_.size()); + const auto allocator = env.GetSharedAllocator(device.GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); + if (!allocator) throw std::runtime_error("Missing shared device allocator"); + device_tensor_ = Ort::Value::CreateTensor(allocator, shape_.data(), shape_.size()); } } else @@ -971,7 +985,7 @@ void TensorBuffer::CopyAsyncToDevice() { if (use_device_io_ && !use_unified_memory_) { - Copy(runner_->env, {&host_tensor_}, {&device_tensor_}, *runner_->compute_stream); + Copy(*env_, {&host_tensor_}, {&device_tensor_}, *compute_stream_); } } @@ -983,7 +997,7 @@ NotificationPtr TensorBuffer::CopyAsyncToDeviceWithNotification() { return NotificationPtr{nullptr}; } - auto notification = NotificationPtr(*runner_->compute_stream); + auto notification = NotificationPtr(*compute_stream_); // Even when the transfer is a unified-memory no-op, record the stream so // callers retain a usable ordering point for Sync(). notification.Record(); @@ -995,7 +1009,7 @@ void TensorBuffer::CopyAsyncToHost() { if (use_device_io_ && !use_unified_memory_) { - Copy(runner_->env, {&device_tensor_}, {&host_tensor_}, *runner_->compute_stream); + Copy(*env_, {&device_tensor_}, {&host_tensor_}, *compute_stream_); } } @@ -1007,7 +1021,7 @@ NotificationPtr TensorBuffer::CopyAsyncToHostWithNotification() { return NotificationPtr{nullptr}; } - auto notification = NotificationPtr(*runner_->compute_stream); + auto notification = NotificationPtr(*compute_stream_); // A unified-memory buffer needs no copy, but the notification still waits // for device writes before its host data is consumed. notification.Record(); @@ -1023,7 +1037,7 @@ void TensorBuffer::CopyFrom(TensorBuffer& other) } if (use_device_io_ && !use_unified_memory_) { - Copy(runner_->env, {&other.device_tensor_}, {&device_tensor_}, *runner_->compute_stream); + Copy(*env_, {&other.device_tensor_}, {&device_tensor_}, *compute_stream_); } else { diff --git a/common/ort_session.h b/common/ort_session.h index 3e0a09f..b0bf293 100644 --- a/common/ort_session.h +++ b/common/ort_session.h @@ -199,6 +199,8 @@ class TensorBuffer { public: TensorBuffer(OrtRunner& runner, const std::vector& shape, bool use_device_io, bool disable_uma = false); + TensorBuffer(Ort::Env& env, Ort::ConstEpDevice device, Ort::SyncStream* stream, + const std::vector& shape, bool use_device_io, bool disable_uma = false); Ort::Value& BindingValue(); @@ -228,7 +230,10 @@ class TensorBuffer [[nodiscard]] size_t ElementCount() const; - OrtRunner* runner_ = nullptr; + // Allocators belong to the environment. External streams must outlive this + // buffer, allowing a model session to be replaced without moving its tensors. + Ort::Env* env_ = nullptr; + Ort::SyncStream* compute_stream_ = nullptr; bool use_device_io_ = false; bool use_unified_memory_ = false; std::vector shape_; diff --git a/common/rng.cpp b/common/rng.cpp new file mode 100644 index 0000000..1db078a --- /dev/null +++ b/common/rng.cpp @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "rng.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace din::common +{ +namespace +{ +float round_bfloat16(float value) +{ + uint32_t bits = std::bit_cast(value); + if ((bits & 0x7fffffffU) > 0x7f800000U) + bits |= 0x00400000U; // Preserve NaNs as quiet NaNs. + else + bits += 0x7fffU + ((bits >> 16) & 1U); + return std::bit_cast(bits & 0xffff0000U); +} + +void box_muller(uint32_t x, uint32_t y, float* result) +{ + constexpr float a = 2.3283064e-10f; + constexpr float b = a * 6.2831855f; + // Match CUDA's fused multiply/add, including integer-to-FP32 conversion. + const float u = std::fma(static_cast(x), a, a / 2); + const float v = std::fma(static_cast(y), b, b / 2); + const float radius = std::sqrt(-2.0f * std::log(u)); + result[0] = radius * std::sin(v); + result[1] = radius * std::cos(v); +} +} + +void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, const NormalOptions& options) +{ + if (rng != Rng::CpuStd && rng != Rng::CpuPhilox) + throw std::invalid_argument("CPU normal generation requires CpuStd or CpuPhilox"); + if (options.philox_stride == 0 || options.philox_stride > std::numeric_limits::max() / 4) + throw std::invalid_argument("Philox stride must be positive and allow four components without overflow"); + const size_t count = output.size(); + if (rng == Rng::CpuStd) + { + std::default_random_engine engine(static_cast(seed)); + std::normal_distribution normal(0.0f, 1.0f); + for (size_t i = 0; i < count; ++i) + { + const float sample = normal(engine); + output[i] = options.round_to_bfloat16 ? round_bfloat16(sample) : sample; + } + return; + } + + using Philox = r123::Philox4x32_R<10>; + const Philox::key_type key = {{static_cast(seed), static_cast(seed >> 32)}}; + const size_t step = 4 * options.philox_stride; + for (size_t thread = 0; thread < std::min(count, options.philox_stride); ++thread) + { + uint64_t iteration = 0; + for (size_t base = thread; base < count; ++iteration) + { + const Philox::ctr_type counter = {{static_cast(iteration), + static_cast(iteration >> 32), static_cast(thread), + static_cast(static_cast(thread) >> 32)}}; + const auto words = Philox{}(counter, key); + float samples[4]; + box_muller(words[0], words[1], samples); + box_muller(words[2], words[3], samples + 2); + for (size_t component = 0; component < 4; ++component) + { + const size_t delta = component * options.philox_stride; + if (delta >= count - base) + break; + const size_t index = base + delta; + output[index] = options.round_to_bfloat16 ? round_bfloat16(samples[component]) : samples[component]; + } + if (count - base <= step) + break; + base += step; + } + } +} +} // namespace din::common diff --git a/common/rng.h b/common/rng.h new file mode 100644 index 0000000..d826ba4 --- /dev/null +++ b/common/rng.h @@ -0,0 +1,31 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include + +namespace din::common +{ +enum class Rng { CpuStd, CudaPhilox, CpuPhilox }; + +// Reference launch geometry: H100 SXM5, 132 SMs * 2048 threads/SM. +inline constexpr size_t kH100PhiloxStride = 270336; + +struct NormalOptions +{ + bool round_to_bfloat16 = false; + // Normal4 components are separated by this many samples. Must be nonzero. + size_t philox_stride = kH100PhiloxStride; +}; + +// Fill a contiguous buffer with N(0, 1) samples, restarting at the seed each call. +// CpuStd preserves std::default_random_engine/normal_distribution behavior, +// including platform dependence and truncation of the seed to unsigned int. +// CpuPhilox uses Philox4x32-10 and cuRAND's Box-Muller convention and ordering. +// CPU math can differ from CUDA, including at BF16 rounding boundaries. +// CudaPhilox must be dispatched through GenerateNormalCuda in rng_cuda.h. +void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, + const NormalOptions& options = {}); +} diff --git a/common/rng_cuda.cu b/common/rng_cuda.cu new file mode 100644 index 0000000..dcee40a --- /dev/null +++ b/common/rng_cuda.cu @@ -0,0 +1,86 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "rng_cuda.h" +#include +#include +#include +#include +#include + +namespace din::common +{ +namespace +{ +constexpr unsigned int kThreadsPerBlock = 256; +void check_cuda(cudaError_t result) +{ + if (result != cudaSuccess) throw std::runtime_error(cudaGetErrorString(result)); +} +// Match ATen DistributionTemplates.h: one Philox subsequence per thread, +// normal4 components spaced by the entire grid, then grid-stride by 4. +__global__ __launch_bounds__(256, 4) void normal_philox_kernel( + uint64_t seed, size_t count, bool round_to_bfloat16, size_t channels, size_t spatial, float* output) +{ + const size_t thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t stride = static_cast(blockDim.x) * gridDim.x; + const size_t rounded = ((count - 1) / (stride * 4) + 1) * stride * 4; + curandStatePhilox4_32_10_t state; + curand_init(seed, thread, 0, &state); + for (size_t base = thread; base < rounded; base += stride * 4) + { + const float4 samples = curand_normal4(&state); + const float values[] = {samples.x, samples.y, samples.z, samples.w}; +#pragma unroll + for (int component = 0; component < 4; ++component) + { + const size_t index = base + stride * component; + if (index < count) + { + const size_t destination = channels == 0 ? index + : ((index / (channels * spatial)) * spatial + index % spatial) * channels + + (index / spatial) % channels; + output[destination] = round_to_bfloat16 + ? __bfloat162float(__float2bfloat16_rn(values[component])) : values[component]; + } + } + } +} + +void generate_normal_cuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, + bool round_to_bfloat16, size_t channels, size_t spatial) +{ + if (count == 0) return; + if (output == nullptr || count > std::numeric_limits::max() / 2) + throw std::invalid_argument("CUDA normal generation requires a valid buffer and supported count"); + int device = 0; + cudaDeviceProp properties{}; + check_cuda(cudaGetDevice(&device)); + check_cuda(cudaGetDeviceProperties(&properties, device)); + // ATen caps the grid by the device's thread capacity, not kernel occupancy. + const size_t capacity = static_cast(properties.multiProcessorCount) * + (properties.maxThreadsPerMultiProcessor / kThreadsPerBlock); + const unsigned int blocks = static_cast( + std::min(capacity, (count + kThreadsPerBlock - 1) / kThreadsPerBlock)); + normal_philox_kernel<<>>( + seed, count, round_to_bfloat16, channels, spatial, output); + check_cuda(cudaPeekAtLastError()); +} + +} + +void GenerateNormalCuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, + bool round_to_bfloat16) +{ + generate_normal_cuda(stream, output, count, seed, round_to_bfloat16, 0, 0); +} + +void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, + size_t channels, size_t spatial, bool round_to_bfloat16) +{ + if (channels == 0 || spatial == 0 || channels > std::numeric_limits::max() / spatial || + count % (channels * spatial) != 0) + throw std::invalid_argument("Packed CUDA normal dimensions must match the buffer"); + generate_normal_cuda(stream, output, count, seed, round_to_bfloat16, channels, spatial); +} + +} // namespace din::common diff --git a/common/rng_cuda.h b/common/rng_cuda.h new file mode 100644 index 0000000..0a1d5e2 --- /dev/null +++ b/common/rng_cuda.h @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include + +namespace din::common +{ +// Asynchronously fill a contiguous device buffer with N(0, 1) samples. +// Fresh Philox seed, offset zero. Uses ATen's device-dependent launch geometry. +// An empty buffer is a no-op. The buffer and stream must remain valid until done. +void GenerateNormalCuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, + bool round_to_bfloat16 = false); + +// Generate the same [batch, channels, spatial] samples directly into contiguous +// [batch, spatial, channels] storage. count must be a multiple of channels * spatial. +void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, + size_t channels, size_t spatial, bool round_to_bfloat16 = false); +} diff --git a/image_gen/flux2/CMakeLists.txt b/image_gen/flux2/CMakeLists.txt index 8689f14..72de3ca 100644 --- a/image_gen/flux2/CMakeLists.txt +++ b/image_gen/flux2/CMakeLists.txt @@ -8,6 +8,7 @@ option(DIN_BUILD_FLUX2_UI "Build the GLFW/OpenGL Flux2 desktop sample" ON) add_library(din_flux2_common STATIC utils.cpp + flux2_models.cpp ) target_include_directories(din_flux2_common @@ -19,6 +20,7 @@ target_include_directories(din_flux2_common target_link_libraries(din_flux2_common PUBLIC din_common_ort + din_common_rng PRIVATE nlohmann_json::nlohmann_json ) @@ -71,7 +73,7 @@ configure_flux2_backend_target(din_flux2_cuda) set_target_properties(din_flux2_cuda PROPERTIES CUDA_ARCHITECTURES "${CUDA_ARCHITECTURES}" ) -target_link_libraries(din_flux2_cuda PRIVATE CUDA::cudart CUDA::nvtx3) +target_link_libraries(din_flux2_cuda PRIVATE din_common_rng_cuda CUDA::cudart CUDA::nvtx3) if(Vulkan_FOUND AND SLANGC) add_subdirectory(shaders) diff --git a/image_gen/flux2/README.md b/image_gen/flux2/README.md index 0647569..4f68f99 100644 --- a/image_gen/flux2/README.md +++ b/image_gen/flux2/README.md @@ -56,6 +56,9 @@ The studio provides prompt editing, encoder/precision/backend selection, seed sw weight-streaming controls, and an image gallery with optional PNG output. It selects the translator when exported; override with `--encoder 4b|translator`. Runtime and compiled-model caching are always enabled; clear affected caches after re-exporting models in place. +The has configuration checks per stage and does partial reload of impacted models upon changes. +Changing RNG, denoising steps, or the runtime weight-streaming budget reuses all +loaded model sessions. For a full execution backend change all models are reloaded. The UI uses GLFW, Dear ImGui, and OpenGL. For Debian/Ubuntu, install `libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev`. diff --git a/image_gen/flux2/cuda_kernels.cu b/image_gen/flux2/cuda_kernels.cu index 104f68b..d32bdc5 100644 --- a/image_gen/flux2/cuda_kernels.cu +++ b/image_gen/flux2/cuda_kernels.cu @@ -5,6 +5,7 @@ #include #include + #include "cu_helper.h" #include "cuda_kernels.h" diff --git a/image_gen/flux2/cuda_kernels.h b/image_gen/flux2/cuda_kernels.h index 46a92d1..072b128 100644 --- a/image_gen/flux2/cuda_kernels.h +++ b/image_gen/flux2/cuda_kernels.h @@ -6,6 +6,7 @@ #include #include +#include void launch_flux_euler_kernel(cudaStream_t stream, float t_curr, float t_next, size_t total_elements, float* hidden_states, const float* transformer_output); diff --git a/image_gen/flux2/flux2_cli.cpp b/image_gen/flux2/flux2_cli.cpp index 83bd212..49b3c80 100644 --- a/image_gen/flux2/flux2_cli.cpp +++ b/image_gen/flux2/flux2_cli.cpp @@ -158,17 +158,17 @@ Flux2Config parse_args(int argc, char* argv[]) .metavar("PATH") .help("Root directory with shared Flux2 ONNX artifacts and transformer_ directories."); parser.add_argument("--precision") - .default_value(config.precision) + .default_value(config.transformer.precision) .nargs(1) .metavar("bf16|fp16|fp8|nvfp4") .help("Transformer precision to load from transformer_."); parser.add_argument("--ep-cache") - .default_value(config.ep_cache_dir.string()) + .default_value(config.text_encoder.ep_cache_dir.string()) .nargs(1) .metavar("PATH") .help("TensorRT RTX runtime cache directory."); parser.add_argument("--ep-context-dir") - .default_value(config.ep_context_dir.string()) + .default_value(config.text_encoder.ep_context_dir.string()) .nargs(1) .metavar("PATH") .help("Directory for ONNX Runtime EP-context models."); @@ -183,6 +183,11 @@ Flux2Config parse_args(int argc, char* argv[]) .metavar("TEXT") .help("Prompt text to encode with the Flux2 tokenizer."); parser.add_argument("--seed").default_value(std::to_string(config.seed)).nargs(1).metavar("N").help("Random seed."); + parser.add_argument("--rng") + .default_value(std::string{"cpu-std"}) + .nargs(1) + .metavar("cpu-std|cuda-philox|cpu-philox") + .help("Noise generator: cpu-std (default), cuda-philox (CUDA only), or cpu-philox (fixed H100 SXM5 ordering)."); parser.add_argument("--encoder").default_value(std::string("4b")).help("4b or translator."); parser.add_argument("--steps").default_value(std::string("4")).help("Denoise steps, 1 through 50."); parser.add_argument("--ws") @@ -206,28 +211,38 @@ Flux2Config parse_args(int argc, char* argv[]) config.processing = parse_processing_backend(parser.get("--processing")); config.provider = parse_execution_provider(parser.get("--provider")); - config.model_dir = parser.get("--model-dir"); - config.precision = parse_precision(parser.get("--precision")); - config.ep_cache_dir = parser.get("--ep-cache"); - config.ep_context_dir = parser.get("--ep-context-dir"); + const std::filesystem::path model_root = parser.get("--model-dir"); + config.transformer.precision = parse_precision(parser.get("--precision")); + config.text_encoder.ep_cache_dir = parser.get("--ep-cache"); + config.text_encoder.ep_context_dir = parser.get("--ep-context-dir"); config.prompt = parser.get("--prompt"); config.seed = parse_uint(parser.get("--seed"), "--seed"); + const auto rng = parser.get("--rng"); + if (rng == "cpu-std") + config.rng = din::common::Rng::CpuStd; + else if (rng == "cuda-philox") + config.rng = din::common::Rng::CudaPhilox; + else if (rng == "cpu-philox") + config.rng = din::common::Rng::CpuPhilox; + else + throw std::invalid_argument("--rng must be cpu-std, cuda-philox, or cpu-philox"); const auto encoder = parser.get("--encoder"); if (encoder != "4b" && encoder != "translator") throw std::invalid_argument("--encoder must be 4b or translator"); - config.text_encoder = encoder == "4b" ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + config.text_encoder.variant = encoder == "4b" ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; const auto steps = parse_uint(parser.get("--steps"), "--steps"); if (steps < 1 || steps > 50) throw std::invalid_argument("--steps must be 1 through 50"); config.steps = static_cast(steps); - config.weight_streaming_budget = parser.get("--ws"); - if (config.weight_streaming_budget == "off") - config.weight_streaming_budget.clear(); + config.transformer.initial_weight_streaming_budget = parser.get("--ws"); + if (config.transformer.initial_weight_streaming_budget == "off") + config.transformer.initial_weight_streaming_budget.clear(); config.num_images = parse_uint(parser.get("--num-images"), "--num-images"); const std::filesystem::path output_dir = parser.get("--output"); std::filesystem::create_directories(output_dir); config.output_path = output_dir / "flux2.png"; + ConfigureFlux2Models(config, model_root, config.text_encoder.ep_cache_dir, config.text_encoder.ep_context_dir); validate_config(config); ValidateFlux2Config(config); @@ -264,16 +279,18 @@ int main(int argc, char* argv[]) try { const Flux2Config config = parse_args(argc, argv); - std::cout << "Model dir: " << config.model_dir.string() << "\n" - << "Precision: " << config.precision << "\n" + std::cout << "Model dir: " << config.text_encoder.model_path.parent_path().parent_path().string() << "\n" + << "Precision: " << config.transformer.precision << "\n" << "Output: " << config.output_path.string() << "\n" << "Prompt: " << config.prompt << "\n" << "Seed: " << config.seed << "\n" + << "RNG: " << (config.rng == din::common::Rng::CpuStd ? "cpu-std" : + config.rng == din::common::Rng::CudaPhilox ? "cuda-philox" : "cpu-philox") << "\n" << "Num images: " << config.num_images << "\n" << "Processing: " << to_string(config.processing) << "\n" << "Provider: " << to_string(config.provider) << "\n" - << "EP cache: " << config.ep_cache_dir.string() << "\n" - << "EP context dir: " << config.ep_context_dir.string() << "\n" + << "EP cache: " << config.text_encoder.ep_cache_dir.string() << "\n" + << "EP context dir: " << config.text_encoder.ep_context_dir.string() << "\n" << std::endl; auto pipeline = CreateFlux2Pipeline(config); diff --git a/image_gen/flux2/flux2_cli.h b/image_gen/flux2/flux2_cli.h index 9561370..996dc24 100644 --- a/image_gen/flux2/flux2_cli.h +++ b/image_gen/flux2/flux2_cli.h @@ -3,10 +3,13 @@ #pragma once +#include "rng.h" + #include #include #include #include +#include #include #include @@ -32,6 +35,15 @@ enum class Flux2TextEncoder Qwen3_06BTranslator }; +inline void ValidateFlux2Rng(din::common::Rng rng, Flux2ProcessingBackend processing) +{ + using din::common::Rng; + if (rng != Rng::CpuStd && rng != Rng::CudaPhilox && rng != Rng::CpuPhilox) + throw std::invalid_argument("Unknown random number generator"); + if (rng == Rng::CudaPhilox && processing != Flux2ProcessingBackend::Cuda) + throw std::invalid_argument("cuda-philox RNG requires --processing cuda"); +} + // Called on the generation thread. Completed denoising steps are in [0, total]. using Flux2Progress = std::function; @@ -45,21 +57,56 @@ struct Flux2Timings double per_step_ms = 0; }; +struct Flux2ModelConfig +{ + std::filesystem::path model_path; + std::filesystem::path ep_cache_dir = "artifacts/flux2/trt_rtx_cache"; + std::filesystem::path ep_context_dir = "artifacts/flux2/ep_context"; + std::string cache_subpath; + bool operator==(const Flux2ModelConfig&) const = default; +}; + +struct Flux2TextEncoderConfig : Flux2ModelConfig +{ + Flux2TextEncoder variant = Flux2TextEncoder::Qwen3_4B; + std::filesystem::path tokenizer_dir; + bool operator==(const Flux2TextEncoderConfig&) const = default; +}; + +struct Flux2TransformerConfig : Flux2ModelConfig +{ + std::string precision = "bf16"; + // Empty keeps a non-streaming engine; -1 enables automatic streaming. + std::string initial_weight_streaming_budget; + bool operator==(const Flux2TransformerConfig&) const = default; +}; + +struct Flux2VaeConfig : Flux2ModelConfig +{ + std::filesystem::path latent_stats_path; + bool operator==(const Flux2VaeConfig&) const = default; +}; + +// Recoverable: the affected stage is unavailable; other sessions remain loaded. +class Flux2StageLoadError : public std::runtime_error +{ +public: + using std::runtime_error::runtime_error; +}; + struct Flux2Config { Flux2ProcessingBackend processing = Flux2ProcessingBackend::Cuda; Flux2ExecutionProvider provider = Flux2ExecutionProvider::TrtRtx; - std::filesystem::path model_dir; - std::string precision = "bf16"; - Flux2TextEncoder text_encoder = Flux2TextEncoder::Qwen3_4B; + Flux2TextEncoderConfig text_encoder; + Flux2TransformerConfig transformer; + Flux2VaeConfig vae; int steps = 4; - // Startup budget: empty keeps a non-streaming engine; -1 enables automatic streaming. - std::string weight_streaming_budget; - std::filesystem::path ep_cache_dir = "artifacts/flux2/trt_rtx_cache"; - std::filesystem::path ep_context_dir = "artifacts/flux2/ep_context"; std::filesystem::path output_path; std::string prompt; unsigned int seed = 42; + // CPU standard noise, CUDA Philox, or CPU Philox with fixed H100 SXM5 ordering. + din::common::Rng rng = din::common::Rng::CpuStd; unsigned int num_images = 5; }; @@ -67,6 +114,9 @@ struct Flux2GenerationOptions { // Omit to restore the startup budget. Applied via transformer session dynamic options only. std::optional weight_streaming_budget; + // Omit to use the startup selections. Overrides apply to this generation only. + std::optional steps; + std::optional rng; }; struct Flux2Image @@ -82,9 +132,12 @@ class Flux2ProcessingPipeline public: virtual ~Flux2ProcessingPipeline() = default; virtual void Initialize() = 0; + virtual void UpdateModels(const Flux2Config& config) = 0; virtual void SetPrompt(std::string prompt) = 0; // Opaque identity for diagnostics; never dereference or retain after pipeline destruction. virtual const void* TransformerSessionIdentity() const = 0; + virtual const void* TextEncoderSessionIdentity() const = 0; + virtual const void* VaeSessionIdentity() const = 0; virtual Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress = {}, const Flux2GenerationOptions& options = {}) = 0; }; @@ -92,4 +145,7 @@ class Flux2ProcessingPipeline std::unique_ptr CreateFlux2Pipeline(const Flux2Config& config); void ValidateFlux2Config(const Flux2Config& config); +// Resolve existing CLI/UI root controls into independent stage settings. +void ConfigureFlux2Models(Flux2Config& config, const std::filesystem::path& model_root, + const std::filesystem::path& cache_root, const std::filesystem::path& context_root); bool IsFlux2BackendAvailable(Flux2ProcessingBackend backend); diff --git a/image_gen/flux2/flux2_cuda_backend.cpp b/image_gen/flux2/flux2_cuda_backend.cpp index b1c1f74..eed4d05 100644 --- a/image_gen/flux2/flux2_cuda_backend.cpp +++ b/image_gen/flux2/flux2_cuda_backend.cpp @@ -25,8 +25,10 @@ #include "cu_helper.h" #include "cuda_kernels.h" +#include "rng_cuda.h" #include "flux2.h" #include "flux2_runtime_context.h" +#include "flux2_models.h" #include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" @@ -377,6 +379,9 @@ struct CudaPipelineState std::unique_ptr vae_decoder_io; std::vector time_schedule; + Flux2LoadedModels loaded; + din::common::ModelProfile model_profile; + std::optional applied_budget; EulerStage euler; PostprocessStage postprocess; }; @@ -477,11 +482,6 @@ Ort::ConstEpDevice find_cpu_device(Ort::Env& env) void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const Flux2Config& config, ExecutionProviderMode provider_mode, SamplingBackend sampling_backend) { - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); - const auto latent_stats = - LoadFlux2LatentStats(model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); - const Flux2ModelCachePaths cache_paths = - MakeFlux2ModelCachePaths(config.precision, "", config.text_encoder, !config.weight_streaming_budget.empty()); const bool trt_rtx_device = is_trt_rtx_device(ep_device); const bool cpu_device = is_cpu_device(ep_device); const bool has_separate_binding = !cpu_device; @@ -513,107 +513,56 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const Ort::SyncStream* compute_stream_ptr = state.compute_stream ? &*state.compute_stream : nullptr; - std::cout << "=== Loading ONNX Models ===" << std::endl; - std::string provider; - if (provider_mode == ExecutionProviderMode::Trt) - { - provider = "trt-rtx"; - } - else - { - provider = "cpu"; - } - din::common::EpContextOptions ep_context; - ep_context.output_dir = config.ep_context_dir.string(); - const std::string cache_dir = config.ep_cache_dir.string(); - - auto make_profile = [&](std::string cache_subpath, bool transformer = false) - { - din::common::ModelProfile profile; - profile.cache_subpath = std::move(cache_subpath); - profile.enable_cuda_graph = false; - profile.embed_ep_context = false; - profile.extra_ep_options.emplace_back("nv_length_aux_stream_array", "0"); - if (transformer && !config.weight_streaming_budget.empty()) - profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.weight_streaming_budget); - return profile; - }; - - state.text_encoder_runner = std::make_unique( - state.env, model_paths.text_encoder_model.string(), provider, cache_dir, ep_context, - make_profile(cache_paths.text_encoder), compute_stream_ptr); - std::cout << " Text encoder loaded" << std::endl; - state.transformer_runner = std::make_unique( - state.env, model_paths.transformer_model.string(), provider, cache_dir, ep_context, - make_profile(cache_paths.transformer, true), compute_stream_ptr); - std::cout << " Transformer loaded" << std::endl; - state.vae_decoder_runner = - std::make_unique(state.env, model_paths.vae_decoder_model.string(), provider, cache_dir, - ep_context, make_profile(cache_paths.vae_decoder), compute_stream_ptr); - std::cout << " VAE decoder loaded" << std::endl; + state.model_profile.enable_cuda_graph = false; + state.model_profile.embed_ep_context = false; + state.model_profile.extra_ep_options.emplace_back("nv_length_aux_stream_array", "0"); std::vector token_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; - state.token = std::make_unique>(*state.text_encoder_runner, token_shape, + state.token = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, token_shape, has_separate_binding); std::vector attn_mask_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; - state.attention_mask = std::make_unique>(*state.text_encoder_runner, + state.attention_mask = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, attn_mask_shape, has_separate_binding); std::vector te_out_shape = {BATCH_SIZE, SEQUENCE_LENGTH, TEXT_ENCODER_EMBED_DIM}; - state.text_encoder_embeds = std::make_unique>(*state.text_encoder_runner, + state.text_encoder_embeds = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, te_out_shape, has_separate_binding); std::vector timestep_shape = {BATCH_SIZE}; - state.timestep = std::make_unique>(*state.transformer_runner, timestep_shape, + state.timestep = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, timestep_shape, has_separate_binding); std::vector hidden_shape = {BATCH_SIZE, TRANSFORMER_HIDDEN_DIM, LATENT_CHANNELS}; - state.hidden_states = std::make_unique>(*state.transformer_runner, hidden_shape, + state.hidden_states = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, hidden_shape, has_separate_binding); std::vector img_ids_shape = {BATCH_SIZE, IMAGE_SEQUENCE, 4}; - state.img_ids = std::make_unique>(*state.transformer_runner, img_ids_shape, + state.img_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, img_ids_shape, has_separate_binding); std::vector txt_ids_shape = {BATCH_SIZE, SEQUENCE_LENGTH, 4}; - state.txt_ids = std::make_unique>(*state.transformer_runner, txt_ids_shape, + state.txt_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, txt_ids_shape, has_separate_binding); - state.transformer_output = std::make_unique>(*state.transformer_runner, + state.transformer_output = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, hidden_shape, has_separate_binding); std::vector dec_latent_shape = {BATCH_SIZE, LATENT_CHANNELS / PATCH_SIZE / PATCH_SIZE, LATENT_HEIGHT * PATCH_SIZE, LATENT_WIDTH * PATCH_SIZE}; - state.decoder_latent = std::make_unique>(*state.vae_decoder_runner, + state.decoder_latent = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, dec_latent_shape, has_separate_binding); std::vector image_shape = {BATCH_SIZE, IMAGE_CHANNELS, IMAGE_HEIGHT, IMAGE_WIDTH}; - state.decoded_image = std::make_unique>(*state.vae_decoder_runner, image_shape, + state.decoded_image = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, image_shape, has_separate_binding); std::vector bn_shape = {LATENT_CHANNELS}; state.bn_mean = - std::make_unique>(*state.transformer_runner, bn_shape, has_separate_binding); + std::make_unique>(state.env, state.ep_device, compute_stream_ptr, bn_shape, has_separate_binding); state.bn_std = - std::make_unique>(*state.transformer_runner, bn_shape, has_separate_binding); + std::make_unique>(state.env, state.ep_device, compute_stream_ptr, bn_shape, has_separate_binding); - state.text_encoder_io = std::make_unique(state.text_encoder_runner->session); - state.text_encoder_io->BindInput("input_ids", state.token->BindingValue()); - state.text_encoder_io->BindInput("attention_mask", state.attention_mask->BindingValue()); - state.text_encoder_io->BindOutput("prompt_embeds", state.text_encoder_embeds->BindingValue()); - - state.transformer_io = std::make_unique(state.transformer_runner->session); - state.transformer_io->BindInput("hidden_states", state.hidden_states->BindingValue()); - state.transformer_io->BindInput("encoder_hidden_states", state.text_encoder_embeds->BindingValue()); - state.transformer_io->BindInput("timestep", state.timestep->BindingValue()); - state.transformer_io->BindInput("img_ids", state.img_ids->BindingValue()); - state.transformer_io->BindInput("txt_ids", state.txt_ids->BindingValue()); - state.transformer_io->BindOutput("sample", state.transformer_output->BindingValue()); - - state.vae_decoder_io = std::make_unique(state.vae_decoder_runner->session); - state.vae_decoder_io->BindInput("latent_sample", state.decoder_latent->BindingValue()); - state.vae_decoder_io->BindOutput("sample", state.decoded_image->BindingValue()); std::vector img_ids_cpu(state.img_ids->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount()); std::vector txt_ids_cpu(state.txt_ids->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount()); @@ -621,26 +570,13 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const std::copy(img_ids_cpu.begin(), img_ids_cpu.end(), state.img_ids->HostData()); std::copy(txt_ids_cpu.begin(), txt_ids_cpu.end(), state.txt_ids->HostData()); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state.token->HostData(), - state.attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); - std::copy_n(latent_stats.mean.data(), LATENT_CHANNELS, state.bn_mean->HostData()); - std::copy_n(latent_stats.std.data(), LATENT_CHANNELS, state.bn_std->HostData()); if (has_separate_binding) { auto img_ids_ready = state.img_ids->CopyAsyncToDeviceWithNotification(); auto txt_ids_ready = state.txt_ids->CopyAsyncToDeviceWithNotification(); - auto token_ready = state.token->CopyAsyncToDeviceWithNotification(); - auto attention_mask_ready = state.attention_mask->CopyAsyncToDeviceWithNotification(); - auto bn_mean_ready = state.bn_mean->CopyAsyncToDeviceWithNotification(); - auto bn_std_ready = state.bn_std->CopyAsyncToDeviceWithNotification(); img_ids_ready.Sync(); txt_ids_ready.Sync(); - token_ready.Sync(); - attention_mask_ready.Sync(); - bn_mean_ready.Sync(); - bn_std_ready.Sync(); } state.time_schedule = MakeFlux2Schedule(config.steps); @@ -656,7 +592,68 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const state.postprocess.BindOutput("latent_sample", *state.decoder_latent); } -Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed, const Flux2Progress& progress) +static void upload_cuda_prompt(CudaPipelineState& state, const Flux2Config& config) +{ + const Flux2ModelPaths model_paths = + MakeFlux2ModelPaths(config); + const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); + FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state.token->HostData(), + state.attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); + + if (state.has_separate_binding) + { + auto token_ready = state.token->CopyAsyncToDeviceWithNotification(); + auto attention_mask_ready = state.attention_mask->CopyAsyncToDeviceWithNotification(); + token_ready.Sync(); + attention_mask_ready.Sync(); + } +} + +static void update_cuda_models(CudaPipelineState& state, const Flux2Config& config) +{ + if (!state.loaded.Matches(config)) + { + if (state.compute_stream) signal_and_wait_on_host(*state.compute_stream); + } + const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; + Ort::SyncStream* stream = state.compute_stream ? &*state.compute_stream : nullptr; + if (state.loaded.text_encoder != config.text_encoder) state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, + state.text_encoder_runner, state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"input_ids", &state.token->BindingValue()}, {"attention_mask", &state.attention_mask->BindingValue()}}, {{"prompt_embeds", &state.text_encoder_embeds->BindingValue()}}); + upload_cuda_prompt(state, config); + return binding; + }); + auto transformer_profile = state.model_profile; + if (!config.transformer.initial_weight_streaming_budget.empty()) + transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", + config.transformer.initial_weight_streaming_budget); + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, + state.transformer_runner, state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) { return BindFlux2Model(session, {{"hidden_states", &state.hidden_states->BindingValue()}, {"encoder_hidden_states", &state.text_encoder_embeds->BindingValue()}, {"timestep", &state.timestep->BindingValue()}, {"img_ids", &state.img_ids->BindingValue()}, {"txt_ids", &state.txt_ids->BindingValue()}}, {{"sample", &state.transformer_output->BindingValue()}}); })) + state.applied_budget.reset(); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, + state.env, provider, stream, state.model_profile, [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_latent->BindingValue()}}, {{"sample", &state.decoded_image->BindingValue()}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + std::copy_n(stats.mean.data(), LATENT_CHANNELS, state.bn_mean->HostData()); + std::copy_n(stats.std.data(), LATENT_CHANNELS, state.bn_std->HostData()); + if (state.has_separate_binding) + { + auto mean_ready = state.bn_mean->CopyAsyncToDeviceWithNotification(); + auto std_ready = state.bn_std->CopyAsyncToDeviceWithNotification(); + mean_ready.Sync(); + std_ready.Sync(); + } + return binding; + }); +} + +Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed, din::common::Rng rng, + const Flux2Progress& progress) { Ort::SyncStream* compute_stream_ptr = state.compute_stream ? &*state.compute_stream : nullptr; @@ -680,8 +677,19 @@ Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed, nvtx3::scoped_range nvtx("initialize_latents"); std::cout << "\n=== Initializing Latent (seed=" << seed << ") ===" << std::endl; const size_t hidden_count = state.hidden_states->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount(); - initialize_latent(state.hidden_states->HostData(), hidden_count, seed); - if (state.binding_is_cuda) + if (rng == din::common::Rng::CudaPhilox) + { + if (!state.binding_is_cuda || compute_stream_ptr == nullptr) + throw std::runtime_error("cuRAND RNG requires device-backed bindings and an ORT CUDA stream"); + const auto stream = reinterpret_cast(compute_stream_ptr->GetHandle()); + din::common::GenerateNormalCudaPacked( + stream, state.hidden_states->BindingValue().GetTensorMutableData(), hidden_count, seed, + LATENT_CHANNELS, LATENT_HEIGHT * LATENT_WIDTH, true); + } + else + initialize_cpu_latent(state.hidden_states->HostData(), hidden_count, seed, rng, + BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); + if (state.binding_is_cuda && rng != din::common::Rng::CudaPhilox) { auto hidden_uploaded = state.hidden_states->CopyAsyncToDeviceWithNotification(); hidden_uploaded.Sync(); @@ -720,6 +728,7 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline { if (state_) { + UpdateModels(config_); return; } @@ -733,16 +742,14 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline throw std::runtime_error("CUDA processing requires --provider trt-rtx"); } - std::cout << "Model dir: " << config_.model_dir.string() << "\n" - << "Execution provider: " << to_string(provider) << "\n" - << "Sampling backend: " << to_string(sampling_backend) << "\n" - << std::endl; state_ = std::make_unique(runtime_.env); if (provider == ExecutionProviderMode::Cpu) { std::cout << "Using CPU execution provider" << std::endl; - initialize_ep(*state_, find_cpu_device(state_->env), config_, provider, sampling_backend); + try { initialize_ep(*state_, find_cpu_device(state_->env), config_, provider, sampling_backend); } + catch (...) { state_.reset(); throw; } + UpdateModels(config_); return; } @@ -753,7 +760,21 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } std::cout << "Using TensorRT RTX execution provider" << std::endl; - initialize_ep(*state_, trt_device, config_, provider, sampling_backend); + try { initialize_ep(*state_, trt_device, config_, provider, sampling_backend); } + catch (...) { state_.reset(); throw; } + UpdateModels(config_); + } + + void UpdateModels(const Flux2Config& requested) override + { + if (!state_) throw std::runtime_error("Pipeline resources are not initialized"); + if (requested.processing != config_.processing || requested.provider != config_.provider) + throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); + // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. + const auto previous_prompt = config_.prompt; + config_ = requested; + config_.prompt = previous_prompt; + update_cuda_models(*state_, requested); } void SetPrompt(std::string prompt) override @@ -767,48 +788,48 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = - MakeFlux2ModelPaths(config_.model_dir, config_.precision, config_.text_encoder); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state_->token->HostData(), - state_->attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); - - if (state_->has_separate_binding) - { - auto token_ready = state_->token->CopyAsyncToDeviceWithNotification(); - auto attention_mask_ready = state_->attention_mask->CopyAsyncToDeviceWithNotification(); - token_ready.Sync(); - attention_mask_ready.Sync(); - } + upload_cuda_prompt(*state_, config_); } const void* TransformerSessionIdentity() const override { - return state_ ? static_cast(state_->transformer_runner->session) : nullptr; + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) : nullptr; + } + + const void* TextEncoderSessionIdentity() const override + { + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) : nullptr; + } + + const void* VaeSessionIdentity() const override + { + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) : nullptr; } Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, const Flux2GenerationOptions& options) override { - if (!state_) + if (!state_ || !state_->loaded.Ready()) { throw std::runtime_error("Flux2 CUDA pipeline was not initialized"); } - const std::string budget = options.weight_streaming_budget.value_or(config_.weight_streaming_budget); - const std::string& previous = applied_budget_ ? *applied_budget_ : config_.weight_streaming_budget; + const din::common::Rng rng = options.rng.value_or(config_.rng); + ValidateFlux2Rng(rng, config_.processing); + state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); + const std::string budget = options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; if (budget != previous) { const char* keys[] = {"nv_weight_streaming_budget"}; const char* values[] = {budget.c_str()}; state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); - applied_budget_ = budget; + state_->applied_budget = budget; } - return run_initialized_pipeline(*state_, seed, progress); + return run_initialized_pipeline(*state_, seed, rng, progress); } private: Flux2Config config_; - std::optional applied_budget_; Flux2RuntimeContext& runtime_; std::unique_ptr state_; }; diff --git a/image_gen/flux2/flux2_dx_backend.cpp b/image_gen/flux2/flux2_dx_backend.cpp index ae06dbe..b413021 100644 --- a/image_gen/flux2/flux2_dx_backend.cpp +++ b/image_gen/flux2/flux2_dx_backend.cpp @@ -19,6 +19,7 @@ #include "dx_helper.h" #include "flux2.h" #include "flux2_runtime_context.h" +#include "flux2_models.h" #include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" @@ -449,6 +450,9 @@ struct DxPipelineState std::unique_ptr transformer_io; std::unique_ptr vae_decoder_io; std::vector time_schedule; + Flux2LoadedModels loaded; + din::common::ModelProfile model_profile; + std::optional applied_budget; }; static void initialize_dx_state(DxPipelineState& state, const Flux2Config& config, Ort::ConstEpDevice trt_device) @@ -458,16 +462,7 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi { throw std::runtime_error("DirectX processing requires --provider trt-rtx"); } - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); - const auto latent_stats = - LoadFlux2LatentStats(model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); - const Flux2ModelCachePaths cache_paths = - MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "dx_cig" : "dx", config.text_encoder, - !config.weight_streaming_budget.empty()); - std::cout << "Model dir: " << model_paths.base_dir.string() << "\n" - << "CIG: " << (state.use_cig ? "enabled" : "disabled") << "\n" - << std::endl; std::cout << "\n=== Initializing D3D12 ===" << std::endl; state.trt_device = trt_device; @@ -500,31 +495,10 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi graphics_ep_options.emplace_back("nv_length_aux_stream_array", "0"); } - std::cout << "=== Loading ONNX Models ===" << std::endl; - din::common::EpContextOptions ep_context; - ep_context.output_dir = config.ep_context_dir.string(); - const std::string cache_dir = config.ep_cache_dir.string(); - auto make_profile = [&](std::string cache_subpath, bool transformer = false) - { - din::common::ModelProfile profile; - profile.cache_subpath = std::move(cache_subpath); - profile.enable_cuda_graph = state.use_cig ? false : true; - profile.embed_ep_context = false; - profile.extra_ep_options = graphics_ep_options; - if (transformer && !config.weight_streaming_budget.empty()) - profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.weight_streaming_budget); - return profile; - }; + state.model_profile.enable_cuda_graph = !state.use_cig; + state.model_profile.embed_ep_context = false; + state.model_profile.extra_ep_options = graphics_ep_options; - state.text_encoder_runner = std::make_unique( - state.env, model_paths.text_encoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.text_encoder), &*state.sync_stream); - state.transformer_runner = std::make_unique( - state.env, model_paths.transformer_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.transformer, true), &*state.sync_stream); - state.vae_decoder_runner = std::make_unique( - state.env, model_paths.vae_decoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.vae_decoder), &*state.sync_stream); std::vector token_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; std::vector attn_mask_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; @@ -583,27 +557,9 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.decoded_image_tensor = state.tensor_importer->create_tensor(state.decoded_image_buf, image_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); - state.text_encoder_io = std::make_unique(state.text_encoder_runner->session); - state.text_encoder_io->BindInput("input_ids", state.token_tensor); - state.text_encoder_io->BindInput("attention_mask", state.attention_mask_tensor); - state.text_encoder_io->BindOutput("prompt_embeds", state.text_encoder_embeds_tensor); - - state.transformer_io = std::make_unique(state.transformer_runner->session); - state.transformer_io->BindInput("hidden_states", state.hidden_states_tensor); - state.transformer_io->BindInput("encoder_hidden_states", state.text_encoder_embeds_tensor); - state.transformer_io->BindInput("timestep", state.timestep_tensor); - state.transformer_io->BindInput("img_ids", state.img_ids_tensor); - state.transformer_io->BindInput("txt_ids", state.txt_ids_tensor); - state.transformer_io->BindOutput("sample", state.transformer_output_tensor); - - state.vae_decoder_io = std::make_unique(state.vae_decoder_runner->session); - state.vae_decoder_io->BindInput("latent_sample", state.decoder_input_tensor); - state.vae_decoder_io->BindOutput("sample", state.decoded_image_tensor); std::cout << "\n=== Uploading Static Inputs (D3D12) ===" << std::endl; state.dx.begin(); - state.dx.record_upload(state.bn_mean_buf, latent_stats.mean.data(), state.bn_mean_buf.size); - state.dx.record_upload(state.bn_std_buf, latent_stats.std.data(), state.bn_std_buf.size); { std::vector img_ids_cpu(shape_numel(img_ids_shape)); @@ -643,24 +599,73 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.time_schedule = MakeFlux2Schedule(config.steps); - std::cout << "\n=== Tokenizing Prompt ===" << std::endl; - { - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); - - std::vector tokens_cpu(shape_numel(token_shape)); - std::vector attn_mask_cpu(shape_numel(attn_mask_shape)); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); - - state.dx.record_upload(state.token_buf, tokens_cpu.data(), state.token_buf.size); - state.dx.record_upload(state.attn_mask_buf, attn_mask_cpu.data(), state.attn_mask_buf.size); - } state.dx.execute(); const uint64_t static_uploads_done = state.sync_fence->signal_d3d12(state.dx.queue.Get()); state.sync_fence->wait_ort(*state.sync_stream, static_uploads_done); } -static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, const Flux2Progress& progress) +static void upload_dx_prompt(DxPipelineState& state, const Flux2Config& config) +{ + const Flux2ModelPaths model_paths = + MakeFlux2ModelPaths(config); + const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); + std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), + BATCH_SIZE, SEQUENCE_LENGTH); + + state.dx.begin(); + state.dx.record_upload(state.token_buf, tokens_cpu.data(), state.token_buf.size); + state.dx.record_upload(state.attn_mask_buf, attn_mask_cpu.data(), state.attn_mask_buf.size); + state.dx.execute(); + const uint64_t prompt_uploaded = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, prompt_uploaded); +} + +static void update_dx_models(DxPipelineState& state, const Flux2Config& config) +{ + if (!state.loaded.Matches(config)) + { + const auto graphics_done = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, graphics_done); + signal_and_wait_on_host(*state.sync_stream); + } + const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; + Ort::SyncStream* stream = &*state.sync_stream; + if (state.loaded.text_encoder != config.text_encoder) state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, + state.text_encoder_runner, state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); + upload_dx_prompt(state, config); + return binding; + }); + auto transformer_profile = state.model_profile; + if (!config.transformer.initial_weight_streaming_budget.empty()) + transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", + config.transformer.initial_weight_streaming_budget); + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, + state.transformer_runner, state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) { return BindFlux2Model(session, {{"hidden_states", &state.hidden_states_tensor}, {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, {"timestep", &state.timestep_tensor}, {"img_ids", &state.img_ids_tensor}, {"txt_ids", &state.txt_ids_tensor}}, {{"sample", &state.transformer_output_tensor}}); })) + state.applied_budget.reset(); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, + state.env, provider, stream, state.model_profile, [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, {{"sample", &state.decoded_image_tensor}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + state.dx.begin(); + state.dx.record_upload(state.bn_mean_buf, stats.mean.data(), state.bn_mean_buf.size); + state.dx.record_upload(state.bn_std_buf, stats.std.data(), state.bn_std_buf.size); + state.dx.execute(); + const auto uploaded = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, uploaded); + signal_and_wait_on_host(*state.sync_stream); + return binding; + }); +} + +static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, din::common::Rng rng, const Flux2Progress& progress) { Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, [&] @@ -670,7 +675,8 @@ static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, const signal_and_wait_on_host(*state.sync_stream); }); std::vector latent_cpu(state.hidden_count); - initialize_latent(latent_cpu.data(), latent_cpu.size(), seed); + initialize_cpu_latent(latent_cpu.data(), latent_cpu.size(), seed, rng, + BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); state.dx.begin(); state.dx.record_upload(state.hidden_states, latent_cpu.data(), state.hidden_states.size); state.dx.execute(); @@ -712,12 +718,27 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline { if (state_) { + UpdateModels(config_); return; } std::cout << "Initializing Flux2 DirectX pipeline" << std::endl; const auto identity = din::common::ResolveOrtGraphicsDeviceIdentity(runtime_.trt_device); state_ = std::make_unique(runtime_.env, identity.luid); - initialize_dx_state(*state_, config_, runtime_.trt_device); + try { initialize_dx_state(*state_, config_, runtime_.trt_device); } + catch (...) { state_.reset(); throw; } + UpdateModels(config_); + } + + void UpdateModels(const Flux2Config& requested) override + { + if (!state_) throw std::runtime_error("Pipeline resources are not initialized"); + if (requested.processing != config_.processing || requested.provider != config_.provider) + throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); + // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. + const auto previous_prompt = config_.prompt; + config_ = requested; + config_.prompt = previous_prompt; + update_dx_models(*state_, requested); } void SetPrompt(std::string prompt) override @@ -731,49 +752,48 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = - MakeFlux2ModelPaths(config_.model_dir, config_.precision, config_.text_encoder); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); - std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); - - state_->dx.begin(); - state_->dx.record_upload(state_->token_buf, tokens_cpu.data(), state_->token_buf.size); - state_->dx.record_upload(state_->attn_mask_buf, attn_mask_cpu.data(), state_->attn_mask_buf.size); - state_->dx.execute(); - const uint64_t prompt_uploaded = state_->sync_fence->signal_d3d12(state_->dx.queue.Get()); - state_->sync_fence->wait_ort(*state_->sync_stream, prompt_uploaded); + upload_dx_prompt(*state_, config_); } const void* TransformerSessionIdentity() const override { - return state_ ? static_cast(state_->transformer_runner->session) : nullptr; + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) : nullptr; + } + + const void* TextEncoderSessionIdentity() const override + { + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) : nullptr; + } + + const void* VaeSessionIdentity() const override + { + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) : nullptr; } Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, const Flux2GenerationOptions& options) override { - if (!state_) + if (!state_ || !state_->loaded.Ready()) { throw std::runtime_error("Flux2 DirectX pipeline was not initialized"); } - const std::string budget = options.weight_streaming_budget.value_or(config_.weight_streaming_budget); - const std::string& previous = applied_budget_ ? *applied_budget_ : config_.weight_streaming_budget; + const din::common::Rng rng = options.rng.value_or(config_.rng); + ValidateFlux2Rng(rng, config_.processing); + state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); + const std::string budget = options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; if (budget != previous) { const char* keys[] = {"nv_weight_streaming_budget"}; const char* values[] = {budget.c_str()}; state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); - applied_budget_ = budget; + state_->applied_budget = budget; } - return run_dx_image(*state_, seed, progress); + return run_dx_image(*state_, seed, rng, progress); } private: Flux2Config config_; - std::optional applied_budget_; Flux2RuntimeContext& runtime_; std::unique_ptr state_; }; diff --git a/image_gen/flux2/flux2_models.cpp b/image_gen/flux2/flux2_models.cpp new file mode 100644 index 0000000..f4998ee --- /dev/null +++ b/image_gen/flux2/flux2_models.cpp @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include "flux2_models.h" + +void ConfigureFlux2Models(Flux2Config& config, const std::filesystem::path& model_root, + const std::filesystem::path& cache_root, const std::filesystem::path& context_root) +{ + const auto root = model_root.empty() ? model_root : std::filesystem::weakly_canonical(model_root); + std::string prefix; + switch (config.processing) + { + case Flux2ProcessingBackend::Dx: prefix = "dx"; break; + case Flux2ProcessingBackend::DxCig: prefix = "dx_cig"; break; + case Flux2ProcessingBackend::Vk: prefix = "vk"; break; + case Flux2ProcessingBackend::VkCig: prefix = "vk_cig"; break; + default: break; + } + if (!prefix.empty()) prefix += '_'; + const std::string encoder = config.text_encoder.variant == Flux2TextEncoder::Qwen3_4B + ? "text_encoder" : "text_encoder_translator"; + const std::string transformer = "transformer_" + config.transformer.precision; + for (Flux2ModelConfig* model : {static_cast(&config.text_encoder), + static_cast(&config.transformer), + static_cast(&config.vae)}) + { + model->ep_cache_dir = std::filesystem::weakly_canonical(cache_root); + model->ep_context_dir = std::filesystem::weakly_canonical(context_root); + } + config.text_encoder.cache_subpath = prefix + encoder; + config.transformer.cache_subpath = prefix + transformer + + (config.transformer.initial_weight_streaming_budget.empty() ? "" : "_ws"); + config.vae.cache_subpath = prefix + "vae_decoder"; + if (root.empty()) return; // The studio can start before a model root is selected. + config.text_encoder.model_path = root / encoder / "model.onnx"; + config.text_encoder.tokenizer_dir = root / "tokenizer"; + config.transformer.model_path = root / transformer / "model.onnx"; + config.vae.model_path = root / "vae_decoder/model.onnx"; + config.vae.latent_stats_path = root / "vae_decoder/latent_stats.json"; +} + +std::unique_ptr BindFlux2Model(Ort::Session& session, + std::initializer_list inputs, std::initializer_list outputs) +{ + auto binding = std::make_unique(session); + Ort::AllocatorWithDefaultOptions allocator; + auto bind = [&](std::initializer_list tensors, bool input) + { + const size_t count = input ? session.GetInputCount() : session.GetOutputCount(); + if (count != tensors.size()) throw std::runtime_error("Unexpected model input/output count"); + for (const auto& tensor : tensors) + { + size_t index = 0; + for (; index < count; ++index) + { + const auto name = input ? session.GetInputNameAllocated(index, allocator) + : session.GetOutputNameAllocated(index, allocator); + if (std::string(name.get()) == tensor.name) break; + } + if (index == count) throw std::runtime_error(std::string("Missing tensor: ") + tensor.name); + const auto type = input ? session.GetInputTypeInfo(index) : session.GetOutputTypeInfo(index); + const auto info = type.GetTensorTypeAndShapeInfo(); + const auto actual = tensor.tensor->GetTensorTypeAndShapeInfo(); + const auto shape = info.GetShape(); + const auto expected = actual.GetShape(); + if (info.GetElementType() != actual.GetElementType() || shape.size() != expected.size()) + throw std::runtime_error(std::string("Incompatible tensor type/rank: ") + tensor.name); + for (size_t axis = 0; axis < shape.size(); ++axis) + if (shape[axis] >= 0 && shape[axis] != expected[axis]) + throw std::runtime_error(std::string("Incompatible tensor shape: ") + tensor.name); + if (input) binding->BindInput(tensor.name, *tensor.tensor); + else binding->BindOutput(tensor.name, *tensor.tensor); + } + }; + bind(inputs, true); + bind(outputs, false); + return binding; +} diff --git a/image_gen/flux2/flux2_models.h b/image_gen/flux2/flux2_models.h new file mode 100644 index 0000000..d44c765 --- /dev/null +++ b/image_gen/flux2/flux2_models.h @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include "flux2_cli.h" +#include "ort_session.h" +#include +#include + +struct Flux2TensorBinding +{ + const char* name; + Ort::Value* tensor; +}; + +std::unique_ptr BindFlux2Model(Ort::Session& session, + std::initializer_list inputs, std::initializer_list outputs); + +struct Flux2LoadedModels +{ + std::optional text_encoder; + std::optional transformer; + std::optional vae; + bool Matches(const Flux2Config& config) const + { + return text_encoder == config.text_encoder && transformer == config.transformer && vae == config.vae; + } + bool Ready() const { return text_encoder && transformer && vae; } +}; + +// The caller drains device work before entering this function. Tensors and the +// externally owned stream remain alive across session replacement. +template +bool ReloadFlux2Stage(const char* name, const Config& requested, std::optional& loaded, + std::unique_ptr& runner, std::unique_ptr& binding, + Ort::Env& env, const char* provider, Ort::SyncStream* stream, + din::common::ModelProfile profile, Bind&& bind) +{ + if (runner && binding && loaded == requested) + { + std::cout << "Reusing " << name << '\n'; + return false; + } + std::cout << "Loading " << name << ": " << requested.model_path.string() << '\n'; + binding.reset(); + runner.reset(); + loaded.reset(); + try + { + profile.cache_subpath = requested.cache_subpath; + runner = std::make_unique(env, requested.model_path.string(), provider, + requested.ep_cache_dir.string(), din::common::EpContextOptions{requested.ep_context_dir.string()}, + profile, stream); + binding = bind(runner->session); + loaded = requested; + } + catch (const std::exception& error) + { + binding.reset(); + runner.reset(); + loaded.reset(); + throw Flux2StageLoadError(std::string("Failed to load ") + name + " (" + + requested.model_path.string() + "): " + error.what()); + } + return true; +} diff --git a/image_gen/flux2/flux2_pipeline.cpp b/image_gen/flux2/flux2_pipeline.cpp index 9c4efe2..aed1da4 100644 --- a/image_gen/flux2/flux2_pipeline.cpp +++ b/image_gen/flux2/flux2_pipeline.cpp @@ -6,6 +6,7 @@ #include #include "flux2_runtime_context.h" +#include "rng.h" namespace { @@ -91,14 +92,15 @@ void ValidateFlux2Config(const Flux2Config& config) throw std::invalid_argument("Denoise steps must be in [1, 50]"); if (config.num_images == 0) throw std::invalid_argument("Image count must be positive"); - if (config.precision != "bf16" && config.precision != "fp16" && config.precision != "fp8" && - config.precision != "nvfp4") + ValidateFlux2Rng(config.rng, config.processing); + if (config.transformer.precision != "bf16" && config.transformer.precision != "fp16" && config.transformer.precision != "fp8" && + config.transformer.precision != "nvfp4") throw std::invalid_argument("Precision must be bf16, fp16, fp8, or nvfp4"); - if (!config.weight_streaming_budget.empty()) + if (!config.transformer.initial_weight_streaming_budget.empty()) { if (config.provider != Flux2ExecutionProvider::TrtRtx) throw std::invalid_argument("Weight streaming requires TensorRT RTX"); - const auto& budget = config.weight_streaming_budget; + const auto& budget = config.transformer.initial_weight_streaming_budget; if (budget == "-1") return; size_t parsed = 0; diff --git a/image_gen/flux2/flux2_ui.cpp b/image_gen/flux2/flux2_ui.cpp index 1ba925c..6303800 100644 --- a/image_gen/flux2/flux2_ui.cpp +++ b/image_gen/flux2/flux2_ui.cpp @@ -98,12 +98,9 @@ class Generator std::vector ready; private: - static bool SameModels(const Flux2Config& a, const Flux2Config& b) + static bool SameBackend(const Flux2Config& a, const Flux2Config& b) { - return a.model_dir == b.model_dir && a.processing == b.processing && a.provider == b.provider && - a.precision == b.precision && a.text_encoder == b.text_encoder && a.steps == b.steps && - a.weight_streaming_budget == b.weight_streaming_budget && a.ep_cache_dir == b.ep_cache_dir && - a.ep_context_dir == b.ep_context_dir; + return a.processing == b.processing && a.provider == b.provider; } void UpdateMemory() @@ -148,14 +145,16 @@ class Generator { const auto& config = request.config; ValidateFlux2Config(config); - if (!loaded || !SameModels(*loaded, config)) + if (!loaded || !SameBackend(*loaded, config)) { pipeline.reset(); // Release old engines before allocating the replacement. loaded.reset(); pipeline = CreateFlux2Pipeline(config); - pipeline->Initialize(); loaded = config; + pipeline->Initialize(); } + pipeline->UpdateModels(config); + loaded = config; pipeline->SetPrompt(config.prompt); for (unsigned int i = 0; i < config.num_images; ++i) { @@ -190,6 +189,12 @@ class Generator completed = total; status = "Done"; } + catch (const Flux2StageLoadError& error) + { + std::lock_guard guard(mutex); + status = std::string("Error: ") + error.what(); + failed = true; + } catch (const std::exception& error) { pipeline.reset(); @@ -240,18 +245,18 @@ int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit Flux2Config config = initial; std::array model{}, output{}, cache{}; bool cache_changed = false; - const auto cache_utf8 = initial.ep_cache_dir.parent_path().u8string(); + const auto cache_utf8 = initial.text_encoder.ep_cache_dir.parent_path().u8string(); std::snprintf(cache.data(), cache.size(), "%s", reinterpret_cast(cache_utf8.c_str())); std::array prompt{}; - const auto model_utf8 = config.model_dir.u8string(); + const auto model_utf8 = config.text_encoder.model_path.parent_path().parent_path().u8string(); std::snprintf(model.data(), model.size(), "%s", reinterpret_cast(model_utf8.c_str())); std::snprintf(output.data(), output.size(), "%s", config.output_path.string().c_str()); std::snprintf(prompt.data(), prompt.size(), "%s", config.prompt.c_str()); int streaming_mode = 0; // Auto / Manual / Off; independent of session configuration. int budget = 50; int count = static_cast(config.num_images); - int encoder = config.text_encoder == Flux2TextEncoder::Qwen3_4B ? 0 : 1; - int precision = config.precision == "bf16" ? 0 : config.precision == "fp16" ? 1 : config.precision == "fp8" ? 2 : 3; + int encoder = config.text_encoder.variant == Flux2TextEncoder::Qwen3_4B ? 0 : 1; + int precision = config.transformer.precision == "bf16" ? 0 : config.transformer.precision == "fp16" ? 1 : config.transformer.precision == "fp8" ? 2 : 3; const char* precisions[] = {"bf16", "fp16", "fp8", "nvfp4"}; const char* backends[] = {"CPU", "CUDA", "DirectX", "DirectX CIG", "Vulkan", "Vulkan CIG"}; int backend = static_cast(config.processing); @@ -332,6 +337,21 @@ int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit } ImGui::EndCombo(); } + if (backend != static_cast(Flux2ProcessingBackend::Cuda) && config.rng == din::common::Rng::CudaPhilox) + config.rng = din::common::Rng::CpuPhilox; + const char* rng_names[] = {"cpu-std", "cuda-philox", "cpu-philox"}; + if (ImGui::BeginCombo("RNG", rng_names[static_cast(config.rng)])) + { + for (int i = 0; i < 3; ++i) + { + const bool disabled = i == 1 && backend != static_cast(Flux2ProcessingBackend::Cuda); + ImGui::BeginDisabled(disabled); + if (ImGui::Selectable(rng_names[i], static_cast(config.rng) == i)) + config.rng = static_cast(i); + ImGui::EndDisabled(); + } + ImGui::EndCombo(); + } ImGui::BeginDisabled(provider == 0); ImGui::Combo("Weight streaming", &streaming_mode, "Auto\0Manual\0Off\0"); if (streaming_mode == 1) @@ -353,23 +373,26 @@ int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit glDeleteTextures(1, &result.texture); results.clear(); selected = -1; - config.model_dir = std::filesystem::u8path(model.data()); + const auto model_root = std::filesystem::u8path(model.data()); config.output_path = output.data(); config.prompt = prompt.data(); - const auto cache_namespace = CacheNamespace(config.model_dir); - config.ep_cache_dir = - (cache_changed ? std::filesystem::u8path(cache.data()) / "runtime" : initial.ep_cache_dir) / + const auto cache_namespace = CacheNamespace(model_root); + const auto cache_root = + (cache_changed ? std::filesystem::u8path(cache.data()) / "runtime" : initial.text_encoder.ep_cache_dir) / cache_namespace; - config.ep_context_dir = - (cache_changed ? std::filesystem::u8path(cache.data()) / "ep_context" : initial.ep_context_dir) / + const auto context_root = + (cache_changed ? std::filesystem::u8path(cache.data()) / "ep_context" : initial.text_encoder.ep_context_dir) / cache_namespace; - config.text_encoder = encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; - config.precision = precisions[precision]; + config.text_encoder.variant = encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + config.transformer.precision = precisions[precision]; config.num_images = static_cast(count); config.processing = static_cast(backend); config.provider = provider == 0 ? Flux2ExecutionProvider::Cpu : Flux2ExecutionProvider::TrtRtx; - config.weight_streaming_budget = provider == 0 ? "" : "-1"; + config.transformer.initial_weight_streaming_budget = provider == 0 ? "" : "-1"; + ConfigureFlux2Models(config, model_root, cache_root, context_root); Flux2GenerationOptions generation; + generation.steps = config.steps; + generation.rng = config.rng; if (provider != 0 && streaming_mode != 0) generation.weight_streaming_budget = streaming_mode == 1 ? std::to_string(budget) + "%" : "0"; generator.Submit({config, generation, random}); @@ -469,24 +492,25 @@ int main(int argc, char** argv) args.add_argument("--exit-after-generation").default_value(false).implicit_value(true); args.parse_args(argc, argv); Flux2Config config; - config.model_dir = args.get("--model-dir"); - if (config.model_dir.empty() && args.get("--auto-generate")) + const std::filesystem::path model_root = args.get("--model-dir"); + if (model_root.empty() && args.get("--auto-generate")) throw std::runtime_error("--auto-generate requires --model-dir"); config.output_path = args.get("--output"); - config.ep_cache_dir = std::filesystem::u8path(args.get("--ep-cache")); - config.ep_context_dir = std::filesystem::u8path(args.get("--ep-context-dir")); - std::filesystem::create_directories(config.ep_cache_dir); - std::filesystem::create_directories(config.ep_context_dir); + config.text_encoder.ep_cache_dir = std::filesystem::u8path(args.get("--ep-cache")); + config.text_encoder.ep_context_dir = std::filesystem::u8path(args.get("--ep-context-dir")); + std::filesystem::create_directories(config.text_encoder.ep_cache_dir); + std::filesystem::create_directories(config.text_encoder.ep_context_dir); config.prompt = args.get("--prompt"); - config.precision = args.get("--precision"); + config.transformer.precision = args.get("--precision"); config.num_images = 1; const auto encoder = args.get("--encoder"); const bool translator = encoder == "translator" || - (encoder == "auto" && !config.model_dir.empty() && - std::filesystem::is_regular_file(config.model_dir / "text_encoder_translator/model.onnx")); - config.text_encoder = translator ? Flux2TextEncoder::Qwen3_06BTranslator : Flux2TextEncoder::Qwen3_4B; - config.weight_streaming_budget = "-1"; + (encoder == "auto" && !model_root.empty() && + std::filesystem::is_regular_file(model_root / "text_encoder_translator/model.onnx")); + config.text_encoder.variant = translator ? Flux2TextEncoder::Qwen3_06BTranslator : Flux2TextEncoder::Qwen3_4B; + config.transformer.initial_weight_streaming_budget = "-1"; + ConfigureFlux2Models(config, model_root, config.text_encoder.ep_cache_dir, config.text_encoder.ep_context_dir); if (!glfwInit()) throw std::runtime_error("GLFW initialization failed"); glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3); diff --git a/image_gen/flux2/flux2_vk_backend.cpp b/image_gen/flux2/flux2_vk_backend.cpp index b1b9559..315be16 100644 --- a/image_gen/flux2/flux2_vk_backend.cpp +++ b/image_gen/flux2/flux2_vk_backend.cpp @@ -27,6 +27,7 @@ #include "flux2.h" #include "flux2_runtime_context.h" +#include "flux2_models.h" #include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" @@ -39,14 +40,6 @@ static const std::string POSTPROCESS_SHADER_PATH = std::string(SHADER_DIR) + "po // Upload latent noise for a new seed (CPU → staging → GPU via Vulkan copy) // ============================================================================ -void initialize_latents(unsigned int seed, size_t hidden_count, void* dest) -{ - nvtx3::scoped_range nvtx("initialize_latents"); - std::cout << "\n=== Initializing Latent (seed=" << seed << ") ===" << std::endl; - - initialize_latent(static_cast(dest), hidden_count, seed); -} - class OrtVulkanGraphicsInteropScope { public: @@ -707,6 +700,9 @@ struct VkPipelineState std::unique_ptr transformer_io; std::unique_ptr vae_decoder_io; std::vector time_schedule; + Flux2LoadedModels loaded; + din::common::ModelProfile model_profile; + std::optional applied_budget; }; static void initialize_vk_state(VkPipelineState& state, const Flux2Config& config, Ort::ConstEpDevice trt_device) @@ -715,17 +711,8 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi { throw std::runtime_error("Vulkan processing requires --provider trt-rtx"); } - const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config.model_dir, config.precision, config.text_encoder); - const auto latent_stats = - LoadFlux2LatentStats(model_paths.vae_decoder_model.parent_path() / "latent_stats.json", LATENT_CHANNELS); state.use_cig = config.processing == Flux2ProcessingBackend::VkCig; - const Flux2ModelCachePaths cache_paths = - MakeFlux2ModelCachePaths(config.precision, state.use_cig ? "vk_cig" : "vk", config.text_encoder, - !config.weight_streaming_budget.empty()); - std::cout << "Model dir: " << model_paths.base_dir.string() << "\n" - << "CIG: " << (state.use_cig ? "enabled" : "disabled") << "\n" - << std::endl; // ----------------------------------------------------------------- // Init Vulkan @@ -755,11 +742,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // ----------------------------------------------------------------- // Load ONNX models // ----------------------------------------------------------------- - std::cout << "=== Loading ONNX Models ===" << std::endl; - auto nvtx_scope_load = nvtx3::start_range("load_onnx_models"); - din::common::EpContextOptions ep_context; - ep_context.output_dir = config.ep_context_dir.string(); - const std::string cache_dir = config.ep_cache_dir.string(); std::vector> graphics_ep_options; if (state.use_cig) { @@ -780,17 +762,9 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // graphics_ep_options.emplace_back("nv_use_sync_gpu_allocator", "1"); } - auto make_profile = [&](std::string cache_subpath, bool transformer = false) - { - din::common::ModelProfile profile; - profile.cache_subpath = std::move(cache_subpath); - profile.embed_ep_context = false; - profile.enable_cuda_graph = !state.use_cig; - profile.extra_ep_options = graphics_ep_options; - if (transformer && !config.weight_streaming_budget.empty()) - profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.weight_streaming_budget); - return profile; - }; + state.model_profile.enable_cuda_graph = !state.use_cig; + state.model_profile.embed_ep_context = false; + state.model_profile.extra_ep_options = graphics_ep_options; // ----------------------------------------------------------------- // Allocate tensors from Vulkan buffers imported through ORT's external @@ -883,20 +857,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi state.bn_mean_buf = state.vk->createExternalBuffer(LATENT_CHANNELS * sizeof(float)); state.bn_std_buf = state.vk->createExternalBuffer(LATENT_CHANNELS * sizeof(float)); - { - VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - memcpy(state.staging.mapped, latent_stats.mean.data(), LATENT_CHANNELS * sizeof(float)); - VkBufferCopy region = {0, 0, state.bn_mean_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.bn_mean_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); - } - { - VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - memcpy(state.staging.mapped, latent_stats.std.data(), LATENT_CHANNELS * sizeof(float)); - VkBufferCopy region = {0, 0, state.bn_std_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.bn_std_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); - } state.postprocess_shader = loadComputeShader(POSTPROCESS_SHADER_PATH, *state.vk, sizeof(PostprocessPushConstants), @@ -916,33 +876,21 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // Sessions are intentionally last: their CIG context sees the fully initialized importer, // external queue, shared buffers, and timeline semaphore. - state.text_encoder_runner = std::make_unique( - state.env, model_paths.text_encoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.text_encoder), &*state.sync_stream); - std::cout << " Text encoder loaded" << std::endl; - - state.transformer_runner = std::make_unique( - state.env, model_paths.transformer_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.transformer, true), &*state.sync_stream); - std::cout << " Transformer loaded" << std::endl; - - state.vae_decoder_runner = std::make_unique( - state.env, model_paths.vae_decoder_model.string(), "trt-rtx", cache_dir, ep_context, - make_profile(cache_paths.vae_decoder), &*state.sync_stream); - std::cout << " VAE decoder loaded" << std::endl; - nvtx3::end_range(nvtx_scope_load); // ----------------------------------------------------------------- // Persistent command buffers (pool already created above) // ----------------------------------------------------------------- // Euler sampler - state.cmd_bufs_euler.resize(config.steps); - state.cmd_bufs_time.resize(config.steps); - state.vk->device.allocateCommandBuffers(state.command_pool, config.steps, state.cmd_bufs_euler.data()); + // Reserve all supported steps so changing the schedule needs no model reload + // or command-buffer allocation while generation is in progress. + constexpr unsigned int max_steps = 50; + state.cmd_bufs_euler.resize(max_steps); + state.cmd_bufs_time.resize(max_steps); + state.vk->device.allocateCommandBuffers(state.command_pool, max_steps, state.cmd_bufs_euler.data()); // Time scheduler - state.vk->device.allocateCommandBuffers(state.command_pool, config.steps, state.cmd_bufs_time.data()); + state.vk->device.allocateCommandBuffers(state.command_pool, max_steps, state.cmd_bufs_time.data()); VkCommandBuffer cmd_bufs[3]; state.vk->device.allocateCommandBuffers(state.command_pool, 3, cmd_bufs); @@ -955,22 +903,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi // ----------------------------------------------------------------- // Create persistent IoBindings (all GPU tensors, reused every run) // ----------------------------------------------------------------- - state.text_encoder_io = std::make_unique(state.text_encoder_runner->session); - state.text_encoder_io->BindInput("input_ids", state.token_tensor); - state.text_encoder_io->BindInput("attention_mask", state.attention_mask_tensor); - state.text_encoder_io->BindOutput("prompt_embeds", state.text_encoder_embeds_tensor); - - state.transformer_io = std::make_unique(state.transformer_runner->session); - state.transformer_io->BindInput("hidden_states", state.hidden_states_tensor); - state.transformer_io->BindInput("encoder_hidden_states", state.text_encoder_embeds_tensor); - state.transformer_io->BindInput("timestep", state.timestep_tensor); - state.transformer_io->BindInput("img_ids", state.img_ids_tensor); - state.transformer_io->BindInput("txt_ids", state.txt_ids_tensor); - state.transformer_io->BindOutput("sample", state.transformer_output_tensor); - - state.vae_decoder_io = std::make_unique(state.vae_decoder_runner->session); - state.vae_decoder_io->BindInput("latent_sample", state.decoder_input_tensor); - state.vae_decoder_io->BindOutput("sample", state.decoded_image_tensor); // ----------------------------------------------------------------- // One-time upload: position IDs (via staging buffer) @@ -1018,39 +950,83 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi } // ----------------------------------------------------------------- - // Time schedule (computed once, constant across images) + // Initial time schedule; generation options can replace it per image. // ----------------------------------------------------------------- state.time_schedule = MakeFlux2Schedule(config.steps); // ----------------------------------------------------------------- // Load prompt tokens & upload to GPU (once for all images) // ----------------------------------------------------------------- - std::cout << "\n=== Tokenizing Prompt ===" << std::endl; - { - nvtx3::scoped_range nvtx("tokenize"); - - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); - - std::vector tokens_cpu(shape_numel(token_shape)); - std::vector attn_mask_cpu(shape_numel(attn_mask_shape)); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); +} - memcpy(state.staging.mapped, tokens_cpu.data(), state.token_buf.size); - VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); - VkBufferCopy region = {0, 0, state.token_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.token_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); +static void upload_vk_prompt(VkPipelineState& state, const Flux2Config& config) +{ + const Flux2ModelPaths model_paths = + MakeFlux2ModelPaths(config); + const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); + std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); + FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), + BATCH_SIZE, SEQUENCE_LENGTH); + + memcpy(state.staging.mapped, tokens_cpu.data(), state.token_buf.size); + VkCommandBuffer cmd = state.vk->beginCommandBuffer(state.command_pool); + VkBufferCopy region = {0, 0, state.token_buf.size}; + vkCmdCopyBuffer(cmd, state.staging.buffer, state.token_buf.buffer, 1, ®ion); + state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); + + memcpy(state.staging.mapped, attn_mask_cpu.data(), state.attn_mask_buf.size); + cmd = state.vk->beginCommandBuffer(state.command_pool); + region = {0, 0, state.attn_mask_buf.size}; + vkCmdCopyBuffer(cmd, state.staging.buffer, state.attn_mask_buf.buffer, 1, ®ion); + state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); +} - memcpy(state.staging.mapped, attn_mask_cpu.data(), state.attn_mask_buf.size); - cmd = state.vk->beginCommandBuffer(state.command_pool); - region = {0, 0, state.attn_mask_buf.size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, state.attn_mask_buf.buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); +static void update_vk_models(VkPipelineState& state, const Flux2Config& config) +{ + if (!state.loaded.Matches(config)) + { + signal_and_wait_on_host(*state.sync_stream); + VK_CHECK(vkDeviceWaitIdle(state.vk->device.raw())); } + const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; + Ort::SyncStream* stream = &*state.sync_stream; + if (state.loaded.text_encoder != config.text_encoder) state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, + state.text_encoder_runner, state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); + upload_vk_prompt(state, config); + return binding; + }); + auto transformer_profile = state.model_profile; + if (!config.transformer.initial_weight_streaming_budget.empty()) + transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", + config.transformer.initial_weight_streaming_budget); + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, + state.transformer_runner, state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) { return BindFlux2Model(session, {{"hidden_states", &state.hidden_states_tensor}, {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, {"timestep", &state.timestep_tensor}, {"img_ids", &state.img_ids_tensor}, {"txt_ids", &state.txt_ids_tensor}}, {{"sample", &state.transformer_output_tensor}}); })) + state.applied_budget.reset(); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, + state.env, provider, stream, state.model_profile, [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, {{"sample", &state.decoded_image_tensor}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + for (const auto& upload : {std::pair{&state.bn_mean_buf, stats.mean.data()}, + std::pair{&state.bn_std_buf, stats.std.data()}}) + { + memcpy(state.staging.mapped, upload.second, upload.first->size); + const auto cmd = state.vk->beginCommandBuffer(state.command_pool); + VkBufferCopy region = {0, 0, upload.first->size}; + vkCmdCopyBuffer(cmd, state.staging.buffer, upload.first->buffer, 1, ®ion); + state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); + } + return binding; + }); } -static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, const Flux2Progress& progress) +static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, din::common::Rng rng, const Flux2Progress& progress) { Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, [&] @@ -1058,7 +1034,8 @@ static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, const signal_and_wait_on_host(*state.sync_stream); VK_CHECK(vkQueueWaitIdle(state.vk->queue)); }); - initialize_latent(static_cast(state.staging.mapped), state.hidden_count, seed); + initialize_cpu_latent(static_cast(state.staging.mapped), state.hidden_count, seed, rng, + BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); const uint64_t vae_done = run_pipeline( state.text_encoder_runner->session, state.transformer_runner->session, state.vae_decoder_runner->session, @@ -1103,11 +1080,26 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline { if (state_) { + UpdateModels(config_); return; } std::cout << "Initializing Flux2 Vulkan pipeline" << std::endl; state_ = std::make_unique(runtime_.env); - initialize_vk_state(*state_, config_, runtime_.trt_device); + try { initialize_vk_state(*state_, config_, runtime_.trt_device); } + catch (...) { state_.reset(); throw; } + UpdateModels(config_); + } + + void UpdateModels(const Flux2Config& requested) override + { + if (!state_) throw std::runtime_error("Pipeline resources are not initialized"); + if (requested.processing != config_.processing || requested.provider != config_.provider) + throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); + // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. + const auto previous_prompt = config_.prompt; + config_ = requested; + config_.prompt = previous_prompt; + update_vk_models(*state_, requested); } void SetPrompt(std::string prompt) override @@ -1121,54 +1113,48 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } state_->prompt_embeds_valid = false; - const Flux2ModelPaths model_paths = - MakeFlux2ModelPaths(config_.model_dir, config_.precision, config_.text_encoder); - const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config_.prompt); - std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); - FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, tokens_cpu.data(), attn_mask_cpu.data(), - BATCH_SIZE, SEQUENCE_LENGTH); - - memcpy(state_->staging.mapped, tokens_cpu.data(), state_->token_buf.size); - VkCommandBuffer cmd = state_->vk->beginCommandBuffer(state_->command_pool); - VkBufferCopy region = {0, 0, state_->token_buf.size}; - vkCmdCopyBuffer(cmd, state_->staging.buffer, state_->token_buf.buffer, 1, ®ion); - state_->vk->endAndSubmitCommandBuffer(cmd, state_->command_pool); - - memcpy(state_->staging.mapped, attn_mask_cpu.data(), state_->attn_mask_buf.size); - cmd = state_->vk->beginCommandBuffer(state_->command_pool); - region = {0, 0, state_->attn_mask_buf.size}; - vkCmdCopyBuffer(cmd, state_->staging.buffer, state_->attn_mask_buf.buffer, 1, ®ion); - state_->vk->endAndSubmitCommandBuffer(cmd, state_->command_pool); + upload_vk_prompt(*state_, config_); } const void* TransformerSessionIdentity() const override { - return state_ ? static_cast(state_->transformer_runner->session) : nullptr; + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) : nullptr; + } + + const void* TextEncoderSessionIdentity() const override + { + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) : nullptr; + } + + const void* VaeSessionIdentity() const override + { + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) : nullptr; } Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, const Flux2GenerationOptions& options) override { - if (!state_) + if (!state_ || !state_->loaded.Ready()) { throw std::runtime_error("Flux2 Vulkan pipeline was not initialized"); } - const std::string budget = options.weight_streaming_budget.value_or(config_.weight_streaming_budget); - const std::string& previous = applied_budget_ ? *applied_budget_ : config_.weight_streaming_budget; + const din::common::Rng rng = options.rng.value_or(config_.rng); + ValidateFlux2Rng(rng, config_.processing); + state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); + const std::string budget = options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; if (budget != previous) { const char* keys[] = {"nv_weight_streaming_budget"}; const char* values[] = {budget.c_str()}; state_->transformer_runner->session.SetEpDynamicOptions(keys, values, 1); - applied_budget_ = budget; + state_->applied_budget = budget; } - return run_vk_image(*state_, seed, progress); + return run_vk_image(*state_, seed, rng, progress); } private: Flux2Config config_; - std::optional applied_budget_; Flux2RuntimeContext& runtime_; std::unique_ptr state_; }; diff --git a/image_gen/flux2/helpers/flux2.h b/image_gen/flux2/helpers/flux2.h index 68c5fc0..7e71b5f 100644 --- a/image_gen/flux2/helpers/flux2.h +++ b/image_gen/flux2/helpers/flux2.h @@ -10,12 +10,13 @@ #include #include #include -#include +#include #include #include #include #include "flux2_cli.h" +#include "rng.h" #include "tokenizer.h" // ============================================================================ @@ -31,39 +32,11 @@ struct Flux2ModelPaths std::filesystem::path tokenizer_dir; }; -struct Flux2ModelCachePaths +inline Flux2ModelPaths MakeFlux2ModelPaths(const Flux2Config& config) { - std::string text_encoder; - std::string transformer; - std::string vae_decoder; -}; - -inline Flux2ModelPaths MakeFlux2ModelPaths(std::filesystem::path model_dir, const std::string& precision, - Flux2TextEncoder encoder = Flux2TextEncoder::Qwen3_4B) -{ - if (model_dir.empty()) - { - throw std::invalid_argument("Model directory must be specified"); - } - return { - model_dir, - model_dir / (encoder == Flux2TextEncoder::Qwen3_4B ? "text_encoder" : "text_encoder_translator") / "model.onnx", - model_dir / ("transformer_" + precision) / "model.onnx", - model_dir / "vae_decoder/model.onnx", - model_dir / "tokenizer", - }; -} - -inline Flux2ModelCachePaths MakeFlux2ModelCachePaths(const std::string& precision, const std::string& prefix = {}, - Flux2TextEncoder encoder = Flux2TextEncoder::Qwen3_4B, - bool weight_streaming = false) -{ - const std::string prefix_separator = prefix.empty() ? "" : prefix + "_"; - return { - prefix_separator + (encoder == Flux2TextEncoder::Qwen3_4B ? "text_encoder" : "text_encoder_translator"), - prefix_separator + "transformer_" + precision + (weight_streaming ? "_ws" : ""), - prefix_separator + "vae_decoder", - }; + return {config.text_encoder.model_path.parent_path().parent_path(), + config.text_encoder.model_path, config.transformer.model_path, + config.vae.model_path, config.text_encoder.tokenizer_dir}; } static const std::filesystem::path TEXT_ENCODER_MODEL_CTX = "text_encoder_model_ctx.onnx"; @@ -177,12 +150,33 @@ inline void FillTextEncoderInputs(const std::vector& token_ids, int64_t // Random latent initialisation // ============================================================================ -inline void initialize_latent(float* latent, size_t num_elements, unsigned int seed = 42) +inline void initialize_cpu_latent(float* output, size_t count, uint64_t seed, din::common::Rng rng, + unsigned int batch, unsigned int channels, unsigned int height, unsigned int width) { - std::default_random_engine rng(seed); - std::normal_distribution dist(0.0f, 1.0f); - for (size_t i = 0; i < num_elements; ++i) - latent[i] = dist(rng); + size_t expected = 1; + for (unsigned int dimension : {batch, channels, height, width}) + { + if (dimension == 0 || expected > std::numeric_limits::max() / dimension) + throw std::invalid_argument("Latent dimensions must be nonempty and fit in size_t"); + expected *= dimension; + } + if (output == nullptr || count != expected) + throw std::invalid_argument("Noise buffer must match the latent dimensions"); + if (rng == din::common::Rng::CpuStd) + { + din::common::GenerateNormalCpu({output, count}, seed, rng); + return; + } + std::vector noise(count); + din::common::GenerateNormalCpu(noise, seed, rng, {.round_to_bfloat16 = true}); + const size_t spatial = static_cast(height) * width; + for (size_t index = 0; index < count; ++index) + { + const size_t b = index / (channels * spatial); + const size_t c = (index / spatial) % channels; + const size_t pixel = index % spatial; + output[(b * spatial + pixel) * channels + c] = noise[index]; + } } // ============================================================================ From e301a62e2441d1f18a49a33c18474deb16d6c288 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Tue, 27 Oct 2026 16:26:37 +0100 Subject: [PATCH 12/14] format using linux --- common/ort_session.cpp | 6 +- common/ort_session.h | 4 +- common/rng.cpp | 15 ++- common/rng.h | 12 +- common/rng_cuda.cu | 50 +++---- common/rng_cuda.h | 7 +- image_gen/flux2/cuda_kernels.cu | 1 - image_gen/flux2/flux2_cli.cpp | 7 +- image_gen/flux2/flux2_cli.h | 4 +- image_gen/flux2/flux2_cuda_backend.cpp | 178 +++++++++++++++---------- image_gen/flux2/flux2_dx_backend.cpp | 111 +++++++++------ image_gen/flux2/flux2_models.cpp | 58 +++++--- image_gen/flux2/flux2_models.h | 32 +++-- image_gen/flux2/flux2_pipeline.cpp | 4 +- image_gen/flux2/flux2_ui.cpp | 27 ++-- image_gen/flux2/flux2_vk_backend.cpp | 114 ++++++++++------ image_gen/flux2/helpers/flux2.h | 11 +- 17 files changed, 381 insertions(+), 260 deletions(-) diff --git a/common/ort_session.cpp b/common/ort_session.cpp index c636e32..9514ebd 100644 --- a/common/ort_session.cpp +++ b/common/ort_session.cpp @@ -934,12 +934,14 @@ TensorBuffer::TensorBuffer(Ort::Env& env, Ort::ConstEpDevice device, Ort::Syn if (!device || stream == nullptr) throw std::invalid_argument("Device tensors require a device and a persistent stream"); const auto pinned = env.GetSharedAllocator(device.GetMemoryInfo(OrtDeviceMemoryType_HOST_ACCESSIBLE)); - if (!pinned) throw std::runtime_error("Missing shared pinned allocator"); + if (!pinned) + throw std::runtime_error("Missing shared pinned allocator"); host_tensor_ = Ort::Value::CreateTensor(pinned, shape_.data(), shape_.size()); if (!use_unified_memory_) { const auto allocator = env.GetSharedAllocator(device.GetMemoryInfo(OrtDeviceMemoryType_DEFAULT)); - if (!allocator) throw std::runtime_error("Missing shared device allocator"); + if (!allocator) + throw std::runtime_error("Missing shared device allocator"); device_tensor_ = Ort::Value::CreateTensor(allocator, shape_.data(), shape_.size()); } } diff --git a/common/ort_session.h b/common/ort_session.h index b0bf293..30c8572 100644 --- a/common/ort_session.h +++ b/common/ort_session.h @@ -199,8 +199,8 @@ class TensorBuffer { public: TensorBuffer(OrtRunner& runner, const std::vector& shape, bool use_device_io, bool disable_uma = false); - TensorBuffer(Ort::Env& env, Ort::ConstEpDevice device, Ort::SyncStream* stream, - const std::vector& shape, bool use_device_io, bool disable_uma = false); + TensorBuffer(Ort::Env& env, Ort::ConstEpDevice device, Ort::SyncStream* stream, const std::vector& shape, + bool use_device_io, bool disable_uma = false); Ort::Value& BindingValue(); diff --git a/common/rng.cpp b/common/rng.cpp index 1db078a..7ed5181 100644 --- a/common/rng.cpp +++ b/common/rng.cpp @@ -2,7 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 #include "rng.h" -#include #include #include #include @@ -10,6 +9,8 @@ #include #include +#include + namespace din::common { namespace @@ -18,7 +19,7 @@ float round_bfloat16(float value) { uint32_t bits = std::bit_cast(value); if ((bits & 0x7fffffffU) > 0x7f800000U) - bits |= 0x00400000U; // Preserve NaNs as quiet NaNs. + bits |= 0x00400000U; // Preserve NaNs as quiet NaNs. else bits += 0x7fffU + ((bits >> 16) & 1U); return std::bit_cast(bits & 0xffff0000U); @@ -35,7 +36,7 @@ void box_muller(uint32_t x, uint32_t y, float* result) result[0] = radius * std::sin(v); result[1] = radius * std::cos(v); } -} +} // namespace void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, const NormalOptions& options) { @@ -64,9 +65,9 @@ void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, const No uint64_t iteration = 0; for (size_t base = thread; base < count; ++iteration) { - const Philox::ctr_type counter = {{static_cast(iteration), - static_cast(iteration >> 32), static_cast(thread), - static_cast(static_cast(thread) >> 32)}}; + const Philox::ctr_type counter = {{static_cast(iteration), static_cast(iteration >> 32), + static_cast(thread), + static_cast(static_cast(thread) >> 32)}}; const auto words = Philox{}(counter, key); float samples[4]; box_muller(words[0], words[1], samples); @@ -85,4 +86,4 @@ void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, const No } } } -} // namespace din::common +} // namespace din::common diff --git a/common/rng.h b/common/rng.h index d826ba4..8ee2233 100644 --- a/common/rng.h +++ b/common/rng.h @@ -8,7 +8,12 @@ namespace din::common { -enum class Rng { CpuStd, CudaPhilox, CpuPhilox }; +enum class Rng +{ + CpuStd, + CudaPhilox, + CpuPhilox +}; // Reference launch geometry: H100 SXM5, 132 SMs * 2048 threads/SM. inline constexpr size_t kH100PhiloxStride = 270336; @@ -26,6 +31,5 @@ struct NormalOptions // CpuPhilox uses Philox4x32-10 and cuRAND's Box-Muller convention and ordering. // CPU math can differ from CUDA, including at BF16 rounding boundaries. // CudaPhilox must be dispatched through GenerateNormalCuda in rng_cuda.h. -void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, - const NormalOptions& options = {}); -} +void GenerateNormalCpu(std::span output, uint64_t seed, Rng rng, const NormalOptions& options = {}); +} // namespace din::common diff --git a/common/rng_cuda.cu b/common/rng_cuda.cu index dcee40a..cba5954 100644 --- a/common/rng_cuda.cu +++ b/common/rng_cuda.cu @@ -1,10 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -#include "rng_cuda.h" +#include + #include #include #include -#include + +#include "rng_cuda.h" #include namespace din::common @@ -14,12 +16,13 @@ namespace constexpr unsigned int kThreadsPerBlock = 256; void check_cuda(cudaError_t result) { - if (result != cudaSuccess) throw std::runtime_error(cudaGetErrorString(result)); + if (result != cudaSuccess) + throw std::runtime_error(cudaGetErrorString(result)); } // Match ATen DistributionTemplates.h: one Philox subsequence per thread, // normal4 components spaced by the entire grid, then grid-stride by 4. -__global__ __launch_bounds__(256, 4) void normal_philox_kernel( - uint64_t seed, size_t count, bool round_to_bfloat16, size_t channels, size_t spatial, float* output) +__global__ __launch_bounds__(256, 4) void normal_philox_kernel(uint64_t seed, size_t count, bool round_to_bfloat16, + size_t channels, size_t spatial, float* output) { const size_t thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const size_t stride = static_cast(blockDim.x) * gridDim.x; @@ -36,20 +39,22 @@ __global__ __launch_bounds__(256, 4) void normal_philox_kernel( const size_t index = base + stride * component; if (index < count) { - const size_t destination = channels == 0 ? index - : ((index / (channels * spatial)) * spatial + index % spatial) * channels - + (index / spatial) % channels; - output[destination] = round_to_bfloat16 - ? __bfloat162float(__float2bfloat16_rn(values[component])) : values[component]; + const size_t destination = + channels == 0 ? index + : ((index / (channels * spatial)) * spatial + index % spatial) * channels + + (index / spatial) % channels; + output[destination] = + round_to_bfloat16 ? __bfloat162float(__float2bfloat16_rn(values[component])) : values[component]; } } } } -void generate_normal_cuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, - bool round_to_bfloat16, size_t channels, size_t spatial) +void generate_normal_cuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, bool round_to_bfloat16, + size_t channels, size_t spatial) { - if (count == 0) return; + if (count == 0) + return; if (output == nullptr || count > std::numeric_limits::max() / 2) throw std::invalid_argument("CUDA normal generation requires a valid buffer and supported count"); int device = 0; @@ -59,23 +64,22 @@ void generate_normal_cuda(cudaStream_t stream, float* output, size_t count, uint // ATen caps the grid by the device's thread capacity, not kernel occupancy. const size_t capacity = static_cast(properties.multiProcessorCount) * (properties.maxThreadsPerMultiProcessor / kThreadsPerBlock); - const unsigned int blocks = static_cast( - std::min(capacity, (count + kThreadsPerBlock - 1) / kThreadsPerBlock)); - normal_philox_kernel<<>>( - seed, count, round_to_bfloat16, channels, spatial, output); + const unsigned int blocks = + static_cast(std::min(capacity, (count + kThreadsPerBlock - 1) / kThreadsPerBlock)); + normal_philox_kernel<<>>(seed, count, round_to_bfloat16, channels, spatial, + output); check_cuda(cudaPeekAtLastError()); } -} +} // namespace -void GenerateNormalCuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, - bool round_to_bfloat16) +void GenerateNormalCuda(cudaStream_t stream, float* output, size_t count, uint64_t seed, bool round_to_bfloat16) { generate_normal_cuda(stream, output, count, seed, round_to_bfloat16, 0, 0); } -void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, - size_t channels, size_t spatial, bool round_to_bfloat16) +void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, size_t channels, + size_t spatial, bool round_to_bfloat16) { if (channels == 0 || spatial == 0 || channels > std::numeric_limits::max() / spatial || count % (channels * spatial) != 0) @@ -83,4 +87,4 @@ void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, generate_normal_cuda(stream, output, count, seed, round_to_bfloat16, channels, spatial); } -} // namespace din::common +} // namespace din::common diff --git a/common/rng_cuda.h b/common/rng_cuda.h index 0a1d5e2..fdc8b65 100644 --- a/common/rng_cuda.h +++ b/common/rng_cuda.h @@ -3,6 +3,7 @@ #pragma once #include + #include #include @@ -16,6 +17,6 @@ void GenerateNormalCuda(cudaStream_t stream, float* output, size_t count, uint64 // Generate the same [batch, channels, spatial] samples directly into contiguous // [batch, spatial, channels] storage. count must be a multiple of channels * spatial. -void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, - size_t channels, size_t spatial, bool round_to_bfloat16 = false); -} +void GenerateNormalCudaPacked(cudaStream_t stream, float* output, size_t count, uint64_t seed, size_t channels, + size_t spatial, bool round_to_bfloat16 = false); +} // namespace din::common diff --git a/image_gen/flux2/cuda_kernels.cu b/image_gen/flux2/cuda_kernels.cu index d32bdc5..104f68b 100644 --- a/image_gen/flux2/cuda_kernels.cu +++ b/image_gen/flux2/cuda_kernels.cu @@ -5,7 +5,6 @@ #include #include - #include "cu_helper.h" #include "cuda_kernels.h" diff --git a/image_gen/flux2/flux2_cli.cpp b/image_gen/flux2/flux2_cli.cpp index 49b3c80..4791c82 100644 --- a/image_gen/flux2/flux2_cli.cpp +++ b/image_gen/flux2/flux2_cli.cpp @@ -284,8 +284,11 @@ int main(int argc, char* argv[]) << "Output: " << config.output_path.string() << "\n" << "Prompt: " << config.prompt << "\n" << "Seed: " << config.seed << "\n" - << "RNG: " << (config.rng == din::common::Rng::CpuStd ? "cpu-std" : - config.rng == din::common::Rng::CudaPhilox ? "cuda-philox" : "cpu-philox") << "\n" + << "RNG: " + << (config.rng == din::common::Rng::CpuStd ? "cpu-std" + : config.rng == din::common::Rng::CudaPhilox ? "cuda-philox" + : "cpu-philox") + << "\n" << "Num images: " << config.num_images << "\n" << "Processing: " << to_string(config.processing) << "\n" << "Provider: " << to_string(config.provider) << "\n" diff --git a/image_gen/flux2/flux2_cli.h b/image_gen/flux2/flux2_cli.h index 996dc24..7e8f775 100644 --- a/image_gen/flux2/flux2_cli.h +++ b/image_gen/flux2/flux2_cli.h @@ -3,8 +3,6 @@ #pragma once -#include "rng.h" - #include #include #include @@ -13,6 +11,8 @@ #include #include +#include "rng.h" + enum class Flux2ProcessingBackend { Cpu, diff --git a/image_gen/flux2/flux2_cuda_backend.cpp b/image_gen/flux2/flux2_cuda_backend.cpp index eed4d05..1ac7d6e 100644 --- a/image_gen/flux2/flux2_cuda_backend.cpp +++ b/image_gen/flux2/flux2_cuda_backend.cpp @@ -25,12 +25,12 @@ #include "cu_helper.h" #include "cuda_kernels.h" -#include "rng_cuda.h" #include "flux2.h" -#include "flux2_runtime_context.h" #include "flux2_models.h" +#include "flux2_runtime_context.h" #include "flux2_telemetry.h" #include "ort_session.h" +#include "rng_cuda.h" #include "utils.h" #include #include @@ -518,51 +518,50 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const state.model_profile.extra_ep_options.emplace_back("nv_length_aux_stream_array", "0"); std::vector token_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; - state.token = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, token_shape, - has_separate_binding); + state.token = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + token_shape, has_separate_binding); std::vector attn_mask_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; - state.attention_mask = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, - attn_mask_shape, has_separate_binding); + state.attention_mask = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, attn_mask_shape, has_separate_binding); std::vector te_out_shape = {BATCH_SIZE, SEQUENCE_LENGTH, TEXT_ENCODER_EMBED_DIM}; - state.text_encoder_embeds = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, - te_out_shape, has_separate_binding); + state.text_encoder_embeds = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, te_out_shape, has_separate_binding); std::vector timestep_shape = {BATCH_SIZE}; - state.timestep = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, timestep_shape, - has_separate_binding); + state.timestep = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + timestep_shape, has_separate_binding); std::vector hidden_shape = {BATCH_SIZE, TRANSFORMER_HIDDEN_DIM, LATENT_CHANNELS}; - state.hidden_states = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, hidden_shape, - has_separate_binding); + state.hidden_states = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, hidden_shape, has_separate_binding); std::vector img_ids_shape = {BATCH_SIZE, IMAGE_SEQUENCE, 4}; - state.img_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, img_ids_shape, - has_separate_binding); + state.img_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + img_ids_shape, has_separate_binding); std::vector txt_ids_shape = {BATCH_SIZE, SEQUENCE_LENGTH, 4}; - state.txt_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, txt_ids_shape, - has_separate_binding); + state.txt_ids = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + txt_ids_shape, has_separate_binding); - state.transformer_output = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, - hidden_shape, has_separate_binding); + state.transformer_output = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, hidden_shape, has_separate_binding); std::vector dec_latent_shape = {BATCH_SIZE, LATENT_CHANNELS / PATCH_SIZE / PATCH_SIZE, LATENT_HEIGHT * PATCH_SIZE, LATENT_WIDTH * PATCH_SIZE}; - state.decoder_latent = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, - dec_latent_shape, has_separate_binding); + state.decoder_latent = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, dec_latent_shape, has_separate_binding); std::vector image_shape = {BATCH_SIZE, IMAGE_CHANNELS, IMAGE_HEIGHT, IMAGE_WIDTH}; - state.decoded_image = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, image_shape, - has_separate_binding); + state.decoded_image = std::make_unique>( + state.env, state.ep_device, compute_stream_ptr, image_shape, has_separate_binding); std::vector bn_shape = {LATENT_CHANNELS}; - state.bn_mean = - std::make_unique>(state.env, state.ep_device, compute_stream_ptr, bn_shape, has_separate_binding); - state.bn_std = - std::make_unique>(state.env, state.ep_device, compute_stream_ptr, bn_shape, has_separate_binding); - + state.bn_mean = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + bn_shape, has_separate_binding); + state.bn_std = std::make_unique>(state.env, state.ep_device, compute_stream_ptr, + bn_shape, has_separate_binding); std::vector img_ids_cpu(state.img_ids->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount()); std::vector txt_ids_cpu(state.txt_ids->BindingValue().GetTensorTypeAndShapeInfo().GetElementCount()); @@ -570,7 +569,6 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const std::copy(img_ids_cpu.begin(), img_ids_cpu.end(), state.img_ids->HostData()); std::copy(txt_ids_cpu.begin(), txt_ids_cpu.end(), state.txt_ids->HostData()); - if (has_separate_binding) { auto img_ids_ready = state.img_ids->CopyAsyncToDeviceWithNotification(); @@ -594,8 +592,7 @@ void initialize_ep(CudaPipelineState& state, Ort::ConstEpDevice ep_device, const static void upload_cuda_prompt(CudaPipelineState& state, const Flux2Config& config) { - const Flux2ModelPaths model_paths = - MakeFlux2ModelPaths(config); + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config); const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); FillTextEncoderInputs(text_inputs.token_ids, text_inputs.pad_token_id, state.token->HostData(), state.attention_mask->HostData(), BATCH_SIZE, SEQUENCE_LENGTH); @@ -613,43 +610,61 @@ static void update_cuda_models(CudaPipelineState& state, const Flux2Config& conf { if (!state.loaded.Matches(config)) { - if (state.compute_stream) signal_and_wait_on_host(*state.compute_stream); + if (state.compute_stream) + signal_and_wait_on_host(*state.compute_stream); } const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; Ort::SyncStream* stream = state.compute_stream ? &*state.compute_stream : nullptr; - if (state.loaded.text_encoder != config.text_encoder) state.prompt_embeds_valid = false; - ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, - state.text_encoder_runner, state.text_encoder_io, state.env, provider, stream, state.model_profile, - [&](Ort::Session& session) - { - auto binding = BindFlux2Model(session, {{"input_ids", &state.token->BindingValue()}, {"attention_mask", &state.attention_mask->BindingValue()}}, {{"prompt_embeds", &state.text_encoder_embeds->BindingValue()}}); - upload_cuda_prompt(state, config); - return binding; - }); + if (state.loaded.text_encoder != config.text_encoder) + state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, state.text_encoder_runner, + state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, + {{"input_ids", &state.token->BindingValue()}, + {"attention_mask", &state.attention_mask->BindingValue()}}, + {{"prompt_embeds", &state.text_encoder_embeds->BindingValue()}}); + upload_cuda_prompt(state, config); + return binding; + }); auto transformer_profile = state.model_profile; if (!config.transformer.initial_weight_streaming_budget.empty()) transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.transformer.initial_weight_streaming_budget); - if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, - state.transformer_runner, state.transformer_io, state.env, provider, stream, transformer_profile, - [&](Ort::Session& session) { return BindFlux2Model(session, {{"hidden_states", &state.hidden_states->BindingValue()}, {"encoder_hidden_states", &state.text_encoder_embeds->BindingValue()}, {"timestep", &state.timestep->BindingValue()}, {"img_ids", &state.img_ids->BindingValue()}, {"txt_ids", &state.txt_ids->BindingValue()}}, {{"sample", &state.transformer_output->BindingValue()}}); })) + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, state.transformer_runner, + state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) + { + return BindFlux2Model( + session, + {{"hidden_states", &state.hidden_states->BindingValue()}, + {"encoder_hidden_states", &state.text_encoder_embeds->BindingValue()}, + {"timestep", &state.timestep->BindingValue()}, + {"img_ids", &state.img_ids->BindingValue()}, + {"txt_ids", &state.txt_ids->BindingValue()}}, + {{"sample", &state.transformer_output->BindingValue()}}); + })) state.applied_budget.reset(); - ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, - state.env, provider, stream, state.model_profile, [&](Ort::Session& session) - { - auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_latent->BindingValue()}}, {{"sample", &state.decoded_image->BindingValue()}}); - const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); - std::copy_n(stats.mean.data(), LATENT_CHANNELS, state.bn_mean->HostData()); - std::copy_n(stats.std.data(), LATENT_CHANNELS, state.bn_std->HostData()); - if (state.has_separate_binding) - { - auto mean_ready = state.bn_mean->CopyAsyncToDeviceWithNotification(); - auto std_ready = state.bn_std->CopyAsyncToDeviceWithNotification(); - mean_ready.Sync(); - std_ready.Sync(); - } - return binding; - }); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, state.env, + provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = + BindFlux2Model(session, {{"latent_sample", &state.decoder_latent->BindingValue()}}, + {{"sample", &state.decoded_image->BindingValue()}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + std::copy_n(stats.mean.data(), LATENT_CHANNELS, state.bn_mean->HostData()); + std::copy_n(stats.std.data(), LATENT_CHANNELS, state.bn_std->HostData()); + if (state.has_separate_binding) + { + auto mean_ready = state.bn_mean->CopyAsyncToDeviceWithNotification(); + auto std_ready = state.bn_std->CopyAsyncToDeviceWithNotification(); + mean_ready.Sync(); + std_ready.Sync(); + } + return binding; + }); } Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed, din::common::Rng rng, @@ -687,8 +702,8 @@ Flux2Image run_initialized_pipeline(CudaPipelineState& state, unsigned int seed, LATENT_CHANNELS, LATENT_HEIGHT * LATENT_WIDTH, true); } else - initialize_cpu_latent(state.hidden_states->HostData(), hidden_count, seed, rng, - BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); + initialize_cpu_latent(state.hidden_states->HostData(), hidden_count, seed, rng, BATCH_SIZE, LATENT_CHANNELS, + LATENT_HEIGHT, LATENT_WIDTH); if (state.binding_is_cuda && rng != din::common::Rng::CudaPhilox) { auto hidden_uploaded = state.hidden_states->CopyAsyncToDeviceWithNotification(); @@ -742,13 +757,19 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline throw std::runtime_error("CUDA processing requires --provider trt-rtx"); } - state_ = std::make_unique(runtime_.env); if (provider == ExecutionProviderMode::Cpu) { std::cout << "Using CPU execution provider" << std::endl; - try { initialize_ep(*state_, find_cpu_device(state_->env), config_, provider, sampling_backend); } - catch (...) { state_.reset(); throw; } + try + { + initialize_ep(*state_, find_cpu_device(state_->env), config_, provider, sampling_backend); + } + catch (...) + { + state_.reset(); + throw; + } UpdateModels(config_); return; } @@ -760,14 +781,22 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } std::cout << "Using TensorRT RTX execution provider" << std::endl; - try { initialize_ep(*state_, trt_device, config_, provider, sampling_backend); } - catch (...) { state_.reset(); throw; } + try + { + initialize_ep(*state_, trt_device, config_, provider, sampling_backend); + } + catch (...) + { + state_.reset(); + throw; + } UpdateModels(config_); } void UpdateModels(const Flux2Config& requested) override { - if (!state_) throw std::runtime_error("Pipeline resources are not initialized"); + if (!state_) + throw std::runtime_error("Pipeline resources are not initialized"); if (requested.processing != config_.processing || requested.provider != config_.provider) throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. @@ -793,17 +822,20 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline const void* TransformerSessionIdentity() const override { - return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) : nullptr; + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) + : nullptr; } const void* TextEncoderSessionIdentity() const override { - return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) : nullptr; + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) + : nullptr; } const void* VaeSessionIdentity() const override { - return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) : nullptr; + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) + : nullptr; } Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, @@ -816,8 +848,10 @@ class CudaFlux2ProcessingPipeline final : public Flux2ProcessingPipeline const din::common::Rng rng = options.rng.value_or(config_.rng); ValidateFlux2Rng(rng, config_.processing); state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); - const std::string budget = options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); - const std::string& previous = state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; + const std::string budget = + options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = + state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; if (budget != previous) { const char* keys[] = {"nv_weight_streaming_budget"}; diff --git a/image_gen/flux2/flux2_dx_backend.cpp b/image_gen/flux2/flux2_dx_backend.cpp index b413021..b2670eb 100644 --- a/image_gen/flux2/flux2_dx_backend.cpp +++ b/image_gen/flux2/flux2_dx_backend.cpp @@ -18,8 +18,8 @@ #include "dx_helper.h" #include "flux2.h" -#include "flux2_runtime_context.h" #include "flux2_models.h" +#include "flux2_runtime_context.h" #include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" @@ -463,7 +463,6 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi throw std::runtime_error("DirectX processing requires --provider trt-rtx"); } - std::cout << "\n=== Initializing D3D12 ===" << std::endl; state.trt_device = trt_device; state.interop_api = &Ort::GetInteropApi(); @@ -499,7 +498,6 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.model_profile.embed_ep_context = false; state.model_profile.extra_ep_options = graphics_ep_options; - std::vector token_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; std::vector attn_mask_shape = {BATCH_SIZE, SEQUENCE_LENGTH}; std::vector te_out_shape = {BATCH_SIZE, SEQUENCE_LENGTH, TEXT_ENCODER_EMBED_DIM}; @@ -557,7 +555,6 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi state.decoded_image_tensor = state.tensor_importer->create_tensor(state.decoded_image_buf, image_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); - std::cout << "\n=== Uploading Static Inputs (D3D12) ===" << std::endl; state.dx.begin(); @@ -606,8 +603,7 @@ static void initialize_dx_state(DxPipelineState& state, const Flux2Config& confi static void upload_dx_prompt(DxPipelineState& state, const Flux2Config& config) { - const Flux2ModelPaths model_paths = - MakeFlux2ModelPaths(config); + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config); const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); @@ -632,40 +628,56 @@ static void update_dx_models(DxPipelineState& state, const Flux2Config& config) } const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; Ort::SyncStream* stream = &*state.sync_stream; - if (state.loaded.text_encoder != config.text_encoder) state.prompt_embeds_valid = false; - ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, - state.text_encoder_runner, state.text_encoder_io, state.env, provider, stream, state.model_profile, - [&](Ort::Session& session) - { - auto binding = BindFlux2Model(session, {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); - upload_dx_prompt(state, config); - return binding; - }); + if (state.loaded.text_encoder != config.text_encoder) + state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, state.text_encoder_runner, + state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model( + session, + {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, + {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); + upload_dx_prompt(state, config); + return binding; + }); auto transformer_profile = state.model_profile; if (!config.transformer.initial_weight_streaming_budget.empty()) transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.transformer.initial_weight_streaming_budget); - if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, - state.transformer_runner, state.transformer_io, state.env, provider, stream, transformer_profile, - [&](Ort::Session& session) { return BindFlux2Model(session, {{"hidden_states", &state.hidden_states_tensor}, {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, {"timestep", &state.timestep_tensor}, {"img_ids", &state.img_ids_tensor}, {"txt_ids", &state.txt_ids_tensor}}, {{"sample", &state.transformer_output_tensor}}); })) + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, state.transformer_runner, + state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) + { + return BindFlux2Model(session, + {{"hidden_states", &state.hidden_states_tensor}, + {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, + {"timestep", &state.timestep_tensor}, + {"img_ids", &state.img_ids_tensor}, + {"txt_ids", &state.txt_ids_tensor}}, + {{"sample", &state.transformer_output_tensor}}); + })) state.applied_budget.reset(); - ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, - state.env, provider, stream, state.model_profile, [&](Ort::Session& session) - { - auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, {{"sample", &state.decoded_image_tensor}}); - const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); - state.dx.begin(); - state.dx.record_upload(state.bn_mean_buf, stats.mean.data(), state.bn_mean_buf.size); - state.dx.record_upload(state.bn_std_buf, stats.std.data(), state.bn_std_buf.size); - state.dx.execute(); - const auto uploaded = state.sync_fence->signal_d3d12(state.dx.queue.Get()); - state.sync_fence->wait_ort(*state.sync_stream, uploaded); - signal_and_wait_on_host(*state.sync_stream); - return binding; - }); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, state.env, + provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, + {{"sample", &state.decoded_image_tensor}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + state.dx.begin(); + state.dx.record_upload(state.bn_mean_buf, stats.mean.data(), state.bn_mean_buf.size); + state.dx.record_upload(state.bn_std_buf, stats.std.data(), state.bn_std_buf.size); + state.dx.execute(); + const auto uploaded = state.sync_fence->signal_d3d12(state.dx.queue.Get()); + state.sync_fence->wait_ort(*state.sync_stream, uploaded); + signal_and_wait_on_host(*state.sync_stream); + return binding; + }); } -static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, din::common::Rng rng, const Flux2Progress& progress) +static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, din::common::Rng rng, + const Flux2Progress& progress) { Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, [&] @@ -675,8 +687,8 @@ static Flux2Image run_dx_image(DxPipelineState& state, unsigned int seed, din::c signal_and_wait_on_host(*state.sync_stream); }); std::vector latent_cpu(state.hidden_count); - initialize_cpu_latent(latent_cpu.data(), latent_cpu.size(), seed, rng, - BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); + initialize_cpu_latent(latent_cpu.data(), latent_cpu.size(), seed, rng, BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, + LATENT_WIDTH); state.dx.begin(); state.dx.record_upload(state.hidden_states, latent_cpu.data(), state.hidden_states.size); state.dx.execute(); @@ -724,14 +736,22 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline std::cout << "Initializing Flux2 DirectX pipeline" << std::endl; const auto identity = din::common::ResolveOrtGraphicsDeviceIdentity(runtime_.trt_device); state_ = std::make_unique(runtime_.env, identity.luid); - try { initialize_dx_state(*state_, config_, runtime_.trt_device); } - catch (...) { state_.reset(); throw; } + try + { + initialize_dx_state(*state_, config_, runtime_.trt_device); + } + catch (...) + { + state_.reset(); + throw; + } UpdateModels(config_); } void UpdateModels(const Flux2Config& requested) override { - if (!state_) throw std::runtime_error("Pipeline resources are not initialized"); + if (!state_) + throw std::runtime_error("Pipeline resources are not initialized"); if (requested.processing != config_.processing || requested.provider != config_.provider) throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. @@ -757,17 +777,20 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline const void* TransformerSessionIdentity() const override { - return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) : nullptr; + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) + : nullptr; } const void* TextEncoderSessionIdentity() const override { - return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) : nullptr; + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) + : nullptr; } const void* VaeSessionIdentity() const override { - return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) : nullptr; + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) + : nullptr; } Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, @@ -780,8 +803,10 @@ class DxFlux2ProcessingPipeline final : public Flux2ProcessingPipeline const din::common::Rng rng = options.rng.value_or(config_.rng); ValidateFlux2Rng(rng, config_.processing); state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); - const std::string budget = options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); - const std::string& previous = state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; + const std::string budget = + options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = + state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; if (budget != previous) { const char* keys[] = {"nv_weight_streaming_budget"}; diff --git a/image_gen/flux2/flux2_models.cpp b/image_gen/flux2/flux2_models.cpp index f4998ee..749ed0f 100644 --- a/image_gen/flux2/flux2_models.cpp +++ b/image_gen/flux2/flux2_models.cpp @@ -9,28 +9,39 @@ void ConfigureFlux2Models(Flux2Config& config, const std::filesystem::path& mode std::string prefix; switch (config.processing) { - case Flux2ProcessingBackend::Dx: prefix = "dx"; break; - case Flux2ProcessingBackend::DxCig: prefix = "dx_cig"; break; - case Flux2ProcessingBackend::Vk: prefix = "vk"; break; - case Flux2ProcessingBackend::VkCig: prefix = "vk_cig"; break; - default: break; + case Flux2ProcessingBackend::Dx: + prefix = "dx"; + break; + case Flux2ProcessingBackend::DxCig: + prefix = "dx_cig"; + break; + case Flux2ProcessingBackend::Vk: + prefix = "vk"; + break; + case Flux2ProcessingBackend::VkCig: + prefix = "vk_cig"; + break; + default: + break; } - if (!prefix.empty()) prefix += '_'; - const std::string encoder = config.text_encoder.variant == Flux2TextEncoder::Qwen3_4B - ? "text_encoder" : "text_encoder_translator"; + if (!prefix.empty()) + prefix += '_'; + const std::string encoder = + config.text_encoder.variant == Flux2TextEncoder::Qwen3_4B ? "text_encoder" : "text_encoder_translator"; const std::string transformer = "transformer_" + config.transformer.precision; - for (Flux2ModelConfig* model : {static_cast(&config.text_encoder), - static_cast(&config.transformer), - static_cast(&config.vae)}) + for (Flux2ModelConfig* model : + {static_cast(&config.text_encoder), static_cast(&config.transformer), + static_cast(&config.vae)}) { model->ep_cache_dir = std::filesystem::weakly_canonical(cache_root); model->ep_context_dir = std::filesystem::weakly_canonical(context_root); } config.text_encoder.cache_subpath = prefix + encoder; - config.transformer.cache_subpath = prefix + transformer + - (config.transformer.initial_weight_streaming_budget.empty() ? "" : "_ws"); + config.transformer.cache_subpath = + prefix + transformer + (config.transformer.initial_weight_streaming_budget.empty() ? "" : "_ws"); config.vae.cache_subpath = prefix + "vae_decoder"; - if (root.empty()) return; // The studio can start before a model root is selected. + if (root.empty()) + return; // The studio can start before a model root is selected. config.text_encoder.model_path = root / encoder / "model.onnx"; config.text_encoder.tokenizer_dir = root / "tokenizer"; config.transformer.model_path = root / transformer / "model.onnx"; @@ -38,15 +49,16 @@ void ConfigureFlux2Models(Flux2Config& config, const std::filesystem::path& mode config.vae.latent_stats_path = root / "vae_decoder/latent_stats.json"; } -std::unique_ptr BindFlux2Model(Ort::Session& session, - std::initializer_list inputs, std::initializer_list outputs) +std::unique_ptr BindFlux2Model(Ort::Session& session, std::initializer_list inputs, + std::initializer_list outputs) { auto binding = std::make_unique(session); Ort::AllocatorWithDefaultOptions allocator; auto bind = [&](std::initializer_list tensors, bool input) { const size_t count = input ? session.GetInputCount() : session.GetOutputCount(); - if (count != tensors.size()) throw std::runtime_error("Unexpected model input/output count"); + if (count != tensors.size()) + throw std::runtime_error("Unexpected model input/output count"); for (const auto& tensor : tensors) { size_t index = 0; @@ -54,9 +66,11 @@ std::unique_ptr BindFlux2Model(Ort::Session& session, { const auto name = input ? session.GetInputNameAllocated(index, allocator) : session.GetOutputNameAllocated(index, allocator); - if (std::string(name.get()) == tensor.name) break; + if (std::string(name.get()) == tensor.name) + break; } - if (index == count) throw std::runtime_error(std::string("Missing tensor: ") + tensor.name); + if (index == count) + throw std::runtime_error(std::string("Missing tensor: ") + tensor.name); const auto type = input ? session.GetInputTypeInfo(index) : session.GetOutputTypeInfo(index); const auto info = type.GetTensorTypeAndShapeInfo(); const auto actual = tensor.tensor->GetTensorTypeAndShapeInfo(); @@ -67,8 +81,10 @@ std::unique_ptr BindFlux2Model(Ort::Session& session, for (size_t axis = 0; axis < shape.size(); ++axis) if (shape[axis] >= 0 && shape[axis] != expected[axis]) throw std::runtime_error(std::string("Incompatible tensor shape: ") + tensor.name); - if (input) binding->BindInput(tensor.name, *tensor.tensor); - else binding->BindOutput(tensor.name, *tensor.tensor); + if (input) + binding->BindInput(tensor.name, *tensor.tensor); + else + binding->BindOutput(tensor.name, *tensor.tensor); } }; bind(inputs, true); diff --git a/image_gen/flux2/flux2_models.h b/image_gen/flux2/flux2_models.h index d44c765..489f7b8 100644 --- a/image_gen/flux2/flux2_models.h +++ b/image_gen/flux2/flux2_models.h @@ -2,19 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#include "flux2_cli.h" -#include "ort_session.h" #include #include +#include "flux2_cli.h" +#include "ort_session.h" + struct Flux2TensorBinding { const char* name; Ort::Value* tensor; }; -std::unique_ptr BindFlux2Model(Ort::Session& session, - std::initializer_list inputs, std::initializer_list outputs); +std::unique_ptr BindFlux2Model(Ort::Session& session, std::initializer_list inputs, + std::initializer_list outputs); struct Flux2LoadedModels { @@ -25,16 +26,19 @@ struct Flux2LoadedModels { return text_encoder == config.text_encoder && transformer == config.transformer && vae == config.vae; } - bool Ready() const { return text_encoder && transformer && vae; } + bool Ready() const + { + return text_encoder && transformer && vae; + } }; // The caller drains device work before entering this function. Tensors and the // externally owned stream remain alive across session replacement. -template +template bool ReloadFlux2Stage(const char* name, const Config& requested, std::optional& loaded, - std::unique_ptr& runner, std::unique_ptr& binding, - Ort::Env& env, const char* provider, Ort::SyncStream* stream, - din::common::ModelProfile profile, Bind&& bind) + std::unique_ptr& runner, std::unique_ptr& binding, + Ort::Env& env, const char* provider, Ort::SyncStream* stream, din::common::ModelProfile profile, + Bind&& bind) { if (runner && binding && loaded == requested) { @@ -48,9 +52,9 @@ bool ReloadFlux2Stage(const char* name, const Config& requested, std::optional(env, requested.model_path.string(), provider, - requested.ep_cache_dir.string(), din::common::EpContextOptions{requested.ep_context_dir.string()}, - profile, stream); + runner = std::make_unique( + env, requested.model_path.string(), provider, requested.ep_cache_dir.string(), + din::common::EpContextOptions{requested.ep_context_dir.string()}, profile, stream); binding = bind(runner->session); loaded = requested; } @@ -59,8 +63,8 @@ bool ReloadFlux2Stage(const char* name, const Config& requested, std::optional(config.num_images); int encoder = config.text_encoder.variant == Flux2TextEncoder::Qwen3_4B ? 0 : 1; - int precision = config.transformer.precision == "bf16" ? 0 : config.transformer.precision == "fp16" ? 1 : config.transformer.precision == "fp8" ? 2 : 3; + int precision = config.transformer.precision == "bf16" ? 0 + : config.transformer.precision == "fp16" ? 1 + : config.transformer.precision == "fp8" ? 2 + : 3; const char* precisions[] = {"bf16", "fp16", "fp8", "nvfp4"}; const char* backends[] = {"CPU", "CUDA", "DirectX", "DirectX CIG", "Vulkan", "Vulkan CIG"}; int backend = static_cast(config.processing); @@ -377,13 +380,14 @@ int Studio(const Flux2Config& initial, bool smoke, bool auto_generate, bool exit config.output_path = output.data(); config.prompt = prompt.data(); const auto cache_namespace = CacheNamespace(model_root); - const auto cache_root = - (cache_changed ? std::filesystem::u8path(cache.data()) / "runtime" : initial.text_encoder.ep_cache_dir) / - cache_namespace; - const auto context_root = - (cache_changed ? std::filesystem::u8path(cache.data()) / "ep_context" : initial.text_encoder.ep_context_dir) / - cache_namespace; - config.text_encoder.variant = encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; + const auto cache_root = (cache_changed ? std::filesystem::u8path(cache.data()) / "runtime" + : initial.text_encoder.ep_cache_dir) / + cache_namespace; + const auto context_root = (cache_changed ? std::filesystem::u8path(cache.data()) / "ep_context" + : initial.text_encoder.ep_context_dir) / + cache_namespace; + config.text_encoder.variant = + encoder == 0 ? Flux2TextEncoder::Qwen3_4B : Flux2TextEncoder::Qwen3_06BTranslator; config.transformer.precision = precisions[precision]; config.num_images = static_cast(count); config.processing = static_cast(backend); @@ -504,10 +508,9 @@ int main(int argc, char** argv) config.transformer.precision = args.get("--precision"); config.num_images = 1; const auto encoder = args.get("--encoder"); - const bool translator = - encoder == "translator" || - (encoder == "auto" && !model_root.empty() && - std::filesystem::is_regular_file(model_root / "text_encoder_translator/model.onnx")); + const bool translator = encoder == "translator" || + (encoder == "auto" && !model_root.empty() && + std::filesystem::is_regular_file(model_root / "text_encoder_translator/model.onnx")); config.text_encoder.variant = translator ? Flux2TextEncoder::Qwen3_06BTranslator : Flux2TextEncoder::Qwen3_4B; config.transformer.initial_weight_streaming_budget = "-1"; ConfigureFlux2Models(config, model_root, config.text_encoder.ep_cache_dir, config.text_encoder.ep_context_dir); diff --git a/image_gen/flux2/flux2_vk_backend.cpp b/image_gen/flux2/flux2_vk_backend.cpp index 315be16..7eb3b97 100644 --- a/image_gen/flux2/flux2_vk_backend.cpp +++ b/image_gen/flux2/flux2_vk_backend.cpp @@ -26,8 +26,8 @@ #endif #include "flux2.h" -#include "flux2_runtime_context.h" #include "flux2_models.h" +#include "flux2_runtime_context.h" #include "flux2_telemetry.h" #include "ort_session.h" #include "utils.h" @@ -713,7 +713,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi } state.use_cig = config.processing == Flux2ProcessingBackend::VkCig; - // ----------------------------------------------------------------- // Init Vulkan // ----------------------------------------------------------------- @@ -857,7 +856,6 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi state.bn_mean_buf = state.vk->createExternalBuffer(LATENT_CHANNELS * sizeof(float)); state.bn_std_buf = state.vk->createExternalBuffer(LATENT_CHANNELS * sizeof(float)); - state.postprocess_shader = loadComputeShader(POSTPROCESS_SHADER_PATH, *state.vk, sizeof(PostprocessPushConstants), {&state.hidden_states, &state.decoder_latent, &state.bn_mean_buf, &state.bn_std_buf}); @@ -961,8 +959,7 @@ static void initialize_vk_state(VkPipelineState& state, const Flux2Config& confi static void upload_vk_prompt(VkPipelineState& state, const Flux2Config& config) { - const Flux2ModelPaths model_paths = - MakeFlux2ModelPaths(config); + const Flux2ModelPaths model_paths = MakeFlux2ModelPaths(config); const Flux2TextEncoderInputs text_inputs = TokenizeFlux2Prompt(model_paths, config.prompt); std::vector tokens_cpu(BATCH_SIZE * SEQUENCE_LENGTH); std::vector attn_mask_cpu(BATCH_SIZE * SEQUENCE_LENGTH); @@ -991,42 +988,58 @@ static void update_vk_models(VkPipelineState& state, const Flux2Config& config) } const char* provider = config.provider == Flux2ExecutionProvider::Cpu ? "cpu" : "trt-rtx"; Ort::SyncStream* stream = &*state.sync_stream; - if (state.loaded.text_encoder != config.text_encoder) state.prompt_embeds_valid = false; - ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, - state.text_encoder_runner, state.text_encoder_io, state.env, provider, stream, state.model_profile, - [&](Ort::Session& session) - { - auto binding = BindFlux2Model(session, {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); - upload_vk_prompt(state, config); - return binding; - }); + if (state.loaded.text_encoder != config.text_encoder) + state.prompt_embeds_valid = false; + ReloadFlux2Stage("text encoder", config.text_encoder, state.loaded.text_encoder, state.text_encoder_runner, + state.text_encoder_io, state.env, provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model( + session, + {{"input_ids", &state.token_tensor}, {"attention_mask", &state.attention_mask_tensor}}, + {{"prompt_embeds", &state.text_encoder_embeds_tensor}}); + upload_vk_prompt(state, config); + return binding; + }); auto transformer_profile = state.model_profile; if (!config.transformer.initial_weight_streaming_budget.empty()) transformer_profile.extra_ep_options.emplace_back("nv_weight_streaming_budget", config.transformer.initial_weight_streaming_budget); - if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, - state.transformer_runner, state.transformer_io, state.env, provider, stream, transformer_profile, - [&](Ort::Session& session) { return BindFlux2Model(session, {{"hidden_states", &state.hidden_states_tensor}, {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, {"timestep", &state.timestep_tensor}, {"img_ids", &state.img_ids_tensor}, {"txt_ids", &state.txt_ids_tensor}}, {{"sample", &state.transformer_output_tensor}}); })) + if (ReloadFlux2Stage("transformer", config.transformer, state.loaded.transformer, state.transformer_runner, + state.transformer_io, state.env, provider, stream, transformer_profile, + [&](Ort::Session& session) + { + return BindFlux2Model(session, + {{"hidden_states", &state.hidden_states_tensor}, + {"encoder_hidden_states", &state.text_encoder_embeds_tensor}, + {"timestep", &state.timestep_tensor}, + {"img_ids", &state.img_ids_tensor}, + {"txt_ids", &state.txt_ids_tensor}}, + {{"sample", &state.transformer_output_tensor}}); + })) state.applied_budget.reset(); - ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, - state.env, provider, stream, state.model_profile, [&](Ort::Session& session) - { - auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, {{"sample", &state.decoded_image_tensor}}); - const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); - for (const auto& upload : {std::pair{&state.bn_mean_buf, stats.mean.data()}, - std::pair{&state.bn_std_buf, stats.std.data()}}) - { - memcpy(state.staging.mapped, upload.second, upload.first->size); - const auto cmd = state.vk->beginCommandBuffer(state.command_pool); - VkBufferCopy region = {0, 0, upload.first->size}; - vkCmdCopyBuffer(cmd, state.staging.buffer, upload.first->buffer, 1, ®ion); - state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); - } - return binding; - }); + ReloadFlux2Stage("VAE", config.vae, state.loaded.vae, state.vae_decoder_runner, state.vae_decoder_io, state.env, + provider, stream, state.model_profile, + [&](Ort::Session& session) + { + auto binding = BindFlux2Model(session, {{"latent_sample", &state.decoder_input_tensor}}, + {{"sample", &state.decoded_image_tensor}}); + const auto stats = LoadFlux2LatentStats(config.vae.latent_stats_path, LATENT_CHANNELS); + for (const auto& upload : {std::pair{&state.bn_mean_buf, stats.mean.data()}, + std::pair{&state.bn_std_buf, stats.std.data()}}) + { + memcpy(state.staging.mapped, upload.second, upload.first->size); + const auto cmd = state.vk->beginCommandBuffer(state.command_pool); + VkBufferCopy region = {0, 0, upload.first->size}; + vkCmdCopyBuffer(cmd, state.staging.buffer, upload.first->buffer, 1, ®ion); + state.vk->endAndSubmitCommandBuffer(cmd, state.command_pool); + } + return binding; + }); } -static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, din::common::Rng rng, const Flux2Progress& progress) +static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, din::common::Rng rng, + const Flux2Progress& progress) { Flux2Telemetry telemetry(static_cast(state.time_schedule.size()) - 1, progress, [&] @@ -1034,8 +1047,8 @@ static Flux2Image run_vk_image(VkPipelineState& state, unsigned int seed, din::c signal_and_wait_on_host(*state.sync_stream); VK_CHECK(vkQueueWaitIdle(state.vk->queue)); }); - initialize_cpu_latent(static_cast(state.staging.mapped), state.hidden_count, seed, rng, - BATCH_SIZE, LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); + initialize_cpu_latent(static_cast(state.staging.mapped), state.hidden_count, seed, rng, BATCH_SIZE, + LATENT_CHANNELS, LATENT_HEIGHT, LATENT_WIDTH); const uint64_t vae_done = run_pipeline( state.text_encoder_runner->session, state.transformer_runner->session, state.vae_decoder_runner->session, @@ -1085,14 +1098,22 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline } std::cout << "Initializing Flux2 Vulkan pipeline" << std::endl; state_ = std::make_unique(runtime_.env); - try { initialize_vk_state(*state_, config_, runtime_.trt_device); } - catch (...) { state_.reset(); throw; } + try + { + initialize_vk_state(*state_, config_, runtime_.trt_device); + } + catch (...) + { + state_.reset(); + throw; + } UpdateModels(config_); } void UpdateModels(const Flux2Config& requested) override { - if (!state_) throw std::runtime_error("Pipeline resources are not initialized"); + if (!state_) + throw std::runtime_error("Pipeline resources are not initialized"); if (requested.processing != config_.processing || requested.provider != config_.provider) throw std::invalid_argument("Changing processing backend or provider requires a new pipeline"); // Keep the previous prompt until SetPrompt decides whether embeddings can be reused. @@ -1118,17 +1139,20 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline const void* TransformerSessionIdentity() const override { - return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) : nullptr; + return state_ && state_->transformer_runner ? static_cast(state_->transformer_runner->session) + : nullptr; } const void* TextEncoderSessionIdentity() const override { - return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) : nullptr; + return state_ && state_->text_encoder_runner ? static_cast(state_->text_encoder_runner->session) + : nullptr; } const void* VaeSessionIdentity() const override { - return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) : nullptr; + return state_ && state_->vae_decoder_runner ? static_cast(state_->vae_decoder_runner->session) + : nullptr; } Flux2Image GenerateImage(unsigned int seed, const Flux2Progress& progress, @@ -1141,8 +1165,10 @@ class VkFlux2ProcessingPipeline final : public Flux2ProcessingPipeline const din::common::Rng rng = options.rng.value_or(config_.rng); ValidateFlux2Rng(rng, config_.processing); state_->time_schedule = MakeFlux2Schedule(options.steps.value_or(config_.steps)); - const std::string budget = options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); - const std::string& previous = state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; + const std::string budget = + options.weight_streaming_budget.value_or(config_.transformer.initial_weight_streaming_budget); + const std::string& previous = + state_->applied_budget ? *state_->applied_budget : config_.transformer.initial_weight_streaming_budget; if (budget != previous) { const char* keys[] = {"nv_weight_streaming_budget"}; diff --git a/image_gen/flux2/helpers/flux2.h b/image_gen/flux2/helpers/flux2.h index 7e71b5f..d5f08bc 100644 --- a/image_gen/flux2/helpers/flux2.h +++ b/image_gen/flux2/helpers/flux2.h @@ -9,8 +9,8 @@ #include #include #include -#include #include +#include #include #include #include @@ -34,9 +34,8 @@ struct Flux2ModelPaths inline Flux2ModelPaths MakeFlux2ModelPaths(const Flux2Config& config) { - return {config.text_encoder.model_path.parent_path().parent_path(), - config.text_encoder.model_path, config.transformer.model_path, - config.vae.model_path, config.text_encoder.tokenizer_dir}; + return {config.text_encoder.model_path.parent_path().parent_path(), config.text_encoder.model_path, + config.transformer.model_path, config.vae.model_path, config.text_encoder.tokenizer_dir}; } static const std::filesystem::path TEXT_ENCODER_MODEL_CTX = "text_encoder_model_ctx.onnx"; @@ -150,8 +149,8 @@ inline void FillTextEncoderInputs(const std::vector& token_ids, int64_t // Random latent initialisation // ============================================================================ -inline void initialize_cpu_latent(float* output, size_t count, uint64_t seed, din::common::Rng rng, - unsigned int batch, unsigned int channels, unsigned int height, unsigned int width) +inline void initialize_cpu_latent(float* output, size_t count, uint64_t seed, din::common::Rng rng, unsigned int batch, + unsigned int channels, unsigned int height, unsigned int width) { size_t expected = 1; for (unsigned int dimension : {batch, channels, height, width}) From 517455607b538755110edc4f432bebfeba8628c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Tue, 27 Oct 2026 16:42:50 +0100 Subject: [PATCH 13/14] clang tools --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 12222c1..96e71cd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,7 +54,7 @@ jobs: run: | set -euxo pipefail apt-get update - apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev + apt-get install -y --no-install-recommends ca-certificates clang clang-tools cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev - name: Check out source uses: actions/checkout@v7 @@ -97,7 +97,7 @@ jobs: run: | set -euxo pipefail apt-get update - apt-get install -y --no-install-recommends ca-certificates clang cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev + apt-get install -y --no-install-recommends ca-certificates clang clang-tools cmake git ninja-build python3 python3-dev libgl1-mesa-dev libx11-dev libxrandr-dev libxinerama-dev libxcursor-dev libxi-dev - name: Check out source uses: actions/checkout@v7 From 900e4c98408cdf3e8fa6b19f8a8506e45f550606 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20M=C3=BCller?= Date: Tue, 27 Oct 2026 16:45:04 +0100 Subject: [PATCH 14/14] trt rtx 1.7 --- CMakePresets.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/CMakePresets.json b/CMakePresets.json index 584a048..eef10b4 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -46,7 +46,7 @@ "cacheVariables": { "CMAKE_SYSTEM_PROCESSOR": "AMD64", "ONNXRUNTIME_ARCH": "x64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-amd64-cuda-13.4-Release-external.zip" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-windows-amd64-cuda-13.4-release-external.zip" }, "vendor": { "jetbrains.com/clion": { @@ -64,7 +64,7 @@ "CMAKE_SYSTEM_NAME": "Windows", "CMAKE_SYSTEM_PROCESSOR": "ARM64", "ONNXRUNTIME_ARCH": "arm64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-arm64-cuda-13.4-Release-external.zip" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-windows-arm64-cuda-13.4-release-external.zip" }, "vendor": { "jetbrains.com/clion": { @@ -81,7 +81,7 @@ "cacheVariables": { "CMAKE_SYSTEM_PROCESSOR": "x86_64", "ONNXRUNTIME_ARCH": "x64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Linux-x86_64-cuda-13.4-Release-external.tar.zst" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-linux-x86_64-cuda-13.4-release-external.tar.zst" }, "vendor": { "jetbrains.com/clion": { @@ -98,7 +98,7 @@ "cacheVariables": { "CMAKE_SYSTEM_PROCESSOR": "aarch64", "ONNXRUNTIME_ARCH": "arm64", - "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Linux-aarch64-cuda-13.4-Release-external.tar.zst" + "TRT_RTX_DOWNLOAD_URL": "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.7/tensorrt-rtx-1.7.1.116-linux-aarch64-cuda-13.4-release-external.tar.zst" }, "vendor": { "jetbrains.com/clion": {