Upload 43 files

.gitattributes

@@ -1,35 +1,39 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+demo_files/comp.gif filter=lfs diff=lfs merge=lfs -text
+demo_files/examples/animal_character_2.png filter=lfs diff=lfs merge=lfs -text
+demo_files/examples/animal_character.png filter=lfs diff=lfs merge=lfs -text
+demo_files/teaser.gif filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,13 @@
----
-title: FLUX.1-dev + Captioner
-emoji: 🐨
-colorFrom: blue
-colorTo: indigo
-sdk: gradio
-sdk_version: 4.37.2
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
+---
+title: FLUX.1-dev + Captioner
+emoji: 🐨
+colorFrom: blue
+colorTo: indigo
+sdk: gradio
+sdk_version: 4.37.2
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,154 +1,131 @@
-import spaces
-import gradio as gr
-import torch
-from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
-from diffusers import DiffusionPipeline
-import random
-import numpy as np
-import os
-import subprocess
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
-# Initialize models
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.bfloat16
-
-huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-
-# FLUX.1-dev model
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, token = huggingface_token).to(device)
-
-# Initialize Florence model
-florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
-florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
-
-# Prompt Enhancer
-enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
-
-# Florence caption function
-@spaces.GPU
-def florence_caption(image):
-    # Convert image to PIL if it's not already
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-    
-    inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
-    generated_ids = florence_model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=1024,
-        early_stopping=False,
-        do_sample=False,
-        num_beams=3,
-    )
-    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = florence_processor.post_process_generation(
-        generated_text,
-        task="<MORE_DETAILED_CAPTION>",
-        image_size=(image.width, image.height)
-    )
-    return parsed_answer["<MORE_DETAILED_CAPTION>"]
-
-# Prompt Enhancer function
-def enhance_prompt(input_prompt):
-    result = enhancer_long("Enhance the description: " + input_prompt)
-    enhanced_text = result[0]['summary_text']
-    return enhanced_text
-
-@spaces.GPU(duration=190)
-def process_workflow(image, text_prompt, use_enhancer, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
-    if image is not None:
-        # Convert image to PIL if it's not already
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        
-        prompt = florence_caption(image)
-        print(prompt)
-    else:
-        prompt = text_prompt
-    
-    if use_enhancer:
-        prompt = enhance_prompt(prompt)
-    
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    image = pipe(
-        prompt=prompt,
-        generator=generator,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale
-    ).images[0]
-    
-    return image, prompt, seed
-
-custom_css = """
-.input-group, .output-group {
-    border: 1px solid #e0e0e0;
-    border-radius: 10px;
-    padding: 20px;
-    margin-bottom: 20px;
-    background-color: #f9f9f9;
-}
-.submit-btn {
-    background-color: #2980b9 !important;
-    color: white !important;
-}
-.submit-btn:hover {
-    background-color: #3498db !important;
-}
-"""
-
-title = """<h1 align="center">FLUX.1-dev with Florence-2 Captioner and Prompt Enhancer</h1>
-<p><center>
-<a href="https://huggingface.co/black-forest-labs/FLUX.1-dev" target="_blank">[FLUX.1-dev Model]</a>
-<a href="https://huggingface.co/microsoft/Florence-2-base" target="_blank">[Florence-2 Model]</a>
-<a href="https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance-Long" target="_blank">[Prompt Enhancer Long]</a>
-<p align="center">Create long prompts from images or enhance your short prompts with prompt enhancer</p>
-</center></p>
-"""
-
-with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray")) as demo:
-    gr.HTML(title)
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            with gr.Group(elem_classes="input-group"):
-                input_image = gr.Image(label="Input Image (Florence-2 Captioner)")
-            
-            with gr.Accordion("Advanced Settings", open=False):
-                text_prompt = gr.Textbox(label="Text Prompt (optional, used if no image is uploaded)")
-                use_enhancer = gr.Checkbox(label="Use Prompt Enhancer", value=False)
-                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
-                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
-                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=15, step=0.1, value=3.5)
-                num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=28)
-            
-            generate_btn = gr.Button("Generate Image", elem_classes="submit-btn")
-        
-        with gr.Column(scale=1):
-            with gr.Group(elem_classes="output-group"):
-                output_image = gr.Image(label="Result", elem_id="gallery", show_label=False)
-                final_prompt = gr.Textbox(label="Final Prompt Used")
-                used_seed = gr.Number(label="Seed Used")
-    
-    generate_btn.click(
-        fn=process_workflow,
-        inputs=[
-            input_image, text_prompt, use_enhancer, seed, randomize_seed,
-            width, height, guidance_scale, num_inference_steps
-        ],
-        outputs=[output_image, final_prompt, used_seed]
-    )
-
-demo.launch(debug=True)
+import os
+import tempfile
+import time
+import gradio as gr
+import torch
+from PIL import Image
+from diffusers import FluxPipeline
+from huggingface_hub import hf_hub_download
+from sf3d.system import SF3D
+import sf3d.utils as sf3d_utils
+from gradio_litmodel3d import LitModel3D
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16
+
+torch.backends.cuda.matmul.allow_tf32 = True
+huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+# Set up environment and cache
+cache_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+
+if not os.path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
+
+# Initialize Flux pipeline
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=huggingface_token)
+pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device="cuda", dtype=torch.bfloat16)
+
+# Initialize SF3D model
+sf3d_model = SF3D.from_pretrained(
+    "stabilityai/stable-fast-3d",
+    config_name="config.yaml",
+    weight_name="model.safetensors",
+    token=huggingface_token
+
+)
+sf3d_model.eval().cuda()
+
+# Constants for SF3D
+COND_WIDTH, COND_HEIGHT = 512, 512
+COND_DISTANCE, COND_FOVY_DEG = 1.6, 40
+BACKGROUND_COLOR = [0.5, 0.5, 0.5]
+
+c2w_cond = sf3d_utils.default_cond_c2w(COND_DISTANCE)
+intrinsic, intrinsic_normed_cond = sf3d_utils.create_intrinsic_from_fov_deg(
+    COND_FOVY_DEG, COND_HEIGHT, COND_WIDTH
+)
+
+def generate_image(prompt, height, width, steps, scales, seed):
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        return pipe(
+            prompt=[prompt],
+            generator=torch.Generator().manual_seed(int(seed)),
+            num_inference_steps=int(steps),
+            guidance_scale=float(scales),
+            height=int(height),
+            width=int(width),
+            max_sequence_length=256
+        ).images[0]
+
+def create_batch(input_image: Image.Image) -> dict:
+    img_cond = torch.from_numpy(
+        np.asarray(input_image.resize((COND_WIDTH, COND_HEIGHT))).astype(np.float32) / 255.0
+    ).float().clip(0, 1)
+    mask_cond = img_cond[:, :, -1:]
+    rgb_cond = torch.lerp(
+        torch.tensor(BACKGROUND_COLOR)[None, None, :], img_cond[:, :, :3], mask_cond
+    )
+
+    batch_elem = {
+        "rgb_cond": rgb_cond,
+        "mask_cond": mask_cond,
+        "c2w_cond": c2w_cond.unsqueeze(0),
+        "intrinsic_cond": intrinsic.unsqueeze(0),
+        "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
+    }
+    return {k: v.unsqueeze(0) for k, v in batch_elem.items()}
+
+def generate_3d_model(input_image):
+    with torch.no_grad():
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            model_batch = create_batch(input_image)
+            model_batch = {k: v.cuda() for k, v in model_batch.items()}
+            trimesh_mesh, _ = sf3d_model.generate_mesh(model_batch, 1024)
+            trimesh_mesh = trimesh_mesh[0]
+
+    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".glb")
+    trimesh_mesh.export(tmp_file.name, file_type="glb", include_normals=True)
+    return tmp_file.name
+
+def process_and_generate(prompt, height, width, steps, scales, seed):
+    # Generate image from prompt
+    generated_image = generate_image(prompt, height, width, steps, scales, seed)
+    
+    # Generate 3D model from the image
+    glb_file = generate_3d_model(generated_image)
+    
+    return generated_image, glb_file
+
+# Gradio interface
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Text-to-3D Model Generator")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            prompt = gr.Textbox(label="Your Image Description", lines=3)
+            with gr.Accordion("Advanced Settings", open=False):
+                height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024)
+                width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024)
+                steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
+                scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
+                seed = gr.Number(label="Seed", value=3413, precision=0)
+            
+            generate_btn = gr.Button("Generate 3D Model", variant="primary")
+
+        with gr.Column(scale=4):
+            output_image = gr.Image(label="Generated Image")
+            output_3d = LitModel3D(label="3D Model", clear_color=[0.0, 0.0, 0.0, 0.0])
+
+    generate_btn.click(
+        process_and_generate,
+        inputs=[prompt, height, width, steps, scales, seed],
+        outputs=[output_image, output_3d]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

flux_lora.py

@@ -0,0 +1,109 @@
+import spaces
+import argparse
+import os
+import time
+from os import path
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+
+import gradio as gr
+import torch
+from diffusers import FluxPipeline
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
+class timer:
+    def __init__(self, method_name="timed process"):
+        self.method = method_name
+    def __enter__(self):
+        self.start = time.time()
+        print(f"{self.method} starts")
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end = time.time()
+        print(f"{self.method} took {str(round(end - self.start, 2))}s")
+
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device="cuda", dtype=torch.bfloat16)
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        <div style="text-align: center; max-width: 650px; margin: 0 auto;">
+            <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">Hyper-FLUX-8steps-LoRA</h1>
+            <p style="font-size: 1rem; margin-bottom: 1.5rem;">AutoML team from ByteDance</p>
+        </div>
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            with gr.Group():
+                prompt = gr.Textbox(
+                    label="Your Image Description",
+                    placeholder="E.g., A serene landscape with mountains and a lake at sunset",
+                    lines=3
+                )
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    with gr.Group():
+                        with gr.Row():
+                            height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024)
+                            width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024)
+                        
+                        with gr.Row():
+                            steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
+                            scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
+                        
+                        seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0)
+                
+                generate_btn = gr.Button("Generate Image", variant="primary", scale=1)
+
+        with gr.Column(scale=4):
+            output = gr.Image(label="Your Generated Image")
+    
+    gr.Markdown(
+        """
+        <div style="max-width: 650px; margin: 2rem auto; padding: 1rem; border-radius: 10px; background-color: #f0f0f0;">
+            <h2 style="font-size: 1.5rem; margin-bottom: 1rem;">How to Use</h2>
+            <ol style="padding-left: 1.5rem;">
+                <li>Enter a detailed description of the image you want to create.</li>
+                <li>Adjust advanced settings if desired (tap to expand).</li>
+                <li>Tap "Generate Image" and wait for your creation!</li>
+            </ol>
+            <p style="margin-top: 1rem; font-style: italic;">Tip: Be specific in your description for best results!</p>
+        </div>
+        """
+    )
+
+    @spaces.GPU
+    def process_image(height, width, steps, scales, prompt, seed):
+        global pipe
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+            return pipe(
+                prompt=[prompt],
+                generator=torch.Generator().manual_seed(int(seed)),
+                num_inference_steps=int(steps),
+                guidance_scale=float(scales),
+                height=int(height),
+                width=int(width),
+                max_sequence_length=256
+            ).images[0]
+
+    generate_btn.click(
+        process_image,
+        inputs=[height, width, steps, scales, prompt, seed],
+        outputs=output
+    )
+
+if __name__ == "__main__":
+    demo.launch()

load/tets/160_tets.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f4be37efc604d28d55a1a78c2aabefeeab7e63149f541aa45f9dd858ee35bb9
+size 15408790

requirements.txt

@@ -1,11 +1,19 @@
-spaces
-huggingface_hub
-accelerate
-git+https://github.com/huggingface/diffusers.git
-torch==2.4.0
-torchvision==0.19.0
-transformers==4.42.4
-xformers
-sentencepiece
-timm
-einops
+torch>=2.1.2
+torchvision>=0.16.2
+einops>=0.7.0
+jaxtyping>=0.2.31
+omegaconf>=2.3.0
+transformers>=4.43.3
+slangtorch>=1.2.2
+open_clip_torch>=2.24.0
+trimesh>=4.4.1
+numpy>=1.26.4
+huggingface-hub>=0.23.4
+rembg[gpu]>=2.0.57
+gradio-litmodel3d>=0.0.1
+accelerate
+diffusers>=0.30.0
+invisible_watermark
+xformers
+sentencepiece
+peft

sf3d/box_uv_unwrap.py

