'
+ '
'
+)
+
+_INTRO_HTML = (
+ 'Modality Forcing for Scalable Spatial Generation
' + '' + ''
+ '1Carnegie Mellon University'
+ '2World Labs'
+ '
'
+ ''
+ '
'
+ '
'
+ '
'
+ ''
+ ''
+ 'Preprint, 2026
'
+ ''
+ 'Modality '
+ 'Forcing turns a pretrained text-to-image diffusion transformer into a '
+ 'joint image-depth generator with a simple post-training recipe.
'
+ 'Type a scene and press Generate, or upload an '
+ 'image to run image→depth mode instead.
'
+)
+
+_EXAMPLE_IMAGES = [
+ ["assets/examples/alley.png"],
+ ["assets/examples/yosemite.png"],
+ ["assets/examples/castle.png"],
+]
+
+_EXAMPLE_PROMPTS = [
+ [DEFAULT_PROMPT],
+ ["A sunlit Scandinavian living room with a linen sofa, a low oak coffee "
+ "table, and tall windows opening onto a snowy courtyard."],
+ ["A misty pine forest at dawn, shafts of golden light cutting between the "
+ "trunks and a narrow dirt trail winding into the distance."],
+ ["A cozy bookshop interior with floor-to-ceiling wooden shelves, a rolling "
+ "ladder, warm pendant lighting, and a worn leather reading chair."],
+ ["A still life on a marble countertop: a bowl of ripe lemons, a ceramic "
+ "pitcher, and a sprig of rosemary lit by soft side light."],
+ ["A coastal cliffside at golden hour overlooking a turquoise bay, with wild "
+ "grass in the foreground and distant sailboats on the water."],
+]
+
+with gr.Blocks(title="Modality Forcing — World Labs") as demo:
+ gr.HTML(_HEADER_HTML, elem_classes="mf-bare")
+ gr.HTML(_INTRO_HTML, elem_classes="mf-bare")
+ gr.HTML('', elem_classes="mf-bare") + + with gr.Row(equal_height=False): + with gr.Column(scale=2, min_width=320) as left_col: + gr.HTML('
Input
', elem_classes="mf-bare")
+ prompt = gr.Textbox(
+ value=DEFAULT_PROMPT,
+ lines=4,
+ label="Scene prompt",
+ placeholder="Describe a scene to generate…",
+ )
+ input_image = gr.Image(
+ label="Optional input image (switches to image→depth mode)",
+ type="numpy", height=200, sources=("upload", "clipboard"),
+ )
+ btn = gr.Button("Generate", variant="primary", size="lg")
+
+ with gr.Accordion("Advanced settings", open=False):
+ with gr.Row():
+ num_steps = gr.Slider(
+ 1, 80, value=50, step=1, label="Sampling steps")
+ cfg_scale = gr.Slider(
+ 1.0, 8.0, value=4.0, step=0.1,
+ label="Guidance (CFG) — text mode only")
+ log2_alpha = gr.Slider(
+ -5.0, 5.0, value=5.0, step=0.5,
+ label="log2(alpha) — depth trajectory (>0 rgb-first, "
+ "cleaner depth; 0 diagonal; <0 depth-first) — text mode")
+ edge_rtol = gr.Slider(
+ 0.0, 0.25, value=0.04, step=0.005,
+ label="Point-cloud depth-edge mask (rtol) — drop pixels at "
+ "depth jumps > this; lower = more aggressive, 0 = off")
+ sor_toggle = gr.Checkbox(
+ value=False,
+ label="Statistical outlier removal (point cloud) — drops "
+ "isolated floaters; can over-trim fine structures")
+ seed = gr.Number(value=0, precision=0, label="Seed")
+ refine_depth = gr.Checkbox(
+ value=True,
+ label="Refine depth: joint (CFG) → image→depth (CFG 1) "
+ "— text mode only",
+ )
+
+ status = gr.Textbox(label="Status", interactive=False)
+
+ with gr.Column(scale=3, min_width=380):
+ gr.HTML('Output
', elem_classes="mf-bare")
+ with gr.Row():
+ rgb_out = gr.Image(
+ label="RGB image", type="numpy", height=320, format="png")
+ depth_out = gr.Image(
+ label="Depth (disparity, magma)",
+ type="numpy", height=320, format="png")
+ with gr.Group():
+ cloud_out = gr.Model3D(
+ label="Interactive 3D point cloud",
+ clear_color=(0.12, 0.12, 0.14, 1.0),
+ zoom_speed=0.5, pan_speed=0.5,
+ height=520,
+ )
+
+ def _run_text_example(prompt_text):
+ # Cached example runs pin the UI defaults so the cache stays valid.
+ return generate(prompt_text, None, 50, 4.0, 0, True, 5.0, 0.04, False)
+
+ def _run_image_example(image):
+ return generate("", image, 50, 4.0, 0, True, 5.0, 0.04, False)
+
+ # Defined after the output components exist, rendered back into the left
+ # column. Cached: clicking an example serves precomputed results instead
+ # of spending GPU time ("lazy" = no cache rebuild at startup; the
+ # committed cache ships with the Space repo).
+ _CACHE_MODE = os.environ.get("EXAMPLES_CACHE_MODE", "lazy")
+ with left_col:
+ gr.Examples(
+ examples=_EXAMPLE_PROMPTS,
+ inputs=[prompt],
+ outputs=[rgb_out, depth_out, cloud_out, status],
+ fn=_run_text_example,
+ cache_examples=True,
+ cache_mode=_CACHE_MODE,
+ label="Example prompts",
+ )
+ gr.Examples(
+ examples=_EXAMPLE_IMAGES,
+ inputs=[input_image],
+ outputs=[rgb_out, depth_out, cloud_out, status],
+ fn=_run_image_example,
+ cache_examples=True,
+ cache_mode=_CACHE_MODE,
+ label="Example images (image → depth)",
+ )
+
+ gr.HTML(
+ '',
+ elem_classes="mf-bare",
+ )
+
+ btn.click(
+ generate,
+ inputs=[prompt, input_image, num_steps, cfg_scale, seed, refine_depth,
+ log2_alpha, edge_rtol, sor_toggle],
+ outputs=[rgb_out, depth_out, cloud_out, status],
+ )
+
+ # Adding an image switches the demo into image→depth mode, where the prompt
+ # is optional — clear it so it defaults to empty. The user can retype one.
+ input_image.change(
+ lambda img: "" if img is not None else gr.update(),
+ inputs=[input_image],
+ outputs=[prompt],
+ )
+
+# The UI is designed light-only (paper-white cards, explicit light CSS). A
+# visitor whose OS is in dark mode otherwise gets gradio's dark text colors on
+# our light backgrounds — an unreadable mix. Redirect to ?__theme=light in
+# , before gradio boots, so every visitor gets the light theme.
+_FORCE_LIGHT_HEAD = """
+
+"""
+
+if __name__ == "__main__":
+ # Gradio 6 takes theme + css on launch() (not the Blocks constructor),
+ # so they must be passed here to actually apply on the Space.
+ demo.queue(max_size=4).launch(theme=_THEME, css=_CSS, head=_FORCE_LIGHT_HEAD)
diff --git a/assets/examples/alley.png b/assets/examples/alley.png
new file mode 100644
index 0000000000000000000000000000000000000000..bafd83e56c8ae253273e665b6a7da93b6f827467
--- /dev/null
+++ b/assets/examples/alley.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12d47280f10a0f7e07077ddb595d18d3636e6ef3cbe63ccf281f193aca7f55e3
+size 1030212
diff --git a/assets/examples/castle.png b/assets/examples/castle.png
new file mode 100644
index 0000000000000000000000000000000000000000..3e594400b074b96a99b85ec56b1944102b1b05b0
--- /dev/null
+++ b/assets/examples/castle.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d94240486e851197441a2b63b02ee70eeab7d2e607d3568e34741c54ca55b94b
+size 1900625
diff --git a/assets/examples/yosemite.png b/assets/examples/yosemite.png
new file mode 100644
index 0000000000000000000000000000000000000000..6d12640ae253c83c3994b156433066dcaa584071
--- /dev/null
+++ b/assets/examples/yosemite.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:edeb3066dd0ac94309cb8ab2372c4b2c3474847cacf2055821dcde036faeca45
+size 1657947
diff --git a/flux_rgbd/__init__.py b/flux_rgbd/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a8c9d585192a6562cce6e514255121592fb85df9
--- /dev/null
+++ b/flux_rgbd/__init__.py
@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""flux_rgbd — joint text → RGB + depth diffusion on top of FLUX.2."""
+
+__version__ = "0.1.0"
+__all__ = ["FluxRGBDRunner"]
+
+
+def __getattr__(name):
+ if name == "FluxRGBDRunner":
+ from flux_rgbd.runner import FluxRGBDRunner
+ return FluxRGBDRunner
+ raise AttributeError(name)
diff --git a/flux_rgbd/_flux2/__init__.py b/flux_rgbd/_flux2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab58edae3e390e9e4d544a88b238b6c6a6bf1cb9
--- /dev/null
+++ b/flux_rgbd/_flux2/__init__.py
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+"""Vendored FLUX.2 reference implementation.
+
+This subpackage holds the unmodified FLUX.2 architecture (Apache-2.0,
+Black Forest Labs upstream). Do not edit these files; they are
+upstream-tracking. All modifications and extensions for RGB+D inference
+live in the parent ``flux_rgbd`` package.
+
+Upstream: https://github.com/black-forest-labs/flux2
+"""
diff --git a/flux_rgbd/_flux2/autoencoder.py b/flux_rgbd/_flux2/autoencoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..9665928ad0be977356a92b9cb71e7f783182750f
--- /dev/null
+++ b/flux_rgbd/_flux2/autoencoder.py
@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FLUX.2 autoencoder.
+
+Adapted from the FLUX.2 codebase:
+ https://github.com/black-forest-labs/flux2
+"""
+
+import math
+from collections.abc import Sequence
+
+import torch
+from einops import rearrange
+from jaxtyping import Float
+from safetensors.torch import load_file as load_safetensors
+from torch import Tensor, nn
+
+__all__ = ["Flux2Encoder", "Flux2Decoder"]
+
+# FLUX.2-dev autoencoder weights, pre-split into encoder + decoder
+# safetensors. (BFL's canonical `ae.safetensors` is a combined file
+# we'd have to filter by key prefix — easier to host the split files
+# inside the model repo. Override via the FLUX_RGBD_LOCAL_AE_{...}
+# env vars if you want to skip the Hub fetch entirely.)
+DEFAULT_AE_REPO = "bartduis/modality_forcing"
+DEFAULT_ENCODER_FILE = "ae_encoder.safetensors"
+DEFAULT_DECODER_FILE = "ae_decoder.safetensors"
+
+DEFAULT_RESOLUTION = 256
+DEFAULT_IN_CHANNELS = 3
+DEFAULT_CH = 128
+DEFAULT_OUT_CHANNELS = 3
+DEFAULT_CH_MULT = (1, 2, 4, 4)
+DEFAULT_NUM_RES_BLOCKS = 2
+DEFAULT_Z_CHANNELS = 32
+
+
+def _hf_download(repo_id: str, filename: str) -> str:
+ """Resolve a HuggingFace Hub file path. Lazy-import hf_hub to keep the
+ base package importable in environments that pin huggingface_hub via
+ extras."""
+ from huggingface_hub import hf_hub_download
+
+ return hf_hub_download(repo_id=repo_id, filename=filename)
+
+
+def _load_weights(repo_id: str, filename: str) -> dict[str, Tensor]:
+ # Allow callers to point at a local file via env var. Useful for tests
+ # and air-gapped environments.
+ import os
+ env_key = f"FLUX_RGBD_LOCAL_{filename.upper().replace('.', '_')}"
+ override = os.environ.get(env_key)
+ if override:
+ if not os.path.isfile(override):
+ raise FileNotFoundError(f"{env_key}={override} is set but not a file")
+ return load_safetensors(override)
+ path = _hf_download(repo_id, filename)
+ return load_safetensors(path)
+
+
+def _swish(x: Tensor) -> Tensor:
+ return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+ """Attention block for FLUX.2 autoencoder."""
+
+ def __init__(self, in_channels: int):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = nn.GroupNorm(
+ num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+ )
+
+ self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+ self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+ self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+ self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+ def attention(self, h_: Tensor) -> Tensor:
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ b, c, h, w = q.shape
+ q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+ k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+ v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+ h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+ return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+ """Resnet block for FLUX.2 autoencoder."""
+
+ def __init__(self, in_channels: int, out_channels: int):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+
+ self.norm1 = nn.GroupNorm(
+ num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+ )
+ self.conv1 = nn.Conv2d(
+ in_channels, out_channels, kernel_size=3, stride=1, padding=1
+ )
+ self.norm2 = nn.GroupNorm(
+ num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
+ )
+ self.conv2 = nn.Conv2d(
+ out_channels, out_channels, kernel_size=3, stride=1, padding=1
+ )
+ if self.in_channels != self.out_channels:
+ self.nin_shortcut = nn.Conv2d(
+ in_channels, out_channels, kernel_size=1, stride=1, padding=0
+ )
+
+ def forward(self, x):
+ h = x
+ h = self.norm1(h)
+ h = _swish(h)
+ h = self.conv1(h)
+
+ h = self.norm2(h)
+ h = _swish(h)
+ h = self.conv2(h)
+
+ if self.in_channels != self.out_channels:
+ x = self.nin_shortcut(x)
+
+ return x + h
+
+
+class Downsample(nn.Module):
+ """Downsample block for FLUX.2 autoencoder."""
+
+ def __init__(self, in_channels: int):
+ super().__init__()
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = nn.Conv2d(
+ in_channels, in_channels, kernel_size=3, stride=2, padding=0
+ )
+
+ def forward(self, x: Tensor):
+ pad = (0, 1, 0, 1)
+ x = nn.functional.pad(x, pad, mode="constant", value=0)
+ x = self.conv(x)
+ return x
+
+
+class Upsample(nn.Module):
+ """Upsample block for FLUX.2 autoencoder."""
+
+ def __init__(self, in_channels: int):
+ super().__init__()
+ self.conv = nn.Conv2d(
+ in_channels, in_channels, kernel_size=3, stride=1, padding=1
+ )
+
+ def forward(self, x: Tensor):
+ x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+ x = self.conv(x)
+ return x
+
+
+class Flux2Encoder(nn.Module):
+ """Encoder for FLUX.2 autoencoder."""
+
+ def __init__(
+ self,
+ resolution: int = DEFAULT_RESOLUTION,
+ in_channels: int = DEFAULT_IN_CHANNELS,
+ ch: int = DEFAULT_CH,
+ ch_mult: Sequence[int] = DEFAULT_CH_MULT,
+ num_res_blocks: int = DEFAULT_NUM_RES_BLOCKS,
+ z_channels: int = DEFAULT_Z_CHANNELS,
+ repo_id: str = DEFAULT_AE_REPO,
+ filename: str = DEFAULT_ENCODER_FILE,
+ ):
+ """Initialize the FLUX.2 encoder.
+
+ Args:
+ resolution: The resolution of the input images.
+ in_channels: The number of channels in the input images.
+ ch: The number of channels in the encoder.
+ ch_mult: The number of channels in the encoder at each resolution.
+ num_res_blocks: The number of ResNet blocks in each downsampling block.
+ z_channels: The number of channels in the latent space.
+ repo_id: HuggingFace Hub repository id from which to load weights.
+ filename: Filename of the safetensors weights within `repo_id`.
+ """
+ super().__init__()
+ self._repo_id = repo_id
+ self._filename = filename
+
+ self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+ self.ch = ch
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.resolution = resolution
+ self.in_channels = in_channels
+ # downsampling
+ self.conv_in = nn.Conv2d(
+ in_channels, self.ch, kernel_size=3, stride=1, padding=1
+ )
+
+ curr_res = resolution
+ in_ch_mult = (1,) + tuple(ch_mult)
+ self.in_ch_mult = in_ch_mult
+ self.down = nn.ModuleList()
+ block_in = self.ch
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = ch * in_ch_mult[i_level]
+ block_out = ch * ch_mult[i_level]
+ for _ in range(self.num_res_blocks):
+ block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+ block_in = block_out
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level != self.num_resolutions - 1:
+ down.downsample = Downsample(block_in)
+ curr_res = curr_res // 2
+ self.down.append(down)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+ self.mid.attn_1 = AttnBlock(block_in)
+ self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+ # end
+ self.norm_out = nn.GroupNorm(
+ num_groups=32, num_channels=block_in, eps=1e-6, affine=True
+ )
+ self.conv_out = nn.Conv2d(
+ block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
+ )
+
+ self.bn_eps = 1e-4
+ self.bn_momentum = 0.1
+ self.ps = [2, 2]
+ self.bn = torch.nn.BatchNorm2d(
+ math.prod(self.ps) * z_channels,
+ eps=self.bn_eps,
+ momentum=self.bn_momentum,
+ affine=False,
+ track_running_stats=True,
+ )
+
+ def load_weights(self):
+ """Fetch and load weights from HuggingFace Hub."""
+ self.load_state_dict(_load_weights(self._repo_id, self._filename))
+
+ def _forward(self, x: Float[Tensor, "n cx hx wx"]) -> Float[Tensor, "n cz hz wz"]:
+ # downsampling
+ hs = [self.conv_in(x)]
+ for i_level in range(self.num_resolutions):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](hs[-1])
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ hs.append(h)
+ if i_level != self.num_resolutions - 1:
+ hs.append(self.down[i_level].downsample(hs[-1]))
+
+ # middle
+ h = hs[-1]
+ h = self.mid.block_1(h)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h)
+ # end
+ h = self.norm_out(h)
+ h = _swish(h)
+ h = self.conv_out(h)
+ h = self.quant_conv(h)
+ return h
+
+ def _normalize(self, z: Float[Tensor, "n cz hz wz"]) -> Float[Tensor, "n cz hz wz"]:
+ self.bn.eval()
+ return self.bn(z)
+
+ def forward(
+ self, x: Float[Tensor, "... hx wx cx"]
+ ) -> Float[Tensor, "... hz wz cz"]:
+ """Encode images to latents.
+
+ Args:
+ x: Images to encode.
+
+ Returns:
+ The encoded latents. The latents will be normalized using the running
+ statistics from training.
+ """
+ *batch_shape, _, _, _ = x.shape
+ # Convert to NCHW.
+ x = rearrange(x, "... h w c -> (...) c h w")
+
+ # Compute latents.
+ moments = self._forward(x)
+ mean = torch.chunk(moments, 2, dim=1)[0]
+ z = rearrange(
+ mean,
+ "... c (i pi) (j pj) -> ... (c pi pj) i j",
+ pi=self.ps[0],
+ pj=self.ps[1],
+ )
+ z = self._normalize(z)
+
+ # Convert back to NHWC and restore batch shape.
+ z = rearrange(z, "... c h w -> ... h w c")
+ return z.reshape(*batch_shape, *z.shape[-3:])
+
+
+class Flux2Decoder(nn.Module):
+ """Decoder for FLUX.2 autoencoder."""
+
+ def __init__(
+ self,
+ ch: int = DEFAULT_CH,
+ out_ch: int = DEFAULT_OUT_CHANNELS,
+ ch_mult: Sequence[int] = DEFAULT_CH_MULT,
+ num_res_blocks: int = DEFAULT_NUM_RES_BLOCKS,
+ in_channels: int = DEFAULT_IN_CHANNELS,
+ resolution: int = DEFAULT_RESOLUTION,
+ z_channels: int = DEFAULT_Z_CHANNELS,
+ repo_id: str = DEFAULT_AE_REPO,
+ filename: str = DEFAULT_DECODER_FILE,
+ ):
+ """Initialize the FLUX.2 decoder.
+
+ Args:
+ ch: The number of channels in the decoder.
+ out_ch: The number of channels in the output.
+ ch_mult: The number of channels in the decoder at each resolution.
+ num_res_blocks: The number of ResNet blocks in each upsampling block.
+ in_channels: The number of channels in the input images.
+ resolution: The resolution of the input images.
+ z_channels: The number of channels in the latent space.
+ repo_id: HuggingFace Hub repository id from which to load weights.
+ filename: Filename of the safetensors weights within `repo_id`.
+ """
+ super().__init__()
+ self._repo_id = repo_id
+ self._filename = filename
+
+ self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
+ self.ch = ch
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.resolution = resolution
+ self.in_channels = in_channels
+ self.ffactor = 2 ** (self.num_resolutions - 1)
+
+ # compute in_ch_mult, block_in and curr_res at lowest res
+ block_in = ch * ch_mult[self.num_resolutions - 1]
+ curr_res = resolution // 2 ** (self.num_resolutions - 1)
+ self.z_shape = (1, z_channels, curr_res, curr_res)
+
+ # z to block_in
+ self.conv_in = nn.Conv2d(
+ z_channels, block_in, kernel_size=3, stride=1, padding=1
+ )
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+ self.mid.attn_1 = AttnBlock(block_in)
+ self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+ # upsampling
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = ch * ch_mult[i_level]
+ for _ in range(self.num_res_blocks + 1):
+ block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+ block_in = block_out
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if i_level != 0:
+ up.upsample = Upsample(block_in)
+ curr_res = curr_res * 2
+ self.up.insert(0, up) # prepend to get consistent order
+
+ # end
+ self.norm_out = nn.GroupNorm(
+ num_groups=32, num_channels=block_in, eps=1e-6, affine=True
+ )
+ self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+ self.bn_eps = 1e-4
+ self.bn_momentum = 0.1
+ self.ps = [2, 2]
+ self.bn = torch.nn.BatchNorm2d(
+ math.prod(self.ps) * z_channels,
+ eps=self.bn_eps,
+ momentum=self.bn_momentum,
+ affine=False,
+ track_running_stats=True,
+ )
+
+ def load_weights(self):
+ """Fetch and load weights from HuggingFace Hub."""
+ self.load_state_dict(_load_weights(self._repo_id, self._filename))
+
+ def _forward(self, z: Float[Tensor, "n cz hz wz"]) -> Float[Tensor, "n cx hx wx"]:
+ z = self.post_quant_conv(z)
+
+ # get dtype for proper tracing
+ upscale_dtype = next(self.up.parameters()).dtype
+
+ # z to block_in
+ h = self.conv_in(z)
+
+ # middle
+ h = self.mid.block_1(h)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h)
+
+ # cast to proper dtype
+ h = h.to(upscale_dtype)
+ # upsampling
+ for i_level in reversed(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](h)
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ if i_level != 0:
+ h = self.up[i_level].upsample(h)
+
+ # end
+ h = self.norm_out(h)
+ h = _swish(h)
+ h = self.conv_out(h)
+ return h
+
+ def _inv_normalize(
+ self, z: Float[Tensor, "n cz hz wz"]
+ ) -> Float[Tensor, "n cz hz wz"]:
+ self.bn.eval()
+ s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps)
+ m = self.bn.running_mean.view(1, -1, 1, 1)
+ return z * s + m
+
+ def forward(
+ self, z: Float[Tensor, "... cz hz wz"]
+ ) -> Float[Tensor, "... cx hx wx"]:
+ """Decode latents to images.
+
+ Args:
+ z: Image latents normalized to have zero mean and unit variance.
+
+ Returns:
+ Decoded images.
+ """
+ *batch_shape, _, _, _ = z.shape
+
+ # Convert to NCHW.
+ z = rearrange(z, "... h w c -> (...) c h w")
+
+ # Denormalize and evaluate decoder.
+ z = self._inv_normalize(z)
+ z = rearrange(
+ z,
+ "... (c pi pj) i j -> ... c (i pi) (j pj)",
+ pi=self.ps[0],
+ pj=self.ps[1],
+ )
+ dec = self._forward(z)
+
+ # Convert back to NHWC and restore batch shape.
+ dec = rearrange(dec, "... c h w -> ... h w c")
+ dec = dec.reshape(*batch_shape, *dec.shape[-3:])
+ return dec
diff --git a/flux_rgbd/_flux2/constants.py b/flux_rgbd/_flux2/constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..872d71ae9ab1d3300f8bb94156e3e819323e6a49
--- /dev/null
+++ b/flux_rgbd/_flux2/constants.py
@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Adapted from the FLUX.2 codebase:
+# https://github.com/black-forest-labs/flux2
+"""
+Architecture parameters and other constants useful for FLUX.2 inference.
+
+This module centralizes all constants used across the FLUX.2 model implementation,
+including model architecture configurations for all public FLUX.2 variants, their
+text encoder parameters, and system prompts for prompt upsampling.
+"""
+
+from types import MappingProxyType
+
+__all__ = [
+ "FLUX2_DEV_PARAMS",
+ "FLUX2_KLEIN_4B_PARAMS",
+ "FLUX2_KLEIN_9B_PARAMS",
+ "OUTPUT_LAYERS_MISTRAL",
+ "OUTPUT_LAYERS_QWEN3",
+ "MAX_LENGTH",
+ "UPSAMPLING_MAX_IMAGE_SIZE",
+ "SYSTEM_MESSAGE",
+ "SYSTEM_MESSAGE_UPSAMPLING_T2I",
+ "SYSTEM_MESSAGE_UPSAMPLING_I2I",
+]
+
+# =============================================================================
+# Model Architecture Parameters
+# =============================================================================
+
+FLUX2_DEV_PARAMS = MappingProxyType(
+ {
+ "in_channels": 128,
+ "context_in_dim": 15360,
+ "hidden_size": 6144,
+ "num_heads": 48,
+ "depth": 8,
+ "depth_single_blocks": 48,
+ "axes_dim": (32, 32, 32, 32),
+ "theta": 2000,
+ "mlp_ratio": 3.0,
+ "use_guidance_embed": True,
+ }
+)
+"""
+Architecture config for FLUX.2 [dev] 32B parameter model (guidance distilled).
+"""
+
+FLUX2_KLEIN_4B_PARAMS = MappingProxyType(
+ {
+ "in_channels": 128,
+ "context_in_dim": 7680,
+ "hidden_size": 3072,
+ "num_heads": 24,
+ "depth": 5,
+ "depth_single_blocks": 20,
+ "axes_dim": (32, 32, 32, 32),
+ "theta": 2000,
+ "mlp_ratio": 3.0,
+ "use_guidance_embed": False,
+ }
+)
+"""
+Architecture config for FLUX.2 [klein] 4B parameter model (size distilled).
+"""
+
+FLUX2_KLEIN_9B_PARAMS = MappingProxyType(
+ {
+ "in_channels": 128,
+ "context_in_dim": 12288,
+ "hidden_size": 4096,
+ "num_heads": 32,
+ "depth": 8,
+ "depth_single_blocks": 24,
+ "axes_dim": (32, 32, 32, 32),
+ "theta": 2000,
+ "mlp_ratio": 3.0,
+ "use_guidance_embed": False,
+ }
+)
+"""
+Immutable architecture configuration for FLUX.2 [klein]-9B model.
+
+This is a mid-size 9B parameter model that balances quality and speed. It uses
+Qwen3-8B as text encoder and provides higher quality than the 4B variant while
+maintaining fast 4-step inference. The model is guidance-distilled.
+
+Key architectural features:
+ - 24 single-stream transformer blocks (between 4B and dev)
+ - 4096 hidden dimensions with 32 attention heads
+ - No guidance embeddings (guidance-distilled into weights)
+ - Context dimension 12288 matches concatenated Qwen3-8B outputs
+
+Default generation parameters (distilled variants):
+ - guidance: 1.0 (fixed, baked into weights)
+ - num_steps: 4 (fixed)
+
+Base variant generation parameters:
+ - guidance: 4.0 (adjustable)
+ - num_steps: 50 (adjustable)
+"""
+
+# =============================================================================
+# Text Encoder Configuration
+# =============================================================================
+
+OUTPUT_LAYERS_MISTRAL = (10, 20, 30)
+"""
+Layer indices to extract hidden states from Mistral-Small-3.2-24B text encoder.
+
+The FLUX.2 [dev] model uses Mistral-Small as its text encoder and concatenates
+hidden states from these three intermediate layers to form rich text embeddings.
+"""
+
+OUTPUT_LAYERS_QWEN3 = (9, 18, 27)
+"""
+Layer indices to extract hidden states from Qwen3 text encoder variants.
+
+The FLUX.2 [klein] models (4B and 9B) use Qwen3 text encoders and concatenate
+hidden states from these three intermediate layers to form text embeddings.
+"""
+
+MAX_LENGTH = 512
+"""
+Maximum sequence length for text encoder tokenization.
+
+Text prompts are truncated or padded to this length during encoding. This
+applies to both Mistral-Small and Qwen3 text encoders.
+"""
+
+UPSAMPLING_MAX_IMAGE_SIZE = 768**2
+"""
+Maximum pixel area for images during prompt upsampling.
+
+When images are provided as conditioning for prompt upsampling (image-to-image
+mode), they are resized to have at most this many pixels while preserving
+aspect ratio. This prevents out-of-memory errors when processing large images.
+"""
+
+# =============================================================================
+# System Messages for Prompt Upsampling
+# =============================================================================
+
+SYSTEM_MESSAGE = """You are an AI that reasons about image descriptions. You give \
+structured responses focusing on object relationships, object attribution and \
+actions without speculation."""
+"""
+Default system message for text encoder when generating image descriptions.
+
+This system message is used for general-purpose text encoding tasks where the
+model needs to reason about images in a structured, factual manner.
+"""
+
+SYSTEM_MESSAGE_UPSAMPLING_T2I = """You are an expert prompt engineer for FLUX.2 \
+by Black Forest Labs. Rewrite user prompts to be more descriptive while \
+strictly preserving their core subject and intent.
+
+Guidelines:
+1. Structure: Keep structured inputs structured (enhance within fields). \
+Convert natural language to detailed paragraphs.
+2. Details: Add concrete visual specifics - form, scale, textures, materials, \
+lighting (quality, direction, color), shadows, spatial relationships, and \
+environmental context.
+3. Text in Images: Put ALL text in quotation marks, matching the prompt's \
+language. Always provide explicit quoted text for objects that would contain \
+text in reality (signs, labels, screens, etc.) - without it, the model \
+generates gibberish.
+
+Output only the revised prompt and nothing else."""
+"""
+System message for text-to-image prompt upsampling.
+
+This prompt engineering template guides the text encoder to expand sparse user
+prompts into detailed, visually-rich descriptions that improve image generation
+quality. It emphasizes adding concrete visual details while preserving the
+original intent.
+"""
+
+SYSTEM_MESSAGE_UPSAMPLING_I2I = """You are FLUX.2 by Black Forest Labs, an \
+image-editing expert. You convert editing requests into one concise \
+instruction (50-80 words, ~30 for brief requests).
+
+Rules:
+- Single instruction only, no commentary
+- Use clear, analytical language (avoid "whimsical," "cascading," etc.)
+- Specify what changes AND what stays the same (face, lighting, composition)
+- Reference actual image elements
+- Turn negatives into positives ("don't change X" → "keep X")
+- Make abstractions concrete ("futuristic" → "glowing cyan neon, metallic panels")
+- Keep content PG-13
+
+Output only the final instruction in plain text and nothing else."""
+"""
+System message for image-to-image prompt upsampling.
+
+This prompt engineering template guides the text encoder to convert free-form
+image editing requests into concise, actionable instructions suitable for
+image-to-image generation tasks.
+"""
diff --git a/flux_rgbd/_flux2/model.py b/flux_rgbd/_flux2/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1b0abfbbbace43c76b8c54f1338bb8fa61c9dff
--- /dev/null
+++ b/flux_rgbd/_flux2/model.py
@@ -0,0 +1,674 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Adapted from the FLUX.2 codebase:
+# https://github.com/black-forest-labs/flux2
+"""
+FLUX.2 diffusion transformer architecture for image generation and editing.
+
+This module implements the core transformer architecture for FLUX.2 models from
+Black Forest Labs. The architecture uses dual-stream and single-stream transformer
+blocks to process text and image latents for text-to-image and image-to-image tasks.
+"""
+
+import math
+
+import einops
+import torch
+from torch import Tensor, nn
+
+
+class Flux2(nn.Module):
+ """
+ FLUX.2 diffusion transformer for image generation and editing.
+
+ This is a flow-matching diffusion model that uses a stack of dual-stream and
+ single-stream transformer blocks for text context and image latents. The
+ model supports text-to-image and image-to-image generation tasks.
+
+ Default parameter values match the FLUX.2 [klein] 4B architecture, which is
+ optimized for fast inference. For other model variants (klein-9B or dev),
+ use the parameters from :mod:`flux_rgbd._flux2.constants`.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 128,
+ context_in_dim: int = 7680,
+ hidden_size: int = 3072,
+ num_heads: int = 24,
+ depth: int = 5,
+ depth_single_blocks: int = 20,
+ axes_dim: tuple[int, int, int, int] = (32, 32, 32, 32),
+ theta: int = 2000,
+ mlp_ratio: float = 3.0,
+ use_guidance_embed: bool = False,
+ ):
+ """
+ Args:
+ in_channels: Number of input channels for image latents. Matches the
+ output dimension of the autoencoder used to encode images.
+ context_in_dim: Dimension of text context embeddings from the text
+ encoder. This should match the concatenated output dimension of
+ the text encoder being used (7680 for Qwen3-4B, 12288 for Qwen3-8B,
+ 15360 for Mistral-Small).
+ hidden_size: Hidden dimension size for transformer blocks. All
+ attention and MLP operations use this dimension internally.
+ num_heads: Number of attention heads in multi-head attention layers.
+ Must evenly divide `hidden_size`.
+ depth: Number of dual-stream transformer blocks. These blocks process
+ text and image streams separately with cross-attention.
+ depth_single_blocks: Number of single-stream transformer blocks. These
+ blocks process the concatenated text+image sequence.
+ axes_dim: Tuple of 4 integers specifying the dimensionality for each
+ axis in rotary embeddings. Must sum to `hidden_size / num_heads`.
+ theta: Base frequency for rotary position embeddings (RoPE). Higher
+ values result in slower position encoding rotation.
+ mlp_ratio: Expansion ratio for MLP hidden dimension relative to
+ `hidden_size`. MLP hidden dim = `hidden_size * mlp_ratio`.
+ use_guidance_embed: Whether to include guidance scale embeddings for
+ classifier-free guidance. Set to False for distilled models where
+ guidance is baked into weights.
+ """
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = in_channels
+ if hidden_size % num_heads != 0:
+ raise ValueError(
+ f"Hidden size {hidden_size} must be divisible by "
+ f"num_heads {num_heads}"
+ )
+ pe_dim = hidden_size // num_heads
+ if sum(axes_dim) != pe_dim:
+ raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}")
+ self.hidden_size = hidden_size
+ self.num_heads = num_heads
+ self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
+ self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=False)
+ self.time_in = MLPEmbedder(
+ in_dim=256, hidden_dim=self.hidden_size, disable_bias=True
+ )
+ self.txt_in = nn.Linear(context_in_dim, self.hidden_size, bias=False)
+
+ self.use_guidance_embed = use_guidance_embed
+ if self.use_guidance_embed:
+ self.guidance_in = MLPEmbedder(
+ in_dim=256, hidden_dim=self.hidden_size, disable_bias=True
+ )
+
+ double_blocks = [
+ DoubleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio)
+ for _ in range(depth)
+ ]
+ self.double_blocks = nn.ModuleList(double_blocks)
+
+ single_blocks = [
+ SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio)
+ for _ in range(depth_single_blocks)
+ ]
+ self.single_blocks = nn.ModuleList(single_blocks)
+
+ self.double_stream_modulation_img = Modulation(
+ self.hidden_size, double=True, disable_bias=True
+ )
+ self.double_stream_modulation_txt = Modulation(
+ self.hidden_size, double=True, disable_bias=True
+ )
+ self.single_stream_modulation = Modulation(
+ self.hidden_size, double=False, disable_bias=True
+ )
+
+ self.final_layer = LastLayer(self.hidden_size, self.out_channels)
+
+ def forward(
+ self,
+ x: Tensor,
+ x_ids: Tensor,
+ timesteps: Tensor,
+ ctx: Tensor,
+ ctx_ids: Tensor,
+ guidance: Tensor | None,
+ ):
+ num_txt_tokens = ctx.shape[1]
+
+ timestep_emb = timestep_embedding(timesteps, 256)
+ vec = self.time_in(timestep_emb)
+ if self.use_guidance_embed:
+ guidance_emb = timestep_embedding(guidance, 256)
+ vec = vec + self.guidance_in(guidance_emb)
+
+ double_block_mod_img = self.double_stream_modulation_img(vec)
+ double_block_mod_txt = self.double_stream_modulation_txt(vec)
+ single_block_mod, _ = self.single_stream_modulation(vec)
+
+ img = self.img_in(x)
+ txt = self.txt_in(ctx)
+
+ pe_x = self.pe_embedder(x_ids)
+ pe_ctx = self.pe_embedder(ctx_ids)
+
+ for block in self.double_blocks:
+ img, txt = block(
+ img,
+ txt,
+ pe_x,
+ pe_ctx,
+ double_block_mod_img,
+ double_block_mod_txt,
+ )
+
+ img = torch.cat((txt, img), dim=1)
+ pe = torch.cat((pe_ctx, pe_x), dim=2)
+
+ for block in self.single_blocks:
+ img = block(
+ img,
+ pe,
+ single_block_mod,
+ )
+
+ img = img[:, num_txt_tokens:, ...]
+
+ img = self.final_layer(img, vec)
+ return img
+
+
+class SelfAttention(nn.Module):
+ """
+ Multi-head self-attention with QK normalization.
+
+ This module computes query, key, and value projections in a single linear layer,
+ applies RMS normalization to queries and keys, then performs attention and projects
+ the output back to the original dimension.
+ """
+
+ def __init__(self, dim: int, num_heads: int = 8):
+ """
+ Args:
+ dim: Hidden dimension size. Must be divisible by `num_heads`.
+ num_heads: Number of parallel attention heads.
+ """
+ super().__init__()
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.qkv = nn.Linear(dim, dim * 3, bias=False)
+
+ self.norm = QKNorm(head_dim)
+ self.proj = nn.Linear(dim, dim, bias=False)
+
+
+class SiLUActivation(nn.Module):
+ """
+ Gated activation using SiLU (Swish) function.
+
+ This module splits the input tensor along the last dimension, applies SiLU to one
+ half, and element-wise multiplies it with the other half. This is commonly known
+ as SwiGLU when used in MLP layers.
+ """
+
+ def __init__(self):
+ super().__init__()
+ self.gate_fn = nn.SiLU()
+
+ def forward(self, x: Tensor) -> Tensor:
+ """
+ Args:
+ x: Input tensor of shape `(..., 2 * dim)` where the last dimension will
+ be split into two equal parts for gating.
+
+ Returns:
+ Gated output of shape `(..., dim)`.
+ """
+ x1, x2 = x.chunk(2, dim=-1)
+ return self.gate_fn(x1) * x2
+
+
+class Modulation(nn.Module):
+ """
+ Adaptive layer normalization (AdaLN) modulation layer.
+
+ This module generates scale, shift, and gate parameters for adaptive normalization
+ from timestep/guidance embeddings. For double-stream blocks, it produces two sets
+ of modulation parameters (one for each stream).
+ """
+
+ def __init__(self, dim: int, double: bool, disable_bias: bool = False):
+ """
+ Args:
+ dim: Hidden dimension size matching the transformer blocks.
+ double: If True, generates parameters for dual-stream blocks (6 params:
+ shift, scale, gate for each stream). If False, generates for
+ single-stream blocks (3 params: shift, scale, gate).
+ disable_bias: If True, the linear layer has no bias term.
+ """
+ super().__init__()
+ self.is_double = double
+ self.multiplier = 6 if double else 3
+ self.lin = nn.Linear(dim, self.multiplier * dim, bias=not disable_bias)
+
+ def forward(self, vec: Tensor):
+ """
+ Args:
+ vec: Timestep/guidance embedding of shape `(batch_size, dim)` or
+ `(batch_size, seq_len, dim)`.
+
+ Returns:
+ Tuple of modulation parameters. For single-stream: `(mod, None)` where
+ `mod` is a 3-tuple of (shift, scale, gate). For double-stream:
+ `(mod1, mod2)` where each is a 3-tuple for different streams.
+ """
+ out = self.lin(nn.functional.silu(vec))
+ if out.ndim == 2:
+ out = out[:, None, :]
+ out = out.chunk(self.multiplier, dim=-1)
+ return out[:3], out[3:] if self.is_double else None
+
+
+class LastLayer(nn.Module):
+ """
+ Final output layer with adaptive layer normalization.
+
+ This module applies AdaLN-modulated normalization followed by a linear projection
+ to map transformer hidden states back to the output space (image latent channels).
+ """
+
+ def __init__(self, hidden_size: int, out_channels: int):
+ """
+ Args:
+ hidden_size: Hidden dimension of transformer blocks.
+ out_channels: Output dimension (number of latent channels).
+ """
+ super().__init__()
+ self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.linear = nn.Linear(hidden_size, out_channels, bias=False)
+ self.adaLN_modulation = nn.Sequential( # pylint: disable=invalid-name
+ nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=False)
+ )
+
+ def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+ """
+ Args:
+ x: Hidden states from transformer as a tensor of shape
+ `(batch_size, seq_len, hidden_size)`.
+ vec: Timestep embedding of shape `(batch_size, hidden_size)`.
+
+ Returns:
+ Output tensor of shape `(batch_size, seq_len, out_channels)`.
+ """
+ mod = self.adaLN_modulation(vec)
+ shift, scale = mod.chunk(2, dim=-1)
+ if shift.ndim == 2:
+ shift = shift[:, None, :]
+ scale = scale[:, None, :]
+ x = (1 + scale) * self.norm_final(x) + shift
+ x = self.linear(x)
+ return x
+
+
+class SingleStreamBlock(nn.Module):
+ """
+ Single-stream transformer block processing concatenated text+image tokens.
+
+ This block applies self-attention and MLP operations to a unified sequence of
+ text and image tokens. Both operations share pre-normalization and use adaptive
+ modulation from timestep embeddings.
+ """
+
+ def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0):
+ """
+ Args:
+ hidden_size: Hidden dimension size for all linear projections.
+ num_heads: Number of attention heads. Must divide `hidden_size` evenly.
+ mlp_ratio: Ratio of MLP hidden dimension to `hidden_size`. The actual
+ MLP hidden dim is `int(hidden_size * mlp_ratio)`.
+ """
+ super().__init__()
+
+ self.hidden_dim = hidden_size
+ self.num_heads = num_heads
+ head_dim = hidden_size // num_heads
+ self.scale = head_dim**-0.5
+ self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ self.mlp_mult_factor = 2
+
+ self.linear1 = nn.Linear(
+ hidden_size,
+ hidden_size * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
+ bias=False,
+ )
+ self.linear2 = nn.Linear(
+ hidden_size + self.mlp_hidden_dim, hidden_size, bias=False
+ )
+ self.norm = QKNorm(head_dim)
+
+ self.hidden_size = hidden_size
+ self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+ self.mlp_act = SiLUActivation()
+
+ def forward(self, x: Tensor, pe: Tensor, mod: tuple[Tensor, Tensor]) -> Tensor:
+ mod_shift, mod_scale, mod_gate = mod
+ x_mod = (1 + mod_scale) * self.pre_norm(x) + mod_shift
+
+ qkv, mlp = torch.split(
+ self.linear1(x_mod),
+ [3 * self.hidden_size, self.mlp_hidden_dim * self.mlp_mult_factor],
+ dim=-1,
+ )
+
+ q, k, v = einops.rearrange(
+ qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+ )
+ q, k = self.norm(q, k, v)
+
+ attn = attention(q, k, v, pe)
+
+ # Compute activation in mlp stream, cat again and run second linear layer.
+ output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+ return x + mod_gate * output
+
+
+class DoubleStreamBlock(nn.Module):
+ """
+ Dual-stream transformer block processing text and image tokens separately.
+
+ This block maintains separate streams for text and image tokens, each with their
+ own self-attention and MLP sublayers. Cross-stream information exchange happens
+ through joint attention where Q, K, V from both streams are concatenated.
+ """
+
+ def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float):
+ """
+ Args:
+ hidden_size: Hidden dimension size for all linear projections.
+ num_heads: Number of attention heads. Must divide `hidden_size` evenly.
+ mlp_ratio: Ratio of MLP hidden dimension to `hidden_size`.
+ """
+ super().__init__()
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ self.num_heads = num_heads
+ assert (
+ hidden_size % num_heads == 0
+ ), f"{hidden_size=} must be divisible by {num_heads=}"
+
+ self.hidden_size = hidden_size
+ self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.mlp_mult_factor = 2
+
+ self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads)
+
+ self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.img_mlp = nn.Sequential(
+ nn.Linear(hidden_size, mlp_hidden_dim * self.mlp_mult_factor, bias=False),
+ SiLUActivation(),
+ nn.Linear(mlp_hidden_dim, hidden_size, bias=False),
+ )
+
+ self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads)
+
+ self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.txt_mlp = nn.Sequential(
+ nn.Linear(
+ hidden_size,
+ mlp_hidden_dim * self.mlp_mult_factor,
+ bias=False,
+ ),
+ SiLUActivation(),
+ nn.Linear(mlp_hidden_dim, hidden_size, bias=False),
+ )
+
+ def forward(
+ self,
+ img: Tensor,
+ txt: Tensor,
+ pe: Tensor,
+ pe_ctx: Tensor,
+ mod_img: tuple[Tensor, Tensor],
+ mod_txt: tuple[Tensor, Tensor],
+ ) -> tuple[Tensor, Tensor]:
+ img_mod1, img_mod2 = mod_img
+ txt_mod1, txt_mod2 = mod_txt
+
+ img_mod1_shift, img_mod1_scale, img_mod1_gate = img_mod1
+ img_mod2_shift, img_mod2_scale, img_mod2_gate = img_mod2
+ txt_mod1_shift, txt_mod1_scale, txt_mod1_gate = txt_mod1
+ txt_mod2_shift, txt_mod2_scale, txt_mod2_gate = txt_mod2
+
+ # Prepare image for attention.
+ img_modulated = self.img_norm1(img)
+ img_modulated = (1 + img_mod1_scale) * img_modulated + img_mod1_shift
+
+ img_qkv = self.img_attn.qkv(img_modulated)
+ img_q, img_k, img_v = einops.rearrange(
+ img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+ )
+ img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+ # Prepare txt for attention.
+ txt_modulated = self.txt_norm1(txt)
+ txt_modulated = (1 + txt_mod1_scale) * txt_modulated + txt_mod1_shift
+
+ txt_qkv = self.txt_attn.qkv(txt_modulated)
+ txt_q, txt_k, txt_v = einops.rearrange(
+ txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+ )
+ txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+ q = torch.cat((txt_q, img_q), dim=2)
+ k = torch.cat((txt_k, img_k), dim=2)
+ v = torch.cat((txt_v, img_v), dim=2)
+
+ pe = torch.cat((pe_ctx, pe), dim=2)
+ attn = attention(q, k, v, pe)
+ txt_attn, img_attn = attn[:, : txt_q.shape[2]], attn[:, txt_q.shape[2] :]
+
+ # Calculate the img blocks.
+ img = img + img_mod1_gate * self.img_attn.proj(img_attn)
+ img = img + img_mod2_gate * self.img_mlp(
+ (1 + img_mod2_scale) * (self.img_norm2(img)) + img_mod2_shift
+ )
+
+ # Calculate the txt blocks.
+ txt = txt + txt_mod1_gate * self.txt_attn.proj(txt_attn)
+ txt = txt + txt_mod2_gate * self.txt_mlp(
+ (1 + txt_mod2_scale) * (self.txt_norm2(txt)) + txt_mod2_shift
+ )
+ return img, txt
+
+
+class MLPEmbedder(nn.Module):
+ """
+ Two-layer MLP for embedding timestep and guidance values.
+
+ This simple MLP transforms scalar timestep or guidance embeddings (after
+ sinusoidal encoding) into the transformer's hidden dimension space.
+ """
+
+ def __init__(self, in_dim: int, hidden_dim: int, disable_bias: bool = False):
+ """
+ Args:
+ in_dim: Input dimension (typically 256 for sinusoidal embeddings).
+ hidden_dim: Output hidden dimension matching transformer blocks.
+ disable_bias: If True, linear layers have no bias terms.
+ """
+ super().__init__()
+ self.in_layer = nn.Linear(in_dim, hidden_dim, bias=not disable_bias)
+ self.silu = nn.SiLU()
+ self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=not disable_bias)
+
+ def forward(self, x: Tensor) -> Tensor:
+ """
+ Args:
+ x: Input embeddings of shape `(batch_size, in_dim)`.
+
+ Returns:
+ Projected embeddings of shape `(batch_size, hidden_dim)`.
+ """
+ return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class EmbedND(nn.Module):
+ """
+ N-dimensional rotary position embeddings (RoPE) for spatial-temporal tokens.
+
+ This module creates rotary embeddings for multi-dimensional position indices
+ (e.g., time, height, width, sequence). Each dimension gets its own embedding
+ component with configurable dimensions.
+ """
+
+ def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+ """
+ Args:
+ dim: Total position embedding dimension. Should equal `sum(axes_dim)`.
+ theta: Base frequency for RoPE. Higher values give slower rotation.
+ axes_dim: Dimension allocation for each position axis.
+ """
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ self.axes_dim = axes_dim
+
+ def forward(self, ids: Tensor) -> Tensor:
+ """
+ Args:
+ ids: Position indices of shape `(..., num_axes)` where `num_axes`
+ matches `len(axes_dim)`.
+
+ Returns:
+ Rotary embeddings of shape `(..., 1, sum(axes_dim), 2, 2)` suitable
+ for applying rotation to query and key tensors.
+ """
+ emb = torch.cat(
+ [
+ rope(ids[..., i], self.axes_dim[i], self.theta)
+ for i in range(len(self.axes_dim))
+ ],
+ dim=-3,
+ )
+
+ return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ t = time_factor * t
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period)
+ * torch.arange(start=0, end=half, device=t.device, dtype=torch.float32)
+ / half
+ )
+
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ if torch.is_floating_point(t):
+ embedding = embedding.to(t)
+ return embedding
+
+
+class RMSNorm(torch.nn.Module):
+ """
+ Root Mean Square Layer Normalization.
+
+ RMSNorm normalizes using only the variance (RMS) without centering by mean,
+ providing a simpler and often equally effective alternative to LayerNorm.
+ """
+
+ def __init__(self, dim: int):
+ """
+ Args:
+ dim: Dimension to normalize over (last dimension of input).
+ """
+ super().__init__()
+ self.scale = nn.Parameter(torch.ones(dim))
+
+ def forward(self, x: Tensor):
+ """
+ Args:
+ x: Input tensor of shape `(..., dim)`.
+
+ Returns:
+ Normalized tensor of same shape as input.
+ """
+ x_dtype = x.dtype
+ x = x.float()
+ rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+ return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+ """
+ Separate RMSNorm for query and key tensors in attention.
+
+ Normalizing queries and keys independently before attention computation improves
+ training stability and can lead to better performance.
+ """
+
+ def __init__(self, dim: int):
+ """
+ Args:
+ dim: Head dimension for queries and keys.
+ """
+ super().__init__()
+ self.query_norm = RMSNorm(dim)
+ self.key_norm = RMSNorm(dim)
+
+ def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+ """
+ Args:
+ q: Query tensor of shape `(..., head_dim)`.
+ k: Key tensor of shape `(..., head_dim)`.
+ v: Value tensor (used only for dtype matching).
+
+ Returns:
+ Tuple of normalized `(query, key)` tensors with dtype matching `v`.
+ """
+ q = self.query_norm(q)
+ k = self.key_norm(k)
+ return q.to(v), k.to(v)
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+ q, k = apply_rope(q, k, pe)
+
+ x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+ x = einops.rearrange(x, "B H L D -> B L (H D)")
+
+ return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+ assert dim % 2 == 0
+ scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+ omega = 1.0 / (theta**scale)
+ out = torch.einsum("...n,d->...nd", pos, omega)
+ out = torch.stack(
+ [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
+ )
+ out = einops.rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+ return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+ xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+ xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+ xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+ xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+ return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
diff --git a/flux_rgbd/depth/__init__.py b/flux_rgbd/depth/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..668d27b8139e3711e9cd7e01739b23257e3eb0df
--- /dev/null
+++ b/flux_rgbd/depth/__init__.py
@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Depth-side preprocessing and schedule helpers for flux_rgbd."""
+
+from flux_rgbd.depth.preprocess import DepthConfig, decode_depth
+from flux_rgbd.depth.schedule import ScheduleConfig, rollout_timesteps
+
+__all__ = ["DepthConfig", "decode_depth", "ScheduleConfig", "rollout_timesteps"]
diff --git a/flux_rgbd/depth/preprocess.py b/flux_rgbd/depth/preprocess.py
new file mode 100644
index 0000000000000000000000000000000000000000..fcf56e6d2358f9ede2e077f3fe7d8dcd1fafc7c3
--- /dev/null
+++ b/flux_rgbd/depth/preprocess.py
@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Depth-side decoding: model tokens → depth pixel map.
+
+Decoding steps:
+ 1. Unpatchify the (tok_h, tok_w, patch_h * patch_w) token grid.
+ 2. Divide out the post-normalization scalar `depth_value_scale`.
+ 3. Apply the inverse of every recorded normalization stage, in reverse.
+
+Only `"contract"` (the mip-NeRF 360 radial squash, arXiv 2111.12077
+Eq. 10) is genuinely invertible — the other stages discard per-sample
+statistics during encoding and pass through unchanged on decode.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+
+import einops
+import torch
+from torch import Tensor
+
+_KNOWN_STAGES = frozenset({"contract", "normal", "mean", "unit_mean", "none"})
+
+
+@dataclasses.dataclass(frozen=True)
+class DepthConfig:
+ """How depth was normalized at training time. Decoder reads these."""
+
+ depth_normalize_mode: tuple[str, ...] | str = ("unit_mean", "contract")
+ patch_size: int = 16
+ depth_value_scale: float = 1.0
+
+ def __post_init__(self) -> None:
+ stages = self.depth_normalize_mode
+ if isinstance(stages, str):
+ object.__setattr__(self, "depth_normalize_mode", (stages,))
+ unknown = set(self.depth_normalize_mode) - _KNOWN_STAGES
+ if unknown:
+ raise ValueError(f"unknown stages {sorted(unknown)}; expected {sorted(_KNOWN_STAGES)}")
+ if self.depth_value_scale <= 0:
+ raise ValueError(f"depth_value_scale must be > 0; got {self.depth_value_scale}")
+ if self.patch_size <= 0:
+ raise ValueError(f"patch_size must be > 0; got {self.patch_size}")
+
+
+def _inv_contract(z: Tensor) -> Tensor:
+ """Inverse of the mip-NeRF 360 radial contract bijection (operates on last dim).
+
+ Contract maps R^n into the ball of radius 2; this undoes it. For |z| <= 1
+ the contract was identity, so this is also identity in that region.
+ """
+ mag_sq = torch.clamp(z.square().sum(dim=-1, keepdim=True), min=1.0)
+ return z / (2.0 * mag_sq.sqrt() - mag_sq)
+
+
+def _undo_stage(depth: Tensor, stage: str) -> Tensor:
+ if stage != "contract": # "normal"/"mean"/"unit_mean"/"none" are pass-through.
+ return depth
+ if depth.shape[-1] == 1:
+ return _inv_contract(depth)
+ return _inv_contract(depth.unsqueeze(-1)).squeeze(-1)
+
+
+def decode_depth(tokens: Tensor, config: DepthConfig) -> Tensor:
+ """Token grid -> depth pixel map.
+
+ Input shape: (..., tok_h, tok_w, patch_h * patch_w)
+ Output shape: (..., tok_h * patch_size, tok_w * patch_size, 1)
+ """
+ p = config.patch_size
+ depth = einops.rearrange(tokens, "... h w (p1 p2 c) -> ... (h p1) (w p2) c", p1=p, p2=p)
+ if config.depth_value_scale != 1.0:
+ depth = depth / config.depth_value_scale
+ for stage in reversed(config.depth_normalize_mode):
+ depth = _undo_stage(depth, stage)
+ return depth
+
+
+def _contract(z: Tensor) -> Tensor:
+ """Forward mip-NeRF 360 radial contract (inverse of `_inv_contract`).
+
+ Identity for |z| <= 1; otherwise squashes radius r into (2 - 1/r).
+ """
+ mag = z.norm(dim=-1, keepdim=True)
+ safe = mag.clamp(min=1e-12)
+ scale = torch.where(mag <= 1.0, torch.ones_like(mag), (2.0 - 1.0 / safe) / safe)
+ return z * scale
+
+
+def _apply_stage(depth: Tensor, stage: str) -> Tensor:
+ """Forward of `_undo_stage`. Only `contract` and `unit_mean` are active.
+
+ `unit_mean` divides by the per-map mean so the depth is scale-normalised
+ the way the model was trained; it is pass-through on decode (the absolute
+ scale is not recoverable), so encode/decode round-trips up to that scale.
+ """
+ if stage == "unit_mean":
+ return depth / depth.mean().clamp(min=1e-12)
+ if stage != "contract": # "normal"/"mean"/"none" are pass-through.
+ return depth
+ if depth.shape[-1] == 1:
+ return _contract(depth)
+ return _contract(depth.unsqueeze(-1)).squeeze(-1)
+
+
+def encode_depth(depth_map: Tensor, config: DepthConfig) -> Tensor:
+ """Depth pixel map -> token grid (inverse of `decode_depth`).
+
+ Input shape: (..., H, W, 1)
+ Output shape: (..., H/patch_size, W/patch_size, patch_size**2)
+
+ Used for ``mode="d2i"`` to turn a depth map into the depth-stream tokens
+ the model conditions on. Because `unit_mean` is scale-normalising, only the
+ relative depth structure matters — the input need not be metric.
+ """
+ p = config.patch_size
+ depth = depth_map
+ for stage in config.depth_normalize_mode:
+ depth = _apply_stage(depth, stage)
+ if config.depth_value_scale != 1.0:
+ depth = depth * config.depth_value_scale
+ return einops.rearrange(depth, "... (h p1) (w p2) c -> ... h w (p1 p2 c)", p1=p, p2=p)
diff --git a/flux_rgbd/depth/schedule.py b/flux_rgbd/depth/schedule.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a7494a14d16249042abd66c8489ae24e111f26f
--- /dev/null
+++ b/flux_rgbd/depth/schedule.py
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Per-modality flow-matching timestep schedules for RGB + depth inference.
+
+Each modality walks from t = t_max down to t = t_min, warped by the
+SD3 / FLUX shift transform
+
+ t' = exp(mu) / (exp(mu) + (1 / t - 1) ** sigma)
+
+so larger mu pushes the schedule toward noisier (later) timesteps.
+
+In asymmetric modes one stream is held at the boundary of its noise
+range:
+ "joint" — both schedules are warped linear walks.
+ "i2d" — RGB held at t = 0 (clean RGB conditioning).
+ "d2i" — depth held at t = 0; RGB optionally held at t = 1.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import math
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+Mode = Literal["joint", "i2d", "d2i"]
+
+
+@dataclasses.dataclass
+class ScheduleConfig:
+ rgb_shift_mu: float = 0.0
+ rgb_shift_sigma: float = 1.0
+ depth_shift_mu: float = 0.0
+ depth_shift_sigma: float = 1.0
+ # Some d2i checkpoints were trained with RGB pinned at full noise; this
+ # flag mirrors that conditioning at sample time.
+ d2i_full_noise_rgb: bool = False
+
+
+def _shift(mu: float, sigma: float, t: Tensor) -> Tensor:
+ return math.exp(mu) / (math.exp(mu) + (1.0 / t - 1.0) ** sigma)
+
+
+def _f_alpha(t: Tensor, alpha: float) -> Tensor:
+ """Latent-forcing Möbius warp f_α(t) = αt / (1 + (α-1)t). Fixes 0 and 1."""
+ return (alpha * t) / (1.0 + (alpha - 1.0) * t)
+
+
+def rollout_timesteps(config: ScheduleConfig, num_steps: int, *,
+ mode: Mode = "joint",
+ log2_alpha: float | None = None,
+ t_max: float = 1.0, t_min: float = 0.0,
+ device: torch.device | str | None = None,
+ dtype: torch.dtype = torch.float32) -> tuple[Tensor, Tensor]:
+ """Return per-modality timestep tensors of shape `(num_steps + 1,)`.
+
+ ``log2_alpha`` (joint mode only) tilts the RGB/depth denoising trajectory:
+ the depth schedule becomes ``f_α(t_rgb)`` with ``α = 2 ** log2_alpha`` on top
+ of the training time-shift density. ``α > 1`` keeps depth noisier for longer
+ so RGB resolves first (rgb-first → cleaner depth); ``α < 1`` is depth-first;
+ ``α = 1`` (or ``None``) is the diagonal joint schedule.
+ """
+ if num_steps < 1:
+ raise ValueError(f"num_steps must be >= 1; got {num_steps}")
+
+ linear = torch.linspace(t_max, t_min, num_steps + 1, device=device, dtype=dtype)
+ t_rgb = _shift(config.rgb_shift_mu, config.rgb_shift_sigma, linear)
+ t_depth = _shift(config.depth_shift_mu, config.depth_shift_sigma, linear)
+
+ if mode == "i2d":
+ t_rgb = torch.zeros_like(t_rgb)
+ elif mode == "d2i":
+ if config.d2i_full_noise_rgb:
+ t_rgb = torch.ones_like(t_rgb)
+ t_depth = torch.zeros_like(t_depth)
+ elif mode == "joint":
+ if log2_alpha is not None:
+ # Warp the depth schedule off the (time-shifted) RGB grid, matching
+ # the final-paper time-shift trajectory: nodes follow the training
+ # sigmoid density, then the Möbius warp tilts depth vs RGB.
+ t_depth = _f_alpha(t_rgb, 2.0 ** log2_alpha)
+ else:
+ raise ValueError(f"unknown mode {mode!r}; expected joint / i2d / d2i")
+
+ return t_rgb, t_depth
diff --git a/flux_rgbd/dit.py b/flux_rgbd/dit.py
new file mode 100644
index 0000000000000000000000000000000000000000..420760e005ef4be05f8d02b1178ce5e318a69264
--- /dev/null
+++ b/flux_rgbd/dit.py
@@ -0,0 +1,407 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Black Forest Labs.
+# Copyright (c) 2026 World Labs. Modifications: depth stream, joint RGBD
+# attention, depth decoder blocks, cross-stream timestep mixing.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FluxRGBDDiT — FLUX.2 dual+single-stream transformer extended for joint RGB+depth.
+
+Derived from the FLUX.2 reference implementation:
+ https://github.com/black-forest-labs/flux2
+
+Architecture, inference-only:
+ * Depth is a peer stream alongside image and text — each gets its own
+ pre-norm, Q/K/V, MLP inside every dual-stream block. Joint attention
+ runs over the concatenated `[txt, img, depth]` sequence.
+ * 8 dual-stream blocks + 24 single-stream blocks (FLUX.2 backbone).
+ * 4-block depth-only decoder before the depth output head.
+ * 4D RoPE position ids (axes_dim=(32,32,32,32), theta=2000) shared via
+ EmbedND. RGB tokens use time_id=0, depth uses time_id=1 so RoPE can
+ distinguish them.
+
+Reuses the FLUX.2 reference building blocks unchanged so img/txt-side
+weights lifted from an RGB-only FLUX.2 [klein-9B] base load 1:1.
+"""
+
+from __future__ import annotations
+
+import einops
+import torch
+from torch import Tensor, nn
+
+from flux_rgbd._flux2 import model as flux2_model
+from flux_rgbd._flux2.model import (
+ EmbedND,
+ LastLayer,
+ MLPEmbedder,
+ Modulation,
+ SelfAttention,
+ SiLUActivation,
+ SingleStreamBlock,
+ timestep_embedding,
+)
+
+
+class _TripleStreamBlock(nn.Module):
+ """Joint-attention block over (img, text, depth) with per-stream AdaLN.
+
+ Submodule names on the img/txt halves match FLUX.2's DoubleStreamBlock
+ so RGB-only base weights load 1:1; the depth half adds new submodules
+ with the same shapes as img.
+ """
+
+ def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float):
+ super().__init__()
+ if hidden_size % num_heads != 0:
+ raise ValueError(f"hidden_size {hidden_size} not divisible by num_heads {num_heads}")
+ mlp_hidden = int(hidden_size * mlp_ratio)
+ self.num_heads = num_heads
+
+ # Same shapes per stream: pre-norm + (Q,K,V,proj) + post-norm + 2-up SiLU MLP.
+ # The `* 2` in the first Linear is the gated-SiLU activation packing FLUX.2 uses.
+ for prefix in ("img", "txt", "depth"):
+ setattr(self, f"{prefix}_norm1", nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6))
+ setattr(self, f"{prefix}_attn", SelfAttention(dim=hidden_size, num_heads=num_heads))
+ setattr(self, f"{prefix}_norm2", nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6))
+ setattr(self, f"{prefix}_mlp", nn.Sequential(
+ nn.Linear(hidden_size, mlp_hidden * 2, bias=False),
+ SiLUActivation(),
+ nn.Linear(mlp_hidden, hidden_size, bias=False),
+ ))
+
+ def _qkv(self, x: Tensor, norm: nn.LayerNorm, attn: SelfAttention,
+ mod1: tuple[Tensor, Tensor, Tensor]):
+ """Pre-norm + AdaLN + Q/K/V projection. Returns (q, k, v, gate)."""
+ shift, scale, gate = mod1
+ qkv = attn.qkv((1 + scale) * norm(x) + shift)
+ q, k, v = einops.rearrange(qkv, "B L (K H D) -> K B H L D",
+ K=3, H=self.num_heads)
+ q, k = attn.norm(q, k, v)
+ return q, k, v, gate
+
+ @staticmethod
+ def _residual(x, attn_out, proj, gate1, mod2, norm2, mlp):
+ """attn residual then MLP residual, both gated. Standard DiT pattern."""
+ x = x + gate1 * proj(attn_out)
+ shift, scale, gate2 = mod2
+ return x + gate2 * mlp((1 + scale) * norm2(x) + shift)
+
+ def forward(self, img, txt, depth, pe_img, pe_txt, pe_depth,
+ mod_img, mod_txt, mod_depth):
+ q_img, k_img, v_img, g_img = self._qkv(img, self.img_norm1, self.img_attn, mod_img[0])
+ q_txt, k_txt, v_txt, g_txt = self._qkv(txt, self.txt_norm1, self.txt_attn, mod_txt[0])
+ q_d, k_d, v_d, g_d = self._qkv(depth, self.depth_norm1, self.depth_attn, mod_depth[0])
+
+ # Joint attention over [txt, img, depth]; BFL's `attention` does
+ # RoPE + scaled_dot_product_attention + rearrange.
+ q = torch.cat([q_txt, q_img, q_d], dim=2)
+ k = torch.cat([k_txt, k_img, k_d], dim=2)
+ v = torch.cat([v_txt, v_img, v_d], dim=2)
+ pe = torch.cat([pe_txt, pe_img, pe_depth], dim=2)
+ out = flux2_model.attention(q, k, v, pe)
+
+ n_txt, n_img = q_txt.shape[2], q_img.shape[2]
+ txt_out = out[:, :n_txt]
+ img_out = out[:, n_txt : n_txt + n_img]
+ depth_out = out[:, n_txt + n_img :]
+
+ img = self._residual(img, img_out, self.img_attn.proj, g_img,
+ mod_img[1], self.img_norm2, self.img_mlp)
+ txt = self._residual(txt, txt_out, self.txt_attn.proj, g_txt,
+ mod_txt[1], self.txt_norm2, self.txt_mlp)
+ depth = self._residual(depth, depth_out, self.depth_attn.proj, g_d,
+ mod_depth[1], self.depth_norm2, self.depth_mlp)
+ return img, txt, depth
+
+
+def _stack_per_token_mod(
+ mod_img: tuple[Tensor, Tensor, Tensor],
+ mod_depth: tuple[Tensor, Tensor, Tensor],
+ n_txt: int, n_img: int, n_depth: int,
+) -> tuple[Tensor, Tensor, Tensor]:
+ """Build the per-token modulation triple consumed by SingleStreamBlock.
+
+ Each ``Modulation`` output is broadcast over the sequence axis. For the
+ triple-stream single stack we need different (shift, scale, gate) for
+ the depth tokens vs the txt+img tokens — so we expand each entry along
+ seq and ``cat``.
+ """
+ n_txt_img = n_txt + n_img
+
+ def _build(slot: int) -> Tensor:
+ head = mod_img[slot].expand(-1, n_txt_img, -1)
+ tail = mod_depth[slot].expand(-1, n_depth, -1)
+ return torch.cat([head, tail], dim=1)
+
+ return _build(0), _build(1), _build(2)
+
+
+class FluxRGBDDiT(nn.Module):
+ """FLUX.2 Klein-9B DiT with peer depth stream.
+
+ Forward: (img, depth, text, RGB timestep, depth timestep, *_ids)
+ → (rgb_latent_out, depth_latent_out).
+
+ Architecture (v1 defaults match the published checkpoint):
+ depth_double=8 triple-stream blocks → full-token concat →
+ depth_single=24 single-stream blocks → split →
+ depth_decoder_num_layers=4 single blocks over depth only →
+ final_layer / depth_final_layer.
+ """
+
+ def __init__(
+ self,
+ *,
+ in_channels: int = 128,
+ depth_channels: int = 256,
+ context_in_dim: int = 12288,
+ hidden_size: int = 4096,
+ num_heads: int = 32,
+ depth_double: int = 8,
+ depth_single: int = 24,
+ depth_decoder_num_layers: int = 4,
+ axes_dim: tuple[int, ...] = (32, 32, 32, 32),
+ theta: int = 2000,
+ mlp_ratio: float = 3.0,
+ use_guidance_embed: bool = False,
+ cross_stream_timestep_mixing: bool = False,
+ ):
+ super().__init__()
+
+ pe_dim = hidden_size // num_heads
+ if sum(axes_dim) != pe_dim:
+ raise ValueError(
+ f"axes_dim sum {sum(axes_dim)} != hidden_size/num_heads {pe_dim}"
+ )
+
+ self.in_channels = in_channels
+ self.depth_channels = depth_channels
+ self.hidden_size = hidden_size
+ self.num_heads = num_heads
+ self.depth_double = depth_double
+ self.depth_single = depth_single
+ self.depth_decoder_num_layers = depth_decoder_num_layers
+
+ # Stream-in projections.
+ self.img_in = nn.Linear(in_channels, hidden_size, bias=False)
+ self.txt_in = nn.Linear(context_in_dim, hidden_size, bias=False)
+ self.depth_in = nn.Linear(depth_channels, hidden_size, bias=False)
+
+ # Timestep embedders. `time_in` carries the RGB timestep (plus optional
+ # guidance scale); `time_in_depth` carries the depth timestep.
+ self.time_in = MLPEmbedder(
+ in_dim=256, hidden_dim=hidden_size, disable_bias=True
+ )
+ self.time_in_depth = MLPEmbedder(
+ in_dim=256, hidden_dim=hidden_size, disable_bias=True
+ )
+ self.use_guidance_embed = use_guidance_embed
+ if use_guidance_embed:
+ self.guidance_in = MLPEmbedder(
+ in_dim=256, hidden_dim=hidden_size, disable_bias=True
+ )
+
+ # Cross-stream timestep mixing: each stream's modulation vector also
+ # sees the *other* stream's noise level, projected through a small MLP
+ # and gated by a learnable scalar (trained from a zero init, so the
+ # path is identity at step 0 and a learned delta thereafter):
+ # vec_img += cross_alpha_img * time_in_depth_to_img(emb(t_depth))
+ # vec_depth += cross_alpha_depth * time_in_rgb_to_depth(emb(t_rgb))
+ # This lets the var-AR checkpoint condition each modality on how noisy
+ # the other one is (e.g. clean RGB in i2d, clean depth in d2i).
+ self.cross_stream_timestep_mixing = cross_stream_timestep_mixing
+ if cross_stream_timestep_mixing:
+ self.time_in_depth_to_img = MLPEmbedder(
+ in_dim=256, hidden_dim=hidden_size, disable_bias=True
+ )
+ self.time_in_rgb_to_depth = MLPEmbedder(
+ in_dim=256, hidden_dim=hidden_size, disable_bias=True
+ )
+ # Shape (1,) not () to mirror the training checkpoint's parameter
+ # layout (kept 1-D for FSDP2's no-scalar-parameter rule).
+ self.cross_alpha_img = nn.Parameter(torch.zeros(1))
+ self.cross_alpha_depth = nn.Parameter(torch.zeros(1))
+
+ # Position embedder shared across streams.
+ self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
+
+ # Dual-stream stack.
+ self.double_blocks = nn.ModuleList(
+ [_TripleStreamBlock(hidden_size, num_heads, mlp_ratio)
+ for _ in range(depth_double)]
+ )
+
+ # Single-stream stack — reuses the FLUX.2 SingleStreamBlock with
+ # per-token modulation built on the fly.
+ self.single_blocks = nn.ModuleList(
+ [SingleStreamBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio)
+ for _ in range(depth_single)]
+ )
+
+ # Modulation heads.
+ self.double_stream_modulation_img = Modulation(
+ hidden_size, double=True, disable_bias=True
+ )
+ self.double_stream_modulation_txt = Modulation(
+ hidden_size, double=True, disable_bias=True
+ )
+ self.double_stream_modulation_depth = Modulation(
+ hidden_size, double=True, disable_bias=True
+ )
+ self.single_stream_modulation = Modulation(
+ hidden_size, double=False, disable_bias=True
+ )
+ self.single_stream_modulation_depth = Modulation(
+ hidden_size, double=False, disable_bias=True
+ )
+
+ # Optional depth-only decoder stack (4 SingleStreamBlocks for v1).
+ if depth_decoder_num_layers > 0:
+ self.depth_decoder_blocks = nn.ModuleList(
+ [SingleStreamBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio)
+ for _ in range(depth_decoder_num_layers)]
+ )
+ self.depth_decoder_modulation = Modulation(
+ hidden_size, double=False, disable_bias=True
+ )
+ else:
+ self.depth_decoder_blocks = nn.ModuleList()
+ self.depth_decoder_modulation = None
+
+ # Final unpatchifiers.
+ self.final_layer = LastLayer(hidden_size, in_channels)
+ self.depth_final_layer = LastLayer(hidden_size, depth_channels)
+
+ def forward(
+ self,
+ img: Tensor,
+ img_ids: Tensor,
+ timesteps: Tensor,
+ ctx: Tensor,
+ ctx_ids: Tensor,
+ depth: Tensor,
+ depth_ids: Tensor,
+ depth_timesteps: Tensor,
+ guidance: Tensor | None = None,
+ ) -> tuple[Tensor, Tensor]:
+ """Triple-stream forward.
+
+ Args:
+ img: `(B, n_img, in_channels)` RGB latent tokens
+ img_ids: `(B, n_img, 4)` 4D position ids
+ timesteps: `(B,)` RGB diffusion time
+ ctx: `(B, n_txt, context_in_dim)` text embeddings
+ ctx_ids: `(B, n_txt, 4)` text position ids
+ depth: `(B, n_depth, depth_channels)` depth latent tokens
+ depth_ids: `(B, n_depth, 4)` depth position ids
+ depth_timesteps: `(B,)` depth diffusion time
+ guidance: `(B,)` optional guidance scale
+
+ Returns:
+ `(rgb_out, depth_out)` patches.
+ """
+ # Body dtype (set by the parameters of the input projections). Inputs
+ # are cast to this so mixed-precision callers (BF16 body + FP32
+ # sampling state) work transparently.
+ body_dtype = self.img_in.weight.dtype
+ img = img.to(body_dtype)
+ ctx = ctx.to(body_dtype)
+ depth = depth.to(body_dtype)
+ timesteps = timesteps.to(body_dtype)
+ depth_timesteps = depth_timesteps.to(body_dtype)
+ if guidance is not None:
+ guidance = guidance.to(body_dtype)
+
+ # Timestep embeddings (separate per modality). `timestep_embedding`
+ # always returns FP32; cast before the BF16 MLPEmbedder weights.
+ rgb_t_emb = timestep_embedding(timesteps, 256).to(body_dtype)
+ depth_t_emb = timestep_embedding(depth_timesteps, 256).to(body_dtype)
+ vec_img = self.time_in(rgb_t_emb)
+ if self.use_guidance_embed and guidance is not None:
+ vec_img = vec_img + self.guidance_in(
+ timestep_embedding(guidance, 256).to(body_dtype)
+ )
+ vec_depth = self.time_in_depth(depth_t_emb)
+
+ # Cross-stream timestep mixing (see __init__): fold the other stream's
+ # noise level into each modulation vector. The scalar gate is part of
+ # the body dtype after `.to(dtype=…)`, so no extra cast is needed.
+ if self.cross_stream_timestep_mixing:
+ vec_img = vec_img + self.cross_alpha_img * self.time_in_depth_to_img(
+ depth_t_emb
+ )
+ vec_depth = vec_depth + self.cross_alpha_depth * self.time_in_rgb_to_depth(
+ rgb_t_emb
+ )
+
+ # Per-stream modulation coefficients.
+ mod_img_double = self.double_stream_modulation_img(vec_img)
+ mod_txt_double = self.double_stream_modulation_txt(vec_img)
+ mod_depth_double = self.double_stream_modulation_depth(vec_depth)
+ mod_img_single, _ = self.single_stream_modulation(vec_img)
+ mod_depth_single, _ = self.single_stream_modulation_depth(vec_depth)
+
+ # Project each stream into the hidden dim.
+ img = self.img_in(img)
+ txt = self.txt_in(ctx)
+ depth = self.depth_in(depth)
+
+ # Position embeddings.
+ pe_img = self.pe_embedder(img_ids)
+ pe_txt = self.pe_embedder(ctx_ids)
+ pe_depth = self.pe_embedder(depth_ids)
+
+ # Dual-stream stack.
+ for block in self.double_blocks:
+ img, txt, depth = block(
+ img, txt, depth,
+ pe_img, pe_txt, pe_depth,
+ mod_img_double, mod_txt_double, mod_depth_double,
+ )
+
+ # Single-stream stack over the concatenated [txt, img, depth] sequence
+ # with per-token modulation.
+ n_txt = txt.shape[1]
+ n_img = img.shape[1]
+ n_depth = depth.shape[1]
+ joint = torch.cat([txt, img, depth], dim=1)
+ pe_joint = torch.cat([pe_txt, pe_img, pe_depth], dim=2)
+ per_token_mod = _stack_per_token_mod(
+ mod_img_single, mod_depth_single, n_txt, n_img, n_depth
+ )
+ for block in self.single_blocks:
+ joint = block(joint, pe_joint, per_token_mod)
+
+ # Split back per stream.
+ rgb_tokens = joint[:, n_txt : n_txt + n_img]
+ depth_tokens = joint[:, n_txt + n_img :]
+
+ # Optional depth-only decoder.
+ if self.depth_decoder_num_layers > 0:
+ depth_decoder_mod, _ = self.depth_decoder_modulation(vec_depth)
+ for block in self.depth_decoder_blocks:
+ depth_tokens = block(depth_tokens, pe_depth, depth_decoder_mod)
+
+ # Each output head may be in a different dtype than the body
+ # (e.g. FP32 depth head over a BF16 body); cast inputs to the
+ # head's own dtype at the boundary.
+ rgb_head_dtype = self.final_layer.linear.weight.dtype
+ depth_head_dtype = self.depth_final_layer.linear.weight.dtype
+ rgb_out = self.final_layer(
+ rgb_tokens.to(rgb_head_dtype), vec_img.to(rgb_head_dtype)
+ )
+ depth_out = self.depth_final_layer(
+ depth_tokens.to(depth_head_dtype), vec_depth.to(depth_head_dtype)
+ )
+ return rgb_out, depth_out
diff --git a/flux_rgbd/factories.py b/flux_rgbd/factories.py
new file mode 100644
index 0000000000000000000000000000000000000000..86111ce12d7dd9fa96bcb27a7c3672796b4d9fc1
--- /dev/null
+++ b/flux_rgbd/factories.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Model variant builders. One function per published checkpoint."""
+
+from __future__ import annotations
+
+from flux_rgbd.depth.preprocess import DepthConfig
+from flux_rgbd.depth.schedule import ScheduleConfig
+from flux_rgbd.dit import FluxRGBDDiT
+from flux_rgbd.model import FluxRGBD
+
+
+def flux_rgbd_9b_v1() -> tuple[FluxRGBD, DepthConfig, ScheduleConfig]:
+ """v1 9B checkpoint: 8 dual + 24 single + 4 depth-decoder blocks."""
+ dit = FluxRGBDDiT(
+ in_channels=128, depth_channels=256, context_in_dim=12_288,
+ hidden_size=4096, num_heads=32,
+ depth_double=8, depth_single=24, depth_decoder_num_layers=4,
+ axes_dim=(32, 32, 32, 32), theta=2000, mlp_ratio=3.0,
+ )
+ return (
+ FluxRGBD(dit),
+ DepthConfig(depth_normalize_mode=("unit_mean", "contract"), patch_size=16),
+ ScheduleConfig(rgb_shift_mu=1.1, depth_shift_mu=1.1),
+ )
+
+
+def flux_rgbd_9b_v2() -> tuple[FluxRGBD, DepthConfig, ScheduleConfig]:
+ """v2 9B checkpoint: v1 architecture + cross-stream timestep mixing.
+
+ Same backbone as v1 (8 dual + 24 single + 4 depth-decoder blocks; ~12B
+ params total — "9B" names the FLUX.2 backbone class), plus a cross-stream
+ timestep-mixing path (``cross_alpha_*`` +
+ ``time_in_{rgb_to_depth,depth_to_img}``): each stream's modulation vector
+ sees a gated embedding of the *other* stream's timestep. Training-time
+ depth-visibility gating was configured as a no-op in the released
+ checkpoint and is not modelled at inference.
+ """
+ dit = FluxRGBDDiT(
+ in_channels=128, depth_channels=256, context_in_dim=12_288,
+ hidden_size=4096, num_heads=32,
+ depth_double=8, depth_single=24, depth_decoder_num_layers=4,
+ axes_dim=(32, 32, 32, 32), theta=2000, mlp_ratio=3.0,
+ cross_stream_timestep_mixing=True,
+ )
+ return (
+ FluxRGBD(dit),
+ DepthConfig(depth_normalize_mode=("unit_mean", "contract"), patch_size=16),
+ ScheduleConfig(rgb_shift_mu=1.1, depth_shift_mu=1.1),
+ )
+
+
+VARIANTS = {
+ "flux_rgbd_9b_v1": flux_rgbd_9b_v1,
+ "flux_rgbd_9b_v2": flux_rgbd_9b_v2,
+}
diff --git a/flux_rgbd/model.py b/flux_rgbd/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d764f19f8114e5c525efa15397b85c9a055144c
--- /dev/null
+++ b/flux_rgbd/model.py
@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""FluxRGBD — pairs the depth-extended DiT with a joint flow-matching sampler.
+
+Public surface:
+ * `model.forward(...)` — one denoiser call.
+ * `model.sample(...)` — N-step Euler rollout over RGB + depth. Three
+ modes: `"joint"` (text → RGBD), `"i2d"` (text+RGB → depth),
+ `"d2i"` (text+depth → RGB).
+"""
+
+from __future__ import annotations
+
+import einops
+import torch
+from torch import Tensor, nn
+
+from flux_rgbd.depth.schedule import Mode, ScheduleConfig, rollout_timesteps
+from flux_rgbd.dit import FluxRGBDDiT
+
+
+class FluxRGBD(nn.Module):
+ def __init__(self, dit: FluxRGBDDiT) -> None:
+ super().__init__()
+ self.dit = dit
+ self.in_channels = dit.in_channels
+ self.depth_channels = dit.depth_channels
+
+ # ------------------------------------------------------------------ ids
+ # FLUX.2's 4D position ids: (time_axis, row, col, seq_idx).
+ # Image tokens use (0, row, col, 0); depth tokens use (1, row, col, 0)
+ # so RoPE can distinguish them. Text tokens use (0, 0, 0, seq_idx).
+
+ @staticmethod
+ def _img_ids(batch: int, h: int, w: int, *, device, time_id: float = 0.0) -> Tensor:
+ ids = torch.zeros(h, w, 4, device=device)
+ ids[..., 0] = time_id
+ ids[..., 1] = torch.arange(h, device=device)[:, None]
+ ids[..., 2] = torch.arange(w, device=device)[None, :]
+ return einops.repeat(ids, "h w d -> b (h w) d", b=batch)
+
+ @staticmethod
+ def _txt_ids(batch: int, seq_len: int, *, device) -> Tensor:
+ ids = torch.zeros(batch, seq_len, 4, device=device)
+ ids[..., 3] = torch.arange(seq_len, device=device)
+ return ids
+
+ # ----------------------------------------------------------------- step
+
+ def forward(self, img, timesteps, ctx, img_height, img_width,
+ depth, depth_timesteps, guidance=None):
+ b, device = img.shape[0], img.device
+ return self.dit(
+ img=img,
+ img_ids=self._img_ids(b, img_height, img_width, device=device),
+ timesteps=timesteps,
+ ctx=ctx,
+ ctx_ids=self._txt_ids(b, ctx.shape[1], device=device),
+ depth=depth,
+ depth_ids=self._img_ids(b, img_height, img_width, device=device, time_id=1.0),
+ depth_timesteps=depth_timesteps,
+ guidance=guidance,
+ )
+
+ # --------------------------------------------------------------- sample
+
+ @torch.no_grad()
+ def sample(self, *, ctx: Tensor, img_height: int, img_width: int,
+ num_steps: int = 50, mode: Mode = "joint",
+ schedule_config: ScheduleConfig | None = None,
+ cfg_scale: float = 1.0, guidance: float = 1.0,
+ seed: int | None = None,
+ rgb_use_x_prediction: bool = False,
+ depth_use_x_prediction: bool = True,
+ x_prediction_t_min: float = 0.05,
+ clean_rgb: Tensor | None = None,
+ clean_depth: Tensor | None = None,
+ null_text_embed: Tensor | None = None,
+ log2_alpha: float | None = None) -> tuple[Tensor, Tensor]:
+ if mode == "i2d" and clean_rgb is None:
+ raise ValueError("clean_rgb is required for mode='i2d'")
+ if mode == "d2i" and clean_depth is None:
+ raise ValueError("clean_depth is required for mode='d2i'")
+
+ device, dtype, batch = ctx.device, ctx.dtype, ctx.shape[0]
+ num_tokens = img_height * img_width
+
+ gen = torch.Generator(device=device).manual_seed(seed) if seed is not None else None
+
+ def randn(*shape):
+ return torch.randn(*shape, device=device, dtype=dtype, generator=gen)
+
+ if mode == "i2d":
+ rgb = clean_rgb.to(dtype=dtype, device=device)
+ depth = randn(batch, num_tokens, self.depth_channels)
+ elif mode == "d2i":
+ rgb = randn(batch, num_tokens, self.in_channels)
+ depth = clean_depth.to(dtype=dtype, device=device)
+ else: # joint
+ rgb = randn(batch, num_tokens, self.in_channels)
+ depth = randn(batch, num_tokens, self.depth_channels)
+
+ schedule_config = schedule_config or ScheduleConfig()
+ t_rgb_sched, t_depth_sched = rollout_timesteps(
+ schedule_config, num_steps, mode=mode, log2_alpha=log2_alpha,
+ device=device, dtype=dtype,
+ )
+
+ guidance_tensor = (torch.full((batch,), guidance, device=device, dtype=dtype)
+ if self.dit.use_guidance_embed else None)
+
+ use_cfg = cfg_scale > 1.0
+ if use_cfg:
+ uncond = null_text_embed if null_text_embed is not None else torch.zeros_like(ctx)
+ if uncond.ndim == 2:
+ uncond = uncond.unsqueeze(0)
+ if uncond.shape[0] == 1 and batch > 1:
+ uncond = uncond.expand(batch, *uncond.shape[1:])
+ uncond = uncond.to(device=device, dtype=dtype)
+
+ for step in range(num_steps):
+ t_rgb_b = t_rgb_sched[step].expand(batch)
+ t_depth_b = t_depth_sched[step].expand(batch)
+ dt_rgb = t_rgb_sched[step + 1] - t_rgb_sched[step]
+ dt_depth = t_depth_sched[step + 1] - t_depth_sched[step]
+
+ pred_rgb, pred_depth = self.forward(
+ img=rgb, timesteps=t_rgb_b, ctx=ctx,
+ img_height=img_height, img_width=img_width,
+ depth=depth, depth_timesteps=t_depth_b, guidance=guidance_tensor,
+ )
+ if use_cfg:
+ pred_rgb_u, pred_depth_u = self.forward(
+ img=rgb, timesteps=t_rgb_b, ctx=uncond,
+ img_height=img_height, img_width=img_width,
+ depth=depth, depth_timesteps=t_depth_b, guidance=guidance_tensor,
+ )
+ pred_rgb = pred_rgb_u + cfg_scale * (pred_rgb - pred_rgb_u)
+ pred_depth = pred_depth_u + cfg_scale * (pred_depth - pred_depth_u)
+
+ # x-prediction → velocity: when the head was trained to predict
+ # the clean x rather than the velocity, convert it before Euler.
+ if rgb_use_x_prediction:
+ clamp = t_rgb_b.view(batch, 1, 1).clamp(min=x_prediction_t_min)
+ pred_rgb = (rgb - pred_rgb) / clamp
+ if depth_use_x_prediction:
+ clamp = t_depth_b.view(batch, 1, 1).clamp(min=x_prediction_t_min)
+ pred_depth = (depth - pred_depth) / clamp
+
+ # Euler step. dt=0 for the frozen modality in i2d/d2i is a no-op.
+ rgb = rgb + dt_rgb * pred_rgb
+ depth = depth + dt_depth * pred_depth
+
+ return rgb, depth
diff --git a/flux_rgbd/pointcloud.py b/flux_rgbd/pointcloud.py
new file mode 100644
index 0000000000000000000000000000000000000000..ce1f3b9a5e45d631c55bd59ce72fe387f287540e
--- /dev/null
+++ b/flux_rgbd/pointcloud.py
@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Point-cloud post-processing utilities."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def statistical_outlier_mask(points: np.ndarray, *, num_neighbors: int = 5,
+ std_ratio: float = 1.0) -> np.ndarray:
+ """Boolean inlier mask via statistical outlier rejection.
+
+ For each point, the mean distance to its ``num_neighbors`` nearest
+ neighbours is computed; a point is an outlier when that mean exceeds
+ ``global_mean + std_ratio * global_std``. This removes the flying-pixel
+ floaters that back-projected depth produces at depth discontinuities.
+ Mirrors Open3D's ``remove_statistical_outlier``.
+
+ Args:
+ points: ``(N, 3)`` float array.
+ num_neighbors: neighbours used to estimate each point's local spacing.
+ std_ratio: smaller is more aggressive (keeps fewer points).
+
+ Returns:
+ ``(N,)`` boolean mask, ``True`` for inliers.
+ """
+ n = len(points)
+ if n <= num_neighbors:
+ return np.ones(n, dtype=bool)
+ # Lazy import: only point-cloud export needs scipy.
+ from scipy.spatial import cKDTree
+ tree = cKDTree(points)
+ # k + 1: the nearest neighbour is the point itself (distance 0); drop it.
+ dists, _ = tree.query(points, k=num_neighbors + 1, workers=-1)
+ mean_dist = dists[:, 1:].mean(axis=1)
+ threshold = mean_dist.mean() + std_ratio * float(mean_dist.std())
+ return mean_dist <= threshold
+
+
+def depth_edge_mask(depth: np.ndarray, *, rtol: float = 0.04,
+ kernel_size: int = 3) -> np.ndarray:
+ """Boolean mask, ``True`` on depth-discontinuity pixels.
+
+ A pixel is an edge when the local depth range (max − min over a
+ ``kernel_size`` window) exceeds ``rtol`` of its depth. Removing these
+ pixels *before* back-projection deletes the occlusion-boundary "veils"
+ that bridge foreground→background — the connected floaters that
+ statistical outlier rejection cannot catch (they have close neighbours).
+ Mirrors MoGe / utils3d ``depth_edge`` (``rtol`` path).
+ """
+ from scipy.ndimage import maximum_filter, minimum_filter
+ local_range = (maximum_filter(depth, size=kernel_size)
+ - minimum_filter(depth, size=kernel_size))
+ return local_range > rtol * np.maximum(depth, 1e-6)
diff --git a/flux_rgbd/runner.py b/flux_rgbd/runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..a35482caf438f63cbca545be0a23c418181214cc
--- /dev/null
+++ b/flux_rgbd/runner.py
@@ -0,0 +1,307 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Public inference API: FluxRGBDRunner.
+
+Bundles the diffusion model, the FLUX.2 VAE decoder, and the Qwen3
+text encoder so callers can go from a text prompt to a numpy RGB image
+plus a depth map in one call.
+"""
+
+from __future__ import annotations
+
+import datetime as _dt
+import json
+import re
+from pathlib import Path
+
+import cv2
+import einops
+import numpy as np
+import torch
+from matplotlib import cm
+from safetensors.torch import load_file as load_safetensors
+from torch import Tensor
+
+from flux_rgbd._flux2.autoencoder import Flux2Decoder, Flux2Encoder
+from flux_rgbd.depth.preprocess import decode_depth, encode_depth
+from flux_rgbd.depth.schedule import Mode
+from flux_rgbd.factories import VARIANTS
+from flux_rgbd.text_encoder import Qwen3Embedder
+
+
+def _load_local_or_hub(repo_id_or_path: str) -> tuple[Path, Path]:
+ """Resolve `config.json` and `model.safetensors` to local paths.
+
+ Accepts a local directory or a HuggingFace repo id.
+ """
+ local = Path(repo_id_or_path)
+ if local.is_dir():
+ return local / "config.json", local / "model.safetensors"
+ from huggingface_hub import hf_hub_download
+ return (
+ Path(hf_hub_download(repo_id=repo_id_or_path, filename="config.json")),
+ Path(hf_hub_download(repo_id=repo_id_or_path, filename="model.safetensors")),
+ )
+
+
+def _materialize_meta_tensors(model: torch.nn.Module, device) -> list[str]:
+ """Give a real (uninitialized) tensor to any parameter/buffer still on the
+ meta device after an ``assign=True`` load.
+
+ A meta-device build defers allocation until weights are assigned, so any
+ tensor the checkpoint does *not* supply remains on `meta` and would crash
+ at first use. A complete checkpoint leaves none — this is a safety net for
+ drift between the model definition and the checkpoint (returns the names it
+ had to materialize so callers can surface the gap).
+ """
+ materialized: list[str] = []
+ for name, tensor in [*model.named_parameters(), *model.named_buffers()]:
+ if not tensor.is_meta:
+ continue
+ *parents, leaf = name.split(".")
+ owner = model
+ for part in parents:
+ owner = getattr(owner, part)
+ real = torch.empty(tensor.shape, dtype=tensor.dtype, device=device)
+ if isinstance(tensor, torch.nn.Parameter):
+ owner._parameters[leaf] = torch.nn.Parameter(real, requires_grad=False)
+ else:
+ owner._buffers[leaf] = real
+ materialized.append(name)
+ return materialized
+
+
+class FluxRGBDRunner:
+ """High-level inference wrapper."""
+
+ def __init__(self, model, decoder, embedder, depth_config, schedule_config,
+ *, device="cuda", img_hw=(512, 512), latent_compression=16):
+ self.model = model
+ self.decoder = decoder
+ self.embedder = embedder
+ self.depth_config = depth_config
+ self.schedule_config = schedule_config
+ self.device = torch.device(device)
+ self.img_hw = img_hw
+ self.latent_compression = latent_compression
+ self.latent_hw = (img_hw[0] // latent_compression, img_hw[1] // latent_compression)
+ self._encoder: Flux2Encoder | None = None # lazy: only needed for i2d
+ self._null_text_embed: Tensor | None = None # lazy CFG uncond, see _null_embed
+
+ @classmethod
+ def from_pretrained(cls, repo_id_or_path, *, device="cuda",
+ dtype=torch.float32, head_dtype: torch.dtype | None = None,
+ text_encoder="Qwen/Qwen3-8B",
+ img_hw=(512, 512), latent_compression=16):
+ """Load model + decoder + embedder. Path or HuggingFace Hub repo id.
+
+ ``head_dtype`` (optional) overrides the dtype of the depth output head.
+ Setting ``dtype=torch.bfloat16, head_dtype=torch.float32`` runs the
+ DiT body in BF16 (fast) but keeps the depth final layer in FP32,
+ which avoids BF16 quantization artifacts under x-prediction.
+ """
+ config_path, weights_path = _load_local_or_hub(repo_id_or_path)
+ config = json.loads(config_path.read_text())
+ builder = VARIANTS[config["variant"]]
+ # Build on the meta device, then assign the checkpoint tensors straight
+ # onto `device`. This skips ~45 s of random initialization for the 9B
+ # model (every parameter is overwritten by the checkpoint anyway) and
+ # the H2D copy of a throwaway randinit model — cutting DiT load from
+ # ~60 s to ~4 s. Bit-identical to the eager construct+load path.
+ with torch.device("meta"):
+ model, depth_cfg, schedule_cfg = builder()
+ state_dict = load_safetensors(str(weights_path), device=str(device))
+ if dtype is not None:
+ state_dict = {k: v.to(dtype) for k, v in state_dict.items()}
+ model.load_state_dict(state_dict, strict=False, assign=True)
+ leftover = _materialize_meta_tensors(model, device)
+ if leftover:
+ raise RuntimeError(
+ f"checkpoint is missing {len(leftover)} tensors (e.g. "
+ f"{leftover[:4]}); refusing to run with uninitialized weights. "
+ f"Re-download the checkpoint or check --model.")
+ model.eval()
+ if head_dtype is not None and head_dtype != dtype:
+ model.dit.depth_final_layer.to(head_dtype)
+
+ decoder = Flux2Decoder().to(device).eval()
+ decoder.load_weights()
+ embedder = Qwen3Embedder(model_spec=text_encoder, device=device).eval()
+ return cls(model, decoder, embedder, depth_cfg, schedule_cfg,
+ device=device, img_hw=img_hw,
+ latent_compression=latent_compression)
+
+ def _ensure_encoder(self) -> Flux2Encoder:
+ if self._encoder is None:
+ enc = Flux2Encoder(filename="ae_encoder.safetensors").to(self.device).eval()
+ enc.load_weights()
+ self._encoder = enc
+ return self._encoder
+
+ @torch.no_grad()
+ def encode_image(self, image_chw_uint8: np.ndarray) -> Tensor:
+ """uint8 (H, W, 3) image → latent tokens (1, num_tokens, in_channels).
+
+ Resizes to ``self.img_hw`` and normalises to [-1, 1] before encoding.
+ Output dtype matches the diffusion model's parameter dtype.
+ """
+ h, w = self.img_hw
+ rgb = cv2.resize(image_chw_uint8, (w, h), interpolation=cv2.INTER_AREA)
+ x = torch.from_numpy(rgb).to(self.device).float() / 255.0 # (H, W, 3)
+ x = x * 2 - 1 # [-1, 1]
+ x = x.unsqueeze(0) # (1, H, W, 3)
+ encoder = self._ensure_encoder()
+ latents_nhwc = encoder(x) # (1, lat_h, lat_w, in_channels)
+ b, lh, lw, c = latents_nhwc.shape
+ # Match the model body dtype so the DiT can consume directly.
+ body_dtype = next(self.model.dit.img_in.parameters()).dtype
+ return latents_nhwc.reshape(b, lh * lw, c).to(body_dtype)
+
+ @torch.no_grad()
+ def encode_depth_map(self, depth_hw: np.ndarray) -> Tensor:
+ """(H, W) depth map → depth-stream tokens (1, num_tokens, depth_channels).
+
+ For ``mode="d2i"``. Resizes to ``self.img_hw`` then patchifies via the
+ same normalisation the model was trained with (see ``encode_depth``).
+ Because that normalisation is scale-invariant (``unit_mean``), the input
+ can be metric or relative depth. Output dtype matches the depth stream.
+ """
+ h, w = self.img_hw
+ d = cv2.resize(depth_hw.astype(np.float32), (w, h), interpolation=cv2.INTER_NEAREST)
+ x = torch.from_numpy(d).to(self.device)[None, ..., None] # (1, H, W, 1)
+ tokens = encode_depth(x, self.depth_config) # (1, lat_h, lat_w, depth_channels)
+ b, lh, lw, c = tokens.shape
+ body_dtype = next(self.model.dit.depth_in.parameters()).dtype
+ return tokens.reshape(b, lh * lw, c).to(body_dtype)
+
+ @torch.no_grad()
+ def _null_embed(self) -> Tensor:
+ """Cached Qwen3 embedding of the empty string, used as the CFG uncond.
+
+ This checkpoint family was trained with the encoder's ``""`` embedding
+ as the unconditional. Letting ``model.sample`` fall back to
+ ``zeros_like(ctx)`` instead gives an out-of-distribution uncond that
+ CFG amplifies by ``(cfg_scale - 1)``, softening the RGB and corrupting
+ depth — so we encode ``""`` with the same embedder used for the prompt
+ (keeping cond/uncond in the same space) and reuse it across calls.
+ """
+ if self._null_text_embed is None:
+ null = self.embedder.forward([""]).to(self.device, torch.float32)
+ if null.ndim == 2:
+ null = null.unsqueeze(0)
+ self._null_text_embed = null
+ return self._null_text_embed
+
+ @torch.no_grad()
+ def generate(self, prompt: str, *, mode: Mode = "joint", num_steps: int = 50,
+ cfg_scale: float = 2.5, seed: int = 0,
+ clean_rgb_image: np.ndarray | None = None,
+ clean_rgb: Tensor | None = None,
+ clean_depth: Tensor | None = None,
+ refine_depth_i2d: bool = False,
+ i2d_cfg_scale: float = 1.0,
+ log2_alpha: float | None = None) -> dict:
+ """Sample one (RGB, depth) pair from `prompt`. Returns a dict with:
+
+ rgb: (H, W, 3) uint8
+ depth: (H, W) float32 depth, relative scale (positive = farther)
+ rgb_latent / depth_latent: raw model outputs (for debugging)
+ metadata: dict of the call parameters
+
+ ``refine_depth_i2d`` (joint mode only) turns the call into a two-stage
+ pipeline: stage 1 joint-samples RGB+depth at ``cfg_scale``, then stage 2
+ re-derives depth via image→depth on the stage-1 RGB *latent* at
+ ``i2d_cfg_scale``. The final RGB is stage 1's; the final depth is stage
+ 2's. Conditioning depth on a fully-formed RGB (instead of a co-evolving
+ noisy one) gives sharper, more RGB-consistent geometry.
+
+ ``log2_alpha`` (joint mode only) tilts the RGB/depth denoising trajectory
+ toward rgb-first (``> 0``) for cleaner depth; see ``rollout_timesteps``.
+ """
+ text_embed = self.embedder.forward([prompt]).to(self.device, torch.float32)
+ if text_embed.ndim == 2:
+ text_embed = text_embed.unsqueeze(0)
+
+ # Convenience: if the caller supplies a clean image instead of token
+ # latents, encode it here.
+ if clean_rgb is None and clean_rgb_image is not None:
+ clean_rgb = self.encode_image(clean_rgb_image)
+
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+
+ lh, lw = self.latent_hw
+ rgb_lat, depth_lat = self.model.sample(
+ ctx=text_embed, img_height=lh, img_width=lw,
+ num_steps=num_steps, mode=mode,
+ schedule_config=self.schedule_config,
+ cfg_scale=cfg_scale, seed=seed,
+ rgb_use_x_prediction=False, depth_use_x_prediction=True,
+ clean_rgb=clean_rgb, clean_depth=clean_depth,
+ null_text_embed=self._null_embed() if cfg_scale > 1.0 else None,
+ log2_alpha=log2_alpha,
+ )
+
+ # Stage 2 (optional): re-derive depth from the stage-1 RGB latent via
+ # i2d. cfg_scale=1.0 here means no guidance, so no null embed is needed.
+ if refine_depth_i2d and mode == "joint":
+ _, depth_lat = self.model.sample(
+ ctx=text_embed, img_height=lh, img_width=lw,
+ num_steps=num_steps, mode="i2d",
+ schedule_config=self.schedule_config,
+ cfg_scale=i2d_cfg_scale, seed=seed,
+ rgb_use_x_prediction=False, depth_use_x_prediction=True,
+ clean_rgb=rgb_lat,
+ null_text_embed=self._null_embed() if i2d_cfg_scale > 1.0 else None,
+ )
+
+ # VAE decode → uint8 RGB.
+ rgb_spatial = einops.rearrange(rgb_lat, "b (h w) c -> b h w c", h=lh, w=lw)
+ rgb_hw3 = (self.decoder(rgb_spatial.float())[0] * 0.5 + 0.5).clamp(0, 1)
+ rgb = (rgb_hw3.float().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+
+ # Contract-bijection decode → depth.
+ depth_spatial = einops.rearrange(depth_lat, "b (h w) c -> b 1 h w c", h=lh, w=lw)
+ depth = decode_depth(depth_spatial, self.depth_config) \
+ .squeeze(0).squeeze(0).squeeze(-1).float().cpu().numpy()
+
+ return {
+ "rgb": rgb, "depth": depth,
+ "rgb_latent": rgb_lat.detach(), "depth_latent": depth_lat.detach(),
+ "metadata": {"prompt": prompt, "mode": mode, "seed": seed,
+ "num_steps": num_steps, "cfg_scale": cfg_scale,
+ "img_hw": self.img_hw},
+ }
+
+ def save(self, result: dict, output_root: str | Path) -> dict[str, str]:
+ """Save a generation into a timestamped subdir of `output_root`.
+
+ Writes: rgb.png, depth_raw.npy (raw depth, relative scale), depth_magma.png
+ (disparity visualization, near = bright), metadata.json.
+ """
+ slug = re.sub(r"[^a-z0-9]+", "_", result["metadata"]["prompt"].lower()).strip("_")[:48]
+ ts = _dt.datetime.now().strftime("%Y%m%d_%H%M%S")
+ d = Path(output_root) / f"{ts}_{slug}"
+ d.mkdir(parents=True, exist_ok=True)
+
+ rgb, depth = result["rgb"], result["depth"]
+ cv2.imwrite(str(d / "rgb.png"), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
+ np.save(d / "depth_raw.npy", depth)
+
+ # Disparity (1/depth) magma, robustly normalized to the 5–95th
+ # percentile so near surfaces read bright and far ones dark.
+ valid = (depth > 0) & np.isfinite(depth)
+ if valid.any():
+ disparity = np.zeros_like(depth, dtype=np.float32)
+ disparity[valid] = 1.0 / depth[valid]
+ lo, hi = np.percentile(disparity[valid], [5, 95])
+ disparity = np.clip((disparity - lo) / max(hi - lo, 1e-8), 0, 1)
+ disparity[~valid] = 0.0
+ magma = (cm.magma(disparity)[..., :3] * 255).astype(np.uint8)
+ cv2.imwrite(str(d / "depth_magma.png"), cv2.cvtColor(magma, cv2.COLOR_RGB2BGR))
+
+ (d / "metadata.json").write_text(json.dumps(result["metadata"], indent=2))
+ return {"run_dir": str(d), "rgb": str(d / "rgb.png"),
+ "depth_raw": str(d / "depth_raw.npy"),
+ "depth_magma": str(d / "depth_magma.png"),
+ "metadata": str(d / "metadata.json")}
diff --git a/flux_rgbd/text_encoder.py b/flux_rgbd/text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..15cdd0859ff01845ffa86e1b7be382d4fbc3cf68
--- /dev/null
+++ b/flux_rgbd/text_encoder.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 World Labs.
+"""Qwen3-based text encoder for the FLUX-RGBD model family.
+
+The encoder turns text prompts into a stack of hidden states drawn from
+three intermediate layers of Qwen3-8B-FP8 (layers 9, 18, 27, mirroring
+the FLUX.2 [klein] recipe). Output shape is
+
+ (batch, MAX_LENGTH=512, 3 * hidden_size = 12288)
+
+Pads to the full ``MAX_LENGTH`` (no stripping) because the FLUX.2 DiT was
+trained against the padded text stream produced by Black Forest Labs.
+
+Loads everything from HuggingFace Hub — no internal mirrors, no auth.
+"""
+
+from __future__ import annotations
+
+import einops
+import torch
+from torch import Tensor, nn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from flux_rgbd._flux2.constants import MAX_LENGTH, OUTPUT_LAYERS_QWEN3
+
+DEFAULT_MODEL_SPEC = "Qwen/Qwen3-8B-FP8"
+
+
+class Qwen3Embedder(nn.Module):
+ """Wrap a Qwen3 causal LM, surface multi-layer hidden states.
+
+ Inference-only. Chat-template policy matches the upstream FLUX.2
+ [klein] recipe: user role, generation-prompt appended, thinking
+ disabled.
+ """
+
+ def __init__(
+ self,
+ model_spec: str = DEFAULT_MODEL_SPEC,
+ *,
+ device: str | torch.device = "cuda",
+ max_length: int = MAX_LENGTH,
+ output_layers: tuple[int, ...] = OUTPUT_LAYERS_QWEN3,
+ ) -> None:
+ super().__init__()
+ self.model_spec = model_spec
+ self.device_ = torch.device(device)
+ self.max_length = max_length
+ self.output_layers = tuple(output_layers)
+
+ self._tokenizer = AutoTokenizer.from_pretrained(model_spec)
+ # Load straight onto the target device. `dtype=None` lets
+ # transformers honor whatever the checkpoint advertises (FP8 for
+ # the Qwen3-8B-FP8 we use by default).
+ self._model = AutoModelForCausalLM.from_pretrained(
+ model_spec, dtype=None, device_map=str(self.device_)
+ )
+ self._model.eval()
+ self._pad_id = self._tokenizer.pad_token_id or 0
+
+ @staticmethod
+ def _apply_chat_template(tok, text: str) -> str:
+ return tok.apply_chat_template(
+ [{"role": "user", "content": text}],
+ tokenize=False,
+ add_generation_prompt=True,
+ enable_thinking=False,
+ )
+
+ def _tokenize(self, prompts: list[str]) -> tuple[Tensor, Tensor]:
+ formatted = [self._apply_chat_template(self._tokenizer, p) for p in prompts]
+ # First tokenize without padding to control truncation; we'll pad
+ # ourselves to `max_length` so the output sequence length is fixed
+ # regardless of the prompt batch.
+ encoded = self._tokenizer(
+ formatted,
+ padding=False,
+ truncation=True,
+ max_length=self.max_length,
+ )
+ rows = encoded["input_ids"]
+ b = len(rows)
+ input_ids = torch.full(
+ (b, self.max_length), fill_value=self._pad_id, dtype=torch.long
+ )
+ attn = torch.zeros((b, self.max_length), dtype=torch.long)
+ for i, row in enumerate(rows):
+ n = min(len(row), self.max_length)
+ input_ids[i, :n] = torch.tensor(row[:n], dtype=torch.long)
+ attn[i, :n] = 1
+ return input_ids.to(self.device_), attn.to(self.device_)
+
+ @torch.inference_mode()
+ def forward(self, prompts: list[str]) -> Tensor:
+ """Encode a batch of prompts.
+
+ Returns:
+ Tensor of shape ``(batch, max_length, len(output_layers) * hidden_size)``.
+ """
+ input_ids, attention_mask = self._tokenize(list(prompts))
+ out = self._model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ output_hidden_states=True,
+ use_cache=False,
+ )
+ # `hidden_states` is a tuple of length num_hidden_layers + 1.
+ layers = torch.stack(
+ [out.hidden_states[k] for k in self.output_layers], dim=1
+ ) # (b, len(output_layers), seq, hidden)
+ return einops.rearrange(layers, "b c l d -> b l (c d)")
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..d1b691a89a81eb7302e2261a0cc658129c0bdaf0
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,16 @@
+torch>=2.5
+einops>=0.8
+jaxtyping>=0.2
+transformers==5.10.2
+safetensors>=0.4
+huggingface_hub>=0.24
+accelerate>=1.0
+kernels>=0.4,<0.15
+numpy>=1.26,<3
+opencv-python-headless>=4.8
+pillow>=10
+matplotlib>=3.8
+scipy>=1.11
+trimesh>=4.4
+gradio>=6,<7
+spaces