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