@@ -0,0 +1,610 @@
+import math
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+from jaxtyping import Float, Integer
+from torch import Tensor
+
+from sf3d.models.utils import dot, triangle_intersection_2d
+
+
+def _box_assign_vertex_to_cube_face(
+    vertex_positions: Float[Tensor, "Nv 3"],
+    vertex_normals: Float[Tensor, "Nv 3"],
+    triangle_idxs: Integer[Tensor, "Nf 3"],
+    bbox: Float[Tensor, "2 3"],
+) -> Tuple[Float[Tensor, "Nf 3 2"], Integer[Tensor, "Nf 3"]]:
+    # Test to not have a scaled model to fit the space better
+    # bbox_min = bbox[:1].mean(-1, keepdim=True)
+    # bbox_max = bbox[1:].mean(-1, keepdim=True)
+    # v_pos_normalized = (vertex_positions - bbox_min) / (bbox_max - bbox_min)
+
+    # Create a [0, 1] normalized vertex position
+    v_pos_normalized = (vertex_positions - bbox[:1]) / (bbox[1:] - bbox[:1])
+    # And to [-1, 1]
+    v_pos_normalized = 2.0 * v_pos_normalized - 1.0
+
+    # Get all vertex positions for each triangle
+    # Now how do we define to which face the triangle belongs? Mean face pos? Max vertex pos?
+    v0 = v_pos_normalized[triangle_idxs[:, 0]]
+    v1 = v_pos_normalized[triangle_idxs[:, 1]]
+    v2 = v_pos_normalized[triangle_idxs[:, 2]]
+    tri_stack = torch.stack([v0, v1, v2], dim=1)
+
+    vn0 = vertex_normals[triangle_idxs[:, 0]]
+    vn1 = vertex_normals[triangle_idxs[:, 1]]
+    vn2 = vertex_normals[triangle_idxs[:, 2]]
+    tri_stack_nrm = torch.stack([vn0, vn1, vn2], dim=1)
+
+    # Just average the normals per face
+    face_normal = F.normalize(torch.sum(tri_stack_nrm, 1), eps=1e-6, dim=-1)
+
+    # Now decide based on the face normal in which box map we project
+    # abs_x, abs_y, abs_z = tri_stack_nrm.abs().unbind(-1)
+    abs_x, abs_y, abs_z = tri_stack.abs().unbind(-1)
+
+    axis = torch.tensor(
+        [
+            [1, 0, 0],  # 0
+            [-1, 0, 0],  # 1
+            [0, 1, 0],  # 2
+            [0, -1, 0],  # 3
+            [0, 0, 1],  # 4
+            [0, 0, -1],  # 5
+        ],
+        device=face_normal.device,
+        dtype=face_normal.dtype,
+    )
+    face_normal_axis = (face_normal[:, None] * axis[None]).sum(-1)
+    index = face_normal_axis.argmax(-1)
+
+    max_axis, uc, vc = (
+        torch.ones_like(abs_x),
+        torch.zeros_like(tri_stack[..., :1]),
+        torch.zeros_like(tri_stack[..., :1]),
+    )
+    mask_pos_x = index == 0
+    max_axis[mask_pos_x] = abs_x[mask_pos_x]
+    uc[mask_pos_x] = tri_stack[mask_pos_x][..., 1:2]
+    vc[mask_pos_x] = -tri_stack[mask_pos_x][..., -1:]
+
+    mask_neg_x = index == 1
+    max_axis[mask_neg_x] = abs_x[mask_neg_x]
+    uc[mask_neg_x] = tri_stack[mask_neg_x][..., 1:2]
+    vc[mask_neg_x] = -tri_stack[mask_neg_x][..., -1:]
+
+    mask_pos_y = index == 2
+    max_axis[mask_pos_y] = abs_y[mask_pos_y]
+    uc[mask_pos_y] = tri_stack[mask_pos_y][..., 0:1]
+    vc[mask_pos_y] = -tri_stack[mask_pos_y][..., -1:]
+
+    mask_neg_y = index == 3
+    max_axis[mask_neg_y] = abs_y[mask_neg_y]
+    uc[mask_neg_y] = tri_stack[mask_neg_y][..., 0:1]
+    vc[mask_neg_y] = -tri_stack[mask_neg_y][..., -1:]
+
+    mask_pos_z = index == 4
+    max_axis[mask_pos_z] = abs_z[mask_pos_z]
+    uc[mask_pos_z] = tri_stack[mask_pos_z][..., 0:1]
+    vc[mask_pos_z] = tri_stack[mask_pos_z][..., 1:2]
+
+    mask_neg_z = index == 5
+    max_axis[mask_neg_z] = abs_z[mask_neg_z]
+    uc[mask_neg_z] = tri_stack[mask_neg_z][..., 0:1]
+    vc[mask_neg_z] = -tri_stack[mask_neg_z][..., 1:2]
+
+    # UC from [-1, 1] to [0, 1]
+    max_dim_div = max_axis.max(dim=0, keepdims=True).values
+    uc = ((uc[..., 0] / max_dim_div + 1.0) * 0.5).clip(0, 1)
+    vc = ((vc[..., 0] / max_dim_div + 1.0) * 0.5).clip(0, 1)
+
+    uv = torch.stack([uc, vc], dim=-1)
+
+    return uv, index
+
+
+def _assign_faces_uv_to_atlas_index(
+    vertex_positions: Float[Tensor, "Nv 3"],
+    triangle_idxs: Integer[Tensor, "Nf 3"],
+    face_uv: Float[Tensor, "Nf 3 2"],
+    face_index: Integer[Tensor, "Nf 3"],
+) -> Integer[Tensor, "Nf"]:  # noqa: F821
+    triangle_pos = vertex_positions[triangle_idxs]
+    # We need to do perform 3 overlap checks.
+    # The first set is placed in the upper two thirds of the UV atlas.
+    # Conceptually, this is the direct visible surfaces from the each cube side
+    # The second set is placed in the lower thirds and the left half of the UV atlas.
+    # This is the first set of occluded surfaces. They will also be saved in the projected fashion
+    # The third pass finds all non assigned faces. They will be placed in the bottom right half of
+    # the UV atlas in scattered fashion.
+    assign_idx = face_index.clone()
+    for overlap_step in range(3):
+        overlapping_indicator = torch.zeros_like(assign_idx, dtype=torch.bool)
+        for i in range(overlap_step * 6, (overlap_step + 1) * 6):
+            mask = assign_idx == i
+            if not mask.any():
+                continue
+            # Get all elements belonging to the projection face
+            uv_triangle = face_uv[mask]
+            cur_triangle_pos = triangle_pos[mask]
+            # Find the center of the uv coordinates
+            center_uv = uv_triangle.mean(dim=1, keepdim=True)
+            # And also the radius of the triangle
+            uv_triangle_radius = (uv_triangle - center_uv).norm(dim=-1).max(-1).values
+
+            potentially_overlapping_mask = (
+                # Find all close triangles
+                (center_uv[None, ...] - center_uv[:, None]).norm(dim=-1)
+                # Do not select the same element by offseting with an large valued identity matrix
+                + torch.eye(
+                    uv_triangle.shape[0],
+                    device=uv_triangle.device,
+                    dtype=uv_triangle.dtype,
+                ).unsqueeze(-1)
+                * 1000
+            )
+            # Mark all potentially overlapping triangles to reduce the number of triangle intersection tests
+            potentially_overlapping_mask = (
+                potentially_overlapping_mask
+                <= (uv_triangle_radius.view(-1, 1, 1) * 3.0)
+            ).squeeze(-1)
+            overlap_coords = torch.stack(torch.where(potentially_overlapping_mask), -1)
+
+            # Only unique triangles (A|B and B|A should be the same)
+            f = torch.min(overlap_coords, dim=-1).values
+            s = torch.max(overlap_coords, dim=-1).values
+            overlap_coords = torch.unique(torch.stack([f, s], dim=1), dim=0)
+            first, second = overlap_coords.unbind(-1)
+
+            # Get the triangles
+            tri_1 = uv_triangle[first]
+            tri_2 = uv_triangle[second]
+
+            # Perform the actual set with the reduced number of potentially overlapping triangles
+            its = triangle_intersection_2d(tri_1, tri_2, eps=1e-6)
+
+            # So we now need to detect which triangles are the occluded ones.
+            # We always assume the first to be the visible one (the others should move)
+            # In the previous step we use a lexigraphical sort to get the unique pairs
+            # In this we use a sort based on the orthographic projection
+            ax = 0 if i < 2 else 1 if i < 4 else 2
+            use_max = i % 2 == 1
+
+            tri1_c = cur_triangle_pos[first].mean(dim=1)
+            tri2_c = cur_triangle_pos[second].mean(dim=1)
+
+            mark_first = (
+                (tri1_c[..., ax] > tri2_c[..., ax])
+                if use_max
+                else (tri1_c[..., ax] < tri2_c[..., ax])
+            )
+            first[mark_first] = second[mark_first]
+
+            # Lastly the same index can be tested multiple times.
+            # If one marks it as overlapping we keep it marked as such.
+            # We do this by testing if it has been marked at least once.
+            unique_idx, rev_idx = torch.unique(first, return_inverse=True)
+
+            add = torch.zeros_like(unique_idx, dtype=torch.float32)
+            add.index_add_(0, rev_idx, its.float())
+            its_mask = add > 0
+
+            # And fill it in the overlapping indicator
+            idx = torch.where(mask)[0][unique_idx]
+            overlapping_indicator[idx] = its_mask
+
+        # Move the index to the overlap regions (shift by 6)
+        assign_idx[overlapping_indicator] += 6
+
+    # We do not care about the correct face placement after the first 2 slices
+    max_idx = 6 * 2
+    return assign_idx.clamp(0, max_idx)
+
+
+def _find_slice_offset_and_scale(
+    index: Integer[Tensor, "Nf"],  # noqa: F821
+) -> Tuple[
+    Float[Tensor, "Nf"], Float[Tensor, "Nf"], Float[Tensor, "Nf"], Float[Tensor, "Nf"]  # noqa: F821
+]:  # noqa: F821
+    # 6 due to the 6 cube faces
+    off = 1 / 3
+    dupl_off = 1 / 6
+
+    # Here, we need to decide how to pack the textures in the case of overlap
+    def x_offset_calc(x, i):
+        offset_calc = i // 6
+        # Initial coordinates - just 3x2 grid
+        if offset_calc == 0:
+            return off * x
+        else:
+            # Smaller 3x2 grid plus eventual shift to right for
+            # second overlap
+            return dupl_off * x + min(offset_calc - 1, 1) * 0.5
+
+    def y_offset_calc(x, i):
+        offset_calc = i // 6
+        # Initial coordinates - just a 3x2 grid
+        if offset_calc == 0:
+            return off * x
+        else:
+            # Smaller coordinates in the lowest row
+            return dupl_off * x + off * 2
+
+    offset_x = torch.zeros_like(index, dtype=torch.float32)
+    offset_y = torch.zeros_like(index, dtype=torch.float32)
+    offset_x_vals = [0, 1, 2, 0, 1, 2]
+    offset_y_vals = [0, 0, 0, 1, 1, 1]
+    for i in range(index.max().item() + 1):
+        mask = index == i
+        if not mask.any():
+            continue
+        offset_x[mask] = x_offset_calc(offset_x_vals[i % 6], i)
+        offset_y[mask] = y_offset_calc(offset_y_vals[i % 6], i)
+
+    div_x = torch.full_like(index, 6 // 2, dtype=torch.float32)
+    # All overlap elements are saved in half scale
+    div_x[index >= 6] = 6
+    div_y = div_x.clone()  # Same for y
+    # Except for the random overlaps
+    div_x[index >= 12] = 2
+    # But the random overlaps are saved in a large block in the lower thirds
+    div_y[index >= 12] = 3
+
+    return offset_x, offset_y, div_x, div_y
+
+
+def rotation_flip_matrix_2d(
+    rad: float, flip_x: bool = False, flip_y: bool = False
+) -> Float[Tensor, "2 2"]:
+    cos = math.cos(rad)
+    sin = math.sin(rad)
+    rot_mat = torch.tensor([[cos, -sin], [sin, cos]], dtype=torch.float32)
+    flip_mat = torch.tensor(
+        [
+            [-1 if flip_x else 1, 0],
+            [0, -1 if flip_y else 1],
+        ],
+        dtype=torch.float32,
+    )
+
+    return flip_mat @ rot_mat
+
+
+def calculate_tangents(
+    vertex_positions: Float[Tensor, "Nv 3"],
+    vertex_normals: Float[Tensor, "Nv 3"],
+    triangle_idxs: Integer[Tensor, "Nf 3"],
+    face_uv: Float[Tensor, "Nf 3 2"],
+) -> Float[Tensor, "Nf 3 4"]:  # noqa: F821
+    vn_idx = [None] * 3
+    pos = [None] * 3
+    tex = face_uv.unbind(1)
+    for i in range(0, 3):
+        pos[i] = vertex_positions[triangle_idxs[:, i]]
+        # t_nrm_idx is always the same as t_pos_idx
+        vn_idx[i] = triangle_idxs[:, i]
+
+    tangents = torch.zeros_like(vertex_normals)
+    tansum = torch.zeros_like(vertex_normals)
+
+    # Compute tangent space for each triangle
+    duv1 = tex[1] - tex[0]
+    duv2 = tex[2] - tex[0]
+    dpos1 = pos[1] - pos[0]
+    dpos2 = pos[2] - pos[0]
+
+    tng_nom = dpos1 * duv2[..., 1:2] - dpos2 * duv1[..., 1:2]
+
+    denom = duv1[..., 0:1] * duv2[..., 1:2] - duv1[..., 1:2] * duv2[..., 0:1]
+
+    # Avoid division by zero for degenerated texture coordinates
+    denom_safe = denom.clip(1e-6)
+    tang = tng_nom / denom_safe
+
+    # Update all 3 vertices
+    for i in range(0, 3):
+        idx = vn_idx[i][:, None].repeat(1, 3)
+        tangents.scatter_add_(0, idx, tang)  # tangents[n_i] = tangents[n_i] + tang
+        tansum.scatter_add_(
+            0, idx, torch.ones_like(tang)
+        )  # tansum[n_i] = tansum[n_i] + 1
+    # Also normalize it. Here we do not normalize the individual triangles first so larger area
+    # triangles influence the tangent space more
+    tangents = tangents / tansum
+
+    # Normalize and make sure tangent is perpendicular to normal
+    tangents = F.normalize(tangents, dim=1)
+    tangents = F.normalize(tangents - dot(tangents, vertex_normals) * vertex_normals)
+
+    return tangents
+
+
+def _rotate_uv_slices_consistent_space(
+    vertex_positions: Float[Tensor, "Nv 3"],
+    vertex_normals: Float[Tensor, "Nv 3"],
+    triangle_idxs: Integer[Tensor, "Nf 3"],
+    uv: Float[Tensor, "Nf 3 2"],
+    index: Integer[Tensor, "Nf"],  # noqa: F821
+):
+    tangents = calculate_tangents(vertex_positions, vertex_normals, triangle_idxs, uv)
+    pos_stack = torch.stack(
+        [
+            -vertex_positions[..., 1],
+            vertex_positions[..., 0],
+            torch.zeros_like(vertex_positions[..., 0]),
+        ],
+        dim=-1,
+    )
+    expected_tangents = F.normalize(
+        torch.linalg.cross(
+            vertex_normals, torch.linalg.cross(pos_stack, vertex_normals)
+        ),
+        -1,
+    )
+
+    actual_tangents = tangents[triangle_idxs]
+    expected_tangents = expected_tangents[triangle_idxs]
+
+    def rotation_matrix_2d(theta):
+        c, s = torch.cos(theta), torch.sin(theta)
+        return torch.tensor([[c, -s], [s, c]])
+
+    # Now find the rotation
+    index_mod = index % 6  # Shouldn't happen. Just for safety
+    for i in range(6):
+        mask = index_mod == i
+        if not mask.any():
+            continue
+
+        actual_mean_tangent = actual_tangents[mask].mean(dim=(0, 1))
+        expected_mean_tangent = expected_tangents[mask].mean(dim=(0, 1))
+
+        dot_product = torch.dot(actual_mean_tangent, expected_mean_tangent)
+        cross_product = (
+            actual_mean_tangent[0] * expected_mean_tangent[1]
+            - actual_mean_tangent[1] * expected_mean_tangent[0]
+        )
+        angle = torch.atan2(cross_product, dot_product)
+
+        rot_matrix = rotation_matrix_2d(angle).to(mask.device)
+        # Center the uv coordinate to be in the range of -1 to 1 and 0 centered
+        uv_cur = uv[mask] * 2 - 1  # Center it first
+        # Rotate it
+        uv[mask] = torch.einsum("ij,nfj->nfi", rot_matrix, uv_cur)
+
+        # Rescale uv[mask] to be within the 0-1 range
+        uv[mask] = (uv[mask] - uv[mask].min()) / (uv[mask].max() - uv[mask].min())
+
+    return uv
+
+
+def _handle_slice_uvs(
+    uv: Float[Tensor, "Nf 3 2"],
+    index: Integer[Tensor, "Nf"],  # noqa: F821
+    island_padding: float,
+    max_index: int = 6 * 2,
+) -> Float[Tensor, "Nf 3 2"]:  # noqa: F821
+    uc, vc = uv.unbind(-1)
+
+    # Get the second slice (The first overlap)
+    index_filter = [index == i for i in range(6, max_index)]
+
+    # Normalize them to always fully fill the atlas patch
+    for i, fi in enumerate(index_filter):
+        if fi.sum() > 0:
+            # Scale the slice but only up to a factor of 2
+            # This keeps the texture resolution with the first slice in line (Half space in UV)
+            uc[fi] = (uc[fi] - uc[fi].min()) / (uc[fi].max() - uc[fi].min()).clip(0.5)
+            vc[fi] = (vc[fi] - vc[fi].min()) / (vc[fi].max() - vc[fi].min()).clip(0.5)
+
+    uc_padded = (uc * (1 - 2 * island_padding) + island_padding).clip(0, 1)
+    vc_padded = (vc * (1 - 2 * island_padding) + island_padding).clip(0, 1)
+
+    return torch.stack([uc_padded, vc_padded], dim=-1)
+
+
+def _handle_remaining_uvs(
+    uv: Float[Tensor, "Nf 3 2"],
+    index: Integer[Tensor, "Nf"],  # noqa: F821
+    island_padding: float,
+) -> Float[Tensor, "Nf 3 2"]:
+    uc, vc = uv.unbind(-1)
+    # Get all remaining elements
+    remaining_filter = index >= 6 * 2
+    squares_left = remaining_filter.sum()
+
+    if squares_left == 0:
+        return uv
+
+    uc = uc[remaining_filter]
+    vc = vc[remaining_filter]
+
+    # Or remaining triangles are distributed in a rectangle
+    # The rectangle takes 0.5 of the entire uv space in width and 1/3 in height
+    ratio = 0.5 * (1 / 3)  # 1.5
+    # sqrt(744/(0.5*(1/3)))
+
+    mult = math.sqrt(squares_left / ratio)
+    num_square_width = int(math.ceil(0.5 * mult))
+    num_square_height = int(math.ceil(squares_left / num_square_width))
+
+    width = 1 / num_square_width
+    height = 1 / num_square_height
+
+    # The idea is again to keep the texture resolution consistent with the first slice
+    # This only occupys half the region in the texture chart but the scaling on the squares
+    # assumes full coverage.
+    clip_val = min(width, height) * 1.5
+    # Now normalize the UVs with taking into account the maximum scaling
+    uc = (uc - uc.min(dim=1, keepdim=True).values) / (
+        uc.amax(dim=1, keepdim=True) - uc.amin(dim=1, keepdim=True)
+    ).clip(clip_val)
+    vc = (vc - vc.min(dim=1, keepdim=True).values) / (
+        vc.amax(dim=1, keepdim=True) - vc.amin(dim=1, keepdim=True)
+    ).clip(clip_val)
+    # Add a small padding
+    uc = (
+        uc * (1 - island_padding * num_square_width * 0.5)
+        + island_padding * num_square_width * 0.25
+    ).clip(0, 1)
+    vc = (
+        vc * (1 - island_padding * num_square_height * 0.5)
+        + island_padding * num_square_height * 0.25
+    ).clip(0, 1)
+
+    uc = uc * width
+    vc = vc * height
+
+    # And calculate offsets for each element
+    idx = torch.arange(uc.shape[0], device=uc.device, dtype=torch.int32)
+    x_idx = idx % num_square_width
+    y_idx = idx // num_square_width
+    # And move each triangle to its own spot
+    uc = uc + x_idx[:, None] * width
+    vc = vc + y_idx[:, None] * height
+
+    uc = (uc * (1 - 2 * island_padding * 0.5) + island_padding * 0.5).clip(0, 1)
+    vc = (vc * (1 - 2 * island_padding * 0.5) + island_padding * 0.5).clip(0, 1)
+
+    uv[remaining_filter] = torch.stack([uc, vc], dim=-1)
+
+    return uv
+
+
+def _distribute_individual_uvs_in_atlas(
+    face_uv: Float[Tensor, "Nf 3 2"],
+    assigned_faces: Integer[Tensor, "Nf"],  # noqa: F821
+    offset_x: Float[Tensor, "Nf"],  # noqa: F821
+    offset_y: Float[Tensor, "Nf"],  # noqa: F821
+    div_x: Float[Tensor, "Nf"],  # noqa: F821
+    div_y: Float[Tensor, "Nf"],  # noqa: F821
+    island_padding: float,
+):
+    # Place the slice first
+    placed_uv = _handle_slice_uvs(face_uv, assigned_faces, island_padding)
+    # Then handle the remaining overlap elements
+    placed_uv = _handle_remaining_uvs(placed_uv, assigned_faces, island_padding)
+
+    uc, vc = placed_uv.unbind(-1)
+    uc = uc / div_x[:, None] + offset_x[:, None]
+    vc = vc / div_y[:, None] + offset_y[:, None]
+
+    uv = torch.stack([uc, vc], dim=-1).view(-1, 2)
+
+    return uv
+
+
+def _get_unique_face_uv(
+    uv: Float[Tensor, "Nf 3 2"],
+) -> Tuple[Float[Tensor, "Utex 3"], Integer[Tensor, "Nf"]]:  # noqa: F821
+    unique_uv, unique_idx = torch.unique(uv, return_inverse=True, dim=0)
+    # And add the face to uv index mapping
+    vtex_idx = unique_idx.view(-1, 3)
+
+    return unique_uv, vtex_idx
+
+
+def _align_mesh_with_main_axis(
+    vertex_positions: Float[Tensor, "Nv 3"], vertex_normals: Float[Tensor, "Nv 3"]
+) -> Tuple[Float[Tensor, "Nv 3"], Float[Tensor, "Nv 3"]]:
+    # Use pca to find the 2 main axis (third is derived by cross product)
+    # Set the random seed so it's repeatable
+    torch.manual_seed(0)
+    _, _, v = torch.pca_lowrank(vertex_positions, q=2)
+    main_axis, seconday_axis = v[:, 0], v[:, 1]
+
+    main_axis: Float[Tensor, "3"] = F.normalize(main_axis, eps=1e-6, dim=-1)
+    # Orthogonalize the second axis
+    seconday_axis: Float[Tensor, "3"] = F.normalize(
+        seconday_axis - dot(seconday_axis, main_axis) * main_axis, eps=1e-6, dim=-1
+    )
+    # Create perpendicular third axis
+    third_axis: Float[Tensor, "3"] = F.normalize(
+        torch.cross(main_axis, seconday_axis), dim=-1, eps=1e-6
+    )
+
+    # Check to which canonical axis each aligns
+    main_axis_max_idx = main_axis.abs().argmax().item()
+    seconday_axis_max_idx = seconday_axis.abs().argmax().item()
+    third_axis_max_idx = third_axis.abs().argmax().item()
+
+    # Now sort the axes based on the argmax so they align with thecanonoical axes
+    # If two axes have the same argmax move one of them
+    all_possible_axis = {0, 1, 2}
+    cur_index = 1
+    while len(set([main_axis_max_idx, seconday_axis_max_idx, third_axis_max_idx])) != 3:
+        # Find missing axis
+        missing_axis = all_possible_axis - set(
+            [main_axis_max_idx, seconday_axis_max_idx, third_axis_max_idx]
+        )
+        missing_axis = missing_axis.pop()
+        # Just assign it to third axis as it had the smallest contribution to the
+        # overall shape
+        if cur_index == 1:
+            third_axis_max_idx = missing_axis
+        elif cur_index == 2:
+            seconday_axis_max_idx = missing_axis
+        else:
+            raise ValueError("Could not find 3 unique axis")
+        cur_index += 1
+
+    if len({main_axis_max_idx, seconday_axis_max_idx, third_axis_max_idx}) != 3:
+        raise ValueError("Could not find 3 unique axis")
+
+    axes = [None] * 3
+    axes[main_axis_max_idx] = main_axis
+    axes[seconday_axis_max_idx] = seconday_axis
+    axes[third_axis_max_idx] = third_axis
+    # Create rotation matrix from the individual axes
+    rot_mat = torch.stack(axes, dim=1).T
+
+    # Now rotate the vertex positions and vertex normals so the mesh aligns with the main axis
+    vertex_positions = torch.einsum("ij,nj->ni", rot_mat, vertex_positions)
+    vertex_normals = torch.einsum("ij,nj->ni", rot_mat, vertex_normals)
+
+    return vertex_positions, vertex_normals
+
+
+def box_projection_uv_unwrap(
+    vertex_positions: Float[Tensor, "Nv 3"],
+    vertex_normals: Float[Tensor, "Nv 3"],
+    triangle_idxs: Integer[Tensor, "Nf 3"],
+    island_padding: float,
+) -> Tuple[Float[Tensor, "Utex 3"], Integer[Tensor, "Nf"]]:  # noqa: F821
+    # Align the mesh with main axis directions first
+    vertex_positions, vertex_normals = _align_mesh_with_main_axis(
+        vertex_positions, vertex_normals
+    )
+
+    bbox: Float[Tensor, "2 3"] = torch.stack(
+        [vertex_positions.min(dim=0).values, vertex_positions.max(dim=0).values], dim=0
+    )
+    # First decide in which cube face the triangle is placed
+    face_uv, face_index = _box_assign_vertex_to_cube_face(
+        vertex_positions, vertex_normals, triangle_idxs, bbox
+    )
+
+    # Rotate the UV islands in a way that they align with the radial z tangent space
+    face_uv = _rotate_uv_slices_consistent_space(
+        vertex_positions, vertex_normals, triangle_idxs, face_uv, face_index
+    )
+
+    # Then find where where the face is placed in the atlas.
+    # This has to detect potential overlaps
+    assigned_atlas_index = _assign_faces_uv_to_atlas_index(
+        vertex_positions, triangle_idxs, face_uv, face_index
+    )
+
+    # Then figure out the final place in the atlas based on the assignment
+    offset_x, offset_y, div_x, div_y = _find_slice_offset_and_scale(
+        assigned_atlas_index
+    )
+
+    # Next distribute the faces in the uv atlas
+    placed_uv = _distribute_individual_uvs_in_atlas(
+        face_uv, assigned_atlas_index, offset_x, offset_y, div_x, div_y, island_padding
+    )
+
+    # And get the unique per-triangle UV coordinates
+    return _get_unique_face_uv(placed_uv)

sf3d/models/camera.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+from typing import List
+
+import torch
+import torch.nn as nn
+
+from sf3d.models.utils import BaseModule
+
+
+class LinearCameraEmbedder(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int = 25
+        out_channels: int = 768
+        conditions: List[str] = field(default_factory=list)
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.linear = nn.Linear(self.cfg.in_channels, self.cfg.out_channels)
+
+    def forward(self, **kwargs):
+        cond_tensors = []
+        for cond_name in self.cfg.conditions:
+            assert cond_name in kwargs
+            cond = kwargs[cond_name]
+            # cond in shape (B, Nv, ...)
+            cond_tensors.append(cond.view(*cond.shape[:2], -1))
+        cond_tensor = torch.cat(cond_tensors, dim=-1)
+        assert cond_tensor.shape[-1] == self.cfg.in_channels
+        embedding = self.linear(cond_tensor)
+        return embedding

sf3d/models/global_estimator/multi_head_estimator.py

@@ -0,0 +1,118 @@
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from sf3d.models.network import get_activation
+from sf3d.models.utils import BaseModule
+
+
+@dataclass
+class HeadSpec:
+    name: str
+    out_channels: int
+    n_hidden_layers: int
+    output_activation: Optional[str] = None
+    output_bias: float = 0.0
+    add_to_decoder_features: bool = False
+    shape: Optional[list[int]] = None
+
+
+class MultiHeadEstimator(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        triplane_features: int = 1024
+
+        n_layers: int = 2
+        hidden_features: int = 512
+        activation: str = "relu"
+
+        pool: str = "max"
+        # Literal["mean", "max"] = "mean"  # noqa: F821
+
+        heads: List[HeadSpec] = field(default_factory=lambda: [])
+
+    cfg: Config
+
+    def configure(self):
+        layers = []
+        cur_features = self.cfg.triplane_features * 3
+        for _ in range(self.cfg.n_layers):
+            layers.append(
+                nn.Conv2d(
+                    cur_features,
+                    self.cfg.hidden_features,
+                    kernel_size=3,
+                    padding=0,
+                    stride=2,
+                )
+            )
+            layers.append(self.make_activation(self.cfg.activation))
+
+            cur_features = self.cfg.hidden_features
+
+        self.layers = nn.Sequential(*layers)
+
+        assert len(self.cfg.heads) > 0
+        heads = {}
+        for head in self.cfg.heads:
+            head_layers = []
+            for i in range(head.n_hidden_layers):
+                head_layers += [
+                    nn.Linear(
+                        self.cfg.hidden_features,
+                        self.cfg.hidden_features,
+                    ),
+                    self.make_activation(self.cfg.activation),
+                ]
+            head_layers += [
+                nn.Linear(
+                    self.cfg.hidden_features,
+                    head.out_channels,
+                ),
+            ]
+            heads[head.name] = nn.Sequential(*head_layers)
+        self.heads = nn.ModuleDict(heads)
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def forward(
+        self,
+        triplane: Float[Tensor, "B 3 F Ht Wt"],
+    ) -> dict[str, Any]:
+        x = self.layers(
+            triplane.reshape(
+                triplane.shape[0], -1, triplane.shape[-2], triplane.shape[-1]
+            )
+        )
+
+        if self.cfg.pool == "max":
+            x = x.amax(dim=[-2, -1])
+        elif self.cfg.pool == "mean":
+            x = x.mean(dim=[-2, -1])
+        else:
+            raise NotImplementedError
+
+        out = {
+            ("decoder_" if head.add_to_decoder_features else "")
+            + head.name: get_activation(head.output_activation)(
+                self.heads[head.name](x) + head.output_bias
+            )
+            for head in self.cfg.heads
+        }
+        for head in self.cfg.heads:
+            if head.shape:
+                head_name = (
+                    "decoder_" if head.add_to_decoder_features else ""
+                ) + head.name
+                out[head_name] = out[head_name].reshape(*head.shape)
+
+        return out

sf3d/models/image_estimator/clip_based_estimator.py

@@ -0,0 +1,168 @@
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+import open_clip
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+from torchvision.transforms import Normalize
+
+from sf3d.models.network import get_activation
+from sf3d.models.utils import BaseModule
+
+
+@dataclass
+class HeadSpec:
+    name: str
+    out_channels: int
+    n_hidden_layers: int
+    output_activation: Optional[str] = None
+    output_bias: float = 0.0
+    add_to_decoder_features: bool = False
+    shape: Optional[list[int]] = None
+
+
+class ClipBasedHeadEstimator(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        model: str = "ViT-B-32"
+        pretrain: str = "laion2b_s34b_b79k"
+
+        distribution: str = "beta"
+
+        # ["mean", "mode", "sample", "sample_mean"]
+        distribution_eval: str = "mode"
+
+        activation: str = "relu"
+        hidden_features: int = 512
+        heads: List[HeadSpec] = field(default_factory=lambda: [])
+
+    cfg: Config
+
+    def configure(self):
+        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+            self.cfg.model, pretrained=self.cfg.pretrain
+        )
+        self.model.eval()
+
+        # Do not add the weights in self.model to the optimizer
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        assert len(self.cfg.heads) > 0
+        heads = {}
+        for head in self.cfg.heads:
+            head_layers = []
+
+            for i in range(head.n_hidden_layers):
+                head_layers += [
+                    nn.Linear(
+                        self.cfg.hidden_features,
+                        self.cfg.hidden_features,
+                    ),
+                    self.make_activation(self.cfg.activation),
+                ]
+
+            head_layers = [nn.Sequential(*head_layers)]
+            head_layers += [
+                nn.Sequential(
+                    nn.Linear(
+                        self.cfg.hidden_features,
+                        self.cfg.hidden_features,
+                    ),
+                    self.make_activation(self.cfg.activation),
+                    nn.Linear(self.cfg.hidden_features, 1),
+                )
+                for _ in range(2)
+            ]
+            heads[head.name] = nn.ModuleList(head_layers)
+        self.heads = nn.ModuleDict(heads)
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def forward(
+        self,
+        cond_image: Float[Tensor, "B 1 H W 3"],
+        sample: bool = True,
+    ) -> dict[str, Any]:
+        # Run the model
+        # Resize cond_image to 224
+        cond_image = nn.functional.interpolate(
+            cond_image.flatten(0, 1).permute(0, 3, 1, 2),
+            size=(224, 224),
+            mode="bilinear",
+            align_corners=False,
+        )
+        cond_image = Normalize(
+            mean=open_clip.constants.OPENAI_DATASET_MEAN,
+            std=open_clip.constants.OPENAI_DATASET_STD,
+        )(cond_image)
+        image_features = self.model.encode_image(cond_image)
+
+        # Run the heads
+        outputs = {}
+
+        for head_dict in self.cfg.heads:
+            head_name = head_dict.name
+            shared_head, d1_h, d2_h = self.heads[head_name]
+            shared_features = shared_head(image_features)
+            d1, d2 = [head(shared_features).squeeze(-1) for head in [d1_h, d2_h]]
+            if self.cfg.distribution == "normal":
+                mean = d1
+                var = d2
+                if mean.shape[-1] == 1:
+                    outputs[head_name] = torch.distributions.Normal(
+                        mean + head_dict.output_bias,
+                        torch.nn.functional.softplus(var),
+                    )
+                else:
+                    outputs[head_name] = torch.distributions.MultivariateNormal(
+                        mean + head_dict.output_bias,
+                        torch.nn.functional.softplus(var).diag_embed(),
+                    )
+            elif self.cfg.distribution == "beta":
+                outputs[head_name] = torch.distributions.Beta(
+                    torch.nn.functional.softplus(d1 + head_dict.output_bias),
+                    torch.nn.functional.softplus(d2 + head_dict.output_bias),
+                )
+            else:
+                raise NotImplementedError
+
+        if sample:
+            for head_dict in self.cfg.heads:
+                head_name = head_dict.name
+                dist = outputs[head_name]
+
+                if self.cfg.distribution_eval == "mean":
+                    out = dist.mean
+                elif self.cfg.distribution_eval == "mode":
+                    out = dist.mode
+                elif self.cfg.distribution_eval == "sample_mean":
+                    out = dist.sample([10]).mean(-1)
+                else:
+                    # use rsample if gradient is needed
+                    out = dist.rsample() if self.training else dist.sample()
+
+                outputs[head_name] = get_activation(head_dict.output_activation)(out)
+                outputs[f"{head_name}_dist"] = dist
+
+        for head in self.cfg.heads:
+            if head.shape:
+                if not sample:
+                    raise ValueError(
+                        "Cannot reshape non-sampled probabilisitic outputs"
+                    )
+                outputs[head.name] = outputs[head.name].reshape(*head.shape)
+
+            if head.add_to_decoder_features:
+                outputs[f"decoder_{head.name}"] = outputs[head.name]
+                del outputs[head.name]
+
+        return outputs

sf3d/models/isosurface.py

@@ -0,0 +1,229 @@
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from jaxtyping import Float, Integer
+from torch import Tensor
+
+from .mesh import Mesh
+
+
+class IsosurfaceHelper(nn.Module):
+    points_range: Tuple[float, float] = (0, 1)
+
+    @property
+    def grid_vertices(self) -> Float[Tensor, "N 3"]:
+        raise NotImplementedError
+
+    @property
+    def requires_instance_per_batch(self) -> bool:
+        return False
+
+
+class MarchingTetrahedraHelper(IsosurfaceHelper):
+    def __init__(self, resolution: int, tets_path: str):
+        super().__init__()
+        self.resolution = resolution
+        self.tets_path = tets_path
+
+        self.triangle_table: Float[Tensor, "..."]
+        self.register_buffer(
+            "triangle_table",
+            torch.as_tensor(
+                [
+                    [-1, -1, -1, -1, -1, -1],
+                    [1, 0, 2, -1, -1, -1],
+                    [4, 0, 3, -1, -1, -1],
+                    [1, 4, 2, 1, 3, 4],
+                    [3, 1, 5, -1, -1, -1],
+                    [2, 3, 0, 2, 5, 3],
+                    [1, 4, 0, 1, 5, 4],
+                    [4, 2, 5, -1, -1, -1],
+                    [4, 5, 2, -1, -1, -1],
+                    [4, 1, 0, 4, 5, 1],
+                    [3, 2, 0, 3, 5, 2],
+                    [1, 3, 5, -1, -1, -1],
+                    [4, 1, 2, 4, 3, 1],
+                    [3, 0, 4, -1, -1, -1],
+                    [2, 0, 1, -1, -1, -1],
+                    [-1, -1, -1, -1, -1, -1],
+                ],
+                dtype=torch.long,
+            ),
+            persistent=False,
+        )
+        self.num_triangles_table: Integer[Tensor, "..."]
+        self.register_buffer(
+            "num_triangles_table",
+            torch.as_tensor(
+                [0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long
+            ),
+            persistent=False,
+        )
+        self.base_tet_edges: Integer[Tensor, "..."]
+        self.register_buffer(
+            "base_tet_edges",
+            torch.as_tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long),
+            persistent=False,
+        )
+
+        tets = np.load(self.tets_path)
+        self._grid_vertices: Float[Tensor, "..."]
+        self.register_buffer(
+            "_grid_vertices",
+            torch.from_numpy(tets["vertices"]).float(),
+            persistent=False,
+        )
+        self.indices: Integer[Tensor, "..."]
+        self.register_buffer(
+            "indices", torch.from_numpy(tets["indices"]).long(), persistent=False
+        )
+
+        self._all_edges: Optional[Integer[Tensor, "Ne 2"]] = None
+
+        center_indices, boundary_indices = self.get_center_boundary_index(
+            self._grid_vertices
+        )
+        self.center_indices: Integer[Tensor, "..."]
+        self.register_buffer("center_indices", center_indices, persistent=False)
+        self.boundary_indices: Integer[Tensor, "..."]
+        self.register_buffer("boundary_indices", boundary_indices, persistent=False)
+
+    def get_center_boundary_index(self, verts):
+        magn = torch.sum(verts**2, dim=-1)
+
+        center_idx = torch.argmin(magn)
+        boundary_neg = verts == verts.max()
+        boundary_pos = verts == verts.min()
+
+        boundary = torch.bitwise_or(boundary_pos, boundary_neg)
+        boundary = torch.sum(boundary.float(), dim=-1)
+
+        boundary_idx = torch.nonzero(boundary)
+        return center_idx, boundary_idx.squeeze(dim=-1)
+
+    def normalize_grid_deformation(
+        self, grid_vertex_offsets: Float[Tensor, "Nv 3"]
+    ) -> Float[Tensor, "Nv 3"]:
+        return (
+            (self.points_range[1] - self.points_range[0])
+            / self.resolution  # half tet size is approximately 1 / self.resolution
+            * torch.tanh(grid_vertex_offsets)
+        )  # FIXME: hard-coded activation
+
+    @property
+    def grid_vertices(self) -> Float[Tensor, "Nv 3"]:
+        return self._grid_vertices
+
+    @property
+    def all_edges(self) -> Integer[Tensor, "Ne 2"]:
+        if self._all_edges is None:
+            # compute edges on GPU, or it would be VERY SLOW (basically due to the unique operation)
+            edges = torch.tensor(
+                [0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3],
+                dtype=torch.long,
+                device=self.indices.device,
+            )
+            _all_edges = self.indices[:, edges].reshape(-1, 2)
+            _all_edges_sorted = torch.sort(_all_edges, dim=1)[0]
+            _all_edges = torch.unique(_all_edges_sorted, dim=0)
+            self._all_edges = _all_edges
+        return self._all_edges
+
+    def sort_edges(self, edges_ex2):
+        with torch.no_grad():
+            order = (edges_ex2[:, 0] > edges_ex2[:, 1]).long()
+            order = order.unsqueeze(dim=1)
+
+            a = torch.gather(input=edges_ex2, index=order, dim=1)
+            b = torch.gather(input=edges_ex2, index=1 - order, dim=1)
+
+        return torch.stack([a, b], -1)
+
+    def _forward(self, pos_nx3, sdf_n, tet_fx4):
+        with torch.no_grad():
+            occ_n = sdf_n > 0
+            occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1, 4)
+            occ_sum = torch.sum(occ_fx4, -1)
+            valid_tets = (occ_sum > 0) & (occ_sum < 4)
+            occ_sum = occ_sum[valid_tets]
+
+            # find all vertices
+            all_edges = tet_fx4[valid_tets][:, self.base_tet_edges].reshape(-1, 2)
+            all_edges = self.sort_edges(all_edges)
+            unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True)
+
+            unique_edges = unique_edges.long()
+            mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
+            mapping = (
+                torch.ones(
+                    (unique_edges.shape[0]), dtype=torch.long, device=pos_nx3.device
+                )
+                * -1
+            )
+            mapping[mask_edges] = torch.arange(
+                mask_edges.sum(), dtype=torch.long, device=pos_nx3.device
+            )
+            idx_map = mapping[idx_map]  # map edges to verts
+
+            interp_v = unique_edges[mask_edges]
+        edges_to_interp = pos_nx3[interp_v.reshape(-1)].reshape(-1, 2, 3)
+        edges_to_interp_sdf = sdf_n[interp_v.reshape(-1)].reshape(-1, 2, 1)
+        edges_to_interp_sdf[:, -1] *= -1
+
+        denominator = edges_to_interp_sdf.sum(1, keepdim=True)
+
+        edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1]) / denominator
+        verts = (edges_to_interp * edges_to_interp_sdf).sum(1)
+
+        idx_map = idx_map.reshape(-1, 6)
+
+        v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=pos_nx3.device))
+        tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1)
+        num_triangles = self.num_triangles_table[tetindex]
+
+        # Generate triangle indices
+        faces = torch.cat(
+            (
+                torch.gather(
+                    input=idx_map[num_triangles == 1],
+                    dim=1,
+                    index=self.triangle_table[tetindex[num_triangles == 1]][:, :3],
+                ).reshape(-1, 3),
+                torch.gather(
+                    input=idx_map[num_triangles == 2],
+                    dim=1,
+                    index=self.triangle_table[tetindex[num_triangles == 2]][:, :6],
+                ).reshape(-1, 3),
+            ),
+            dim=0,
+        )
+
+        return verts, faces
+
+    def forward(
+        self,
+        level: Float[Tensor, "N3 1"],
+        deformation: Optional[Float[Tensor, "N3 3"]] = None,
+    ) -> Mesh:
+        if deformation is not None:
+            grid_vertices = self.grid_vertices + self.normalize_grid_deformation(
+                deformation
+            )
+        else:
+            grid_vertices = self.grid_vertices
+
+        v_pos, t_pos_idx = self._forward(grid_vertices, level, self.indices)
+
+        mesh = Mesh(
+            v_pos=v_pos,
+            t_pos_idx=t_pos_idx,
+            # extras
+            grid_vertices=grid_vertices,
+            tet_edges=self.all_edges,
+            grid_level=level,
+            grid_deformation=deformation,
+        )
+
+        return mesh

