From a1cd298a64d635af548325913ed18b30bcaf3639 Mon Sep 17 00:00:00 2001 From: hlky Date: Sat, 10 Oct 2026 21:20:15 +0100 Subject: [PATCH] TripoSplat --- docs/source/en/_toctree.yml | 4 + docs/source/en/api/models/triposplat.md | 43 ++ docs/source/en/api/pipelines/triposplat.md | 119 ++++ scripts/convert_triposplat_to_diffusers.py | 148 +++++ src/diffusers/__init__.py | 37 ++ src/diffusers/guiders/__init__.py | 1 + .../triposplat_classifier_free_guidance.py | 51 ++ src/diffusers/image_processor.py | 82 +++ src/diffusers/loaders/single_file_model.py | 12 + src/diffusers/loaders/single_file_utils.py | 27 + src/diffusers/models/__init__.py | 4 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_triposplat.py | 482 +++++++++++++++ src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_triposplat.py | 355 +++++++++++ src/diffusers/modular_pipelines/__init__.py | 2 + .../modular_pipelines/modular_pipeline.py | 4 + .../modular_pipelines/triposplat/__init__.py | 57 ++ .../triposplat/before_denoise.py | 153 +++++ .../modular_pipelines/triposplat/decoders.py | 84 +++ .../modular_pipelines/triposplat/denoise.py | 201 +++++++ .../modular_pipelines/triposplat/encoders.py | 216 +++++++ .../triposplat/modular_blocks_triposplat.py | 162 +++++ .../triposplat/modular_pipeline.py | 20 + src/diffusers/pipelines/__init__.py | 2 + .../pipelines/triposplat/__init__.py | 75 +++ .../pipelines/triposplat/modeling_birefnet.py | 556 ++++++++++++++++++ .../pipelines/triposplat/pipeline_output.py | 43 ++ .../triposplat/pipeline_triposplat.py | 376 ++++++++++++ src/diffusers/utils/__init__.py | 4 + src/diffusers/utils/dummy_pt_objects.py | 60 ++ .../dummy_torch_and_torchvision_objects.py | 17 + .../dummy_torch_and_transformers_objects.py | 60 ++ src/diffusers/utils/gaussian_splat.py | 146 +++++ src/diffusers/utils/import_utils.py | 7 + .../test_models_autoencoder_triposplat.py | 107 ++++ tests/models/test_models_birefnet.py | 95 +++ .../test_models_transformer_triposplat.py | 116 ++++ .../modular_pipelines/triposplat/__init__.py | 1 + .../test_modular_pipeline_triposplat.py | 225 +++++++ tests/others/test_image_processor.py | 48 +- tests/pipelines/triposplat/__init__.py | 1 + .../triposplat/test_pipeline_triposplat.py | 167 ++++++ tests/pipelines/triposplat/testing_utils.py | 75 +++ .../test_models_triposplat_single_file.py | 73 +++ 45 files changed, 4519 insertions(+), 1 deletion(-) create mode 100644 docs/source/en/api/models/triposplat.md create mode 100644 docs/source/en/api/pipelines/triposplat.md create mode 100644 scripts/convert_triposplat_to_diffusers.py create mode 100644 src/diffusers/guiders/triposplat_classifier_free_guidance.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_triposplat.py create mode 100644 src/diffusers/models/transformers/transformer_triposplat.py create mode 100644 src/diffusers/modular_pipelines/triposplat/__init__.py create mode 100644 src/diffusers/modular_pipelines/triposplat/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/triposplat/decoders.py create mode 100644 src/diffusers/modular_pipelines/triposplat/denoise.py create mode 100644 src/diffusers/modular_pipelines/triposplat/encoders.py create mode 100644 src/diffusers/modular_pipelines/triposplat/modular_blocks_triposplat.py create mode 100644 src/diffusers/modular_pipelines/triposplat/modular_pipeline.py create mode 100644 src/diffusers/pipelines/triposplat/__init__.py create mode 100644 src/diffusers/pipelines/triposplat/modeling_birefnet.py create mode 100644 src/diffusers/pipelines/triposplat/pipeline_output.py create mode 100644 src/diffusers/pipelines/triposplat/pipeline_triposplat.py create mode 100644 src/diffusers/utils/dummy_torch_and_torchvision_objects.py create mode 100644 src/diffusers/utils/gaussian_splat.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_triposplat.py create mode 100644 tests/models/test_models_birefnet.py create mode 100644 tests/models/transformers/test_models_transformer_triposplat.py create mode 100644 tests/modular_pipelines/triposplat/__init__.py create mode 100644 tests/modular_pipelines/triposplat/test_modular_pipeline_triposplat.py create mode 100644 tests/pipelines/triposplat/__init__.py create mode 100644 tests/pipelines/triposplat/test_pipeline_triposplat.py create mode 100644 tests/pipelines/triposplat/testing_utils.py create mode 100644 tests/single_file/test_models_triposplat_single_file.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 3abfb9052cd2..bd9c71325dc0 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -423,6 +423,8 @@ title: Transformer2DModel - local: api/models/transformer_temporal title: TransformerTemporalModel + - local: api/models/triposplat + title: TripoSplat models - local: api/models/wan_animate_2_transformer_3d title: WanAnimate2Transformer3DModel - local: api/models/wan_animate_transformer_3d @@ -687,6 +689,8 @@ title: Stable Diffusion - local: api/pipelines/stable_unclip title: Stable unCLIP + - local: api/pipelines/triposplat + title: TripoSplat - local: api/pipelines/value_guided_sampling title: Value-guided sampling - local: api/pipelines/visualcloze diff --git a/docs/source/en/api/models/triposplat.md b/docs/source/en/api/models/triposplat.md new file mode 100644 index 000000000000..b83b2de08aee --- /dev/null +++ b/docs/source/en/api/models/triposplat.md @@ -0,0 +1,43 @@ +# TripoSplat models + +[TripoSplat](https://github.com/VAST-AI-Research/TripoSplat) generates 3D Gaussian splats from an image. Its transformer +jointly denoises Gaussian and camera latents, conditioned on DINOv3 features and packed Flux2 VAE image latents. +The Gaussian decoder samples an octree and predicts splat parameters at the requested density. + +## TripoSplatTransformer3DModel + +[[autodoc]] TripoSplatTransformer3DModel + - all + - forward + +## TripoSplatGaussianDecoder + +[[autodoc]] TripoSplatGaussianDecoder + - all + - forward + +The decoder returns a tensor of shape `(batch, num_gaussians, 14)`. Columns contain xyz position (3), degree-zero +spherical harmonic color (3), positive scale (3), wxyz rotation (4), and opacity (1). Positions use the model's +coordinate system; export helpers apply the reference viewer transform. + +The decoder's adaptive octree uses data-dependent tensor sizes. Compile its repeated transformer blocks with +`decoder.compile_repeated_blocks()`; compiling the whole decoder as a static graph is not supported. + +## Original single-file checkpoints + +The transformer, Gaussian decoder, and optional `BiRefNetModel` support `from_single_file`. Supply the component's +local Diffusers config explicitly; a default converted Hub repository has not been published yet. + +```python +import torch +from diffusers import TripoSplatTransformer3DModel + +transformer = TripoSplatTransformer3DModel.from_single_file( + "./ckpts/diffusion_models/triposplat_fp16.safetensors", + config="./triposplat-diffusers/transformer", + dtype=torch.float16, +) +``` + +Use the same loading method with `TripoSplatGaussianDecoder` and its `decoder` config, or `BiRefNetModel` and its +`background_remover` config. The checkpoint key mappings are shared with the conversion script. diff --git a/docs/source/en/api/pipelines/triposplat.md b/docs/source/en/api/pipelines/triposplat.md new file mode 100644 index 000000000000..07872b19fb5d --- /dev/null +++ b/docs/source/en/api/pipelines/triposplat.md @@ -0,0 +1,119 @@ +# TripoSplat + +[TripoSplat](https://github.com/VAST-AI-Research/TripoSplat) generates 3D Gaussian splats from foreground images. +Both the standard and modular APIs use the same converted components. DINOv3 and the image VAE use bfloat16 in +the reference workflow; the transformer, Gaussian decoder, and optional BiRefNet model use float16. + +The pipeline uses `FlowMatchEulerDiscreteScheduler` with `shift=3.0` and supplies the reference sigma grid. +Predictions are cast to float32 for scheduler updates so Gaussian and camera latents remain in float32 throughout +denoising. Sigma storage and velocity updates use float32, which can produce numerical differences from the reference. + +## Convert the checkpoint + +The [official checkpoint](https://huggingface.co/VAST-AI/TripoSplat) must first be converted to Diffusers format. +Run the conversion script from a Diffusers checkout: + +```bash +python scripts/convert_triposplat_to_diffusers.py \ + --output_dir ./triposplat-diffusers \ + --include_background_remover +``` + +The default `--dtype auto` preserves the reference component dtypes. An explicit `--dtype float16`, +`--dtype bfloat16`, or `--dtype float32` applies that dtype to every component. + +Omit `--include_background_remover` to use images with a foreground alpha mask, or prepared RGB images with +`is_preprocessed=True`. BiRefNet is an optional component implemented under the TripoSplat pipeline package and +requires torchvision for deformable convolutions. It loads through the normal Diffusers model loading API. +Its Swin attention uses additive relative-position bias, so select a backend that supports additive masks, such +as native PyTorch SDPA. + +## Standard pipeline + +```python +import torch +from PIL import Image +from diffusers import TripoSplatPipeline +from diffusers.utils import export_to_gaussian_ply, export_to_splat + +dtype = {"default": torch.float16, "image_encoder": torch.bfloat16, "vae": torch.bfloat16} +pipe = TripoSplatPipeline.from_pretrained("./triposplat-diffusers", dtype=dtype).to("cuda") +output = pipe( + image=Image.open("object.png"), + num_inference_steps=20, + guidance_scale=3.0, + num_gaussians=32768, + generator=torch.Generator("cuda").manual_seed(42), + decoder_generator=torch.Generator("cuda").manual_seed(1234), +) +export_to_gaussian_ply(output.gaussians[0], "object.ply") +export_to_splat(output.gaussians[0], "object.splat") +``` + +Use `num_gaussians=[32768, 65536]` to decode several densities after one denoising run. The output contains one +batched Gaussian tensor per requested density. Counts range from 32768 to 262144 and are rounded to multiples of +32 for the released decoder. Set `output_type="latent"` to return the denoised Gaussian and camera latents. + +Tensor image inputs use `(channels, height, width)` or `(batch, channels, height, width)` layout. NumPy image +arrays use channel-last layout. Float image inputs use values between zero and one. + +`TripoSplatImageProcessor.prepare_foreground()` crops RGBA foregrounds and returns RGB PIL images on black +canvases. The pipeline predicts missing alpha masks with the optional BiRefNet component before this step. +The processor inherits `VaeImageProcessor.preprocess()` to convert the prepared RGB images to tensors in `[0, 1]`; +the DINO and VAE encoding steps apply their respective normalization. + +## Modular pipeline + +```python +import torch +from PIL import Image +from diffusers import ModularPipeline + +pipe = ModularPipeline.from_pretrained("./triposplat-diffusers") +pipe.load_components(dtype={"default": torch.float16, "image_encoder": torch.bfloat16, "vae": torch.bfloat16}) +pipe.to("cuda") +pipe.update_components(guider=pipe.guider.new(guidance_scale=3.0)) +state = pipe( + image=Image.open("object.png"), + num_inference_steps=20, + num_gaussians=32768, + generator=torch.Generator("cuda").manual_seed(42), + decoder_generator=torch.Generator("cuda").manual_seed(1234), +) +gaussians = state.get("gaussians") +``` + +The default blockset has `preprocess`, `image_encoder`, `vae_encoder`, `denoise`, and `decode` steps. +The denoise sequence separates conditioning expansion, timestep setup, noise preparation, and the denoising loop. +Each step can be removed or run independently. For example, remove `decode` from a `TripoSplatAutoBlocks` +instance before calling `init_pipeline()` to obtain latents, then reuse them with the decode block at different +densities. Configure guidance on the guider component. + +Updated `from_config` components, including `image_processor`, are not stored in the modular index by the current +framework. After reloading, reapply a customized processor with `update_components(image_processor=...)`. +The default 1024-pixel processor is unaffected. + +`TripoSplatClassifierFreeGuidance` evaluates guidance as `scale * conditional - (scale - 1) * unconditional`. +The standard guider uses an equivalent formula with a different operation order, which changes float16 rounding. +TripoSplat uses only the conditional prediction for guidance scales at or below one. + +## TripoSplatPipeline + +[[autodoc]] TripoSplatPipeline + - all + - __call__ + +## TripoSplatPipelineOutput + +[[autodoc]] TripoSplatPipelineOutput + +## TripoSplatModularPipeline + +[[autodoc]] TripoSplatModularPipeline + - all + +## BiRefNetModel + +[[autodoc]] BiRefNetModel + - all + - forward diff --git a/scripts/convert_triposplat_to_diffusers.py b/scripts/convert_triposplat_to_diffusers.py new file mode 100644 index 000000000000..1491854409d1 --- /dev/null +++ b/scripts/convert_triposplat_to_diffusers.py @@ -0,0 +1,148 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import argparse +from pathlib import Path + +import torch +from accelerate import init_empty_weights +from huggingface_hub import snapshot_download +from safetensors.torch import load_file +from transformers import DINOv3ViTConfig, DINOv3ViTModel + +from diffusers import ( + AutoencoderKLFlux2, + FlowMatchEulerDiscreteScheduler, + TripoSplatAutoBlocks, + TripoSplatPipeline, +) +from diffusers.loaders.single_file_utils import ( + convert_triposplat_birefnet_checkpoint_to_diffusers, + convert_triposplat_decoder_checkpoint_to_diffusers, + convert_triposplat_transformer_checkpoint_to_diffusers, +) +from diffusers.models.autoencoders.autoencoder_triposplat import TripoSplatGaussianDecoder +from diffusers.models.transformers.transformer_triposplat import TripoSplatTransformer3DModel +from diffusers.pipelines.triposplat.modeling_birefnet import BiRefNetModel + + +def _load_converted(model_class, state_dict, **kwargs): + with init_empty_weights(): + model = model_class(**kwargs) + model.load_state_dict(state_dict, strict=True, assign=True) + return model.eval() + + +def convert_triposplat_transformer(state_dict): + return _load_converted( + TripoSplatTransformer3DModel, convert_triposplat_transformer_checkpoint_to_diffusers(state_dict) + ) + + +def convert_triposplat_decoder(state_dict): + return _load_converted(TripoSplatGaussianDecoder, convert_triposplat_decoder_checkpoint_to_diffusers(state_dict)) + + +def convert_triposplat_birefnet(state_dict): + return _load_converted(BiRefNetModel, convert_triposplat_birefnet_checkpoint_to_diffusers(state_dict)) + + +def convert_triposplat_image_encoder(state_dict): + config = DINOv3ViTConfig( + hidden_size=1280, + intermediate_size=5120, + num_hidden_layers=32, + num_attention_heads=20, + patch_size=16, + num_register_tokens=4, + layer_norm_eps=1e-5, + query_bias=True, + key_bias=False, + value_bias=True, + use_gated_mlp=True, + hidden_act="silu", + rope_theta=100.0, + ) + with init_empty_weights(): + model = DINOv3ViTModel(config) + target_keys = model.state_dict() + converted = { + (f"model.{key}" if key.startswith("layer.") and f"model.{key}" in target_keys else key): value + for key, value in state_dict.items() + } + model.load_state_dict(converted, strict=True, assign=True) + return model.eval() + + +def convert_triposplat_image_vae(state_dict): + return _load_converted(AutoencoderKLFlux2, state_dict, batch_norm_eps=1e-5) + + +def main(): + parser = argparse.ArgumentParser(description="Convert the official TripoSplat inference checkpoints to Diffusers.") + parser.add_argument("--checkpoint_dir", type=Path) + parser.add_argument("--output_dir", type=Path, required=True) + parser.add_argument("--revision", default="56a96e603204ec410c4da60c13ea4fa09a2169a9") + parser.add_argument( + "--dtype", + choices=("auto", "float32", "float16", "bfloat16"), + default="auto", + help="auto uses the reference component dtypes; an explicit dtype applies to every component.", + ) + parser.add_argument("--include_background_remover", action="store_true") + args = parser.parse_args() + files = { + "transformer": ("diffusion_models/triposplat_fp16.safetensors", convert_triposplat_transformer), + "decoder": ("vae/triposplat_vae_decoder_fp16.safetensors", convert_triposplat_decoder), + "image_encoder": ("clip_vision/dino_v3_vit_h.safetensors", convert_triposplat_image_encoder), + "vae": ("vae/flux2-vae.safetensors", convert_triposplat_image_vae), + } + if args.include_background_remover: + files["background_remover"] = ("background_removal/birefnet.safetensors", convert_triposplat_birefnet) + checkpoint_dir = args.checkpoint_dir + if checkpoint_dir is None: + checkpoint_dir = Path( + snapshot_download( + "VAST-AI/TripoSplat", revision=args.revision, allow_patterns=[item[0] for item in files.values()] + ) + ) + components = {} + for name, (filename, convert) in files.items(): + model = convert(load_file(str(checkpoint_dir / filename))) + if args.dtype == "auto": + dtype = torch.bfloat16 if name in ("image_encoder", "vae") else torch.float16 + else: + dtype = getattr(torch, args.dtype) + model.to(dtype=dtype).save_pretrained(args.output_dir / name) + reloaded = type(model).from_pretrained(args.output_dir / name, torch_dtype=dtype) + if set(model.state_dict()) != set(reloaded.state_dict()): + raise RuntimeError(f"The {name} checkpoint did not round-trip.") + for key, value in model.state_dict().items(): + if not torch.equal(value, reloaded.state_dict()[key]): + raise RuntimeError(f"The {name} checkpoint changed {key} on reload.") + components[name] = reloaded + del model + print(f"Converted {name}", flush=True) + components.setdefault("background_remover", None) + pipeline = TripoSplatPipeline(**components, scheduler=FlowMatchEulerDiscreteScheduler(shift=3.0)) + pipeline.save_pretrained(args.output_dir) + modular = TripoSplatAutoBlocks().init_pipeline() + modular.update_components(**pipeline.components) + modular.save_pretrained(str(args.output_dir), overwrite_modular_index=True) + print(f"Saved standard and modular pipelines to {args.output_dir}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 599237a6f2a0..eb4b482bab88 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -22,6 +22,7 @@ is_torch_available, is_torchao_available, is_torchsde_available, + is_torchvision_available, is_transformers_available, is_transformers_version, ) @@ -194,6 +195,7 @@ "SkipLayerGuidance", "SmoothedEnergyGuidance", "TangentialClassifierFreeGuidance", + "TripoSplatClassifierFreeGuidance", ] ) _import_structure["hooks"].extend( @@ -222,6 +224,7 @@ "InpaintProcessor", "IPAdapterMaskProcessor", "PixArtImageProcessor", + "TripoSplatImageProcessor", "VaeImageProcessor", "VaeImageProcessorLDM3D", ] @@ -361,6 +364,8 @@ "TensorParallelConfig", "Transformer2DModel", "TransformerTemporalModel", + "TripoSplatGaussianDecoder", + "TripoSplatTransformer3DModel", "UNet1DModel", "UNet2DConditionModel", "UNet2DModel", @@ -509,6 +514,18 @@ else: _import_structure["schedulers"].extend(["CosineDPMSolverMultistepScheduler", "DPMSolverSDEScheduler"]) +try: + if not (is_torch_available() and is_torchvision_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_torch_and_torchvision_objects + + _import_structure["utils.dummy_torch_and_torchvision_objects"] = [ + name for name in dir(dummy_torch_and_torchvision_objects) if not name.startswith("_") + ] +else: + _import_structure["pipelines.triposplat.modeling_birefnet"] = ["BiRefNetModel"] + try: if not (is_torch_available() and is_transformers_available()): raise OptionalDependencyNotAvailable() @@ -578,6 +595,8 @@ "StableDiffusion3ModularPipeline", "StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline", + "TripoSplatAutoBlocks", + "TripoSplatModularPipeline", "Wan22Blocks", "Wan22Image2VideoBlocks", "Wan22Image2VideoModularPipeline", @@ -876,6 +895,8 @@ "TextToVideoSDPipeline", "TextToVideoZeroPipeline", "TextToVideoZeroSDXLPipeline", + "TripoSplatPipeline", + "TripoSplatPipelineOutput", "UnCLIPImageVariationPipeline", "UnCLIPPipeline", "UniDiffuserModel", @@ -1091,6 +1112,7 @@ SkipLayerGuidance, SmoothedEnergyGuidance, TangentialClassifierFreeGuidance, + TripoSplatClassifierFreeGuidance, ) from .hooks import ( FasterCacheConfig, @@ -1116,6 +1138,7 @@ InpaintProcessor, IPAdapterMaskProcessor, PixArtImageProcessor, + TripoSplatImageProcessor, VaeImageProcessor, VaeImageProcessorLDM3D, ) @@ -1253,6 +1276,8 @@ TensorParallelConfig, Transformer2DModel, TransformerTemporalModel, + TripoSplatGaussianDecoder, + TripoSplatTransformer3DModel, UNet1DModel, UNet2DConditionModel, UNet2DModel, @@ -1387,6 +1412,14 @@ else: from .schedulers import CosineDPMSolverMultistepScheduler, DPMSolverSDEScheduler + try: + if not (is_torch_available() and is_torchvision_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_torch_and_torchvision_objects import * # noqa F403 + else: + from .pipelines.triposplat.modeling_birefnet import BiRefNetModel + try: if not (is_torch_available() and is_transformers_available()): raise OptionalDependencyNotAvailable() @@ -1450,6 +1483,8 @@ StableDiffusion3ModularPipeline, StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline, + TripoSplatAutoBlocks, + TripoSplatModularPipeline, Wan22Blocks, Wan22Image2VideoBlocks, Wan22Image2VideoModularPipeline, @@ -1742,6 +1777,8 @@ TextToVideoSDPipeline, TextToVideoZeroPipeline, TextToVideoZeroSDXLPipeline, + TripoSplatPipeline, + TripoSplatPipelineOutput, UnCLIPImageVariationPipeline, UnCLIPPipeline, UniDiffuserModel, diff --git a/src/diffusers/guiders/__init__.py b/src/diffusers/guiders/__init__.py index fe2a07858e71..fcd3799fa479 100644 --- a/src/diffusers/guiders/__init__.py +++ b/src/diffusers/guiders/__init__.py @@ -30,3 +30,4 @@ from .skip_layer_guidance import SkipLayerGuidance from .smoothed_energy_guidance import SmoothedEnergyGuidance from .tangential_classifier_free_guidance import TangentialClassifierFreeGuidance + from .triposplat_classifier_free_guidance import TripoSplatClassifierFreeGuidance diff --git a/src/diffusers/guiders/triposplat_classifier_free_guidance.py b/src/diffusers/guiders/triposplat_classifier_free_guidance.py new file mode 100644 index 000000000000..28db2a8997a9 --- /dev/null +++ b/src/diffusers/guiders/triposplat_classifier_free_guidance.py @@ -0,0 +1,51 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import torch + +from ..configuration_utils import register_to_config +from .classifier_free_guidance import ClassifierFreeGuidance +from .guider_utils import GuiderOutput, rescale_noise_cfg + + +class TripoSplatClassifierFreeGuidance(ClassifierFreeGuidance): + """Classifier-free guidance with TripoSplat's prediction arithmetic and conditional-only scales below one.""" + + @register_to_config + def __init__( + self, + guidance_scale: float = 3.0, + guidance_rescale: float = 0.0, + use_original_formulation: bool = False, + start: float = 0.0, + stop: float = 1.0, + enabled: bool = True, + ): + super().__init__(guidance_scale, guidance_rescale, use_original_formulation, start, stop, enabled) + + def _is_cfg_enabled(self) -> bool: + if not self.use_original_formulation and self.guidance_scale <= 1: + return False + return super()._is_cfg_enabled() + + def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput: + if not self._is_cfg_enabled(): + pred = pred_cond + else: + scale = self.guidance_scale + int(self.use_original_formulation) + pred = scale * pred_cond - (scale - 1) * pred_uncond + if self.guidance_rescale > 0: + pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale) + return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond) diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 6dccc02dd5ef..17a22506ab7f 100644 --- a/src/diffusers/image_processor.py +++ b/src/diffusers/image_processor.py @@ -1473,3 +1473,85 @@ def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: in samples = samples[:, :, start_y:end_y, start_x:end_x] return samples + + +class TripoSplatImageProcessor(VaeImageProcessor): + """Prepare foreground images and convert prepared RGB images to tensors in `[0, 1]`. + + Args: + canvas_size (`int`, defaults to `1024`): + Side length of the prepared RGB images. + """ + + @register_to_config + def __init__(self, canvas_size: int = 1024) -> None: + super().__init__(do_resize=False, do_normalize=False, vae_scale_factor=16) + + def to_pil(self, image: PipelineImageInput) -> list[Image.Image]: + """Convert an image or batch to a list of PIL images.""" + if isinstance(image, (np.ndarray, torch.Tensor)) and image.ndim == 4: + image = list(image) + images = image if isinstance(image, list) else [image] + converted = [] + for item in images: + if isinstance(item, Image.Image): + item = item.copy() + else: + if isinstance(item, torch.Tensor): + if item.ndim != 3: + raise ValueError("An individual tensor image must have three dimensions.") + dtype = item.dtype + item = self.pt_to_numpy(item.detach().unsqueeze(0))[0] + if dtype == torch.uint8: + item = item.astype(np.uint8) + if not isinstance(item, np.ndarray): + raise TypeError(f"Unsupported image type: {type(item)}") + if item.dtype != np.uint8: + item = (item.clip(0, 1) * 255).astype(np.uint8) + item = Image.fromarray(item) + converted.append(item) + return converted + + def resize_shortest_edge(self, images: list[Image.Image]) -> list[Image.Image]: + """Resize PIL images so their shorter side matches the canvas size.""" + resized = [] + for image in images: + scale = self.config.canvas_size / min(image.size) + resized.append(self.resize(image, height=round(image.height * scale), width=round(image.width * scale))) + return resized + + def prepare_foreground( + self, image: PipelineImageInput, erode_radius: int = 1, is_preprocessed: bool = False + ) -> list[Image.Image]: + """Crop RGBA foregrounds or resize prepared RGB images onto black square canvases.""" + if not isinstance(erode_radius, int) or erode_radius < 0: + raise ValueError("erode_radius must be a nonnegative integer.") + images = self.to_pil(image) + if not is_preprocessed: + images = self.resize_shortest_edge(images) + prepared = [] + for item in images: + size = self.config.canvas_size + if is_preprocessed: + item = self.resize(item.convert("RGB"), height=size, width=size) + else: + if item.mode != "RGBA": + raise ValueError( + "Foreground images require an alpha mask; prepare RGB images through the pipeline." + ) + if erode_radius: + item.putalpha(item.getchannel("A").filter(ImageFilter.MinFilter(2 * erode_radius + 1))) + ys, xs = np.nonzero(np.asarray(item.getchannel("A"))) + if len(xs) == 0: + raise ValueError("The foreground mask is empty.") + cx, cy = (xs.min() + xs.max()) / 2, (ys.min() + ys.max()) / 2 + half = max(xs.max() - xs.min(), ys.max() - ys.min()) / 2 * 1.2 + if half == 0: + raise ValueError("The foreground mask must span more than one pixel.") + item = item.crop([int(cx - half), int(cy - half), int(cx + half), int(cy + half)]) + item = self.resize(item, height=size, width=size) + background = Image.new("RGB", (size, size), (0, 0, 0)) + background.paste(item, mask=item.getchannel("A")) + item = background + prepared.append(item) + return prepared diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 1556227673f4..66f128825031 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -57,6 +57,9 @@ convert_sana_transformer_to_diffusers, convert_sd3_transformer_checkpoint_to_diffusers, convert_stable_cascade_unet_single_file_to_diffusers, + convert_triposplat_birefnet_checkpoint_to_diffusers, + convert_triposplat_decoder_checkpoint_to_diffusers, + convert_triposplat_transformer_checkpoint_to_diffusers, convert_wan_animate_2_transformer_to_diffusers, convert_wan_transformer_to_diffusers, convert_wan_vae_to_diffusers, @@ -85,6 +88,15 @@ _LOW_CPU_MEM_USAGE_DEFAULT = False SINGLE_FILE_LOADABLE_CLASSES = { + "TripoSplatTransformer3DModel": { + "checkpoint_mapping_fn": convert_triposplat_transformer_checkpoint_to_diffusers, + }, + "TripoSplatGaussianDecoder": { + "checkpoint_mapping_fn": convert_triposplat_decoder_checkpoint_to_diffusers, + }, + "BiRefNetModel": { + "checkpoint_mapping_fn": convert_triposplat_birefnet_checkpoint_to_diffusers, + }, "StableCascadeUNet": { "checkpoint_mapping_fn": convert_stable_cascade_unet_single_file_to_diffusers, }, diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index 3e160be02b4b..737813a95bdb 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -2795,6 +2795,33 @@ def convert_mochi_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): return converted_state_dict +def convert_triposplat_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): + converted = {} + for key, value in checkpoint.items(): + key = key.replace(".mlp.mlp.0.", ".mlp.net.0.proj.").replace(".mlp.mlp.2.", ".mlp.net.2.") + key = key.replace("cam_refiner.mlp.", "cam_refiner.") + converted[key.replace(".attn.qkv.", ".attn.to_qkv.").replace(".attn.out.", ".attn.to_out.")] = value + return converted + + +def convert_triposplat_decoder_checkpoint_to_diffusers(checkpoint, **kwargs): + converted = {} + for key, value in checkpoint.items(): + key = key.replace(".q_rms_norm.", ".q_norm.").replace(".k_rms_norm.", ".k_norm.") + key = key.replace(".mlp.mlp.0.", ".mlp.net.0.proj.").replace(".mlp.mlp.2.", ".mlp.net.2.") + converted[key] = value + return converted + + +def convert_triposplat_birefnet_checkpoint_to_diffusers(checkpoint, **kwargs): + training_heads = ("decoder.conv_ms_spvn_", "decoder.gdt_convs_pred_") + return { + key.replace(".regular_conv.", "."): value + for key, value in checkpoint.items() + if not key.startswith(training_heads) + } + + def convert_hunyuan_video_transformer_to_diffusers(checkpoint, **kwargs): def remap_norm_scale_shift_(key, state_dict): weight = state_dict.pop(key) diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 476b80f94d56..44bac5fb88ab 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -60,6 +60,7 @@ _import_structure["autoencoders.autoencoder_rae"] = ["AutoencoderRAE"] _import_structure["autoencoders.autoencoder_same"] = ["AutoencoderSAME"] _import_structure["autoencoders.autoencoder_tiny"] = ["AutoencoderTiny"] + _import_structure["autoencoders.autoencoder_triposplat"] = ["TripoSplatGaussianDecoder"] _import_structure["autoencoders.autoencoder_vidtok"] = ["AutoencoderVidTok"] _import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"] _import_structure["autoencoders.ltx2_diffusion_decoder"] = ["LTX2VideoDiffusionDecoderModel"] @@ -157,6 +158,7 @@ _import_structure["transformers.transformer_skyreels_v2"] = ["SkyReelsV2Transformer3DModel"] _import_structure["transformers.transformer_stable_audio3"] = ["StableAudio3DiTModel"] _import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"] + _import_structure["transformers.transformer_triposplat"] = ["TripoSplatTransformer3DModel"] _import_structure["transformers.transformer_wan"] = ["WanTransformer3DModel"] _import_structure["transformers.transformer_wan_animate"] = ["WanAnimateTransformer3DModel"] _import_structure["transformers.transformer_wan_animate_2"] = ["WanAnimate2Transformer3DModel"] @@ -219,6 +221,7 @@ MiniMaxMusic3Vocoder, MMAudioVAE, MMAudioVocoder, + TripoSplatGaussianDecoder, VQModel, ) from .cache_utils import CacheMixin @@ -319,6 +322,7 @@ WanVACETransformer3DModel, ZImageTransformer2DModel, ) + from .transformers.transformer_triposplat import TripoSplatTransformer3DModel from .unets import ( DreamLiteUNetModel, I2VGenXLUNet, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index bb4abb06928b..39a78aa4f793 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -30,6 +30,7 @@ from .autoencoder_rae import AutoencoderRAE from .autoencoder_same import AutoencoderSAME from .autoencoder_tiny import AutoencoderTiny +from .autoencoder_triposplat import TripoSplatGaussianDecoder from .autoencoder_vidtok import AutoencoderVidTok from .consistency_decoder_vae import ConsistencyDecoderVAE from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel diff --git a/src/diffusers/models/autoencoders/autoencoder_triposplat.py b/src/diffusers/models/autoencoders/autoencoder_triposplat.py new file mode 100644 index 000000000000..33d5af59fe6f --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_triposplat.py @@ -0,0 +1,482 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import math +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin +from ...utils import BaseOutput +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin +from ..normalization import FP32LayerNorm + + +@dataclass +class TripoSplatGaussianDecoderOutput(BaseOutput): + """Gaussian parameters shaped `(batch, num_gaussians, 14)`. + + Args: + sample (`torch.Tensor`): + Columns contain xyz position, degree-zero SH color, positive scale, wxyz rotation, and opacity. + """ + + sample: torch.Tensor + + +class TripoSplatDecoderAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "TripoSplatDecoderAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + ) -> torch.Tensor: + if encoder_hidden_states is None: + query, key, value = attn.to_qkv(hidden_states).unflatten(-1, (3, attn.num_heads, attn.head_dim)).unbind(2) + else: + query = attn.to_q(hidden_states).unflatten(-1, (attn.num_heads, attn.head_dim)) + key, value = attn.to_kv(encoder_hidden_states).unflatten(-1, (2, attn.num_heads, attn.head_dim)).unbind(2) + query = attn.q_norm(query) + key = attn.k_norm(key) + hidden_states = dispatch_attention_fn( + query, key, value, backend=self._attention_backend, parallel_config=self._parallel_config + ) + return attn.to_out(hidden_states.flatten(-2)) + + +class TripoSplatDecoderAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = TripoSplatDecoderAttnProcessor + _available_processors = [TripoSplatDecoderAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, channels: int, num_heads: int, ctx_channels: int | None = None) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = channels // num_heads + if ctx_channels is None: + self.to_qkv = nn.Linear(channels, 3 * channels, bias=True) + else: + self.to_q = nn.Linear(channels, channels, bias=True) + self.to_kv = nn.Linear(ctx_channels, 2 * channels, bias=True) + self.q_norm = TripoSplatDecoderMultiHeadRMSNorm(self.head_dim, num_heads) + self.k_norm = TripoSplatDecoderMultiHeadRMSNorm(self.head_dim, num_heads) + self.to_out = nn.Linear(channels, channels) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None) -> torch.Tensor: + return self.processor(self, hidden_states, encoder_hidden_states) + + +class TripoSplatDecoderMultiHeadRMSNorm(nn.Module): + def __init__(self, dim: int, heads: int) -> None: + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(heads, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + origin_dtype = x.dtype + return (F.normalize(x.float(), dim=-1) * self.gamma.float() * self.scale).to(origin_dtype) + + +class TripoSplatDecoderPcdAbsolutePositionEmbedderV2(nn.Module): + def __init__(self, channels: int, in_channels: int = 3, max_res: int = 10) -> None: + super().__init__() + self.channels = channels + self.in_channels = in_channels + self.max_res = max_res + self.freq_dim = channels // in_channels // 2 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + orig_dtype = x.dtype + x = x.float() + N, D = x.shape + logs = torch.linspace(0.0, float(self.max_res), steps=self.freq_dim, device=x.device, dtype=x.dtype) + frequencies = torch.pow(2.0, logs) + ang = x.unsqueeze(-1) * frequencies * torch.pi + embed = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1).reshape(N, -1) + if embed.shape[1] < self.channels: + embed = torch.cat( + [embed, torch.zeros(N, self.channels - embed.shape[1], device=embed.device, dtype=embed.dtype)], dim=-1 + ) + return embed.to(orig_dtype) + + +class TripoSplatDecoderLevelEmbedder(nn.Module): + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(256, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequencies = torch.exp(-math.log(1024) * torch.arange(128, dtype=torch.float32, device="cpu") / 128) + + def forward(self, levels: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + angles = levels[:, None].float() * self.frequencies.to(levels.device)[None] * 2 * torch.pi + embeddings = torch.cat([torch.cos(angles), torch.sin(angles)], dim=-1) + return self.mlp(embeddings.to(dtype)) + + +class TripoSplatDecoderModulatedTransformerCrossOnlyBlock(nn.Module): + def __init__(self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0) -> None: + super().__init__() + self.norm1 = FP32LayerNorm(channels, elementwise_affine=False, eps=1e-06) + self.norm2 = FP32LayerNorm(channels, elementwise_affine=False, eps=1e-06) + self.cross_attn = TripoSplatDecoderAttention(channels, ctx_channels=ctx_channels, num_heads=num_heads) + self.mlp = FeedForward(channels, mult=mlp_ratio, activation_fn="gelu-approximate") + + def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1) + h = self.norm1(x) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + x = x + self.cross_attn(h, context) * gate_msa.unsqueeze(1) + h = self.norm2(x) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + x = x + self.mlp(h) * gate_mlp.unsqueeze(1) + return x + + +class TripoSplatDecoderOctreeProbabilityFixedlenDecoder(nn.Module): + def __init__( + self, + model_channels: int, + cond_channels: int, + num_blocks: int, + num_heads: int, + mlp_ratio: float, + max_voxel_level: int, + ) -> None: + super().__init__() + self.model_channels = model_channels + self.max_voxel_level = max_voxel_level + self.input_layer = nn.Linear(model_channels, model_channels) + self.l_embedder = TripoSplatDecoderLevelEmbedder(model_channels) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True)) + self.blocks = nn.ModuleList( + [ + TripoSplatDecoderModulatedTransformerCrossOnlyBlock( + model_channels, ctx_channels=cond_channels, num_heads=num_heads, mlp_ratio=mlp_ratio + ) + for _ in range(num_blocks) + ] + ) + self.out_proj = nn.Linear(model_channels, 8) + self.in_proj = nn.Linear(3, model_channels) + self.pos_embedder = TripoSplatDecoderPcdAbsolutePositionEmbedderV2(channels=model_channels, in_channels=3) + + def forward( + self, cond: torch.Tensor, num_points: int, generator: torch.Generator | None = None + ) -> dict[str, torch.Tensor]: + B = cond.shape[0] + device = cond.device + child_offset = torch.tensor( + [[i, j, k] for k in [0, 1] for j in [0, 1] for i in [0, 1]], dtype=torch.long, device=device + ) + prev_coords_int = torch.zeros(B, 1, 3, dtype=torch.long, device=device) + prev_counts = torch.full((B, 1), num_points, dtype=torch.long, device=device) + batch_indices_range = torch.arange(B, device=device).unsqueeze(1) + for lv in range(1, self.max_voxel_level + 1): + res_p = 1 << lv - 1 + res = 1 << lv + parent_coords_norm = (prev_coords_int.float() + 0.5) / res_p + res_tensor = torch.full((B,), res, dtype=torch.long, device=device) + hidden_states = self.in_proj(parent_coords_norm.to(cond.dtype)) + position_embeds = self.pos_embedder(parent_coords_norm.reshape(-1, 3)).reshape(B, -1, self.model_channels) + hidden_states = hidden_states + position_embeds.to(device=hidden_states.device, dtype=cond.dtype) + hidden_states = self.input_layer(hidden_states) + level_embeds = self.l_embedder(res_tensor, hidden_states.dtype) + modulation = self.adaLN_modulation(level_embeds) + for block in self.blocks: + hidden_states = block(hidden_states, modulation, cond) + hidden_states = F.layer_norm(hidden_states.float(), hidden_states.shape[-1:]).to(cond.dtype) + pred_logits = self.out_proj(hidden_states) + pred_probs = torch.softmax(pred_logits, dim=-1).to(device) + counts = prev_counts.flatten().to(device=device, dtype=torch.long) + probabilities = pred_probs.reshape(counts.numel(), -1).float().clamp_min_(0) + probabilities = probabilities / probabilities.sum(1, keepdim=True).clamp_min_(1) + sampled = torch.zeros_like(probabilities, dtype=torch.long) + cdf = probabilities.cumsum(dim=1).clamp(max=1.0 - 1e-12) + unique_counts, inverse_indices = counts.unique(sorted=False, return_inverse=True) + random_device = device if generator is None else generator.device + for index, count in enumerate(unique_counts.tolist()): + if count == 0: + continue + rows = (inverse_indices == index).nonzero(as_tuple=False).squeeze(1) + start = torch.rand( + (rows.numel(), 1), generator=generator, device=random_device, dtype=probabilities.dtype + ).to(device) / float(count) + grid = torch.arange(count, device=device, dtype=probabilities.dtype)[None, :] / float(count) + samples = (start + grid).clamp(max=1.0 - 1e-12) + child_indices = torch.searchsorted(cdf.index_select(0, rows), samples).clamp_max( + probabilities.shape[1] - 1 + ) + child_counts = torch.zeros( + rows.numel(), probabilities.shape[1], dtype=probabilities.dtype, device=device + ) + child_counts.scatter_add_(1, child_indices, torch.ones_like(child_indices, dtype=child_counts.dtype)) + sampled.index_copy_(0, rows, child_counts.to(torch.long)) + sampled = sampled.reshape(*prev_counts.shape, -1).flatten(1, 2) + child_coords_int = (prev_coords_int[:, :, None, :] * 2 + child_offset[None, None, :, :]).flatten(1, 2) + mask = sampled > 0 + max_valid = mask.sum(dim=1).max().item() + scatter_indices = mask.cumsum(dim=1) - 1 + valid_scatter_indices = scatter_indices[mask] + valid_batch_indices = batch_indices_range.expand_as(mask)[mask] + next_prev_coords_int = torch.zeros(B, max_valid, 3, dtype=child_coords_int.dtype, device=device) + next_prev_coords_int[valid_batch_indices, valid_scatter_indices] = child_coords_int[mask] + next_prev_counts = torch.zeros(B, max_valid, dtype=sampled.dtype, device=device) + next_prev_counts[valid_batch_indices, valid_scatter_indices] = sampled[mask] + prev_coords_int = next_prev_coords_int + prev_counts = next_prev_counts + res = 1 << self.max_voxel_level + coords_int = torch.repeat_interleave(prev_coords_int.flatten(0, 1), prev_counts.flatten(0, 1), dim=0).reshape( + B, num_points, -1 + ) + coords_norm = coords_int.float() + random_device = device if generator is None else generator.device + jitter = torch.rand(coords_norm.shape, generator=generator, device=random_device, dtype=coords_norm.dtype).to( + device + ) + coords_norm = (coords_norm + jitter) / res + return {"points": coords_norm} + + +class TripoSplatDecoderTransformerCrossBlock(nn.Module): + def __init__(self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0) -> None: + super().__init__() + self.norm1 = FP32LayerNorm(channels, elementwise_affine=False, eps=1e-06) + self.norm2 = FP32LayerNorm(channels, elementwise_affine=True, eps=1e-06) + self.norm3 = FP32LayerNorm(channels, elementwise_affine=False, eps=1e-06) + self.self_attn = TripoSplatDecoderAttention(channels, num_heads=num_heads) + self.cross_attn = TripoSplatDecoderAttention(channels, ctx_channels=ctx_channels, num_heads=num_heads) + self.mlp = FeedForward(channels, mult=mlp_ratio, activation_fn="gelu-approximate") + + def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor: + x = x + self.self_attn(self.norm1(x)) + x = x + self.cross_attn(self.norm2(x), context) + x = x + self.mlp(self.norm3(x)) + return x + + +class TripoSplatDecoderElasticGaussianFixedlenDecoder(nn.Module): + def __init__( + self, + model_channels: int, + cond_channels: int, + num_blocks: int, + num_heads: int, + mlp_ratio: float, + gaussians_per_point: int, + ) -> None: + super().__init__() + self.model_channels = model_channels + self.layout = { + "_xyz": {"shape": (gaussians_per_point, 3), "size": gaussians_per_point * 3}, + "_features_dc": {"shape": (gaussians_per_point, 1, 3), "size": gaussians_per_point * 3}, + "_scaling": {"shape": (gaussians_per_point, 3), "size": gaussians_per_point * 3}, + "_rotation": {"shape": (gaussians_per_point, 4), "size": gaussians_per_point * 4}, + "_opacity": {"shape": (gaussians_per_point, 1), "size": gaussians_per_point}, + "_offset_scale": {"shape": (gaussians_per_point, 1), "size": gaussians_per_point}, + } + out_channels = 0 + for layout in self.layout.values(): + layout["range"] = (out_channels, out_channels + layout["size"]) + out_channels += layout["size"] + self.input_layer = nn.Linear(model_channels, model_channels) + self.blocks = nn.ModuleList( + [ + TripoSplatDecoderTransformerCrossBlock( + model_channels, ctx_channels=cond_channels, num_heads=num_heads, mlp_ratio=mlp_ratio + ) + for _ in range(num_blocks) + ] + ) + self.in_proj = nn.Linear(3, model_channels) + self.pos_embedder = TripoSplatDecoderPcdAbsolutePositionEmbedderV2(channels=model_channels, in_channels=3) + self.out_proj = nn.Linear(model_channels, out_channels) + perturbation = [] + for index in range(gaussians_per_point): + point = [index / gaussians_per_point] + for base in (2, 3): + value = 0.0 + inverse_base = 1.0 / base + factor = inverse_base + remaining = index + while remaining > 0: + value += (remaining % base) * factor + remaining //= base + factor *= inverse_base + point.append(value) + perturbation.append(point) + perturbation = torch.tensor(perturbation).float() + perturbation = torch.atanh((perturbation * 2 - 1) / 1.5) + self.register_buffer("points_offset_perturbation", perturbation) + base = torch.tensor(0.05) + self.register_buffer("base_offset_scale", torch.log(torch.exp(base) - 1.0)) + + def forward(self, points: torch.Tensor, cond: torch.Tensor) -> dict[str, torch.Tensor]: + batch, length, _ = points.shape + hidden_states = self.in_proj(points.to(cond.dtype)) + position_embeds = self.pos_embedder(points.reshape(-1, 3)).reshape(batch, length, -1) + hidden_states = hidden_states + position_embeds.to(device=hidden_states.device, dtype=cond.dtype) + hidden_states = self.input_layer(hidden_states) + for block in self.blocks: + hidden_states = block(hidden_states, cond) + hidden_states = F.layer_norm(hidden_states.float(), hidden_states.shape[-1:]).to(cond.dtype) + features = self.out_proj(hidden_states) + start, end = self.layout["_offset_scale"]["range"] + offset_scale = F.softplus( + features[:, :, start:end].reshape(batch, -1, *self.layout["_offset_scale"]["shape"]) + + self.base_offset_scale.to(features.device) + ) + start, end = self.layout["_xyz"]["range"] + offset = features[:, :, start:end].reshape(batch, -1, *self.layout["_xyz"]["shape"]) + offset = offset + self.points_offset_perturbation.to(features.device) + offset = torch.tanh(offset) * 0.5 * 1.5 + offset = offset * offset_scale + return {"features": features, "offset": offset} + + +class TripoSplatGaussianDecoder(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): + """Decode Gaussian latents by sampling an octree and predicting splat parameters. + + Parameters: + model_channels (`int`, defaults to `1024`): + Decoder hidden width. + cond_channels (`int`, defaults to `16`): + Width of the denoised Gaussian latents. + num_octree_blocks (`int`, defaults to `4`): + Number of octree probability transformer blocks. + num_gaussian_blocks (`int`, defaults to `16`): + Number of Gaussian parameter transformer blocks. + num_heads (`int`, defaults to `16`): + Number of attention heads. + mlp_ratio (`float`, defaults to `4.0`): + Feed-forward hidden width relative to `model_channels`. + gaussians_per_point (`int`, defaults to `32`): + Number of Gaussians predicted for each sampled point. + max_voxel_level (`int`, defaults to `8`): + Number of octree subdivision levels. + """ + + _skip_layerwise_casting_patterns = ["norm", "input_layer", "in_proj", "l_embedder"] + _no_split_modules = [ + "TripoSplatDecoderModulatedTransformerCrossOnlyBlock", + "TripoSplatDecoderTransformerCrossBlock", + ] + _repeated_blocks = [ + "TripoSplatDecoderModulatedTransformerCrossOnlyBlock", + "TripoSplatDecoderTransformerCrossBlock", + ] + + @register_to_config + def __init__( + self, + model_channels: int = 1024, + cond_channels: int = 16, + num_octree_blocks: int = 4, + num_gaussian_blocks: int = 16, + num_heads: int = 16, + mlp_ratio: float = 4.0, + gaussians_per_point: int = 32, + max_voxel_level: int = 8, + ) -> None: + super().__init__() + self.octree = TripoSplatDecoderOctreeProbabilityFixedlenDecoder( + model_channels=model_channels, + cond_channels=cond_channels, + num_blocks=num_octree_blocks, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + max_voxel_level=max_voxel_level, + ) + self.gs = TripoSplatDecoderElasticGaussianFixedlenDecoder( + model_channels=model_channels, + cond_channels=cond_channels, + num_blocks=num_gaussian_blocks, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + gaussians_per_point=gaussians_per_point, + ) + scaling_bias = torch.tensor(0.004, dtype=torch.float32, device="cpu") + self.scaling_bias = scaling_bias + torch.log(-torch.expm1(-scaling_bias)) + opacity_bias = torch.tensor(0.1, dtype=torch.float32, device="cpu") + self.opacity_bias = torch.log(opacity_bias / (1 - opacity_bias)) + self.identity_rotation = torch.tensor([1, 0, 0, 0], dtype=torch.float32, device="cpu") + + def forward( + self, + hidden_states: torch.Tensor, + num_gaussians: int = 262144, + generator: torch.Generator | list[torch.Generator] | None = None, + return_dict: bool = True, + ) -> TripoSplatGaussianDecoderOutput | tuple[torch.Tensor]: + """ + Args: + hidden_states (`torch.Tensor`): + Denoised Gaussian latents of shape `(batch, latent_tokens, cond_channels)`. + num_gaussians (`int`, defaults to `262144`): + Number of Gaussians to decode, divisible by `gaussians_per_point`. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Generator or one generator per batch item for octree sampling and point jitter. + return_dict (`bool`, defaults to `True`): + Whether to return a structured output or a tuple. + + Returns: + `TripoSplatGaussianDecoderOutput` or `tuple`: Parameters of shape `(batch, num_gaussians, 14)` in xyz, + degree-zero spherical harmonic color, scale, wxyz rotation, and opacity order. + """ + if num_gaussians <= 0 or num_gaussians % self.config.gaussians_per_point: + raise ValueError("num_gaussians must be a positive multiple of gaussians_per_point.") + hidden_states = hidden_states.to(self.dtype) + batch_size = hidden_states.shape[0] + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError("Pass one generator per batch item.") + points = [] + for index in range(batch_size): + item_generator = generator[index] if isinstance(generator, list) else generator + prediction = self.octree( + hidden_states[index : index + 1], + num_points=num_gaussians // self.config.gaussians_per_point, + generator=item_generator, + ) + points.append(prediction["points"]) + points = torch.cat(points, dim=0) + prediction = self.gs(points, hidden_states) + features = prediction["features"] + positions = (points.to(features.device)[:, :, None] + prediction["offset"]).flatten(1, 2) - 0.5 + parameters = {} + for name in ("_features_dc", "_scaling", "_rotation", "_opacity"): + layout = self.gs.layout[name] + start, end = layout["range"] + parameters[name] = features[..., start:end].reshape(batch_size, -1, *layout["shape"]).flatten(1, 2) + colors = parameters["_features_dc"].squeeze(2) + scaling_bias = self.scaling_bias.to(features.device) + scales = F.softplus(parameters["_scaling"] + scaling_bias) + scales = (scales.square() + 0.0009**2).sqrt() + rotations = parameters["_rotation"] * 0.1 + rotations = rotations + self.identity_rotation.to(features.device) + opacity_bias = self.opacity_bias.to(features.device) + opacity = torch.sigmoid(parameters["_opacity"] + opacity_bias) + sample = torch.cat([positions, colors, scales, rotations, opacity], dim=-1) + if not return_dict: + return (sample,) + return TripoSplatGaussianDecoderOutput(sample=sample) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index d29d49d16857..346b6e1a50aa 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -68,6 +68,7 @@ from .transformer_skyreels_v2 import SkyReelsV2Transformer3DModel from .transformer_stable_audio3 import StableAudio3DiTModel from .transformer_temporal import TransformerTemporalModel + from .transformer_triposplat import TripoSplatTransformer3DModel from .transformer_wan import WanTransformer3DModel from .transformer_wan_animate import WanAnimateTransformer3DModel from .transformer_wan_animate_2 import WanAnimate2Transformer3DModel diff --git a/src/diffusers/models/transformers/transformer_triposplat.py b/src/diffusers/models/transformers/transformer_triposplat.py new file mode 100644 index 000000000000..d555c92b6c85 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_triposplat.py @@ -0,0 +1,355 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import math +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin +from ...utils import BaseOutput +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin +from ..normalization import FP32LayerNorm + + +@dataclass +class TripoSplatTransformer3DModelOutput(BaseOutput): + """Joint Gaussian and camera velocity predictions. + + Args: + sample (`torch.Tensor`): + Gaussian velocity shaped `(batch, query_tokens, in_channels)`. + camera_sample (`torch.Tensor`, *optional*): + Camera velocity shaped `(batch, 1, cam_channels)`. + """ + + sample: torch.Tensor + camera_sample: torch.Tensor | None = None + + +class TripoSplatRMSNorm(nn.Module): + def __init__(self, head_dim: int, num_heads: int) -> None: + super().__init__() + self.scale = head_dim**0.5 + self.gamma = nn.Parameter(torch.ones(num_heads, head_dim)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return (F.normalize(hidden_states.float(), dim=-1) * self.gamma.float() * self.scale).to(hidden_states.dtype) + + +class TripoSplatAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, attn: "TripoSplatAttention", hidden_states: torch.Tensor, rotary_emb: torch.Tensor + ) -> torch.Tensor: + query, key, value = attn.to_qkv(hidden_states).unflatten(-1, (3, attn.num_heads, attn.head_dim)).unbind(2) + # TripoSplat predicts separate complex rotation phases for each token and attention head. + query = torch.view_as_complex(query.float().unflatten(-1, (-1, 2))) * rotary_emb + key = torch.view_as_complex(key.float().unflatten(-1, (-1, 2))) * rotary_emb + query = torch.view_as_real(query).flatten(-2).to(hidden_states.dtype) + key = torch.view_as_real(key).flatten(-2).to(hidden_states.dtype) + query = attn.q_norm(query) + key = attn.k_norm(key) + hidden_states = dispatch_attention_fn( + query, key, value, backend=self._attention_backend, parallel_config=self._parallel_config + ) + return attn.to_out(hidden_states.flatten(-2)) + + +class TripoSplatAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = TripoSplatAttnProcessor + _available_processors = [TripoSplatAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, channels: int, num_heads: int) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = channels // num_heads + self.to_qkv = nn.Linear(channels, 3 * channels) + self.q_norm = TripoSplatRMSNorm(self.head_dim, num_heads) + self.k_norm = TripoSplatRMSNorm(self.head_dim, num_heads) + self.to_out = nn.Linear(channels, channels) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> torch.Tensor: + return self.processor(self, hidden_states, rotary_emb) + + +class TripoSplatRotaryEmbedding(nn.Module): + def __init__(self, model_channels: int, num_heads: int, head_dim: int) -> None: + super().__init__() + self.num_heads = num_heads + hidden_size = int(model_channels * 0.125) + self.norm = FP32LayerNorm(model_channels) + self.gate_map = nn.Linear(model_channels, hidden_size, bias=False) + self.content_map = nn.Linear(model_channels, hidden_size, bias=False) + self.act = nn.SiLU() + self.final_map = nn.Linear(hidden_size, 3 * num_heads, bias=False) + dims = [2 * (head_dim // 6), 2 * (head_dim // 6), head_dim - 4 * (head_dim // 6)] + self.freqs_0 = nn.Parameter(torch.linspace(1.0, 16.0, dims[0] // 2)) + self.freqs_1 = nn.Parameter(torch.linspace(1.0, 16.0, dims[1] // 2)) + self.freqs_2 = nn.Parameter(torch.linspace(1.0, 16.0, dims[2] // 2)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(hidden_states) + features = self.act(self.gate_map(hidden_states)) * self.content_map(hidden_states) + positions = self.final_map(features).unflatten(-1, (self.num_heads, 3)) + angles = [] + for axis, frequencies in enumerate((self.freqs_0, self.freqs_1, self.freqs_2)): + position = positions[..., axis, None] + frequency_tanh = frequencies.tanh() + angle = position * frequency_tanh + position.detach() * (frequencies - frequency_tanh) + angles.append(angle * torch.pi) + angles = torch.cat(angles, dim=-1).float() + return torch.polar(torch.ones_like(angles), angles) + + +class TripoSplatPositionEmbedder(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.channels = channels + self.freq_dim = channels // 3 // 2 + + def forward(self, positions: torch.Tensor) -> torch.Tensor: + frequencies = torch.arange(16, dtype=positions.dtype, device=positions.device) + residual_dim = max(0, self.freq_dim - 16) + residual = ( + torch.arange(residual_dim, dtype=positions.dtype, device=positions.device) / max(residual_dim, 1) * 16 + ) + frequencies = torch.pow(2.0, torch.cat([frequencies, residual])[: self.freq_dim]) + angles = torch.outer(positions.float().reshape(-1), frequencies) * 2 * torch.pi + embeddings = torch.cat([angles.sin(), angles.cos()], dim=-1).reshape(*positions.shape[:-1], -1) + return F.pad(embeddings, (0, self.channels - embeddings.shape[-1])).to(positions.dtype) + + +class TripoSplatTransformerBlock(nn.Module): + def __init__(self, channels: int, num_heads: int, mlp_ratio: float, modulation: bool) -> None: + super().__init__() + self.modulation = modulation + self.norm1 = FP32LayerNorm(channels, elementwise_affine=not modulation, eps=1e-6) + self.norm2 = FP32LayerNorm(channels, elementwise_affine=not modulation, eps=1e-6) + self.attn = TripoSplatAttention(channels, num_heads) + self.mlp = FeedForward(channels, mult=mlp_ratio, activation_fn="gelu-approximate") + if modulation: + self.shift_table = nn.Parameter(torch.randn(1, 6 * channels) / channels**0.5) + + def forward( + self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor, modulation: torch.Tensor | None = None + ) -> torch.Tensor: + if self.modulation: + modulation = modulation + self.shift_table.to(modulation.dtype) + shift_attn, scale_attn, gate_attn, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=1) + residual = self.norm1(hidden_states) * (1 + scale_attn[:, None]) + shift_attn[:, None] + hidden_states = hidden_states + self.attn(residual, rotary_emb) * gate_attn[:, None] + residual = self.norm2(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + hidden_states = hidden_states + self.mlp(residual) * gate_mlp[:, None] + else: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), rotary_emb) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class TripoSplatTimestepEmbedder(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.mlp = nn.Sequential(nn.Linear(256, channels), nn.SiLU(), nn.Linear(channels, channels)) + self.frequencies = torch.exp(-math.log(10000) * torch.arange(128, dtype=torch.float32, device="cpu") / 128) + + def forward(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + frequencies = self.frequencies.to(timestep.device) + angles = timestep[:, None].float() * frequencies[None] + embeddings = torch.cat([torch.cos(angles), torch.sin(angles)], dim=-1) + return self.mlp(embeddings.to(dtype)) + + +class TripoSplatTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): + """Image-conditioned flow transformer for Gaussian and camera latents. + + Parameters: + q_token_length (`int`, defaults to `8192`): + Number of Gaussian latent tokens. + in_channels (`int`, defaults to `16`): + Number of channels in each Gaussian latent token. + cam_channels (`int`, defaults to `5`): + Number of camera latent channels. + out_channels (`int`, defaults to `16`): + Number of predicted Gaussian velocity channels. + model_channels (`int`, defaults to `1024`): + Transformer hidden width. + cond_channels (`int`, defaults to `1280`): + DINOv3 feature width. + cond2_channels (`int`, defaults to `128`): + Packed image VAE feature width. + num_refiner_blocks (`int`, defaults to `2`): + Number of Gaussian and image conditioning refinement blocks. + num_blocks (`int`, defaults to `24`): + Number of joint transformer blocks. + num_heads (`int`, defaults to `16`): + Number of attention heads. + mlp_ratio (`float`, defaults to `4.0`): + Feed-forward hidden width relative to `model_channels`. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["TripoSplatTransformerBlock", "TripoSplatRotaryEmbedding"] + _repeated_blocks = ["TripoSplatTransformerBlock"] + _skip_layerwise_casting_patterns = ["norm", "embedder", "input_layer", "shift_table"] + + @register_to_config + def __init__( + self, + q_token_length: int = 8192, + in_channels: int = 16, + cam_channels: int = 5, + out_channels: int = 16, + model_channels: int = 1024, + cond_channels: int = 1280, + cond2_channels: int = 128, + num_refiner_blocks: int = 2, + num_blocks: int = 24, + num_heads: int = 16, + mlp_ratio: float = 4.0, + ) -> None: + super().__init__() + self.gradient_checkpointing = False + self.t_embedder = TripoSplatTimestepEmbedder(model_channels) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 6 * model_channels)) + self.input_layer = nn.Linear(in_channels, model_channels) + self.cond_embedder = nn.Linear(cond_channels, model_channels) + self.cond_embedder2 = nn.Linear(cond2_channels, model_channels) + # The checkpoint uses a fixed scrambled Sobol lattice for its query positions. + sobol = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123) + positions = sobol.draw(q_token_length, dtype=torch.float32) + # Sobol's first point uses the default dtype, which can be float16 during checkpoint loading. + positions[0] = sobol.shift.to(torch.float32) / 2**sobol.MAXBIT + self.pos_pe = positions.unsqueeze(0) + self.pos_embedder = TripoSplatPositionEmbedder(model_channels) + self.noise_repo_layers = nn.ModuleList( + [ + TripoSplatRotaryEmbedding(model_channels, num_heads, model_channels // num_heads) + for _ in range(num_refiner_blocks) + ] + ) + self.context_repo_layers = nn.ModuleList( + [ + TripoSplatRotaryEmbedding(model_channels, num_heads, model_channels // num_heads) + for _ in range(num_refiner_blocks) + ] + ) + self.repo_layers = nn.ModuleList( + [ + TripoSplatRotaryEmbedding(model_channels, num_heads, model_channels // num_heads) + for _ in range(num_blocks) + ] + ) + self.noise_refiner = nn.ModuleList( + [TripoSplatTransformerBlock(model_channels, num_heads, mlp_ratio, True) for _ in range(num_refiner_blocks)] + ) + self.context_refiner = nn.ModuleList( + [ + TripoSplatTransformerBlock(model_channels, num_heads, mlp_ratio, False) + for _ in range(num_refiner_blocks) + ] + ) + camera_layers = [] + for index in range(num_refiner_blocks - 1): + camera_layers.extend( + [ + nn.Linear(cam_channels if index == 0 else model_channels, model_channels), + nn.GELU(approximate="tanh"), + ] + ) + camera_layers.append(nn.Linear(cam_channels if num_refiner_blocks == 1 else model_channels, model_channels)) + self.cam_refiner = nn.Sequential(*camera_layers) + self.blocks = nn.ModuleList( + [TripoSplatTransformerBlock(model_channels, num_heads, mlp_ratio, True) for _ in range(num_blocks)] + ) + self.shift_table = nn.Parameter(torch.randn(1, 2, model_channels) / model_channels**0.5) + self.out_layer = nn.Linear(model_channels, out_channels) + self.cam_out_layer = nn.Linear(model_channels, cam_channels) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + image_latents: torch.Tensor, + camera_latents: torch.Tensor, + return_dict: bool = True, + ) -> TripoSplatTransformer3DModelOutput | tuple[torch.Tensor, torch.Tensor]: + """ + Args: + hidden_states (`torch.Tensor`): + Gaussian latents of shape `(batch, q_token_length, in_channels)`. + timestep (`torch.Tensor`): + Timesteps of shape `(batch,)`, scaled to the scheduler's training timestep range. + encoder_hidden_states (`torch.Tensor`): + Normalized DINOv3 features of shape `(batch, image_tokens, cond_channels)`. + image_latents (`torch.Tensor`): + Packed VAE features of shape `(batch, image_tokens, cond2_channels)`. + camera_latents (`torch.Tensor`): + Camera latents of shape `(batch, 1, cam_channels)`. + return_dict (`bool`, defaults to `True`): + Whether to return a structured output or a tuple. + + Returns: + `TripoSplatTransformer3DModelOutput` or `tuple`: Gaussian and camera velocity predictions. + """ + latent_length = hidden_states.shape[1] + hidden_states = self.input_layer(hidden_states.to(self.dtype)) + conditioning = self.cond_embedder(encoder_hidden_states.to(self.dtype)) + conditioning = conditioning + self.cond_embedder2(image_latents.to(self.dtype)).to(conditioning.device) + timestep_emb = self.t_embedder(timestep, hidden_states.dtype) + modulation = self.adaLN_modulation(timestep_emb) + hidden_states = hidden_states + self.pos_embedder(self.pos_pe.to(hidden_states.device)).to(self.dtype) + for block, rotary_layer in zip(self.noise_refiner, self.noise_repo_layers): + rotary_emb = rotary_layer(hidden_states) + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(block, hidden_states, rotary_emb, modulation) + else: + hidden_states = block(hidden_states, rotary_emb, modulation) + for block, rotary_layer in zip(self.context_refiner, self.context_repo_layers): + rotary_emb = rotary_layer(conditioning) + if torch.is_grad_enabled() and self.gradient_checkpointing: + conditioning = self._gradient_checkpointing_func(block, conditioning, rotary_emb) + else: + conditioning = block(conditioning, rotary_emb) + camera = self.cam_refiner(camera_latents.to(self.dtype)) + hidden_states = torch.cat( + [hidden_states, conditioning.to(hidden_states.device), camera.to(hidden_states.device)], dim=1 + ) + for block, rotary_layer in zip(self.blocks, self.repo_layers): + rotary_emb = rotary_layer(hidden_states) + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(block, hidden_states, rotary_emb, modulation) + else: + hidden_states = block(hidden_states, rotary_emb, modulation) + camera = F.layer_norm(hidden_states[:, -camera.shape[1] :].float(), hidden_states.shape[-1:]).to(self.dtype) + hidden_states = F.layer_norm(hidden_states[:, :latent_length].float(), hidden_states.shape[-1:]).to(self.dtype) + shift, scale = ( + self.shift_table.to(hidden_states.device) + timestep_emb.to(hidden_states.device)[:, None] + ).chunk(2, dim=1) + hidden_states = self.out_layer(hidden_states * (1 + scale) + shift) + camera = self.cam_out_layer(camera * (1 + scale) + shift).to(hidden_states.device) + if not return_dict: + return hidden_states, camera + return TripoSplatTransformer3DModelOutput(sample=hidden_states, camera_sample=camera) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index aaa37918956a..88c1e5834f0b 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -73,6 +73,7 @@ "FluxKontextAutoBlocks", "FluxKontextModularPipeline", ] + _import_structure["triposplat"] = ["TripoSplatAutoBlocks", "TripoSplatModularPipeline"] _import_structure["flux2"] = [ "Flux2AutoBlocks", "Flux2KleinAutoBlocks", @@ -228,6 +229,7 @@ ) from .stable_diffusion_3 import StableDiffusion3AutoBlocks, StableDiffusion3ModularPipeline from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline + from .triposplat import TripoSplatAutoBlocks, TripoSplatModularPipeline from .wan import ( Wan22Blocks, Wan22Image2VideoBlocks, diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index c501d7b1fc42..7577f8d9e715 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -182,6 +182,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("wan-vace", _create_default_map_fn("Wan22VaceModularPipeline")), ("flux", _create_default_map_fn("FluxModularPipeline")), ("flux-kontext", _create_default_map_fn("FluxKontextModularPipeline")), + ("triposplat", _create_default_map_fn("TripoSplatModularPipeline")), ("flux2", _create_default_map_fn("Flux2ModularPipeline")), ("flux2-klein", _flux2_klein_map_fn), ("ideogram4", _create_default_map_fn("Ideogram4ModularPipeline")), @@ -1825,6 +1826,9 @@ def __init__( for name, value in config_dict.items(): if name in self._component_specs and isinstance(value, (tuple, list)) and len(value) == 2: library, class_name = value + if library is None and class_name is None: + self._component_specs[name].pretrained_model_name_or_path = None + continue component_spec_dict = { "repo": pretrained_model_name_or_path, "subfolder": name, diff --git a/src/diffusers/modular_pipelines/triposplat/__init__.py b/src/diffusers/modular_pipelines/triposplat/__init__.py new file mode 100644 index 000000000000..531aa33dd5ff --- /dev/null +++ b/src/diffusers/modular_pipelines/triposplat/__init__.py @@ -0,0 +1,57 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_torch_available() and is_transformers_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_triposplat"] = ["TripoSplatAutoBlocks"] + _import_structure["modular_pipeline"] = ["TripoSplatModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_torch_available() and is_transformers_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_triposplat import TripoSplatAutoBlocks + from .modular_pipeline import TripoSplatModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/triposplat/before_denoise.py b/src/diffusers/modular_pipelines/triposplat/before_denoise.py new file mode 100644 index 000000000000..3c3b54e6ed8e --- /dev/null +++ b/src/diffusers/modular_pipelines/triposplat/before_denoise.py @@ -0,0 +1,153 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import numpy as np +import torch + +from ...models.transformers.transformer_triposplat import TripoSplatTransformer3DModel +from ...schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipeline, ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +class TripoSplatImageInputStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def description(self): + return "Expand per-image conditioning for the requested number of objects." + + @property + def inputs(self): + return [ + InputParam( + "encoder_hidden_states", required=True, type_hint=torch.Tensor, description="DINOv3 image features." + ), + InputParam( + "image_latents", required=True, type_hint=torch.Tensor, description="Packed VAE image latents." + ), + InputParam.template("num_images_per_prompt", default=1), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam("encoder_hidden_states", type_hint=torch.Tensor, description="Expanded DINOv3 conditioning."), + OutputParam("image_latents", type_hint=torch.Tensor, description="Expanded VAE conditioning."), + ] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + if not isinstance(block_state.num_images_per_prompt, int) or block_state.num_images_per_prompt < 1: + raise ValueError("num_images_per_prompt must be a positive integer.") + if block_state.encoder_hidden_states.shape[:2] != block_state.image_latents.shape[:2]: + raise ValueError("The DINO and VAE features must have matching batch and token dimensions.") + block_state.encoder_hidden_states = block_state.encoder_hidden_states.repeat_interleave( + block_state.num_images_per_prompt, 0 + ) + block_state.image_latents = block_state.image_latents.repeat_interleave(block_state.num_images_per_prompt, 0) + self.set_block_state(state, block_state) + return components, state + + +class TripoSplatSetTimestepsStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self): + return "Configure the shifted Euler schedule." + + @property + def inputs(self): + return [InputParam.template("num_inference_steps", default=20)] + + @property + def intermediate_outputs(self): + return [OutputParam("timesteps", type_hint=torch.Tensor, description="Shifted Euler timesteps.")] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + sigmas = np.linspace(1.0, 0.0, block_state.num_inference_steps + 1)[:-1] + components.scheduler.set_timesteps(sigmas=sigmas, device=components._execution_device) + block_state.timesteps = components.scheduler.timesteps + self.set_block_state(state, block_state) + return components, state + + +class TripoSplatPrepareLatentsStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + return [ComponentSpec("transformer", TripoSplatTransformer3DModel)] + + @property + def description(self): + return "Prepare Gaussian and camera noise." + + @property + def inputs(self): + return [ + InputParam( + "encoder_hidden_states", required=True, type_hint=torch.Tensor, description="Expanded image features." + ), + InputParam.template("generator", type_hint=torch.Generator | list[torch.Generator]), + InputParam.template("latents"), + InputParam( + "camera_latents", type_hint=torch.Tensor, description="Initial camera noise, with shape (batch, 1, 5)." + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Initial Gaussian noise."), + OutputParam("camera_latents", type_hint=torch.Tensor, description="Initial camera noise."), + ] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + device = components._execution_device + batch_size = block_state.encoder_hidden_states.shape[0] + if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size: + raise ValueError("Pass one generator per generated sample.") + shape = (batch_size, components.transformer.config.q_token_length, components.transformer.config.in_channels) + camera_shape = (batch_size, 1, components.transformer.config.cam_channels) + if block_state.latents is None: + block_state.latents = randn_tensor( + shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + elif tuple(block_state.latents.shape) != shape: + raise ValueError(f"Expected latents with shape {shape}, got {tuple(block_state.latents.shape)}.") + if block_state.camera_latents is None: + block_state.camera_latents = randn_tensor( + camera_shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) + elif tuple(block_state.camera_latents.shape) != camera_shape: + raise ValueError( + f"Expected camera_latents with shape {camera_shape}, got {tuple(block_state.camera_latents.shape)}." + ) + block_state.latents = block_state.latents.to(device=device, dtype=torch.float32) + block_state.camera_latents = block_state.camera_latents.to(device=device, dtype=torch.float32) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/triposplat/decoders.py b/src/diffusers/modular_pipelines/triposplat/decoders.py new file mode 100644 index 000000000000..34b404fcd250 --- /dev/null +++ b/src/diffusers/modular_pipelines/triposplat/decoders.py @@ -0,0 +1,84 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import numpy as np +import torch + +from ...models.autoencoders.autoencoder_triposplat import TripoSplatGaussianDecoder +from ..modular_pipeline import ModularPipeline, ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +class TripoSplatGaussianDecodeStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + return [ComponentSpec("decoder", TripoSplatGaussianDecoder)] + + @property + def description(self): + return "Decode denoised latents into Gaussian splats at the requested densities." + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam.template("generator", type_hint=torch.Generator | list[torch.Generator]), + InputParam( + "num_gaussians", + default=262144, + type_hint=int | list[int], + description="Gaussian counts between 32768 and 262144.", + ), + InputParam( + "decoder_generator", + type_hint=torch.Generator | list, + description="Random generator or per-sample generators for octree sampling.", + ), + InputParam("output_type", default="pt", type_hint=str, description="Output format: 'pt' or 'np'."), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "gaussians", + type_hint=torch.Tensor | np.ndarray | list, + description="Gaussian parameters in xyz, SH color, scale, wxyz rotation, opacity order.", + ) + ] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + if block_state.output_type not in ("pt", "np"): + raise ValueError("output_type must be 'pt' or 'np'; remove the decode block to obtain latents.") + counts = ( + block_state.num_gaussians if isinstance(block_state.num_gaussians, list) else [block_state.num_gaussians] + ) + if not counts or any(not isinstance(count, int) or not 32768 <= count <= 262144 for count in counts): + raise ValueError( + "num_gaussians must be an integer or a nonempty list of integers between 32768 and 262144." + ) + gaussians_per_point = components.decoder.config.gaussians_per_point + counts = [round(count / gaussians_per_point) * gaussians_per_point for count in counts] + generator = block_state.generator if block_state.decoder_generator is None else block_state.decoder_generator + gaussians = [components.decoder(block_state.latents, count, generator=generator).sample for count in counts] + if block_state.output_type == "np": + gaussians = [item.cpu().numpy() for item in gaussians] + block_state.gaussians = gaussians if isinstance(block_state.num_gaussians, list) else gaussians[0] + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/triposplat/denoise.py b/src/diffusers/modular_pipelines/triposplat/denoise.py new file mode 100644 index 000000000000..b7726e57722b --- /dev/null +++ b/src/diffusers/modular_pipelines/triposplat/denoise.py @@ -0,0 +1,201 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import torch + +from ...configuration_utils import FrozenDict +from ...guiders.triposplat_classifier_free_guidance import TripoSplatClassifierFreeGuidance +from ...models.transformers.transformer_triposplat import TripoSplatTransformer3DModel +from ...schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipeline, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +class TripoSplatLoopDenoiser(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + return [ + ComponentSpec("transformer", TripoSplatTransformer3DModel), + ComponentSpec( + "guider", + TripoSplatClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ] + + @property + def description(self): + return "Predict guided Gaussian and camera velocities." + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam("camera_latents", required=True, type_hint=torch.Tensor, description="Current camera latents."), + InputParam( + "encoder_hidden_states", + required=True, + type_hint=torch.Tensor, + description="DINOv3 image conditioning.", + ), + InputParam( + "image_latents", required=True, type_hint=torch.Tensor, description="Packed VAE image conditioning." + ), + InputParam.template("num_inference_steps", default=20), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam("noise_pred", type_hint=torch.Tensor, description="Guided Gaussian velocity."), + OutputParam("camera_pred", type_hint=torch.Tensor, description="Guided camera velocity."), + ] + + @torch.no_grad() + def __call__( + self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> tuple[ModularPipeline, BlockState]: + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs( + { + "encoder_hidden_states": ( + block_state.encoder_hidden_states, + torch.zeros_like(block_state.encoder_hidden_states), + ), + "image_latents": (block_state.image_latents, torch.zeros_like(block_state.image_latents)), + } + ) + for condition in guider_state: + components.guider.prepare_models(components.transformer) + noise_pred, camera_pred = components.transformer( + block_state.latents, + t.expand(block_state.latents.shape[0]), + condition.encoder_hidden_states.to(block_state.latents.device), + condition.image_latents.to(block_state.latents.device), + block_state.camera_latents, + return_dict=False, + ) + condition.noise_pred = torch.cat([noise_pred.flatten(1), camera_pred.flatten(1)], dim=1) + components.guider.cleanup_models(components.transformer) + prediction = components.guider(guider_state)[0] + latent_size = block_state.latents[0].numel() + block_state.noise_pred = prediction[:, :latent_size].reshape_as(block_state.latents) + block_state.camera_pred = prediction[:, latent_size:].reshape_as(block_state.camera_latents) + return components, block_state + + +class TripoSplatLoopSchedulerStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def description(self): + return "Advance Gaussian and camera latents with one Euler step." + + @property + def inputs(self): + return [ + InputParam.template("latents", required=True), + InputParam("camera_latents", required=True, type_hint=torch.Tensor, description="Current camera latents."), + InputParam("noise_pred", required=True, type_hint=torch.Tensor, description="Gaussian velocity."), + InputParam("camera_pred", required=True, type_hint=torch.Tensor, description="Camera velocity."), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("camera_latents", type_hint=torch.Tensor, description="Updated camera latents."), + ] + + @torch.no_grad() + def __call__( + self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + ) -> tuple[ModularPipeline, BlockState]: + latent_size = block_state.latents[0].numel() + prediction = torch.cat([block_state.noise_pred.flatten(1), block_state.camera_pred.flatten(1)], dim=1) + sample = torch.cat([block_state.latents.flatten(1), block_state.camera_latents.flatten(1)], dim=1) + sample = components.scheduler.step(prediction.float(), t, sample, return_dict=False)[0] + block_state.latents = sample[:, :latent_size].reshape_as(block_state.latents) + block_state.camera_latents = sample[:, latent_size:].reshape_as(block_state.camera_latents) + return components, block_state + + +# auto_docstring +class TripoSplatDenoiseStep(LoopSequentialPipelineBlocks): + """ + Denoise Gaussian and camera latents over the configured schedule. + + Components: + transformer (`TripoSplatTransformer3DModel`) guider (`TripoSplatClassifierFreeGuidance`) scheduler + (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + timesteps (`Tensor`): + Euler timesteps. + num_inference_steps (`int`, *optional*, defaults to 20): + The number of denoising steps. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + camera_latents (`Tensor`): + Current camera latents. + encoder_hidden_states (`Tensor`): + DINOv3 image conditioning. + image_latents (`Tensor`): + Packed VAE image conditioning. + + Outputs: + latents (`Tensor`): + Denoised latents. + camera_latents (`Tensor`): + Updated camera latents. + """ + + model_name = "triposplat" + block_classes = [TripoSplatLoopDenoiser, TripoSplatLoopSchedulerStep] + block_names = ["denoiser", "scheduler"] + + @property + def description(self): + return "Denoise Gaussian and camera latents over the configured schedule." + + @property + def loop_inputs(self): + return [ + InputParam("timesteps", required=True, type_hint=torch.Tensor, description="Euler timesteps."), + InputParam.template("num_inference_steps", default=20), + ] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + progress_bar.update() + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/triposplat/encoders.py b/src/diffusers/modular_pipelines/triposplat/encoders.py new file mode 100644 index 000000000000..2e4d82f1644e --- /dev/null +++ b/src/diffusers/modular_pipelines/triposplat/encoders.py @@ -0,0 +1,216 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import torch +from PIL import Image +from torch.nn import functional as F +from transformers import DINOv3ViTModel + +from ...configuration_utils import FrozenDict +from ...image_processor import TripoSplatImageProcessor +from ...models import AutoencoderKLFlux2 +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipeline, ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +def processor_spec(): + return ComponentSpec( + "image_processor", + TripoSplatImageProcessor, + config=FrozenDict({"canvas_size": 1024}), + default_creation_method="from_config", + ) + + +class TripoSplatImagePreprocessStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + from diffusers import BiRefNetModel + + return [processor_spec(), ComponentSpec("background_remover", BiRefNetModel)] + + @property + def description(self): + return "Prepare foreground images on a black square canvas." + + @property + def inputs(self): + return [ + InputParam.template("image", required=True), + InputParam("erode_radius", default=1, type_hint=int, description="Radius of the foreground mask erosion."), + InputParam( + "is_preprocessed", + default=False, + type_hint=bool, + description="Whether inputs are already cropped RGB images composited on black.", + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "preprocessed_images", type_hint=list, description="Prepared RGB images on black square canvases." + ) + ] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + images = components.image_processor.to_pil(block_state.image) + if not block_state.is_preprocessed: + images = components.image_processor.resize_shortest_edge(images) + for index, item in enumerate(images): + if item.mode == "RGBA" and item.getchannel("A").getextrema()[0] < 255: + continue + if components.background_remover is None: + raise ValueError( + "RGB inputs require background_remover, or set is_preprocessed=True for prepared RGB images." + ) + pixels = components.image_processor.preprocess([item.convert("RGB")]) + pixels = F.interpolate( + pixels, + size=(components.background_remover.config.sample_size,) * 2, + mode="bilinear", + align_corners=True, + ) + mean = pixels.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = pixels.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + pixels = ((pixels - mean) / std).to( + device=components._execution_device, dtype=components.background_remover.dtype + ) + alpha = components.background_remover(pixels).sample + alpha = F.interpolate( + alpha.float(), size=(item.height, item.width), mode="bilinear", align_corners=True + )[0, 0] + alpha = (alpha.clamp(0, 1) * 255).to(torch.uint8).cpu().numpy() + item = item.convert("RGB") + item.putalpha(Image.fromarray(alpha)) + images[index] = item + block_state.preprocessed_images = components.image_processor.prepare_foreground( + images, erode_radius=block_state.erode_radius, is_preprocessed=block_state.is_preprocessed + ) + self.set_block_state(state, block_state) + return components, state + + +class TripoSplatImageEncoderStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + return [processor_spec(), ComponentSpec("image_encoder", DINOv3ViTModel)] + + @property + def description(self): + return "Encode prepared images into DINOv3 features." + + @property + def inputs(self): + return [InputParam("preprocessed_images", required=True, type_hint=list, description="Prepared RGB images.")] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "encoder_hidden_states", type_hint=torch.Tensor, description="Normalized DINOv3 image features." + ) + ] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + pixels = components.image_processor.preprocess(block_state.preprocessed_images).to( + device=components._execution_device + ) + mean = pixels.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = pixels.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + pixels = ((pixels - mean) / std).to(components.image_encoder.dtype) + original_inv_freq = components.image_encoder.rope_embeddings.inv_freq + try: + components.image_encoder.rope_embeddings.inv_freq = original_inv_freq.to( + dtype=components.image_encoder.dtype + ) + features = components.image_encoder(pixel_values=pixels).last_hidden_state + finally: + components.image_encoder.rope_embeddings.inv_freq = original_inv_freq + block_state.encoder_hidden_states = F.layer_norm(features.float(), features.shape[-1:]) + self.set_block_state(state, block_state) + return components, state + + +class TripoSplatVaeEncoderStep(ModularPipelineBlocks): + model_name = "triposplat" + + @property + def expected_components(self): + return [processor_spec(), ComponentSpec("vae", AutoencoderKLFlux2)] + + @property + def expected_configs(self): + from ..modular_pipeline_utils import ConfigSpec + + return [ConfigSpec("num_prefix_tokens", default=5)] + + @property + def description(self): + return "Sample and normalize packed Flux2 VAE image latents." + + @property + def inputs(self): + return [ + InputParam("preprocessed_images", required=True, type_hint=list, description="Prepared RGB images."), + InputParam.template("generator", type_hint=torch.Generator | list[torch.Generator]), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "image_latents", type_hint=torch.Tensor, description="Packed VAE latents with zero prefix tokens." + ) + ] + + @torch.no_grad() + def __call__(self, components: ModularPipeline, state: PipelineState) -> tuple[ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + generator = block_state.generator + if isinstance(generator, list): + if len(generator) % len(block_state.preprocessed_images): + raise ValueError("The generator list must contain one generator per image or generated sample.") + generator = generator[:: len(generator) // len(block_state.preprocessed_images)] + pixels = components.image_processor.preprocess(block_state.preprocessed_images).to( + device=components._execution_device, dtype=components.vae.dtype + ) + pixels = components.image_processor.normalize(pixels) + parameters = components.vae.encode(pixels).latent_dist.parameters + # TripoSplat samples the posterior variance without the distribution's log-variance clamp. + mean, logvar = parameters.chunk(2, dim=1) + noise = randn_tensor(mean.shape, generator=generator, device=components._execution_device, dtype=mean.dtype) + latents = mean + torch.exp(0.5 * logvar) * noise + batch, channels, height, width = latents.shape + latents = latents.view(batch, channels, height // 2, 2, width // 2, 2).permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch, channels * 4, height // 2, width // 2) + bn_mean = components.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + bn_std = torch.sqrt(components.vae.bn.running_var.view(1, -1, 1, 1) + components.vae.bn.eps).to( + latents.device, latents.dtype + ) + latents = ((latents - bn_mean) / bn_std).float().flatten(2).transpose(1, 2).contiguous() + block_state.image_latents = F.pad(latents, (0, 0, components.config.num_prefix_tokens, 0)) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/triposplat/modular_blocks_triposplat.py b/src/diffusers/modular_pipelines/triposplat/modular_blocks_triposplat.py new file mode 100644 index 000000000000..88dde400906f --- /dev/null +++ b/src/diffusers/modular_pipelines/triposplat/modular_blocks_triposplat.py @@ -0,0 +1,162 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import numpy as np +import torch + +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import TripoSplatImageInputStep, TripoSplatPrepareLatentsStep, TripoSplatSetTimestepsStep +from .decoders import TripoSplatGaussianDecodeStep +from .denoise import TripoSplatDenoiseStep +from .encoders import TripoSplatImageEncoderStep, TripoSplatImagePreprocessStep, TripoSplatVaeEncoderStep + + +TripoSplatDenoiseBlocks = InsertableDict( + [ + ("inputs", TripoSplatImageInputStep()), + ("timesteps", TripoSplatSetTimestepsStep()), + ("prepare", TripoSplatPrepareLatentsStep()), + ("loop", TripoSplatDenoiseStep()), + ] +) + + +# auto_docstring +class TripoSplatCoreDenoiseStep(SequentialPipelineBlocks): + """ + Prepare noise and denoise with image conditioning. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`TripoSplatTransformer3DModel`) guider + (`TripoSplatClassifierFreeGuidance`) + + Inputs: + encoder_hidden_states (`Tensor`): + DINOv3 image features. + image_latents (`Tensor`): + Packed VAE image latents. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + num_inference_steps (`int`, *optional*, defaults to 20): + The number of denoising steps. + generator (`Generator | list`, *optional*): + Torch generator for deterministic generation. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + camera_latents (`Tensor`, *optional*): + Initial camera noise, with shape (batch, 1, 5). + + Outputs: + latents (`Tensor`): + Denoised latents. + camera_latents (`Tensor`): + Denoised camera latents. + """ + + model_name = "triposplat" + block_classes = TripoSplatDenoiseBlocks.values() + block_names = TripoSplatDenoiseBlocks.keys() + + @property + def description(self): + return "Prepare noise and denoise with image conditioning." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("camera_latents", type_hint=torch.Tensor, description="Denoised camera latents."), + ] + + +TripoSplatBlocks = InsertableDict( + [ + ("preprocess", TripoSplatImagePreprocessStep()), + ("image_encoder", TripoSplatImageEncoderStep()), + ("vae_encoder", TripoSplatVaeEncoderStep()), + ("denoise", TripoSplatCoreDenoiseStep()), + ("decode", TripoSplatGaussianDecodeStep()), + ] +) + + +# auto_docstring +class TripoSplatAutoBlocks(SequentialPipelineBlocks): + """ + Generate Gaussian splats from foreground images. + + Components: + image_processor (`TripoSplatImageProcessor`) background_remover (`BiRefNetModel`) image_encoder + (`DINOv3ViTModel`) vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer + (`TripoSplatTransformer3DModel`) guider (`TripoSplatClassifierFreeGuidance`) decoder + (`TripoSplatGaussianDecoder`) + + Configs: + num_prefix_tokens (default: 5) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + erode_radius (`int`, *optional*, defaults to 1): + Radius of the foreground mask erosion. + is_preprocessed (`bool`, *optional*, defaults to False): + Whether inputs are already cropped RGB images composited on black. + generator (`Generator | list`, *optional*): + Torch generator for deterministic generation. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + num_inference_steps (`int`, *optional*, defaults to 20): + The number of denoising steps. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + camera_latents (`Tensor`, *optional*): + Initial camera noise, with shape (batch, 1, 5). + num_gaussians (`int | list`, *optional*, defaults to 262144): + Gaussian counts between 32768 and 262144. + decoder_generator (`Generator | list`, *optional*): + Random generator or per-sample generators for octree sampling. + output_type (`str`, *optional*, defaults to pt): + Output format: 'pt' or 'np'. + + Outputs: + gaussians (`Tensor | ndarray | list`): + Decoded Gaussian parameters. + latents (`Tensor`): + Denoised latents. + camera_latents (`Tensor`): + Denoised camera latents. + preprocessed_images (`list`): + Prepared RGB images. + """ + + model_name = "triposplat" + block_classes = TripoSplatBlocks.values() + block_names = TripoSplatBlocks.keys() + + @property + def description(self): + return "Generate Gaussian splats from foreground images." + + @property + def outputs(self): + return [ + OutputParam( + "gaussians", type_hint=torch.Tensor | np.ndarray | list, description="Decoded Gaussian parameters." + ), + OutputParam.template("latents"), + OutputParam("camera_latents", type_hint=torch.Tensor, description="Denoised camera latents."), + OutputParam("preprocessed_images", type_hint=list, description="Prepared RGB images."), + ] diff --git a/src/diffusers/modular_pipelines/triposplat/modular_pipeline.py b/src/diffusers/modular_pipelines/triposplat/modular_pipeline.py new file mode 100644 index 000000000000..6d25c3c0d6da --- /dev/null +++ b/src/diffusers/modular_pipelines/triposplat/modular_pipeline.py @@ -0,0 +1,20 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +from ..modular_pipeline import ModularPipeline + + +class TripoSplatModularPipeline(ModularPipeline): + default_blocks_name = "TripoSplatAutoBlocks" diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index b8cb626d8283..ce47118d9fe8 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -170,6 +170,7 @@ ] _import_structure["bria"] = ["BriaPipeline"] _import_structure["bria_fibo"] = ["BriaFiboPipeline", "BriaFiboEditPipeline"] + _import_structure["triposplat"] = ["TripoSplatPipeline", "TripoSplatPipelineOutput"] _import_structure["flux2"] = [ "Flux2Pipeline", "Flux2KleinPipeline", @@ -927,6 +928,7 @@ StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline, ) + from .triposplat import TripoSplatPipeline, TripoSplatPipelineOutput from .visualcloze import VisualClozeGenerationPipeline, VisualClozePipeline from .wan import ( WanAnimatePipeline, diff --git a/src/diffusers/pipelines/triposplat/__init__.py b/src/diffusers/pipelines/triposplat/__init__.py new file mode 100644 index 000000000000..fc76602d7931 --- /dev/null +++ b/src/diffusers/pipelines/triposplat/__init__.py @@ -0,0 +1,75 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_torchvision_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_torch_available() and is_transformers_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_output"] = ["TripoSplatPipelineOutput"] + _import_structure["pipeline_triposplat"] = ["TripoSplatPipeline"] + +try: + if not (is_torch_available() and is_torchvision_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_torchvision_objects + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_torchvision_objects)) +else: + _import_structure["modeling_birefnet"] = ["BiRefNetModel"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_torch_available() and is_transformers_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .pipeline_output import TripoSplatPipelineOutput + from .pipeline_triposplat import TripoSplatPipeline + try: + if not (is_torch_available() and is_torchvision_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_torchvision_objects import * # noqa F403 + else: + from .modeling_birefnet import BiRefNetModel +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/triposplat/modeling_birefnet.py b/src/diffusers/pipelines/triposplat/modeling_birefnet.py new file mode 100644 index 000000000000..ae82bf233092 --- /dev/null +++ b/src/diffusers/pipelines/triposplat/modeling_birefnet.py @@ -0,0 +1,556 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +import math +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin +from ...models.attention import AttentionMixin, AttentionModuleMixin +from ...models.attention_dispatch import dispatch_attention_fn +from ...models.modeling_utils import ModelMixin +from ...utils import BaseOutput, requires_backends + + +@dataclass +class BiRefNetOutput(BaseOutput): + """Foreground probabilities predicted by BiRefNet. + + Args: + sample (`torch.Tensor`): + Probabilities between zero and one, shaped `(batch, 1, height, width)`. + """ + + sample: torch.Tensor + + +class BiRefNetAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, attn: "BiRefNetWindowAttention", hidden_states: torch.Tensor, mask: torch.Tensor | None = None + ) -> torch.Tensor: + batch, length, channels = hidden_states.shape + query, key, value = ( + attn.qkv(hidden_states).reshape(batch, length, 3, attn.num_heads, channels // attn.num_heads).unbind(2) + ) + bias = attn.relative_position_bias_table[attn.relative_position_index.reshape(-1)] + bias = bias.reshape(length, length, attn.num_heads).permute(2, 0, 1)[None] + if mask is not None: + bias = bias + mask[:, None] + bias = bias.repeat(batch // mask.shape[0], 1, 1, 1) + hidden_states = dispatch_attention_fn( + query, key, value, attn_mask=bias, backend=self._attention_backend, parallel_config=self._parallel_config + ) + return attn.proj(hidden_states.reshape(batch, length, channels)) + + +class BiRefNetSwinMlp(nn.Module): + def __init__(self, in_features: int, hidden_features: int) -> None: + super().__init__() + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden_features, in_features) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2(self.act(self.fc1(x))) + + +class BiRefNetWindowAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = BiRefNetAttnProcessor + _available_processors = [BiRefNetAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, dim: int, window_size: tuple[int, int], num_heads: int) -> None: + super().__init__() + self.dim = dim + self.window_size = window_size + self.num_heads = num_heads + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) + coords_h = torch.arange(window_size[0]) + coords_w = torch.arange(window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) + coords_flatten = torch.flatten(coords, 1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += window_size[0] - 1 + relative_coords[:, :, 1] += window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * window_size[1] - 1 + self.register_buffer("relative_position_index", relative_coords.sum(-1)) + self.qkv = nn.Linear(dim, dim * 3, bias=True) + self.proj = nn.Linear(dim, dim) + nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + return self.processor(self, hidden_states, mask) + + +class BiRefNetSwinBlock(nn.Module): + def __init__(self, dim: int, num_heads: int, window_size: int, shift_size: int, mlp_ratio: float = 4.0) -> None: + super().__init__() + self.dim = dim + self.window_size = window_size + self.shift_size = shift_size + self.norm1 = nn.LayerNorm(dim) + self.attn = BiRefNetWindowAttention(dim, (window_size, window_size), num_heads) + self.norm2 = nn.LayerNorm(dim) + self.mlp = BiRefNetSwinMlp(dim, int(dim * mlp_ratio)) + + def forward(self, x: torch.Tensor, mask_matrix: torch.Tensor, H: int, W: int) -> torch.Tensor: + B, L, C = x.shape + shortcut = x + x = self.norm1(x).view(B, H, W, C) + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b)) + _, Hp, Wp, _ = x.shape + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + x_windows = shifted_x.view( + B, Hp // self.window_size, self.window_size, Wp // self.window_size, self.window_size, C + ) + x_windows = x_windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.window_size**2, C) + attn_windows = self.attn(x_windows, mask=attn_mask).view(-1, self.window_size, self.window_size, C) + shifted_x = attn_windows.view( + B, Hp // self.window_size, Wp // self.window_size, self.window_size, self.window_size, C + ) + shifted_x = shifted_x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, C) + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + x = x.view(B, H * W, C) + x = shortcut + x + x = x + self.mlp(self.norm2(x)) + return x + + +class BiRefNetPatchMerging(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = nn.LayerNorm(4 * dim) + + def forward(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor: + B, L, C = x.shape + x = x.view(B, H, W, C) + if H % 2 == 1 or W % 2 == 1: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + x0 = x[:, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, :] + x2 = x[:, 0::2, 1::2, :] + x3 = x[:, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3], -1).view(B, -1, 4 * C) + return self.reduction(self.norm(x)) + + +class BiRefNetSwinBasicLayer(nn.Module): + def __init__( + self, dim: int, depth: int, num_heads: int, window_size: int, mlp_ratio: float = 4.0, downsample: bool = True + ) -> None: + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.blocks = nn.ModuleList( + [ + BiRefNetSwinBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if i % 2 == 0 else window_size // 2, + mlp_ratio=mlp_ratio, + ) + for i in range(depth) + ] + ) + self.downsample = BiRefNetPatchMerging(dim) if downsample else None + + def forward(self, x: torch.Tensor, H: int, W: int) -> tuple[torch.Tensor, int, int, torch.Tensor, int, int]: + Hp = int(math.ceil(H / self.window_size)) * self.window_size + Wp = int(math.ceil(W / self.window_size)) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + mask_windows = img_mask.view( + 1, Hp // self.window_size, self.window_size, Wp // self.window_size, self.window_size, 1 + ) + mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.window_size**2) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = ( + attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype) + ) + for blk in self.blocks: + x = blk(x, attn_mask, H, W) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = ((H + 1) // 2, (W + 1) // 2) + return (x, H, W, x_down, Wh, Ww) + return (x, H, W, x, H, W) + + +class BiRefNetSwinPatchEmbed(nn.Module): + def __init__(self, patch_size: int = 4, in_channels: int = 3, embed_dim: int = 192) -> None: + super().__init__() + self.patch_size = (patch_size, patch_size) + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + self.norm = nn.LayerNorm(embed_dim) + self.embed_dim = embed_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, _, H, W = x.shape + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + x = self.proj(x) + Wh, Ww = (x.size(2), x.size(3)) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + return x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + +class BiRefNetSwinLarge(nn.Module): + def __init__(self, embed_dim: int, depths: tuple[int, ...], num_heads: tuple[int, ...], window_size: int) -> None: + super().__init__() + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.patch_embed = BiRefNetSwinPatchEmbed(patch_size=4, in_channels=3, embed_dim=embed_dim) + self.layers = nn.ModuleList( + [ + BiRefNetSwinBasicLayer( + dim=int(embed_dim * 2**i), + depth=depths[i], + num_heads=num_heads[i], + window_size=window_size, + downsample=i < self.num_layers - 1, + ) + for i in range(self.num_layers) + ] + ) + num_features = [int(embed_dim * 2**i) for i in range(self.num_layers)] + self.num_features = num_features + for i in range(self.num_layers): + self.add_module(f"norm{i}", nn.LayerNorm(num_features[i])) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]: + x = self.patch_embed(x) + Wh, Ww = (x.size(2), x.size(3)) + x = x.flatten(2).transpose(1, 2) + outs = [] + for i in range(self.num_layers): + x_out, H, W, x, Wh, Ww = self.layers[i](x, Wh, Ww) + norm_layer = getattr(self, f"norm{i}") + x_out = norm_layer(x_out) + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + return tuple(outs) + + +class BiRefNetDeformableConv2d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int] = 3, + stride: int | tuple[int, int] = 1, + padding: int | tuple[int, int] = 1, + bias: bool = False, + ) -> None: + super().__init__() + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size) + self.stride = (stride, stride) if isinstance(stride, int) else stride + self.padding = padding + self.offset_conv = nn.Conv2d( + in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=True, + ) + self.modulator_conv = nn.Conv2d( + in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=True, + ) + convolution = nn.Conv2d( + in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias + ) + self.weight = convolution.weight + self.bias = convolution.bias + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from torchvision.ops import deform_conv2d + + offset = self.offset_conv(x) + modulator = 2.0 * torch.sigmoid(self.modulator_conv(x)) + dtype = x.dtype + weight, bias = self.weight, self.bias + if dtype == torch.bfloat16 or (x.device.type == "cpu" and dtype == torch.float16): + x, offset, modulator, weight = x.float(), offset.float(), modulator.float(), weight.float() + bias = bias.float() if bias is not None else None + return deform_conv2d( + input=x, offset=offset, weight=weight, bias=bias, padding=self.padding, mask=modulator, stride=self.stride + ).to(dtype) + + +class BiRefNetASPPModuleDeformable(nn.Module): + def __init__(self, in_channels: int, planes: int, kernel_size: int, padding: int) -> None: + super().__init__() + self.atrous_conv = BiRefNetDeformableConv2d( + in_channels, planes, kernel_size=kernel_size, stride=1, padding=padding, bias=False + ) + self.bn = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.bn(self.atrous_conv(x))) + + +class BiRefNetASPPDeformable(nn.Module): + def __init__(self, in_channels: int) -> None: + super().__init__() + out_channels = in_channels + inter = 256 + self.aspp1 = BiRefNetASPPModuleDeformable(in_channels, inter, 1, padding=0) + self.aspp_deforms = nn.ModuleList( + [BiRefNetASPPModuleDeformable(in_channels, inter, k, padding=k // 2) for k in (1, 3, 7)] + ) + self.global_avg_pool = nn.Sequential( + nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, inter, 1, stride=1, bias=False), + nn.BatchNorm2d(inter), + nn.ReLU(inplace=True), + ) + self.conv1 = nn.Conv2d(inter * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x1 = self.aspp1(x) + x_aspp_deforms = [m(x) for m in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode="bilinear", align_corners=True) + y = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + return self.relu(self.bn1(self.conv1(y))) + + +class BiRefNetBasicDecBlk(nn.Module): + def __init__(self, in_channels: int, out_channels: int, inter_channels: int = 64) -> None: + super().__init__() + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) + self.relu_in = nn.ReLU(inplace=True) + self.dec_att = BiRefNetASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.relu_in(self.bn_in(self.conv_in(x))) + x = self.dec_att(x) + x = self.bn_out(self.conv_out(x)) + return x + + +class BiRefNetBasicLatBlk(nn.Module): + def __init__(self, in_channels: int, out_channels: int) -> None: + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class BiRefNetSimpleConvs(nn.Module): + def __init__(self, in_channels: int, out_channels: int, inter_channels: int = 64) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv_out(self.conv1(x)) + + +def BiRefNetimage2patches(image: torch.Tensor, patch_ref: torch.Tensor) -> torch.Tensor: + """Stack non-overlapping image patches along the channel axis to match `patch_ref` spatial dimensions.""" + b, c, h_full, w_full = image.shape + hg, wg = (h_full // patch_ref.shape[-2], w_full // patch_ref.shape[-1]) + h, w = (h_full // hg, w_full // wg) + return image.view(b, c, hg, h, wg, w).permute(0, 1, 2, 4, 3, 5).reshape(b, c * hg * wg, h, w) + + +class BiRefNetBiRefNetDecoder(nn.Module): + def __init__(self, channels: tuple[int, ...] = (3072, 1536, 768, 384)) -> None: + super().__init__() + c = channels + self.ipt_blk5 = BiRefNetSimpleConvs(2**10 * 3, c[0] // 8, inter_channels=64) + self.ipt_blk4 = BiRefNetSimpleConvs(2**8 * 3, c[0] // 8, inter_channels=64) + self.ipt_blk3 = BiRefNetSimpleConvs(2**6 * 3, c[1] // 8, inter_channels=64) + self.ipt_blk2 = BiRefNetSimpleConvs(2**4 * 3, c[2] // 8, inter_channels=64) + self.ipt_blk1 = BiRefNetSimpleConvs(2**0 * 3, c[3] // 8, inter_channels=64) + self.decoder_block4 = BiRefNetBasicDecBlk(c[0] + c[0] // 8, c[1]) + self.decoder_block3 = BiRefNetBasicDecBlk(c[1] + c[0] // 8, c[2]) + self.decoder_block2 = BiRefNetBasicDecBlk(c[2] + c[1] // 8, c[3]) + self.decoder_block1 = BiRefNetBasicDecBlk(c[3] + c[2] // 8, c[3] // 2) + self.conv_out1 = nn.Sequential(nn.Conv2d(c[3] // 2 + c[3] // 8, 1, 1, 1, 0)) + self.lateral_block4 = BiRefNetBasicLatBlk(c[1], c[1]) + self.lateral_block3 = BiRefNetBasicLatBlk(c[2], c[2]) + self.lateral_block2 = BiRefNetBasicLatBlk(c[3], c[3]) + _N = 16 + + def _gdt_branch(in_c): + return nn.Sequential(nn.Conv2d(in_c, _N, 3, 1, 1), nn.BatchNorm2d(_N), nn.ReLU(inplace=True)) + + self.gdt_convs_4 = _gdt_branch(c[1]) + self.gdt_convs_3 = _gdt_branch(c[2]) + self.gdt_convs_2 = _gdt_branch(c[3]) + + def _head_1x1(): + return nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = _head_1x1() + self.gdt_convs_attn_3 = _head_1x1() + self.gdt_convs_attn_2 = _head_1x1() + + def forward( + self, x: torch.Tensor, x1: torch.Tensor, x2: torch.Tensor, x3: torch.Tensor, x4: torch.Tensor + ) -> torch.Tensor: + x4 = torch.cat((x4, self.ipt_blk5(BiRefNetimage2patches(x, x4)).to(x4.device)), 1) + p4 = self.decoder_block4(x4) + p4 = p4 * self.gdt_convs_attn_4(self.gdt_convs_4(p4)).sigmoid().to(p4.device) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode="bilinear", align_corners=True) + _p3 = _p4 + self.lateral_block4(x3).to(_p4.device) + _p3 = torch.cat((_p3, self.ipt_blk4(BiRefNetimage2patches(x, _p3)).to(_p3.device)), 1) + p3 = self.decoder_block3(_p3) + p3 = p3 * self.gdt_convs_attn_3(self.gdt_convs_3(p3)).sigmoid().to(p3.device) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode="bilinear", align_corners=True) + _p2 = _p3 + self.lateral_block3(x2).to(_p3.device) + _p2 = torch.cat((_p2, self.ipt_blk3(BiRefNetimage2patches(x, _p2)).to(_p2.device)), 1) + p2 = self.decoder_block2(_p2) + p2 = p2 * self.gdt_convs_attn_2(self.gdt_convs_2(p2)).sigmoid().to(p2.device) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode="bilinear", align_corners=True) + _p1 = _p2 + self.lateral_block2(x1).to(_p2.device) + _p1 = torch.cat((_p1, self.ipt_blk2(BiRefNetimage2patches(x, _p1)).to(_p1.device)), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode="bilinear", align_corners=True) + _p1 = torch.cat((_p1, self.ipt_blk1(BiRefNetimage2patches(x, _p1)).to(_p1.device)), 1) + return self.conv_out1(_p1) + + +class BiRefNetModel(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): + """Optional BiRefNet foreground segmentation model used by TripoSplat. + + Parameters: + embed_dim (`int`, defaults to `192`): + Initial Swin backbone hidden width. + depths (`tuple[int, ...]`, defaults to `(2, 2, 18, 2)`): + Number of blocks in each Swin stage. + num_heads (`tuple[int, ...]`, defaults to `(6, 12, 24, 48)`): + Number of attention heads in each Swin stage. + window_size (`int`, defaults to `12`): + Side length of each attention window. + sample_size (`int`, defaults to `1024`): + Image side length used for foreground mask prediction. + """ + + _skip_layerwise_casting_patterns = ["norm", "patch_embed", "bn"] + _repeated_blocks = ["BiRefNetSwinBlock"] + _no_split_modules = ["BiRefNetSwinBlock", "BiRefNetBasicDecBlk", "BiRefNetDeformableConv2d"] + + @register_to_config + def __init__( + self, + embed_dim: int = 192, + depths: tuple[int, ...] = (2, 2, 18, 2), + num_heads: tuple[int, ...] = (6, 12, 24, 48), + window_size: int = 12, + sample_size: int = 1024, + ) -> None: + super().__init__() + requires_backends(self, ["torchvision"]) + self.bb = BiRefNetSwinLarge(embed_dim, depths, num_heads, window_size) + self._CHANNELS = (embed_dim * 16, embed_dim * 8, embed_dim * 4, embed_dim * 2) + cxt = list(self._CHANNELS[1:][::-1][-3:]) + self.squeeze_module = nn.Sequential(BiRefNetBasicDecBlk(self._CHANNELS[0] + sum(cxt), self._CHANNELS[0])) + self.decoder = BiRefNetBiRefNetDecoder(channels=self._CHANNELS) + self.eval() + + def forward(self, sample: torch.Tensor, return_dict: bool = True) -> BiRefNetOutput | tuple[torch.Tensor]: + """ + Args: + sample (`torch.Tensor`): + Normalized RGB images of shape `(batch, 3, height, width)`. + return_dict (`bool`, defaults to `True`): + Whether to return a structured output or a tuple. + + Returns: + `BiRefNetOutput` or `tuple`: Foreground probabilities of shape `(batch, 1, height, width)`. + """ + x = sample + x1, x2, x3, x4 = self.bb(x) + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H // 2, W // 2), mode="bilinear", align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode="bilinear", align_corners=True)], 1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode="bilinear", align_corners=True)], 1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode="bilinear", align_corners=True)], 1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode="bilinear", align_corners=True)], 1) + x4 = torch.cat( + [ + F.interpolate(x1, size=x4.shape[2:], mode="bilinear", align_corners=True).to(x4.device), + F.interpolate(x2, size=x4.shape[2:], mode="bilinear", align_corners=True).to(x4.device), + F.interpolate(x3, size=x4.shape[2:], mode="bilinear", align_corners=True).to(x4.device), + x4, + ], + 1, + ) + x4 = self.squeeze_module(x4) + logits = self.decoder(sample, x1, x2, x3, x4) + alpha = torch.sigmoid(logits) + if not return_dict: + return (alpha,) + return BiRefNetOutput(sample=alpha) diff --git a/src/diffusers/pipelines/triposplat/pipeline_output.py b/src/diffusers/pipelines/triposplat/pipeline_output.py new file mode 100644 index 000000000000..caa6e6d063a3 --- /dev/null +++ b/src/diffusers/pipelines/triposplat/pipeline_output.py @@ -0,0 +1,43 @@ +# Copyright 2026 The Hugging Face Team. All rights reserved. +# +# 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. + +from dataclasses import dataclass + +import numpy as np +import torch +from PIL import Image + +from ...utils import BaseOutput + + +@dataclass +class TripoSplatPipelineOutput(BaseOutput): + """Generated Gaussian parameters and the conditioning images. + + Args: + gaussians (`torch.Tensor`, `numpy.ndarray`, or `list`, *optional*): + Parameters shaped `(batch, num_gaussians, 14)`, or one such batch per requested density. `None` for latent + output. Columns contain xyz position, degree-zero SH color, scale, wxyz rotation, and opacity. + latents (`torch.Tensor`, *optional*): + Denoised Gaussian latents. + camera_latents (`torch.Tensor`, *optional*): + Denoised camera latents. + preprocessed_images (`list[PIL.Image.Image]`, *optional*): + Cropped RGB images composited onto black canvases. + """ + + gaussians: torch.Tensor | np.ndarray | list[torch.Tensor] | list[np.ndarray] | None + latents: torch.Tensor | None = None + camera_latents: torch.Tensor | None = None + preprocessed_images: list[Image.Image] | None = None diff --git a/src/diffusers/pipelines/triposplat/pipeline_triposplat.py b/src/diffusers/pipelines/triposplat/pipeline_triposplat.py new file mode 100644 index 000000000000..e84b7e8978c9 --- /dev/null +++ b/src/diffusers/pipelines/triposplat/pipeline_triposplat.py @@ -0,0 +1,376 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +from typing import Callable + +import numpy as np +import torch +from PIL import Image +from torch.nn import functional as F +from transformers import DINOv3ViTModel + +from ...image_processor import PipelineImageInput, TripoSplatImageProcessor +from ...models import AutoencoderKLFlux2 +from ...models.autoencoders.autoencoder_triposplat import TripoSplatGaussianDecoder +from ...models.transformers.transformer_triposplat import TripoSplatTransformer3DModel +from ...schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from ...utils import replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .modeling_birefnet import BiRefNetModel +from .pipeline_output import TripoSplatPipelineOutput + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from PIL import Image + >>> from diffusers import TripoSplatPipeline + >>> from diffusers.utils import export_to_gaussian_ply + + >>> dtype = {"default": torch.float16, "image_encoder": torch.bfloat16, "vae": torch.bfloat16} + >>> pipe = TripoSplatPipeline.from_pretrained("./triposplat-diffusers", dtype=dtype).to("cuda") + >>> image = Image.open("object.png") + >>> output = pipe(image, num_gaussians=32768, generator=torch.Generator("cuda").manual_seed(42)) + >>> export_to_gaussian_ply(output.gaussians[0], "object.ply") + ``` +""" + + +class TripoSplatPipeline(DiffusionPipeline): + """Generate 3D Gaussian splats from foreground images. + + Args: + transformer (`TripoSplatTransformer3DModel`): + Image-conditioned Gaussian and camera flow transformer. + decoder (`TripoSplatGaussianDecoder`): + Octree and Gaussian parameter decoder. + image_encoder (`transformers.DINOv3ViTModel`): + DINOv3 vision encoder. + vae (`AutoencoderKLFlux2`): + Image VAE used to produce packed conditioning latents. + scheduler (`FlowMatchEulerDiscreteScheduler`): + Euler flow scheduler with the reference sigma shift. + background_remover (`BiRefNetModel`, *optional*): + Foreground segmentation model for RGB inputs. + canvas_size (`int`, defaults to `1024`): + Side length of prepared input images. + """ + + model_cpu_offload_seq = "background_remover->image_encoder->vae->transformer->decoder" + _optional_components = ["background_remover"] + _callback_tensor_inputs = ["latents", "camera_latents", "encoder_hidden_states", "image_latents"] + + def __init__( + self, + transformer: TripoSplatTransformer3DModel, + decoder: TripoSplatGaussianDecoder, + image_encoder: DINOv3ViTModel, + vae: AutoencoderKLFlux2, + scheduler: FlowMatchEulerDiscreteScheduler, + background_remover: BiRefNetModel | None = None, + canvas_size: int = 1024, + ) -> None: + super().__init__() + self.register_modules( + transformer=transformer, + decoder=decoder, + image_encoder=image_encoder, + vae=vae, + scheduler=scheduler, + background_remover=background_remover, + ) + self.register_to_config(canvas_size=canvas_size) + self.image_processor = TripoSplatImageProcessor(canvas_size=canvas_size) + self.num_prefix_tokens = ( + self.image_encoder.config.num_register_tokens + 1 if getattr(self, "image_encoder", None) else 5 + ) + self.q_token_length = self.transformer.config.q_token_length if getattr(self, "transformer", None) else 8192 + self.latent_channels = self.transformer.config.in_channels if getattr(self, "transformer", None) else 16 + self.camera_channels = self.transformer.config.cam_channels if getattr(self, "transformer", None) else 5 + self.gaussians_per_point = self.decoder.config.gaussians_per_point if getattr(self, "decoder", None) else 32 + + def check_inputs( + self, + num_images_per_prompt: int, + num_gaussians: int | list[int], + output_type: str, + callback_on_step_end_tensor_inputs: list[str], + ) -> None: + """Validate output, batching, density, and callback arguments.""" + if output_type not in ("pt", "np", "latent"): + raise ValueError("output_type must be 'pt', 'np', or 'latent'.") + if not isinstance(num_images_per_prompt, int) or num_images_per_prompt < 1: + raise ValueError("num_images_per_prompt must be a positive integer.") + if not set(callback_on_step_end_tensor_inputs).issubset(self._callback_tensor_inputs): + raise ValueError(f"Callback inputs must be drawn from {self._callback_tensor_inputs}.") + counts = num_gaussians if isinstance(num_gaussians, list) else [num_gaussians] + if not counts or any(not isinstance(count, int) or not 32768 <= count <= 262144 for count in counts): + raise ValueError( + "num_gaussians must be an integer or a nonempty list of integers between 32768 and 262144." + ) + + def prepare_image( + self, image: PipelineImageInput, erode_radius: int = 1, is_preprocessed: bool = False + ) -> list[Image.Image]: + """Prepare foreground images, predicting alpha masks for RGB inputs when BiRefNet is supplied.""" + images = self.image_processor.to_pil(image) + if not is_preprocessed: + images = self.image_processor.resize_shortest_edge(images) + for index, item in enumerate(images): + if item.mode == "RGBA" and item.getchannel("A").getextrema()[0] < 255: + continue + if self.background_remover is None: + raise ValueError( + "RGB inputs require background_remover, or set is_preprocessed=True for prepared RGB images." + ) + pixels = self.image_processor.preprocess([item.convert("RGB")]) + pixels = F.interpolate( + pixels, size=(self.background_remover.config.sample_size,) * 2, mode="bilinear", align_corners=True + ) + mean = pixels.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = pixels.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + pixels = ((pixels - mean) / std).to(device=self._execution_device, dtype=self.background_remover.dtype) + alpha = self.background_remover(pixels).sample + alpha = F.interpolate( + alpha.float(), size=(item.height, item.width), mode="bilinear", align_corners=True + )[0, 0] + alpha = (alpha.clamp(0, 1) * 255).to(torch.uint8).cpu().numpy() + item = item.convert("RGB") + item.putalpha(Image.fromarray(alpha)) + images[index] = item + return self.image_processor.prepare_foreground( + images, erode_radius=erode_radius, is_preprocessed=is_preprocessed + ) + + def encode_image(self, images: list[Image.Image], device: torch.device | str) -> torch.Tensor: + """Encode prepared RGB images into normalized DINOv3 features.""" + pixels = self.image_processor.preprocess(images).to(device=device) + mean = pixels.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = pixels.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + pixels = ((pixels - mean) / std).to(self.image_encoder.dtype) + original_inv_freq = self.image_encoder.rope_embeddings.inv_freq + try: + self.image_encoder.rope_embeddings.inv_freq = original_inv_freq.to(dtype=self.image_encoder.dtype) + features = self.image_encoder(pixel_values=pixels).last_hidden_state + finally: + self.image_encoder.rope_embeddings.inv_freq = original_inv_freq + return F.layer_norm(features.float(), features.shape[-1:]) + + def encode_vae_image( + self, + images: list[Image.Image], + device: torch.device | str, + generator: torch.Generator | list[torch.Generator] | None = None, + ) -> torch.Tensor: + """Sample and normalize packed VAE conditioning latents for each input image.""" + if isinstance(generator, list): + if len(generator) % len(images): + raise ValueError("The generator list must contain one generator per image or generated sample.") + generator = generator[:: len(generator) // len(images)] + pixels = self.image_processor.preprocess(images).to(device=device, dtype=self.vae.dtype) + pixels = self.image_processor.normalize(pixels) + parameters = self.vae.encode(pixels).latent_dist.parameters + # TripoSplat samples the posterior variance without the distribution's log-variance clamp. + mean, logvar = parameters.chunk(2, dim=1) + noise = randn_tensor(mean.shape, generator=generator, device=device, dtype=mean.dtype) + latents = mean + torch.exp(0.5 * logvar) * noise + batch, channels, height, width = latents.shape + latents = latents.view(batch, channels, height // 2, 2, width // 2, 2).permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch, channels * 4, height // 2, width // 2) + bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.bn.eps).to( + latents.device, latents.dtype + ) + latents = ((latents - bn_mean) / bn_std).float().flatten(2).transpose(1, 2).contiguous() + return F.pad(latents, (0, 0, self.num_prefix_tokens, 0)) + + def prepare_latents( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device | str, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + camera_latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare Gaussian and camera noise for the denoising loop.""" + shape = (batch_size, self.q_token_length, self.latent_channels) + camera_shape = (batch_size, 1, self.camera_channels) + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + elif tuple(latents.shape) != shape: + raise ValueError(f"Expected latents with shape {shape}, got {tuple(latents.shape)}.") + if camera_latents is None: + camera_latents = randn_tensor(camera_shape, generator=generator, device=device, dtype=dtype) + elif tuple(camera_latents.shape) != camera_shape: + raise ValueError(f"Expected camera_latents with shape {camera_shape}, got {tuple(camera_latents.shape)}.") + return latents.to(device=device, dtype=dtype), camera_latents.to(device=device, dtype=dtype) + + @property + def num_timesteps(self) -> int: + return self._num_timesteps + + @property + def guidance_scale(self) -> float: + return self._guidance_scale + + @property + def do_classifier_free_guidance(self) -> bool: + return self.guidance_scale > 1 + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + image: PipelineImageInput, + num_inference_steps: int = 20, + guidance_scale: float = 3.0, + num_gaussians: int | list[int] = 262144, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + decoder_generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + camera_latents: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + image_latents: torch.Tensor | None = None, + erode_radius: int = 1, + is_preprocessed: bool = False, + output_type: str = "pt", + return_dict: bool = True, + callback_on_step_end: Callable | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + ) -> TripoSplatPipelineOutput | tuple: + """ + Args: + image (`PipelineImageInput`): + Input image or batch. Tensors use channel-first layout; arrays use channel-last layout. + num_inference_steps (`int`, defaults to `20`): + Number of Euler denoising steps. + guidance_scale (`float`, defaults to `3.0`): + Classifier-free guidance scale. Values at or below one use only the conditional prediction. + num_gaussians (`int` or `list[int]`, defaults to `262144`): + Gaussian counts between 32768 and 262144, rounded to multiples of the decoder's Gaussians per point. + num_images_per_prompt (`int`, defaults to `1`): + Number of generated objects per input image. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Generator or one generator per generated sample for VAE sampling, initial noise, and decoding. + decoder_generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Override for octree sampling and point jitter. Uses `generator` when omitted. + latents (`torch.Tensor`, *optional*): + Initial Gaussian noise. Generated when omitted. + camera_latents (`torch.Tensor`, *optional*): + Initial camera noise. Generated when omitted. + encoder_hidden_states (`torch.Tensor`, *optional*): + Precomputed DINOv3 features, before expansion by `num_images_per_prompt`. + image_latents (`torch.Tensor`, *optional*): + Precomputed packed VAE features, before expansion by `num_images_per_prompt`. + erode_radius (`int`, defaults to `1`): + Foreground mask erosion radius. + is_preprocessed (`bool`, defaults to `False`): + Whether inputs are already cropped RGB images composited on black. + output_type (`str`, defaults to `"pt"`): + Gaussian output format: `"pt"` or `"np"`. Use `"latent"` to skip decoding. + return_dict (`bool`, defaults to `True`): + Whether to return a structured output or a tuple. + callback_on_step_end (`Callable`, *optional*): + Called with the pipeline, step index, timestep, and requested tensors after each denoising step. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): + Tensor names to pass to the callback. + + Returns: + `TripoSplatPipelineOutput` or `tuple`: Gaussian parameters, denoised latents, camera latents, and prepared + input images. A list of densities returns one batched Gaussian tensor per density. + + Examples: + """ + self.check_inputs(num_images_per_prompt, num_gaussians, output_type, callback_on_step_end_tensor_inputs) + counts = num_gaussians if isinstance(num_gaussians, list) else [num_gaussians] + counts = [round(count / self.gaussians_per_point) * self.gaussians_per_point for count in counts] + self._guidance_scale = guidance_scale + device = self._execution_device + images = self.prepare_image(image, erode_radius, is_preprocessed) + batch_size = len(images) * num_images_per_prompt + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError("Pass one generator per generated sample.") + if encoder_hidden_states is None: + encoder_hidden_states = self.encode_image(images, device) + if image_latents is None: + image_latents = self.encode_vae_image(images, device, generator) + if encoder_hidden_states.shape[:2] != image_latents.shape[:2] or encoder_hidden_states.shape[0] != len(images): + raise ValueError("The DINO and VAE features must have matching batch and token dimensions.") + encoder_hidden_states = encoder_hidden_states.to(device).repeat_interleave(num_images_per_prompt, dim=0) + image_latents = image_latents.to(device).repeat_interleave(num_images_per_prompt, dim=0) + latents, camera_latents = self.prepare_latents( + batch_size, torch.float32, device, generator, latents, camera_latents + ) + sigmas = np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1] + self.scheduler.set_timesteps(sigmas=sigmas, device=device) + self._num_timesteps = len(self.scheduler.timesteps) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for index, timestep in enumerate(self.scheduler.timesteps): + noise_pred, camera_pred = self.transformer( + latents, + timestep.expand(batch_size), + encoder_hidden_states, + image_latents, + camera_latents, + return_dict=False, + ) + if self.do_classifier_free_guidance: + negative_pred, negative_camera = self.transformer( + latents, + timestep.expand(batch_size), + torch.zeros_like(encoder_hidden_states), + torch.zeros_like(image_latents), + camera_latents, + return_dict=False, + ) + noise_pred = self.guidance_scale * noise_pred - (self.guidance_scale - 1) * negative_pred + camera_pred = self.guidance_scale * camera_pred - (self.guidance_scale - 1) * negative_camera + latent_size = latents[0].numel() + prediction = torch.cat([noise_pred.flatten(1), camera_pred.flatten(1)], dim=1) + sample = torch.cat([latents.flatten(1), camera_latents.flatten(1)], dim=1) + sample = self.scheduler.step(prediction.float(), timestep, sample, return_dict=False)[0] + latents = sample[:, :latent_size].reshape_as(latents) + camera_latents = sample[:, latent_size:].reshape_as(camera_latents) + if callback_on_step_end is not None: + tensors = { + "latents": latents, + "camera_latents": camera_latents, + "encoder_hidden_states": encoder_hidden_states, + "image_latents": image_latents, + } + updates = callback_on_step_end( + self, index, timestep, {name: tensors[name] for name in callback_on_step_end_tensor_inputs} + ) + latents = updates.pop("latents", latents) + camera_latents = updates.pop("camera_latents", camera_latents) + encoder_hidden_states = updates.pop("encoder_hidden_states", encoder_hidden_states) + image_latents = updates.pop("image_latents", image_latents) + progress_bar.update() + gaussians = None + if output_type != "latent": + decoder_generator = generator if decoder_generator is None else decoder_generator + decoded = [self.decoder(latents, count, generator=decoder_generator).sample for count in counts] + if output_type == "np": + decoded = [item.cpu().numpy() for item in decoded] + gaussians = decoded if isinstance(num_gaussians, list) else decoded[0] + self.maybe_free_model_hooks() + output = TripoSplatPipelineOutput( + gaussians=gaussians, latents=latents, camera_latents=camera_latents, preprocessed_images=images + ) + return output if return_dict else output.to_tuple() diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 5f89eb318d92..7db3cf34d3f3 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -176,3 +176,7 @@ def check_min_version(min_version): error_message = f"This example requires a minimum version of {min_version}," error_message += f" but the version found is {__version__}.\n" raise ImportError(error_message) + + +if is_torch_available(): + from .gaussian_splat import export_to_gaussian_ply, export_to_splat diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 2d89f6f955aa..810e4920510b 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -182,6 +182,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class TripoSplatClassifierFreeGuidance(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class FasterCacheConfig(metaclass=DummyObject): _backends = ["torch"] @@ -409,6 +424,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class TripoSplatImageProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class VaeImageProcessor(metaclass=DummyObject): _backends = ["torch"] @@ -2434,6 +2464,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class TripoSplatGaussianDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class TripoSplatTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class UNet1DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_torchvision_objects.py b/src/diffusers/utils/dummy_torch_and_torchvision_objects.py new file mode 100644 index 000000000000..0bbec2ed2337 --- /dev/null +++ b/src/diffusers/utils/dummy_torch_and_torchvision_objects.py @@ -0,0 +1,17 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class BiRefNetModel(metaclass=DummyObject): + _backends = ["torch", "torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "torchvision"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "torchvision"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "torchvision"]) diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 9ede7a7543f6..97690696a801 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -842,6 +842,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class TripoSplatAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class TripoSplatModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Wan22Blocks(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -5192,6 +5222,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class TripoSplatPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class TripoSplatPipelineOutput(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class UnCLIPImageVariationPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/src/diffusers/utils/gaussian_splat.py b/src/diffusers/utils/gaussian_splat.py new file mode 100644 index 000000000000..dea8f67032db --- /dev/null +++ b/src/diffusers/utils/gaussian_splat.py @@ -0,0 +1,146 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + + +from pathlib import Path + +import numpy as np +import torch + + +def _transform_gaussians(gaussians, transform): + parameters = ( + gaussians.detach().float().cpu().numpy() + if isinstance(gaussians, torch.Tensor) + else np.asarray(gaussians, dtype=np.float32) + ) + if parameters.ndim != 2 or parameters.shape[-1] != 14: + raise ValueError("Pass one Gaussian set with shape (num_gaussians, 14).") + positions = parameters[:, :3] + rotations = parameters[:, 9:13] + transform = np.asarray([[1, 0, 0], [0, 0, -1], [0, 1, 0]] if transform is None else transform, dtype=np.float32) + if transform.shape != (3, 3): + raise ValueError("transform must have shape (3, 3).") + positions = positions @ transform.T + quaternion = rotations / np.linalg.norm(rotations, axis=-1, keepdims=True) + w, x, y, z = quaternion.T + matrix = np.stack( + [ + 1 - 2 * (y * y + z * z), + 2 * (x * y - w * z), + 2 * (x * z + w * y), + 2 * (x * y + w * z), + 1 - 2 * (x * x + z * z), + 2 * (y * z - w * x), + 2 * (x * z - w * y), + 2 * (y * z + w * x), + 1 - 2 * (x * x + y * y), + ], + axis=-1, + ).reshape(-1, 3, 3) + matrix = transform @ matrix + trace = np.trace(matrix, axis1=1, axis2=2) + result = np.zeros_like(rotations) + scale = np.sqrt(np.maximum(trace + 1, 0)) * 2 + result[:, 0] = 0.25 * scale + denominator = np.where(scale != 0, scale, 1) + result[:, 1] = (matrix[:, 2, 1] - matrix[:, 1, 2]) / denominator + result[:, 2] = (matrix[:, 0, 2] - matrix[:, 2, 0]) / denominator + result[:, 3] = (matrix[:, 1, 0] - matrix[:, 0, 1]) / denominator + for axis in range(3): + other1, other2 = (axis + 1) % 3, (axis + 2) % 3 + mask = ( + (scale == 0) + & (matrix[:, axis, axis] >= matrix[:, other1, other1]) + & (matrix[:, axis, axis] >= matrix[:, other2, other2]) + ) + for previous in range(axis): + mask &= matrix[:, axis, axis] > matrix[:, previous, previous] + axis_scale = ( + np.sqrt(np.maximum(1 + matrix[:, axis, axis] - matrix[:, other1, other1] - matrix[:, other2, other2], 0)) + * 2 + ) + axis_denominator = np.where(axis_scale != 0, axis_scale, 1) + result[mask, 0] = (matrix[mask, other2, other1] - matrix[mask, other1, other2]) / axis_denominator[mask] + result[mask, axis + 1] = 0.25 * axis_scale[mask] + result[mask, other1 + 1] = (matrix[mask, other1, axis] + matrix[mask, axis, other1]) / axis_denominator[mask] + result[mask, other2 + 1] = (matrix[mask, other2, axis] + matrix[mask, axis, other2]) / axis_denominator[mask] + result /= np.linalg.norm(result, axis=-1, keepdims=True) + return parameters, positions, result + + +def export_to_gaussian_ply(gaussians, output_path, transform=None) -> str: + """Save `(num_gaussians, 14)` parameters as a binary Gaussian PLY file. + + Parameters contain xyz, SH DC color, physical scale, wxyz rotation, and opacity. The default transform maps xyz to + `(x, -z, y)`. Pass a 3-by-3 rotation matrix to use a different orientation. + """ + parameters, positions, rotations = _transform_gaussians(gaussians, transform) + colors = parameters[:, 3:6] + opacity = parameters[:, 13:14] + scales = parameters[:, 6:9] + with np.errstate(divide="ignore"): + logits = np.log(opacity / (1 - opacity)) + values = np.concatenate([positions, np.zeros_like(positions), colors, logits, np.log(scales), rotations], axis=-1) + names = [ + "x", + "y", + "z", + "nx", + "ny", + "nz", + "f_dc_0", + "f_dc_1", + "f_dc_2", + "opacity", + "scale_0", + "scale_1", + "scale_2", + "rot_0", + "rot_1", + "rot_2", + "rot_3", + ] + header = f"ply\nformat binary_little_endian 1.0\nelement vertex {len(values)}\n" + header += "".join(f"property float {name}\n" for name in names) + "end_header\n" + output_path = Path(output_path) + output_path.write_bytes(header.encode("ascii") + values.astype(" str: + """Save `(num_gaussians, 14)` parameters as 32-byte records for Gaussian splat viewers. + + Parameters contain xyz, SH DC color, physical scale, wxyz rotation, and opacity. The default transform maps xyz to + `(x, -z, y)`. Pass a 3-by-3 rotation matrix to use a different orientation. + """ + parameters, positions, rotations = _transform_gaussians(gaussians, transform) + scales = parameters[:, 6:9] + opacity = parameters[:, 13:14] + rgb = np.clip((parameters[:, 3:6] * 0.28209479177387814 + 0.5) * 255, 0, 255).astype(np.uint8) + rgba = np.concatenate([rgb, np.clip(opacity * 255, 0, 255).astype(np.uint8)], axis=-1) + rotation_bytes = np.clip(rotations * 128 + 128, 0, 255).astype(np.uint8) + order = np.argsort(-opacity[:, 0] * np.prod(scales, axis=-1)) + records = np.concatenate( + [ + positions.astype(" dict[str, int | list[int]]: + return { + "model_channels": 32, + "cond_channels": 16, + "num_octree_blocks": 1, + "num_gaussian_blocks": 1, + "num_heads": 2, + "mlp_ratio": 2.0, + "gaussians_per_point": 32, + "max_voxel_level": 2, + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 16, 16), generator=self.generator, device=torch_device), + "num_gaussians": 64, + "generator": self.generator, + } + + @property + def input_shape(self) -> tuple[int, ...]: + return (16, 16) + + @property + def output_shape(self) -> tuple[int, ...]: + return (64, 14) + + +class TestTripoSplatGaussianDecoderModel(TripoSplatGaussianDecoderTesterConfig, ModelTesterMixin): + pass + + +class TestTripoSplatGaussianDecoderMemory(TripoSplatGaussianDecoderTesterConfig, MemoryTesterMixin): + def get_dummy_inputs(self): + inputs = super().get_dummy_inputs() + inputs["generator"] = torch.manual_seed(0) + return inputs + + +class TestTripoSplatGaussianDecoderTorchCompile(TripoSplatGaussianDecoderTesterConfig, TorchCompileTesterMixin): + @pytest.mark.skip(reason="Adaptive octree sampling uses data-dependent tensor sizes; compile the repeated blocks.") + def test_torch_compile_recompilation_and_graph_break(self): + pass + + @pytest.mark.skip(reason="Adaptive octree sampling cannot be exported as one static graph.") + def test_compile_works_with_aot(self, tmp_path): + pass + + @pytest.mark.skip(reason="Adaptive octree sampling uses data-dependent tensor sizes; compile the repeated blocks.") + def test_compile_on_different_shapes(self): + pass + + +class TestTripoSplatGaussianDecoderAttention(TripoSplatGaussianDecoderTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/models/test_models_birefnet.py b/tests/models/test_models_birefnet.py new file mode 100644 index 000000000000..c1517008d51d --- /dev/null +++ b/tests/models/test_models_birefnet.py @@ -0,0 +1,95 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import pytest +import torch + +from diffusers import BiRefNetModel +from diffusers.utils.torch_utils import randn_tensor + +from ..testing_utils import enable_full_determinism, torch_device +from .testing_utils import ( + AttentionTesterMixin, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) + + +enable_full_determinism() + + +class BiRefNetTesterConfig: + main_input_name = "sample" + + @property + def model_class(self): + return BiRefNetModel + + @property + def pretrained_model_name_or_path(self): + return None + + @property + def pretrained_model_kwargs(self): + return {} + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + return {"embed_dim": 8, "depths": (1, 1, 1, 1), "num_heads": (1, 1, 2, 4), "window_size": 2, "sample_size": 64} + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return {"sample": randn_tensor((1, 3, 64, 64), generator=self.generator, device=torch_device)} + + @property + def input_shape(self) -> tuple[int, ...]: + return (3, 64, 64) + + @property + def output_shape(self) -> tuple[int, ...]: + return (1, 64, 64) + + +class TestBiRefNetModel(BiRefNetTesterConfig, ModelTesterMixin): + def test_torchvision_dependency_error(self, monkeypatch): + from diffusers.utils import import_utils + + monkeypatch.setitem( + import_utils.BACKENDS_MAPPING, "torchvision", (lambda: False, import_utils.TORCHVISION_IMPORT_ERROR) + ) + with pytest.raises(ImportError, match="requires the torchvision library"): + BiRefNetModel(**self.get_init_dict()) + + +class TestBiRefNetMemory(BiRefNetTesterConfig, MemoryTesterMixin): + @pytest.mark.skip(reason="This BiRefNet adapter contains inference heads and is not intended for training.") + def test_layerwise_casting_training(self): + pass + + +class TestBiRefNetTorchCompile(BiRefNetTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(64, 64), (64, 128), (128, 128)] + + def get_dummy_inputs(self, height: int = 64, width: int = 64) -> dict[str, torch.Tensor]: + return {"sample": randn_tensor((1, 3, height, width), generator=self.generator, device=torch_device)} + + +class TestBiRefNetAttention(BiRefNetTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/models/transformers/test_models_transformer_triposplat.py b/tests/models/transformers/test_models_transformer_triposplat.py new file mode 100644 index 000000000000..31732657cca9 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_triposplat.py @@ -0,0 +1,116 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import torch + +from diffusers import TripoSplatTransformer3DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class TripoSplatTransformerTesterConfig: + main_input_name = "hidden_states" + + @property + def model_class(self): + return TripoSplatTransformer3DModel + + @property + def pretrained_model_name_or_path(self): + return None + + @property + def pretrained_model_kwargs(self): + return {} + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + return { + "q_token_length": 16, + "in_channels": 16, + "out_channels": 16, + "cam_channels": 5, + "model_channels": 32, + "cond_channels": 32, + "cond2_channels": 8, + "num_refiner_blocks": 1, + "num_blocks": 1, + "num_heads": 2, + "mlp_ratio": 2.0, + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 16, 16), generator=self.generator, device=torch_device), + "timestep": torch.tensor([500.0], device=torch_device), + "encoder_hidden_states": randn_tensor((1, 9, 32), generator=self.generator, device=torch_device), + "image_latents": randn_tensor((1, 9, 8), generator=self.generator, device=torch_device), + "camera_latents": randn_tensor((1, 1, 5), generator=self.generator, device=torch_device), + } + + @property + def input_shape(self) -> tuple[int, ...]: + return (16, 16) + + @property + def output_shape(self) -> tuple[int, ...]: + return (16, 16) + + +class TestTripoSplatTransformerModel(TripoSplatTransformerTesterConfig, ModelTesterMixin): + pass + + +class TestTripoSplatTransformerMemory(TripoSplatTransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestTripoSplatTransformerTorchCompile(TripoSplatTransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + return { + "hidden_states": randn_tensor((1, 16, 16), generator=self.generator, device=torch_device), + "timestep": torch.tensor([500.0], device=torch_device), + "encoder_hidden_states": randn_tensor( + (1, height * width, 32), generator=self.generator, device=torch_device + ), + "image_latents": randn_tensor((1, height * width, 8), generator=self.generator, device=torch_device), + "camera_latents": randn_tensor((1, 1, 5), generator=self.generator, device=torch_device), + } + + +class TestTripoSplatTransformerTraining(TripoSplatTransformerTesterConfig, TrainingTesterMixin): + pass + + +class TestTripoSplatTransformerAttention(TripoSplatTransformerTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/modular_pipelines/triposplat/__init__.py b/tests/modular_pipelines/triposplat/__init__.py new file mode 100644 index 000000000000..8b137891791f --- /dev/null +++ b/tests/modular_pipelines/triposplat/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/modular_pipelines/triposplat/test_modular_pipeline_triposplat.py b/tests/modular_pipelines/triposplat/test_modular_pipeline_triposplat.py new file mode 100644 index 000000000000..d2a2dc923fca --- /dev/null +++ b/tests/modular_pipelines/triposplat/test_modular_pipeline_triposplat.py @@ -0,0 +1,225 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import pytest +import torch +from PIL import Image + +from diffusers import ( + ComponentsManager, + ModularPipeline, + TripoSplatAutoBlocks, + TripoSplatClassifierFreeGuidance, + TripoSplatImageProcessor, + TripoSplatModularPipeline, + TripoSplatPipeline, +) + +from ...pipelines.triposplat.testing_utils import get_triposplat_dummy_components +from ...testing_utils import require_accelerator, torch_device +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularGuiderTesterMixin, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, +) + + +MODULAR_PROCESSOR_CONFIG_XFAIL = pytest.mark.xfail( + strict=True, reason="ModularPipeline does not serialize updated from_config processor components." +) + + +class TripoSplatModularPipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = TripoSplatModularPipeline + pipeline_blocks_class = TripoSplatAutoBlocks + params = frozenset(["image", "num_gaussians", "decoder_generator", "is_preprocessed"]) + batch_params = frozenset(["image"]) + output_name = "gaussians" + expected_workflow_blocks = { + None: [ + ("preprocess", "TripoSplatImagePreprocessStep"), + ("image_encoder", "TripoSplatImageEncoderStep"), + ("vae_encoder", "TripoSplatVaeEncoderStep"), + ("denoise", "TripoSplatCoreDenoiseStep"), + ("decode", "TripoSplatGaussianDecodeStep"), + ] + } + + @pytest.fixture(scope="class", autouse=True) + def tiny_checkpoint(self, tmp_path_factory, request): + path = tmp_path_factory.mktemp("triposplat") + standard = TripoSplatPipeline(**get_triposplat_dummy_components()) + standard.save_pretrained(path) + modular = TripoSplatAutoBlocks().init_pipeline(str(path)) + modular.update_components(**standard.components, image_processor=TripoSplatImageProcessor(canvas_size=32)) + modular.save_pretrained(str(path), overwrite_modular_index=True) + request.cls._checkpoint_path = str(path) + + @property + def pretrained_model_name_or_path(self): + return self._checkpoint_path + + def get_pipeline(self, components_manager=None, dtype=torch.float32): + pipe = super().get_pipeline(components_manager=components_manager, dtype=dtype) + pipe.update_components(image_processor=TripoSplatImageProcessor(canvas_size=32)) + return pipe + + def get_dummy_inputs(self, seed=0): + return { + "image": Image.new("RGB", (32, 32), (128, 64, 192)), + "is_preprocessed": True, + "num_inference_steps": 2, + "num_gaussians": 32768, + "output_type": "pt", + "generator": self.get_generator(seed), + } + + +class TestTripoSplatModularPipeline(TripoSplatModularPipelineTesterConfig, ModularPipelineTesterMixin): + def test_default_blocks(self): + blocks = self.pipeline_blocks_class() + actual = [(name, type(block).__name__) for name, block in blocks.sub_blocks.items()] + assert actual == self.expected_workflow_blocks[None] + + def test_float32_latents_with_half_precision_transformer(self): + components = get_triposplat_dummy_components() + components["transformer"].to(dtype=torch.float16) + pipe = TripoSplatAutoBlocks().sub_blocks["denoise"].init_pipeline() + pipe.update_components(transformer=components["transformer"], scheduler=components["scheduler"]) + pipe.to(torch_device) + output = pipe( + encoder_hidden_states=torch.randn(1, 9, 32, generator=self.get_generator(0), device=torch_device), + image_latents=torch.randn(1, 9, 8, generator=self.get_generator(1), device=torch_device), + generator=self.get_generator(2), + num_inference_steps=2, + ) + assert output.get("latents").dtype == torch.float32 + assert output.get("camera_latents").dtype == torch.float32 + assert output.get("latents").isfinite().all() + assert output.get("camera_latents").isfinite().all() + + def test_standalone_decode(self): + blocks = TripoSplatAutoBlocks() + pipe = blocks.sub_blocks["decode"].init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components() + latents = torch.randn(1, 16, 16, generator=self.get_generator(0)) + output = pipe(latents=latents, num_gaussians=32768, decoder_generator=self.get_generator(1)) + assert output.get("gaussians").shape == (1, 32768, 14) + + @pytest.mark.parametrize( + "block_name,output_name", [("image_encoder", "encoder_hidden_states"), ("vae_encoder", "image_latents")] + ) + def test_standalone_encoder(self, block_name, output_name): + pipe = self.pipeline_blocks_class().sub_blocks[block_name].init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=torch.float32) + pipe.update_components(image_processor=TripoSplatImageProcessor(canvas_size=32)) + pipe.to(torch_device) + images = [self.get_dummy_inputs()["image"]] + actual = pipe(preprocessed_images=images, generator=self.get_generator(0), output=output_name) + standard = TripoSplatPipeline(**get_triposplat_dummy_components()).to(torch_device) + if block_name == "image_encoder": + expected = standard.encode_image(images, torch_device) + else: + expected = standard.encode_vae_image(images, torch_device, self.get_generator(0)) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + def test_standalone_preprocess(self): + pipe = self.pipeline_blocks_class().sub_blocks["preprocess"].init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components() + pipe.update_components(image_processor=TripoSplatImageProcessor(canvas_size=32)) + image = Image.new("RGBA", (16, 32), (0, 0, 0, 0)) + image.paste((128, 64, 192, 255), (4, 8, 12, 24)) + actual = pipe(image=image, erode_radius=0, output="preprocessed_images") + expected = TripoSplatImageProcessor(canvas_size=32).prepare_foreground(image, erode_radius=0) + assert len(actual) == len(expected) == 1 + assert actual[0].tobytes() == expected[0].tobytes() + + @pytest.mark.parametrize("batch_size", [1, 2]) + def test_standalone_prepare_latents(self, batch_size): + block = self.pipeline_blocks_class().sub_blocks["denoise"].sub_blocks["prepare"] + pipe = block.init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=torch.float32) + pipe.to(torch_device) + conditioning = torch.zeros(batch_size, 9, 32, device=torch_device) + actual = pipe( + encoder_hidden_states=conditioning, generator=[self.get_generator(index) for index in range(batch_size)] + ) + standard = TripoSplatPipeline(**get_triposplat_dummy_components()).to(torch_device) + expected_latents, expected_camera = standard.prepare_latents( + batch_size, + torch.float32, + torch_device, + generator=[self.get_generator(index) for index in range(batch_size)], + ) + torch.testing.assert_close(actual.get("latents"), expected_latents, rtol=0, atol=0) + torch.testing.assert_close(actual.get("camera_latents"), expected_camera, rtol=0, atol=0) + reused = pipe( + encoder_hidden_states=conditioning, + latents=actual.get("latents"), + camera_latents=actual.get("camera_latents"), + ) + torch.testing.assert_close(reused.get("latents"), expected_latents, rtol=0, atol=0) + torch.testing.assert_close(reused.get("camera_latents"), expected_camera, rtol=0, atol=0) + + +class TestTripoSplatModularLoading(TripoSplatModularPipelineTesterConfig, ModularLoadingTesterMixin): + @MODULAR_PROCESSOR_CONFIG_XFAIL + def test_save_from_pretrained(self, tmp_path, base_pipe_output): + pipe = self.get_pipeline() + pipe.save_pretrained(str(tmp_path), overwrite_modular_index=True) + reloaded = ModularPipeline.from_pretrained(str(tmp_path)) + reloaded.load_components(dtype=torch.float32) + assert reloaded.image_processor.config == pipe.image_processor.config + reloaded.to(torch_device) + actual = reloaded(**self.get_dummy_inputs(), output=self.output_name) + torch.testing.assert_close(actual, base_pipe_output, atol=1e-5, rtol=1e-5) + + def test_optional_component_is_not_loaded(self): + pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path) + pipe.load_components() + assert pipe.background_remover is None + + def test_base_pipeline_dispatch(self, base_pipe_output): + pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path) + assert isinstance(pipe, self.pipeline_class) + pipe.load_components(dtype=torch.float32) + pipe.update_components(image_processor=TripoSplatImageProcessor(canvas_size=32)) + pipe.to(torch_device) + actual = pipe(**self.get_dummy_inputs(), output=self.output_name) + torch.testing.assert_close(actual, base_pipe_output, atol=1e-5, rtol=1e-5) + + +class TestTripoSplatModularMemory(TripoSplatModularPipelineTesterConfig, ModularMemoryTesterMixin): + @require_accelerator + def test_components_auto_cpu_offload_inference_consistent(self, base_pipe_output): + manager = ComponentsManager() + manager.enable_auto_cpu_offload(device=torch_device) + pipe = self.get_pipeline(components_manager=manager) + actual = pipe(**self.get_dummy_inputs(), output=self.output_name) + torch.testing.assert_close(actual, base_pipe_output, atol=1e-5, rtol=1e-5) + + +class TestTripoSplatModularGuider(TripoSplatModularPipelineTesterConfig, ModularGuiderTesterMixin): + @pytest.mark.parametrize("scale", [0.0, 0.5, 1.0, 3.0]) + def test_reference_guidance_arithmetic(self, scale): + guider = TripoSplatClassifierFreeGuidance(guidance_scale=scale) + guider.set_state(step=0, num_inference_steps=20, timestep=torch.tensor(1000.0)) + conditional = torch.randn(2, 11, generator=torch.Generator().manual_seed(0)).half() + unconditional = torch.randn(2, 11, generator=torch.Generator().manual_seed(1)).half() + actual = guider.forward(conditional, unconditional).pred + expected = scale * conditional - (scale - 1) * unconditional if scale > 1 else conditional + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index 05cbb8a27a2d..ae863625fd03 100644 --- a/tests/others/test_image_processor.py +++ b/tests/others/test_image_processor.py @@ -18,7 +18,53 @@ import pytest import torch -from diffusers.image_processor import InpaintProcessor, VaeImageProcessor +from diffusers.image_processor import InpaintProcessor, TripoSplatImageProcessor, VaeImageProcessor + + +class TestTripoSplatImageProcessor: + @pytest.mark.parametrize("input_type", ["pil", "np", "pt"]) + def test_preprocess_preserves_rgb_tensor_contract(self, input_type): + pixels = np.arange(12 * 20 * 3, dtype=np.uint8).reshape(12, 20, 3) + pil_image = PIL.Image.fromarray(pixels) + image = pil_image + if input_type == "np": + image = pixels.astype(np.float32) / 255 + elif input_type == "pt": + image = torch.from_numpy(pixels.astype(np.float32) / 255).permute(2, 0, 1) + processor = TripoSplatImageProcessor(canvas_size=32) + actual = processor.preprocess(image) + expected = VaeImageProcessor(do_resize=False, do_normalize=False).preprocess(pil_image) + assert actual.shape == (1, 3, 12, 20) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + @pytest.mark.parametrize("input_type", ["pil", "np", "pt", "np_uint8", "pt_uint8", "pt_bfloat16"]) + @pytest.mark.parametrize("batched", [False, True]) + def test_prepare_foreground_preserves_alpha_and_input(self, input_type, batched): + image = PIL.Image.new("RGBA", (40, 24), (0, 0, 0, 0)) + image.paste((255, 128, 64, 255), (12, 4, 28, 20)) + pixels = np.asarray(image).copy() + value = image + if input_type.startswith("np"): + value = pixels if input_type == "np_uint8" else pixels.astype(np.float32) / 255 + if batched: + value = np.stack([value, value]) + elif input_type.startswith("pt"): + value = torch.from_numpy(pixels).permute(2, 0, 1) + if input_type != "pt_uint8": + value = value.to(torch.bfloat16 if input_type == "pt_bfloat16" else torch.float32) / 255 + if batched: + value = torch.stack([value, value]) + elif batched: + value = [image, image] + processor = TripoSplatImageProcessor(canvas_size=32) + expected = processor.prepare_foreground(image, erode_radius=1)[0] + actual = processor.prepare_foreground(value, erode_radius=1) + assert len(actual) == (2 if batched else 1) + for prepared in actual: + assert prepared.mode == "RGB" + assert prepared.size == (32, 32) + np.testing.assert_array_equal(np.asarray(prepared), np.asarray(expected)) + np.testing.assert_array_equal(np.asarray(image), pixels) class TestImageProcessor: diff --git a/tests/pipelines/triposplat/__init__.py b/tests/pipelines/triposplat/__init__.py new file mode 100644 index 000000000000..8b137891791f --- /dev/null +++ b/tests/pipelines/triposplat/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/pipelines/triposplat/test_pipeline_triposplat.py b/tests/pipelines/triposplat/test_pipeline_triposplat.py new file mode 100644 index 000000000000..72c858a4d4d7 --- /dev/null +++ b/tests/pipelines/triposplat/test_pipeline_triposplat.py @@ -0,0 +1,167 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import numpy as np +import pytest +import torch +from PIL import Image + +from diffusers import ( + FlowMatchEulerDiscreteScheduler, + TripoSplatAutoBlocks, + TripoSplatImageProcessor, + TripoSplatPipeline, +) +from diffusers.utils import export_to_gaussian_ply, export_to_splat + +from ...testing_utils import assert_tensors_close +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from .testing_utils import get_triposplat_dummy_components + + +class TripoSplatPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = TripoSplatPipeline + required_input_params_in_call_signature = frozenset(["image", "guidance_scale", "num_gaussians"]) + batch_input_params = frozenset(["image"]) + output_shape = (32768, 14) + + def get_dummy_components(self): + return get_triposplat_dummy_components() + + def get_dummy_inputs(self): + return { + "image": Image.new("RGB", (32, 32), (128, 64, 192)), + "is_preprocessed": True, + "num_inference_steps": 2, + "num_gaussians": 32768, + "guidance_scale": 3.0, + "generator": self.get_generator(0), + "output_type": "pt", + } + + +class TestTripoSplatPipeline(TripoSplatPipelineTesterConfig, PipelineTesterMixin): + def test_gaussian_parameters_and_exports(self, tmp_path): + pipe = self.get_pipeline() + output = pipe(**self.get_dummy_inputs()) + gaussians = output.gaussians[0] + assert gaussians.shape == self.output_shape + assert gaussians.isfinite().all() + assert (gaussians[:, 6:9] > 0).all() + assert ((gaussians[:, 13] >= 0) & (gaussians[:, 13] <= 1)).all() + export_to_gaussian_ply(gaussians, tmp_path / "object.ply") + export_to_splat(gaussians, tmp_path / "object.splat") + assert (tmp_path / "object.ply").read_bytes().startswith(b"ply\nformat binary_little_endian 1.0\n") + assert (tmp_path / "object.splat").stat().st_size == 32 * gaussians.shape[0] + + def test_rgba_and_optional_background_remover(self): + processor = TripoSplatImageProcessor(canvas_size=32) + image = Image.new("RGBA", (32, 32), (0, 0, 0, 0)) + image.paste((255, 128, 64, 255), (8, 8, 24, 24)) + prepared = processor.prepare_foreground(image) + assert prepared[0].mode == "RGB" + assert prepared[0].size == (32, 32) + with pytest.raises(ValueError, match="require background_remover"): + self.get_pipeline().prepare_image(Image.new("RGB", (32, 32))) + with pytest.raises(ValueError, match="empty"): + processor.prepare_foreground(Image.new("RGBA", (32, 32))) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_flow_match_schedule_and_float32_latents(self, dtype): + pipe = self.get_pipeline() + pipe.transformer.to(dtype=dtype) + inputs = self.get_dummy_inputs() + inputs["output_type"] = "latent" + inputs["num_inference_steps"] = 20 + + def callback(pipeline, index, timestep, tensors): + assert tensors["latents"].dtype == torch.float32 + assert tensors["camera_latents"].dtype == torch.float32 + return tensors + + inputs["callback_on_step_end"] = callback + inputs["callback_on_step_end_tensor_inputs"] = ["latents", "camera_latents"] + output = pipe(**inputs) + assert isinstance(pipe.scheduler, FlowMatchEulerDiscreteScheduler) + sigmas = np.linspace(1.0, 0.0, 21) + sigmas = 3.0 * sigmas / (1.0 + 2.0 * sigmas) + np.testing.assert_allclose(pipe.scheduler.sigmas.cpu().numpy(), sigmas, rtol=1e-6, atol=1e-7) + assert output.latents.shape == (1, pipe.q_token_length, pipe.latent_channels) + assert output.camera_latents.shape == (1, 1, pipe.camera_channels) + assert output.latents.dtype == torch.float32 + assert output.camera_latents.dtype == torch.float32 + assert output.latents.isfinite().all() + assert output.camera_latents.isfinite().all() + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_encode_image_restores_shared_rope_buffer(self, dtype): + pipe = self.get_pipeline() + pipe.image_encoder.to(dtype=dtype) + original = pipe.image_encoder.rope_embeddings.inv_freq.float() + pipe.image_encoder.rope_embeddings.inv_freq = original + images = [Image.new("RGB", (32, 32), (128, 64, 192))] + actual = pipe.encode_image(images, pipe._execution_device) + assert actual.isfinite().all() + assert actual.requires_grad + assert pipe.image_encoder.rope_embeddings.inv_freq is original + assert original.dtype == torch.float32 + + def test_multiple_densities_denoise_once(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + steps = [] + + def callback(pipeline, index, timestep, tensors): + steps.append(index) + return tensors + + inputs.update(num_gaussians=[32768, 65536], callback_on_step_end=callback) + output = pipe(**inputs) + assert len(steps) == inputs["num_inference_steps"] + assert [item.shape[1] for item in output.gaussians] == [32768, 65536] + + @pytest.mark.parametrize("batch_size,num_images_per_prompt", [(1, 1), (2, 2)]) + def test_standard_modular_latent_parity(self, batch_size, num_images_per_prompt): + components = self.get_dummy_components() + standard = TripoSplatPipeline(**components) + blocks = TripoSplatAutoBlocks() + blocks.sub_blocks.pop("decode") + modular = blocks.init_pipeline() + modular.update_components( + **{key: value for key, value in components.items() if key != "canvas_size"}, + image_processor=TripoSplatImageProcessor(canvas_size=32), + ) + + def make_inputs(): + inputs = self.get_dummy_inputs() + inputs["num_images_per_prompt"] = num_images_per_prompt + if batch_size > 1: + inputs["image"] = [inputs["image"]] * batch_size + inputs["generator"] = [self.get_generator(seed) for seed in range(batch_size * num_images_per_prompt)] + return inputs + + inputs = make_inputs() + inputs["output_type"] = "latent" + expected = standard(**inputs) + inputs = make_inputs() + for key in ("guidance_scale", "num_gaussians", "decoder_generator", "output_type"): + inputs.pop(key, None) + actual = modular(**inputs) + assert_tensors_close(actual.get("latents"), expected.latents, atol=1e-5, rtol=1e-5) + assert_tensors_close(actual.get("camera_latents"), expected.camera_latents, atol=1e-5, rtol=1e-5) + + +class TestTripoSplatPipelineMemory(TripoSplatPipelineTesterConfig, MemoryTesterMixin): + pass diff --git a/tests/pipelines/triposplat/testing_utils.py b/tests/pipelines/triposplat/testing_utils.py new file mode 100644 index 000000000000..893c3bee582e --- /dev/null +++ b/tests/pipelines/triposplat/testing_utils.py @@ -0,0 +1,75 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import torch +from transformers import DINOv3ViTConfig, DINOv3ViTModel + +from diffusers import ( + AutoencoderKLFlux2, + FlowMatchEulerDiscreteScheduler, + TripoSplatGaussianDecoder, + TripoSplatTransformer3DModel, +) + + +def get_triposplat_dummy_components(): + torch.manual_seed(0) + transformer = TripoSplatTransformer3DModel( + q_token_length=16, + model_channels=32, + cond_channels=32, + cond2_channels=8, + num_refiner_blocks=1, + num_blocks=1, + num_heads=2, + mlp_ratio=2.0, + ) + torch.manual_seed(0) + decoder = TripoSplatGaussianDecoder( + model_channels=32, num_octree_blocks=1, num_gaussian_blocks=1, num_heads=2, mlp_ratio=2.0, max_voxel_level=2 + ) + torch.manual_seed(0) + image_encoder = DINOv3ViTModel( + DINOv3ViTConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + num_register_tokens=4, + patch_size=16, + use_gated_mlp=True, + hidden_act="silu", + layer_norm_eps=1e-5, + pos_embed_rescale=None, + ) + ) + torch.manual_seed(0) + vae = AutoencoderKLFlux2( + block_out_channels=(4, 8, 8, 8), + layers_per_block=1, + latent_channels=2, + norm_num_groups=1, + sample_size=32, + batch_norm_eps=1e-5, + ) + return { + "transformer": transformer, + "decoder": decoder, + "image_encoder": image_encoder, + "vae": vae, + "scheduler": FlowMatchEulerDiscreteScheduler(shift=3.0), + "background_remover": None, + "canvas_size": 32, + } diff --git a/tests/single_file/test_models_triposplat_single_file.py b/tests/single_file/test_models_triposplat_single_file.py new file mode 100644 index 000000000000..85b5bf87038f --- /dev/null +++ b/tests/single_file/test_models_triposplat_single_file.py @@ -0,0 +1,73 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + +import pytest +import torch +from safetensors.torch import save_file + +from diffusers import BiRefNetModel +from diffusers.utils import is_torchvision_available + +from ..pipelines.triposplat.testing_utils import get_triposplat_dummy_components + + +class TestTripoSplatSingleFile: + @pytest.mark.parametrize( + "component_name", + [ + "transformer", + "decoder", + pytest.param( + "background_remover", + marks=pytest.mark.skipif(not is_torchvision_available(), reason="BiRefNet requires torchvision"), + ), + ], + ) + @pytest.mark.parametrize("original_format", [False, True]) + @pytest.mark.parametrize("low_cpu_mem_usage", [False, True]) + def test_single_file_loading(self, tmp_path, component_name, original_format, low_cpu_mem_usage): + if component_name == "background_remover": + torch.manual_seed(0) + model = BiRefNetModel( + embed_dim=8, depths=(1, 1, 1, 1), num_heads=(1, 1, 2, 4), window_size=2, sample_size=64 + ) + else: + model = get_triposplat_dummy_components()[component_name] + model.eval() + model.save_config(tmp_path) + checkpoint = {} + for key, value in model.state_dict().items(): + if original_format: + if component_name == "transformer": + key = key.replace(".mlp.net.0.proj.", ".mlp.mlp.0.").replace(".mlp.net.2.", ".mlp.mlp.2.") + key = key.replace("cam_refiner.", "cam_refiner.mlp.") + key = key.replace(".attn.to_qkv.", ".attn.qkv.").replace(".attn.to_out.", ".attn.out.") + elif component_name == "decoder": + key = key.replace(".q_norm.", ".q_rms_norm.").replace(".k_norm.", ".k_rms_norm.") + key = key.replace(".mlp.net.0.proj.", ".mlp.mlp.0.").replace(".mlp.net.2.", ".mlp.mlp.2.") + elif key.endswith(".atrous_conv.weight"): + key = key.replace(".atrous_conv.weight", ".atrous_conv.regular_conv.weight") + checkpoint[key] = value + if original_format and component_name == "background_remover": + checkpoint["decoder.conv_ms_spvn_2.weight"] = torch.ones(1, 8, 1, 1) + checkpoint["decoder.gdt_convs_pred_2.0.weight"] = torch.ones(1, 8, 1, 1) + path = tmp_path / "model.safetensors" + save_file(checkpoint, path) + loaded = type(model).from_single_file( + str(path), config=str(tmp_path), dtype=torch.float32, low_cpu_mem_usage=low_cpu_mem_usage + ) + assert not loaded.training + assert set(loaded.state_dict()) == set(model.state_dict()) + for key, value in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], value, rtol=0, atol=0)