Skip to main content

Crate rlx_models

Crate rlx_models 

Source
Expand description

RLX model loading — parse configs, load weights, build IR graphs.

This crate is a thin facade over per-model workspace members (rlx-qwen3, rlx-sam, …). Depend on a specific model crate directly when you only need one family.

Modules§

arch_registry
Architecture registry (plan #82).
config
Model configuration structs — parsed from HuggingFace config.json.
dataprocessing
Reusable batch-prep utilities (plan #83).
embed
flow_bridge
Bridge between rlx-models loaders/runtime and rlx-flow.
flow_util
Shared helpers for tier-0 model flow migration.
flux2
gguf_resolve
Pluggable GGUF tensor-name resolution per general.architecture.
gguf_support
Shared GGUF helpers for LM runners (architecture checks, path resolution).
lm
Shared causal-LM flow helpers — re-export tier-0 surface for model authors.
qwen35
qwen35_synth
Synthetic configs/weights for tests and qwen35_inference bench. Synthetic Qwen3.5 weights for integration tests and criterion benches.
vision_ops_ir
Shared HIR builders for NCHW vision ops (Conv, ConvTranspose2d, LayerNorm2d, bias broadcast). Used by SAM / SAM2 / SAM3.
weight_loader
Pluggable weight loader trait (plan #56).
weight_map
Safetensors weight loading — standalone, no framework dependency.
weight_registry
Extensible weight-format registry — register custom loaders for new extensions.
weights
Model-agnostic weight I/O — paths, formats, drain policy only.

Structs§

BertConfig
BERT model configuration.
BertTokenizer
Wrapper around HuggingFace tokenizer configured for BERT-style encoding.
BuiltModel
Result of assembling a model flow.
ChatMessage
One turn in a ChatML conversation.
CompileProfile
Tier-1 compile configuration. Load from *.rlx.toml or use Rust presets.
Flux2CfgCombineFlow
Tier-0 CFG combine: neg + scale * (pos - neg).
Flux2CfgCombineGraph
Flux2Checkpoint
Resolved on-disk layout for a FLUX.2 HF repo (diffusers-style tree).
Flux2Config
FLUX.2 rectified-flow transformer (denoiser) configuration.
Flux2Flow
Tier-0 FLUX.2 dual-stream flow builder.
Flux2ForwardBuilt
Full forward build product (includes non-f32 typed param blobs).
Flux2ForwardGraph
Flux2ForwardInput
Inputs for one transformer forward (noise prediction).
Flux2Output
Noise prediction from Flux2Runner::forward.
Flux2PromptOutput
Flux2Runner
FLUX.2 denoiser runner — native CPU or compiled HIR on any Device.
Flux2RunnerBuilder
Builder for Flux2Runner.
Flux2Session
One loaded FLUX.2 pipeline — cheap to clone via Arc.
Flux2SessionCache
Process-wide cache of Flux2Runner instances (CLI --reuse-session / serve mode).
Flux2SessionKey
Cache key for deduplicating loaded runners.
Flux2TextEncoderBuilt
Flux2TextEncoderFlow
Tier-0 FLUX.2 text encoder flow (Qwen3-shaped causal LM trunk).
Flux2VaeConfig
Flux2VaeDecoderFlow
Tier-0 FLUX.2 VAE decoder flow.
Flux2VaeEncoderFlow
Tier-0 FLUX.2 VAE encoder flow.
Flux2VaeGraph
Flux2VaeWeights
Flux2Weights
GgufDirGuide
LlamaFamilyGgufResolver
HF model.layers.N.* ↔ GGUF blk.N.* (Llama, Qwen3, Qwen35, …).
LoadWeightsOptions
Options for load_weights_resolved — prefer crate::weights::LoadOpts presets at call sites.
ModelInfo
Metadata for an embedding model.
NomicBertConfig
NomicBERT model configuration.
NomicVisionConfig
NomicVision model configuration.
Qwen35Config
Qwen3.5 model config — fields covering both the per-layer Mamba+ Attention block and the MTP head.
Qwen35FullAttnLayer
Standard full-attention trunk layer (interspersed every full_attention_interval blocks). Per qwen35.cpp::load_block_trunk non-recurrent branch.
Qwen35LinearLayer
Gated DeltaNet (“linear attention”) trunk layer. Mirrors qwen35.cpp::load_block_trunk for the is_recurrent(il) branch.
Qwen35MoeFfn
MoE FFN tensors for one decoder layer (trunk or MTP).
Qwen35MtpLayer
One MTP (NextN) layer. Per qwen35.cpp::load_block_mtp.
Qwen35NativeGgufResolver
Qwen3.5 native blk.N.* names; also accept HF aliases via the Llama mapper.
Qwen35PrefillOutput
Qwen35Runner
Qwen35RunnerBuilder
Qwen35Weights
Top-level Qwen3.5 / Qwen3.6 weight bundle.
RegisteredFormat
One registered on-disk format (built-in or custom).
ResolveWeightsOptions
Options for resolve_weights_file_with_options.
RlxBertModel
RLX-compiled BERT model ready for inference.
RlxEmbed
High-level embedding model — auto-detects BERT / NomicBERT / NomicVision.
RlxNomicModel
RLX-compiled NomicBERT with shape-bucketed compile cache.
RlxVisionModel
RLX-compiled NomicVision encoder (patch preprocess host-side, trunk on RLX).
TokenizedBatch
Output of batch tokenization: token IDs, attention masks, and token type IDs.
WeightFormatRegistration
Describes one on-disk weight format.
WeightMap
Map of tensor name → (f32 data, shape).
WeightMapSource
Adapt in-memory WeightMap to WeightSource.

Enums§

Arch
Detected embedding architecture from config.json.
ChatRole
Conversation role for ChatMessage.
EmbedGgufKind
BERT vs NomicBERT discriminator from GGUF metadata.
EmbeddingModel
Supported text embedding models.
GgufModelFamily
LM families in this workspace that load .gguf weights.
ImageEmbeddingModel
Supported image embedding models.
LoadedWeights
Result of resolving and opening weights.
MatWeight
Storage variant for matmul weight tensors. The big projections (qkv / gate / ffn / lm_head) dominate the load footprint; the Packed variant keeps GGUF K-quant bytes in-place so the graph can emit Op::DequantMatMul instead of a full F32 dequant.
ModelArch
Model architecture type.
Pooling
Pooling strategy for reducing token hidden states to one vector per sequence.
Qwen35LayerFfn
Per-layer feed-forward: dense SwiGLU or MoE (routed + gated shared expert).
Qwen35TrunkLayer
One trunk-layer tensor bundle. Either a gated-DeltaNet “linear attention” block or a standard full-attention block.
WeightDrainPolicy
How WeightMap::drain_loader / WeightMap::from_weight_loader handle leftovers.

Constants§

DEFAULT_TEXT_ENCODER_LAYERS
Default hidden-state indices for FLUX.2 Klein (matches mflux).
STANDARD_DEVICE_NAMES
CLI / help string for --device.

Traits§

FlowBuildExt
GgufTensorNameResolver
Resolve a builder-requested tensor name to the name stored in a GGUF file.
WeightLoader
Common interface every weight format must satisfy. Mirrors the existing WeightMap API so the safetensors impl is a one-line adapter.

Functions§

apply_compile_profile
Apply tier-1 profile options to runtime compile options.
assemble_vision_hidden
Assemble encoder input [batch, seq, hidden] from NCHW pixels + preprocess weights.
assert_gguf_family
Open the file and ensure general.architecture matches expected.
build_flux2_cfg_combine_hir
HIR graph: neg + guidance_scale * (pos - neg) in f32.
build_flux2_forward_graph
build_flux2_forward_hir
Build the full denoiser forward graph in HIR.
build_flux2_minimal_graph
Lower minimal HIR to legacy Graph (MIR inner) for Session::compile.
build_flux2_minimal_hir
Build a compile-minimal HIR module: x_embedder(hidden)proj_out.
build_flux2_text_encoder_hir
build_graph
Build via flow and lower to MIR graph + params.
build_qwen35_decode_graph
Single-token decode graph at prefix length past_seq.
build_qwen35_decode_hir_dynamic_ext
Decode HIR with symbolic past length (sym::PAST_SEQ) for dynamic compile cache.
build_qwen35_graph_sized
Build the Qwen3.5 forward IR.
build_qwen35_graph_sized_ext
Forward graph with optional runtime MRoPE inputs (rope_cos/rope_sin).
build_qwen35_graph_sized_stub
Legacy redirect — qwen35 forward is implemented via build_qwen35_graph_sized. Kept so older call sites get a clear message instead of a missing-symbol error.
build_qwen35_prefill_cache_graph
Prefill graph that seeds super::cache::Qwen35DecodeCache.
build_qwen35_prefill_cache_graph_ext
Prefill-cache graph with optional runtime MRoPE inputs (multimodal).
build_qwen35_prefill_cache_hir_dynamic_ext
Prefill-cache HIR with symbolic seq dim (sym::SEQ) for dynamic compile cache.
built_from_graph
built_from_hir
built_from_hir_with_profile
cfg_combine
Native CFG blend in float32.
compile_built
Compile a BuiltModel on the given device using its embedded profile.
compile_built_cpu
Compile a BuiltModel on CPU with default options (embedding quick-check tests).
compile_flux2_cfg_combine
compile_flux2_forward
compile_flux2_forward_via_flow
Compile denoiser via tier-0 Flux2Flow wrapper (same numerics as super::hir_builder::compile_flux2_forward).
compile_flux2_minimal
Compile minimal HIR on CPU (HIR → MIR → LIR).
compile_flux2_text_encoder_hir
compile_graph_encoder
Bidirectional encoder defaults (BERT, DINOv2, Wav2Vec2, vision towers).
compile_graph_encoder_with_params
CompileProfile::encoder + params.
compile_graph_legacy
Unprofiled compile (parity probes / bisect tests).
compile_graph_llama32_decode
Llama 3.2 decode graphs.
compile_graph_llama32_prefill
Llama 3.2 prefill graphs.
compile_graph_profile
Lower a graph with a tier-1 profile and attach params (tests / examples).
compile_graph_qwen3_decode
Qwen3 single-token decode graphs.
compile_graph_qwen3_prefill
Qwen3 prefill / full-sequence graphs.
compile_graph_qwen3_prefill_with_params
CompileProfile::qwen3_prefill + params.
compile_graph_qwen35_decode
Qwen3.5 decode-step graphs.
compile_graph_qwen35_decode_with_params
CompileProfile::qwen35_decode + params.
compile_graph_qwen35_prefill
Qwen3.5 prefill-cache / predict graphs.
compile_graph_qwen35_prefill_with_params
CompileProfile::qwen35_prefill + params.
compile_graph_sam
Compile a SAM/SAM2/SAM3 vision subgraph with tier-1 encoder profile options.
compile_graph_sam_with_params
CompileProfile::sam_encoder + params.
compile_graph_with_profile
Compile a vision subgraph with explicit tier-1 profile options.
compile_model
Compile an embedding graph for the given batch/seq on device.
decode_step_feeds
Build host feeds for a single decode step from cache.
detect_arch
Detect architecture from config.json fields.
download_flux2_repo
embed_with_rlx
Embed texts with a compiled BERT model: tokenize, forward, pool, L2-normalize.
encode_chat_auto
Resolve tokenizer next to weights and encode a chat conversation.
encode_flux2_prompt
End-to-end: tokenize (optional) + text encoder → embeddings + text ids.
encode_prompt_embeds_default_layers
Encode with default Klein layer indices (9, 18, 27).
encode_prompt_padded
Encode and pad/truncate to fixed seq_len (pad token id 0).
extract_flux2_vae_weights
extract_flux2_weights
extract_text_encoder_weights
flux2_decode_packed_latents
Full post-denoise decode: packed transformer latents → 8-bit RGB planar [batch, 3, H, W].
flux2_prefers_compiled_hir
True when the denoiser / VAE should use compiled HIR (non-CPU backends).
flux2_prefers_compiled_te
Text encoder HIR on CUDA compiles a full Qwen3 trunk and can take hours + fill VRAM while the denoiser is still resident. Native CPU encode once, then drop TE.
flux2_rgb_to_u8
Planar RGB [-1,1] → interleaved u8 HWC for PNG.
flux2_transformer_forward
Run the FLUX.2 transformer and return noise prediction [batch, img_seq, patch_size² * out_channels].
format_chatml
Format messages as a ChatML prompt ending with an open assistant turn.
format_for_extension
Extension → format id (last registration wins).
gguf_architecture_str
general.architecture string from GGUF metadata, if present.
gguf_dir_guide
Numbered .gguf listing + resolve hints for a directory (CLI / errors).
gguf_f32_bytes_estimate
Rough F32 dequant footprint (every tensor × 4 bytes).
gguf_family_for_arch
Map a GGUF architecture tag to the runner family that should load it.
gguf_runner_hint
Suggested runner / crate for a GGUF architecture tag (for CLI and errors).
graph_from_built
Build a flow and return (Graph, params) — preferred compile entry point.
graph_from_hir
Lower an existing HIR module through BuiltModel (utility for HIR-first builders).
host_temb
Host-side temb for compiled forward (timestep × 1000, optional guidance × 1000).
into_compile_parts
Split built flow for compile — no Graph/HIR imports needed at call site.
is_standard_device
True when device is in STANDARD_DEVICES.
list_registered_formats
All registered formats (built-ins first, then custom registrations).
load_and_apply_flux2_lora
Load LoRA from safetensors and merge into base.
load_compile_profile
Load a tier-1 profile from disk; fall back to default when missing or invalid.
load_flux2_vae_weights
load_flux2_weights
load_from_path
Dispatch on the file extension via crate::weight_registry.
load_rgb_planar
Load an image, resize to (width, height), return planar NCHW f32 in [-1, 1].
load_text_encoder_weights
load_weight_map_resolved
Convenience: resolve + drain to F32 WeightMap.
load_weights_resolved
Resolve a file or directory, enforce GGUF arch policy, open via registry, optionally drain.
messages_from_prompt
Convenience: system (optional) + user prompt → ChatML messages.
models_map
Get the global model registry.
mrope_prefill_feeds
Flattened [seq * head_half] cos/sin for runtime MRoPE graph inputs.
mrope_row_for_sections
Build one MRoPE cos/sin row from explicit per-section positions.
mrope_slice_at_pos
Slice MRoPE cos/sin at absolute text position pos.
mtp_draft_vocab_size
MTP LM head output width: full vocab, or trimmed for FastMTP draft speed.
open_map
Resolve + drain to F32 WeightMap.
open_map_with
Resolve + drain with options.
open_weights
Resolve + open (live WeightLoader).
open_with
Resolve + open with options.
pack_input_ids
Pack per-row prompts into [batch, max_seq] row-major F32 ids (zero-pad).
parse_lora_scale
Parse --lora-scale style input; rejects NaN/inf.
parse_messages_json
Parse a JSON array of { "role": "...", "content": "..." }.
prepare_latent_ids
FLUX.2 latent position ids [batch, h*w, 4] with (t=0, h, w, l=0).
prepare_text_ids
Build FLUX.2-style text position ids [batch, seq, 4] flattened as [seq*4].
prepare_weight_map
Full load-time adaptation pipeline.
profile_near_weights
Load profile_file next to weights (parent directory); fall back to default.
recurrent_output_count
Number of extra graph outputs after logits (and optional MTP).
register_gguf_tensor_resolver
Register a custom resolver (call before first GGUF load). Later registrations win among resolvers that match the same architecture.
register_weight_format
Register a custom weight format (call before the first load). Later entries override built-ins when the same extension is registered twice.
resolve_text_encoder_dir
Resolve text_encoder/ next to a transformer weights file or model root.
resolve_tokenizer_path
Resolve tokenizer path: explicit, tokenizer/tokenizer.json, or sibling tokenizer.json.
resolve_transformer_config
Resolve transformer config.json from explicit override or sibling search.
resolve_vae_dir
Resolve vae/ next to a transformer weights file or model root.
resolve_weights_file
Resolve --weights to a single file: pass-through for files, or pick one .gguf / model.safetensors inside a directory.
resolve_weights_file_with_options
Resolve with optional GGUF file selection inside a directory.
seed_cache_from_outputs
Parse prefill-cache graph outputs into logits/hidden + Qwen35DecodeCache. When trunk_is_hidden, the first output is [batch × hidden_size] not logits.
supports_multimodal_mrope
True when the checkpoint declares a non-zero 4th MRoPE section (vision).
text_section_pos
Text-modality default: [p, p, p, 0] per llama.cpp token batches.
tiny_text_encoder_config
Qwen3Config sized for super::weights::synthetic_text_encoder_weights tests.
validate_device
Validate that device is in the workspace standard backend set (CPU, Metal, MLX, CUDA, ROCm, WGPU, Vulkan). Build with all-backends on rlx-qwen35 to link every native runtime backend into the rlx-qwen35 binary.
validate_sam_device
SAM v1 also documents tpu on [rlx_sam::Sam::from_safetensors_on].
validate_standard_device
Fail fast on exotic runtime devices (TPU, ANE, OpenGL, …).
zero_recurrent_inputs
Zero-initialized recurrent inputs for a prefill-cache seed graph.

Type Aliases§

Flux2GraphParams
Param tensors keyed by name for rlx_runtime::CompiledGraph::set_param.
LoadOpts
Alias for LoadWeightsOptions.
PassThroughGgufResolver
Alias for PrefixStripGgufResolver (older name).
ResolveOpts
Alias for ResolveWeightsOptions.