sf3d/models/mesh.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn.functional as F
+from jaxtyping import Float, Integer
+from torch import Tensor
+
+from sf3d.box_uv_unwrap import box_projection_uv_unwrap
+from sf3d.models.utils import dot
+
+
+class Mesh:
+    def __init__(
+        self, v_pos: Float[Tensor, "Nv 3"], t_pos_idx: Integer[Tensor, "Nf 3"], **kwargs
+    ) -> None:
+        self.v_pos: Float[Tensor, "Nv 3"] = v_pos
+        self.t_pos_idx: Integer[Tensor, "Nf 3"] = t_pos_idx
+        self._v_nrm: Optional[Float[Tensor, "Nv 3"]] = None
+        self._v_tng: Optional[Float[Tensor, "Nv 3"]] = None
+        self._v_tex: Optional[Float[Tensor, "Nt 3"]] = None
+        self._edges: Optional[Integer[Tensor, "Ne 2"]] = None
+        self.extras: Dict[str, Any] = {}
+        for k, v in kwargs.items():
+            self.add_extra(k, v)
+
+    def add_extra(self, k, v) -> None:
+        self.extras[k] = v
+
+    @property
+    def requires_grad(self):
+        return self.v_pos.requires_grad
+
+    @property
+    def v_nrm(self):
+        if self._v_nrm is None:
+            self._v_nrm = self._compute_vertex_normal()
+        return self._v_nrm
+
+    @property
+    def v_tng(self):
+        if self._v_tng is None:
+            self._v_tng = self._compute_vertex_tangent()
+        return self._v_tng
+
+    @property
+    def v_tex(self):
+        if self._v_tex is None:
+            self.unwrap_uv()
+        return self._v_tex
+
+    @property
+    def edges(self):
+        if self._edges is None:
+            self._edges = self._compute_edges()
+        return self._edges
+
+    def _compute_vertex_normal(self):
+        i0 = self.t_pos_idx[:, 0]
+        i1 = self.t_pos_idx[:, 1]
+        i2 = self.t_pos_idx[:, 2]
+
+        v0 = self.v_pos[i0, :]
+        v1 = self.v_pos[i1, :]
+        v2 = self.v_pos[i2, :]
+
+        face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+
+        # Splat face normals to vertices
+        v_nrm = torch.zeros_like(self.v_pos)
+        v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
+        v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
+        v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
+
+        # Normalize, replace zero (degenerated) normals with some default value
+        v_nrm = torch.where(
+            dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
+        )
+        v_nrm = F.normalize(v_nrm, dim=1)
+
+        if torch.is_anomaly_enabled():
+            assert torch.all(torch.isfinite(v_nrm))
+
+        return v_nrm
+
+    def _compute_vertex_tangent(self):
+        vn_idx = [None] * 3
+        pos = [None] * 3
+        tex = [None] * 3
+        for i in range(0, 3):
+            pos[i] = self.v_pos[self.t_pos_idx[:, i]]
+            tex[i] = self.v_tex[self.t_pos_idx[:, i]]
+            # t_nrm_idx is always the same as t_pos_idx
+            vn_idx[i] = self.t_pos_idx[:, i]
+
+        tangents = torch.zeros_like(self.v_nrm)
+        tansum = torch.zeros_like(self.v_nrm)
+
+        # Compute tangent space for each triangle
+        duv1 = tex[1] - tex[0]
+        duv2 = tex[2] - tex[0]
+        dpos1 = pos[1] - pos[0]
+        dpos2 = pos[2] - pos[0]
+
+        tng_nom = dpos1 * duv2[..., 1:2] - dpos2 * duv1[..., 1:2]
+
+        denom = duv1[..., 0:1] * duv2[..., 1:2] - duv1[..., 1:2] * duv2[..., 0:1]
+
+        # Avoid division by zero for degenerated texture coordinates
+        denom_safe = denom.clip(1e-6)
+        tang = tng_nom / denom_safe
+
+        # Update all 3 vertices
+        for i in range(0, 3):
+            idx = vn_idx[i][:, None].repeat(1, 3)
+            tangents.scatter_add_(0, idx, tang)  # tangents[n_i] = tangents[n_i] + tang
+            tansum.scatter_add_(
+                0, idx, torch.ones_like(tang)
+            )  # tansum[n_i] = tansum[n_i] + 1
+        # Also normalize it. Here we do not normalize the individual triangles first so larger area
+        # triangles influence the tangent space more
+        tangents = tangents / tansum
+
+        # Normalize and make sure tangent is perpendicular to normal
+        tangents = F.normalize(tangents, dim=1)
+        tangents = F.normalize(tangents - dot(tangents, self.v_nrm) * self.v_nrm)
+
+        if torch.is_anomaly_enabled():
+            assert torch.all(torch.isfinite(tangents))
+
+        return tangents
+
+    @torch.no_grad()
+    def unwrap_uv(
+        self,
+        island_padding: float = 0.02,
+    ) -> Mesh:
+        uv, indices = box_projection_uv_unwrap(
+            self.v_pos, self.v_nrm, self.t_pos_idx, island_padding
+        )
+
+        # Do store per vertex UVs.
+        # This means we need to duplicate some vertices at the seams
+        individual_vertices = self.v_pos[self.t_pos_idx].reshape(-1, 3)
+        individual_faces = torch.arange(
+            individual_vertices.shape[0],
+            device=individual_vertices.device,
+            dtype=self.t_pos_idx.dtype,
+        ).reshape(-1, 3)
+        uv_flat = uv[indices].reshape((-1, 2))
+        # uv_flat[:, 1] = 1 - uv_flat[:, 1]
+
+        self.v_pos = individual_vertices
+        self.t_pos_idx = individual_faces
+        self._v_tex = uv_flat
+        self._v_nrm = self._compute_vertex_normal()
+        self._v_tng = self._compute_vertex_tangent()
+
+    def _compute_edges(self):
+        # Compute edges
+        edges = torch.cat(
+            [
+                self.t_pos_idx[:, [0, 1]],
+                self.t_pos_idx[:, [1, 2]],
+                self.t_pos_idx[:, [2, 0]],
+            ],
+            dim=0,
+        )
+        edges = edges.sort()[0]
+        edges = torch.unique(edges, dim=0)
+        return edges

sf3d/models/network.py

@@ -0,0 +1,195 @@
+from dataclasses import dataclass, field
+from typing import Callable, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+from torch.autograd import Function
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+from sf3d.models.utils import BaseModule, normalize
+
+
+class PixelShuffleUpsampleNetwork(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int = 1024
+        out_channels: int = 40
+        scale_factor: int = 4
+
+        conv_layers: int = 4
+        conv_kernel_size: int = 3
+
+    cfg: Config
+
+    def configure(self) -> None:
+        layers = []
+        output_channels = self.cfg.out_channels * self.cfg.scale_factor**2
+
+        in_channels = self.cfg.in_channels
+        for i in range(self.cfg.conv_layers):
+            cur_out_channels = (
+                in_channels if i != self.cfg.conv_layers - 1 else output_channels
+            )
+            layers.append(
+                nn.Conv2d(
+                    in_channels,
+                    cur_out_channels,
+                    self.cfg.conv_kernel_size,
+                    padding=(self.cfg.conv_kernel_size - 1) // 2,
+                )
+            )
+            if i != self.cfg.conv_layers - 1:
+                layers.append(nn.ReLU(inplace=True))
+
+        layers.append(nn.PixelShuffle(self.cfg.scale_factor))
+
+        self.upsample = nn.Sequential(*layers)
+
+    def forward(
+        self, triplanes: Float[Tensor, "B 3 Ci Hp Wp"]
+    ) -> Float[Tensor, "B 3 Co Hp2 Wp2"]:
+        return rearrange(
+            self.upsample(
+                rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
+            ),
+            "(B Np) Co Hp Wp -> B Np Co Hp Wp",
+            Np=3,
+        )
+
+
+class _TruncExp(Function):  # pylint: disable=abstract-method
+    # Implementation from torch-ngp:
+    # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, x):  # pylint: disable=arguments-differ
+        ctx.save_for_backward(x)
+        return torch.exp(x)
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, g):  # pylint: disable=arguments-differ
+        x = ctx.saved_tensors[0]
+        return g * torch.exp(torch.clamp(x, max=15))
+
+
+trunc_exp = _TruncExp.apply
+
+
+def get_activation(name) -> Callable:
+    if name is None:
+        return lambda x: x
+    name = name.lower()
+    if name == "none" or name == "linear" or name == "identity":
+        return lambda x: x
+    elif name == "lin2srgb":
+        return lambda x: torch.where(
+            x > 0.0031308,
+            torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) * 1.055 - 0.055,
+            12.92 * x,
+        ).clamp(0.0, 1.0)
+    elif name == "exp":
+        return lambda x: torch.exp(x)
+    elif name == "shifted_exp":
+        return lambda x: torch.exp(x - 1.0)
+    elif name == "trunc_exp":
+        return trunc_exp
+    elif name == "shifted_trunc_exp":
+        return lambda x: trunc_exp(x - 1.0)
+    elif name == "sigmoid":
+        return lambda x: torch.sigmoid(x)
+    elif name == "tanh":
+        return lambda x: torch.tanh(x)
+    elif name == "shifted_softplus":
+        return lambda x: F.softplus(x - 1.0)
+    elif name == "scale_-11_01":
+        return lambda x: x * 0.5 + 0.5
+    elif name == "negative":
+        return lambda x: -x
+    elif name == "normalize_channel_last":
+        return lambda x: normalize(x)
+    elif name == "normalize_channel_first":
+        return lambda x: normalize(x, dim=1)
+    else:
+        try:
+            return getattr(F, name)
+        except AttributeError:
+            raise ValueError(f"Unknown activation function: {name}")
+
+
+@dataclass
+class HeadSpec:
+    name: str
+    out_channels: int
+    n_hidden_layers: int
+    output_activation: Optional[str] = None
+    out_bias: float = 0.0
+
+
+class MaterialMLP(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int = 120
+        n_neurons: int = 64
+        activation: str = "silu"
+        heads: List[HeadSpec] = field(default_factory=lambda: [])
+
+    cfg: Config
+
+    def configure(self) -> None:
+        assert len(self.cfg.heads) > 0
+        heads = {}
+        for head in self.cfg.heads:
+            head_layers = []
+            for i in range(head.n_hidden_layers):
+                head_layers += [
+                    nn.Linear(
+                        self.cfg.in_channels if i == 0 else self.cfg.n_neurons,
+                        self.cfg.n_neurons,
+                    ),
+                    self.make_activation(self.cfg.activation),
+                ]
+            head_layers += [
+                nn.Linear(
+                    self.cfg.n_neurons,
+                    head.out_channels,
+                ),
+            ]
+            heads[head.name] = nn.Sequential(*head_layers)
+        self.heads = nn.ModuleDict(heads)
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def keys(self):
+        return self.heads.keys()
+
+    def forward(
+        self, x, include: Optional[List] = None, exclude: Optional[List] = None
+    ):
+        if include is not None and exclude is not None:
+            raise ValueError("Cannot specify both include and exclude.")
+        if include is not None:
+            heads = [h for h in self.cfg.heads if h.name in include]
+        elif exclude is not None:
+            heads = [h for h in self.cfg.heads if h.name not in exclude]
+        else:
+            heads = self.cfg.heads
+
+        out = {
+            head.name: get_activation(head.output_activation)(
+                self.heads[head.name](x) + head.out_bias
+            )
+            for head in heads
+        }
+
+        return out

sf3d/models/tokenizers/dinov2.py

@@ -0,0 +1,1196 @@
+# coding=utf-8
+# Copyright 2023 Meta AI and The HuggingFace Inc. 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.
+"""PyTorch DINOv2 model."""
+
+import collections.abc
+import math
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (
+    BackboneOutput,
+    BaseModelOutput,
+    BaseModelOutputWithPooling,
+    ImageClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.dinov2.configuration_dinov2 import Dinov2Config
+from transformers.pytorch_utils import (
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+from transformers.utils import (
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils.backbone_utils import BackboneMixin
+
+logger = logging.get_logger(__name__)
+
+# General docstring
+_CONFIG_FOR_DOC = "Dinov2Config"
+
+# Base docstring
+_CHECKPOINT_FOR_DOC = "facebook/dinov2-base"
+_EXPECTED_OUTPUT_SHAPE = [1, 257, 768]
+
+# Image classification docstring
+_IMAGE_CLASS_CHECKPOINT = "facebook/dinov2-base"
+
+
+DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "facebook/dinov2-base",
+    # See all DINOv2 models at https://huggingface.co/models?filter=dinov2
+]
+
+
+class Dinov2Embeddings(nn.Module):
+    """
+    Construct the CLS token, mask token, position and patch embeddings.
+    """
+
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        # register as mask token as it's not used in optimization
+        # to avoid the use of find_unused_parameters_true
+        # self.mask_token = nn.Parameter(torch.zeros(1, config.hidden_size))
+        self.register_buffer("mask_token", torch.zeros(1, config.hidden_size))
+        self.patch_embeddings = Dinov2PatchEmbeddings(config)
+        num_patches = self.patch_embeddings.num_patches
+        self.position_embeddings = nn.Parameter(
+            torch.randn(1, num_patches + 1, config.hidden_size)
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.config = config
+
+    def interpolate_pos_encoding(
+        self, embeddings: torch.Tensor, height: int, width: int
+    ) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+
+        num_patches = embeddings.shape[1] - 1
+        num_positions = self.position_embeddings.shape[1] - 1
+        if num_patches == num_positions and height == width:
+            return self.position_embeddings
+        class_pos_embed = self.position_embeddings[:, 0]
+        patch_pos_embed = self.position_embeddings[:, 1:]
+        dim = embeddings.shape[-1]
+        height = height // self.config.patch_size
+        width = width // self.config.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        height, width = height + 0.1, width + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(
+            1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim
+        )
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            scale_factor=(
+                height / math.sqrt(num_positions),
+                width / math.sqrt(num_positions),
+            ),
+            mode="bicubic",
+            align_corners=False,
+        )
+        if (
+            int(height) != patch_pos_embed.shape[-2]
+            or int(width) != patch_pos_embed.shape[-1]
+        ):
+            raise ValueError(
+                "Width or height does not match with the interpolated position embeddings"
+            )
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        bool_masked_pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch_size, _, height, width = pixel_values.shape
+        patch_embeddings = self.patch_embeddings(pixel_values)
+        embeddings = patch_embeddings
+
+        if bool_masked_pos is not None:
+            embeddings = torch.where(
+                bool_masked_pos.unsqueeze(-1),
+                self.mask_token.to(embeddings.dtype).unsqueeze(0),
+                embeddings,
+            )
+
+        # add the [CLS] token to the embedded patch tokens
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        embeddings = torch.cat((cls_tokens, embeddings), dim=1)
+
+        # add positional encoding to each token
+        embeddings = embeddings + self.interpolate_pos_encoding(
+            embeddings, height, width
+        )
+
+        embeddings = self.dropout(embeddings)
+
+        return embeddings
+
+
+class Dinov2PatchEmbeddings(nn.Module):
+    """
+    This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+    `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+    Transformer.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        image_size, patch_size = config.image_size, config.patch_size
+        num_channels, hidden_size = config.num_channels, config.hidden_size
+
+        image_size = (
+            image_size
+            if isinstance(image_size, collections.abc.Iterable)
+            else (image_size, image_size)
+        )
+        patch_size = (
+            patch_size
+            if isinstance(patch_size, collections.abc.Iterable)
+            else (patch_size, patch_size)
+        )
+        num_patches = (image_size[1] // patch_size[1]) * (
+            image_size[0] // patch_size[0]
+        )
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+
+        self.projection = nn.Conv2d(
+            num_channels, hidden_size, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        num_channels = pixel_values.shape[1]
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+                f" Expected {self.num_channels} but got {num_channels}."
+            )
+        """
+        embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
+        return embeddings
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->Dinov2
+class Dinov2SelfAttention(nn.Module):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
+            config, "embedding_size"
+        ):
+            raise ValueError(
+                f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
+                f"heads {config.num_attention_heads}."
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.attention_probs_dropout_prob = config.attention_probs_dropout_prob
+
+        self.query = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.qkv_bias
+        )
+        self.key = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.qkv_bias
+        )
+        self.value = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.qkv_bias
+        )
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+        new_x_shape = x.size()[:-1] + (
+            self.num_attention_heads,
+            self.attention_head_size,
+        )
+        x = x.view(new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        mixed_query_layer = self.query(hidden_states)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            assert head_mask is None and not output_attentions
+            new_size = hidden_states.size()[:-1] + (
+                self.num_attention_heads,
+                self.attention_head_size,
+            )
+            key_layer = self.key(hidden_states).reshape(new_size).transpose(1, 2)
+            value_layer = self.value(hidden_states).reshape(new_size).transpose(1, 2)
+            query_layer = mixed_query_layer.reshape(new_size).transpose(1, 2)
+            context_layer = F.scaled_dot_product_attention(
+                query_layer,
+                key_layer,
+                value_layer,
+                dropout_p=self.attention_probs_dropout_prob,
+                is_causal=False,
+            )
+            context_layer = context_layer.transpose(1, 2).reshape(
+                *hidden_states.size()[:-1], -1
+            )
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            query_layer = self.transpose_for_scores(mixed_query_layer)
+
+            # Take the dot product between "query" and "key" to get the raw attention scores.
+            attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+            attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+            # Normalize the attention scores to probabilities.
+            attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+            # This is actually dropping out entire tokens to attend to, which might
+            # seem a bit unusual, but is taken from the original Transformer paper.
+            attention_probs = self.dropout(attention_probs)
+
+            # Mask heads if we want to
+            if head_mask is not None:
+                attention_probs = attention_probs * head_mask
+
+            context_layer = torch.matmul(attention_probs, value_layer)
+
+            context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+            new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+            context_layer = context_layer.view(new_context_layer_shape)
+
+        outputs = (
+            (context_layer, attention_probs) if output_attentions else (context_layer,)
+        )
+
+        return outputs
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->Dinov2
+class Dinov2SelfOutput(nn.Module):
+    """
+    The residual connection is defined in Dinov2Layer instead of here (as is the case with other models), due to the
+    layernorm applied before each block.
+    """
+
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self, hidden_states: torch.Tensor, input_tensor: torch.Tensor
+    ) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        return hidden_states
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Dinov2
+class Dinov2Attention(nn.Module):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        self.attention = Dinov2SelfAttention(config)
+        self.output = Dinov2SelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads: Set[int]) -> None:
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads,
+            self.attention.num_attention_heads,
+            self.attention.attention_head_size,
+            self.pruned_heads,
+        )
+
+        # Prune linear layers
+        self.attention.query = prune_linear_layer(self.attention.query, index)
+        self.attention.key = prune_linear_layer(self.attention.key, index)
+        self.attention.value = prune_linear_layer(self.attention.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.attention.num_attention_heads = self.attention.num_attention_heads - len(
+            heads
+        )
+        self.attention.all_head_size = (
+            self.attention.attention_head_size * self.attention.num_attention_heads
+        )
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        self_outputs = self.attention(hidden_states, head_mask, output_attentions)
+
+        attention_output = self.output(self_outputs[0], hidden_states)
+
+        outputs = (attention_output,) + self_outputs[
+            1:
+        ]  # add attentions if we output them
+        return outputs
+
+
+class Dinov2LayerScale(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.lambda1 = nn.Parameter(
+            config.layerscale_value * torch.ones(config.hidden_size)
+        )
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        return hidden_state * self.lambda1
+
+
+# Copied from transformers.models.beit.modeling_beit.drop_path
+def drop_path(
+    input: torch.Tensor, drop_prob: float = 0.0, training: bool = False
+) -> torch.Tensor:
+    """
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,
+    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the
+    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the
+    argument.
+    """
+    if drop_prob == 0.0 or not training:
+        return input
+    keep_prob = 1 - drop_prob
+    shape = (input.shape[0],) + (1,) * (
+        input.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(
+        shape, dtype=input.dtype, device=input.device
+    )
+    random_tensor.floor_()  # binarize
+    output = input.div(keep_prob) * random_tensor
+    return output
+
+
+# Copied from transformers.models.beit.modeling_beit.BeitDropPath
+class Dinov2DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob: Optional[float] = None) -> None:
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return drop_path(hidden_states, self.drop_prob, self.training)
+
+    def extra_repr(self) -> str:
+        return "p={}".format(self.drop_prob)
+
+
+class Dinov2MLP(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        in_features = out_features = config.hidden_size
+        hidden_features = int(config.hidden_size * config.mlp_ratio)
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=True)
+        if isinstance(config.hidden_act, str):
+            self.activation = ACT2FN[config.hidden_act]
+        else:
+            self.activation = config.hidden_act
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=True)
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.fc1(hidden_state)
+        hidden_state = self.activation(hidden_state)
+        hidden_state = self.fc2(hidden_state)
+        return hidden_state
+
+
+class Dinov2SwiGLUFFN(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        in_features = out_features = config.hidden_size
+        hidden_features = int(config.hidden_size * config.mlp_ratio)
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+
+        self.weights_in = nn.Linear(in_features, 2 * hidden_features, bias=True)
+        self.weights_out = nn.Linear(hidden_features, out_features, bias=True)
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.weights_in(hidden_state)
+        x1, x2 = hidden_state.chunk(2, dim=-1)
+        hidden = nn.functional.silu(x1) * x2
+        return self.weights_out(hidden)
+
+
+class Dinov2Layer(nn.Module):
+    """This corresponds to the Block class in the original implementation."""
+
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+
+        self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.norm1_modulation = None
+        self.attention = Dinov2Attention(config)
+        self.layer_scale1 = Dinov2LayerScale(config)
+        self.drop_path1 = (
+            Dinov2DropPath(config.drop_path_rate)
+            if config.drop_path_rate > 0.0
+            else nn.Identity()
+        )
+
+        self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.norm2_modulation = None
+
+        if config.use_swiglu_ffn:
+            self.mlp = Dinov2SwiGLUFFN(config)
+        else:
+            self.mlp = Dinov2MLP(config)
+        self.layer_scale2 = Dinov2LayerScale(config)
+        self.drop_path2 = (
+            Dinov2DropPath(config.drop_path_rate)
+            if config.drop_path_rate > 0.0
+            else nn.Identity()
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        modulation_cond: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
+        hidden_states_norm = self.norm1(hidden_states)
+        if self.norm1_modulation is not None:
+            assert modulation_cond is not None
+            hidden_states_norm = self.norm1_modulation(
+                hidden_states_norm, modulation_cond
+            )
+        self_attention_outputs = self.attention(
+            hidden_states_norm,  # in Dinov2, layernorm is applied before self-attention
+            head_mask,
+            output_attentions=output_attentions,
+        )
+        attention_output = self_attention_outputs[0]
+
+        attention_output = self.layer_scale1(attention_output)
+        outputs = self_attention_outputs[
+            1:
+        ]  # add self attentions if we output attention weights
+
+        # first residual connection
+        hidden_states = attention_output + hidden_states
+
+        # in Dinov2, layernorm is also applied after self-attention
+        layer_output = self.norm2(hidden_states)
+        if self.norm2_modulation is not None:
+            assert modulation_cond is not None
+            layer_output = self.norm2_modulation(layer_output, modulation_cond)
+        layer_output = self.mlp(layer_output)
+        layer_output = self.layer_scale2(layer_output)
+
+        # second residual connection
+        layer_output = layer_output + hidden_states
+
+        outputs = (layer_output,) + outputs
+
+        return outputs
+
+    def register_ada_norm_modulation(self, norm1_mod: nn.Module, norm2_mod: nn.Module):
+        self.norm1_modulation = norm1_mod
+        self.norm2_modulation = norm2_mod
+
+
+# Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->Dinov2
+class Dinov2Encoder(nn.Module):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList(
+            [Dinov2Layer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        modulation_cond: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[tuple, BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    layer_head_mask,
+                    modulation_cond,
+                    use_reentrant=False,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states, layer_head_mask, modulation_cond, output_attentions
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, all_hidden_states, all_self_attentions]
+                if v is not None
+            )
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+class Dinov2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = Dinov2Config
+    base_model_prefix = "dinov2"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid
+            # `trunc_normal_cpu` not implemented in `half` issues
+            module.weight.data = nn.init.trunc_normal_(
+                module.weight.data.to(torch.float32),
+                mean=0.0,
+                std=self.config.initializer_range,
+            ).to(module.weight.dtype)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, Dinov2Embeddings):
+            module.position_embeddings.data = nn.init.trunc_normal_(
+                module.position_embeddings.data.to(torch.float32),
+                mean=0.0,
+                std=self.config.initializer_range,
+            ).to(module.position_embeddings.dtype)
+
+            module.cls_token.data = nn.init.trunc_normal_(
+                module.cls_token.data.to(torch.float32),
+                mean=0.0,
+                std=self.config.initializer_range,
+            ).to(module.cls_token.dtype)
+
+    def _set_gradient_checkpointing(
+        self, module: Dinov2Encoder, value: bool = False
+    ) -> None:
+        if isinstance(module, Dinov2Encoder):
+            module.gradient_checkpointing = value
+
+
+DINOV2_START_DOCSTRING = r"""
+    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
+    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
+    behavior.
+
+    Parameters:
+        config ([`Dinov2Config`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+DINOV2_BASE_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
+            [`BitImageProcessor.preprocess`] for details.
+
+        bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`):
+            Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for
+            pre-training.
+
+        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+DINOV2_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
+            [`BitImageProcessor.preprocess`] for details.
+
+        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@dataclass
+class CustomBaseModelOutputWithPooling(BaseModelOutputWithPooling):
+    patch_embeddings: Optional[torch.FloatTensor] = None
+
+
+@add_start_docstrings(
+    "The bare DINOv2 Model transformer outputting raw hidden-states without any specific head on top.",
+    DINOV2_START_DOCSTRING,
+)
+class Dinov2Model(Dinov2PreTrainedModel):
+    def __init__(self, config: Dinov2Config):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = Dinov2Embeddings(config)
+        self.encoder = Dinov2Encoder(config)
+
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> Dinov2PatchEmbeddings:
+        return self.embeddings.patch_embeddings
+
+    def expand_input_channels(self, extra_input_channels: int) -> None:
+        if extra_input_channels == 0:
+            return
+        conv_old = self.embeddings.patch_embeddings.projection
+        conv_new = nn.Conv2d(
+            self.config.num_channels + extra_input_channels,
+            self.config.hidden_size,
+            kernel_size=self.config.patch_size,
+            stride=self.config.patch_size,
+        ).to(self.device)
+        with torch.no_grad():
+            conv_new.weight[:, :3] = conv_old.weight
+            conv_new.bias = conv_old.bias
+        self.embeddings.patch_embeddings.projection = conv_new
+        del conv_old
+
+    def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    @add_start_docstrings_to_model_forward(DINOV2_BASE_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPooling,
+        config_class=_CONFIG_FOR_DOC,
+        modality="vision",
+        expected_output=_EXPECTED_OUTPUT_SHAPE,
+    )
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        bool_masked_pos: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        modulation_cond: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if pixel_values is None:
+            raise ValueError("You have to specify pixel_values")
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embedding_output = self.embeddings(
+            pixel_values, bool_masked_pos=bool_masked_pos
+        )
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            head_mask=head_mask,
+            modulation_cond=modulation_cond,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        sequence_output = self.layernorm(sequence_output)
+        pooled_output = sequence_output[:, 0, :]
+
+        if not return_dict:
+            head_outputs = (sequence_output, pooled_output)
+            return head_outputs + encoder_outputs[1:]
+
+        return CustomBaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            patch_embeddings=embedding_output,
+        )
+
+    def set_gradient_checkpointing(self, value: bool = False) -> None:
+        self._set_gradient_checkpointing(self.encoder, value)
+
+
+@add_start_docstrings(
+    """
+    Dinov2 Model transformer with an image classification head on top (a linear layer on top of the final hidden state
+    of the [CLS] token) e.g. for ImageNet.
+    """,
+    DINOV2_START_DOCSTRING,
+)
+class Dinov2ForImageClassification(Dinov2PreTrainedModel):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.dinov2 = Dinov2Model(config)
+
+        # Classifier head
+        self.classifier = (
+            nn.Linear(config.hidden_size * 2, config.num_labels)
+            if config.num_labels > 0
+            else nn.Identity()
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_IMAGE_CLASS_CHECKPOINT,
+        output_type=ImageClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[tuple, ImageClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.dinov2(
+            pixel_values,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]  # batch_size, sequence_length, hidden_size
+
+        cls_token = sequence_output[:, 0]
+        patch_tokens = sequence_output[:, 1:]
+
+        linear_input = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1)
+
+        logits = self.classifier(linear_input)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    Dinov2 backbone, to be used with frameworks like DETR and MaskFormer.
+    """,
+    DINOV2_START_DOCSTRING,
+)
+class Dinov2Backbone(Dinov2PreTrainedModel, BackboneMixin):
+    def __init__(self, config):
+        super().__init__(config)
+        super()._init_backbone(config)
+
+        self.num_features = [
+            config.hidden_size for _ in range(config.num_hidden_layers + 1)
+        ]
+        self.embeddings = Dinov2Embeddings(config)
+        self.encoder = Dinov2Encoder(config)
+
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> Dinov2PatchEmbeddings:
+        return self.embeddings.patch_embeddings
+
+    @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> BackboneOutput:
+        """
+        Returns:
+
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, AutoBackbone
+        >>> import torch
+        >>> from PIL import Image
+        >>> import requests
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
+        >>> model = AutoBackbone.from_pretrained(
+        ...     "facebook/dinov2-base", out_features=["stage2", "stage5", "stage8", "stage11"]
+        ... )
+
+        >>> inputs = processor(image, return_tensors="pt")
+
+        >>> outputs = model(**inputs)
+        >>> feature_maps = outputs.feature_maps
+        >>> list(feature_maps[-1].shape)
+        [1, 768, 16, 16]
+        ```"""
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+
+        embedding_output = self.embeddings(pixel_values)
+
+        outputs = self.encoder(
+            embedding_output,
+            output_hidden_states=True,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs.hidden_states if return_dict else outputs[1]
+
+        feature_maps = ()
+        for stage, hidden_state in zip(self.stage_names, hidden_states):
+            if stage in self.out_features:
+                if self.config.apply_layernorm:
+                    hidden_state = self.layernorm(hidden_state)
+                if self.config.reshape_hidden_states:
+                    batch_size, _, height, width = pixel_values.shape
+                    patch_size = self.config.patch_size
+                    hidden_state = hidden_state[:, 1:, :].reshape(
+                        batch_size, width // patch_size, height // patch_size, -1
+                    )
+                    hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
+                feature_maps += (hidden_state,)
+
+        if not return_dict:
+            if output_hidden_states:
+                output = (feature_maps,) + outputs[1:]
+            else:
+                output = (feature_maps,) + outputs[2:]
+            return output
+
+        return BackboneOutput(
+            feature_maps=feature_maps,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            attentions=outputs.attentions if output_attentions else None,
+        )
+
+
+class CustomPatchEmbeddings(nn.Module):
+    """
+    This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+    `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+    Transformer.
+    """
+
+    def __init__(
+        self, image_size: int, patch_size: int, num_channels: int, hidden_size: int
+    ):
+        super().__init__()
+
+        image_size = (
+            image_size
+            if isinstance(image_size, collections.abc.Iterable)
+            else (image_size, image_size)
+        )
+        patch_size = (
+            patch_size
+            if isinstance(patch_size, collections.abc.Iterable)
+            else (patch_size, patch_size)
+        )
+        num_patches = (image_size[1] // patch_size[1]) * (
+            image_size[0] // patch_size[0]
+        )
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+
+        self.projection = nn.Conv2d(
+            num_channels, hidden_size, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        num_channels = pixel_values.shape[1]
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+                f" Expected {self.num_channels} but got {num_channels}."
+            )
+        embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
+        return embeddings
+
+
+class CustomEmbeddings(nn.Module):
+    """
+    Construct the CLS token, mask token, position and patch embeddings.
+    """
+
+    def __init__(
+        self, image_size: int, patch_size: int, num_channels: int, hidden_size: int
+    ) -> None:
+        super().__init__()
+
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.hidden_size = hidden_size
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, self.hidden_size))
+
+        self.patch_embeddings = CustomPatchEmbeddings(
+            image_size, patch_size, num_channels, hidden_size
+        )
+        num_patches = self.patch_embeddings.num_patches
+        self.position_embeddings = nn.Parameter(
+            torch.randn(1, num_patches + 1, self.hidden_size)
+        )
+
+    def interpolate_pos_encoding(
+        self, embeddings: torch.Tensor, height: int, width: int
+    ) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+
+        num_patches = embeddings.shape[1] - 1
+        num_positions = self.position_embeddings.shape[1] - 1
+        if num_patches == num_positions and height == width:
+            return self.position_embeddings
+        class_pos_embed = self.position_embeddings[:, 0]
+        patch_pos_embed = self.position_embeddings[:, 1:]
+        dim = embeddings.shape[-1]
+        height = height // self.patch_size
+        width = width // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        height, width = height + 0.1, width + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(
+            1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim
+        )
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            scale_factor=(
+                height / math.sqrt(num_positions),
+                width / math.sqrt(num_positions),
+            ),
+            mode="bicubic",
+            align_corners=False,
+        )
+        if (
+            int(height) != patch_pos_embed.shape[-2]
+            or int(width) != patch_pos_embed.shape[-1]
+        ):
+            raise ValueError(
+                "Width or height does not match with the interpolated position embeddings"
+            )
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size, _, height, width = pixel_values.shape
+        patch_embeddings = self.patch_embeddings(pixel_values)
+        embeddings = patch_embeddings
+
+        # add the [CLS] token to the embedded patch tokens
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        embeddings = torch.cat((cls_tokens, embeddings), dim=1)
+
+        # add positional encoding to each token
+        embeddings = embeddings + self.interpolate_pos_encoding(
+            embeddings, height, width
+        )
+
+        return embeddings

sf3d/models/tokenizers/image.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+
+from sf3d.models.tokenizers.dinov2 import Dinov2Model
+from sf3d.models.transformers.attention import Modulation
+from sf3d.models.utils import BaseModule
+
+
+class DINOV2SingleImageTokenizer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pretrained_model_name_or_path: str = "facebook/dinov2-large"
+        width: int = 512
+        height: int = 512
+        modulation_cond_dim: int = 768
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.model = Dinov2Model.from_pretrained(self.cfg.pretrained_model_name_or_path)
+
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+        self.model.eval()
+
+        self.model.set_gradient_checkpointing(False)
+
+        # add modulation
+        modulations = []
+        for layer in self.model.encoder.layer:
+            norm1_modulation = Modulation(
+                self.model.config.hidden_size,
+                self.cfg.modulation_cond_dim,
+                zero_init=True,
+                single_layer=True,
+            )
+            norm2_modulation = Modulation(
+                self.model.config.hidden_size,
+                self.cfg.modulation_cond_dim,
+                zero_init=True,
+                single_layer=True,
+            )
+            layer.register_ada_norm_modulation(norm1_modulation, norm2_modulation)
+            modulations += [norm1_modulation, norm2_modulation]
+        self.modulations = nn.ModuleList(modulations)
+
+        self.register_buffer(
+            "image_mean",
+            torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+        self.register_buffer(
+            "image_std",
+            torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+
+    def forward(
+        self,
+        images: Float[Tensor, "B *N C H W"],
+        modulation_cond: Optional[Float[Tensor, "B *N Cc"]],
+        **kwargs,
+    ) -> Float[Tensor, "B *N Ct Nt"]:
+        model = self.model
+
+        packed = False
+        if images.ndim == 4:
+            packed = True
+            images = images.unsqueeze(1)
+            if modulation_cond is not None:
+                assert modulation_cond.ndim == 2
+                modulation_cond = modulation_cond.unsqueeze(1)
+
+        batch_size, n_input_views = images.shape[:2]
+        images = (images - self.image_mean) / self.image_std
+        out = model(
+            rearrange(images, "B N C H W -> (B N) C H W"),
+            modulation_cond=rearrange(modulation_cond, "B N Cc -> (B N) Cc")
+            if modulation_cond is not None
+            else None,
+        )
+        local_features = out.last_hidden_state
+        local_features = local_features.permute(0, 2, 1)
+        local_features = rearrange(
+            local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size
+        )
+        if packed:
+            local_features = local_features.squeeze(1)
+
+        return local_features
+
+    def detokenize(self, *args, **kwargs):
+        raise NotImplementedError

sf3d/models/tokenizers/triplane.py

@@ -0,0 +1,49 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from jaxtyping import Float
+from torch import Tensor
+
+from sf3d.models.utils import BaseModule
+
+
+class TriplaneLearnablePositionalEmbedding(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        plane_size: int = 96
+        num_channels: int = 1024
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.embeddings = nn.Parameter(
+            torch.randn(
+                (3, self.cfg.num_channels, self.cfg.plane_size, self.cfg.plane_size),
+                dtype=torch.float32,
+            )
+            * 1
+            / math.sqrt(self.cfg.num_channels)
+        )
+
+    def forward(self, batch_size: int) -> Float[Tensor, "B Ct Nt"]:
+        return rearrange(
+            repeat(self.embeddings, "Np Ct Hp Wp -> B Np Ct Hp Wp", B=batch_size),
+            "B Np Ct Hp Wp -> B Ct (Np Hp Wp)",
+        )
+
+    def detokenize(
+        self, tokens: Float[Tensor, "B Ct Nt"]
+    ) -> Float[Tensor, "B 3 Ct Hp Wp"]:
+        batch_size, Ct, Nt = tokens.shape
+        assert Nt == self.cfg.plane_size**2 * 3
+        assert Ct == self.cfg.num_channels
+        return rearrange(
+            tokens,
+            "B Ct (Np Hp Wp) -> B Np Ct Hp Wp",
+            Np=3,
+            Hp=self.cfg.plane_size,
+            Wp=self.cfg.plane_size,
+        )

sf3d/models/transformers/attention.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+
+
+class Modulation(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        condition_dim: int,
+        zero_init: bool = False,
+        single_layer: bool = False,
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        if single_layer:
+            self.linear1 = nn.Identity()
+        else:
+            self.linear1 = nn.Linear(condition_dim, condition_dim)
+
+        self.linear2 = nn.Linear(condition_dim, embedding_dim * 2)
+
+        # Only zero init the last linear layer
+        if zero_init:
+            nn.init.zeros_(self.linear2.weight)
+            nn.init.zeros_(self.linear2.bias)
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        emb = self.linear2(self.silu(self.linear1(condition)))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return x

sf3d/models/transformers/backbone.py

@@ -0,0 +1,515 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from sf3d.models.utils import BaseModule
+
+
+class GEGLU(nn.Module):
+    r"""
+    A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+        if gate.device.type != "mps":
+            return F.gelu(gate)
+        # mps: gelu is not implemented for float16
+        return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
+
+    def forward(self, hidden_states, scale: float = 1.0):
+        args = ()
+        hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1)
+        return hidden_states * self.gelu(gate)
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        kv_dim=None,
+        num_heads=16,
+        qkv_bias=False,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+        kv_dim = dim if not kv_dim else kv_dim
+        self.wq = nn.Linear(dim, dim, bias=qkv_bias)
+        self.wk = nn.Linear(kv_dim, dim, bias=qkv_bias)
+        self.wv = nn.Linear(kv_dim, dim, bias=qkv_bias)
+        self.attn_drop = attn_drop
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x_q, x_kv):
+        B, N_q, C = x_q.shape
+        B, N_kv, _ = x_kv.shape
+        # [B, N_q, C] -> [B, N_q, H, C/H]
+        q = self.wq(x_q).reshape(B, N_q, self.num_heads, C // self.num_heads)
+        # [B, N_kv, C] -> [B, N_kv, H, C/H]
+        k = self.wk(x_kv).reshape(B, N_kv, self.num_heads, C // self.num_heads)
+        v = self.wv(x_kv).reshape(B, N_kv, self.num_heads, C // self.num_heads)
+
+        #  attention
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q.permute(0, 2, 1, 3),
+            k.permute(0, 2, 1, 3),
+            v.permute(0, 2, 1, 3),
+            attn_mask=None,
+            dropout_p=self.attn_drop,
+            scale=self.scale,
+        ).permute(0, 2, 1, 3)
+
+        # [B, N_q, H, C/H] -> [B, N_q, C]
+        x = x.reshape(B, N_q, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        act_fn = GEGLU(dim, inner_dim)
+        self.net = nn.ModuleList([])
+        self.net.append(act_fn)
+        self.net.append(nn.Dropout(dropout))
+        self.net.append(nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            x = module(x)
+        return x
+
+
+class BasicBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        kv_dim: Optional[int] = None,
+        num_heads: int = 16,
+        qkv_bias: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        ff_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn1 = CrossAttention(
+            dim,
+            kv_dim=dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+        )
+        self.norm2 = nn.LayerNorm(dim)
+        self.attn2 = CrossAttention(
+            dim,
+            kv_dim=kv_dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+        )
+        self.norm3 = nn.LayerNorm(dim)
+        self.ff = FeedForward(dim, dropout=ff_drop)
+
+    def forward(self, z, x):
+        z_norm = self.norm1(z)
+        z = z + self.attn1(z_norm, z_norm)
+        # TODO: do we need to have the second attention when x is None?
+        z_norm = self.norm2(z)
+        z = z + self.attn2(z_norm, x if x is not None else z_norm)
+        z_norm = self.norm3(z)
+        z = z + self.ff(z_norm)
+        return z
+
+
+class SingleStreamTransformer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        num_attention_heads: int = 16
+        attention_head_dim: int = 88
+        in_channels: Optional[int] = None
+        out_channels: Optional[int] = None
+        num_layers: int = 16
+        dropout: float = 0.0
+        norm_num_groups: int = 32
+        cross_attention_dim: Optional[int] = None
+        attention_bias: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.num_attention_heads = self.cfg.num_attention_heads
+        self.attention_head_dim = self.cfg.attention_head_dim
+        inner_dim = self.num_attention_heads * self.attention_head_dim
+
+        # Define input layers
+        self.norm = torch.nn.GroupNorm(
+            num_groups=self.cfg.norm_num_groups,
+            num_channels=self.cfg.in_channels,
+            eps=1e-6,
+            affine=True,
+        )
+        self.proj_in = nn.Linear(self.cfg.in_channels, inner_dim)
+
+        # Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicBlock(
+                    inner_dim,
+                    kv_dim=self.cfg.cross_attention_dim,
+                    num_heads=self.num_attention_heads,
+                    qkv_bias=self.cfg.attention_bias,
+                    proj_drop=self.cfg.dropout,
+                    ff_drop=self.cfg.dropout,
+                )
+                for d in range(self.cfg.num_layers)
+            ]
+        )
+
+        # 4. Define output layers
+        self.proj_out = nn.Linear(inner_dim, self.cfg.in_channels)
+
+    def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+        residual = hidden_states
+        hidden_states = self.norm(hidden_states)
+        hidden_states = hidden_states.permute(0, 2, 1)
+        hidden_states = self.proj_in(hidden_states)
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, encoder_hidden_states)
+        hidden_states = self.proj_out(hidden_states).permute(0, 2, 1).contiguous()
+        # TODO: do we really need to add the residual?
+        hidden_states = hidden_states + residual
+        return hidden_states
+
+
+class FuseBlock(nn.Module):
+    """
+    Fuse X in to Z with cross attention
+    """
+
+    def __init__(
+        self,
+        dim_z: int,
+        dim_x: int,
+        num_heads: int = 16,
+        qkv_bias: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        ff_drop: float = 0.0,
+        norm_x_input: bool = True,
+    ):
+        super().__init__()
+        self.norm_x_input = norm_x_input
+        if self.norm_x_input:
+            self.norm_x = nn.LayerNorm(dim_x)
+        self.attn = CrossAttention(
+            dim_z,
+            kv_dim=dim_x,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+        )
+        self.norm_z1 = nn.LayerNorm(dim_z)
+        self.norm_z2 = nn.LayerNorm(dim_z)
+        self.ff = FeedForward(dim_z, dropout=ff_drop)
+
+    def forward(self, z, x):
+        # TODO: do we need to normalize x?
+        z = z + self.attn(self.norm_z1(z), self.norm_x(x) if self.norm_x_input else x)
+        z = z + self.ff(self.norm_z2(z))
+        return z
+
+
+@torch.no_grad()
+def get_triplane_attention_mask(res):
+    N = 3 * res * res
+    attn_mask = torch.zeros(3, res, res, 3, res, res)
+
+    i, j = torch.meshgrid(torch.arange(res), torch.arange(res))
+
+    attn_mask[0, i, j, 1, i, :] = 1.0
+    attn_mask[0, i, j, 2, j, :] = 1.0
+    attn_mask[1, i, j, 0, i, :] = 1.0
+    attn_mask[1, i, j, 2, :, j] = 1.0
+    attn_mask[2, i, j, 0, :, i] = 1.0
+    attn_mask[2, i, j, 1, :, j] = 1.0
+    attn_mask = attn_mask.bool()
+
+    attn_bias = torch.empty_like(attn_mask, dtype=torch.float)
+    attn_bias.masked_fill_(attn_mask, 0.0)
+    attn_bias.masked_fill_(~attn_mask, float("-inf"))
+
+    return attn_bias.reshape(N, N)
+
+
+class TriplaneAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        resolution: int,
+        num_heads: int = 16,
+        qkv_bias: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        full_attention: bool = False,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+        self.wq = nn.Linear(dim, dim, bias=qkv_bias)
+        self.wk = nn.Linear(dim, dim, bias=qkv_bias)
+        self.wv = nn.Linear(dim, dim, bias=qkv_bias)
+        self.attn_drop = attn_drop
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.resolution = resolution
+        self.full_attention = full_attention
+        self.attn_mask = (
+            get_triplane_attention_mask(resolution) if not full_attention else None
+        )
+
+    def forward(self, x):
+        B, N, C = x.shape
+        # [B, N, C] -> [B, N, H, C/H]
+        q = self.wq(x).reshape(B, N, self.num_heads, C // self.num_heads)
+        k = self.wk(x).reshape(B, N, self.num_heads, C // self.num_heads)
+        v = self.wv(x).reshape(B, N, self.num_heads, C // self.num_heads)
+
+        # detokenize the planes
+        assert N == self.resolution**2 * 3
+        attn_bias = (
+            self.attn_mask.to(q)
+            .unsqueeze(0)
+            .unsqueeze(0)
+            .expand(B, self.num_heads, -1, -1)
+            if not self.full_attention
+            else None
+        )
+
+        # full attention
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q.permute(0, 2, 1, 3),
+            k.permute(0, 2, 1, 3),
+            v.permute(0, 2, 1, 3),
+            attn_mask=attn_bias,
+            dropout_p=self.attn_drop,
+            scale=self.scale,
+        ).permute(0, 2, 1, 3)
+
+        # [B, N_q, H, C/H] -> [B, N_q, C]
+        x = x.reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class TwoStreamBlock(nn.Module):
+    def __init__(
+        self,
+        dim_latent: int,
+        dim_input: int,
+        num_basic_blocks: int = 4,
+        num_heads: int = 16,
+        qkv_bias: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        ff_drop: float = 0.0,
+        norm_x_input: bool = True,
+        dim_cross: Optional[int] = None,
+    ):
+        super().__init__()
+
+        # Define the fuse block that fuse the input into the latent
+        self.fuse_block_in = FuseBlock(
+            dim_latent,
+            dim_input,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            ff_drop=ff_drop,
+            norm_x_input=norm_x_input,
+        )
+
+        # Define the transformer block that process the latent
+        self.transformer_block = nn.ModuleList(
+            [
+                BasicBlock(
+                    dim_latent,
+                    kv_dim=dim_cross,
+                    num_heads=num_heads,
+                    qkv_bias=qkv_bias,
+                    proj_drop=proj_drop,
+                    ff_drop=ff_drop,
+                )
+                for _ in range(num_basic_blocks)
+            ]
+        )
+
+        # Define the fuse block that fuse the latent into the input
+        self.fuse_block_out = FuseBlock(
+            dim_input,
+            dim_latent,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            ff_drop=ff_drop,
+            norm_x_input=norm_x_input,
+        )
+
+    def forward(self, latent, input, cross_input):
+        latent = self.fuse_block_in(latent, input)
+        for block in self.transformer_block:
+            latent = block(latent, cross_input)
+        input = self.fuse_block_out(input, latent)
+        return latent, input
+
+
+class TwoStreamInterleaveTransformer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        num_attention_heads: int = 16
+        attention_head_dim: int = 64
+        raw_triplane_channels: int = 1024
+        triplane_channels: int = 1024
+        raw_image_channels: int = 1024
+        num_latents: int = 1792
+        num_blocks: int = 4
+        num_basic_blocks: int = 3
+        dropout: float = 0.0
+        latent_init_std: float = 0.02
+        norm_num_groups: int = 32
+        attention_bias: bool = False
+        norm_x_input: bool = False
+        cross_attention_dim: int = 1024
+        mix_latent: bool = True
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.mix_latent = self.cfg.mix_latent
+
+        # Define the dimensions
+        self.num_attention_heads = self.cfg.num_attention_heads
+        self.attention_head_dim = self.cfg.attention_head_dim
+        self.num_latents = self.cfg.num_latents
+        self.latent_dim = self.num_attention_heads * self.attention_head_dim
+
+        # Define input layers
+        if self.cfg.norm_num_groups > 0:
+            self.norm_triplane = torch.nn.GroupNorm(
+                num_groups=self.cfg.norm_num_groups,
+                num_channels=self.cfg.raw_triplane_channels,
+                eps=1e-6,
+                affine=True,
+            )
+        else:
+            self.norm_triplane = nn.LayerNorm(self.cfg.raw_triplane_channels)
+        self.proj_triplane = nn.Linear(
+            self.cfg.raw_triplane_channels, self.cfg.triplane_channels
+        )
+        if self.mix_latent:
+            self.norm_image = nn.LayerNorm(self.cfg.raw_image_channels)
+            self.proj_image = nn.Linear(self.cfg.raw_image_channels, self.latent_dim)
+        self.norm_latent = nn.LayerNorm(self.latent_dim)
+        self.proj_latent = nn.Linear(self.latent_dim, self.latent_dim)
+
+        # Define the latents
+        self.latent_init = nn.Parameter(
+            torch.zeros(1, self.num_latents, self.latent_dim)
+        )
+        nn.init.normal_(self.latent_init, std=self.cfg.latent_init_std)
+
+        # Define the transformer blocks
+        self.main_blocks = nn.ModuleList(
+            [
+                TwoStreamBlock(
+                    self.latent_dim,
+                    self.cfg.triplane_channels,
+                    num_basic_blocks=self.cfg.num_basic_blocks,
+                    num_heads=self.num_attention_heads,
+                    qkv_bias=self.cfg.attention_bias,
+                    proj_drop=self.cfg.dropout,
+                    ff_drop=self.cfg.dropout,
+                    norm_x_input=self.cfg.norm_x_input,
+                    dim_cross=self.cfg.cross_attention_dim,
+                )
+                for _ in range(self.cfg.num_blocks)
+            ]
+        )
+
+        # 4. Define output layers
+        self.proj_out = nn.Linear(
+            self.cfg.triplane_channels, self.cfg.raw_triplane_channels
+        )
+
+    def forward(self, hidden_states, encoder_hidden_states, **kwargs):
+        # hidden_states: [B, triplane_dim, N_triplane] is triplane tokens
+        # encoder_hidden_states: [B, N_image, image_dim] is the image tokens
+        if isinstance(self.norm_triplane, nn.GroupNorm):
+            triplane_tokens = self.norm_triplane(hidden_states)
+            triplane_tokens = triplane_tokens.permute(
+                0, 2, 1
+            )  # [B, N_triplane, triplane_dim]
+        elif isinstance(self.norm_triplane, nn.LayerNorm):
+            triplane_tokens = self.norm_triplane(hidden_states.permute(0, 2, 1))
+        else:
+            raise ValueError("Unknown normalization layer")
+        triplane_tokens = self.proj_triplane(triplane_tokens)
+        if self.mix_latent:
+            image_tokens = self.norm_image(
+                encoder_hidden_states
+            )  # [B, N_image, image_dim]
+            image_tokens = self.proj_image(image_tokens)
+        init_latents = self.latent_init.expand(
+            hidden_states.shape[0], -1, -1
+        )  # [B, N_latent_init, latent_dim]
+        init_latents = self.norm_latent(init_latents)
+        init_latents = self.proj_latent(init_latents)
+        if self.mix_latent:
+            latent_tokens = torch.cat(
+                [image_tokens, init_latents], dim=1
+            )  # [B, N_latent, latent_dim]
+        else:
+            latent_tokens = init_latents
+
+        # forward the main blocks
+        for block in self.main_blocks:
+            latent_tokens, triplane_tokens = block(
+                latent_tokens, triplane_tokens, encoder_hidden_states
+            )
+
+        # project the triplane tokens back to the original dimension
+        triplane_tokens = self.proj_out(triplane_tokens).permute(0, 2, 1).contiguous()
+        triplane_tokens = triplane_tokens + hidden_states
+        return triplane_tokens

sf3d/models/utils.py

@@ -0,0 +1,292 @@
+import dataclasses
+import importlib
+import math
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple, Union
+
+import numpy as np
+import PIL
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from jaxtyping import Bool, Float, Int, Num
+from omegaconf import DictConfig, OmegaConf
+from torch import Tensor
+
+
+class BaseModule(nn.Module):
+    @dataclass
+    class Config:
+        pass
+
+    cfg: Config  # add this to every subclass of BaseModule to enable static type checking
+
+    def __init__(
+        self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
+    ) -> None:
+        super().__init__()
+        self.cfg = parse_structured(self.Config, cfg)
+        self.configure(*args, **kwargs)
+
+    def configure(self, *args, **kwargs) -> None:
+        raise NotImplementedError
+
+
+def find_class(cls_string):
+    module_string = ".".join(cls_string.split(".")[:-1])
+    cls_name = cls_string.split(".")[-1]
+    module = importlib.import_module(module_string, package=None)
+    cls = getattr(module, cls_name)
+    return cls
+
+
+def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
+    # Check if cfg.keys are in fields
+    cfg_ = cfg.copy()
+    keys = list(cfg_.keys())
+
+    field_names = {f.name for f in dataclasses.fields(fields)}
+    for key in keys:
+        # This is helpful when swapping out modules from CLI
+        if key not in field_names:
+            print(f"Ignoring {key} as it's not supported by {fields}")
+            cfg_.pop(key)
+    scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg_)
+    return scfg
+
+
+EPS_DTYPE = {
+    torch.float16: 1e-4,
+    torch.bfloat16: 1e-4,
+    torch.float32: 1e-7,
+    torch.float64: 1e-8,
+}
+
+
+def dot(x, y, dim=-1):
+    return torch.sum(x * y, dim, keepdim=True)
+
+
+def reflect(x, n):
+    return x - 2 * dot(x, n) * n
+
+
+def normalize(x, dim=-1, eps=None):
+    if eps is None:
+        eps = EPS_DTYPE[x.dtype]
+    return F.normalize(x, dim=dim, p=2, eps=eps)
+
+
+def tri_winding(tri: Float[Tensor, "*B 3 2"]) -> Float[Tensor, "*B 3 3"]:
+    # One pad for determinant
+    tri_sq = F.pad(tri, (0, 1), "constant", 1.0)
+    det_tri = torch.det(tri_sq)
+    tri_rev = torch.cat(
+        (tri_sq[..., 0:1, :], tri_sq[..., 2:3, :], tri_sq[..., 1:2, :]), -2
+    )
+    tri_sq[det_tri < 0] = tri_rev[det_tri < 0]
+    return tri_sq
+
+
+def triangle_intersection_2d(
+    t1: Float[Tensor, "*B 3 2"],
+    t2: Float[Tensor, "*B 3 2"],
+    eps=1e-12,
+) -> Float[Tensor, "*B"]:  # noqa: F821
+    """Returns True if triangles collide, False otherwise"""
+
+    def chk_edge(x: Float[Tensor, "*B 3 3"]) -> Bool[Tensor, "*B"]:  # noqa: F821
+        logdetx = torch.logdet(x.double())
+        if eps is None:
+            return ~torch.isfinite(logdetx)
+        return ~(torch.isfinite(logdetx) & (logdetx > math.log(eps)))
+
+    t1s = tri_winding(t1)
+    t2s = tri_winding(t2)
+
+    # Assume the triangles do not collide in the begging
+    ret = torch.zeros(t1.shape[0], dtype=torch.bool, device=t1.device)
+    for i in range(3):
+        edge = torch.roll(t1s, i, dims=1)[:, :2, :]
+        # Check if all points of triangle 2 lay on the external side of edge E.
+        # If this is the case the triangle do not collide
+        upd = (
+            chk_edge(torch.cat((edge, t2s[:, 0:1]), 1))
+            & chk_edge(torch.cat((edge, t2s[:, 1:2]), 1))
+            & chk_edge(torch.cat((edge, t2s[:, 2:3]), 1))
+        )
+        # Here no collision is still True due to inversion
+        ret = ret | upd
+
+    for i in range(3):
+        edge = torch.roll(t2s, i, dims=1)[:, :2, :]
+
+        upd = (
+            chk_edge(torch.cat((edge, t1s[:, 0:1]), 1))
+            & chk_edge(torch.cat((edge, t1s[:, 1:2]), 1))
+            & chk_edge(torch.cat((edge, t1s[:, 2:3]), 1))
+        )
+        # Here no collision is still True due to inversion
+        ret = ret | upd
+
+    return ~ret  # Do the inversion
+
+
+ValidScale = Union[Tuple[float, float], Num[Tensor, "2 D"]]
+
+
+def scale_tensor(
+    dat: Num[Tensor, "... D"], inp_scale: ValidScale, tgt_scale: ValidScale
+):
+    if inp_scale is None:
+        inp_scale = (0, 1)
+    if tgt_scale is None:
+        tgt_scale = (0, 1)
+    if isinstance(tgt_scale, Tensor):
+        assert dat.shape[-1] == tgt_scale.shape[-1]
+    dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0])
+    dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0]
+    return dat
+
+
+def dilate_fill(img, mask, iterations=10):
+    oldMask = mask.float()
+    oldImg = img
+
+    mask_kernel = torch.ones(
+        (1, 1, 3, 3),
+        dtype=oldMask.dtype,
+        device=oldMask.device,
+    )
+
+    for i in range(iterations):
+        newMask = torch.nn.functional.max_pool2d(oldMask, 3, 1, 1)
+
+        # Fill the extension with mean color of old valid regions
+        img_unfold = F.unfold(oldImg, (3, 3)).view(1, 3, 3 * 3, -1)
+        mask_unfold = F.unfold(oldMask, (3, 3)).view(1, 1, 3 * 3, -1)
+        new_mask_unfold = F.unfold(newMask, (3, 3)).view(1, 1, 3 * 3, -1)
+
+        # Average color of the valid region
+        mean_color = (img_unfold.sum(dim=2) / mask_unfold.sum(dim=2).clip(1)).unsqueeze(
+            2
+        )
+        # Extend it to the new region
+        fill_color = (mean_color * new_mask_unfold).view(1, 3 * 3 * 3, -1)
+
+        mask_conv = F.conv2d(
+            newMask, mask_kernel, padding=1
+        )  # Get the sum for each kernel patch
+        newImg = F.fold(
+            fill_color, (img.shape[-2], img.shape[-1]), (3, 3)
+        ) / mask_conv.clamp(1)
+
+        diffMask = newMask - oldMask
+
+        oldMask = newMask
+        oldImg = torch.lerp(oldImg, newImg, diffMask)
+
+    return oldImg
+
+
+def float32_to_uint8_np(
+    x: Float[np.ndarray, "*B H W C"],
+    dither: bool = True,
+    dither_mask: Optional[Float[np.ndarray, "*B H W C"]] = None,
+    dither_strength: float = 1.0,
+) -> Int[np.ndarray, "*B H W C"]:
+    if dither:
+        dither = (
+            dither_strength * np.random.rand(*x[..., :1].shape).astype(np.float32) - 0.5
+        )
+        if dither_mask is not None:
+            dither = dither * dither_mask
+        return np.clip(np.floor((256.0 * x + dither)), 0, 255).astype(np.uint8)
+    return np.clip(np.floor((256.0 * x)), 0, 255).astype(torch.uint8)
+
+
+def convert_data(data):
+    if data is None:
+        return None
+    elif isinstance(data, np.ndarray):
+        return data
+    elif isinstance(data, torch.Tensor):
+        if data.dtype in [torch.float16, torch.bfloat16]:
+            data = data.float()
+        return data.detach().cpu().numpy()
+    elif isinstance(data, list):
+        return [convert_data(d) for d in data]
+    elif isinstance(data, dict):
+        return {k: convert_data(v) for k, v in data.items()}
+    else:
+        raise TypeError(
+            "Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting",
+            type(data),
+        )
+
+
+class ImageProcessor:
+    def convert_and_resize(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+        size: int,
+    ):
+        if isinstance(image, PIL.Image.Image):
+            image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
+        elif isinstance(image, np.ndarray):
+            if image.dtype == np.uint8:
+                image = torch.from_numpy(image.astype(np.float32) / 255.0)
+            else:
+                image = torch.from_numpy(image)
+        elif isinstance(image, torch.Tensor):
+            pass
+
+        batched = image.ndim == 4
+
+        if not batched:
+            image = image[None, ...]
+        image = F.interpolate(
+            image.permute(0, 3, 1, 2),
+            (size, size),
+            mode="bilinear",
+            align_corners=False,
+            antialias=True,
+        ).permute(0, 2, 3, 1)
+        if not batched:
+            image = image[0]
+        return image
+
+    def __call__(
+        self,
+        image: Union[
+            PIL.Image.Image,
+            np.ndarray,
+            torch.FloatTensor,
+            List[PIL.Image.Image],
+            List[np.ndarray],
+            List[torch.FloatTensor],
+        ],
+        size: int,
+    ) -> Any:
+        if isinstance(image, (np.ndarray, torch.FloatTensor)) and image.ndim == 4:
+            image = self.convert_and_resize(image, size)
+        else:
+            if not isinstance(image, list):
+                image = [image]
+            image = [self.convert_and_resize(im, size) for im in image]
+            image = torch.stack(image, dim=0)
+        return image
+
+
+def get_intrinsic_from_fov(fov, H, W, bs=-1):
+    focal_length = 0.5 * H / np.tan(0.5 * fov)
+    intrinsic = np.identity(3, dtype=np.float32)
+    intrinsic[0, 0] = focal_length
+    intrinsic[1, 1] = focal_length
+    intrinsic[0, 2] = W / 2.0
+    intrinsic[1, 2] = H / 2.0
+
+    if bs > 0:
+        intrinsic = intrinsic[None].repeat(bs, axis=0)
+
+    return torch.from_numpy(intrinsic)

sf3d/system.py

@@ -0,0 +1,482 @@
+import os
+from dataclasses import dataclass, field
+from typing import Any, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from jaxtyping import Float
+from omegaconf import OmegaConf
+from PIL import Image
+from safetensors.torch import load_model
+from torch import Tensor
+
+from sf3d.models.isosurface import MarchingTetrahedraHelper
+from sf3d.models.mesh import Mesh
+from sf3d.models.utils import (
+    BaseModule,
+    ImageProcessor,
+    convert_data,
+    dilate_fill,
+    dot,
+    find_class,
+    float32_to_uint8_np,
+    normalize,
+    scale_tensor,
+)
+from sf3d.utils import create_intrinsic_from_fov_deg, default_cond_c2w
+
+from .texture_baker import TextureBaker
+
+
+class SF3D(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        cond_image_size: int
+        isosurface_resolution: int
+        isosurface_threshold: float = 10.0
+        radius: float = 1.0
+        background_color: list[float] = field(default_factory=lambda: [0.5, 0.5, 0.5])
+        default_fovy_deg: float = 40.0
+        default_distance: float = 1.6
+
+        camera_embedder_cls: str = ""
+        camera_embedder: dict = field(default_factory=dict)
+
+        image_tokenizer_cls: str = ""
+        image_tokenizer: dict = field(default_factory=dict)
+
+        tokenizer_cls: str = ""
+        tokenizer: dict = field(default_factory=dict)
+
+        backbone_cls: str = ""
+        backbone: dict = field(default_factory=dict)
+
+        post_processor_cls: str = ""
+        post_processor: dict = field(default_factory=dict)
+
+        decoder_cls: str = ""
+        decoder: dict = field(default_factory=dict)
+
+        image_estimator_cls: str = ""
+        image_estimator: dict = field(default_factory=dict)
+
+        global_estimator_cls: str = ""
+        global_estimator: dict = field(default_factory=dict)
+
+    cfg: Config
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: str, config_name: str, weight_name: str
+    ):
+        if os.path.isdir(pretrained_model_name_or_path):
+            config_path = os.path.join(pretrained_model_name_or_path, config_name)
+            weight_path = os.path.join(pretrained_model_name_or_path, weight_name)
+        else:
+            config_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path, filename=config_name
+            )
+            weight_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path, filename=weight_name
+            )
+
+        cfg = OmegaConf.load(config_path)
+        OmegaConf.resolve(cfg)
+        model = cls(cfg)
+        load_model(model, weight_path)
+        return model
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def configure(self):
+        self.image_tokenizer = find_class(self.cfg.image_tokenizer_cls)(
+            self.cfg.image_tokenizer
+        )
+        self.tokenizer = find_class(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
+        self.camera_embedder = find_class(self.cfg.camera_embedder_cls)(
+            self.cfg.camera_embedder
+        )
+        self.backbone = find_class(self.cfg.backbone_cls)(self.cfg.backbone)
+        self.post_processor = find_class(self.cfg.post_processor_cls)(
+            self.cfg.post_processor
+        )
+        self.decoder = find_class(self.cfg.decoder_cls)(self.cfg.decoder)
+        self.image_estimator = find_class(self.cfg.image_estimator_cls)(
+            self.cfg.image_estimator
+        )
+        self.global_estimator = find_class(self.cfg.global_estimator_cls)(
+            self.cfg.global_estimator
+        )
+
+        self.bbox: Float[Tensor, "2 3"]
+        self.register_buffer(
+            "bbox",
+            torch.as_tensor(
+                [
+                    [-self.cfg.radius, -self.cfg.radius, -self.cfg.radius],
+                    [self.cfg.radius, self.cfg.radius, self.cfg.radius],
+                ],
+                dtype=torch.float32,
+            ),
+        )
+        self.isosurface_helper = MarchingTetrahedraHelper(
+            self.cfg.isosurface_resolution,
+            os.path.join(
+                os.path.dirname(__file__),
+                "..",
+                "load",
+                "tets",
+                f"{self.cfg.isosurface_resolution}_tets.npz",
+            ),
+        )
+
+        self.baker = TextureBaker()
+        self.image_processor = ImageProcessor()
+
+    def triplane_to_meshes(
+        self, triplanes: Float[Tensor, "B 3 Cp Hp Wp"]
+    ) -> list[Mesh]:
+        meshes = []
+        for i in range(triplanes.shape[0]):
+            triplane = triplanes[i]
+            grid_vertices = scale_tensor(
+                self.isosurface_helper.grid_vertices.to(triplanes.device),
+                self.isosurface_helper.points_range,
+                self.bbox,
+            )
+
+            values = self.query_triplane(grid_vertices, triplane)
+            decoded = self.decoder(values, include=["vertex_offset", "density"])
+            sdf = decoded["density"] - self.cfg.isosurface_threshold
+
+            deform = decoded["vertex_offset"].squeeze(0)
+
+            mesh: Mesh = self.isosurface_helper(
+                sdf.view(-1, 1), deform.view(-1, 3) if deform is not None else None
+            )
+            mesh.v_pos = scale_tensor(
+                mesh.v_pos, self.isosurface_helper.points_range, self.bbox
+            )
+
+            meshes.append(mesh)
+
+        return meshes
+
+    def query_triplane(
+        self,
+        positions: Float[Tensor, "*B N 3"],
+        triplanes: Float[Tensor, "*B 3 Cp Hp Wp"],
+    ) -> Float[Tensor, "*B N F"]:
+        batched = positions.ndim == 3
+        if not batched:
+            # no batch dimension
+            triplanes = triplanes[None, ...]
+            positions = positions[None, ...]
+        assert triplanes.ndim == 5 and positions.ndim == 3
+
+        positions = scale_tensor(
+            positions, (-self.cfg.radius, self.cfg.radius), (-1, 1)
+        )
+
+        indices2D: Float[Tensor, "B 3 N 2"] = torch.stack(
+            (positions[..., [0, 1]], positions[..., [0, 2]], positions[..., [1, 2]]),
+            dim=-3,
+        ).to(triplanes.dtype)
+        out: Float[Tensor, "B3 Cp 1 N"] = F.grid_sample(
+            rearrange(triplanes, "B Np Cp Hp Wp -> (B Np) Cp Hp Wp", Np=3).float(),
+            rearrange(indices2D, "B Np N Nd -> (B Np) () N Nd", Np=3).float(),
+            align_corners=True,
+            mode="bilinear",
+        )
+        out = rearrange(out, "(B Np) Cp () N -> B N (Np Cp)", Np=3)
+
+        return out
+
+    def get_scene_codes(self, batch) -> Float[Tensor, "B 3 C H W"]:
+        # if batch[rgb_cond] is only one view, add a view dimension
+        if len(batch["rgb_cond"].shape) == 4:
+            batch["rgb_cond"] = batch["rgb_cond"].unsqueeze(1)
+            batch["mask_cond"] = batch["mask_cond"].unsqueeze(1)
+            batch["c2w_cond"] = batch["c2w_cond"].unsqueeze(1)
+            batch["intrinsic_cond"] = batch["intrinsic_cond"].unsqueeze(1)
+            batch["intrinsic_normed_cond"] = batch["intrinsic_normed_cond"].unsqueeze(1)
+        batch_size, n_input_views = batch["rgb_cond"].shape[:2]
+
+        camera_embeds: Optional[Float[Tensor, "B Nv Cc"]]
+        camera_embeds = self.camera_embedder(**batch)
+
+        input_image_tokens: Float[Tensor, "B Nv Cit Nit"] = self.image_tokenizer(
+            rearrange(batch["rgb_cond"], "B Nv H W C -> B Nv C H W"),
+            modulation_cond=camera_embeds,
+        )
+
+        input_image_tokens = rearrange(
+            input_image_tokens, "B Nv C Nt -> B (Nv Nt) C", Nv=n_input_views
+        )
+
+        tokens: Float[Tensor, "B Ct Nt"] = self.tokenizer(batch_size)
+
+        tokens = self.backbone(
+            tokens,
+            encoder_hidden_states=input_image_tokens,
+            modulation_cond=None,
+        )
+
+        direct_codes = self.tokenizer.detokenize(tokens)
+        scene_codes = self.post_processor(direct_codes)
+        return scene_codes, direct_codes
+
+    def run_image(
+        self,
+        image: Image,
+        bake_resolution: int,
+        estimate_illumination: bool = False,
+    ) -> Tuple[trimesh.Trimesh, dict[str, Any]]:
+        if image.mode != "RGBA":
+            raise ValueError("Image must be in RGBA mode")
+        img_cond = (
+            torch.from_numpy(
+                np.asarray(
+                    image.resize((self.cfg.cond_image_size, self.cfg.cond_image_size))
+                ).astype(np.float32)
+                / 255.0
+            )
+            .float()
+            .clip(0, 1)
+            .to(self.device)
+        )
+        mask_cond = img_cond[:, :, -1:]
+        rgb_cond = torch.lerp(
+            torch.tensor(self.cfg.background_color, device=self.device)[None, None, :],
+            img_cond[:, :, :3],
+            mask_cond,
+        )
+
+        c2w_cond = default_cond_c2w(self.cfg.default_distance).to(self.device)
+        intrinsic, intrinsic_normed_cond = create_intrinsic_from_fov_deg(
+            self.cfg.default_fovy_deg,
+            self.cfg.cond_image_size,
+            self.cfg.cond_image_size,
+        )
+
+        batch = {
+            "rgb_cond": rgb_cond,
+            "mask_cond": mask_cond,
+            "c2w_cond": c2w_cond.unsqueeze(0),
+            "intrinsic_cond": intrinsic.to(self.device).unsqueeze(0),
+            "intrinsic_normed_cond": intrinsic_normed_cond.to(self.device).unsqueeze(0),
+        }
+
+        meshes, global_dict = self.generate_mesh(
+            batch, bake_resolution, estimate_illumination
+        )
+        return meshes[0], global_dict
+
+    def generate_mesh(
+        self,
+        batch,
+        bake_resolution: int,
+        estimate_illumination: bool = False,
+    ) -> Tuple[List[trimesh.Trimesh], dict[str, Any]]:
+        batch["rgb_cond"] = self.image_processor(
+            batch["rgb_cond"], self.cfg.cond_image_size
+        )
+        batch["mask_cond"] = self.image_processor(
+            batch["mask_cond"], self.cfg.cond_image_size
+        )
+        scene_codes, non_postprocessed_codes = self.get_scene_codes(batch)
+
+        global_dict = {}
+        if self.image_estimator is not None:
+            global_dict.update(
+                self.image_estimator(batch["rgb_cond"] * batch["mask_cond"])
+            )
+        if self.global_estimator is not None and estimate_illumination:
+            global_dict.update(self.global_estimator(non_postprocessed_codes))
+
+        with torch.no_grad():
+            with torch.autocast(device_type="cuda", enabled=False):
+                meshes = self.triplane_to_meshes(scene_codes)
+
+                rets = []
+                for i, mesh in enumerate(meshes):
+                    # Check for empty mesh
+                    if mesh.v_pos.shape[0] == 0:
+                        rets.append(trimesh.Trimesh())
+                        continue
+
+                    mesh.unwrap_uv()
+
+                    # Build textures
+                    rast = self.baker.rasterize(
+                        mesh.v_tex, mesh.t_pos_idx, bake_resolution
+                    )
+                    bake_mask = self.baker.get_mask(rast)
+
+                    pos_bake = self.baker.interpolate(
+                        mesh.v_pos,
+                        rast,
+                        mesh.t_pos_idx,
+                        mesh.v_tex,
+                    )
+                    gb_pos = pos_bake[bake_mask]
+
+                    tri_query = self.query_triplane(gb_pos, scene_codes[i])[0]
+                    decoded = self.decoder(
+                        tri_query, exclude=["density", "vertex_offset"]
+                    )
+
+                    nrm = self.baker.interpolate(
+                        mesh.v_nrm,
+                        rast,
+                        mesh.t_pos_idx,
+                        mesh.v_tex,
+                    )
+                    gb_nrm = F.normalize(nrm[bake_mask], dim=-1)
+                    decoded["normal"] = gb_nrm
+
+                    # Check if any keys in global_dict start with decoded_
+                    for k, v in global_dict.items():
+                        if k.startswith("decoder_"):
+                            decoded[k.replace("decoder_", "")] = v[i]
+
+                    mat_out = {
+                        "albedo": decoded["features"],
+                        "roughness": decoded["roughness"],
+                        "metallic": decoded["metallic"],
+                        "normal": normalize(decoded["perturb_normal"]),
+                        "bump": None,
+                    }
+
+                    for k, v in mat_out.items():
+                        if v is None:
+                            continue
+                        if v.shape[0] == 1:
+                            # Skip and directly add a single value
+                            mat_out[k] = v[0]
+                        else:
+                            f = torch.zeros(
+                                bake_resolution,
+                                bake_resolution,
+                                v.shape[-1],
+                                dtype=v.dtype,
+                                device=v.device,
+                            )
+                            if v.shape == f.shape:
+                                continue
+                            if k == "normal":
+                                # Use un-normalized tangents here so that larger smaller tris
+                                # Don't effect the tangents that much
+                                tng = self.baker.interpolate(
+                                    mesh.v_tng,
+                                    rast,
+                                    mesh.t_pos_idx,
+                                    mesh.v_tex,
+                                )
+                                gb_tng = tng[bake_mask]
+                                gb_tng = F.normalize(gb_tng, dim=-1)
+                                gb_btng = F.normalize(
+                                    torch.cross(gb_tng, gb_nrm, dim=-1), dim=-1
+                                )
+                                normal = F.normalize(mat_out["normal"], dim=-1)
+
+                                bump = torch.cat(
+                                    # Check if we have to flip some things
+                                    (
+                                        dot(normal, gb_tng),
+                                        dot(normal, gb_btng),
+                                        dot(normal, gb_nrm).clip(
+                                            0.3, 1
+                                        ),  # Never go below 0.3. This would indicate a flipped (or close to one) normal
+                                    ),
+                                    -1,
+                                )
+                                bump = (bump * 0.5 + 0.5).clamp(0, 1)
+
+                                f[bake_mask] = bump.view(-1, 3)
+                                mat_out["bump"] = f
+                            else:
+                                f[bake_mask] = v.view(-1, v.shape[-1])
+                                mat_out[k] = f
+
+                    def uv_padding(arr):
+                        if arr.ndim == 1:
+                            return arr
+                        return (
+                            dilate_fill(
+                                arr.permute(2, 0, 1)[None, ...],
+                                bake_mask.unsqueeze(0).unsqueeze(0),
+                                iterations=bake_resolution // 150,
+                            )
+                            .squeeze(0)
+                            .permute(1, 2, 0)
+                        )
+
+                    verts_np = convert_data(mesh.v_pos)
+                    faces = convert_data(mesh.t_pos_idx)
+                    uvs = convert_data(mesh.v_tex)
+
+                    basecolor_tex = Image.fromarray(
+                        float32_to_uint8_np(convert_data(uv_padding(mat_out["albedo"])))
+                    ).convert("RGB")
+                    basecolor_tex.format = "JPEG"
+
+                    metallic = mat_out["metallic"].squeeze().cpu().item()
+                    roughness = mat_out["roughness"].squeeze().cpu().item()
+
+                    if "bump" in mat_out and mat_out["bump"] is not None:
+                        bump_np = convert_data(uv_padding(mat_out["bump"]))
+                        bump_up = np.ones_like(bump_np)
+                        bump_up[..., :2] = 0.5
+                        bump_up[..., 2:] = 1
+                        bump_tex = Image.fromarray(
+                            float32_to_uint8_np(
+                                bump_np,
+                                dither=True,
+                                # Do not dither if something is perfectly flat
+                                dither_mask=np.all(
+                                    bump_np == bump_up, axis=-1, keepdims=True
+                                ).astype(np.float32),
+                            )
+                        ).convert("RGB")
+                        bump_tex.format = (
+                            "JPEG"  # PNG would be better but the assets are larger
+                        )
+                    else:
+                        bump_tex = None
+
+                    material = trimesh.visual.material.PBRMaterial(
+                        baseColorTexture=basecolor_tex,
+                        roughnessFactor=roughness,
+                        metallicFactor=metallic,
+                        normalTexture=bump_tex,
+                    )
+
+                    tmesh = trimesh.Trimesh(
+                        vertices=verts_np,
+                        faces=faces,
+                        visual=trimesh.visual.texture.TextureVisuals(
+                            uv=uvs, material=material
+                        ),
+                    )
+                    rot = trimesh.transformations.rotation_matrix(
+                        np.radians(-90), [1, 0, 0]
+                    )
+                    tmesh.apply_transform(rot)
+                    tmesh.apply_transform(
+                        trimesh.transformations.rotation_matrix(
+                            np.radians(90), [0, 1, 0]
+                        )
+                    )
+
+                    tmesh.invert()
+
+                    rets.append(tmesh)
+
+        return rets, global_dict

sf3d/texture_baker.py

@@ -0,0 +1,87 @@
+import os
+
+import slangtorch
+import torch
+import torch.nn as nn
+from jaxtyping import Bool, Float
+from torch import Tensor
+
+
+class TextureBaker(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.baker = slangtorch.loadModule(
+            os.path.join(os.path.dirname(__file__), "texture_baker.slang")
+        )
+
+    def rasterize(
+        self,
+        uv: Float[Tensor, "Nv 2"],
+        face_indices: Float[Tensor, "Nf 3"],
+        bake_resolution: int,
+    ) -> Float[Tensor, "bake_resolution bake_resolution 4"]:
+        if not face_indices.is_cuda or not uv.is_cuda:
+            raise ValueError("All input tensors must be on cuda")
+
+        face_indices = face_indices.to(torch.int32)
+        uv = uv.to(torch.float32)
+
+        rast_result = torch.empty(
+            bake_resolution, bake_resolution, 4, device=uv.device, dtype=torch.float32
+        )
+
+        block_size = 16
+        grid_size = bake_resolution // block_size
+        self.baker.bake_uv(uv=uv, indices=face_indices, output=rast_result).launchRaw(
+            blockSize=(block_size, block_size, 1), gridSize=(grid_size, grid_size, 1)
+        )
+
+        return rast_result
+
+    def get_mask(
+        self, rast: Float[Tensor, "bake_resolution bake_resolution 4"]
+    ) -> Bool[Tensor, "bake_resolution bake_resolution"]:
+        return rast[..., -1] >= 0
+
+    def interpolate(
+        self,
+        attr: Float[Tensor, "Nv 3"],
+        rast: Float[Tensor, "bake_resolution bake_resolution 4"],
+        face_indices: Float[Tensor, "Nf 3"],
+        uv: Float[Tensor, "Nv 2"],
+    ) -> Float[Tensor, "bake_resolution bake_resolution 3"]:
+        # Make sure all input tensors are on torch
+        if not attr.is_cuda or not face_indices.is_cuda or not rast.is_cuda:
+            raise ValueError("All input tensors must be on cuda")
+
+        attr = attr.to(torch.float32)
+        face_indices = face_indices.to(torch.int32)
+        uv = uv.to(torch.float32)
+
+        pos_bake = torch.zeros(
+            rast.shape[0],
+            rast.shape[1],
+            3,
+            device=attr.device,
+            dtype=attr.dtype,
+        )
+
+        block_size = 16
+        grid_size = rast.shape[0] // block_size
+        self.baker.interpolate(
+            attr=attr, indices=face_indices, rast=rast, output=pos_bake
+        ).launchRaw(
+            blockSize=(block_size, block_size, 1), gridSize=(grid_size, grid_size, 1)
+        )
+
+        return pos_bake
+
+    def forward(
+        self,
+        attr: Float[Tensor, "Nv 3"],
+        uv: Float[Tensor, "Nv 2"],
+        face_indices: Float[Tensor, "Nf 3"],
+        bake_resolution: int,
+    ) -> Float[Tensor, "bake_resolution bake_resolution 3"]:
+        rast = self.rasterize(uv, face_indices, bake_resolution)
+        return self.interpolate(attr, rast, face_indices, uv)

sf3d/texture_baker.slang

@@ -0,0 +1,93 @@
+// xy: 2D test position
+// v1: vertex position 1
+// v2: vertex position 2
+// v3: vertex position 3
+//
+bool barycentric_coordinates(float2 xy, float2 v1, float2 v2, float2 v3, out float u, out float v, out float w)
+{
+    // Return true if the point (x,y) is inside the triangle defined by the vertices v1, v2, v3.
+    // If the point is inside the triangle, the barycentric coordinates are stored in u, v, and w.
+    float2 v1v2 = v2 - v1;
+    float2 v1v3 = v3 - v1;
+    float2 xyv1 = xy - v1;
+
+    float d00 = dot(v1v2, v1v2);
+    float d01 = dot(v1v2, v1v3);
+    float d11 = dot(v1v3, v1v3);
+    float d20 = dot(xyv1, v1v2);
+    float d21 = dot(xyv1, v1v3);
+
+    float denom = d00 * d11 - d01 * d01;
+    v = (d11 * d20 - d01 * d21) / denom;
+    w = (d00 * d21 - d01 * d20) / denom;
+    u = 1.0 - v - w;
+
+    return (v >= 0.0) && (w >= 0.0) && (v + w <= 1.0);
+}
+
+[AutoPyBindCUDA]
+[CUDAKernel]
+void interpolate(
+    TensorView<float3> attr,
+    TensorView<int3> indices,
+    TensorView<float4> rast,
+    TensorView<float3> output)
+{
+    // Interpolate the attr into output based on the rast result (barycentric coordinates, + triangle idx)
+
+    uint3 dispatch_id = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();
+
+    if (dispatch_id.x > output.size(0) || dispatch_id.y > output.size(1))
+        return;
+
+    float4 barycentric = rast[dispatch_id.x, dispatch_id.y];
+    int triangle_idx = int(barycentric.w);
+
+    if (triangle_idx < 0) {
+        output[dispatch_id.x, dispatch_id.y] = float3(0.0, 0.0, 0.0);
+        return;
+    }
+
+    float3 v1 = attr[indices[triangle_idx].x];
+    float3 v2 = attr[indices[triangle_idx].y];
+    float3 v3 = attr[indices[triangle_idx].z];
+
+    output[dispatch_id.x, dispatch_id.y] = v1 * barycentric.x + v2 * barycentric.y + v3 * barycentric.z;
+}
+
+[AutoPyBindCUDA]
+[CUDAKernel]
+void bake_uv(
+    TensorView<float2> uv,
+    TensorView<int3> indices,
+    TensorView<float4> output)
+{
+    uint3 dispatch_id = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();
+
+    if (dispatch_id.y > output.size(0) || dispatch_id.x > output.size(1))
+        return;
+
+    // We index x,y but the orginal coords are HW. So swap them
+    float2 pixel_coord = float2(dispatch_id.y, dispatch_id.x);
+    // Normalize to [0, 1]
+    pixel_coord /= float2(output.size(1), output.size(0));
+    pixel_coord = clamp(pixel_coord, 0.0, 1.0);
+    // Flip x-axis
+    pixel_coord.y = 1 - pixel_coord.y;
+
+    for (int i = 0; i < indices.size(0); i++) {
+        float2 v1 = float2(uv[indices[i].x].x, uv[indices[i].x].y);
+        float2 v2 = float2(uv[indices[i].y].x, uv[indices[i].y].y);
+        float2 v3 = float2(uv[indices[i].z].x, uv[indices[i].z].y);
+
+        float u, v, w;
+        bool hit = barycentric_coordinates(pixel_coord, v1, v2, v3, u, v, w);
+
+        if (hit){
+            output[dispatch_id.x, dispatch_id.y] = float4(u, v, w, i);
+            return;
+        }
+    }
+
+    output[dispatch_id.x, dispatch_id.y] = float4(0.0, 0.0, 0.0, -1);
+}

sf3d/utils.py

@@ -0,0 +1,91 @@
+from typing import Any
+
+import numpy as np
+import rembg
+import torch
+from PIL import Image
+
+import sf3d.models.utils as sf3d_utils
+
+
+def create_intrinsic_from_fov_deg(fov_deg: float, cond_height: int, cond_width: int):
+    intrinsic = sf3d_utils.get_intrinsic_from_fov(
+        np.deg2rad(fov_deg),
+        H=cond_height,
+        W=cond_width,
+    )
+    intrinsic_normed_cond = intrinsic.clone()
+    intrinsic_normed_cond[..., 0, 2] /= cond_width
+    intrinsic_normed_cond[..., 1, 2] /= cond_height
+    intrinsic_normed_cond[..., 0, 0] /= cond_width
+    intrinsic_normed_cond[..., 1, 1] /= cond_height
+
+    return intrinsic, intrinsic_normed_cond
+
+
+def default_cond_c2w(distance: float):
+    c2w_cond = torch.as_tensor(
+        [
+            [0, 0, 1, distance],
+            [1, 0, 0, 0],
+            [0, 1, 0, 0],
+            [0, 0, 0, 1],
+        ]
+    ).float()
+    return c2w_cond
+
+
+def remove_background(
+    image: Image,
+    rembg_session: Any = None,
+    force: bool = False,
+    **rembg_kwargs,
+) -> Image:
+    do_remove = True
+    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
+        do_remove = False
+    do_remove = do_remove or force
+    if do_remove:
+        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
+    return image
+
+
+def resize_foreground(
+    image: Image,
+    ratio: float,
+) -> Image:
+    image = np.array(image)
+    assert image.shape[-1] == 4
+    alpha = np.where(image[..., 3] > 0)
+    y1, y2, x1, x2 = (
+        alpha[0].min(),
+        alpha[0].max(),
+        alpha[1].min(),
+        alpha[1].max(),
+    )
+    # crop the foreground
+    fg = image[y1:y2, x1:x2]
+    # pad to square
+    size = max(fg.shape[0], fg.shape[1])
+    ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
+    ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
+    new_image = np.pad(
+        fg,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+
+    # compute padding according to the ratio
+    new_size = int(new_image.shape[0] / ratio)
+    # pad to size, double side
+    ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
+    ph1, pw1 = new_size - size - ph0, new_size - size - pw0
+    new_image = np.pad(
+        new_image,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+    new_image = Image.fromarray(new_image, mode="RGBA")
+    return new_image

stable_fast.py

@@ -0,0 +1,355 @@
+import os
+import tempfile
+import time
+from functools import lru_cache
+from typing import Any
+
+import gradio as gr
+import numpy as np
+import rembg
+import torch
+from gradio_litmodel3d import LitModel3D
+from PIL import Image
+
+import sf3d.utils as sf3d_utils
+from sf3d.system import SF3D
+
+rembg_session = rembg.new_session()
+
+COND_WIDTH = 512
+COND_HEIGHT = 512
+COND_DISTANCE = 1.6
+COND_FOVY_DEG = 40
+BACKGROUND_COLOR = [0.5, 0.5, 0.5]
+
+# Cached. Doesn't change
+c2w_cond = sf3d_utils.default_cond_c2w(COND_DISTANCE)
+intrinsic, intrinsic_normed_cond = sf3d_utils.create_intrinsic_from_fov_deg(
+    COND_FOVY_DEG, COND_HEIGHT, COND_WIDTH
+)
+
+
+model = SF3D.from_pretrained(
+    "stabilityai/stable-fast-3d",
+    config_name="config.yaml",
+    weight_name="model.safetensors",
+)
+model.eval().cuda()
+
+example_files = [
+    os.path.join("demo_files/examples", f) for f in os.listdir("demo_files/examples")
+]
+
+
+def run_model(input_image):
+    start = time.time()
+    with torch.no_grad():
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            model_batch = create_batch(input_image)
+            model_batch = {k: v.cuda() for k, v in model_batch.items()}
+            trimesh_mesh, _glob_dict = model.generate_mesh(model_batch, 1024)
+            trimesh_mesh = trimesh_mesh[0]
+
+    # Create new tmp file
+    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".glb")
+
+    trimesh_mesh.export(tmp_file.name, file_type="glb", include_normals=True)
+
+    print("Generation took:", time.time() - start, "s")
+
+    return tmp_file.name
+
+
+def create_batch(input_image: Image) -> dict[str, Any]:
+    img_cond = (
+        torch.from_numpy(
+            np.asarray(input_image.resize((COND_WIDTH, COND_HEIGHT))).astype(np.float32)
+            / 255.0
+        )
+        .float()
+        .clip(0, 1)
+    )
+    mask_cond = img_cond[:, :, -1:]
+    rgb_cond = torch.lerp(
+        torch.tensor(BACKGROUND_COLOR)[None, None, :], img_cond[:, :, :3], mask_cond
+    )
+
+    batch_elem = {
+        "rgb_cond": rgb_cond,
+        "mask_cond": mask_cond,
+        "c2w_cond": c2w_cond.unsqueeze(0),
+        "intrinsic_cond": intrinsic.unsqueeze(0),
+        "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
+    }
+    # Add batch dim
+    batched = {k: v.unsqueeze(0) for k, v in batch_elem.items()}
+    return batched
+
+
+@lru_cache
+def checkerboard(squares: int, size: int, min_value: float = 0.5):
+    base = np.zeros((squares, squares)) + min_value
+    base[1::2, ::2] = 1
+    base[::2, 1::2] = 1
+
+    repeat_mult = size // squares
+    return (
+        base.repeat(repeat_mult, axis=0)
+        .repeat(repeat_mult, axis=1)[:, :, None]
+        .repeat(3, axis=-1)
+    )
+
+
+def remove_background(input_image: Image) -> Image:
+    return rembg.remove(input_image, session=rembg_session)
+
+
+def resize_foreground(
+    image: Image,
+    ratio: float,
+) -> Image:
+    image = np.array(image)
+    assert image.shape[-1] == 4
+    alpha = np.where(image[..., 3] > 0)
+    y1, y2, x1, x2 = (
+        alpha[0].min(),
+        alpha[0].max(),
+        alpha[1].min(),
+        alpha[1].max(),
+    )
+    # crop the foreground
+    fg = image[y1:y2, x1:x2]
+    # pad to square
+    size = max(fg.shape[0], fg.shape[1])
+    ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
+    ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
+    new_image = np.pad(
+        fg,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+
+    # compute padding according to the ratio
+    new_size = int(new_image.shape[0] / ratio)
+    # pad to size, double side
+    ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
+    ph1, pw1 = new_size - size - ph0, new_size - size - pw0
+    new_image = np.pad(
+        new_image,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+    new_image = Image.fromarray(new_image, mode="RGBA").resize(
+        (COND_WIDTH, COND_HEIGHT)
+    )
+    return new_image
+
+
+def square_crop(input_image: Image) -> Image:
+    # Perform a center square crop
+    min_size = min(input_image.size)
+    left = (input_image.size[0] - min_size) // 2
+    top = (input_image.size[1] - min_size) // 2
+    right = (input_image.size[0] + min_size) // 2
+    bottom = (input_image.size[1] + min_size) // 2
+    return input_image.crop((left, top, right, bottom)).resize(
+        (COND_WIDTH, COND_HEIGHT)
+    )
+
+
+def show_mask_img(input_image: Image) -> Image:
+    img_numpy = np.array(input_image)
+    alpha = img_numpy[:, :, 3] / 255.0
+    chkb = checkerboard(32, 512) * 255
+    new_img = img_numpy[..., :3] * alpha[:, :, None] + chkb * (1 - alpha[:, :, None])
+    return Image.fromarray(new_img.astype(np.uint8), mode="RGB")
+
+
+def run_button(run_btn, input_image, background_state, foreground_ratio):
+    if run_btn == "Run":
+        glb_file: str = run_model(background_state)
+
+        return (
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update(value=glb_file, visible=True),
+            gr.update(visible=True),
+        )
+    elif run_btn == "Remove Background":
+        rem_removed = remove_background(input_image)
+
+        sqr_crop = square_crop(rem_removed)
+        fr_res = resize_foreground(sqr_crop, foreground_ratio)
+
+        return (
+            gr.update(value="Run", visible=True),
+            sqr_crop,
+            fr_res,
+            gr.update(value=show_mask_img(fr_res), visible=True),
+            gr.update(value=None, visible=False),
+            gr.update(visible=False),
+        )
+
+
+def requires_bg_remove(image, fr):
+    if image is None:
+        return (
+            gr.update(visible=False, value="Run"),
+            None,
+            None,
+            gr.update(value=None, visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    alpha_channel = np.array(image.getchannel("A"))
+    min_alpha = alpha_channel.min()
+
+    if min_alpha == 0:
+        print("Already has alpha")
+        sqr_crop = square_crop(image)
+        fr_res = resize_foreground(sqr_crop, fr)
+        return (
+            gr.update(value="Run", visible=True),
+            sqr_crop,
+            fr_res,
+            gr.update(value=show_mask_img(fr_res), visible=True),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    return (
+        gr.update(value="Remove Background", visible=True),
+        None,
+        None,
+        gr.update(value=None, visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False),
+    )
+
+
+def update_foreground_ratio(img_proc, fr):
+    foreground_res = resize_foreground(img_proc, fr)
+    return (
+        foreground_res,
+        gr.update(value=show_mask_img(foreground_res)),
+    )
+
+
+with gr.Blocks() as demo:
+    img_proc_state = gr.State()
+    background_remove_state = gr.State()
+    gr.Markdown("""
+    # SF3D: Stable Fast 3D Mesh Reconstruction with UV-unwrapping and Illumination Disentanglement
+
+    **SF3D** is a state-of-the-art method for 3D mesh reconstruction from a single image.
+    This demo allows you to upload an image and generate a 3D mesh model from it.
+
+    **Tips**
+    1. If the image already has an alpha channel, you can skip the background removal step.
+    2. You can adjust the foreground ratio to control the size of the foreground object. This can influence the shape
+    3. You can upload your own HDR environment map to light the 3D model.
+    """)
+    with gr.Row(variant="panel"):
+        with gr.Column():
+            with gr.Row():
+                input_img = gr.Image(
+                    type="pil", label="Input Image", sources="upload", image_mode="RGBA"
+                )
+                preview_removal = gr.Image(
+                    label="Preview Background Removal",
+                    type="pil",
+                    image_mode="RGB",
+                    interactive=False,
+                    visible=False,
+                )
+
+            foreground_ratio = gr.Slider(
+                label="Foreground Ratio",
+                minimum=0.5,
+                maximum=1.0,
+                value=0.85,
+                step=0.05,
+            )
+
+            foreground_ratio.change(
+                update_foreground_ratio,
+                inputs=[img_proc_state, foreground_ratio],
+                outputs=[background_remove_state, preview_removal],
+            )
+
+            run_btn = gr.Button("Run", variant="primary", visible=False)
+
+        with gr.Column():
+            output_3d = LitModel3D(
+                label="3D Model",
+                visible=False,
+                clear_color=[0.0, 0.0, 0.0, 0.0],
+                tonemapping="aces",
+                contrast=1.0,
+                scale=1.0,
+            )
+            with gr.Column(visible=False, scale=1.0) as hdr_row:
+                gr.Markdown("""## HDR Environment Map
+
+                Select an HDR environment map to light the 3D model. You can also upload your own HDR environment maps.
+                """)
+
+                with gr.Row():
+                    hdr_illumination_file = gr.File(
+                        label="HDR Env Map", file_types=[".hdr"], file_count="single"
+                    )
+                    example_hdris = [
+                        os.path.join("demo_files/hdri", f)
+                        for f in os.listdir("demo_files/hdri")
+                    ]
+                    hdr_illumination_example = gr.Examples(
+                        examples=example_hdris,
+                        inputs=hdr_illumination_file,
+                    )
+
+                    hdr_illumination_file.change(
+                        lambda x: gr.update(env_map=x.name if x is not None else None),
+                        inputs=hdr_illumination_file,
+                        outputs=[output_3d],
+                    )
+
+    examples = gr.Examples(
+        examples=example_files,
+        inputs=input_img,
+    )
+
+    input_img.change(
+        requires_bg_remove,
+        inputs=[input_img, foreground_ratio],
+        outputs=[
+            run_btn,
+            img_proc_state,
+            background_remove_state,
+            preview_removal,
+            output_3d,
+            hdr_row,
+        ],
+    )
+
+    run_btn.click(
+        run_button,
+        inputs=[
+            run_btn,
+            input_img,
+            background_remove_state,
+            foreground_ratio,
+        ],
+        outputs=[
+            run_btn,
+            img_proc_state,
+            background_remove_state,
+            preview_removal,
+            output_3d,
+            hdr_row,
+        ],
+    )
+
+demo.launch()