init

LICENSE

@@ -0,0 +1,35 @@
+S-Lab License 1.0
+
+Copyright 2025 S-Lab
+
+Redistribution and use for non-commercial purpose in source and 
+binary forms, with or without modification, are permitted provided 
+that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright 
+   notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright 
+   notice, this list of conditions and the following disclaimer in 
+   the documentation and/or other materials provided with the 
+   distribution.
+
+3. Neither the name of the copyright holder nor the names of its 
+   contributors may be used to endorse or promote products derived 
+   from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+In the event that redistribution and/or use for commercial purpose in 
+source or binary forms, with or without modification is required, 
+please contact the contributor(s) of the work.

app.py

@@ -0,0 +1,227 @@
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+import spaces  # Import spaces for ZeroGPU compatibility
+from einops import rearrange
+
+import math
+import torch
+import argparse
+from PIL import Image
+from einops import rearrange
+from mmengine.config import Config
+from xtuner.registry import BUILDER
+from xtuner.model.utils import guess_load_checkpoint
+
+from scripts.camera.cam_dataset import Cam_Generator
+
+##### load model
+config = "configs/pipelines/stage_2_base.py"
+config = Config.fromfile(config)
+model = BUILDER.build(config.model).eval()
+checkpoint_path = "checkpoints/Puffin-Base.pth"
+state_dict = guess_load_checkpoint(checkpoint_path)
+model.load_state_dict(state_dict, strict=False)
+
+if torch.cuda.is_available():
+    model = model.to(torch.bfloat16).cuda()
+else:
+    model = model.to(torch.float32)
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+@torch.inference_mode()
+@spaces.GPU(duration=120) 
+# Multimodal Understanding function
+def multimodal_understanding(image, question, seed, top_p, temperature, progress=gr.Progress(track_tqdm=True)):
+    # Clear CUDA cache before generating
+    torch.cuda.empty_cache()
+    
+    # set seed
+    # torch.manual_seed(seed)
+    # np.random.seed(seed)
+    # torch.cuda.manual_seed(seed)
+    print(torch.cuda.is_available())
+
+    max_new_tokens = 512
+    image_size = 512
+    '''
+    assert image_size == 512
+    image = Image.fromarray(image).convert('RGB')
+    image = expand2square(
+        image, (127, 127, 127))
+    image = image.resize(size=(image_size, image_size))
+    image = torch.from_numpy(np.array(image)).to(dtype=model.dtype, device=model.device)
+    image = rearrange(image, 'h w c -> c h w')[None]
+    image = 2 * (image / 255) - 1
+
+    prompt = PROMPT_TEMPLATE['INSTRUCTION'].format(input="<image>\n" + question)
+    assert '<image>' in prompt
+    image_length = (image_size // 16) ** 2 + harmon_model.mar.buffer_size
+    prompt = prompt.replace('<image>', '<image>' * image_length)
+    input_ids = harmon_tokenizer.encode(
+        prompt, add_special_tokens=True, return_tensors='pt').to(harmon_model.device)
+    _, z_enc = harmon_model.extract_visual_feature(harmon_model.encode(image))
+    inputs_embeds = z_enc.new_zeros(*input_ids.shape, harmon_model.llm.config.hidden_size)
+    inputs_embeds[input_ids == image_token_idx] = z_enc.flatten(0, 1)
+    inputs_embeds[input_ids != image_token_idx] = harmon_model.llm.get_input_embeddings()(
+        input_ids[input_ids != image_token_idx]
+    )
+    output = harmon_model.llm.generate(inputs_embeds=inputs_embeds,
+                                       eos_token_id=harmon_tokenizer.eos_token_id,
+                                       pad_token_id=harmon_tokenizer.pad_token_id
+                                       if harmon_tokenizer.pad_token_id is not None else
+                                       harmon_tokenizer.eos_token_id,
+                                       max_new_tokens=max_new_tokens,
+                                       do_sample=False,  # if temperature == 0 else True,
+                                       use_cache=True,
+                                       # temperature=temperature,
+                                       # top_p=top_p
+                                       )
+    '''
+    return 1#harmon_tokenizer.decode(output[0],  skip_special_tokens=True)
+
+
+@torch.inference_mode()
+@spaces.GPU(duration=120)  # Specify a duration to avoid timeout
+def generate_image(prompt_scene,
+                   seed=42,
+                   roll=3,
+                   pitch=1.0,
+                   fov=1.0,
+                   progress=gr.Progress(track_tqdm=True)):
+    # Clear CUDA cache and avoid tracking gradients
+    torch.cuda.empty_cache()
+    # Set the seed for reproducible results
+    # if seed is not None:
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    np.random.seed(seed)
+    print(torch.cuda.is_available())
+    
+    generator = torch.Generator().manual_seed(seed)
+    prompt_camera = (
+        "The camera parameters (roll, pitch, and field-of-view) are: "
+        f"{roll:.4f}, {pitch:.4f}, {fov:.4f}."
+    )
+    gen = Cam_Generator()
+    cam_map = gen.get_cam(prompt_camera).to(model.device)
+    cam_map = cam_map / (math.pi / 2)
+    
+    prompt = prompt_scene + " " + prompt_camera
+    
+    
+    bsz = 4
+    with torch.no_grad():
+        images, output_reasoning = model.generate(
+            prompt=[prompt]*bsz,
+            cfg_prompt=[""]*bsz,
+            pixel_values_init=None,
+            cfg_scale=4.5,
+            num_steps=50,
+            cam_values=[[cam_map]]*bsz,
+            progress_bar=False,
+            reasoning=False,
+            prompt_reasoning=[""]*bsz,
+            generator=generator,
+            height=512,
+            width=512
+        )
+
+        images = rearrange(images, 'b c h w -> b h w c')
+        images = torch.clamp(127.5 * images + 128.0, 0, 255).to("cpu", dtype=torch.uint8).numpy()
+        ret_images = [Image.fromarray(image) for image in images]
+        return ret_images
+
+
+
+# Gradio interface
+css = '''
+.gradio-container {max-width: 960px !important}
+'''
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("# Puffin")
+
+    with gr.Tab("Camera-controllable Image Generation"):
+        gr.Markdown(value="## Camera-controllable Image Generation")
+
+        prompt_input = gr.Textbox(label="Prompt.")
+
+        with gr.Accordion("Camera Parameters", open=True):
+            with gr.Row():
+                roll = gr.Slider(minimum=-0.7854, maximum=0.7854, value=0.1000, step=0.1000, label="roll value")
+                pitch = gr.Slider(minimum=-0.7854, maximum=0.7854, value=-0.1000, step=0.1000, label="pitch value")
+                fov = gr.Slider(minimum=0.3491, maximum=1.8326, value=1.5000, step=0.1000, label="fov value")
+            seed_input = gr.Number(label="Seed (Optional)", precision=0, value=1234)
+            
+        generation_button = gr.Button("Generate Images")
+    
+        image_output = gr.Gallery(label="Generated Images", columns=4, rows=1)
+    
+        examples_t2i = gr.Examples(
+            label="Prompt examples.",
+            examples=[
+                "A sunny day casts light on two warmly colored buildings—yellow with green accents and deeper orange—framed by a lush green tree, with a blue sign and street lamp adding details in the foreground.",
+                "A high-vantage-point view of lush, autumn-colored mountains blanketed in green and gold, set against a clear blue sky with scattered white clouds, offering a tranquil and breathtaking vista of a serene valley below.",
+                "A grand, historic castle with pointed spires and elaborate stone structures stands against a clear blue sky, flanked by a circular fountain, vibrant red flowers, and neatly trimmed hedges in a beautifully landscaped garden.",
+                "A serene aerial view of a coastal landscape at sunrise/sunset, featuring warm pink and orange skies transitioning to cool blues, with calm waters stretching to rugged, snow-capped mountains in the background, creating a tranquil and picturesque scene.",
+                "A worn, light-yellow walls room with herringbone terracotta floors and three large arched windows framed in pink trim and white panes, showcasing signs of age and disrepair, overlooks a residential area through glimpses of greenery and neighboring buildings.",
+            ],
+            inputs=prompt_input,
+        )
+
+    with gr.Tab("Multimodal Understanding"):
+        gr.Markdown(value="## Multimodal Understanding")
+        image_input = gr.Image()
+        with gr.Column():
+            question_input = gr.Textbox(label="Question")
+
+        understanding_button = gr.Button("Chat")
+        understanding_output = gr.Textbox(label="Response")
+
+        with gr.Accordion("Advanced options", open=False):
+            und_seed_input = gr.Number(label="Seed", precision=0, value=42)
+            top_p = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="top_p")
+            temperature = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.05, label="temperature")
+
+        examples_inpainting = gr.Examples(
+            label="Multimodal Understanding examples",
+            examples=[
+                [
+                    "Is the picture taken in winter?",
+                    "view.jpg",
+                ],
+                [
+                    "Briefly describe the image.",
+                    "view.jpg",
+                ],
+            ],
+            inputs=[question_input, image_input],
+        )
+
+    generation_button.click(
+        fn=generate_image,
+        inputs=[prompt_input, seed_input, roll, pitch, fov],
+        outputs=image_output
+    )
+
+    understanding_button.click(
+        multimodal_understanding,
+        inputs=[image_input, question_input, und_seed_input, top_p, temperature],
+        outputs=understanding_output
+    )
+
+demo.launch(share=True)

configs/models/qwen2_5_1_5b_radio_sd3_dynamic_puffin.py

@@ -0,0 +1,76 @@
+import torch
+from src.models.puffin.model import Qwen2p5RadioStableDiffusion3HFDynamic
+from src.models.stable_diffusion3.transformer_sd3_dynamic import SD3Transformer2DModel
+from src.models.radiov3.hf_model import RADIOModel
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+llm_name_or_path = 'Qwen/Qwen2.5-1.5B-Instruct'
+sd3_model_name_or_path = "stabilityai/stable-diffusion-3-medium-diffusers"
+
+prompt_template = dict(
+    SYSTEM=('<|im_start|>system\n{system}<|im_end|>\n'),
+    INSTRUCTION=('<|im_start|>user\n{input}<|im_end|>\n'
+                 '<|im_start|>assistant\n'),
+    SUFFIX='<|im_end|>',
+    IMG_START_TOKEN='<|vision_start|>',
+    IMG_END_TOKEN='<|vision_end|>',
+    IMG_CONTEXT_TOKEN='<|image_pad|>',
+    GENERATION='Generate an image: {input}',
+    GENERATION_CROSS='Generate a target image given an initial view: {input}',
+    SUFFIX_AS_EOS=True,
+    SEP='\n',
+    STOP_WORDS=['<|im_end|>', '<|endoftext|>']
+)
+
+model = dict(type=Qwen2p5RadioStableDiffusion3HFDynamic,
+             num_queries=64,
+             connector_1=dict(
+                 hidden_size=1024,
+                 intermediate_size=4096,
+                 num_hidden_layers=6,
+                 _attn_implementation='flash_attention_2',
+                 num_attention_heads=16, ),
+             connector_2=dict(
+                 hidden_size=1024,
+                 intermediate_size=4096,
+                 num_hidden_layers=6,
+                 _attn_implementation='flash_attention_2',
+                 num_attention_heads=16, ),
+             transformer=dict(
+                 type=SD3Transformer2DModel.from_pretrained,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="transformer",
+                 torch_dtype=torch.bfloat16),
+             test_scheduler=dict(
+                 type=FlowMatchEulerDiscreteScheduler.from_pretrained,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="scheduler"),
+             train_scheduler=dict(
+                 type=FlowMatchEulerDiscreteScheduler.from_pretrained,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="scheduler"),
+             vae=dict(
+                 type=AutoencoderKL.from_pretrained,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="vae",
+                 torch_dtype=torch.bfloat16),
+             freeze_visual_encoder=True,
+             freeze_llm=True,
+             llm=dict(
+                 type=AutoModelForCausalLM.from_pretrained,
+                 pretrained_model_name_or_path=llm_name_or_path,
+                 torch_dtype=torch.bfloat16,
+                 attn_implementation='flash_attention_2',
+             ),
+             tokenizer=dict(
+                 type=AutoTokenizer.from_pretrained,
+                 pretrained_model_name_or_path=llm_name_or_path),
+             prompt_template=prompt_template,
+             pretrained_pth=None,
+             use_activation_checkpointing=False,
+             visual_encoder=dict(
+                 type=RADIOModel.from_pretrained,
+                 pretrained_model_name_or_path="nvidia/C-RADIOv3-H",
+                 torch_dtype=torch.bfloat16,),
+             )

configs/pipelines/stage_2_base.py

@@ -0,0 +1,10 @@
+from mmengine.config import read_base
+
+with read_base():
+    from ..models.qwen2_5_1_5b_radio_sd3_dynamic_puffin import model
+
+model.freeze_visual_encoder = False
+model.freeze_llm = False
+model.freeze_transformer = False
+model.use_activation_checkpointing = True
+model.visual_encoder_grad_scale = 0.1

configs/pipelines/stage_3_thinking.py

@@ -0,0 +1,11 @@
+from mmengine.config import read_base
+
+with read_base():
+    from ..models.qwen2_5_1_5b_radio_sd3_dynamic_puffin import model
+
+model.freeze_visual_encoder = False
+model.freeze_llm = False
+model.freeze_transformer = False
+model.use_activation_checkpointing = True
+model.visual_encoder_grad_scale = 0.1
+#model.pretrained_pth = 'work_dirs/stage_2_base/iter_30000.pth'

configs/pipelines/stage_4_instruction_tuning.py

@@ -0,0 +1,9 @@
+from mmengine.config import read_base
+with read_base():
+    from ..models.qwen2_5_1_5b_radio_sd3_dynamic_puffin import model
+
+model.freeze_visual_encoder = True
+model.freeze_llm = False
+model.freeze_transformer = False
+model.use_activation_checkpointing = True
+model.unconditional_cross_view=0.1

requirements.txt

@@ -0,0 +1,48 @@
+accelerate==1.6.0
+click==8.1.8
+decorator==5.2.1
+deepspeed==0.16.7
+diffusers==0.34.0
+einops==0.8.1
+feedparser==6.0.11
+flash_attn==2.3.4
+huggingface-hub==0.31.1
+hyperframe==6.1.0
+idna==3.10
+imageio==2.37.0
+importlib_metadata==8.7.0
+json5==0.12.0
+lazy_loader==0.4
+lightning-utilities==0.14.3
+matplotlib==3.10.1
+matplotlib-inline==0.1.7
+mmengine==0.10.7
+networkx==3.4.2
+ninja==1.11.1.4
+numpy==2.2.5
+opencv-python==4.11.0.86
+opencv-python-headless==4.12.0.88
+openpyxl==3.1.5
+pandas==2.2.3
+peft==0.15.2
+pillow==11.2.1
+pytz==2025.2
+PyYAML==6.0.2
+safetensors==0.5.3
+scikit-image==0.25.2
+scipy==1.15.2
+six==1.17.0
+timm==0.9.12
+tokenizers==0.21.2
+torch==2.7.0
+torch-fidelity==0.3.0
+torchmetrics==1.7.2
+torchvision==0.22.0
+tornado==6.4.2
+tqdm==4.67.1
+transformers==4.49.0
+transformers-stream-generator==0.0.5
+triton==3.3.0
+xtuner==0.1.23
+yarl==1.20.0
+

scripts/camera/cam_dataset.py

@@ -0,0 +1,107 @@
+import os
+import json
+import argparse
+import torch
+from tqdm import tqdm
+
+from scripts.camera.geometry.camera import SimpleRadial
+from scripts.camera.geometry.gravity import Gravity
+from scripts.camera.geometry.perspective_fields import get_perspective_field
+from scripts.camera.utils.conversions import fov2focal
+from scripts.camera.utils.text import parse_camera_params
+
+class Cam_Generator:
+    def __init__(self, mode="base"):
+        self.mode = mode
+
+    def _load_text(self, caption, h=512, w=512, k1=0, k2=0):
+        # Parse camera params from caption
+        roll, pitch, vfov = parse_camera_params(caption, self.mode)
+        
+        # Convert vertical FoV to focal length
+        f = fov2focal(torch.tensor(vfov), h)
+        px, py = w / 2, h / 2
+        params = torch.tensor([w, h, f, f, px, py, k1, k2]).float()
+        gravity = torch.tensor([roll, pitch]).float()
+        return params, gravity
+
+    def _read_param(self, parameters, gravity):
+        # Build camera and gravity objects
+        camera = SimpleRadial(parameters).float()
+        roll, pitch = gravity.unbind(-1)
+        gravity_obj = Gravity.from_rp(roll, pitch)
+        camera = camera.scale(torch.Tensor([1, 1]))
+        return {"camera": camera, "gravity": gravity_obj}
+
+    def _get_perspective(self, data):
+        # Generate up and latitude fields
+        camera = data["camera"]
+        gravity_obj = data["gravity"]
+        up_field, lat_field = get_perspective_field(
+            camera, gravity_obj, use_up=True, use_latitude=True
+        )
+        del camera, gravity_obj
+        return torch.cat([up_field[0], lat_field[0]], dim=0)
+
+    def get_cam(self, caption):
+        params, gravity = self._load_text(caption)
+        data = self._read_param(params, gravity)
+        return self._get_perspective(data)
+
+def process_folders(input_root, output_root, start_idx=0, num_folders=None, mode="base"):
+    gen = Cam_Generator(mode=mode)
+    all_dirs = sorted([
+        d for d in os.listdir(input_root)
+        if os.path.isdir(os.path.join(input_root, d))
+    ])
+    if num_folders is None:
+        num_folders = len(all_dirs) - start_idx
+    selected = all_dirs[start_idx:start_idx + num_folders]
+
+    for sub in tqdm(selected, desc="Subfolders"):
+        in_sub = os.path.join(input_root, sub)
+        out_sub = os.path.join(output_root, sub)
+        os.makedirs(out_sub, exist_ok=True)
+
+        json_files = sorted([
+            f for f in os.listdir(in_sub)
+            if f.lower().endswith('.json')
+        ])
+
+        for jf in tqdm(json_files, desc=f"Processing {sub}", leave=False):
+            in_path = os.path.join(in_sub, jf)
+            with open(in_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+            caption = data.get('caption', '')
+            cam = gen.get_cam(caption)
+            out_name = os.path.splitext(jf)[0] + '.pt'
+            out_path = os.path.join(out_sub, out_name)
+            torch.save(cam, out_path)
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Batch process the captions to the camera maps and save as .pt"
+    )
+    parser.add_argument('--input_root', type=str,
+                        help='Root directory of JSON subfolders')
+    parser.add_argument('--output_root', type=str,
+                        help='Root directory to save .pt files')
+    parser.add_argument('--start_idx', type=int, default=0,
+                        help='Start index of subfolders (0-based, default=0)')
+    parser.add_argument('--num_folders', type=int, default=None,
+                        help='Number of subfolders to process (default: all)')
+    parser.add_argument('--mode', type=str, default='base',
+                        help='parse_camera_params mode')
+    args = parser.parse_args()
+
+    process_folders(
+        args.input_root,
+        args.output_root,
+        start_idx=args.start_idx,
+        num_folders=args.num_folders,
+        mode=args.mode
+    )
+
+
+if __name__ == '__main__':
+    main()

scripts/camera/geometry/base_camera.py

@@ -0,0 +1,518 @@
+"""Convenience classes a for camera models.
+
+Based on PyTorch tensors: differentiable, batched, with GPU support.
+Adapted from https://github.com/cvg/GeoCalib
+"""
+
+from abc import abstractmethod
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+from torch.func import jacfwd, vmap
+from torch.nn import functional as F
+
+from scripts.camera.geometry.gravity import Gravity
+from scripts.camera.utils.conversions import deg2rad, focal2fov, fov2focal, rad2rotmat
+from scripts.camera.utils.tensor import TensorWrapper, autocast
+
+# mypy: ignore-errors
+
+
+class BaseCamera(TensorWrapper):
+    """Camera tensor class."""
+
+    eps = 1e-3
+
+    @autocast
+    def __init__(self, data: torch.Tensor):
+        """Camera parameters with shape (..., {w, h, fx, fy, cx, cy, *dist}).
+
+        Tensor convention: (..., {w, h, fx, fy, cx, cy, pitch, roll, *dist}) where
+        - w, h: image size in pixels
+        - fx, fy: focal lengths in pixels
+        - cx, cy: principal points in normalized image coordinates
+        - dist: distortion parameters
+
+        Args:
+            data (torch.Tensor): Camera parameters with shape (..., {6, 7, 8}).
+        """
+        # w, h, fx, fy, cx, cy, dist
+        assert data.shape[-1] in {6, 7, 8}, data.shape
+
+        pad = data.new_zeros(data.shape[:-1] + (8 - data.shape[-1],))
+        data = torch.cat([data, pad], -1) if data.shape[-1] != 8 else data
+        super().__init__(data)
+
+    @classmethod
+    def from_dict(cls, param_dict: Dict[str, torch.Tensor]) -> "BaseCamera":
+        """Create a Camera object from a dictionary of parameters.
+
+        Args:
+            param_dict (Dict[str, torch.Tensor]): Dictionary of parameters.
+
+        Returns:
+            Camera: Camera object.
+        """
+        for key, value in param_dict.items():
+            if not isinstance(value, torch.Tensor):
+                param_dict[key] = torch.tensor(value)
+
+        h, w = param_dict["height"], param_dict["width"]
+        cx, cy = param_dict.get("cx", w / 2), param_dict.get("cy", h / 2)
+
+        vfov = param_dict.get("vfov")
+        f = param_dict.get("f", fov2focal(vfov, h))
+
+        if "dist" in param_dict:
+            k1, k2 = param_dict["dist"][..., 0], param_dict["dist"][..., 1]
+        elif "k1_hat" in param_dict:
+            k1 = param_dict["k1_hat"] * (f / h) ** 2
+
+            k2 = param_dict.get("k2", torch.zeros_like(k1))
+        else:
+            k1 = param_dict.get("k1", torch.zeros_like(f))
+            k2 = param_dict.get("k2", torch.zeros_like(f))
+
+        fx, fy = f, f
+        if "scales" in param_dict:
+            scales = param_dict["scales"]
+            fx = fx * scales[..., 0] / scales[..., 1]
+
+        params = torch.stack([w, h, fx, fy, cx, cy, k1, k2], dim=-1)
+        return cls(params)
+
+    def pinhole(self):
+        """Return the pinhole camera model."""
+        return self.__class__(self._data[..., :6])
+
+    @property
+    def size(self) -> torch.Tensor:
+        """Size (width height) of the images, with shape (..., 2)."""
+        return self._data[..., :2]
+
+    @property
+    def f(self) -> torch.Tensor:
+        """Focal lengths (fx, fy) with shape (..., 2)."""
+        return self._data[..., 2:4]
+
+    @property
+    def vfov(self) -> torch.Tensor:
+        """Vertical field of view in radians."""
+        return focal2fov(self.f[..., 1], self.size[..., 1])
+
+    @property
+    def hfov(self) -> torch.Tensor:
+        """Horizontal field of view in radians."""
+        return focal2fov(self.f[..., 0], self.size[..., 0])
+
+    @property
+    def c(self) -> torch.Tensor:
+        """Principal points (cx, cy) with shape (..., 2)."""
+        return self._data[..., 4:6]
+
+    @property
+    def K(self) -> torch.Tensor:
+        """Returns the self intrinsic matrix with shape (..., 3, 3)."""
+        shape = self.shape + (3, 3)
+        K = self._data.new_zeros(shape)
+        K[..., 0, 0] = self.f[..., 0]
+        K[..., 1, 1] = self.f[..., 1]
+        K[..., 0, 2] = self.c[..., 0]
+        K[..., 1, 2] = self.c[..., 1]
+        K[..., 2, 2] = 1
+        return K
+
+    def update_focal(self, delta: torch.Tensor, as_log: bool = False):
+        """Update the self parameters after changing the focal length."""
+        f = torch.exp(torch.log(self.f) + delta) if as_log else self.f + delta
+
+        # clamp focal length to a reasonable range for stability during training
+        min_f = fov2focal(self.new_ones(self.shape[0]) * deg2rad(150), self.size[..., 1])
+        max_f = fov2focal(self.new_ones(self.shape[0]) * deg2rad(5), self.size[..., 1])
+        min_f = min_f.unsqueeze(-1).expand(-1, 2)
+        max_f = max_f.unsqueeze(-1).expand(-1, 2)
+        f = f.clamp(min=min_f, max=max_f)
+
+        # make sure focal ration stays the same (avoid inplace operations)
+        fx = f[..., 1] * self.f[..., 0] / self.f[..., 1]
+        f = torch.stack([fx, f[..., 1]], -1)
+
+        dist = self.dist if hasattr(self, "dist") else self.new_zeros(self.f.shape)
+        return self.__class__(torch.cat([self.size, f, self.c, dist], -1))
+
+    def scale(self, scales: Union[float, int, Tuple[Union[float, int]]]):
+        """Update the self parameters after resizing an image."""
+        scales = (scales, scales) if isinstance(scales, (int, float)) else scales
+        s = scales if isinstance(scales, torch.Tensor) else self.new_tensor(scales)
+
+        dist = self.dist if hasattr(self, "dist") else self.new_zeros(self.f.shape)
+        return self.__class__(torch.cat([self.size * s, self.f * s, self.c * s, dist], -1))
+
+    def crop(self, pad: Tuple[float]):
+        """Update the self parameters after cropping an image."""
+        pad = pad if isinstance(pad, torch.Tensor) else self.new_tensor(pad)
+        size = self.size + pad.to(self.size)
+        c = self.c + pad.to(self.c) / 2
+
+        dist = self.dist if hasattr(self, "dist") else self.new_zeros(self.f.shape)
+        return self.__class__(torch.cat([size, self.f, c, dist], -1))
+
+    def undo_scale_crop(self, data: Dict[str, torch.Tensor]):
+        """Undo transforms done during scaling and cropping."""
+        camera = self.crop(-data["crop_pad"]) if "crop_pad" in data else self
+        return camera.scale(1.0 / data["scales"])
+
+    @autocast
+    def in_image(self, p2d: torch.Tensor):
+        """Check if 2D points are within the image boundaries."""
+        assert p2d.shape[-1] == 2
+        size = self.size.unsqueeze(-2)
+        return torch.all((p2d >= 0) & (p2d <= (size - 1)), -1)
+
+    @autocast
+    def project(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Project 3D points into the self plane and check for visibility."""
+        z = p3d[..., -1]
+        valid = z > self.eps
+        z = z.clamp(min=self.eps)
+        p2d = p3d[..., :-1] / z.unsqueeze(-1)
+        return p2d, valid
+
+    def J_project(self, p3d: torch.Tensor):
+        """Jacobian of the projection function."""
+        x, y, z = p3d[..., 0], p3d[..., 1], p3d[..., 2]
+        zero = torch.zeros_like(z)
+        z = z.clamp(min=self.eps)
+        J = torch.stack([1 / z, zero, -x / z**2, zero, 1 / z, -y / z**2], dim=-1)
+        J = J.reshape(p3d.shape[:-1] + (2, 3))
+        return J  # N x 2 x 3
+
+    @abstractmethod
+    def distort(self, pts: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates and check for validity of the distortion model."""
+        raise NotImplementedError("distort() must be implemented.")
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the distortion function."""
+        if wrt == "scale2pts":  # (..., 2)
+            J = [
+                vmap(jacfwd(lambda x: self[idx].distort(x, return_scale=True)[0]))(p2d[idx])[None]
+                for idx in range(p2d.shape[0])
+            ]
+
+            return torch.cat(J, dim=0).squeeze(-3, -2)
+
+        elif wrt == "scale2dist":  # (..., 1)
+            J = []
+            for idx in range(p2d.shape[0]):  # loop to batch pts dimension
+
+                def func(x):
+                    params = torch.cat([self._data[idx, :6], x[None]], -1)
+                    return self.__class__(params).distort(p2d[idx], return_scale=True)[0]
+
+                J.append(vmap(jacfwd(func))(self[idx].dist))
+
+            return torch.cat(J, dim=0)
+
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    @abstractmethod
+    def undistort(self, pts: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates and check for validity of the distortion model."""
+        raise NotImplementedError("undistort() must be implemented.")
+
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        if wrt == "pts":  # (..., 2, 2)
+            J = [
+                vmap(jacfwd(lambda x: self[idx].undistort(x)[0]))(p2d[idx])[None]
+                for idx in range(p2d.shape[0])
+            ]
+
+            return torch.cat(J, dim=0).squeeze(-3)
+
+        elif wrt == "dist":  # (..., 1)
+            J = []
+            for batch_idx in range(p2d.shape[0]):  # loop to batch pts dimension
+
+                def func(x):
+                    params = torch.cat([self._data[batch_idx, :6], x[None]], -1)
+                    return self.__class__(params).undistort(p2d[batch_idx])[0]
+
+                J.append(vmap(jacfwd(func))(self[batch_idx].dist))
+
+            return torch.cat(J, dim=0)
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    @autocast
+    def up_projection_offset(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Compute the offset for the up-projection."""
+        return self.J_distort(p2d, wrt="scale2pts")  # (B, N, 2)
+
+    def J_up_projection_offset(self, p2d: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+        """Jacobian of the distortion offset for up-projection."""
+        if wrt == "uv":  # (B, N, 2, 2)
+            J = [
+                vmap(jacfwd(lambda x: self[idx].up_projection_offset(x)[0, 0]))(p2d[idx])[None]
+                for idx in range(p2d.shape[0])
+            ]
+
+            return torch.cat(J, dim=0)
+
+        elif wrt == "dist":  # (B, N, 2)
+            J = []
+            for batch_idx in range(p2d.shape[0]):  # loop to batch pts dimension
+
+                def func(x):
+                    params = torch.cat([self._data[batch_idx, :6], x[None]], -1)[None]
+                    return self.__class__(params).up_projection_offset(p2d[batch_idx][None])
+
+                J.append(vmap(jacfwd(func))(self[batch_idx].dist))
+
+            return torch.cat(J, dim=0).squeeze(1)
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    @autocast
+    def denormalize(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Convert normalized 2D coordinates into pixel coordinates."""
+        return p2d * self.f.unsqueeze(-2) + self.c.unsqueeze(-2)
+
+    def J_denormalize(self):
+        """Jacobian of the denormalization function."""
+        return torch.diag_embed(self.f)  # ..., 2 x 2
+
+    @autocast
+    def normalize(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Convert pixel coordinates into normalized 2D coordinates."""
+        return (p2d - self.c.unsqueeze(-2)) / (self.f.unsqueeze(-2))
+
+    def J_normalize(self, p2d: torch.Tensor, wrt: str = "f"):
+        """Jacobian of the normalization function."""
+        # ... x N x 2 x 2
+        if wrt == "f":
+            J_f = -(p2d - self.c.unsqueeze(-2)) / ((self.f.unsqueeze(-2)) ** 2)
+            return torch.diag_embed(J_f)
+        elif wrt == "pts":
+            J_pts = 1 / self.f
+            return torch.diag_embed(J_pts)
+        else:
+            raise NotImplementedError(f"Jacobian not implemented for wrt={wrt}")
+
+    def pixel_coordinates(self) -> torch.Tensor:
+        """Pixel coordinates in self frame.
+
+        Returns:
+            torch.Tensor: Pixel coordinates as a tensor of shape (B, h * w, 2).
+        """
+        w, h = self.size[0].unbind(-1)
+        h, w = h.round().to(int), w.round().to(int)
+
+        # create grid
+        x = torch.arange(0, w, dtype=self.dtype, device=self.device)
+        y = torch.arange(0, h, dtype=self.dtype, device=self.device)
+        x, y = torch.meshgrid(x, y, indexing="xy")
+        xy = torch.stack((x, y), dim=-1).reshape(-1, 2)  # shape (h * w, 2)
+
+        # add batch dimension (normalize() would broadcast but we make it explicit)
+        B = self.shape[0]
+        xy = xy.unsqueeze(0).expand(B, -1, -1)  # if B > 0 else xy
+
+        return xy.to(self.device).to(self.dtype)
+
+    def normalized_image_coordinates(self) -> torch.Tensor:
+        """Normalized image coordinates in self frame.
+
+        Returns:
+            torch.Tensor: Normalized image coordinates as a tensor of shape (B, h * w, 3).
+        """
+        xy = self.pixel_coordinates()
+        uv1, _ = self.image2world(xy)
+
+        B = self.shape[0]
+        uv1 = uv1.reshape(B, -1, 3)
+        return uv1.to(self.device).to(self.dtype)
+
+    @autocast
+    def pixel_bearing_many(self, p3d: torch.Tensor) -> torch.Tensor:
+        """Get the bearing vectors of pixel coordinates.
+
+        Args:
+            p2d (torch.Tensor): Pixel coordinates as a tensor of shape (..., 3).
+
+        Returns:
+            torch.Tensor: Bearing vectors as a tensor of shape (..., 3).
+        """
+        return F.normalize(p3d, dim=-1)
+
+    @autocast
+    def world2image(self, p3d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Transform 3D points into 2D pixel coordinates."""
+        p2d, visible = self.project(p3d)
+        p2d, mask = self.distort(p2d)
+        p2d = self.denormalize(p2d)
+        valid = visible & mask & self.in_image(p2d)
+        return p2d, valid
+
+    @autocast
+    def J_world2image(self, p3d: torch.Tensor):
+        """Jacobian of the world2image function."""
+        p2d_proj, valid = self.project(p3d)
+
+        J_dnorm = self.J_denormalize()
+        J_dist = self.J_distort(p2d_proj)
+        J_proj = self.J_project(p3d)
+
+        J = torch.einsum("...ij,...jk,...kl->...il", J_dnorm, J_dist, J_proj)
+        return J, valid
+
+    @autocast
+    def image2world(self, p2d: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Transform point in the image plane to 3D world coordinates."""
+        p2d = self.normalize(p2d)
+        p2d, valid = self.undistort(p2d)
+        ones = p2d.new_ones(p2d.shape[:-1] + (1,))
+        p3d = torch.cat([p2d, ones], -1)
+        return p3d, valid
+
+    @autocast
+    def J_image2world(self, p2d: torch.Tensor, wrt: str = "f") -> Tuple[torch.Tensor, torch.Tensor]:
+        """Jacobian of the image2world function."""
+        if wrt == "dist":
+            p2d_norm = self.normalize(p2d)
+            return self.J_undistort(p2d_norm, wrt)
+        elif wrt == "f":
+            J_norm2f = self.J_normalize(p2d, wrt)
+            p2d_norm = self.normalize(p2d)
+            J_dist2norm = self.J_undistort(p2d_norm, "pts")
+
+            return torch.einsum("...ij,...jk->...ik", J_dist2norm, J_norm2f)
+        else:
+            raise ValueError(f"Unknown wrt: {wrt}")
+
+    @autocast
+    def undistort_image(self, img: torch.Tensor) -> torch.Tensor:
+        """Undistort an image using the distortion model."""
+        assert self.shape[0] == 1, "Batch size must be 1."
+        W, H = self.size.unbind(-1)
+        H, W = H.int().item(), W.int().item()
+
+        x, y = torch.arange(0, W), torch.arange(0, H)
+        x, y = torch.meshgrid(x, y, indexing="xy")
+        coords = torch.stack((x, y), dim=-1).reshape(-1, 2)
+
+        p3d, _ = self.pinhole().image2world(coords.to(self.device).to(self.dtype))
+        p2d, _ = self.world2image(p3d)
+
+        mapx, mapy = p2d[..., 0].reshape((1, H, W)), p2d[..., 1].reshape((1, H, W))
+        grid = torch.stack((mapx, mapy), dim=-1)
+        grid = 2.0 * grid / torch.tensor([W - 1, H - 1]).to(grid) - 1
+        return F.grid_sample(img, grid, align_corners=True)
+
+    def get_img_from_pano(
+        self,
+        pano_img: torch.Tensor,
+        gravity: Gravity,
+        yaws: torch.Tensor = 0.0,
+        resize_factor: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Render an image from a panorama.
+
+        Args:
+            pano_img (torch.Tensor): Panorama image of shape (3, H, W) in [0, 1].
+            gravity (Gravity): Gravity direction of the camera.
+            yaws (torch.Tensor | list, optional): Yaw angle in radians. Defaults to 0.0.
+            resize_factor (torch.Tensor, optional): Resize the panorama to be a multiple of the
+            field of view. Defaults to 1.
+
+        Returns:
+            torch.Tensor: Image rendered from the panorama.
+        """
+        B = self.shape[0]
+        if B > 0:
+            assert self.size[..., 0].unique().shape[0] == 1, "All images must have the same width."
+            assert self.size[..., 1].unique().shape[0] == 1, "All images must have the same height."
+
+        w, h = self.size[0].unbind(-1)
+        h, w = h.round().to(int), w.round().to(int)
+
+        if isinstance(yaws, (int, float)):
+            yaws = [yaws]
+        if isinstance(resize_factor, (int, float)):
+            resize_factor = [resize_factor]
+
+        yaws = (
+            yaws.to(self.dtype).to(self.device)
+            if isinstance(yaws, torch.Tensor)
+            else self.new_tensor(yaws)
+        )
+
+        if isinstance(resize_factor, torch.Tensor):
+            resize_factor = resize_factor.to(self.dtype).to(self.device)
+        elif resize_factor is not None:
+            resize_factor = self.new_tensor(resize_factor)
+
+        assert isinstance(pano_img, torch.Tensor), "Panorama image must be a torch.Tensor."
+        pano_img = pano_img if pano_img.dim() == 4 else pano_img.unsqueeze(0)  # B x 3 x H x W
+
+        pano_imgs = []
+        for i, yaw in enumerate(yaws):
+            if resize_factor is not None:
+                # resize the panorama such that the fov of the panorama has the same height as the
+                # image
+                vfov = self.vfov[i] if B != 0 else self.vfov
+                scale = torch.pi / float(vfov) * float(h) / pano_img.shape[-2] * resize_factor[i]
+                pano_shape = (int(pano_img.shape[-2] * scale), int(pano_img.shape[-1] * scale))
+
+                mode = "bicubic" if scale >= 1 else "area"
+                resized_pano = F.interpolate(pano_img, size=pano_shape, mode=mode)
+            else:
+                # make sure to copy: resized_pano = pano_img
+                resized_pano = pano_img
+                pano_shape = pano_img.shape[-2:][::-1]
+
+            pano_imgs.append((resized_pano, pano_shape))
+
+        xy = self.pixel_coordinates()
+        uv1, valid = self.image2world(xy)
+        bearings = self.pixel_bearing_many(uv1)
+
+        # rotate bearings
+        R_yaw = rad2rotmat(self.new_zeros(yaw.shape), self.new_zeros(yaw.shape), yaws)
+        rotated_bearings = bearings @ gravity.R @ R_yaw
+
+        # spherical coordinates
+        lon = torch.atan2(rotated_bearings[..., 0], rotated_bearings[..., 2])
+        lat = torch.atan2(
+            rotated_bearings[..., 1], torch.norm(rotated_bearings[..., [0, 2]], dim=-1)
+        )
+
+        images = []
+        for idx, (resized_pano, pano_shape) in enumerate(pano_imgs):
+            min_lon, max_lon = -torch.pi, torch.pi
+            min_lat, max_lat = -torch.pi / 2.0, torch.pi / 2.0
+            min_x, max_x = 0, pano_shape[0] - 1.0
+            min_y, max_y = 0, pano_shape[1] - 1.0
+
+            # map Spherical Coordinates to Panoramic Coordinates
+            nx = (lon[idx] - min_lon) / (max_lon - min_lon) * (max_x - min_x) + min_x
+            ny = (lat[idx] - min_lat) / (max_lat - min_lat) * (max_y - min_y) + min_y
+
+            # reshape and cast to numpy for remap
+            mapx = nx.reshape((1, h, w))
+            mapy = ny.reshape((1, h, w))
+
+            grid = torch.stack((mapx, mapy), dim=-1)  # Add batch dimension
+            # Normalize to [-1, 1]
+            grid = 2.0 * grid / torch.tensor([pano_shape[-2] - 1, pano_shape[-1] - 1]).to(grid) - 1
+            # Apply grid sample
+            image = F.grid_sample(resized_pano, grid, align_corners=True)#True
+            images.append(image)
+
+        return torch.concatenate(images, 0) if B > 0 else images[0]
+
+    def __repr__(self):
+        """Print the Camera object."""
+        return f"{self.__class__.__name__} {self.shape} {self.dtype} {self.device}"

scripts/camera/geometry/camera.py

@@ -0,0 +1,281 @@
+"""Implementation of the pinhole, simple radial, and simple divisional camera models."""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+from typing import Tuple
+
+import torch
+
+from scripts.camera.geometry.base_camera import BaseCamera
+from scripts.camera.utils.tensor import autocast
+
+# flake8: noqa: E741
+
+# mypy: ignore-errors
+
+
+class Pinhole(BaseCamera):
+    """Implementation of the pinhole camera model."""
+
+    def distort(self, p2d: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates."""
+        if return_scale:
+            return p2d.new_ones(p2d.shape[:-1] + (1,))
+
+        return p2d, p2d.new_ones((p2d.shape[0], 1)).bool()
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the distortion function."""
+        if wrt == "pts":
+            return torch.eye(2, device=p2d.device, dtype=p2d.dtype).expand(p2d.shape[:-1] + (2, 2))
+        else:
+            raise ValueError(f"Unknown wrt: {wrt}")
+
+    def undistort(self, pts: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates."""
+        return pts, pts.new_ones((pts.shape[0], 1)).bool()
+
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        if wrt == "pts":
+            return torch.eye(2, device=p2d.device, dtype=p2d.dtype).expand(p2d.shape[:-1] + (2, 2))
+        else:
+            raise ValueError(f"Unknown wrt: {wrt}")
+
+
+class SimpleRadial(BaseCamera):
+    """Implementation of the simple radial camera model."""
+
+    @property
+    def dist(self) -> torch.Tensor:
+        """Distortion parameters, with shape (..., 1)."""
+        return self._data[..., 6:]
+
+    @property
+    def k1(self) -> torch.Tensor:
+        """Distortion parameters, with shape (...)."""
+        return self._data[..., 6]
+
+    @property
+    def k1_hat(self) -> torch.Tensor:
+        """Distortion parameters, with shape (...)."""
+        return self.k1 / (self.f[..., 1] / self.size[..., 1]) ** 2
+
+    def update_dist(self, delta: torch.Tensor, dist_range: Tuple[float, float] = (-0.7, 0.7)):
+        """Update the self parameters after changing the k1 distortion parameter."""
+        delta_dist = self.new_ones(self.dist.shape) * delta
+        dist = (self.dist + delta_dist).clamp(*dist_range)
+        data = torch.cat([self.size, self.f, self.c, dist], -1)
+        return self.__class__(data)
+
+    @autocast
+    def check_valid(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Check if the distorted points are valid."""
+        return p2d.new_ones(p2d.shape[:-1]).bool()
+
+    def distort(self, p2d: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates and check for validity of the distortion model."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        radial = 1 + self.k1[..., None, None] * r2
+
+        if return_scale:
+            return radial, None
+
+        return p2d * radial, self.check_valid(p2d)
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts"):
+        """Jacobian of the distortion function."""
+        k1 = self.k1[..., None, None]
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        if wrt == "pts":  # (..., 2, 2)
+            radial = 1 + k1 * r2
+            ppT = torch.einsum("...i,...j->...ij", p2d, p2d)  # (..., 2, 2)
+            return (2 * k1 * ppT) + torch.diag_embed(radial.expand(radial.shape[:-1] + (2,)))
+        elif wrt == "dist":  # (..., 2)
+            return r2 * p2d
+        elif wrt == "scale2dist":  # (..., 1)
+            return r2
+        elif wrt == "scale2pts":  # (..., 2)
+            return 2 * k1 * p2d
+        else:
+            return super().J_distort(p2d, wrt)
+
+    @autocast
+    def undistort(self, p2d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates and check for validity of the distortion model."""
+        b1 = -self.k1[..., None, None]
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        radial = 1 + b1 * r2
+        return p2d * radial, self.check_valid(p2d)
+
+    @autocast
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        b1 = -self.k1[..., None, None]
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        if wrt == "dist":
+            return -r2 * p2d
+        elif wrt == "pts":
+            radial = 1 + b1 * r2
+            ppT = torch.einsum("...i,...j->...ij", p2d, p2d)  # (..., 2, 2)
+            return (2 * b1[..., None] * ppT) + torch.diag_embed(
+                radial.expand(radial.shape[:-1] + (2,))
+            )
+        else:
+            return super().J_undistort(p2d, wrt)
+
+    def J_up_projection_offset(self, p2d: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+        """Jacobian of the up-projection offset."""
+        if wrt == "uv":  # (..., 2, 2)
+            return torch.diag_embed((2 * self.k1[..., None, None]).expand(p2d.shape[:-1] + (2,)))
+        elif wrt == "dist":
+            return 2 * p2d  # (..., 2)
+        else:
+            return super().J_up_projection_offset(p2d, wrt)
+
+
+class SimpleDivisional(BaseCamera):
+    """Implementation of the simple divisional camera model."""
+
+    @property
+    def dist(self) -> torch.Tensor:
+        """Distortion parameters, with shape (..., 1)."""
+        return self._data[..., 6:]
+
+    @property
+    def k1(self) -> torch.Tensor:
+        """Distortion parameters, with shape (...)."""
+        return self._data[..., 6]
+
+    def update_dist(self, delta: torch.Tensor, dist_range: Tuple[float, float] = (-3.0, 3.0)):
+        """Update the self parameters after changing the k1 distortion parameter."""
+        delta_dist = self.new_ones(self.dist.shape) * delta
+        dist = (self.dist + delta_dist).clamp(*dist_range)
+        data = torch.cat([self.size, self.f, self.c, dist], -1)
+        return self.__class__(data)
+
+    @autocast
+    def check_valid(self, p2d: torch.Tensor) -> torch.Tensor:
+        """Check if the distorted points are valid."""
+        return p2d.new_ones(p2d.shape[:-1]).bool()
+
+    def distort(self, p2d: torch.Tensor, return_scale: bool = False) -> Tuple[torch.Tensor]:
+        """Distort normalized 2D coordinates and check for validity of the distortion model."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        radial = 1 - torch.sqrt((1 - 4 * self.k1[..., None, None] * r2).clamp(min=0))
+        denom = 2 * self.k1[..., None, None] * r2
+
+        ones = radial.new_ones(radial.shape)
+        radial = torch.where(denom == 0, ones, radial / denom.masked_fill(denom == 0, 1e6))
+
+        if return_scale:
+            return radial, None
+
+        return p2d * radial, self.check_valid(p2d)
+
+    def J_distort(self, p2d: torch.Tensor, wrt: str = "pts"):
+        """Jacobian of the distortion function."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        t0 = torch.sqrt((1 - 4 * self.k1[..., None, None] * r2).clamp(min=1e-6))
+        if wrt == "scale2pts":  # (B, N, 2)
+            d1 = t0 * 2 * r2
+            d2 = self.k1[..., None, None] * r2**2
+            denom = d1 * d2
+            return p2d * (4 * d2 - (1 - t0) * d1) / denom.masked_fill(denom == 0, 1e6)
+
+        elif wrt == "scale2dist":
+            d1 = 2 * self.k1[..., None, None] * t0
+            d2 = 2 * r2 * self.k1[..., None, None] ** 2
+            denom = d1 * d2
+            return (2 * d2 - (1 - t0) * d1) / denom.masked_fill(denom == 0, 1e6)
+
+        else:
+            return super().J_distort(p2d, wrt)
+
+    @autocast
+    def undistort(self, p2d: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Undistort normalized 2D coordinates and check for validity of the distortion model."""
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        denom = 1 + self.k1[..., None, None] * r2
+        radial = 1 / denom.masked_fill(denom == 0, 1e6)
+        return p2d * radial, self.check_valid(p2d)
+
+    def J_undistort(self, p2d: torch.Tensor, wrt: str = "pts") -> torch.Tensor:
+        """Jacobian of the undistortion function."""
+        # return super().J_undistort(p2d, wrt)
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        k1 = self.k1[..., None, None]
+        if wrt == "dist":
+            denom = (1 + k1 * r2) ** 2
+            return -r2 / denom.masked_fill(denom == 0, 1e6) * p2d
+        elif wrt == "pts":
+            t0 = 1 + k1 * r2
+            t0 = t0.masked_fill(t0 == 0, 1e6)
+            ppT = torch.einsum("...i,...j->...ij", p2d, p2d)  # (..., 2, 2)
+            J = torch.diag_embed((1 / t0).expand(p2d.shape[:-1] + (2,)))
+            return J - 2 * k1[..., None] * ppT / t0[..., None] ** 2  # (..., N, 2, 2)
+
+        else:
+            return super().J_undistort(p2d, wrt)
+
+    def J_up_projection_offset(self, p2d: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+        """Jacobian of the up-projection offset.
+
+        func(uv, dist) = 4 / (2 * norm2(uv)^2 * (1-4*k1*norm2(uv)^2)^0.5) * uv
+        - (1-(1-4*k1*norm2(uv)^2)^0.5) / (k1 * norm2(uv)^4) * uv
+        """
+        k1 = self.k1[..., None, None]
+        r2 = torch.sum(p2d**2, -1, keepdim=True)
+        t0 = (1 - 4 * k1 * r2).clamp(min=1e-6)
+        t1 = torch.sqrt(t0)
+        if wrt == "dist":
+            denom = 4 * t0 ** (3 / 2)
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = 16 / denom
+
+            denom = r2 * t1 * k1
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - 2 / denom
+
+            denom = (r2 * k1) ** 2
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J + (1 - t1) / denom
+
+            return J * p2d
+        elif wrt == "uv":
+            # ! unstable (gradient checker might fail), rewrite to use single division (by denom)
+            ppT = torch.einsum("...i,...j->...ij", p2d, p2d)  # (..., 2, 2)
+
+            denom = 2 * r2 * t1
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = torch.diag_embed((4 / denom).expand(p2d.shape[:-1] + (2,)))
+
+            denom = 4 * t1 * r2**2
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - 16 / denom[..., None] * ppT
+
+            denom = 4 * r2 * t0 ** (3 / 2)
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J + (32 * k1[..., None]) / denom[..., None] * ppT
+
+            denom = r2**2 * t1
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - 4 / denom[..., None] * ppT
+
+            denom = k1 * r2**3
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J + (4 * (1 - t1) / denom)[..., None] * ppT
+
+            denom = k1 * r2**2
+            denom = denom.masked_fill(denom == 0, 1e6)
+            J = J - torch.diag_embed(((1 - t1) / denom).expand(p2d.shape[:-1] + (2,)))
+
+            return J
+        else:
+            return super().J_up_projection_offset(p2d, wrt)
+
+
+camera_models = {
+    "pinhole": Pinhole,
+    "simple_radial": SimpleRadial,
+    "simple_divisional": SimpleDivisional,
+}

scripts/camera/geometry/gravity.py

@@ -0,0 +1,129 @@
+"""Tensor class for gravity vector in camera frame."""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+import torch
+from torch.nn import functional as F
+
+from scripts.camera.geometry.manifolds import EuclideanManifold, SphericalManifold
+from scripts.camera.utils.conversions import rad2rotmat
+from scripts.camera.utils.tensor import TensorWrapper, autocast
+
+# mypy: ignore-errors
+
+
+class Gravity(TensorWrapper):
+    """Gravity vector in camera frame."""
+
+    eps = 1e-4
+
+    @autocast
+    def __init__(self, data: torch.Tensor) -> None:
+        """Create gravity vector from data.
+
+        Args:
+            data (torch.Tensor): gravity vector as 3D vector in camera frame.
+        """
+        assert data.shape[-1] == 3, data.shape
+
+        data = F.normalize(data, dim=-1)
+
+        super().__init__(data)
+
+    @classmethod
+    def from_rp(cls, roll: torch.Tensor, pitch: torch.Tensor) -> "Gravity":
+        """Create gravity vector from roll and pitch angles."""
+        if not isinstance(roll, torch.Tensor):
+            roll = torch.tensor(roll)
+        if not isinstance(pitch, torch.Tensor):
+            pitch = torch.tensor(pitch)
+
+        sr, cr = torch.sin(roll), torch.cos(roll)
+        sp, cp = torch.sin(pitch), torch.cos(pitch)
+        return cls(torch.stack([-sr * cp, -cr * cp, sp], dim=-1))
+
+    @property
+    def vec3d(self) -> torch.Tensor:
+        """Return the gravity vector in the representation."""
+        return self._data
+
+    @property
+    def x(self) -> torch.Tensor:
+        """Return first component of the gravity vector."""
+        return self._data[..., 0]
+
+    @property
+    def y(self) -> torch.Tensor:
+        """Return second component of the gravity vector."""
+        return self._data[..., 1]
+
+    @property
+    def z(self) -> torch.Tensor:
+        """Return third component of the gravity vector."""
+        return self._data[..., 2]
+
+    @property
+    def roll(self) -> torch.Tensor:
+        """Return the roll angle of the gravity vector."""
+        roll = torch.asin(-self.x / (torch.sqrt(1 - self.z**2) + self.eps))
+        offset = -torch.pi * torch.sign(self.x)
+        return torch.where(self.y < 0, roll, -roll + offset)
+
+    def J_roll(self) -> torch.Tensor:
+        """Return the Jacobian of the roll angle of the gravity vector."""
+        cp, _ = torch.cos(self.pitch), torch.sin(self.pitch)
+        cr, sr = torch.cos(self.roll), torch.sin(self.roll)
+        Jr = self.new_zeros(self.shape + (3,))
+        Jr[..., 0] = -cr * cp
+        Jr[..., 1] = sr * cp
+        return Jr
+
+    @property
+    def pitch(self) -> torch.Tensor:
+        """Return the pitch angle of the gravity vector."""
+        return torch.asin(self.z)
+
+    def J_pitch(self) -> torch.Tensor:
+        """Return the Jacobian of the pitch angle of the gravity vector."""
+        cp, sp = torch.cos(self.pitch), torch.sin(self.pitch)
+        cr, sr = torch.cos(self.roll), torch.sin(self.roll)
+
+        Jp = self.new_zeros(self.shape + (3,))
+        Jp[..., 0] = sr * sp
+        Jp[..., 1] = cr * sp
+        Jp[..., 2] = cp
+        return Jp
+
+    @property
+    def rp(self) -> torch.Tensor:
+        """Return the roll and pitch angles of the gravity vector."""
+        return torch.stack([self.roll, self.pitch], dim=-1)
+
+    def J_rp(self) -> torch.Tensor:
+        """Return the Jacobian of the roll and pitch angles of the gravity vector."""
+        return torch.stack([self.J_roll(), self.J_pitch()], dim=-1)
+
+    @property
+    def R(self) -> torch.Tensor:
+        """Return the rotation matrix from the gravity vector."""
+        return rad2rotmat(roll=self.roll, pitch=self.pitch)
+
+    def J_R(self) -> torch.Tensor:
+        """Return the Jacobian of the rotation matrix from the gravity vector."""
+        raise NotImplementedError
+
+    def update(self, delta: torch.Tensor, spherical: bool = False) -> "Gravity":
+        """Update the gravity vector by adding a delta."""
+        if spherical:
+            data = SphericalManifold.plus(self.vec3d, delta)
+            return self.__class__(data)
+
+        data = EuclideanManifold.plus(self.rp, delta)
+        return self.from_rp(data[..., 0], data[..., 1])
+
+    def J_update(self, spherical: bool = False) -> torch.Tensor:
+        """Return the Jacobian of the update."""
+        return SphericalManifold if spherical else EuclideanManifold
+
+    def __repr__(self):
+        """Print the Camera object."""
+        return f"{self.__class__.__name__} {self.shape} {self.dtype} {self.device}"

scripts/camera/geometry/jacobians.py

@@ -0,0 +1,63 @@
+"""Jacobians for optimization."""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+import torch
+
+
+@torch.jit.script
+def J_vecnorm(vec: torch.Tensor) -> torch.Tensor:
+    """Compute the jacobian of vec / norm2(vec).
+
+    Args:
+        vec (torch.Tensor): [..., D] tensor.
+
+    Returns:
+        torch.Tensor: [..., D, D] Jacobian.
+    """
+    D = vec.shape[-1]
+    norm_x = torch.norm(vec, dim=-1, keepdim=True).unsqueeze(-1)  # (..., 1, 1)
+
+    if (norm_x == 0).any():
+        norm_x = norm_x + 1e-6
+
+    xxT = torch.einsum("...i,...j->...ij", vec, vec)  # (..., D, D)
+    identity = torch.eye(D, device=vec.device, dtype=vec.dtype)  # (D, D)
+
+    return identity / norm_x - (xxT / norm_x**3)  # (..., D, D)
+
+
+@torch.jit.script
+def J_focal2fov(focal: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+    """Compute the jacobian of the focal2fov function."""
+    return -4 * h / (4 * focal**2 + h**2)
+
+
+@torch.jit.script
+def J_up_projection(uv: torch.Tensor, abc: torch.Tensor, wrt: str = "uv") -> torch.Tensor:
+    """Compute the jacobian of the up-vector projection.
+
+    Args:
+        uv (torch.Tensor): Normalized image coordinates of shape (..., 2).
+        abc (torch.Tensor): Gravity vector of shape (..., 3).
+        wrt (str, optional): Parameter to differentiate with respect to. Defaults to "uv".
+
+    Raises:
+        ValueError: If the wrt parameter is unknown.
+
+    Returns:
+        torch.Tensor: Jacobian with respect to the parameter.
+    """
+    if wrt == "uv":
+        c = abc[..., 2][..., None, None, None]
+        return -c * torch.eye(2, device=uv.device, dtype=uv.dtype).expand(uv.shape[:-1] + (2, 2))
+
+    elif wrt == "abc":
+        J = uv.new_zeros(uv.shape[:-1] + (2, 3))
+        J[..., 0, 0] = 1
+        J[..., 1, 1] = 1
+        J[..., 0, 2] = -uv[..., 0]
+        J[..., 1, 2] = -uv[..., 1]
+        return J
+
+    else:
+        raise ValueError(f"Unknown wrt: {wrt}")

scripts/camera/geometry/manifolds.py

@@ -0,0 +1,113 @@
+"""Implementation of manifolds."""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class EuclideanManifold:
+    """Simple euclidean manifold."""
+
+    @staticmethod
+    def J_plus(x: torch.Tensor) -> torch.Tensor:
+        """Plus operator Jacobian."""
+        return torch.eye(x.shape[-1]).to(x)
+
+    @staticmethod
+    def plus(x: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
+        """Plus operator."""
+        return x + delta
+
+
+class SphericalManifold:
+    """Implementation of the spherical manifold.
+
+    Following the derivation from 'Integrating Generic Sensor Fusion Algorithms with Sound State
+    Representations through Encapsulation of Manifolds' by Hertzberg et al. (B.2, p. 25).
+
+    Householder transformation following Algorithm 5.1.1 (p. 210) from 'Matrix Computations' by
+    Golub et al.
+    """
+
+    @staticmethod
+    def householder_vector(x: torch.Tensor) -> torch.Tensor:
+        """Return the Householder vector and beta.
+
+        Algorithm 5.1.1 (p. 210) from 'Matrix Computations' by Golub et al. (Johns Hopkins Studies
+        in Mathematical Sciences) but using the nth element of the input vector as pivot instead of
+        first.
+
+        This computes the vector v with v(n) = 1 and beta such that H = I - beta * v * v^T is
+        orthogonal and H * x = ||x||_2 * e_n.
+
+        Args:
+            x (torch.Tensor): [..., n] tensor.
+
+        Returns:
+            torch.Tensor: v of shape [..., n]
+            torch.Tensor: beta of shape [...]
+        """
+        sigma = torch.sum(x[..., :-1] ** 2, -1)
+        xpiv = x[..., -1]
+        norm = torch.norm(x, dim=-1)
+        if torch.any(sigma < 1e-7):
+            sigma = torch.where(sigma < 1e-7, sigma + 1e-7, sigma)
+            logger.warning("sigma < 1e-7")
+
+        vpiv = torch.where(xpiv < 0, xpiv - norm, -sigma / (xpiv + norm))
+        beta = 2 * vpiv**2 / (sigma + vpiv**2)
+        v = torch.cat([x[..., :-1] / vpiv[..., None], torch.ones_like(vpiv)[..., None]], -1)
+        return v, beta
+
+    @staticmethod
+    def apply_householder(y: torch.Tensor, v: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
+        """Apply Householder transformation.
+
+        Args:
+            y (torch.Tensor): Vector to transform of shape [..., n].
+            v (torch.Tensor): Householder vector of shape [..., n].
+            beta (torch.Tensor): Householder beta of shape [...].
+
+        Returns:
+            torch.Tensor: Transformed vector of shape [..., n].
+        """
+        return y - v * (beta * torch.einsum("...i,...i->...", v, y))[..., None]
+
+    @classmethod
+    def J_plus(cls, x: torch.Tensor) -> torch.Tensor:
+        """Plus operator Jacobian."""
+        v, beta = cls.householder_vector(x)
+        H = -torch.einsum("..., ...k, ...l->...kl", beta, v, v)
+        H = H + torch.eye(H.shape[-1]).to(H)
+        return H[..., :-1]  # J
+
+    @classmethod
+    def plus(cls, x: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
+        """Plus operator.
+
+        Equation 109 (p. 25) from 'Integrating Generic Sensor Fusion Algorithms with Sound State
+        Representations through Encapsulation of Manifolds' by Hertzberg et al. but using the nth
+        element of the input vector as pivot instead of first.
+
+        Args:
+            x: point on the manifold
+            delta: tangent vector
+        """
+        eps = 1e-7
+        # keep norm is not equal to 1
+        nx = torch.norm(x, dim=-1, keepdim=True)
+        nd = torch.norm(delta, dim=-1, keepdim=True)
+
+        # make sure we don't divide by zero in backward as torch.where computes grad for both
+        # branches
+        nd_ = torch.where(nd < eps, nd + eps, nd)
+        sinc = torch.where(nd < eps, nd.new_ones(nd.shape), torch.sin(nd_) / nd_)
+
+        # cos is applied to last dim instead of first
+        exp_delta = torch.cat([sinc * delta, torch.cos(nd)], -1)
+
+        v, beta = cls.householder_vector(x)
+        return nx * cls.apply_householder(exp_delta, v, beta)

scripts/camera/geometry/perspective_fields.py

@@ -0,0 +1,379 @@
+"""Implementation of perspective fields.
+
+Adapted from https://github.com/jinlinyi/PerspectiveFields/blob/main/perspective2d/utils/panocam.py
+"""
+
+from typing import Tuple
+
+import torch
+from torch.nn import functional as F
+
+from scripts.camera.geometry.base_camera import BaseCamera
+from scripts.camera.geometry.gravity import Gravity
+from scripts.camera.geometry.jacobians import J_up_projection, J_vecnorm
+from scripts.camera.geometry.manifolds import SphericalManifold
+
+# flake8: noqa: E266
+
+
+def get_horizon_line(camera: BaseCamera, gravity: Gravity, relative: bool = True) -> torch.Tensor:
+    """Get the horizon line from the camera parameters.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        relative (bool, optional): Whether to normalize horizon line by img_h. Defaults to True.
+
+    Returns:
+        torch.Tensor: In image frame, fraction of image left/right border intersection with
+        respect to image height.
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    # project horizon midpoint to image plane
+    horizon_midpoint = camera.new_tensor([0, 0, 1])
+    horizon_midpoint = camera.K @ gravity.R @ horizon_midpoint
+    midpoint = horizon_midpoint[:2] / horizon_midpoint[2]
+
+    # compute left and right offset to borders
+    left_offset = midpoint[0] * torch.tan(gravity.roll)
+    right_offset = (camera.size[0] - midpoint[0]) * torch.tan(gravity.roll)
+    left, right = midpoint[1] + left_offset, midpoint[1] - right_offset
+
+    horizon = camera.new_tensor([left, right])
+    return horizon / camera.size[1] if relative else horizon
+
+
+def get_up_field(camera: BaseCamera, gravity: Gravity, normalize: bool = True) -> torch.Tensor:
+    """Get the up vector field from the camera parameters.
+
+    Args:
+        camera (Camera): Camera parameters.
+        normalize (bool, optional): Whether to normalize the up vector. Defaults to True.
+
+    Returns:
+        torch.Tensor: up vector field as tensor of shape (..., h, w, 2).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    uv = camera.normalize(camera.pixel_coordinates())
+
+    # projected up is (a, b) - c * (u, v)
+    abc = gravity.vec3d
+    projected_up2d = abc[..., None, :2] - abc[..., 2, None, None] * uv  # (..., N, 2)
+
+    if hasattr(camera, "dist"):
+        d_uv = camera.distort(uv, return_scale=True)[0]  # (..., N, 1)
+        d_uv = torch.diag_embed(d_uv.expand(d_uv.shape[:-1] + (2,)))  # (..., N, 2, 2)
+        offset = camera.up_projection_offset(uv)  # (..., N, 2)
+        offset = torch.einsum("...i,...j->...ij", offset, uv)  # (..., N, 2, 2)
+
+        # (..., N, 2)
+        projected_up2d = torch.einsum("...Nij,...Nj->...Ni", d_uv + offset, projected_up2d)
+
+    if normalize:
+        projected_up2d = F.normalize(projected_up2d, dim=-1)  # (..., N, 2)
+        
+    try:
+        del uv, abc, d_uv, offset
+    except NameError:
+        pass
+
+    return projected_up2d.reshape(camera.shape[0], h, w, 2)
+
+
+def J_up_field(
+    camera: BaseCamera, gravity: Gravity, spherical: bool = False, log_focal: bool = False
+) -> torch.Tensor:
+    """Get the jacobian of the up field.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        spherical (bool, optional): Whether to use spherical coordinates. Defaults to False.
+        log_focal (bool, optional): Whether to use log-focal length. Defaults to False.
+
+    Returns:
+        torch.Tensor: Jacobian of the up field as a tensor of shape (..., h, w, 2, 2, 3).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    # Forward
+    xy = camera.pixel_coordinates()
+    uv = camera.normalize(xy)
+
+    projected_up2d = gravity.vec3d[..., None, :2] - gravity.vec3d[..., 2, None, None] * uv
+
+    # Backward
+    J = []
+
+    # (..., N, 2, 2)
+    J_norm2proj = J_vecnorm(
+        get_up_field(camera, gravity, normalize=False).reshape(camera.shape[0], -1, 2)
+    )
+
+    # distortion values
+    if hasattr(camera, "dist"):
+        d_uv = camera.distort(uv, return_scale=True)[0]  # (..., N, 1)
+        d_uv = torch.diag_embed(d_uv.expand(d_uv.shape[:-1] + (2,)))  # (..., N, 2, 2)
+        offset = camera.up_projection_offset(uv)  # (..., N, 2)
+        offset_uv = torch.einsum("...i,...j->...ij", offset, uv)  # (..., N, 2, 2)
+
+    ######################
+    ## Gravity Jacobian ##
+    ######################
+
+    J_proj2abc = J_up_projection(uv, gravity.vec3d, wrt="abc")  # (..., N, 2, 3)
+
+    if hasattr(camera, "dist"):
+        # (..., N, 2, 3)
+        J_proj2abc = torch.einsum("...Nij,...Njk->...Nik", d_uv + offset_uv, J_proj2abc)
+
+    J_abc2delta = SphericalManifold.J_plus(gravity.vec3d) if spherical else gravity.J_rp()
+    J_proj2delta = torch.einsum("...Nij,...jk->...Nik", J_proj2abc, J_abc2delta)
+    J_up2delta = torch.einsum("...Nij,...Njk->...Nik", J_norm2proj, J_proj2delta)
+    J.append(J_up2delta)
+
+    ######################
+    ### Focal Jacobian ###
+    ######################
+
+    J_proj2uv = J_up_projection(uv, gravity.vec3d, wrt="uv")  # (..., N, 2, 2)
+
+    if hasattr(camera, "dist"):
+        J_proj2up = torch.einsum("...Nij,...Njk->...Nik", d_uv + offset_uv, J_proj2uv)
+        J_proj2duv = torch.einsum("...i,...j->...ji", offset, projected_up2d)
+
+        inner = (uv * projected_up2d).sum(-1)[..., None, None]
+        J_proj2offset1 = inner * camera.J_up_projection_offset(uv, wrt="uv")
+        J_proj2offset2 = torch.einsum("...i,...j->...ij", offset, projected_up2d)  # (..., N, 2, 2)
+        J_proj2uv = (J_proj2duv + J_proj2offset1 + J_proj2offset2) + J_proj2up
+
+    J_uv2f = camera.J_normalize(xy)  # (..., N, 2, 2)
+
+    if log_focal:
+        J_uv2f = J_uv2f * camera.f[..., None, None, :]  # (..., N, 2, 2)
+
+    J_uv2f = J_uv2f.sum(-1)  # (..., N, 2)
+
+    J_proj2f = torch.einsum("...ij,...j->...i", J_proj2uv, J_uv2f)  # (..., N, 2)
+    J_up2f = torch.einsum("...Nij,...Nj->...Ni", J_norm2proj, J_proj2f)[..., None]  # (..., N, 2, 1)
+    J.append(J_up2f)
+
+    ######################
+    ##### K1 Jacobian ####
+    ######################
+
+    if hasattr(camera, "dist"):
+        J_duv = camera.J_distort(uv, wrt="scale2dist")
+        J_duv = torch.diag_embed(J_duv.expand(J_duv.shape[:-1] + (2,)))  # (..., N, 2, 2)
+        J_offset = torch.einsum(
+            "...i,...j->...ij", camera.J_up_projection_offset(uv, wrt="dist"), uv
+        )
+        J_proj2k1 = torch.einsum("...Nij,...Nj->...Ni", J_duv + J_offset, projected_up2d)
+        J_k1 = torch.einsum("...Nij,...Nj->...Ni", J_norm2proj, J_proj2k1)[..., None]
+        J.append(J_k1)
+
+    n_params = sum(j.shape[-1] for j in J)
+    return torch.cat(J, axis=-1).reshape(camera.shape[0], h, w, 2, n_params)
+
+
+def get_latitude_field(camera: BaseCamera, gravity: Gravity) -> torch.Tensor:
+    """Get the latitudes of the camera pixels in radians.
+
+    Latitudes are defined as the angle between the ray and the up vector.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+
+    Returns:
+        torch.Tensor: Latitudes in radians as a tensor of shape (..., h, w, 1).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    uv1, _ = camera.image2world(camera.pixel_coordinates())
+    rays = camera.pixel_bearing_many(uv1)
+
+    lat = torch.einsum("...Nj,...j->...N", rays, gravity.vec3d)
+
+    eps = 1e-6
+    lat_asin = torch.asin(lat.clamp(min=-1 + eps, max=1 - eps))
+    
+    try:
+        del uv1, rays
+    except NameError:
+        pass
+
+    return lat_asin.reshape(camera.shape[0], h, w, 1)
+
+
+def J_latitude_field(
+    camera: BaseCamera, gravity: Gravity, spherical: bool = False, log_focal: bool = False
+) -> torch.Tensor:
+    """Get the jacobian of the latitude field.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        spherical (bool, optional): Whether to use spherical coordinates. Defaults to False.
+        log_focal (bool, optional): Whether to use log-focal length. Defaults to False.
+
+    Returns:
+        torch.Tensor: Jacobian of the latitude field as a tensor of shape (..., h, w, 1, 3).
+    """
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    # Forward
+    xy = camera.pixel_coordinates()
+    uv1, _ = camera.image2world(xy)
+    uv1_norm = camera.pixel_bearing_many(uv1)  # (..., N, 3)
+
+    # Backward
+    J = []
+    J_norm2w_to_img = J_vecnorm(uv1)[..., :2]  # (..., N, 2)
+
+    ######################
+    ## Gravity Jacobian ##
+    ######################
+
+    J_delta = SphericalManifold.J_plus(gravity.vec3d) if spherical else gravity.J_rp()
+    J_delta = torch.einsum("...Ni,...ij->...Nj", uv1_norm, J_delta)  # (..., N, 2)
+    J.append(J_delta)
+
+    ######################
+    ### Focal Jacobian ###
+    ######################
+
+    J_w_to_img2f = camera.J_image2world(xy, "f")  # (..., N, 2, 2)
+    if log_focal:
+        J_w_to_img2f = J_w_to_img2f * camera.f[..., None, None, :]
+    J_w_to_img2f = J_w_to_img2f.sum(-1)  # (..., N, 2)
+
+    J_norm2f = torch.einsum("...Nij,...Nj->...Ni", J_norm2w_to_img, J_w_to_img2f)  # (..., N, 3)
+    J_f = torch.einsum("...Ni,...i->...N", J_norm2f, gravity.vec3d).unsqueeze(-1)  # (..., N, 1)
+    J.append(J_f)
+
+    ######################
+    ##### K1 Jacobian ####
+    ######################
+
+    if hasattr(camera, "dist"):
+        J_w_to_img2k1 = camera.J_image2world(xy, "dist")  # (..., N, 2)
+        # (..., N, 2)
+        J_norm2k1 = torch.einsum("...Nij,...Nj->...Ni", J_norm2w_to_img, J_w_to_img2k1)
+        # (..., N, 1)
+        J_k1 = torch.einsum("...Ni,...i->...N", J_norm2k1, gravity.vec3d).unsqueeze(-1)
+        J.append(J_k1)
+
+    n_params = sum(j.shape[-1] for j in J)
+    return torch.cat(J, axis=-1).reshape(camera.shape[0], h, w, 1, n_params)
+
+
+def get_perspective_field(
+    camera: BaseCamera,
+    gravity: Gravity,
+    use_up: bool = True,
+    use_latitude: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Get the perspective field from the camera parameters.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        use_up (bool, optional): Whether to include the up vector field. Defaults to True.
+        use_latitude (bool, optional): Whether to include the latitude field. Defaults to True.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Up and latitude fields as tensors of shape
+        (..., 2, h, w) and (..., 1, h, w).
+    """
+    assert use_up or use_latitude, "At least one of use_up or use_latitude must be True."
+
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    if use_up:
+        permute = (0, 3, 1, 2)
+        # (..., 2, h, w)
+        up = get_up_field(camera, gravity).permute(permute)
+    else:
+        shape = (camera.shape[0], 2, h, w)
+        up = camera.new_zeros(shape)
+
+    if use_latitude:
+        permute = (0, 3, 1, 2)
+        # (..., 1, h, w)
+        lat = get_latitude_field(camera, gravity).permute(permute)
+    else:
+        shape = (camera.shape[0], 1, h, w)
+        lat = camera.new_zeros(shape)
+        
+    torch.cuda.empty_cache()
+
+    return up, lat
+
+
+def J_perspective_field(
+    camera: BaseCamera,
+    gravity: Gravity,
+    use_up: bool = True,
+    use_latitude: bool = True,
+    spherical: bool = False,
+    log_focal: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Get the jacobian of the perspective field.
+
+    Args:
+        camera (Camera): Camera parameters.
+        gravity (Gravity): Gravity vector.
+        use_up (bool, optional): Whether to include the up vector field. Defaults to True.
+        use_latitude (bool, optional): Whether to include the latitude field. Defaults to True.
+        spherical (bool, optional): Whether to use spherical coordinates. Defaults to False.
+        log_focal (bool, optional): Whether to use log-focal length. Defaults to False.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Up and latitude jacobians as tensors of shape
+        (..., h, w, 2, 4) and (..., h, w, 1, 4).
+    """
+    assert use_up or use_latitude, "At least one of use_up or use_latitude must be True."
+
+    camera = camera.unsqueeze(0) if len(camera.shape) == 0 else camera
+    gravity = gravity.unsqueeze(0) if len(gravity.shape) == 0 else gravity
+
+    w, h = camera.size[0].unbind(-1)
+    h, w = h.round().to(int), w.round().to(int)
+
+    if use_up:
+        J_up = J_up_field(camera, gravity, spherical, log_focal)  # (..., h, w, 2, 4)
+    else:
+        shape = (camera.shape[0], h, w, 2, 4)
+        J_up = camera.new_zeros(shape)
+
+    if use_latitude:
+        J_lat = J_latitude_field(camera, gravity, spherical, log_focal)  # (..., h, w, 1, 4)
+    else:
+        shape = (camera.shape[0], h, w, 1, 4)
+        J_lat = camera.new_zeros(shape)
+
+    return J_up, J_lat

scripts/camera/utils/conversions.py

@@ -0,0 +1,150 @@
+"""Utility functions for conversions between different representations."""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+from typing import Optional
+
+import torch
+
+
+def skew_symmetric(v: torch.Tensor) -> torch.Tensor:
+    """Create a skew-symmetric matrix from a (batched) vector of size (..., 3).
+
+    Args:
+        (torch.Tensor): Vector of size (..., 3).
+
+    Returns:
+        (torch.Tensor): Skew-symmetric matrix of size (..., 3, 3).
+    """
+    z = torch.zeros_like(v[..., 0])
+    return torch.stack(
+        [
+            z,
+            -v[..., 2],
+            v[..., 1],
+            v[..., 2],
+            z,
+            -v[..., 0],
+            -v[..., 1],
+            v[..., 0],
+            z,
+        ],
+        dim=-1,
+    ).reshape(v.shape[:-1] + (3, 3))
+
+
+def rad2rotmat(
+    roll: torch.Tensor, pitch: torch.Tensor, yaw: Optional[torch.Tensor] = None
+) -> torch.Tensor:
+    """Convert (batched) roll, pitch, yaw angles (in radians) to rotation matrix.
+
+    Args:
+        roll (torch.Tensor): Roll angle in radians.
+        pitch (torch.Tensor): Pitch angle in radians.
+        yaw (torch.Tensor, optional): Yaw angle in radians. Defaults to None.
+
+    Returns:
+        torch.Tensor: Rotation matrix of shape (..., 3, 3).
+    """
+    if yaw is None:
+        yaw = roll.new_zeros(roll.shape)
+
+    Rx = pitch.new_zeros(pitch.shape + (3, 3))
+    Rx[..., 0, 0] = 1
+    Rx[..., 1, 1] = torch.cos(pitch)
+    Rx[..., 1, 2] = torch.sin(pitch)
+    Rx[..., 2, 1] = -torch.sin(pitch)
+    Rx[..., 2, 2] = torch.cos(pitch)
+
+    Ry = yaw.new_zeros(yaw.shape + (3, 3))
+    Ry[..., 0, 0] = torch.cos(yaw)
+    Ry[..., 0, 2] = -torch.sin(yaw)
+    Ry[..., 1, 1] = 1
+    Ry[..., 2, 0] = torch.sin(yaw)
+    Ry[..., 2, 2] = torch.cos(yaw)
+
+    Rz = roll.new_zeros(roll.shape + (3, 3))
+    Rz[..., 0, 0] = torch.cos(roll)
+    Rz[..., 0, 1] = torch.sin(roll)
+    Rz[..., 1, 0] = -torch.sin(roll)
+    Rz[..., 1, 1] = torch.cos(roll)
+    Rz[..., 2, 2] = 1
+
+    return Rz @ Rx @ Ry
+
+
+def fov2focal(fov: torch.Tensor, size: torch.Tensor) -> torch.Tensor:
+    """Compute focal length from (vertical/horizontal) field of view.
+
+    Args:
+        fov (torch.Tensor): Field of view in radians.
+        size (torch.Tensor): Image height / width in pixels.
+
+    Returns:
+        torch.Tensor: Focal length in pixels.
+    """
+    return size / 2 / torch.tan(fov / 2)
+
+
+def focal2fov(focal: torch.Tensor, size: torch.Tensor) -> torch.Tensor:
+    """Compute (vertical/horizontal) field of view from focal length.
+
+    Args:
+        focal (torch.Tensor): Focal length in pixels.
+        size (torch.Tensor): Image height / width in pixels.
+
+    Returns:
+        torch.Tensor: Field of view in radians.
+    """
+    return 2 * torch.arctan(size / (2 * focal))
+
+
+def pitch2rho(pitch: torch.Tensor, f: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+    """Compute the distance from principal point to the horizon.
+
+    Args:
+        pitch (torch.Tensor): Pitch angle in radians.
+        f (torch.Tensor): Focal length in pixels.
+        h (torch.Tensor): Image height in pixels.
+
+    Returns:
+        torch.Tensor: Relative distance to the horizon.
+    """
+    return torch.tan(pitch) * f / h
+
+
+def rho2pitch(rho: torch.Tensor, f: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+    """Compute the pitch angle from the distance to the horizon.
+
+    Args:
+        rho (torch.Tensor): Relative distance to the horizon.
+        f (torch.Tensor): Focal length in pixels.
+        h (torch.Tensor): Image height in pixels.
+
+    Returns:
+        torch.Tensor: Pitch angle in radians.
+    """
+    return torch.atan(rho * h / f)
+
+
+def rad2deg(rad: torch.Tensor) -> torch.Tensor:
+    """Convert radians to degrees.
+
+    Args:
+        rad (torch.Tensor): Angle in radians.
+
+    Returns:
+        torch.Tensor: Angle in degrees.
+    """
+    return rad / torch.pi * 180
+
+
+def deg2rad(deg: torch.Tensor) -> torch.Tensor:
+    """Convert degrees to radians.
+
+    Args:
+        deg (torch.Tensor): Angle in degrees.
+
+    Returns:
+        torch.Tensor: Angle in radians.
+    """
+    return deg / 180 * torch.pi

scripts/camera/utils/image.py

@@ -0,0 +1,182 @@
+"""Image preprocessing utilities."""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+import collections.abc as collections
+from pathlib import Path
+from typing import Optional, Tuple
+
+import cv2
+import kornia
+import numpy as np
+import torch
+import torchvision
+from omegaconf import OmegaConf
+from PIL import Image
+
+from tensor import fit_features_to_multiple
+
+# mypy: ignore-errors
+
+
+class ImagePreprocessor:
+    """Preprocess images for calibration."""
+
+    default_conf = {
+        "resize": None,  # target edge length (320), None for no resizing
+        "edge_divisible_by": None,
+        "side": "short",
+        "interpolation": "bilinear",
+        "align_corners": None,
+        "antialias": True,
+        "square_crop": False,
+        "add_padding_mask": False,
+        "resize_backend": "kornia",  # torchvision, kornia
+    }
+
+    def __init__(self, conf) -> None:
+        """Initialize the image preprocessor."""
+        super().__init__()
+        default_conf = OmegaConf.create(self.default_conf)
+        OmegaConf.set_struct(default_conf, True)
+        self.conf = OmegaConf.merge(default_conf, conf)
+
+    def __call__(self, img: torch.Tensor, interpolation: Optional[str] = None) -> dict:
+        """Resize and preprocess an image, return image and resize scale."""
+        h, w = img.shape[-2:]
+        size = h, w
+
+        if self.conf.square_crop:
+            min_size = min(h, w)
+            offset = (h - min_size) // 2, (w - min_size) // 2
+            img = img[:, offset[0] : offset[0] + min_size, offset[1] : offset[1] + min_size]
+            size = img.shape[-2:]
+
+        if self.conf.resize is not None:
+            if interpolation is None:
+                interpolation = self.conf.interpolation
+            size = self.get_new_image_size(h, w)
+            img = self.resize(img, size, interpolation)
+
+        scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img)
+        T = np.diag([scale[0].cpu(), scale[1].cpu(), 1])
+
+        data = {
+            "scales": scale,
+            "image_size": np.array(size[::-1]),
+            "transform": T,
+            "original_image_size": np.array([w, h]),
+        }
+
+        if self.conf.edge_divisible_by is not None:
+            # crop to make the edge divisible by a number
+            w_, h_ = img.shape[-1], img.shape[-2]
+            img, _ = fit_features_to_multiple(img, self.conf.edge_divisible_by, crop=True)
+            crop_pad = torch.Tensor([img.shape[-1] - w_, img.shape[-2] - h_]).to(img)
+            data["crop_pad"] = crop_pad
+            data["image_size"] = np.array([img.shape[-1], img.shape[-2]])
+
+        data["image"] = img
+        return data
+
+    def resize(self, img: torch.Tensor, size: Tuple[int, int], interpolation: str) -> torch.Tensor:
+        """Resize an image using the specified backend."""
+        if self.conf.resize_backend == "kornia":
+            return kornia.geometry.transform.resize(
+                img,
+                size,
+                side=self.conf.side,
+                antialias=self.conf.antialias,
+                align_corners=self.conf.align_corners,
+                interpolation=interpolation,
+            )
+        elif self.conf.resize_backend == "PIL":
+            device = img.device
+            imgs = []
+            has_batch_dim = img.ndim == 4
+            img = img if has_batch_dim else img[None]
+            for im in img:
+                im = (im.permute(1, 2, 0) * 255).cpu().numpy().astype(np.uint8)
+                im = Image.fromarray(im).resize(size[::-1], Image.BILINEAR)
+                im = torch.tensor(np.array(im)).permute(2, 0, 1) / 255.0
+                imgs.append(im.to(device))
+            imgs = torch.stack(imgs)
+            return imgs if has_batch_dim else imgs[0]
+
+        elif self.conf.resize_backend == "torchvision":
+            return torchvision.transforms.Resize(size, antialias=self.conf.antialias)(img)
+        else:
+            raise ValueError(f"{self.conf.resize_backend} not implemented.")
+
+    def load_image(self, image_path: Path) -> dict:
+        """Load an image from a path and preprocess it."""
+        return self(load_image(image_path))
+
+    def get_new_image_size(self, h: int, w: int) -> Tuple[int, int]:
+        """Get the new image size after resizing."""
+        side = self.conf.side
+        if isinstance(self.conf.resize, collections.Iterable):
+            assert len(self.conf.resize) == 2
+            return tuple(self.conf.resize)
+        side_size = self.conf.resize
+        aspect_ratio = w / h
+        if side not in ("short", "long", "vert", "horz"):
+            raise ValueError(
+                f"side can be one of 'short', 'long', 'vert', and 'horz'. Got '{side}'"
+            )
+        return (
+            (side_size, int(side_size * aspect_ratio))
+            if side == "vert" or (side != "horz" and (side == "short") ^ (aspect_ratio < 1.0))
+            else (int(side_size / aspect_ratio), side_size)
+        )
+
+
+def numpy_image_to_torch(image: np.ndarray) -> torch.Tensor:
+    """Normalize the image tensor and reorder the dimensions."""
+    if image.ndim == 3:
+        image = image.transpose((2, 0, 1))  # HxWxC to CxHxW
+    elif image.ndim == 2:
+        image = image[None]  # add channel axis
+    else:
+        raise ValueError(f"Not an image: {image.shape}")
+    return torch.tensor(image / 255.0, dtype=torch.float)
+
+
+def torch_image_to_numpy(image: torch.Tensor) -> np.ndarray:
+    """Normalize and reorder the dimensions of an image tensor."""
+    if image.ndim == 3:
+        image = image.permute((1, 2, 0))  # CxHxW to HxWxC
+    elif image.ndim == 2:
+        image = image[None]  # add channel axis
+    else:
+        raise ValueError(f"Not an image: {image.shape}")
+    return (image.cpu().detach().numpy() * 255).astype(np.uint8)
+
+
+def read_image(path: Path, grayscale: bool = False) -> np.ndarray:
+    """Read an image from path as RGB or grayscale."""
+    if not Path(path).exists():
+        raise FileNotFoundError(f"No image at path {path}.")
+    mode = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR
+    image = cv2.imread(str(path), mode)
+    if image is None:
+        raise IOError(f"Could not read image at {path}.")
+    if not grayscale:
+        image = image[..., ::-1]
+    return image
+
+
+def write_image(img: torch.Tensor, path: Path):
+    """Write an image tensor to a file."""
+    img = torch_image_to_numpy(img) if isinstance(img, torch.Tensor) else img
+    cv2.imwrite(str(path), img[..., ::-1])
+
+
+def load_image(path: Path, grayscale: bool = False, return_tensor: bool = True) -> torch.Tensor:
+    """Load an image from a path and return as a tensor."""
+    image = read_image(path, grayscale=grayscale)
+    if return_tensor:
+        return numpy_image_to_torch(image)
+
+    assert image.ndim in [2, 3], f"Not an image: {image.shape}"
+    image = image[None] if image.ndim == 2 else image
+    return torch.tensor(image.copy(), dtype=torch.uint8)

scripts/camera/utils/tensor.py

@@ -0,0 +1,249 @@
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+import collections.abc as collections
+import functools
+import inspect
+from typing import Callable, List, Tuple
+
+import numpy as np
+import torch
+
+# flake8: noqa
+# mypy: ignore-errors
+
+
+string_classes = (str, bytes)
+
+
+def autocast(func: Callable) -> Callable:
+    """Cast the inputs of a TensorWrapper method to PyTorch tensors if they are numpy arrays.
+
+    Use the device and dtype of the wrapper.
+
+    Args:
+        func (Callable): Method of a TensorWrapper class.
+
+    Returns:
+        Callable: Wrapped method.
+    """
+
+    @functools.wraps(func)
+    def wrap(self, *args):
+        device = torch.device("cpu")
+        dtype = None
+        if isinstance(self, TensorWrapper):
+            if self._data is not None:
+                device = self.device
+                dtype = self.dtype
+        elif not inspect.isclass(self) or not issubclass(self, TensorWrapper):
+            raise ValueError(self)
+
+        cast_args = []
+        for arg in args:
+            if isinstance(arg, np.ndarray):
+                arg = torch.from_numpy(arg)
+                arg = arg.to(device=device, dtype=dtype)
+            cast_args.append(arg)
+        return func(self, *cast_args)
+
+    return wrap
+
+
+class TensorWrapper:
+    """Wrapper for PyTorch tensors."""
+
+    _data = None
+
+    @autocast
+    def __init__(self, data: torch.Tensor):
+        """Wrapper for PyTorch tensors."""
+        self._data = data
+
+    @property
+    def shape(self) -> torch.Size:
+        """Shape of the underlying tensor."""
+        return self._data.shape[:-1]
+
+    @property
+    def device(self) -> torch.device:
+        """Get the device of the underlying tensor."""
+        return self._data.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Get the dtype of the underlying tensor."""
+        return self._data.dtype
+
+    def __getitem__(self, index) -> torch.Tensor:
+        """Get the underlying tensor."""
+        return self.__class__(self._data[index])
+
+    def __setitem__(self, index, item):
+        """Set the underlying tensor."""
+        self._data[index] = item.data
+
+    def to(self, *args, **kwargs):
+        """Move the underlying tensor to a new device."""
+        return self.__class__(self._data.to(*args, **kwargs))
+
+    def cpu(self):
+        """Move the underlying tensor to the CPU."""
+        return self.__class__(self._data.cpu())
+
+    def cuda(self):
+        """Move the underlying tensor to the GPU."""
+        return self.__class__(self._data.cuda())
+
+    def pin_memory(self):
+        """Pin the underlying tensor to memory."""
+        return self.__class__(self._data.pin_memory())
+
+    def float(self):
+        """Cast the underlying tensor to float."""
+        return self.__class__(self._data.float())
+
+    def double(self):
+        """Cast the underlying tensor to double."""
+        return self.__class__(self._data.double())
+
+    def detach(self):
+        """Detach the underlying tensor."""
+        return self.__class__(self._data.detach())
+
+    def numpy(self):
+        """Convert the underlying tensor to a numpy array."""
+        return self._data.detach().cpu().numpy()
+
+    def new_tensor(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_tensor(*args, **kwargs)
+
+    def new_zeros(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_zeros(*args, **kwargs)
+
+    def new_ones(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_ones(*args, **kwargs)
+
+    def new_full(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_full(*args, **kwargs)
+
+    def new_empty(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self._data.new_empty(*args, **kwargs)
+
+    def unsqueeze(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self.__class__(self._data.unsqueeze(*args, **kwargs))
+
+    def squeeze(self, *args, **kwargs):
+        """Create a new tensor of the same type and device."""
+        return self.__class__(self._data.squeeze(*args, **kwargs))
+
+    @classmethod
+    def stack(cls, objects: List, dim=0, *, out=None):
+        """Stack a list of objects with the same type and shape."""
+        data = torch.stack([obj._data for obj in objects], dim=dim, out=out)
+        return cls(data)
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        """Support torch functions."""
+        if kwargs is None:
+            kwargs = {}
+        return cls.stack(*args, **kwargs) if func is torch.stack else NotImplemented
+
+
+def map_tensor(input_, func):
+    if isinstance(input_, string_classes):
+        return input_
+    elif isinstance(input_, collections.Mapping):
+        return {k: map_tensor(sample, func) for k, sample in input_.items()}
+    elif isinstance(input_, collections.Sequence):
+        return [map_tensor(sample, func) for sample in input_]
+    elif input_ is None:
+        return None
+    else:
+        return func(input_)
+
+
+def batch_to_numpy(batch):
+    return map_tensor(batch, lambda tensor: tensor.cpu().numpy())
+
+
+def batch_to_device(batch, device, non_blocking=True, detach=False):
+    def _func(tensor):
+        t = tensor.to(device=device, non_blocking=non_blocking, dtype=torch.float32)
+        return t.detach() if detach else t
+
+    return map_tensor(batch, _func)
+
+
+def remove_batch_dim(data: dict) -> dict:
+    """Remove batch dimension from elements in data"""
+    return {
+        k: v[0] if isinstance(v, (torch.Tensor, np.ndarray, list)) else v for k, v in data.items()
+    }
+
+
+def add_batch_dim(data: dict) -> dict:
+    """Add batch dimension to elements in data"""
+    return {
+        k: v[None] if isinstance(v, (torch.Tensor, np.ndarray, list)) else v
+        for k, v in data.items()
+    }
+
+
+def fit_to_multiple(x: torch.Tensor, multiple: int, mode: str = "center", crop: bool = False):
+    """Get padding to make the image size a multiple of the given number.
+
+    Args:
+        x (torch.Tensor): Input tensor.
+        multiple (int, optional): Multiple.
+        crop (bool, optional): Whether to crop or pad. Defaults to False.
+
+    Returns:
+        torch.Tensor: Padding.
+    """
+    h, w = x.shape[-2:]
+
+    if crop:
+        pad_w = (w // multiple) * multiple - w
+        pad_h = (h // multiple) * multiple - h
+    else:
+        pad_w = (multiple - w % multiple) % multiple
+        pad_h = (multiple - h % multiple) % multiple
+
+    if mode == "center":
+        pad_l = pad_w // 2
+        pad_r = pad_w - pad_l
+        pad_t = pad_h // 2
+        pad_b = pad_h - pad_t
+    elif mode == "left":
+        pad_l = 0
+        pad_r = pad_w
+        pad_t = 0
+        pad_b = pad_h
+    else:
+        raise ValueError(f"Unknown mode {mode}")
+
+    return (pad_l, pad_r, pad_t, pad_b)
+
+
+def fit_features_to_multiple(
+    features: torch.Tensor, multiple: int = 32, crop: bool = False
+) -> Tuple[torch.Tensor, Tuple[int, int]]:
+    """Pad image to a multiple of the given number.
+
+    Args:
+        features (torch.Tensor): Input features.
+        multiple (int, optional): Multiple. Defaults to 32.
+        crop (bool, optional): Whether to crop or pad. Defaults to False.
+
+    Returns:
+        Tuple[torch.Tensor, Tuple[int, int]]: Padded features and padding.
+    """
+    pad = fit_to_multiple(features, multiple, crop=crop)
+    return torch.nn.functional.pad(features, pad, mode="reflect"), pad

scripts/camera/utils/text.py

@@ -0,0 +1,47 @@
+import re
+from typing import Tuple
+
+def parse_camera_params(
+    text: str,
+    mode: str = "base"
+) -> Tuple[float, float, float]:
+    """
+    Extract roll, pitch, fov from text using one of two patterns:
+      - 'base'   mode: ... are: roll, pitch, fov.
+      - 'cot'    mode: <answer>roll, pitch, fov</answer>
+
+    Args:
+        text: The full text to search.
+        mode: One of {"base", "cot"}.
+
+    Returns:
+        roll, pitch, fov as floats.
+
+    Raises:
+        ValueError if the chosen pattern is not found, or mode is invalid.
+    """
+    # compile both regexes
+    pat_base = re.compile(
+        r"are:\s*([+-]?\d+(?:\.\d+)?)\s*,\s*"
+        r"([+-]?\d+(?:\.\d+)?)\s*,\s*"
+        r"([+-]?\d+(?:\.\d+)?)[\.\s]*$"
+    )
+    pat_cot = re.compile(
+        r"<answer>\s*([+-]?\d+(?:\.\d+)?)\s*,\s*"
+        r"([+-]?\d+(?:\.\d+)?)\s*,\s*"
+        r"([+-]?\d+(?:\.\d+)?)\s*</answer>"
+    )
+
+    m = None
+    if mode == "base":
+        m = pat_base.search(text)
+    elif mode == "cot":
+        m = pat_cot.search(text)
+    else:
+        raise ValueError(f"Invalid mode: {mode!r}. Choose 'base', 'cot', or 'auto'.")
+
+    if not m:
+        raise ValueError(f"No camera parameters found using mode '{mode}'.")
+
+    roll_s, pitch_s, fov_s = m.group(1), m.group(2), m.group(3)
+    return float(roll_s), float(pitch_s), float(fov_s)

scripts/camera/visualization/visualize_batch.py

@@ -0,0 +1,188 @@
+"""Visualization of predicted and ground truth for a single batch."""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+from typing import Any, Dict
+
+import numpy as np
+import torch
+
+from scripts.camera.geometry.perspective_fields import get_latitude_field
+from scripts.camera.utils.conversions import rad2deg
+from scripts.camera.utils.tensor import batch_to_device
+from scripts.camera.visualization.viz2d import (
+    plot_confidences,
+    plot_heatmaps,
+    plot_image_grid,
+    plot_latitudes,
+    plot_vector_fields,
+)
+
+
+def make_up_figure(
+    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
+) -> Dict[str, Any]:
+    """Get predicted and ground truth up fields and errors.
+
+    Args:
+        pred (Dict[str, torch.Tensor]): Predicted up field.
+        data (Dict[str, torch.Tensor]): Ground truth up field.
+        n_pairs (int): Number of pairs to visualize.
+
+    Returns:
+        Dict[str, Any]: Dictionary with figure.
+    """
+    pred = batch_to_device(pred, "cpu", detach=True)
+    data = batch_to_device(data, "cpu", detach=True)
+
+    n_pairs = min(n_pairs, len(data["image"]))
+
+    if "up_field" not in pred.keys():
+        return {}
+
+    up_fields = []
+    for i in range(n_pairs):
+        row = [data["up_field"][i]]
+        titles = ["Up GT"]
+
+        if "up_confidence" in pred.keys():
+            row += [pred["up_confidence"][i]]
+            titles += ["Up Confidence"]
+
+        row = [r.float().numpy() if isinstance(r, torch.Tensor) else r for r in row]
+        up_fields.append(row)
+
+    # create figure
+    N, M = len(up_fields), len(up_fields[0]) + 1
+    imgs = [[data["image"][i].permute(1, 2, 0).cpu().clip(0, 1)] * M for i in range(n_pairs)]
+    fig, ax = plot_image_grid(imgs, return_fig=True, set_lim=True)
+    ax = np.array(ax)
+
+    for i in range(n_pairs):
+        plot_vector_fields([up_fields[i][0]], axes=ax[i, [1]])
+        #plot_heatmaps([up_fields[i][2]], cmap="turbo", colorbar=True, axes=ax[i, [3]])
+
+        if "up_confidence" in pred.keys():
+            plot_confidences([up_fields[i][3]], axes=ax[i, [4]])
+
+    return {"up": fig}
+
+
+def make_latitude_figure(
+    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
+) -> Dict[str, Any]:
+    """Get predicted and ground truth latitude fields and errors.
+
+    Args:
+        pred (Dict[str, torch.Tensor]): Predicted latitude field.
+        data (Dict[str, torch.Tensor]): Ground truth latitude field.
+        n_pairs (int, optional): Number of pairs to visualize. Defaults to 2.
+
+    Returns:
+        Dict[str, Any]: Dictionary with figure.
+    """
+    pred = batch_to_device(pred, "cpu", detach=True)
+    data = batch_to_device(data, "cpu", detach=True)
+
+    n_pairs = min(n_pairs, len(data["image"]))
+    latitude_fields = []
+
+    if "latitude_field" not in pred.keys():
+        return {}
+
+    for i in range(n_pairs):
+        row = [
+            rad2deg(data["latitude_field"][i][0]),
+            #rad2deg(pred["latitude_field"][i][0]),
+            #errors[i],
+        ]
+        titles = ["Latitude GT"]
+
+        if "latitude_confidence" in pred.keys():
+            row += [pred["latitude_confidence"][i]]
+            titles += ["Latitude Confidence"]
+
+        row = [r.float().numpy() if isinstance(r, torch.Tensor) else r for r in row]
+        latitude_fields.append(row)
+
+    # create figure
+    N, M = len(latitude_fields), len(latitude_fields[0]) + 1
+    imgs = [[data["image"][i].permute(1, 2, 0).cpu().clip(0, 1)] * M for i in range(n_pairs)]
+    fig, ax = plot_image_grid(imgs, return_fig=True, set_lim=True)
+    ax = np.array(ax)
+
+    for i in range(n_pairs):
+        plot_latitudes([latitude_fields[i][0]], is_radians=False, axes=ax[i, [1]])
+        #plot_heatmaps([latitude_fields[i][2]], cmap="turbo", colorbar=True, axes=ax[i, [3]])
+
+        if "latitude_confidence" in pred.keys():
+            plot_confidences([latitude_fields[i][3]], axes=ax[i, [4]])
+
+    return {"latitude": fig}
+
+
+def make_camera_figure(
+    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
+) -> Dict[str, Any]:
+    """Get predicted and ground truth camera parameters.
+
+    Args:
+        pred (Dict[str, torch.Tensor]): Predicted camera parameters.
+        data (Dict[str, torch.Tensor]): Ground truth camera parameters.
+        n_pairs (int, optional): Number of pairs to visualize. Defaults to 2.
+
+    Returns:
+        Dict[str, Any]: Dictionary with figure.
+    """
+    pred = batch_to_device(pred, "cpu", detach=True)
+    data = batch_to_device(data, "cpu", detach=True)
+
+    n_pairs = min(n_pairs, len(data["image"]))
+
+    if "camera" not in pred.keys():
+        return {}
+
+    latitudes = []
+    for i in range(n_pairs):
+        titles = ["Cameras GT"]
+        row = [get_latitude_field(data["camera"][i], data["gravity"][i])]
+
+        if "camera" in pred.keys() and "gravity" in pred.keys():
+            row += [get_latitude_field(pred["camera"][i], pred["gravity"][i])]
+            titles += ["Cameras Pred"]
+
+        row = [rad2deg(r).squeeze(-1).float().numpy()[0] for r in row]
+        latitudes.append(row)
+
+    # create figure
+    N, M = len(latitudes), len(latitudes[0]) + 1
+    imgs = [[data["image"][i].permute(1, 2, 0).cpu().clip(0, 1)] * M for i in range(n_pairs)]
+    fig, ax = plot_image_grid(imgs, titles=[["Image"] + titles] * N, return_fig=True, set_lim=True)
+    ax = np.array(ax)
+
+    for i in range(n_pairs):
+        plot_latitudes(latitudes[i], is_radians=False, axes=ax[i, 1:])
+
+    return {"camera": fig}
+
+
+def make_perspective_figures(
+    pred: Dict[str, torch.Tensor], data: Dict[str, torch.Tensor], n_pairs: int = 2
+) -> Dict[str, Any]:
+    """Get predicted and ground truth perspective fields.
+
+    Args:
+        pred (Dict[str, torch.Tensor]): Predicted perspective fields.
+        data (Dict[str, torch.Tensor]): Ground truth perspective fields.
+        n_pairs (int, optional): Number of pairs to visualize. Defaults to 2.
+
+    Returns:
+        Dict[str, Any]: Dictionary with figure.
+    """
+    n_pairs = min(n_pairs, len(data["image"]))
+    figures = make_up_figure(pred, data, n_pairs)
+    figures |= make_latitude_figure(pred, data, n_pairs)
+    #figures |= make_camera_figure(pred, data, n_pairs)
+
+    {f.tight_layout() for f in figures.values()}
+
+    return figures

scripts/camera/visualization/viz2d.py

@@ -0,0 +1,521 @@
+"""
+2D visualization primitives based on Matplotlib.
+1) Plot images with `plot_images`.
+2) Call TODO: add functions
+3) Optionally: save a .png or .pdf plot (nice in papers!) with `save_plot`.
+"""
+"""Adapted from https://github.com/cvg/GeoCalib"""
+
+import matplotlib.patheffects as path_effects
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from scripts.camera.geometry.perspective_fields import get_perspective_field
+from scripts.camera.utils.conversions import rad2deg
+
+# flake8: noqa
+# mypy: ignore-errors
+
+
+def cm_ranking(sc, ths=None):
+    if ths is None:
+        ths = [512, 1024, 2048, 4096]
+
+    ls = sc.shape[0]
+    colors = ["red", "yellow", "lime", "cyan", "blue"]
+    out = ["gray"] * ls
+    for i in range(ls):
+        for c, th in zip(colors[: len(ths) + 1], ths + [ls]):
+            if i < th:
+                out[i] = c
+                break
+    sid = np.argsort(sc, axis=0).flip(0)
+    return np.array(out)[sid]
+
+
+def cm_RdBl(x):
+    """Custom colormap: red (0) -> yellow (0.5) -> green (1)."""
+    x = np.clip(x, 0, 1)[..., None] * 2
+    c = x * np.array([[0, 0, 1.0]]) + (2 - x) * np.array([[1.0, 0, 0]])
+    return np.clip(c, 0, 1)
+
+
+def cm_RdGn(x):
+    """Custom colormap: red (0) -> yellow (0.5) -> green (1)."""
+    x = np.clip(x, 0, 1)[..., None] * 2
+    c = x * np.array([[0, 1.0, 0]]) + (2 - x) * np.array([[1.0, 0, 0]])
+    return np.clip(c, 0, 1)
+
+
+def cm_BlRdGn(x_):
+    """Custom colormap: blue (-1) -> red (0.0) -> green (1)."""
+    x = np.clip(x_, 0, 1)[..., None] * 2
+    c = x * np.array([[0, 1.0, 0, 1.0]]) + (2 - x) * np.array([[1.0, 0, 0, 1.0]])
+
+    xn = -np.clip(x_, -1, 0)[..., None] * 2
+    cn = xn * np.array([[0, 1.0, 0, 1.0]]) + (2 - xn) * np.array([[1.0, 0, 0, 1.0]])
+    return np.clip(np.where(x_[..., None] < 0, cn, c), 0, 1)
+
+
+def plot_images(imgs, titles=None, cmaps="gray", dpi=200, pad=0.5, adaptive=True):
+    """Plot a list of images.
+
+    Args:
+        imgs (List[np.ndarray]): List of images to plot.
+        titles (List[str], optional): Titles. Defaults to None.
+        cmaps (str, optional): Colormaps. Defaults to "gray".
+        dpi (int, optional): Dots per inch. Defaults to 200.
+        pad (float, optional): Padding. Defaults to 0.5.
+        adaptive (bool, optional): Whether to adapt the aspect ratio. Defaults to True.
+
+    Returns:
+        plt.Figure: Figure of the images.
+    """
+    n = len(imgs)
+    if not isinstance(cmaps, (list, tuple)):
+        cmaps = [cmaps] * n
+
+    ratios = [i.shape[1] / i.shape[0] for i in imgs] if adaptive else [4 / 3] * n
+    figsize = [sum(ratios) * 4.5, 4.5]
+    fig, axs = plt.subplots(1, n, figsize=figsize, dpi=dpi, gridspec_kw={"width_ratios": ratios})
+    if n == 1:
+        axs = [axs]
+    for i, (img, ax) in enumerate(zip(imgs, axs)):
+        ax.imshow(img, cmap=plt.get_cmap(cmaps[i]))
+        ax.set_axis_off()
+        if titles:
+            ax.set_title(titles[i])
+    fig.tight_layout(pad=pad)
+
+    return fig
+
+
+def plot_image_grid(
+    imgs,
+    titles=None,
+    cmaps="gray",
+    dpi=100,
+    pad=0.5,
+    fig=None,
+    adaptive=True,
+    figs=3.0,
+    return_fig=False,
+    set_lim=False,
+) -> plt.Figure:
+    """Plot a grid of images.
+
+    Args:
+        imgs (List[np.ndarray]): List of images to plot.
+        titles (List[str], optional): Titles. Defaults to None.
+        cmaps (str, optional): Colormaps. Defaults to "gray".
+        dpi (int, optional): Dots per inch. Defaults to 100.
+        pad (float, optional): Padding. Defaults to 0.5.
+        fig (_type_, optional): Figure to plot on. Defaults to None.
+        adaptive (bool, optional): Whether to adapt the aspect ratio. Defaults to True.
+        figs (float, optional): Figure size. Defaults to 3.0.
+        return_fig (bool, optional): Whether to return the figure. Defaults to False.
+        set_lim (bool, optional): Whether to set the limits. Defaults to False.
+
+    Returns:
+        plt.Figure: Figure and axes or just axes.
+    """
+    nr, n = len(imgs), len(imgs[0])
+    if not isinstance(cmaps, (list, tuple)):
+        cmaps = [cmaps] * n
+
+    if adaptive:
+        ratios = [i.shape[1] / i.shape[0] for i in imgs[0]]  # W / H
+    else:
+        ratios = [4 / 3] * n
+
+    figsize = [sum(ratios) * figs, nr * figs]
+    if fig is None:
+        fig, axs = plt.subplots(
+            nr, n, figsize=figsize, dpi=dpi, gridspec_kw={"width_ratios": ratios}
+        )
+    else:
+        axs = fig.subplots(nr, n, gridspec_kw={"width_ratios": ratios})
+        fig.figure.set_size_inches(figsize)
+
+    if nr == 1 and n == 1:
+        axs = [[axs]]
+    elif n == 1:
+        axs = axs[:, None]
+    elif nr == 1:
+        axs = [axs]
+
+    for j in range(nr):
+        for i in range(n):
+            ax = axs[j][i]
+            ax.imshow(imgs[j][i], cmap=plt.get_cmap(cmaps[i]))
+            ax.set_axis_off()
+            if set_lim:
+                ax.set_xlim([0, imgs[j][i].shape[1]])
+                ax.set_ylim([imgs[j][i].shape[0], 0])
+            if titles:
+                ax.set_title(titles[j][i])
+    if isinstance(fig, plt.Figure):
+        fig.tight_layout(pad=pad)
+    return (fig, axs) if return_fig else axs
+
+
+def add_text(
+    idx,
+    text,
+    pos=(0.01, 0.99),
+    fs=15,
+    color="w",
+    lcolor="k",
+    lwidth=4,
+    ha="left",
+    va="top",
+    axes=None,
+    **kwargs,
+):
+    """Add text to a plot.
+
+    Args:
+        idx (int): Index of the axes.
+        text (str): Text to add.
+        pos (tuple, optional): Text position. Defaults to (0.01, 0.99).
+        fs (int, optional): Font size. Defaults to 15.
+        color (str, optional): Text color. Defaults to "w".
+        lcolor (str, optional): Line color. Defaults to "k".
+        lwidth (int, optional): Line width. Defaults to 4.
+        ha (str, optional): Horizontal alignment. Defaults to "left".
+        va (str, optional): Vertical alignment. Defaults to "top".
+        axes (List[plt.Axes], optional): Axes to put text on. Defaults to None.
+
+    Returns:
+        plt.Text: Text object.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    ax = axes[idx]
+
+    t = ax.text(
+        *pos,
+        text,
+        fontsize=fs,
+        ha=ha,
+        va=va,
+        color=color,
+        transform=ax.transAxes,
+        zorder=5,
+        **kwargs,
+    )
+    if lcolor is not None:
+        t.set_path_effects(
+            [
+                path_effects.Stroke(linewidth=lwidth, foreground=lcolor),
+                path_effects.Normal(),
+            ]
+        )
+    return t
+
+
+def plot_heatmaps(
+    heatmaps,
+    vmin=-1e-6,  # include negative zero
+    vmax=None,
+    cmap="Spectral",
+    a=0.5,
+    axes=None,
+    contours_every=None,
+    contour_style="solid",
+    colorbar=False,
+):
+    """Plot heatmaps with optional contours.
+
+    To plot latitude field, set vmin=-90, vmax=90 and contours_every=15.
+
+    Args:
+        heatmaps (List[np.ndarray | torch.Tensor]): List of 2D heatmaps.
+        vmin (float, optional): Min Value. Defaults to -1e-6.
+        vmax (float, optional): Max Value. Defaults to None.
+        cmap (str, optional): Colormap. Defaults to "Spectral".
+        a (float, optional): Alpha value. Defaults to 0.5.
+        axes (List[plt.Axes], optional): Axes to plot on. Defaults to None.
+        contours_every (int, optional): If not none, will draw contours. Defaults to None.
+        contour_style (str, optional): Style of the contours. Defaults to "solid".
+        colorbar (bool, optional): Whether to show colorbar. Defaults to False.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+    artists = []
+
+    for i in range(len(axes)):
+        a_ = a if isinstance(a, float) else a[i]
+
+        if isinstance(heatmaps[i], torch.Tensor):
+            heatmaps[i] = heatmaps[i].detach().cpu().numpy()
+
+        alpha = a_
+        # Plot the heatmap
+        art = axes[i].imshow(
+            heatmaps[i],
+            alpha=alpha,
+            vmin=vmin,
+            vmax=vmax,
+            cmap=cmap,
+        )
+        if colorbar:
+            cmax = vmax or np.percentile(heatmaps[i], 99)
+            art.set_clim(vmin, cmax)
+            cbar = plt.colorbar(art, ax=axes[i])
+            artists.append(cbar)
+
+        artists.append(art)
+
+        if contours_every is not None:
+            # Add contour lines to the heatmap
+            contour_data = np.arange(vmin, vmax + contours_every, contours_every)
+
+            # Get the colormap colors for contour lines
+            contour_colors = [
+                plt.colormaps.get_cmap(cmap)(plt.Normalize(vmin=vmin, vmax=vmax)(level))
+                for level in contour_data
+            ]
+            contours = axes[i].contour(
+                heatmaps[i],
+                levels=contour_data,
+                linewidths=2,
+                colors=contour_colors,
+                linestyles=contour_style,
+            )
+
+            contours.set_clim(vmin, vmax)
+
+            fmt = {
+                level: f"{label}°"
+                for level, label in zip(contour_data, contour_data.astype(int).astype(str))
+            }
+            t = axes[i].clabel(contours, inline=True, fmt=fmt, fontsize=16, colors="white")
+
+            for label in t:
+                label.set_path_effects(
+                    [
+                        path_effects.Stroke(linewidth=1, foreground="k"),
+                        path_effects.Normal(),
+                    ]
+                )
+            artists.append(contours)
+
+    return artists
+
+
+def plot_horizon_lines(
+    cameras, gravities, line_colors="orange", lw=2, styles="solid", alpha=1.0, ax=None
+):
+    """Plot horizon lines on the perspective field.
+
+    Args:
+        cameras (List[Camera]): List of cameras.
+        gravities (List[Gravity]): Gravities.
+        line_colors (str, optional): Line Colors. Defaults to "orange".
+        lw (int, optional): Line width. Defaults to 2.
+        styles (str, optional): Line styles. Defaults to "solid".
+        alpha (float, optional): Alphas. Defaults to 1.0.
+        ax (List[plt.Axes], optional): Axes to draw horizon line on. Defaults to None.
+    """
+    if not isinstance(line_colors, list):
+        line_colors = [line_colors] * len(cameras)
+
+    if not isinstance(styles, list):
+        styles = [styles] * len(cameras)
+
+    fig = plt.gcf()
+    ax = fig.gca() if ax is None else ax
+
+    if isinstance(ax, plt.Axes):
+        ax = [ax] * len(cameras)
+
+    assert len(ax) == len(cameras), f"{len(ax)}, {len(cameras)}"
+
+    for i in range(len(cameras)):
+        _, lat = get_perspective_field(cameras[i], gravities[i])
+        # horizon line is zero level of the latitude field
+        lat = lat[0, 0].cpu().numpy()
+        contours = ax[i].contour(lat, levels=[0], linewidths=lw, colors=line_colors[i])
+        for contour_line in contours.collections:
+            contour_line.set_linestyle(styles[i])
+
+
+def plot_vector_fields(
+    vector_fields,
+    cmap="lime",
+    subsample=15,
+    scale=None,
+    lw=None,
+    alphas=0.8,
+    axes=None,
+):
+    """Plot vector fields.
+
+    Args:
+        vector_fields (List[torch.Tensor]): List of vector fields of shape (2, H, W).
+        cmap (str, optional): Color of the vectors. Defaults to "lime".
+        subsample (int, optional): Subsample the vector field. Defaults to 15.
+        scale (float, optional): Scale of the vectors. Defaults to None.
+        lw (float, optional): Line width of the vectors. Defaults to None.
+        alphas (float | np.ndarray, optional): Alpha per vector or global. Defaults to 0.8.
+        axes (List[plt.Axes], optional): List of axes to draw on. Defaults to None.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    vector_fields = [v.cpu().numpy() if isinstance(v, torch.Tensor) else v for v in vector_fields]
+
+    artists = []
+
+    H, W = vector_fields[0].shape[-2:]
+    if scale is None:
+        scale = subsample / min(H, W)
+
+    if lw is None:
+        lw = 0.1 / subsample
+
+    if alphas is None:
+        alphas = np.ones_like(vector_fields[0][0])
+        alphas = np.stack([alphas] * len(vector_fields), 0)
+    elif isinstance(alphas, float):
+        alphas = np.ones_like(vector_fields[0][0]) * alphas
+        alphas = np.stack([alphas] * len(vector_fields), 0)
+    else:
+        alphas = np.array(alphas)
+
+    subsample = min(W, H) // subsample
+    offset_x = ((W % subsample) + subsample) // 2
+
+    samples_x = np.arange(offset_x, W, subsample)
+    samples_y = np.arange(int(subsample * 0.9), H, subsample)
+
+    x_grid, y_grid = np.meshgrid(samples_x, samples_y)
+
+    for i in range(len(axes)):
+        # vector field of shape (2, H, W) with vectors of norm == 1
+        vector_field = vector_fields[i]
+
+        a = alphas[i][samples_y][:, samples_x]
+        x, y = vector_field[:, samples_y][:, :, samples_x]
+
+        c = cmap
+        if not isinstance(cmap, str):
+            c = cmap[i][samples_y][:, samples_x].reshape(-1, 3)
+
+        s = scale * min(H, W)
+        arrows = axes[i].quiver(
+            x_grid,
+            y_grid,
+            x,
+            y,
+            scale=s,
+            scale_units="width" if H > W else "height",
+            units="width" if H > W else "height",
+            alpha=a,
+            color=c,
+            angles="xy",
+            antialiased=True,
+            width=lw,
+            headaxislength=3.5,
+            zorder=5,
+        )
+
+        artists.append(arrows)
+
+    return artists
+
+
+def plot_latitudes(
+    latitude,
+    is_radians=True,
+    vmin=-90,
+    vmax=90,
+    cmap="seismic",
+    contours_every=15,
+    alpha=0.4,
+    axes=None,
+    **kwargs,
+):
+    """Plot latitudes.
+
+    Args:
+        latitude (List[torch.Tensor]): List of latitudes.
+        is_radians (bool, optional): Whether the latitudes are in radians. Defaults to True.
+        vmin (int, optional): Min value to clip to. Defaults to -90.
+        vmax (int, optional): Max value to clip to. Defaults to 90.
+        cmap (str, optional): Colormap. Defaults to "seismic".
+        contours_every (int, optional): Contours every. Defaults to 15.
+        alpha (float, optional): Alpha value. Defaults to 0.4.
+        axes (List[plt.Axes], optional): Axes to plot on. Defaults to None.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    assert len(axes) == len(latitude), f"{len(axes)}, {len(latitude)}"
+    lat = [rad2deg(lat) for lat in latitude] if is_radians else latitude
+    return plot_heatmaps(
+        lat,
+        vmin=vmin,
+        vmax=vmax,
+        cmap=cmap,
+        a=alpha,
+        axes=axes,
+        contours_every=contours_every,
+        **kwargs,
+    )
+
+
+def plot_confidences(
+    confidence,
+    as_log=True,
+    vmin=-4,
+    vmax=0,
+    cmap="turbo",
+    alpha=0.4,
+    axes=None,
+    **kwargs,
+):
+    """Plot confidences.
+
+    Args:
+        confidence (List[torch.Tensor]): Confidence maps.
+        as_log (bool, optional): Whether to plot in log scale. Defaults to True.
+        vmin (int, optional): Min value to clip to. Defaults to -4.
+        vmax (int, optional): Max value to clip to. Defaults to 0.
+        cmap (str, optional): Colormap. Defaults to "turbo".
+        alpha (float, optional): Alpha value. Defaults to 0.4.
+        axes (List[plt.Axes], optional): Axes to plot on. Defaults to None.
+
+    Returns:
+        List[plt.Artist]: List of artists.
+    """
+    if axes is None:
+        axes = plt.gcf().axes
+
+    confidence = [c.cpu() if isinstance(c, torch.Tensor) else torch.tensor(c) for c in confidence]
+
+    assert len(axes) == len(confidence), f"{len(axes)}, {len(confidence)}"
+
+    if as_log:
+        confidence = [torch.log10(c.clip(1e-5)).clip(vmin, vmax) for c in confidence]
+
+    # normalize to [0, 1]
+    confidence = [(c - c.min()) / (c.max() - c.min()) for c in confidence]
+    return plot_heatmaps(confidence, vmin=0, vmax=1, cmap=cmap, a=alpha, axes=axes, **kwargs)
+
+
+def save_plot(path, **kw):
+    """Save the current figure without any white margin."""
+    plt.savefig(path, bbox_inches="tight", pad_inches=0, **kw)

src/datasets/utils.py

@@ -0,0 +1,162 @@
+import copy
+import random
+from xtuner.dataset.utils import get_bos_eos_token_ids
+from xtuner.utils import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX, IMAGE_TOKEN_INDEX
+import json
+
+INPUT_IMAGE_TOKEN_INDEX = IMAGE_TOKEN_INDEX
+OUTPUT_IMAGE_TOKEN_INDEX = -300
+QUERY_TOKEN_INDEX = -400
+QUERY_TOKEN = '<query>'
+
+def crop2square(pil_img):
+    width, height = pil_img.width, pil_img.height
+
+    if width > height:
+        y0, y1 = 0, height
+        x0 = random.randint(0, width - height)
+        x1 = x0 + height 
+    else:
+        x0, x1 = 0, width
+        y0 = random.randint(0, height - width)
+        y1 = y0 + width
+
+    return pil_img.crop(box=(x0, y0, x1, y1))
+
+def load_jsonl(json_file):
+    with open(json_file) as f:
+        lines = f.readlines()
+    data = []
+    for line in lines:
+        data.append(json.loads(line))
+    return data
+
+
+def encode_fn(example,
+              tokenizer,
+              max_length=None,
+              image_length=1,
+              query_length=1,
+              input_ids_with_output=True,
+              with_image_token=False,
+              prompt_template=None,
+              truncation='right'):
+    """Only support the following three scenarios:
+
+    1. Incremental pretraining dataset.
+        example['conversation'] = [
+                {
+                    'input': '',
+                    'output': '### Human: Can you write xxx'
+                }
+            ]
+
+    2. Single-turn conversation dataset.
+        example['conversation'] = [
+                {
+                    'input': 'Give three tips for staying healthy.',
+                    'output': '1.Eat a balanced diet xxx'
+                }
+            ]
+
+    3. Multi-turn conversation dataset.
+        example['conversation'] = [
+                {
+                    'input': 'Give three tips for staying healthy.',
+                    'output': '1.Eat a balanced diet xxx'
+                },
+                {
+                    'input': 'Please expand on the second point.',
+                    'output': 'Here is an expanded explanation of the xxx'
+                }
+            ]
+    """
+    
+    bos_token_id, eos_token_id = get_bos_eos_token_ids(tokenizer)
+    is_multi_turn_conversation = len(example['conversation']) > 1
+    if is_multi_turn_conversation:
+        assert input_ids_with_output
+
+    input_ids, labels = [], []
+    next_needs_bos_token = True
+    for single_turn_conversation in example['conversation']:
+        input = single_turn_conversation['input']
+        if DEFAULT_IMAGE_TOKEN in input and with_image_token:
+            chunk_encode = [
+                tokenizer.encode(chunk, add_special_tokens=False)
+                for chunk in input.split(DEFAULT_IMAGE_TOKEN)
+            ]
+            assert len(chunk_encode) == 2
+            input_encode = []
+            for idx, cur_chunk_encode in enumerate(chunk_encode):
+                input_encode.extend(cur_chunk_encode)
+                if idx != len(chunk_encode) - 1:
+                    input_encode += [INPUT_IMAGE_TOKEN_INDEX] * image_length
+        else:
+            input_encode = tokenizer.encode(input, add_special_tokens=False)
+        if next_needs_bos_token:
+            input_ids += bos_token_id
+            labels += [IGNORE_INDEX] * len(bos_token_id)
+        input_ids += input_encode
+        labels += [IGNORE_INDEX] * len(input_encode)
+        if input_ids_with_output and 'output' in single_turn_conversation:
+            # Add output
+            output_with_loss = single_turn_conversation.get(
+                'output_with_loss', True)
+            output = single_turn_conversation['output']
+            if DEFAULT_IMAGE_TOKEN in output and with_image_token:
+                chunk_encode = [
+                    tokenizer.encode(chunk, add_special_tokens=False)
+                    for chunk in output.split(DEFAULT_IMAGE_TOKEN)
+                ]
+                assert len(chunk_encode) == 2
+                output_encode = []
+                for idx, cur_chunk_encode in enumerate(chunk_encode):
+                    output_encode.extend(cur_chunk_encode)
+                    if idx != len(chunk_encode) - 1:
+                        output_encode += [OUTPUT_IMAGE_TOKEN_INDEX] * image_length
+            elif QUERY_TOKEN in output:
+                chunk_encode = [
+                    tokenizer.encode(chunk, add_special_tokens=False)
+                    for chunk in output.split(QUERY_TOKEN)
+                ]
+                assert len(chunk_encode) == 2
+                output_encode = []
+                for idx, cur_chunk_encode in enumerate(chunk_encode):
+                    output_encode.extend(cur_chunk_encode)
+                    if idx != len(chunk_encode) - 1:
+                        output_encode += [QUERY_TOKEN_INDEX] * query_length
+            else:
+                output_encode = tokenizer.encode(output, add_special_tokens=False)
+            input_ids += output_encode
+            if output_with_loss:
+                labels += copy.deepcopy(output_encode)
+            else:
+                labels += [IGNORE_INDEX] * len(output_encode)
+            # Add EOS_TOKEN (with loss)
+            if single_turn_conversation.get('need_eos_token', True):
+                next_needs_bos_token = True
+                input_ids += eos_token_id
+                if output_with_loss:
+                    labels += copy.deepcopy(eos_token_id)
+                else:
+                    labels += [IGNORE_INDEX] * len(eos_token_id)
+            else:
+                next_needs_bos_token = False
+            # Add SEP (without loss)
+            sep = single_turn_conversation.get('sep', '')
+            if sep != '':
+                sep_encode = tokenizer.encode(sep, add_special_tokens=False)
+                input_ids += sep_encode
+                labels += [IGNORE_INDEX] * len(sep_encode)
+
+    if max_length is not None and len(input_ids) > max_length:
+        if truncation == 'right':
+            input_ids = input_ids[:max_length]
+            labels = labels[:max_length]
+        elif truncation == 'left':
+            input_ids = input_ids[-max_length:]
+            labels = labels[-max_length:]
+        else:
+            assert truncation is None
+    return {'input_ids': input_ids, 'labels': labels}

src/models/connector/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_connector import ConnectorConfig
+from .modeling_connector import ConnectorEncoder

src/models/connector/configuration_connector.py

@@ -0,0 +1,27 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class ConnectorConfig(PretrainedConfig):
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        hidden_act="gelu_pytorch_tanh",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act

src/models/connector/modeling_connector.py

@@ -0,0 +1,507 @@
+# coding=utf-8
+# Copyright 2024 Google AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Connector model."""
+
+import math
+import warnings
+from typing import Any, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn.init import _calculate_fan_in_and_fan_out
+
+from transformers.activations import ACT2FN
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    ModelOutput,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+    torch_int,
+)
+from .configuration_connector import ConnectorConfig
+
+
+if is_flash_attn_2_available():
+    from transformers.modeling_flash_attention_utils import _flash_attention_forward
+
+
+logger = logging.get_logger(__name__)
+
+
+def init_weights(module):
+    """Initialize the weights"""
+    if isinstance(module, nn.Embedding):
+        default_flax_embed_init(module.weight)
+    elif isinstance(module, ConnectorAttention):
+        nn.init.xavier_uniform_(module.q_proj.weight)
+        nn.init.xavier_uniform_(module.k_proj.weight)
+        nn.init.xavier_uniform_(module.v_proj.weight)
+        nn.init.xavier_uniform_(module.out_proj.weight)
+        nn.init.zeros_(module.q_proj.bias)
+        nn.init.zeros_(module.k_proj.bias)
+        nn.init.zeros_(module.v_proj.bias)
+        nn.init.zeros_(module.out_proj.bias)
+    elif isinstance(module, ConnectorMLP):
+        nn.init.xavier_uniform_(module.fc1.weight)
+        nn.init.xavier_uniform_(module.fc2.weight)
+        nn.init.normal_(module.fc1.bias, std=1e-6)
+        nn.init.normal_(module.fc2.bias, std=1e-6)
+    elif isinstance(module, (nn.Linear, nn.Conv2d)):
+        lecun_normal_(module.weight)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+
+def _trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    # Values are generated by using a truncated uniform distribution and
+    # then using the inverse CDF for the normal distribution.
+    # Get upper and lower cdf values
+    l = norm_cdf((a - mean) / std)
+    u = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [l, u], then translate to
+    # [2l-1, 2u-1].
+    tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated
+    # standard normal
+    tensor.erfinv_()
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.0))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    tensor.clamp_(min=a, max=b)
+
+
+def trunc_normal_tf_(
+    tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0
+) -> torch.Tensor:
+    """Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \\leq \text{mean} \\leq b`.
+
+    NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
+    bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
+    and the result is subsequently scaled and shifted by the mean and std args.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    """
+    with torch.no_grad():
+        _trunc_normal_(tensor, 0, 1.0, a, b)
+        tensor.mul_(std).add_(mean)
+
+
+def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    if mode == "fan_in":
+        denom = fan_in
+    elif mode == "fan_out":
+        denom = fan_out
+    elif mode == "fan_avg":
+        denom = (fan_in + fan_out) / 2
+
+    variance = scale / denom
+
+    if distribution == "truncated_normal":
+        # constant is stddev of standard normal truncated to (-2, 2)
+        trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
+    elif distribution == "normal":
+        with torch.no_grad():
+            tensor.normal_(std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        with torch.no_grad():
+            tensor.uniform_(-bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")
+
+
+def default_flax_embed_init(tensor):
+    variance_scaling_(tensor, mode="fan_in", distribution="normal")
+
+
+class ConnectorAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Input shape: Batch x Time x Channel"""
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        k_v_seq_len = key_states.shape[-2]
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
+
+        if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class ConnectorFlashAttention2(ConnectorAttention):
+    """
+    ConnectorAttention flash attention module. This module inherits from `ConnectorAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    is_causal = False
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    # Adapted from transformers.models.llama.modeling_llama.LlamaFlashAttention2.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        output_attentions = False
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32.
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = _flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=dropout_rate,
+            is_causal=self.is_causal,
+            use_top_left_mask=self._flash_attn_uses_top_left_mask,
+        )
+
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim).contiguous()
+        attn_output = self.out_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights
+
+
+class ConnectorSdpaAttention(ConnectorAttention):
+    """
+    Connector attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `ConnectorAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    is_causal = False
+
+    # Adapted from ConnectorAttention.forward and transformers.models.llama.modeling_llama.LlamaSdpaAttention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "ConnectorModel is using ConnectorSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+            )
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+        # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+        is_causal = True if self.is_causal and q_len > 1 else False
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None
+
+
+CONNECTOR_ATTENTION_CLASSES = {
+    "eager": ConnectorAttention,
+    "flash_attention_2": ConnectorFlashAttention2,
+    "sdpa": ConnectorSdpaAttention,
+}
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Connector
+class ConnectorMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class ConnectorEncoderLayer(nn.Module):
+    def __init__(self, config: ConnectorConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = CONNECTOR_ATTENTION_CLASSES[config._attn_implementation](config=config)
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = ConnectorMLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+    # Ignore copy
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            attention_mask (`torch.FloatTensor`):
+                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->Connector
+class ConnectorEncoder(nn.Module):
+    def __init__(self, config: ConnectorConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([ConnectorEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+        self.apply(init_weights)
+
+    def forward(self, inputs_embeds):
+        hidden_states = inputs_embeds
+        for encoder_layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    encoder_layer.__call__,
+                    hidden_states,
+                    None,
+                    False,
+                    use_reentrant=False
+                )
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                    None,
+                    output_attentions=False,
+                )
+
+            hidden_states = layer_outputs[0]
+
+        return hidden_states

src/models/connector/modeling_qwen2.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+from transformers import Qwen2PreTrainedModel, Qwen2Config
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm, Qwen2DecoderLayer
+
+
+class Qwen2Connector(Qwen2PreTrainedModel):
+    def __init__(self, config: Qwen2Config):
+        super().__init__(config)
+        self.layers = nn.ModuleList(
+            [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+
+        for layer in self.layers:
+            layer.self_attn.is_causal = False
+
+        self._attn_implementation = config._attn_implementation
+        assert self._attn_implementation == 'flash_attention_2'
+        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(self, inputs_embeds):
+        position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)
+        position_ids = position_ids.expand(inputs_embeds.shape[0], -1)
+        hidden_states = inputs_embeds
+
+        for encoder_layer in self.layers:
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    encoder_layer.__call__,
+                    hidden_states,
+                    None,
+                    position_ids,
+                    use_reentrant=False
+                )
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                    attention_mask=None,
+                    position_ids=position_ids,
+                )
+
+            hidden_states = layer_outputs[0]
+
+        hidden_states = self.norm(hidden_states)
+
+        return hidden_states

src/models/puffin/model.py

@@ -0,0 +1,790 @@
+import random
+import torch
+import math
+from tqdm import tqdm
+from einops import rearrange
+from copy import deepcopy
+from six.moves import zip
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd.function import Function
+from torch.nn.utils.rnn import pad_sequence
+from mmengine.logging import print_log
+from mmengine.model import BaseModel
+from xtuner.utils import IGNORE_INDEX
+from xtuner.registry import BUILDER
+from xtuner.model.utils import guess_load_checkpoint
+from xtuner.dataset.map_fns.template_map_fn import template_map_fn
+from transformers.cache_utils import DynamicCache
+from diffusers.training_utils import compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
+
+from src.models.connector import ConnectorConfig, ConnectorEncoder
+from src.models.stable_diffusion3.pipeline_stable_diffusion_3_dynamic import StableDiffusion3Pipeline
+from src.datasets.utils import encode_fn, QUERY_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, INPUT_IMAGE_TOKEN_INDEX
+
+class _ScaleGradient(Function):
+    @staticmethod
+    def forward(ctx, input, scale):
+        ctx.scale = scale
+        return input
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output * ctx.scale, None
+
+def build_mlp(hidden_size, projector_dim, z_dim):
+    return nn.Sequential(
+        nn.Linear(hidden_size, projector_dim),
+        nn.SiLU(),
+        nn.Linear(projector_dim, z_dim),)
+
+def pad_an_image_tensor(image, pad_value=0):
+    h, w = image.shape[-2:]
+    if h > w:
+        pad_left = (h - w) // 2
+        pad_right = h - w - pad_left
+        p2d = (pad_left, pad_right, 0, 0)
+    else:
+        pad_top = (h - w) // 2
+        pad_bottom = h - w - pad_top
+        p2d = (0, 0, pad_top, pad_bottom)
+
+    image = F.pad(image, p2d, "constant", pad_value)
+
+    return image
+
+class Qwen2p5RadioStableDiffusion3HFDynamic(BaseModel):
+    def __init__(self,
+                 llm,
+                 tokenizer,
+                 prompt_template,
+                 visual_encoder,
+                 vae,
+                 transformer,
+                 train_scheduler,
+                 test_scheduler,
+                 connector_1,
+                 connector_2,
+                 num_queries=64,
+                 freeze_transformer=True,
+                 max_length=256,
+                 freeze_visual_encoder=True,
+                 freeze_llm=True,
+                 visual_encoder_grad_scale=0.1,
+                 fold_size=2,
+                 unconditional=0.1,
+                 unconditional_cross_view=0.1,
+                 pretrained_pth=None,
+                 use_activation_checkpointing=False,
+                 *args, **kwargs):
+        super().__init__()
+        
+        # basic settings
+        self.max_length = max_length
+        self.fold_size = fold_size
+        self.prompt_template = prompt_template
+        self.unconditional = unconditional
+        self.unconditional_cross_view = unconditional_cross_view
+        
+        # networks building
+        # understanding branch
+        self.visual_encoder = BUILDER.build(visual_encoder)
+        self.llm = BUILDER.build(llm)
+        self.tokenizer = BUILDER.build(tokenizer)
+        self.projector = build_mlp(hidden_size=self.visual_encoder.model.embed_dim*fold_size**2,
+                                   projector_dim=self.llm.config.hidden_size,
+                                   z_dim=self.llm.config.hidden_size)
+        self.image_token_id = self.tokenizer.convert_tokens_to_ids(prompt_template['IMG_CONTEXT_TOKEN'])
+        
+        # generation branch
+        self.vae = BUILDER.build(vae)
+        self.vae.requires_grad_(False)
+        self.transformer = BUILDER.build(transformer)
+        self.num_queries = num_queries
+        self.connector_1 = ConnectorEncoder(ConnectorConfig(**connector_1))
+        self.connector_2 = ConnectorEncoder(ConnectorConfig(**connector_2))
+
+        self.llm2connector_1 = nn.Linear(self.llm.config.hidden_size, self.connector_1.config.hidden_size)
+        self.llm2connector_2 = nn.Linear(self.llm.config.hidden_size, self.connector_2.config.hidden_size)
+        self.projector_1 = nn.Linear(self.connector_1.config.hidden_size, self.transformer.config.pooled_projection_dim)
+        self.projector_2 = nn.Linear(self.connector_2.config.hidden_size, self.transformer.config.joint_attention_dim)
+        nn.init.zeros_(self.projector_1.weight)
+        nn.init.zeros_(self.projector_2.weight)
+        nn.init.zeros_(self.projector_1.bias)
+        nn.init.zeros_(self.projector_2.bias)
+
+        self.meta_queries = nn.Parameter(
+            torch.zeros(num_queries, self.llm.config.hidden_size))
+        nn.init.normal_(self.meta_queries, std=1 / math.sqrt(self.llm.config.hidden_size))
+        
+        # networks and training initialization
+        if freeze_visual_encoder:
+            self.visual_encoder.requires_grad_(False)
+        self.freeze_visual_encoder = freeze_visual_encoder
+        if freeze_llm:
+            self.llm.requires_grad_(False)
+        self.freeze_llm = freeze_llm
+        if freeze_transformer:
+            self.transformer.requires_grad_(False)
+        self.freeze_transformer = freeze_transformer
+        
+        self.visual_encoder_grad_scale = visual_encoder_grad_scale
+        self.train_scheduler = BUILDER.build(train_scheduler)
+        self.test_scheduler = BUILDER.build(test_scheduler)
+
+        self.use_activation_checkpointing = use_activation_checkpointing
+        if use_activation_checkpointing:
+            self.llm.enable_input_require_grads()
+            self.gradient_checkpointing_enable()
+
+        if pretrained_pth is not None:
+            pretrained_state_dict = guess_load_checkpoint(pretrained_pth)
+            info = self.load_state_dict(pretrained_state_dict, strict=False)
+            print_log(f'Load pretrained weight from {pretrained_pth}')
+            
+    @property
+    def device(self):
+        return self.llm.device
+
+    @property
+    def dtype(self):
+        return self.llm.dtype
+
+    def gradient_checkpointing_enable(self):
+        self.activation_checkpointing_enable()
+
+    def activation_checkpointing_enable(self):
+        self.llm.gradient_checkpointing_enable()
+        self.transformer.enable_gradient_checkpointing()
+        self.connector_1.gradient_checkpointing = True
+        self.connector_2.gradient_checkpointing = True
+        
+    def gradient_checkpointing_disable(self):
+        self.activation_checkpointing_disable()
+
+    def activation_checkpointing_disable(self):
+        self.llm.gradient_checkpointing_disable()
+        self.transformer.disable_gradient_checkpointing()
+        self.connector_1.gradient_checkpointing = False
+        self.connector_2.gradient_checkpointing = False
+        
+    def forward(self, data, data_samples=None, mode='loss'):
+        if mode == 'loss':
+            return self.compute_loss(data_dict=data)
+        else:
+            raise NotImplementedError
+
+    def extract_visual_features(self, pixel_values):
+        pixel_values = (pixel_values + 1.0) / 2     # [0, 1]
+        height, width = pixel_values.shape[-2:]
+        summary, features = self.visual_encoder(pixel_values)
+        patch_size = int((height * width // features.shape[1]) ** 0.5)
+        height, width = height // (patch_size * self.fold_size), width // (patch_size * self.fold_size)
+        features = rearrange(features, 'b (h p w q) d -> b (h w) (p q d)',
+                             h=height, w=width, p=self.fold_size, q=self.fold_size)
+        
+        return features
+
+    def llm2dit(self, x):
+        x_1 = self.connector_1(self.llm2connector_1(x))
+        x_1 = self.projector_1(x_1.mean(1))
+        x_2 = self.connector_2(self.llm2connector_2(x))
+        x_2 = self.projector_2(x_2)
+        
+        return x_1, x_2
+    
+    
+    @torch.no_grad()
+    def prepare_gen_prompts(self, texts, data_type='text2image', num_refs=None, ref_lens=None, gen_type='GENERATION_CROSS'):
+        if data_type == 'text2image':
+            prompts = [self.prompt_template['GENERATION'].format(input=text) for text in texts]
+            prompts = [self.prompt_template['INSTRUCTION'].format(input=text) for text in prompts]
+
+        elif data_type == 'image2image':
+            assert num_refs is not None and ref_lens is not None, "num_refs and ref_lens are required for image2image"
+            prompts = []
+            cnt = 0
+            for text, num_ref in zip(texts, num_refs):
+                image_tokens = ''
+                for _ in range(num_ref):
+                    image_tokens += (
+                        self.prompt_template['IMG_START_TOKEN'] +
+                        self.prompt_template['IMG_CONTEXT_TOKEN'] * ref_lens[cnt] +
+                        self.prompt_template['IMG_END_TOKEN']
+                    )
+                    cnt += 1
+
+                text = self.prompt_template[gen_type].format(input=text)
+                prompt = self.prompt_template['INSTRUCTION'].format(input=f'{image_tokens}\n{text}')
+                prompts.append(prompt)
+        else:
+            raise ValueError(f"Unsupported data_type: {data_type}")
+
+        return self.tokenizer(
+            prompts, add_special_tokens=True, return_tensors='pt', padding=True, padding_side='left').to(self.device)
+
+
+    @torch.no_grad()
+    def prepare_und_prompts(self, conversations, data_type='image2text', image_lengths=None, input_ids_with_output=True):
+        input_ids, labels, input_lengths = [], [], []
+
+        if data_type == 'image2text':
+            assert image_lengths is not None, "`image_lengths` must be provided for image2text"
+            if isinstance(image_lengths, int):
+                image_lengths = [image_lengths] * len(conversations)
+        elif data_type == 'text2text':
+            image_lengths = [None] * len(conversations)
+        else:
+            raise ValueError(f"Unsupported data_type: {data_type}")
+
+        for conv, image_len in zip(conversations, image_lengths):
+            data_dict = template_map_fn(example=dict(conversation=deepcopy(conv)), template=self.prompt_template)
+            data_dict.update(encode_fn(data_dict,
+                                      tokenizer=self.tokenizer,
+                                      max_length=None,
+                                      input_ids_with_output=input_ids_with_output,
+                                      with_image_token=(data_type == 'image2text'),
+                                      image_length=image_len,
+                                      prompt_template=self.prompt_template))
+
+            input_ids.append(torch.tensor(data_dict['input_ids'], dtype=torch.long, device=self.device))
+            labels.append(torch.tensor(data_dict['labels'], dtype=torch.long, device=self.device))
+            input_lengths.append(len(data_dict['input_ids']))
+
+        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=0, padding_side='left')
+        labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX, padding_side='left')
+
+        attention_mask = torch.zeros_like(input_ids).bool()
+        for i in range(len(input_ids)):
+            attention_mask[i, -input_lengths[i]:] = True
+
+        position_ids = torch.cumsum(attention_mask, dim=1) - 1
+        position_ids[position_ids < 0] = 0
+
+        return dict(input_ids=input_ids, attention_mask=attention_mask, labels=labels, position_ids=position_ids)
+
+    def train(self, mode=True):
+        super().train(mode=mode)
+        self.vae.train(mode=False)
+        if not mode:
+            self.gradient_checkpointing_disable()
+
+        return self
+
+    @torch.no_grad()
+    def pixels_to_latents(self, x):
+        z = self.vae.encode(x).latent_dist.sample()
+        z = (z - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        return z
+
+    @torch.no_grad()
+    def latents_to_pixels(self, z):
+        z = (z / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        x_rec = self.vae.decode(z).sample
+        return x_rec
+
+    def prepare_forward_input(self,
+                              query_embeds,
+                              input_ids=None,
+                              image_embeds=None,
+                              attention_mask=None,
+                              past_key_values=None,
+                              append_queries=True):
+        b, l, _ = query_embeds.shape
+        assert l > 0
+        attention_mask = attention_mask.to(device=self.device, dtype=torch.bool)
+        assert l == self.num_queries
+
+        if append_queries:
+            input_ids = torch.cat([
+                input_ids, input_ids.new_full(size=(b, l), fill_value=QUERY_TOKEN_INDEX)], dim=1)
+            attention_mask = torch.cat([attention_mask, attention_mask.new_ones(b, l)], dim=1)
+
+        position_ids = torch.cumsum(attention_mask, dim=1) - 1
+        position_ids[position_ids < 0] = 0
+
+        # prepare context
+        if past_key_values is not None:
+            inputs_embeds = query_embeds
+            position_ids = position_ids[..., -l:]
+        else:
+            inputs_embeds = torch.zeros(*input_ids.shape, self.llm.config.hidden_size,
+                                        device=self.device, dtype=self.dtype)
+            if image_embeds is not None:
+                inputs_embeds[input_ids == self.image_token_id] = \
+                    image_embeds.contiguous().view(-1, self.llm.config.hidden_size)
+
+            inputs_embeds[input_ids == QUERY_TOKEN_INDEX] = \
+                query_embeds.contiguous().view(-1, self.llm.config.hidden_size)
+
+            text_places = torch.logical_and(input_ids != self.image_token_id, input_ids != QUERY_TOKEN_INDEX)
+
+            inputs_embeds[text_places] = self.llm.get_input_embeddings()(input_ids[text_places])
+
+        inputs = dict(inputs_embeds=inputs_embeds,
+                      attention_mask=attention_mask,
+                      position_ids=position_ids,
+                      past_key_values=past_key_values)
+
+        return inputs
+
+    def get_sigmas(self, timesteps, n_dim=4):
+        sigmas = self.train_scheduler.sigmas.to(device=self.device, dtype=self.dtype)
+        schedule_timesteps = self.train_scheduler.timesteps.to(self.device)
+        timesteps = timesteps.to(self.device)
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < n_dim:
+            sigma = sigma.unsqueeze(-1)
+        return sigma
+
+    def diff_loss(self, model_input, pooled_prompt_embeds, prompt_embeds, cond_input=None):
+        noise = [torch.randn_like(x) for x in model_input]
+        bsz = len(model_input)
+
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme='none',
+            batch_size=bsz,
+            logit_mean=0.0,
+            logit_std=1.0,
+        )
+        indices = (u * self.train_scheduler.config.num_train_timesteps).long()
+        timesteps = self.train_scheduler.timesteps[indices].to(device=self.device)
+
+        # Add noise according to flow matching
+        sigmas = self.get_sigmas(timesteps, n_dim=model_input[0].ndim + 1)
+        noisy_model_input = [(1.0 - x) * y + x * z  for x, y, z in zip(sigmas, model_input, noise)]
+
+        # Predict the noise residual
+        model_pred = self.transformer(
+            hidden_states=noisy_model_input,
+            cond_hidden_states=cond_input,
+            encoder_hidden_states=prompt_embeds,
+            pooled_projections=pooled_prompt_embeds,
+            timestep=timesteps,
+            return_dict=False,
+        )[0]
+
+        weighting = compute_loss_weighting_for_sd3(weighting_scheme='none', sigmas=sigmas)
+
+        # flow matching loss
+        target = [x - y for x, y in zip(noise, model_input)]
+
+        loss = [(x.float() * (y.float() - z.float()) ** 2).mean() for x, y, z in zip(weighting, model_pred, target)]
+        loss = sum(loss) / len(loss)
+
+        return loss
+
+    '''text-to-image generation (single-view)'''
+    def text2image_loss(self, data_dict):
+        pixel_values = [p.to(dtype=self.dtype, device=self.device) for p in data_dict['pixel_values']]
+        image_latents = [self.pixels_to_latents(p[None])[0] for p in pixel_values]
+
+        b = len(image_latents)
+
+        texts = ['' if random.uniform(0, 1) < self.unconditional else text
+                 for text in data_dict['texts']]
+
+        text_inputs = self.prepare_gen_prompts(texts)
+        hidden_states = self.meta_queries[None].expand(b, self.num_queries, -1)
+
+        inputs = self.prepare_forward_input(query_embeds=hidden_states, **text_inputs)
+
+        max_length = self.max_length + self.num_queries
+        inputs_embeds = inputs['inputs_embeds'][:, -max_length:]
+        attention_mask = inputs['attention_mask'][:, -max_length:]
+        position_ids = inputs['position_ids'][:, -max_length:]
+
+        output = self.llm.model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            return_dict=True)
+
+        hidden_states = output.last_hidden_state[:, -self.num_queries:]
+        pooled_prompt_embeds, prompt_embeds = self.llm2dit(hidden_states)
+
+        loss_diff = self.diff_loss(model_input=image_latents,
+                                   pooled_prompt_embeds=pooled_prompt_embeds,
+                                   prompt_embeds=prompt_embeds)
+
+        return loss_diff
+    
+    '''text-to-image generation (single-view) with camera map'''
+    def cam2image_loss(self, data_dict):
+        pixel_values = [p.to(dtype=self.dtype, device=self.device) for p in data_dict['pixel_values']]
+        image_latents = [self.pixels_to_latents(p[None])[0] for p in pixel_values]
+        b = len(image_latents)
+        # camera map as condition for the diffusion model
+        cam_values = [[img.to(dtype=self.dtype, device=self.device) for img in ref_images]
+                            for ref_images in data_dict['cam_values']]
+        cam_latents = [[self.pixels_to_latents(img[None])[0] for img in ref_images]
+                            for ref_images in cam_values]
+
+        texts = ['' if random.uniform(0, 1) < self.unconditional else text
+                for text in data_dict['texts']]
+
+        text_inputs = self.prepare_gen_prompts(texts)
+        hidden_states = self.meta_queries[None].expand(b, self.num_queries, -1)
+
+        inputs = self.prepare_forward_input(query_embeds=hidden_states, **text_inputs)
+
+        max_length = self.max_length + self.num_queries
+        inputs_embeds = inputs['inputs_embeds'][:, -max_length:]
+        attention_mask = inputs['attention_mask'][:, -max_length:]
+        position_ids = inputs['position_ids'][:, -max_length:]
+
+        output = self.llm.model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            return_dict=True)
+
+        hidden_states = output.last_hidden_state[:, -self.num_queries:]
+        pooled_prompt_embeds, prompt_embeds = self.llm2dit(hidden_states)
+
+        loss_diff = self.diff_loss(model_input=image_latents,
+                                   pooled_prompt_embeds=pooled_prompt_embeds,
+                                   prompt_embeds=prompt_embeds,
+                                   cond_input=cam_latents)
+        
+        return loss_diff
+    
+    '''image-to-image (cross-view) generation'''
+    def image2image_loss(self, data_dict):
+        # condition for the diffusion model (concat the camera map and the initial view)
+        cam_values = [[img.to(dtype=self.dtype, device=self.device) for img in ref_images]
+                            for ref_images in data_dict['cam_values']]
+        cam_latents = [[self.pixels_to_latents(img[None])[0] for img in ref_images]
+                            for ref_images in cam_values]
+        pixel_values_init = [[img.to(dtype=self.dtype, device=self.device) for img in ref_images]
+                            for ref_images in data_dict['pixel_values_init']]
+        image_latents_init = [[self.pixels_to_latents(img[None])[0] for img in ref_images]
+                            for ref_images in pixel_values_init]
+        mix_latents = [cam + img for cam, img in zip(cam_latents, image_latents_init)]
+        
+        # condition embedding for querying the LLM (only initial view)
+        num_refs = [len(ref_images) for ref_images in pixel_values_init]
+        image_embeds = self.extract_visual_features(
+            torch.stack([pad_an_image_tensor(img) for ref_images in pixel_values_init for img in ref_images]))
+
+        image_embeds = self.projector(image_embeds)
+        ref_lens = [len(x) for x in image_embeds]
+        text_inputs = self.prepare_gen_prompts(data_dict['texts'], data_type='image2image', 
+                                                    num_refs=num_refs, ref_lens=ref_lens)
+        
+        # input for the diffusion model
+        pixel_values = [p.to(dtype=self.dtype, device=self.device) for p in data_dict['pixel_values']]
+        image_latents = [self.pixels_to_latents(p[None])[0] for p in pixel_values]
+
+        # querying the LLM
+        b = len(image_latents)
+        hidden_states = self.meta_queries[None].expand(b, self.num_queries, -1)
+        inputs = self.prepare_forward_input(query_embeds=hidden_states, image_embeds=image_embeds, **text_inputs)
+
+        max_length = self.max_length + max(num_refs) * max(ref_lens) + self.num_queries
+        inputs_embeds = inputs['inputs_embeds'][:, -max_length:]
+        attention_mask = inputs['attention_mask'][:, -max_length:]
+        position_ids = inputs['position_ids'][:, -max_length:]
+
+        output = self.llm.model(inputs_embeds=inputs_embeds,
+                          attention_mask=attention_mask,
+                          position_ids=position_ids,
+                          return_dict=True)
+        hidden_states = output.last_hidden_state[:, -self.num_queries:]
+        pooled_prompt_embeds, prompt_embeds = self.llm2dit(hidden_states)
+        loss_diff = self.diff_loss(model_input=image_latents,
+                                   pooled_prompt_embeds=pooled_prompt_embeds,
+                                   prompt_embeds=prompt_embeds,
+                                   cond_input=mix_latents)
+        
+        return loss_diff
+    
+    '''image-to-text(camera) understanding, mixed base, thinking, and instruction tuning'''
+    def image2text_loss(self, data_dict):
+        pixel_values = [pad_an_image_tensor(img) for img in data_dict['pixel_values']]
+        pixel_values = torch.stack(pixel_values).to(dtype=self.dtype, device=self.device)
+        image_embeds = self.extract_visual_features(pixel_values)
+
+        if not self.freeze_visual_encoder:
+            image_embeds = _ScaleGradient.apply(image_embeds, self.visual_encoder_grad_scale)
+
+        image_embeds = self.projector(image_embeds)
+        text_inputs = self.prepare_und_prompts(conversations=data_dict['conversations'],
+                                               data_type='image2text',
+                                               image_lengths=image_embeds.shape[1])
+
+        labels, input_ids, attention_mask, position_ids = \
+            text_inputs['labels'], text_inputs['input_ids'], text_inputs['attention_mask'], text_inputs['position_ids']
+
+
+        inputs_embeds = torch.zeros(*input_ids.shape, self.llm.config.hidden_size,
+                                    device=self.device, dtype=self.dtype)
+        inputs_embeds[input_ids == INPUT_IMAGE_TOKEN_INDEX] = image_embeds.flatten(0, 1)
+        inputs_embeds[input_ids != INPUT_IMAGE_TOKEN_INDEX] = \
+            self.llm.get_input_embeddings()(input_ids[input_ids != INPUT_IMAGE_TOKEN_INDEX])
+
+        max_length = self.max_length + image_embeds.shape[1]
+        inputs_embeds = inputs_embeds[:, -max_length:]
+        attention_mask = attention_mask[:, -max_length:]
+        position_ids = position_ids[:, -max_length:]
+        labels = labels[:, -max_length:]
+
+        output = self.llm.model(inputs_embeds=inputs_embeds,
+                                attention_mask=attention_mask,
+                                position_ids=position_ids,
+                                return_dict=True)
+
+        hidden_states = output.last_hidden_state[:, :-1]
+        labels = labels[:, 1:]
+        hidden_states = hidden_states[labels >= 0]
+        labels = labels[labels >= 0]
+
+        logits = self.llm.get_output_embeddings()(hidden_states)
+        loss = F.cross_entropy(input=logits, target=labels)
+
+        return loss
+    
+    '''text-to-text understanding, offering the enhanced caption for the generation'''
+    def text2text_loss(self, data_dict):
+        text_inputs = self.prepare_und_prompts(conversations=data_dict['conversations'], data_type='text2text')
+        labels, input_ids, attention_mask, position_ids = \
+            text_inputs['labels'], text_inputs['input_ids'], text_inputs['attention_mask'], text_inputs['position_ids']
+
+        inputs_embeds = self.llm.get_input_embeddings()(input_ids)
+        max_length = self.max_length
+        inputs_embeds = inputs_embeds[:, -max_length:]
+        attention_mask = attention_mask[:, -max_length:]
+        position_ids = position_ids[:, -max_length:]
+        labels = labels[:, -max_length:]
+
+        output = self.llm.model(inputs_embeds=inputs_embeds,
+                                attention_mask=attention_mask,
+                                position_ids=position_ids,
+                                return_dict=True)
+
+        hidden_states = output.last_hidden_state[:, :-1]
+        labels = labels[:, 1:]
+        hidden_states = hidden_states[labels >= 0]
+        labels = labels[labels >= 0]
+
+        logits = self.llm.get_output_embeddings()(hidden_states)
+        loss = F.cross_entropy(input=logits, target=labels)
+
+        return loss
+    
+    '''distribute different losses for each task'''
+    def compute_loss(self, data_dict):
+        loss_fn_map = {
+            'text2image': self.text2image_loss,
+            'cam2image': self.cam2image_loss,
+            'image2text': self.image2text_loss,
+            'text2text': self.text2text_loss,
+            'image2image': self.image2image_loss,
+            'image2text_cross_view': self.image2text_loss,
+        }
+
+        losses = {}
+        for data_type, batch_data in data_dict.items():
+            if data_type not in loss_fn_map:
+                raise ValueError(f"Unsupported data_type: {data_type}")
+            loss_fn = loss_fn_map[data_type]
+            loss = loss_fn(batch_data)
+            losses[f'loss_{data_type}'] = loss
+        return losses
+
+    @torch.no_grad()
+    def generate(self,
+                 prompt,
+                 cfg_prompt,
+                 cam_values=None,
+                 pixel_values_init=None,
+                 cfg_scale=4.5,
+                 num_steps=50,
+                 generator=None,
+                 height=512,
+                 width=512,
+                 max_new_tokens=512,
+                 reasoning=False,
+                 prompt_reasoning=None,
+                 progress_bar=True):
+        assert len(prompt) == len(cfg_prompt)
+        b = len(prompt)
+        output_reasoning = [''] * b
+        
+        if reasoning:
+            # enrich the prompt if required reasoning generation
+            assert prompt_reasoning is not None, \
+                "prompt_reasoning must be provided for reasoning generation"
+            if isinstance(prompt_reasoning, str):
+                prompt_reasoning = [prompt_reasoning]
+            if isinstance(prompt, str):
+                prompt = [prompt]
+
+            conversations = [[{'input': f"{p1} {p2}",}] 
+                                    for p1, p2 in zip(prompt_reasoning, prompt)]
+
+            text_inputs = self.prepare_und_prompts(
+                conversations=conversations, data_type="text2text", input_ids_with_output=False)
+            input_ids, attention_mask, position_ids = \
+                text_inputs['input_ids'], text_inputs['attention_mask'], text_inputs['position_ids']
+
+            inputs_embeds = self.llm.get_input_embeddings()(input_ids)
+            past_key_values = DynamicCache.from_legacy_cache()
+
+            output_ids = []
+            for _ in tqdm(range(max_new_tokens), disable=not progress_bar):
+                output = self.llm.model(
+                    inputs_embeds=inputs_embeds,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                    return_dict=True)
+                logits = self.llm.get_output_embeddings()(output.last_hidden_state[:, -1:])
+                input_ids = torch.argmax(logits, dim=-1)   # b 1
+                if len(output_ids) > 0:
+                    input_ids = torch.where(output_ids[-1] == self.tokenizer.eos_token_id,
+                                            output_ids[-1], input_ids)
+                output_ids.append(input_ids)
+
+                if (input_ids == self.tokenizer.eos_token_id).all():
+                    break
+
+                inputs_embeds =  self.llm.get_input_embeddings()(input_ids)
+                attention_mask = torch.cat([attention_mask, attention_mask.new_ones(b, 1)], dim=1)
+                position_ids = torch.max(position_ids, dim=1, keepdim=True).values + 1
+                past_key_values = output.past_key_values
+
+            output_ids = torch.cat(output_ids, dim=1)
+            output_reasoning = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+            prompt = [f"{p} {o}" for p, o in zip(prompt, output_reasoning)]
+        
+        if cam_values is not None:
+            # for the generation with the camera map
+            cam_values = [[img.to(dtype=self.dtype, device=self.device) for img in ref_images]
+                                for ref_images in cam_values]
+            cond_latents = [[self.pixels_to_latents(img[None])[0] for img in ref_images]
+                                for ref_images in cam_values]
+            text_inputs = self.prepare_gen_prompts(prompt + cfg_prompt)
+            if pixel_values_init is not None:
+                # for the generation with the camera map and initial view (cross-view generation)
+                num_refs = [len(ref_images) for ref_images in pixel_values_init]
+                pixel_values_init = [[img.to(dtype=self.dtype, device=self.device) for img in ref_images]
+                                    for ref_images in pixel_values_init]
+                image_embeds = self.extract_visual_features(
+                    torch.stack([pad_an_image_tensor(img) for ref_images in pixel_values_init for img in ref_images]))
+                image_embeds = self.projector(image_embeds)
+
+                ref_lens = [len(x) for x in image_embeds]
+                text_inputs = self.prepare_gen_prompts(prompt + cfg_prompt, data_type='image2image', num_refs=num_refs*2, ref_lens=ref_lens*2)
+                text_inputs.update(image_embeds=torch.cat([image_embeds]*2))
+                
+                cond_latents_init = [[self.pixels_to_latents(img[None])[0] for img in ref_imgs]
+                                for ref_imgs in pixel_values_init]
+                cond_latents = [cam + img for cam, img in zip(cond_latents, cond_latents_init)]
+            
+            cond_latents = cond_latents * 2
+        else:
+            # for the text2image generation
+            text_inputs = self.prepare_gen_prompts(prompt + cfg_prompt)
+            cond_latents = None
+
+        hidden_states = self.meta_queries[None].expand(2*b, self.num_queries, -1)
+        inputs = self.prepare_forward_input(query_embeds=hidden_states, **text_inputs)
+
+        output = self.llm.model(**inputs, return_dict=True)
+        hidden_states = output.last_hidden_state[:, -self.num_queries:]
+        pooled_prompt_embeds, prompt_embeds = self.llm2dit(hidden_states)
+
+        pipeline = StableDiffusion3Pipeline(
+            transformer=self.transformer,
+            scheduler=self.test_scheduler,
+            vae=self.vae,
+            text_encoder=None,
+            tokenizer=None,
+            text_encoder_2=None,
+            tokenizer_2=None,
+            text_encoder_3=None,
+            tokenizer_3=None,
+        )
+
+        pipeline.set_progress_bar_config(disable=not progress_bar)
+
+        samples = pipeline(
+            height=height,
+            width=width,
+            guidance_scale=cfg_scale,
+            num_inference_steps=num_steps,
+            prompt_embeds=prompt_embeds[:b],
+            pooled_prompt_embeds=pooled_prompt_embeds[:b],
+            negative_prompt_embeds=prompt_embeds[b:],
+            negative_pooled_prompt_embeds=pooled_prompt_embeds[b:],
+            generator=generator,
+            output_type='latent',
+            cond_latents=cond_latents
+        ).images.to(self.dtype)
+
+        return self.latents_to_pixels(samples), output_reasoning
+    
+    @torch.no_grad()
+    def understand(self, prompt, pixel_values, max_new_tokens=512, progress_bar=True):
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if isinstance(pixel_values, torch.Tensor):
+            pixel_values = [pixel_values]
+
+        bsz = len(prompt)
+        assert len(pixel_values) == bsz
+
+        pixel_values = [pad_an_image_tensor(img) for img in pixel_values]
+        pixel_values = torch.stack(pixel_values).to(dtype=self.dtype, device=self.device)
+        image_embeds = self.extract_visual_features(pixel_values)
+        image_embeds = self.projector(image_embeds)
+
+        conversations = [[{'input': f"{DEFAULT_IMAGE_TOKEN}\n{p}",}] for p in prompt]
+
+        text_inputs = self.prepare_und_prompts(conversations=conversations, image_lengths=image_embeds.shape[1], 
+                                                input_ids_with_output=False)
+
+        input_ids, attention_mask, position_ids = \
+            text_inputs['input_ids'], text_inputs['attention_mask'], text_inputs['position_ids']
+
+        inputs_embeds = torch.zeros(*input_ids.shape, self.llm.config.hidden_size,
+                                    device=self.device, dtype=self.dtype)
+        inputs_embeds[input_ids == INPUT_IMAGE_TOKEN_INDEX] = image_embeds.flatten(0, 1)
+        inputs_embeds[input_ids != INPUT_IMAGE_TOKEN_INDEX] = \
+            self.llm.get_input_embeddings()(input_ids[input_ids != INPUT_IMAGE_TOKEN_INDEX])
+
+        past_key_values = DynamicCache.from_legacy_cache()
+
+        output_ids = []
+
+        for _ in tqdm(range(max_new_tokens), disable=not progress_bar):
+            output = self.llm.model(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=True,
+                return_dict=True)
+            logits = self.llm.get_output_embeddings()(output.last_hidden_state[:, -1:])
+            input_ids = torch.argmax(logits, dim=-1)   # b 1
+            if len(output_ids) > 0:
+                input_ids = torch.where(output_ids[-1] == self.tokenizer.eos_token_id,
+                                        output_ids[-1], input_ids)
+            output_ids.append(input_ids)
+
+            if (input_ids == self.tokenizer.eos_token_id).all():
+                break
+
+            inputs_embeds =  self.llm.get_input_embeddings()(input_ids)
+            attention_mask = torch.cat([attention_mask, attention_mask.new_ones(bsz, 1)], dim=1)
+            position_ids = torch.max(position_ids, dim=1, keepdim=True).values + 1
+            past_key_values = output.past_key_values
+
+        output_ids = torch.cat(output_ids, dim=1)
+        output_text = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+
+        return output_text

src/models/radiov3/adaptor_base.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+from typing import NamedTuple, Optional
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class AdaptorInput(NamedTuple):
+    images: torch.Tensor
+    summary: torch.Tensor
+    features: torch.Tensor
+    feature_fmt: str
+    patch_size: int
+
+
+class RadioOutput(NamedTuple):
+    summary: torch.Tensor
+    features: torch.Tensor
+
+    def to(self, *args, **kwargs):
+        return RadioOutput(
+            self.summary.to(*args, **kwargs) if self.summary is not None else None,
+            self.features.to(*args, **kwargs) if self.features is not None else None,
+        )
+
+
+class AdaptorBase(nn.Module):
+    def forward(self, input: AdaptorInput) -> RadioOutput:
+        raise NotImplementedError("Subclasses must implement this!")

src/models/radiov3/adaptor_generic.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .adaptor_base import AdaptorBase, AdaptorInput, RadioOutput
+from .adaptor_mlp import create_mlp_from_state, create_mlp_from_config
+
+
+class GenericAdaptor(AdaptorBase):
+    def __init__(self, main_config: Namespace, adaptor_config, state, mlp_config=None):
+        super().__init__()
+
+        extra_args = dict()
+        ups = None
+        ups_rank = None
+        if adaptor_config is not None:
+            ups = adaptor_config.get('fd_upsample_factor', None)
+            ups_rank = adaptor_config.get('fd_upsample_rank', None)
+        elif mlp_config is not None:
+            ups = mlp_config["feature"].get('upsample_factor', None)
+            ups_rank = mlp_config["feature"].get('upsample_rank', None)
+        if ups is not None:
+            extra_args['upsample_factor'] = ups
+            extra_args['upsample_rank'] = ups_rank
+
+        if state is not None:
+            spectral_heads = getattr(main_config, 'spectral_heads', False)
+            self.head_mlp = create_mlp_from_state(main_config.mlp_version, state, 'summary.', spectral_weights=spectral_heads)
+            self.feat_mlp = create_mlp_from_state(main_config.mlp_version, state, 'feature.', spectral_weights=spectral_heads, **extra_args)
+        else:
+            assert mlp_config is not None, "Config must not be None if state is None"
+
+            self.head_mlp =  create_mlp_from_config(
+                main_config.mlp_version,
+                mlp_config["summary"]["input_dim"],
+                mlp_config["summary"]["hidden_dim"],
+                mlp_config["summary"]["output_dim"],
+                mlp_config["summary"]["num_inner"],
+            )
+            self.feat_mlp = create_mlp_from_config(
+                main_config.mlp_version,
+                mlp_config["feature"]["input_dim"],
+                mlp_config["feature"]["hidden_dim"],
+                mlp_config["feature"]["output_dim"],
+                mlp_config["feature"]["num_inner"],
+                **extra_args
+            )
+
+    def forward(self, input: AdaptorInput) -> RadioOutput:
+        # Convert input'd type to the type of the first parameter of the adaptor.
+        first_param = next(self.parameters())
+        summary = self.head_mlp(input.summary.to(dtype=first_param.dtype)).to(dtype=input.summary.dtype)
+        feat = self.feat_mlp(input.features.to(dtype=first_param.dtype), images=input.images, patch_size=input.patch_size).to(dtype=input.features.dtype)
+
+        if input.feature_fmt == 'NCHW':
+            feat = (feat.reshape(feat.shape[0], input.images.shape[-2] // input.patch_size * self.feat_mlp.upsample_factor, input.images.shape[-1] // input.patch_size * self.feat_mlp.upsample_factor, feat.shape[2])
+                        .permute(0, 3, 1, 2)
+            )
+
+        return RadioOutput(summary, feat)

src/models/radiov3/adaptor_mlp.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+import math
+from typing import Dict, Optional
+
+import torch
+from torch import nn
+
+from einops import rearrange
+from timm.models.vision_transformer import Block
+
+from .enable_spectral_reparam import disable_spectral_reparam, enable_spectral_reparam
+
+
+class MLP(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, output_size: int,
+                 num_inner: int = 0, device: torch.device = None, **kwargs):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size, device=device)
+        self.norm = nn.LayerNorm(hidden_size, device=device)
+        self.relu = nn.ReLU()
+
+        inner = []
+        for _ in range(num_inner):
+            inner.extend([
+                nn.Linear(hidden_size, hidden_size, device=device),
+                nn.LayerNorm(hidden_size, device=device),
+                nn.ReLU(),
+            ])
+        if inner:
+            self.inner = nn.Sequential(*inner)
+        else:
+            self.inner = nn.Identity()
+
+        self.fc2 = nn.Linear(hidden_size, output_size, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.norm(x)
+        x = self.relu(x)
+        x = self.inner(x)
+        x = self.fc2(x)
+        return x
+
+
+class MLP2(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, output_size: int,
+                 num_inner: int = 0,
+                 pre_norm: bool = False, device: torch.device = None,
+                 upsample_factor: int = 1,
+                 upsample_rank: int = None,
+                 from_config: bool = False,
+                 **kwargs):
+        super().__init__()
+
+        self.pre_norm = nn.Sequential(
+            nn.LayerNorm(input_size),
+            nn.GELU(),
+        ) if pre_norm else nn.Identity()
+
+        self.upsample_factor = upsample_factor
+        sq_ups = upsample_factor ** 2
+
+        self._real_output_dim = output_size // sq_ups
+
+        # hidden_size *= upsample_factor
+        # output_size *= (upsample_factor ** 2)
+
+        self.fc1 = nn.Linear(input_size, hidden_size, device=device)
+
+        blocks = []
+        for _ in range(num_inner):
+            blocks.append(nn.Sequential(
+                nn.LayerNorm(hidden_size, device=device),
+                nn.GELU(),
+                nn.Linear(hidden_size, hidden_size, device=device),
+            ))
+        self.blocks = nn.ModuleList(blocks)
+
+        self.final = nn.Sequential(
+            nn.LayerNorm(hidden_size, device=device),
+            nn.GELU(),
+            nn.Linear(hidden_size, output_size, device=device),
+        )
+
+    def forward(self, x: torch.Tensor, images: Optional[torch.Tensor] = None, patch_size: Optional[int] = None) -> torch.Tensor:
+        x = self.pre_norm(x)
+        x = self.fc1(x)
+        for block in self.blocks:
+            x = x + block(x)
+        x = self.final(x)
+
+        if self.upsample_factor > 1:
+            if images is None:
+                raise ValueError(f'`images` cannot be `None` when the head\'s `upsample_factor > 1`!')
+            if patch_size is None:
+                raise ValueError(f'`patch_size` cannot be `None` when the head\'s `upsample_factor > 1`!')
+            h, w = tuple(d // patch_size for d in images.shape[-2:])
+            x = rearrange(x, 'b (h w) (u1 u2 c) -> b (h u1 w u2) c',
+                          h=h, w=w, u1=self.upsample_factor, u2=self.upsample_factor,
+                          c=self._real_output_dim)
+
+        return x
+
+
+MLP_FACTORY = {
+    'v1': MLP,
+    'v2': MLP2,
+}
+
+
+def strip_prefix(state: Dict[str, torch.Tensor], prefix: str):
+    state = {
+        k[len(prefix):]: v
+        for k, v in state.items()
+        if k.startswith(prefix)
+    }
+    return state
+
+
+def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False):
+    state = strip_prefix(state, prefix)
+
+    weight_suffix = 'weight' if not spectral_weights else 'parametrizations.weight.original'
+
+    if version == 'v1':
+        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
+        output_dim = state[f'fc2.{weight_suffix}'].shape[0]
+
+        for num_inner in range(1000):
+            k = f'inner.{num_inner}.0.weight'
+            if k not in state:
+                break
+    elif version == 'v2':
+        hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
+        output_dim = state[f'final.2.{weight_suffix}'].shape[0]
+
+        for num_inner in range(1000):
+            k = f'blocks.{num_inner}.0.weight'
+            if k not in state:
+                break
+    else:
+        raise ValueError(f'Unsupported MLP version: {version}')
+
+    return input_dim, hidden_dim, output_dim, num_inner
+
+
+def create_mlp_from_config(version: str, input_dim: int, hidden_dim: int, output_dim: int, num_inner: int, **kwargs):
+    ret: nn.Module = MLP_FACTORY[version](input_dim, hidden_dim, output_dim, num_inner, from_config=True, **kwargs)
+
+    return ret
+
+
+def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False, **kwargs):
+    state = strip_prefix(state, prefix)
+
+    input_dim, hidden_dim, output_dim, num_inner = get_mlp_info_from_state(version, state, spectral_weights=spectral_weights)
+
+    ret: nn.Module = create_mlp_from_config(version, input_dim, hidden_dim, output_dim, num_inner, **kwargs)
+
+    if spectral_weights:
+        enable_spectral_reparam(ret, init_norm_to_current=False, state_dict_guidance=state)
+
+    ret.load_state_dict(state)
+
+    if spectral_weights:
+        disable_spectral_reparam(ret)
+
+    return ret

src/models/radiov3/adaptor_registry.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+from typing import Dict, Any
+
+import torch
+
+from .adaptor_generic import GenericAdaptor, AdaptorBase
+
+dict_t = Dict[str, Any]
+state_t = Dict[str, torch.Tensor]
+
+
+class AdaptorRegistry:
+    def __init__(self):
+        self._registry = {}
+
+    def register_adaptor(self, name):
+        def decorator(factory_function):
+            if name in self._registry:
+                raise ValueError(f"Model '{name}' already registered")
+            self._registry[name] = factory_function
+            return factory_function
+        return decorator
+
+    def create_adaptor(self, name, main_config: Namespace, adaptor_config: dict_t, state: state_t) -> AdaptorBase:
+        if name not in self._registry:
+            return GenericAdaptor(main_config, adaptor_config, state)
+        return self._registry[name](main_config, adaptor_config, state)
+
+# Creating an instance of the registry
+adaptor_registry = AdaptorRegistry()

src/models/radiov3/cls_token.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class ClsToken(nn.Module):
+    def __init__(self, ndim: int,
+                 num_tokens: int = 1,
+                 enabled: bool = True,
+                 register_multiple: Optional[int] = None,
+                 num_registers: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.ndim = ndim
+        self.enabled = enabled
+        self.num_registers = 0
+        self.num_tokens = num_tokens
+        if enabled:
+            if num_registers:
+                self.num_registers = num_registers
+            elif register_multiple:
+                self.num_registers = register_multiple - (num_tokens % register_multiple)
+
+            scale = ndim ** -0.5
+            self.token = nn.Parameter(torch.randn(num_tokens + self.num_registers, ndim) * scale)
+        else:
+            self.token = None
+
+        self.num_patches = self.num_tokens + self.num_registers
+
+    def disable(self):
+        self.token = None
+        self.enabled = False
+
+    def forward(self, x: torch.Tensor):
+        if self.token is None:
+            return x
+
+        token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)
+        x = torch.cat([
+            token,
+            x,
+        ], dim=1)
+
+        return x
+
+    def no_weight_decay(self):
+        return [
+            'token',
+        ]

src/models/radiov3/common.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from dataclasses import dataclass
+from typing import Optional
+
+from .radio_model import Resolution
+
+
+@dataclass
+class RadioResource:
+    url: str
+    patch_size: int
+    max_resolution: int
+    preferred_resolution: Resolution
+    supports_vitdet: bool = True
+    vitdet_num_windowed: Optional[int] = None
+    vitdet_num_global: Optional[int] = None
+
+
+RESOURCE_MAP = {
+    # RADIOv2.5
+    "radio_v2.5-b": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-b_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-l": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h-norm": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h-norm.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-g.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1792,
+        preferred_resolution=(896, 896),
+        vitdet_num_global=8,
+    ),
+    # RADIO
+    "radio_v2.1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.1_bf16.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v1.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1050,
+        preferred_resolution=Resolution(378, 378),
+    ),
+    # E-RADIO
+    "e-radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/eradio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+    # C-RADIO
+    "c-radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/C-RADIOv2-g/resolve/main/c-radio_v2-g_half.pth.tar",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=8,
+    ),
+    "c-radio_v3-b": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-L
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-B/resolve/main/c-radio-v3_b_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+        supports_vitdet=False,
+    ),
+    "c-radio_v3-l": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-L
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-L/resolve/main/c-radio-v3_l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+    "c-radio_v3-h": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-H
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-H/resolve/main/c-radio_v3-h_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+    "c-radio_v3-g": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-G
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-G/resolve/main/c-radio-v3_g_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+}
+
+DEFAULT_VERSION = "c-radio_v3-h"

src/models/radiov3/dinov2_arch.py

@@ -0,0 +1,1016 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+# Nvidia
+# NOTE: We re-define this model architecture primarily so that we don't have to worry about version compatibility breaking,
+# but also because Huggingface does a string replace of `gamma` to something else when loading the model state,
+# and this breaks loading of this model.
+
+from enum import Enum
+from functools import partial
+import logging
+import math
+import os
+import sys
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.init import trunc_normal_
+
+_torch_has_sdpa = hasattr(F, 'scaled_dot_product_attention')
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import fmha, scaled_index_add, index_select_cat, SwiGLU, memory_efficient_attention, unbind
+
+        XFORMERS_AVAILABLE = True
+    else:
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        if _torch_has_sdpa:
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                is_causal=False,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+                scale=self.scale,
+            )
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+if not XFORMERS_AVAILABLE:
+    SwiGLU = SwiGLUFFN
+
+
+class SwiGLUFFNFused(SwiGLU):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, torch.Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.grandma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.grandma) if self.inplace else x * self.grandma
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        # Huggingface is absurd and it will rename strings that contain `gamma`, which means that the normal DINO implementation
+        # of LayerScale won't work with HFHub. So we rename the variable to 'grandma', and support loading checkpoints in either
+        # format
+        key_a = f'{prefix}gamma'
+        key_b = f'{prefix}grandma'
+        if key_a in state_dict:
+            gamma = state_dict[key_a]
+        elif key_b in state_dict:
+            gamma = state_dict[key_b]
+        else:
+            if strict:
+                raise KeyError(f"Couldn't find the key {key_a} nor {key_b} in the state dict!")
+            else:
+                missing_keys.append(key_a)
+                missing_keys.append(key_b)
+                unexpected_keys.extend(state_dict.keys())
+                gamma = None
+
+        if gamma is not None:
+            self.grandma.data.copy_(gamma)
+
+        # return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def attn_residual_func(x: torch.Tensor) -> torch.Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: torch.Tensor) -> torch.Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+class NestedTensorBlock(Block):
+    def forward_nested(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.grandma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.grandma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, torch.Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            if not XFORMERS_AVAILABLE:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError
+
+
+def drop_add_residual_stochastic_depth(
+    x: torch.Tensor,
+    residual_func: Callable[[torch.Tensor], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> torch.Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[torch.Tensor],
+    residual_func: Callable[[torch.Tensor, Any], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> torch.Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    def forward(self, x):
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        ffn_bias=True,
+        proj_bias=True,
+        drop_path_rate=0.0,
+        drop_path_uniform=False,
+        init_values=None,  # for layerscale: None or 0 => no layerscale
+        embed_layer=PatchEmbed,
+        act_layer=nn.GELU,
+        block_fn=Block,
+        ffn_layer="mlp",
+        block_chunks=1,
+        num_register_tokens=0,
+        interpolate_antialias=False,
+        interpolate_offset=0.1,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            proj_bias (bool): enable bias for proj in attn if True
+            ffn_bias (bool): enable bias for ffn if True
+            drop_path_rate (float): stochastic depth rate
+            drop_path_uniform (bool): apply uniform drop rate across blocks
+            weight_init (str): weight init scheme
+            init_values (float): layer-scale init values
+            embed_layer (nn.Module): patch embedding layer
+            act_layer (nn.Module): MLP activation layer
+            block_fn (nn.Module): transformer block class
+            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+            num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+        """
+        super().__init__()
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.interpolate_antialias = interpolate_antialias
+        self.interpolate_offset = interpolate_offset
+
+        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        assert num_register_tokens >= 0
+        self.register_tokens = (
+            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+        )
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == "mlp":
+            ffn_layer = Mlp
+        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+            ffn_layer = SwiGLUFFNFused
+        elif ffn_layer == "identity":
+            def f(*args, **kwargs):
+                return nn.Identity()
+
+            ffn_layer = f
+        else:
+            raise NotImplementedError
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+    def interpolate_pos_encoding(self, x, w, h):
+        previous_dtype = x.dtype
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        pos_embed = self.pos_embed.float()
+        class_pos_embed = pos_embed[:, 0]
+        patch_pos_embed = pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        M = int(math.sqrt(N))  # Recover the number of patches in each dimension
+        assert N == M * M
+        kwargs = {}
+        if self.interpolate_offset:
+            # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8
+            # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors
+            sx = float(w0 + self.interpolate_offset) / M
+            sy = float(h0 + self.interpolate_offset) / M
+            kwargs["scale_factor"] = (sx, sy)
+        else:
+            # Simply specify an output size instead of a scale factor
+            kwargs["size"] = (w0, h0)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2),
+            mode="bicubic",
+            antialias=self.interpolate_antialias,
+            **kwargs,
+        )
+        assert (w0, h0) == patch_pos_embed.shape[-2:]
+        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).to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        if self.register_tokens is not None:
+            x = torch.cat(
+                (
+                    x[:, :1],
+                    self.register_tokens.expand(x.shape[0], -1, -1),
+                    x[:, 1:],
+                ),
+                dim=1,
+            )
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        return {
+            "x_norm_clstoken": x_norm[:, 0],
+            "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+            "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+            "x_prenorm": x,
+            "masks": masks,
+        }
+
+    def _get_intermediate_layers_not_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm=True,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
+        if reshape:
+            B, _, w, h = x.shape
+            outputs = [
+                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+    def forward(self, *args, is_training=False, **kwargs):
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
+    """
+    Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+    """
+    model = DinoVisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1536,
+        depth=40,
+        num_heads=24,
+        mlp_ratio=4,
+        block_fn=partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=num_register_tokens,
+        **kwargs,
+    )
+    return model
+
+
+class Weights(Enum):
+    LVD142M = "LVD142M"
+
+
+def _make_dinov2_model(
+    *,
+    arch_name: str = "vit_large",
+    img_size: int = 518,
+    patch_size: int = 14,
+    init_values: float = 1.0,
+    ffn_layer: str = "mlp",
+    block_chunks: int = 0,
+    num_register_tokens: int = 0,
+    interpolate_antialias: bool = False,
+    interpolate_offset: float = 0.1,
+    weights: Union[Weights, str] = Weights.LVD142M,
+    **kwargs,
+):
+    if isinstance(weights, str):
+        try:
+            weights = Weights[weights]
+        except KeyError:
+            raise AssertionError(f"Unsupported weights: {weights}")
+
+    vit_kwargs = dict(
+        img_size=img_size,
+        patch_size=patch_size,
+        init_values=init_values,
+        ffn_layer=ffn_layer,
+        block_chunks=block_chunks,
+        num_register_tokens=num_register_tokens,
+        interpolate_antialias=interpolate_antialias,
+        interpolate_offset=interpolate_offset,
+    )
+    vit_kwargs.update(**kwargs)
+    model = sys.modules[__name__].__dict__[arch_name](**vit_kwargs)
+
+    return model
+
+
+def dinov2_vits14(**kwargs):
+    """
+    DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_small", **kwargs)
+
+
+def dinov2_vitb14(**kwargs):
+    """
+    DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_base", **kwargs)
+
+
+def dinov2_vitl14(**kwargs):
+    """
+    DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(arch_name="vit_large", **kwargs)
+
+
+def dinov2_vitg14(**kwargs):
+    """
+    DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_giant2",
+        ffn_layer="swiglufused",
+        **kwargs,
+    )
+
+
+def dinov2_vits14_reg(**kwargs):
+    """
+    DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_small",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitb14_reg(**kwargs):
+    """
+    DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_base",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitl14_reg(**kwargs):
+    """
+    DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_large",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )
+
+
+def dinov2_vitg14_reg(**kwargs):
+    """
+    DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
+    """
+    return _make_dinov2_model(
+        arch_name="vit_giant2",
+        ffn_layer="swiglufused",
+        num_register_tokens=4,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+        **kwargs,
+    )

src/models/radiov3/dual_hybrid_vit.py

@@ -0,0 +1,213 @@
+from logging import getLogger
+from typing import Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from timm.models import register_model
+from timm.models import vision_transformer as tvit
+from timm.models import convnext as tconv
+
+from einops import rearrange
+
+from . import extra_timm_models as et
+
+
+class Fuser(nn.Module):
+    def __init__(self, src_dim: int, tgt_dim: int, gated: bool = True):
+        super().__init__()
+        self.gated = gated
+
+        mid_dim = max(src_dim, tgt_dim) * 2
+
+        self.fwd = nn.Sequential(
+            nn.Conv2d(src_dim, mid_dim, kernel_size=3, stride=1, padding=1),
+            nn.GELU(),
+            nn.Conv2d(mid_dim, tgt_dim * (2 if gated else 1), kernel_size=3, stride=1, padding=1),
+        )
+
+    def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
+        if src.ndim == 3:
+            shape = tgt.shape[-2:]
+        else:
+            shape = src.shape[-2:]
+
+        nd = shape[0] * shape[1]
+
+        if src.ndim == 3:
+            src = src[:, -nd:].reshape(src.shape[0], src.shape[2], *shape)
+
+        if tgt.ndim == 3:
+            tgt_pre = tgt[:, :-nd]
+            tgt = tgt[:, -nd:].reshape(tgt.shape[0], tgt.shape[2], *shape)
+        else:
+            tgt_pre = None
+
+        pred = self.fwd(src)
+
+        if self.gated:
+            g, pred = torch.chunk(pred, 2, dim=1)
+
+            g = F.sigmoid(g)
+
+            pred = g * pred
+
+        tgt = tgt + pred
+
+        if tgt_pre is not None:
+            tgt = rearrange(tgt, 'b c h w -> b (h w) c')
+            tgt = torch.cat([tgt_pre, tgt], dim=1)
+
+        return tgt
+
+
+class AttnDownsample(nn.Module):
+    def __init__(self, dim: int, window_size: int, num_heads: int = 16):
+        super().__init__()
+        self.q = nn.Parameter(torch.randn(1, num_heads, 1, dim // num_heads) * 0.01)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+        self.window_size = window_size
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+    def forward(self, x: torch.Tensor, twod_shape: Tuple[int, int]) -> torch.Tensor:
+        ntok = twod_shape[0] * twod_shape[1]
+        x_pre = x[:, :-ntok]
+
+        B = x.shape[0]
+        ds_hw = tuple(s // self.window_size for s in twod_shape)
+
+        x_spat = rearrange(
+            x[:, -ntok:],
+            'b (h d1 w d2) c -> (b h w) (d1 d2) c',
+            h=ds_hw[0], w=ds_hw[1],
+            d1=self.window_size, d2=self.window_size,
+        )
+
+        B, N, C = x_spat.shape
+
+        k, v = self.kv(x_spat).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+
+        q = (self.q * self.scale).expand(B, -1, -1, -1)
+        attn = q @ k.transpose(-2, -1)
+        attn = F.softmax(attn, dim=-1)
+        x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, C)
+        x = self.proj(x)
+
+        x = rearrange(x, '(b h w) c -> b (h w) c', b=x_pre.shape[0], h=ds_hw[0], w=ds_hw[1])
+
+        x = torch.cat([x_pre, x], dim=1)
+        return x
+
+
+class HybridModel(nn.Module):
+    def __init__(self, vit: tvit.VisionTransformer, conv: tconv.ConvNeXt, pretrained: bool = False,
+                 concatenate: bool = False, **kwargs):
+        super().__init__()
+        self.conv = conv
+        self.vit = vit
+        self.concatenate = concatenate
+
+        conv.stages = nn.ModuleList(conv.stages)
+        vit.blocks = nn.ModuleList(vit.blocks)
+
+        self._half_vit_idx = len(vit.blocks) // 2 + 1
+
+        self._half_conv_idx = None
+        x = torch.empty(1, 3, 256, 256)
+        x = self.conv.stem(x)
+        for i in range(len(conv.stages)):
+            x = conv.stages[i](x)
+            if self._half_conv_idx is None and x.shape[-2:] == (16, 16):
+                self._half_conv_idx = i + 1
+                half_conv_dim = x.shape[1]
+            final_conv_dim = x.shape[1]
+
+        self.vit_to_conv_fusion = Fuser(vit.embed_dim, half_conv_dim)
+        self.conv_to_vit_fusion = Fuser(half_conv_dim, vit.embed_dim)
+        self.vit_ds = AttnDownsample(vit.embed_dim, window_size=2)
+
+        embed_dim = vit.embed_dim + (final_conv_dim if concatenate else 0)
+        if not concatenate:
+            self.final_fuse = Fuser(final_conv_dim, vit.embed_dim, gated=False)
+        self.final_block = tvit.Block(embed_dim, num_heads=16)
+
+        self.embed_dim = embed_dim
+
+    @property
+    def patch_size(self):
+        return 32
+
+    @property
+    def no_fsdp_wrap_types(self):
+        return {tvit.VisionTransformer, tconv.ConvNeXt}
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_features(x)
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        y_vit = self.vit.patch_generator(x)
+
+        for i in range(self._half_vit_idx):
+            y_vit = self.vit.blocks[i](y_vit)
+
+        y_conv = self.conv.stem(x)
+        for i in range(self._half_conv_idx):
+            y_conv = self.conv.stages[i](y_conv)
+
+        y_vit, y_conv = self.conv_to_vit_fusion(y_conv, y_vit), self.vit_to_conv_fusion(y_vit, y_conv)
+
+        y_vit = self.vit_ds(y_vit, y_conv.shape[-2:])
+
+        for i in range(self._half_vit_idx, len(self.vit.blocks)):
+            y_vit = self.vit.blocks[i](y_vit)
+
+        for i in range(self._half_conv_idx, len(self.conv.stages)):
+            y_conv = self.conv.stages[i](y_conv)
+
+        if self.concatenate:
+            y_conv = rearrange(y_conv, 'b c h w -> b (h w) c')
+            # Average pool across the board, and replicate for each cls/register token
+            conv_summary = y_conv.mean(dim=1, keepdim=True).expand(-1, self.vit.patch_generator.num_cls_patches, -1)
+            y_conv = torch.cat([conv_summary, y_conv], dim=1)
+            y = torch.cat([y_vit, y_conv], dim=2)
+        else:
+            y = self.final_fuse(y_conv, y_vit)
+        y = self.final_block(y)
+
+        summary = y[:, :self.vit.patch_generator.num_cls_tokens]
+        features = y[:, self.vit.patch_generator.num_cls_patches:]
+
+        return summary, features
+
+
+@register_model
+def hybrid_base(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_base(pretrained=pretrained, **cfg)
+    vit = tvit.vit_base_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)
+
+
+@register_model
+def hybrid_large(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_large(pretrained=pretrained, **cfg)
+    vit = tvit.vit_large_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)
+
+
+@register_model
+def hybrid_huge(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
+    cfg = dict(num_classes=0, **kwargs)
+    conv = tconv.convnextv2_huge(pretrained=pretrained, **cfg)
+    vit = et.vit_huge_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
+
+    return HybridModel(vit, conv, pretrained, concatenate=concatenate)

src/models/radiov3/enable_cpe_support.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from typing import List, Optional, Set, Tuple, Union
+from types import MethodType
+
+import torch
+from torch import nn
+
+from timm.models import VisionTransformer, checkpoint_seq
+
+from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermediateFeatureNormalizer
+
+from .extra_models import DinoWrapper
+from .vit_patch_generator import ViTPatchGenerator
+from .forward_intermediates import forward_intermediates
+from .dual_hybrid_vit import HybridModel
+
+
+def _forward_cpe(self: VisionTransformer, x: torch.Tensor) -> torch.Tensor:
+    x = self.patch_generator(x)
+    if getattr(self, 'grad_checkpointing', False) and not torch.jit.is_scripting():
+        x = checkpoint_seq(self.blocks, x)
+    else:
+        x = self.blocks(x)
+    x = self.norm(x)
+    return x
+
+
+def _take_indices(
+        num_blocks: int,
+        n: Optional[Union[int, List[int], Tuple[int]]],
+) -> Tuple[Set[int], int]:
+    if isinstance(n, int):
+        assert n >= 0
+        take_indices = {x for x in range(num_blocks - n, num_blocks)}
+    else:
+        take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}
+    return take_indices, max(take_indices)
+
+
+def _forward_intermediates_cpe(
+        self,
+        x: torch.Tensor,
+        norm: bool = False,
+        **kwargs,
+) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+    return forward_intermediates(
+        self,
+        patch_extractor=self.patch_generator,
+        num_summary_tokens=self.patch_generator.num_skip,
+        num_cls_tokens=self.patch_generator.num_cls_tokens,
+        norm=self.norm if norm else lambda y: y,
+        x=x,
+        **kwargs,
+    )
+
+
+def _forward_cpe_dinov2(self: DinoWrapper, x: torch.Tensor) -> torch.Tensor:
+    y = _forward_cpe(self.inner, x)
+
+    return y[:, 0], y[:, self.num_summary_tokens:]
+
+
+def _forward_intermediates_cpe_dinov2(self: DinoWrapper, *args, **kwargs):
+    return _forward_intermediates_cpe(self.inner, *args, **kwargs)
+
+
+def _enable_cpe_for_timm_vit(model: VisionTransformer,
+                             max_img_size: Union[int, Tuple[int, int]] = 1024,
+                             num_cls_tokens: int = 1,
+                             pos_dropout: float = 0.1,
+                             register_multiple: int = Optional[None],
+                             num_registers: int = Optional[None],
+):
+    if not isinstance(model, VisionTransformer):
+        raise ValueError("CPE only support for VisionTransformer models!")
+
+    patch_size = model.patch_embed.patch_size[0]
+    embed_dim = model.embed_dim
+    input_dims = model.patch_embed.img_size
+    normalize_patches = not isinstance(model.patch_embed.norm, nn.Identity)
+    cls_token = model.cls_token is not None
+
+    max_img_size = int(round(max_img_size / patch_size) * patch_size)
+
+    patch_generator = ViTPatchGenerator(
+        patch_size=patch_size,
+        embed_dim=embed_dim,
+        input_dims=input_dims,
+        normalize_patches=normalize_patches,
+        cls_token=cls_token,
+        max_input_dims=max_img_size,
+        pos_dropout=pos_dropout,
+        num_cls_tokens=num_cls_tokens,
+        register_multiple=register_multiple,
+        num_registers=num_registers,
+    )
+
+    model.patch_generator = patch_generator
+    model.patch_embed = None
+    model.cls_token = None
+    model.pos_embed = None
+    model.pos_drop = None
+    model.patch_size = patch_size
+    model.num_cls_tokens = num_cls_tokens
+    model.num_registers = patch_generator.num_registers
+
+    model.forward_features = MethodType(_forward_cpe, model)
+    model.forward_intermediates = MethodType(_forward_intermediates_cpe, model)
+
+
+def _enable_cpe_for_dv2_reg_vit(model: DinoWrapper,
+                                max_img_size: Union[int, Tuple[int, int]] = 1024,
+                                num_cls_tokens: int = 1,
+                                pos_dropout: float = 0.1,
+                                register_multiple: int = Optional[None],
+                                num_registers: int = Optional[None],
+):
+    patch_size = model.patch_size
+    embed_dim = model.embed_dim
+    input_dims = model.inner.patch_embed.patches_resolution
+    normalize_patches = not isinstance(model.inner.patch_embed.norm, nn.Identity)
+    cls_token = True
+
+    max_img_size = int(round(max_img_size / patch_size) * patch_size)
+
+    patch_generator = ViTPatchGenerator(
+        patch_size=patch_size,
+        embed_dim=embed_dim,
+        input_dims=input_dims,
+        normalize_patches=normalize_patches,
+        cls_token=cls_token,
+        max_input_dims=max_img_size,
+        pos_dropout=pos_dropout,
+        num_cls_tokens=num_cls_tokens,
+        register_multiple=register_multiple,
+        num_registers=num_registers,
+        patch_bias=True,
+    )
+
+    inner = model.inner
+    inner.patch_generator = patch_generator
+    inner.patch_embed = None
+    inner.cls_token = None
+    inner.pos_embed = None
+    inner.register_tokens = None
+    inner.patch_size = patch_size
+
+    model.forward_features = MethodType(_forward_cpe_dinov2, model)
+    model.forward_intermediates = MethodType(_forward_intermediates_cpe_dinov2, model)
+
+
+def enable_cpe(model: nn.Module,
+               *args,
+               **kwargs,
+):
+    if isinstance(model, VisionTransformer):
+        _enable_cpe_for_timm_vit(model, *args, **kwargs)
+    elif isinstance(model, DinoWrapper):
+        _enable_cpe_for_dv2_reg_vit(model, *args, **kwargs)
+    elif isinstance(model, HybridModel):
+        _enable_cpe_for_timm_vit(model.vit, *args, **kwargs)
+    else:
+        raise ValueError(f'CPE not supported for this model type: {type(model)}')

src/models/radiov3/enable_spectral_reparam.py

@@ -0,0 +1,277 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from logging import getLogger
+import math
+import os
+from typing import Dict, List, Optional, Union, Tuple
+from types import MethodType
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import parametrize
+from torch.nn.utils.parametrizations import _SpectralNorm
+
+from timm.models.vision_transformer import Attention, Mlp
+
+_EPS = 1e-5
+
+
+class _SNReweight(_SpectralNorm):
+    def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, alpha: float = 0.05, version: int = 2, **kwargs):
+        super().__init__(weight, *args, **kwargs)
+
+        self.alpha = alpha
+        self.version = version
+        self.register_buffer('_sn_version', torch.tensor(version))
+
+        if init_norm_to_current:
+            # This will set the numerator to match the denominator, which should preserve the original values
+            init_scale = self._get_sigma(weight, n_power_iterations=20).item()
+        else:
+            init_scale = 1.0
+
+        if version == 1:
+            init_value = init_scale
+        elif version == 2:
+            t = init_scale - alpha
+            if t < _EPS:
+                getLogger("spectral_reparam").warn(f'The initialized spectral norm {init_scale} is too small to be represented. Setting to {_EPS} instead.')
+                t = _EPS
+
+            init_value = math.log(math.exp(t) - 1)
+        else:
+            raise ValueError(f'Unsupported version: {version}')
+
+        # Make 2D so that weight decay gets applied
+        self.scale = nn.Parameter(torch.tensor([[init_value]], dtype=torch.float32, device=weight.device))
+
+    # Re-implementing this because we need to make division by sigma safe
+    def _get_sigma(self, weight: torch.Tensor, n_power_iterations: int = None) -> torch.Tensor:
+        if not n_power_iterations:
+            n_power_iterations = self.n_power_iterations
+        if weight.ndim == 1:
+            # Faster and more exact path, no need to approximate anything
+            sigma = weight.norm()
+        else:
+            weight_mat = self._reshape_weight_to_matrix(weight)
+            if self.training:
+                self._power_method(weight_mat, n_power_iterations)
+            # See above on why we need to clone
+            u = self._u.clone(memory_format=torch.contiguous_format)
+            v = self._v.clone(memory_format=torch.contiguous_format)
+            # The proper way of computing this should be through F.bilinear, but
+            # it seems to have some efficiency issues:
+            # https://github.com/pytorch/pytorch/issues/58093
+            sigma = torch.dot(u, torch.mv(weight_mat, v))
+
+        return sigma + self.eps
+
+    def forward(self, weight: torch.Tensor, *args, **kwargs):
+        dtype = weight.dtype
+        sigma = self._get_sigma(weight, *args, **kwargs)
+
+        if self.version == 1:
+            scale = self.scale
+        elif self.version == 2:
+            scale = F.softplus(self.scale) + self.alpha
+        else:
+            raise ValueError(f'Unsupported version: {self.version}')
+
+        scale = scale.float() / sigma.float()
+
+        y = weight * scale
+
+        if dtype in (torch.float16, torch.bfloat16):
+            y = y.to(dtype)
+        return y
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        version_key = f'{prefix}_sn_version'
+        if version_key not in state_dict:
+            self.version = 1
+            state_dict[version_key] = torch.tensor(1)
+        return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+
+class _ChunkedSNReweight(nn.Module):
+    def __init__(self, weight: torch.Tensor, num_chunks: int, *args, init_norm_to_current: bool = False, **kwargs):
+        super().__init__()
+
+        self.num_chunks = num_chunks
+        parts = weight.split(weight.shape[0] // num_chunks, dim=0)
+
+        self.parts = nn.ModuleList([
+            _SNReweight(p, *args, init_norm_to_current=init_norm_to_current, **kwargs)
+            for p in parts
+        ])
+
+    def forward(self, weight: torch.Tensor, *args, **kwargs):
+        parts = weight.split(weight.shape[0] // self.num_chunks, dim=0)
+
+        parts = [
+            fn(p)
+            for fn, p in zip(self.parts, parts)
+        ]
+
+        return torch.cat(parts, dim=0)
+
+
+class _AttnSNReweight(_ChunkedSNReweight):
+    def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, renorm_values: bool = False, **kwargs):
+        super().__init__(weight, 3, *args, init_norm_to_current=init_norm_to_current, **kwargs)
+
+        if not renorm_values:
+            self.parts[2] = nn.Identity()
+
+
+def enable_spectral_reparam(model: Union[nn.Module, List[nn.Module]],
+                            n_power_iterations: int = 1,
+                            eps: float = 1e-6,
+                            init_norm_to_current: bool = False,
+                            renorm_values: bool = True,
+                            renorm_mlp: bool = True,
+                            state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):
+    if isinstance(model, (list, tuple)):
+        for i, sub in enumerate(model):
+            sub_sd = state_dict_guidance[i] if isinstance(state_dict_guidance, (list, tuple)) else state_dict_guidance
+            enable_spectral_reparam(sub, n_power_iterations=n_power_iterations, eps=eps,
+                                    init_norm_to_current=init_norm_to_current, renorm_values=renorm_values,
+                                    renorm_mlp=renorm_mlp, state_dict_guidance=sub_sd)
+        return
+
+    print('Enabling spectral reparametrization')
+    args = dict(n_power_iterations=n_power_iterations, dim=0, eps=eps, init_norm_to_current=init_norm_to_current)
+    visited_prefixes = set()
+
+    def is_guidance_parametrized(name: str):
+        if state_dict_guidance is None:
+            return True
+
+        p_name = f'{name}.parametrizations'
+        is_prm = any(k for k in state_dict_guidance if k.startswith(p_name) and k.endswith('_sn_version'))
+        return is_prm
+
+    def parametrize_linear(linear: nn.Linear):
+        parametrize.register_parametrization(
+            linear,
+            'weight',
+            _SNReweight(linear.weight, **args)
+        )
+
+    for name, mod in model.named_modules():
+        pref = '.'.join(name.split('.')[:-1])
+        if pref in visited_prefixes:
+            continue
+
+        if isinstance(mod, Attention) or name.endswith('.attn'):
+            if is_guidance_parametrized(f'{name}.qkv'):
+                parametrize.register_parametrization(
+                    mod.qkv,
+                    'weight',
+                    _AttnSNReweight(mod.qkv.weight, renorm_values=renorm_values, **args),
+                )
+            if hasattr(mod, 'proj') and is_guidance_parametrized(f'{name}.proj'):
+                parametrize_linear(mod.proj)
+            visited_prefixes.add(name)
+        elif name.endswith('mlp') and renorm_mlp and hasattr(mod, 'w12'):
+            if is_guidance_parametrized(f'{name}.w12'):
+                parametrize.register_parametrization(
+                    mod.w12,
+                    'weight',
+                    _ChunkedSNReweight(mod.w12.weight, num_chunks=2, **args),
+                )
+            if is_guidance_parametrized(f'{name}.w3'):
+                parametrize_linear(mod.w3)
+            visited_prefixes.add(name)
+        elif isinstance(mod, nn.Linear) and 'patch_generator' not in name and is_guidance_parametrized(name):
+            parametrize_linear(mod)
+
+
+def configure_spectral_reparam_from_args(model: nn.Module, args, state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):
+    spectral_reparam = getattr(args, 'spectral_reparam', False)
+    if isinstance(spectral_reparam, bool) and spectral_reparam:
+        enable_spectral_reparam(model, init_norm_to_current=True, state_dict_guidance=state_dict_guidance)
+    elif isinstance(spectral_reparam, dict):
+        enable_spectral_reparam(
+            model,
+            n_power_iterations=spectral_reparam.get('n_power_iterations', 1),
+            eps=spectral_reparam.get('eps', 1e-12),
+            init_norm_to_current=True,
+            state_dict_guidance=state_dict_guidance,
+        )
+
+
+def disable_spectral_reparam(model: nn.Module):
+    print('Disabling spectral reparametrization')
+    for name, mod in model.named_modules():
+        if parametrize.is_parametrized(mod):
+            parametrize.remove_parametrizations(mod, 'weight')
+            pass
+
+
+
+if __name__ == '__main__':
+    import argparse
+    from . import radio_model as create_model
+
+    parser = argparse.ArgumentParser(description='Remove parametrization from state dict')
+    parser.add_argument('--checkpoint', type=str, required=True, help='The checkpoint to load')
+    parser.add_argument('--output', type=str, default='', help='Where to store the checkpoint')
+    parser.add_argument('--release', default=False, action='store_true', help='Prune extraneous checkpoint fields')
+    parser.add_argument('--strict', default=False, action='store_true', help='Strictly load the state dict')
+
+    args = parser.parse_args()
+
+    if not args.output:
+        chk_dir, chk_name = os.path.split(args.checkpoint)
+        args.output = os.path.join(chk_dir, f'clean_{chk_name}')
+        print(f'Set output to "{args.output}"')
+
+    chk = torch.load(args.checkpoint, map_location='cpu', mmap=True)
+
+    model = create_model.create_model_from_args(chk['args'])
+
+    key = 'base_model.'
+    mod_state = dict()
+    extra_state = dict()
+    for k, v in chk['state_dict'].items():
+        if k.startswith(key):
+            mod_state[k[len(key):]] = v
+        else:
+            extra_state[k] = v
+
+    chk_load_info = model.load_state_dict(mod_state, strict=args.strict)
+    if chk_load_info.unexpected_keys or chk_load_info.missing_keys:
+        print(chk_load_info)
+
+    if chk['args'].spectral_reparam:
+        disable_spectral_reparam(model)
+
+    if hasattr(chk['args'], 'dtype'):
+        model.to(dtype=chk['args'].dtype)
+
+    mod_state = model.state_dict()
+    final_state = dict()
+    final_state.update({f'{key}{k}': v for k, v in mod_state.items()})
+    final_state.update(extra_state)
+
+    chk['state_dict'] = final_state
+    chk['args'].spectral_reparam = False
+
+    if args.release:
+        chk = {
+            'arch': chk['arch'],
+            'epoch': chk['epoch'],
+            'state_dict': chk['state_dict'],
+            'args': chk['args'],
+        }
+
+    torch.save(chk, args.output)
+    pass

src/models/radiov3/eradio_model.py

@@ -0,0 +1,1392 @@
+#!/usr/bin/env python3
+
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# E-RADIO model from
+# Mike Ranzinger, Greg Heinrich, Jan Kautz, and Pavlo Molchanov. "AM-RADIO: Agglomerative Model--Reduce All Domains Into One." arXiv preprint arXiv:2312.06709 (2023).
+
+# based on FasterViT, Swin Transformer, YOLOv8
+
+# FasterViT:
+# Ali Hatamizadeh, Greg Heinrich, Hongxu Yin, Andrew Tao, Jose M. Alvarez, Jan Kautz, and Pavlo Molchanov. "FasterViT: Fast Vision Transformers with Hierarchical Attention." arXiv preprint arXiv:2306.06189 (2023).
+
+import timm
+import torch
+import torch.nn as nn
+from timm.models.registry import register_model
+
+from timm.models.layers import trunc_normal_, DropPath, LayerNorm2d
+import numpy as np
+import torch.nn.functional as F
+import math
+import warnings
+
+#######################
+## Codebase from YOLOv8
+## BEGINNING
+#######################
+
+class C2f(nn.Module):
+    """Faster Implementation of CSP Bottleneck with 2 convolutions."""
+    """From YOLOv8 codebase"""
+    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, drop_path=None):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        if drop_path is None:
+            drop_path = [0.0] * n
+
+        self.c = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
+        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
+        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0, drop_path=drop_path[i]) for i in range(n))
+
+    def forward(self, x):
+        """Forward pass through C2f layer."""
+        y = list(self.cv1(x).chunk(2, 1))
+        y.extend(m(y[-1]) for m in self.m)
+        return self.cv2(torch.cat(y, 1))
+
+    def forward_split(self, x):
+        """Forward pass using split() instead of chunk()."""
+        y = list(self.cv1(x).split((self.c, self.c), 1))
+        y.extend(m(y[-1]) for m in self.m)
+        return self.cv2(torch.cat(y, 1))
+
+class Bottleneck(nn.Module):
+    """Standard bottleneck."""
+
+    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5, drop_path=0.0):  # ch_in, ch_out, shortcut, groups, kernels, expand
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, k[0], 1)
+        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
+        self.add = shortcut and c1 == c2
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        """'forward()' applies the YOLOv5 FPN to input data."""
+        return x + self.drop_path1(self.cv2(self.cv1(x))) if self.add else self.cv2(self.cv1(x))
+
+
+class Conv(nn.Module):
+    """Modified to support layer fusion"""
+    default_act = nn.SiLU()  # default activation
+
+    def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, dilation=1, bn_weight_init=1, bias=False, act=True):
+        super().__init__()
+
+        self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, autopad(kernel_size, padding, dilation), dilation, g, bias=False)
+        if 1:
+            self.bn = torch.nn.BatchNorm2d(b)
+            torch.nn.init.constant_(self.bn.weight, bn_weight_init)
+            torch.nn.init.constant_(self.bn.bias, 0)
+        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
+
+
+    def forward(self,x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.act(x)
+        return x
+
+    @torch.no_grad()
+    def switch_to_deploy(self):
+        # return 1
+        if not isinstance(self.bn, nn.Identity):
+            c, bn = self.conv, self.bn
+            w = bn.weight / (bn.running_var + bn.eps) ** 0.5
+            w = c.weight * w[:, None, None, None]
+            b = bn.bias - bn.running_mean * bn.weight / \
+                (bn.running_var + bn.eps)**0.5
+
+            self.conv.weight.data.copy_(w)
+            self.conv.bias = nn.Parameter(b)
+
+            self.bn = nn.Identity()
+
+def autopad(k, p=None, d=1):  # kernel, padding, dilation
+    """Pad to 'same' shape outputs."""
+    if d > 1:
+        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
+    return p
+
+
+#######################
+## Codebase from YOLOv8
+## END
+#######################
+
+def pixel_unshuffle(data, factor=2):
+    # performs nn.PixelShuffle(factor) in reverse, torch has some bug for ONNX and TRT, so doing it manually
+    B, C, H, W = data.shape
+    return data.view(B, C, factor, H//factor, factor, W//factor).permute(0,1,2,4,3,5).reshape(B, -1, H//factor, W//factor)
+
+class SwiGLU(nn.Module):
+    # should be more advanced, but doesnt improve results so far
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1)
+        return F.silu(gate) * x
+
+
+def window_partition(x, window_size):
+    """
+    Function for partitioning image into windows and later do windowed attention
+    Args:
+        x: (B, C, H, W)
+        window_size: window size
+    Returns:
+        windows - local window features (num_windows*B, window_size*window_size, C)
+        (Hp, Wp) -  the size of the padded image
+    """
+    B, C, H, W = x.shape
+
+    if window_size == 0 or (window_size==H and window_size==W):
+        windows = x.flatten(2).transpose(1, 2)
+        Hp, Wp = H, W
+    else:
+        pad_h = (window_size - H % window_size) % window_size
+        pad_w = (window_size - W % window_size) % window_size
+        if pad_h > 0 or pad_w > 0:
+            x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect")
+        Hp, Wp = H + pad_h, W + pad_w
+
+        x = x.view(B, C, Hp // window_size, window_size, Wp // window_size, window_size)
+        windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size*window_size, C)
+
+    return windows, (Hp, Wp)
+
+class Conv2d_BN(nn.Module):
+    '''
+    Conv2d + BN layer with folding capability to speed up inference
+    Can be merged with Conv() function with additional arguments
+    '''
+    def __init__(self, a, b, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bn_weight_init=1, bias=False):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, padding, dilation, groups, bias=False)
+        if 1:
+            self.bn = torch.nn.BatchNorm2d(b)
+            torch.nn.init.constant_(self.bn.weight, bn_weight_init)
+            torch.nn.init.constant_(self.bn.bias, 0)
+
+    def forward(self,x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+    @torch.no_grad()
+    def switch_to_deploy(self):
+        if not isinstance(self.bn, nn.Identity):
+            c, bn = self.conv, self.bn
+            w = bn.weight / (bn.running_var + bn.eps) ** 0.5
+            w = c.weight * w[:, None, None, None]
+            b = bn.bias - bn.running_mean * bn.weight / \
+                (bn.running_var + bn.eps)**0.5
+            self.conv.weight.data.copy_(w)
+            self.conv.bias = nn.Parameter(b)
+            self.bn = nn.Identity()
+
+
+
+def window_reverse(windows, window_size, H, W, pad_hw):
+    """
+    Windows to the full feature map
+    Args:
+        windows: local window features (num_windows*B, window_size, window_size, C)
+        window_size: Window size
+        H: Height of image
+        W: Width of image
+        pad_w - a tuple of image passing used in windowing step
+    Returns:
+        x: (B, C, H, W)
+
+    """
+    # print(f"window_reverse, windows.shape {windows.shape}")
+    Hp, Wp = pad_hw
+    if window_size == 0 or (window_size==H and window_size==W):
+        B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
+        x = windows.transpose(1, 2).view(B, -1, H, W)
+    else:
+        B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
+        x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
+        x = x.permute(0, 5, 1, 3, 2, 4).reshape(B,windows.shape[2], Hp, Wp)
+
+        if Hp > H or Wp > W:
+            x = x[:, :, :H, :W, ].contiguous()
+
+    return x
+
+
+
+class PosEmbMLPSwinv2D(nn.Module):
+    """
+    2D positional embedding from Swin Transformer v2
+    Added functionality to store the positional embedding in the model and not recompute it every time
+    """
+    def __init__(
+        self, window_size, pretrained_window_size, num_heads, seq_length, no_log=False, cpb_mlp_hidden=512,
+    ):
+        super().__init__()
+        self.window_size = window_size
+        self.num_heads = num_heads
+        # mlp to generate continuous relative position bias
+        self.cpb_mlp = nn.Sequential(
+            nn.Linear(2, cpb_mlp_hidden, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(cpb_mlp_hidden, num_heads, bias=False),
+        )
+
+        self.grid_exists = False
+        self.seq_length = seq_length
+        self.deploy = False
+        self.num_heads = num_heads
+        self.no_log = no_log
+        self.pretrained_window_size = pretrained_window_size
+        self.relative_bias_window_size = window_size
+
+        relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(window_size, num_heads,
+                                                                                                     pretrained_window_size, seq_length,
+                                                                                                     no_log)
+
+        self.register_buffer("relative_coords_table", relative_coords_table)
+        self.register_buffer("relative_position_index", relative_position_index)
+        self.register_buffer("relative_bias", relative_bias)  # for EMA
+
+    def relative_bias_initialization(self, window_size, num_heads, pretrained_window_size, seq_length, no_log):
+        # as in separate function to support window size chage after model weights loading
+        relative_coords_h = torch.arange(
+            -(window_size[0] - 1), window_size[0], dtype=torch.float32
+        )
+        relative_coords_w = torch.arange(
+            -(window_size[1] - 1), window_size[1], dtype=torch.float32
+        )
+        relative_coords_table = (
+            torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w]))
+            .permute(1, 2, 0)
+            .contiguous()
+            .unsqueeze(0)
+        )  # 1, 2*Wh-1, 2*Ww-1, 2
+        if pretrained_window_size[0] > 0:
+            relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1
+            relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1
+        else:
+            relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1
+            relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1
+
+        if not no_log:
+            relative_coords_table *= 8  # normalize to -8, 8
+            relative_coords_table = (
+                torch.sign(relative_coords_table)
+                * torch.log2(torch.abs(relative_coords_table) + 1.0)
+                / np.log2(8)
+            )
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = (
+            coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        )  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(
+            1, 2, 0
+        ).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+
+        relative_bias = torch.zeros(1, num_heads, seq_length, seq_length)
+
+        self.relative_bias_window_size = window_size
+
+        return relative_coords_table, relative_position_index, relative_bias
+
+
+    def switch_to_deploy(self):
+        self.deploy = True
+        self.grid_exists = True
+
+    def forward(self, input_tensor):
+        # for efficiency, we want this forward to be folded into a single operation (sum)
+        # if resolution stays the same, then we dont need to recompute MLP layers
+
+        if not self.deploy or self.training:
+            self.grid_exists = False
+
+        #compare if all elements in self.window_size list match those in self.relative_bias_window_size
+        if not all([self.window_size[i] == self.relative_bias_window_size[i] for i in range(len(self.window_size))]):
+            relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(self.window_size, self.num_heads,
+                                                                                                        self.pretrained_window_size, self.seq_length,
+                                                                                                        self.no_log)
+
+            self.relative_coords_table = relative_coords_table.to(self.relative_coords_table.device)
+            self.relative_position_index = relative_position_index.to(self.relative_position_index.device)
+            self.relative_bias = relative_bias.to(self.relative_bias.device)
+
+        if self.deploy and self.grid_exists:
+            input_tensor = input_tensor + self.relative_bias
+            return input_tensor
+
+        if 1:
+            self.grid_exists = True
+
+            relative_position_bias_table = self.cpb_mlp(
+                self.relative_coords_table
+            ).view(-1, self.num_heads)
+            relative_position_bias = relative_position_bias_table[
+                self.relative_position_index.view(-1)
+            ].view(
+                self.window_size[0] * self.window_size[1],
+                self.window_size[0] * self.window_size[1],
+                -1,
+            )  # Wh*Ww,Wh*Ww,nH
+
+            relative_position_bias = relative_position_bias.permute(
+                2, 0, 1
+            ).contiguous()  # nH, Wh*Ww, Wh*Ww
+            relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
+
+            self.relative_bias = relative_position_bias.unsqueeze(0)
+
+        input_tensor = input_tensor + self.relative_bias
+        return input_tensor
+
+
+class GRAAttentionBlock(nn.Module):
+    def __init__(self, window_size, dim_in, dim_out,
+                 num_heads, drop_path=0., qk_scale=None, qkv_bias=False,
+                 norm_layer=nn.LayerNorm, layer_scale=None,
+                  use_swiglu=True,
+                  subsample_ratio=1, dim_ratio=1, conv_base=False,
+                  do_windowing=True, multi_query=False, use_shift=0,
+                  cpb_mlp_hidden=512, conv_groups_ratio=0):
+        '''
+        Global Resolution Attention Block , see README for details
+        Attention with subsampling to get a bigger receptive field for attention
+        conv_base - use conv2d instead of avgpool2d for downsample / upsample
+
+
+        '''
+        super().__init__()
+
+        self.shift_size=window_size//2 if use_shift else 0
+
+        self.do_windowing = do_windowing
+        self.subsample_ratio = subsample_ratio
+
+
+
+        if do_windowing:
+            if conv_base:
+                    self.downsample_op = nn.Conv2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+
+
+                    self.downsample_mixer = nn.Identity()
+                    self.upsample_mixer = nn.Identity()
+                    self.upsample_op = nn.ConvTranspose2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+            else:
+                self.downsample_op = nn.AvgPool2d(kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+                self.downsample_mixer = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1) if subsample_ratio > 1 else nn.Identity()
+                self.upsample_mixer = nn.Upsample(scale_factor=subsample_ratio, mode='nearest') if subsample_ratio > 1 else nn.Identity()
+                self.upsample_op = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=False) if subsample_ratio > 1 else nn.Identity()
+
+
+        # in case there is no downsampling conv we want to have it separately
+        # will help with information propagation between windows
+        if subsample_ratio == 1:
+            # conv_groups_ratio=0
+            self.pre_conv = Conv2d_BN(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)
+            # self.pre_conv = nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)
+            # self.pre_conv_act = nn.ReLU6()
+            #for simplicity:
+            self.pre_conv_act = nn.Identity()
+            if conv_groups_ratio == -1:
+                self.pre_conv = nn.Identity()
+                self.pre_conv_act = nn.Identity()
+
+        self.window_size = window_size
+
+        self.norm1 = norm_layer(dim_in)
+
+        self.attn = WindowAttention(
+            dim_in,
+            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            resolution=window_size,
+            seq_length=window_size**2, dim_out=dim_in, multi_query=multi_query,
+            shift_size=self.shift_size, cpb_mlp_hidden=cpb_mlp_hidden)
+
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]
+        self.gamma1 = nn.Parameter(layer_scale * torch.ones(dim_in))  if use_layer_scale else 1
+
+        ### mlp layer
+        mlp_ratio = 4
+        self.norm2 = norm_layer(dim_in)
+        mlp_hidden_dim = int(dim_in * mlp_ratio)
+
+        activation = nn.GELU if not use_swiglu else SwiGLU
+        mlp_hidden_dim = int((4 * dim_in * 1 / 2) / 64) * 64 if use_swiglu else mlp_hidden_dim
+
+        self.mlp = Mlp(in_features=dim_in, hidden_features=mlp_hidden_dim, act_layer=activation, use_swiglu=use_swiglu)
+
+        self.gamma2 = nn.Parameter(layer_scale * torch.ones(dim_in)) if layer_scale else 1
+        self.drop_path2=DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+
+    def forward(self, x):
+        skip_connection = x
+        attn_mask = None
+
+        # in case there is no downsampling conv we want to have it separately
+        # will help with information propagation
+        if self.subsample_ratio == 1:
+            x = self.pre_conv_act(self.pre_conv(x)) + skip_connection
+
+        if self.do_windowing:
+            # performing windowing if required
+            x = self.downsample_op(x)
+            x = self.downsample_mixer(x)
+
+            if self.window_size>0:
+                H, W = x.shape[2], x.shape[3]
+
+            if self.shift_size > 0 and H>self.window_size and W>self.window_size:
+                # @swin like cyclic shift, doesnt show better performance
+                x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(2, 3))
+
+            x, pad_hw = window_partition(x, self.window_size)
+
+            if self.shift_size > 0 and H>self.window_size and W>self.window_size:
+                # set atten matrix to have -100 and the top right square
+                # attn[:, :, :-self.shift_size, -self.shift_size:] = -100.0
+                # calculate attention mask for SW-MSA
+                # not used in final version, can be useful for some cases especially for high res
+                H, W = pad_hw
+                img_mask = torch.zeros((1, H, W, 1), device=x.device)  # 1 H W 1
+                h_slices = (slice(0, -self.window_size),
+                            slice(-self.window_size, -self.shift_size),
+                            slice(-self.shift_size, None))
+                w_slices = (slice(0, -self.window_size),
+                            slice(-self.window_size, -self.shift_size),
+                            slice(-self.shift_size, None))
+                cnt = 0
+                for h in h_slices:
+                    for w in w_slices:
+                        img_mask[:, h, w, :] = cnt
+                        cnt += 1
+                img_mask = img_mask.transpose(1,2).transpose(1,3)
+                mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+
+                mask_windows = mask_windows[0].view(-1, self.window_size * self.window_size)
+                attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+                attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+        # window attention
+        x = x + self.drop_path1(self.gamma1*self.attn(self.norm1(x), attn_mask=attn_mask)) # or pass H,W
+        # mlp layer
+        x = x + self.drop_path2(self.gamma2*self.mlp(self.norm2(x)))
+
+        if self.do_windowing:
+            if self.window_size > 0:
+                x = window_reverse(x, self.window_size, H, W, pad_hw)
+
+            # reverse cyclic shift
+            if self.shift_size > 0 and H>self.window_size and W>self.window_size:
+                # @swin like cyclic shift, not tested
+                x = torch.roll(x, shifts=(self.shift_size, self.shift_size), dims=(2, 3))
+
+            x = self.upsample_mixer(x)
+            x = self.upsample_op(x)
+
+
+            if x.shape[2] != skip_connection.shape[2] or x.shape[3] != skip_connection.shape[3]:
+                x = torch.nn.functional.pad(x, ( 0, -x.shape[3] + skip_connection.shape[3], 0, -x.shape[2] + skip_connection.shape[2]), mode="reflect")
+        # need to add skip connection because downsampling and upsampling will break residual connection
+        # 0.5 is needed to make sure that the skip connection is not too strong
+        # in case of no downsample / upsample we can show that 0.5 compensates for the residual connection
+        x = 0.5 * x + 0.5 * skip_connection
+        return x
+
+
+
+
+class MultiResolutionAttention(nn.Module):
+    """
+    MultiResolutionAttention (MRA) module
+    The idea is to use multiple attention blocks with different resolution
+    Feature maps are downsampled / upsampled for each attention block on different blocks
+    Every attention block supports windowing
+    """
+
+    def __init__(self, window_size, sr_ratio,
+                 dim, dim_ratio, num_heads,
+                 do_windowing=True,
+                 layer_scale=1e-5, norm_layer=nn.LayerNorm,
+                 drop_path = 0, qkv_bias=False, qk_scale=1.0,
+                 use_swiglu=True, multi_query=False, conv_base=False,
+                 use_shift=0, cpb_mlp_hidden=512, conv_groups_ratio=0) -> None:
+        """
+        Args:
+            input_resolution: input image resolution
+            window_size: window size
+            compression_ratio: compression ratio
+            max_depth: maximum depth of the GRA module
+            use_shift: do window shifting
+        """
+        super().__init__()
+
+        depth = len(sr_ratio)
+
+        self.attention_blocks = nn.ModuleList()
+
+
+        for i in range(depth):
+            subsample_ratio = sr_ratio[i]
+            if len(window_size) > i:
+                window_size_local = window_size[i]
+            else:
+                window_size_local = window_size[0]
+
+            self.attention_blocks.append(GRAAttentionBlock(window_size=window_size_local,
+                                            dim_in=dim, dim_out=dim, num_heads=num_heads,
+                                            qkv_bias=qkv_bias, qk_scale=qk_scale, norm_layer=norm_layer,
+                                            layer_scale=layer_scale, drop_path=drop_path,
+                                            use_swiglu=use_swiglu, subsample_ratio=subsample_ratio, dim_ratio=dim_ratio,
+                                            do_windowing=do_windowing, multi_query=multi_query, conv_base=conv_base,
+                                            use_shift=use_shift, cpb_mlp_hidden=cpb_mlp_hidden, conv_groups_ratio=conv_groups_ratio),
+                                        )
+
+    def forward(self, x):
+
+        for attention_block in self.attention_blocks:
+            x = attention_block(x)
+
+        return x
+
+
+
+class Mlp(nn.Module):
+    """
+    Multi-Layer Perceptron (MLP) block
+    """
+
+    def __init__(self,
+                 in_features,
+                 hidden_features=None,
+                 out_features=None,
+                 act_layer=nn.GELU,
+                 use_swiglu=True,
+                 drop=0.):
+        """
+        Args:
+            in_features: input features dimension.
+            hidden_features: hidden features dimension.
+            out_features: output features dimension.
+            act_layer: activation function.
+            drop: dropout rate.
+        """
+
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features * (2 if use_swiglu else 1), bias=False)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=False)
+
+    def forward(self, x):
+        x_size = x.size()
+        x = x.view(-1, x_size[-1])
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        x = x.view(x_size)
+        return x
+
+class Downsample(nn.Module):
+    """
+    Down-sampling block
+    Pixel Unshuffle is used for down-sampling, works great accuracy - wise but takes 10% more TRT time
+    """
+
+    def __init__(self,
+                 dim,
+                 shuffle = False,
+                 ):
+        """
+        Args:
+            dim: feature size dimension.
+            shuffle: idea with
+            keep_dim: bool argument for maintaining the resolution.
+        """
+
+        super().__init__()
+        dim_out = 2 * dim
+
+        if shuffle:
+            self.norm = lambda x: pixel_unshuffle(x, factor=2)
+            self.reduction = Conv2d_BN(dim*4, dim_out, 1, 1, 0, bias=False)
+            # pixel unshuffleging works well but doesnt provide any speedup
+        else:
+            # removed layer norm for better, in this formulation we are getting 10% better speed
+            # LayerNorm for high resolution inputs will be a pain as it pools over the entire spatial dimension
+            # therefore we remove it compared to the original implementation in FasterViT
+            self.norm = nn.Identity()
+            self.reduction = Conv2d_BN(dim, dim_out, 3, 2, 1, bias=False)
+
+
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.reduction(x)
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    Patch embedding block
+    Used to convert image into an initial set of feature maps with lower resolution
+    """
+
+    def __init__(self, in_chans=3, in_dim=64, dim=96, shuffle_down=False):
+        """
+        Args:
+            in_chans: number of input channels.
+            in_dim: intermediate feature size dimension to speed up stem.
+            dim: final stem channel number
+            shuffle_down: use PixelUnshuffle for down-sampling, effectively increases the receptive field
+        """
+
+        super().__init__()
+        # shuffle_down = False
+        if not shuffle_down:
+            self.proj = nn.Identity()
+            self.conv_down = nn.Sequential(
+                Conv2d_BN(in_chans, in_dim, 3, 2, 1, bias=False),
+                nn.ReLU(),
+                Conv2d_BN(in_dim, dim, 3, 2, 1, bias=False),
+                nn.ReLU()
+                )
+        else:
+            self.proj = lambda x: pixel_unshuffle(x, factor=4)
+            self.conv_down = nn.Sequential(Conv2d_BN(in_chans*16, dim, 3, 1, 1),
+                                           nn.ReLU(),
+                                           )
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = self.conv_down(x)
+        return x
+
+
+
+class ConvBlock(nn.Module):
+    """
+    Convolutional block, used in first couple of stages
+    Experimented with plan resnet-18 like modules, they are the best in terms of throughput
+    Finally, YOLOv8 idea seem to work fine (resnet-18 like block with squeezed feature dimension, and feature concatendation at the end)
+    """
+    def __init__(self, dim,
+                 drop_path=0.,
+                 layer_scale=None,
+                 kernel_size=3,
+                 ):
+        super().__init__()
+
+        self.conv1 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
+        self.act1 = nn.GELU()
+
+        self.conv2 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
+
+        self.layer_scale = layer_scale
+        if layer_scale is not None and type(layer_scale) in [int, float]:
+            self.gamma = nn.Parameter(layer_scale * torch.ones(dim))
+            self.layer_scale = True
+        else:
+            self.layer_scale = False
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+
+        x = self.conv1(x)
+        x = self.act1(x)
+        x = self.conv2(x)
+
+        if self.layer_scale:
+            x = x * self.gamma.view(1, -1, 1, 1)
+        x = input + self.drop_path(x)
+        return x
+
+
+class WindowAttention(nn.Module):
+    # Windowed Attention from SwinV2
+    # use a MLP trick to deal with various input image resolutions, then fold it to improve speed
+
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, resolution=0,
+                 seq_length=0, dim_out=None, multi_query=False, shift_size=0, cpb_mlp_hidden=512):
+        # taken from EdgeViT and tweaked with attention bias.
+        super().__init__()
+        if not dim_out: dim_out = dim
+        self.shift_size = shift_size
+        self.multi_query = multi_query
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.head_dim = dim // num_heads
+
+        self.dim_internal = dim
+
+        self.scale = qk_scale or head_dim ** -0.5
+        if not multi_query:
+            self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        else:
+            self.qkv = nn.Linear(dim, dim + 2*self.head_dim, bias=qkv_bias)
+
+        self.proj = nn.Linear(dim, dim_out, bias=False)
+        # attention positional bias
+        self.pos_emb_funct = PosEmbMLPSwinv2D(window_size=[resolution, resolution],
+                                              pretrained_window_size=[resolution, resolution],
+                                              num_heads=num_heads,
+                                              seq_length=seq_length,
+                                              cpb_mlp_hidden=cpb_mlp_hidden)
+
+        self.resolution = resolution
+
+    def forward(self, x, attn_mask = None):
+        B, N, C = x.shape
+
+        if not self.multi_query:
+            qkv = self.qkv(x).reshape(B, -1, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+            q, k, v = qkv[0], qkv[1], qkv[2]
+        else:
+            qkv = self.qkv(x)
+            (q, k, v) = qkv.split([self.dim_internal, self.head_dim, self.head_dim], dim=2)
+
+            q = q.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+            k = k.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
+            v = v.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+
+        attn = self.pos_emb_funct(attn)
+
+        #add window shift
+        if attn_mask is not None:
+            nW = attn_mask.shape[0]
+            attn = attn.view(B // nW, nW, self.num_heads, N, N) + attn_mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = attn.softmax(dim=-1)
+        x = (attn @ v).transpose(1, 2).reshape(B, -1, C)
+        x = self.proj(x)
+        return x
+
+
+
+class ERADIOLayer(nn.Module):
+    """
+    E-RADIO Layer
+    """
+
+    def __init__(self,
+                 dim,
+                 depth,
+                 num_heads,
+                 window_size,
+                 conv=False,
+                 downsample=True,
+                 mlp_ratio=4.,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 norm_layer=nn.LayerNorm,
+                 drop_path=0.,
+                 layer_scale=None,
+                 layer_scale_conv=None,
+                 sr_dim_ratio=1,
+                 sr_ratio=1,
+                 multi_query=False,
+                 use_swiglu=True,
+                 yolo_arch=False,
+                 downsample_shuffle=False,
+                 conv_base=False,
+                 use_shift=False,
+                 cpb_mlp_hidden=512,
+                 conv_groups_ratio=0,
+                 verbose: bool = True,
+
+    ):
+        """
+        Args:
+            dim: feature size dimension.
+            depth: number of layers in each stage.
+            input_resolution: input image resolution.
+            window_size: window size in each stage.
+            downsample: bool argument for down-sampling.
+            mlp_ratio: MLP ratio.
+            num_heads: number of heads in each stage.
+            qkv_bias: bool argument for query, key, value learnable bias.
+            qk_scale: bool argument to scaling query, key.
+            drop: dropout rate.
+            attn_drop: attention dropout rate.
+            drop_path: drop path rate.
+            norm_layer: normalization layer.
+            layer_scale: layer scaling coefficient.
+            use_shift: SWIN like window shifting for half the window size for every alternating layer (considering multi-resolution)
+            conv_groups_ratio: group ratio for conv when no subsampling in multi-res attention
+        """
+
+        super().__init__()
+        self.conv = conv
+        self.yolo_arch=False
+        self.verbose = verbose
+        if conv:
+            if not yolo_arch:
+                self.blocks = nn.ModuleList([
+                    ConvBlock(dim=dim,
+                            drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                            layer_scale=layer_scale_conv)
+                    for i in range(depth)])
+                self.blocks = nn.Sequential(*self.blocks)
+            else:
+                self.blocks = C2f(dim,dim,n=depth,shortcut=True,e=0.5)
+                self.yolo_arch=True
+        else:
+            if not isinstance(window_size, list): window_size = [window_size]
+            self.window_size = window_size[0]
+            self.do_single_windowing = True
+            if not isinstance(sr_ratio, list): sr_ratio = [sr_ratio]
+            self.sr_ratio = sr_ratio
+            if any([sr!=1 for sr in sr_ratio]) or len(set(window_size))>1:
+                self.do_single_windowing = False
+                do_windowing = True
+            else:
+                self.do_single_windowing = True
+                do_windowing = False
+
+            #for v2_2
+            if conv_groups_ratio != -1:
+                self.do_single_windowing = False
+                do_windowing = True
+
+            self.blocks = nn.ModuleList()
+            for i in range(depth):
+                self.blocks.append(
+                    MultiResolutionAttention(window_size=window_size,
+                                             sr_ratio=sr_ratio,
+                                             dim=dim,
+                                             dim_ratio = sr_dim_ratio,
+                                             num_heads=num_heads,
+                                             norm_layer=norm_layer,
+                                             drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                             layer_scale=layer_scale,
+                                             qkv_bias=qkv_bias,
+                                             qk_scale=qk_scale,
+                                             use_swiglu=use_swiglu,
+                                             do_windowing=do_windowing,
+                                             multi_query=multi_query,
+                                             conv_base=conv_base,
+                                             cpb_mlp_hidden=cpb_mlp_hidden,
+                                             use_shift =0 if ((not use_shift) or ((i) % 2 == 0)) else True    ,
+                                             conv_groups_ratio=conv_groups_ratio,
+                    ))
+            self.blocks = nn.Sequential(*self.blocks)
+
+        self.transformer = not conv
+        self.downsample = None if not downsample else Downsample(dim=dim, shuffle=downsample_shuffle)
+
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+
+        # do padding for transforemr
+        interpolate = True
+        if self.transformer and interpolate:
+            # Windowed Attention will split feature map into windows with the size of window_size x window_size
+            # if the resolution is not divisible by window_size, we need to interpolate the feature map
+            # can be done via padding, but doing so after training hurts the model performance.
+            # interpolation affects the performance as well, but not as much as padding
+            if isinstance(self.window_size, list) or isinstance(self.window_size, tuple):
+                current_max_window_size = max(self.window_size)
+            else:
+                current_max_window_size = self.window_size
+
+            max_window_size = max([res_upsample*current_max_window_size for res_upsample in self.sr_ratio])
+            if H % max_window_size != 0 or W % max_window_size != 0:
+                new_h = int(np.ceil(H/max_window_size)*max_window_size)
+                new_w = int(np.ceil(W/max_window_size)*max_window_size)
+                x = F.interpolate(x, size=(new_h, new_w), mode='nearest')
+                if self.verbose:
+                    warnings.warn(f"Choosen window size is not optimal for given resolution. Interpolation of features maps will be done and it can affect the performance. Max window size is {max_window_size}, feature map size is {H}x{W}, interpolated feature map size is {new_h}x{new_w}.")
+
+
+        if self.transformer and self.do_single_windowing:
+            H, W = x.shape[2], x.shape[3]
+            x, pad_hw = window_partition(x, self.window_size)
+
+        #run main blocks
+        x = self.blocks(x)
+
+        if self.transformer and self.do_single_windowing:
+            x = window_reverse(x, self.window_size, H, W, pad_hw)
+
+        if self.transformer and interpolate:
+            #lets keep original resolution, might be not ideal, but for the upsampling tower we need to keep the expected resolution.
+            x = F.interpolate(x, size=(H, W), mode='nearest')
+
+        if self.downsample is None:
+            return x, x
+
+        return self.downsample(x), x  # changing to output pre downsampled features
+
+
+class InterpolateLayer(nn.Module):
+    def __init__(self, size=None, scale_factor=None, mode='nearest'):
+        super(InterpolateLayer, self).__init__()
+        self.size = size
+        self.scale_factor = scale_factor
+        self.mode = mode
+
+    def forward(self, x):
+        return F.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode)
+
+
+class HiResNeck(nn.Module):
+    """
+    The block is used to output dense features from all stages
+    Otherwise, by default, only the last stage features are returned with E-RADIO
+    """
+    def __init__(self, dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled):
+
+        '''
+        Hi Resolution neck to support output of high res features that are useful for dense tasks.
+        depths - total number of layers in the base model
+        neck_start_stage - when to start the neck, 0 - start from the first stage, 1 - start from the second stage etc.
+                            earlier layers result in higher resolution features at the cost of compute
+        full_features_head_dim - number of channels in the dense features head
+        '''
+        super().__init__()
+        # create feature projection layers for segmentation output
+        self.neck_features_proj = nn.ModuleList()
+        self.neck_start_stage = neck_start_stage
+        upsample_ratio = 1
+        for i in range(len(depths)):
+            level_n_features_output = int(dim * 2 ** i)
+
+            if self.neck_start_stage > i: continue
+
+            if (upsample_ratio > 1) or full_features_head_dim!=level_n_features_output:
+                feature_projection = nn.Sequential()
+                if False:
+                    feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output)) #fast, but worse
+                    feature_projection.add_module("dconv", nn.ConvTranspose2d(level_n_features_output,
+                                                                            full_features_head_dim, kernel_size=upsample_ratio, stride=upsample_ratio))
+                else:
+                    # B, in_channels, H, W -> B, in_channels, H*upsample_ratio, W*upsample_ratio
+                    # print("upsample ratio", upsample_ratio, level_n_features_output, level_n_features_output)
+                    feature_projection.add_module("upsample", InterpolateLayer(scale_factor=upsample_ratio, mode='nearest'))
+                    feature_projection.add_module("conv1", nn.Conv2d(level_n_features_output, level_n_features_output, kernel_size=3, stride=1, padding=1, groups=level_n_features_output))
+                    feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output))
+                    # B, in_channels, H*upsample_ratio, W*upsample_ratio -> B, full_features_head_dim, H*upsample_ratio, W*upsample_ratio
+                    feature_projection.add_module("conv2", nn.Conv2d(level_n_features_output, full_features_head_dim, kernel_size=1, stride=1, padding=0))
+            else:
+                feature_projection = nn.Sequential()
+
+            self.neck_features_proj.append(feature_projection)
+
+            if i>0 and downsample_enabled[i]:
+                upsample_ratio *= 2
+
+    def forward(self, x, il_level=-1, full_features=None):
+        if self.neck_start_stage > il_level:
+            return full_features
+
+        if full_features is None:
+            full_features = self.neck_features_proj[il_level - self.neck_start_stage](x)
+        else:
+            #upsample torch tensor x to match full_features size, and add to full_features
+            feature_projection = self.neck_features_proj[il_level - self.neck_start_stage](x)
+            if feature_projection.shape[2] != full_features.shape[2] or feature_projection.shape[3] != full_features.shape[3]:
+                feature_projection = torch.nn.functional.pad(feature_projection, ( 0, -feature_projection.shape[3] + full_features.shape[3], 0, -feature_projection.shape[2] + full_features.shape[2]))
+            full_features = full_features + feature_projection
+        return full_features
+
+class ERADIO(nn.Module):
+    """
+    Efficient RADIO
+    """
+
+    def __init__(self,
+                 dim,
+                 in_dim,
+                 depths,
+                 window_size,
+                 mlp_ratio,
+                 num_heads,
+                 drop_path_rate=0.2,
+                 in_chans=3,
+                 num_classes=1000,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 layer_scale=None,
+                 layer_scale_conv=None,
+                 layer_norm_last=False,
+                 sr_ratio = [1, 1, 1, 1],
+                 max_depth = -1,
+                 conv_base=False,
+                 use_swiglu=False,
+                 multi_query=False,
+                 norm_layer=nn.LayerNorm,
+                 drop_uniform=False,
+                 yolo_arch=False,
+                 shuffle_down=False,
+                 downsample_shuffle=False,
+                 return_full_features=False,
+                 full_features_head_dim=128,
+                 neck_start_stage=1,
+                 use_neck=False,
+                 use_shift=False,
+                 cpb_mlp_hidden=512,
+                 conv_groups_ratio=0,
+                 verbose: bool = False,
+                 **kwargs):
+        """
+        Args:
+            dim: feature size dimension.
+            depths: number of layers in each stage.
+            window_size: window size in each stage.
+            mlp_ratio: MLP ratio.
+            num_heads: number of heads in each stage.
+            drop_path_rate: drop path rate.
+            in_chans: number of input channels.
+            num_classes: number of classes.
+            qkv_bias: bool argument for query, key, value learnable bias.
+            qk_scale: bool argument to scaling query, key.
+            drop_rate: dropout rate.
+            attn_drop_rate: attention dropout rate.
+            norm_layer: normalization layer.
+            layer_scale: layer scaling coefficient.
+            return_full_features: output dense features as well as logits
+            full_features_head_dim: number of channels in the dense features head
+            neck_start_stage: a stage id to start full feature neck. Model has 4 stages, indix starts with 0
+                                for 224 resolution, the output of the stage before downsample:
+                                stage 0: 56x56, stage 1: 28x28, stage 2: 14x14, stage 3: 7x7
+            use_neck: even for summarization embedding use neck
+            use_shift: SWIN like window shifting but without masking attention
+            conv_groups_ratio: will be used for conv blocks where there is no multires attention,
+                                if 0 then normal conv,
+                                if 1 then channels are independent,
+                                if -1 then no conv at all
+
+        """
+        super().__init__()
+
+        num_features = int(dim * 2 ** (len(depths) - 1))
+        self.num_classes = num_classes
+        self.patch_embed = PatchEmbed(in_chans=in_chans, in_dim=in_dim, dim=dim, shuffle_down=shuffle_down)
+        # set return_full_features true if we want to return full features from all stages
+        self.return_full_features = return_full_features
+        self.use_neck = use_neck
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        if drop_uniform:
+            dpr = [drop_path_rate for x in range(sum(depths))]
+
+        if not isinstance(max_depth, list): max_depth = [max_depth] * len(depths)
+
+        self.levels = nn.ModuleList()
+        for i in range(len(depths)):
+            conv = True if (i == 0 or i == 1) else False
+
+            level = ERADIOLayer(dim=int(dim * 2 ** i),
+                                   depth=depths[i],
+                                   num_heads=num_heads[i],
+                                   window_size=window_size[i],
+                                   mlp_ratio=mlp_ratio,
+                                   qkv_bias=qkv_bias,
+                                   qk_scale=qk_scale,
+                                   conv=conv,
+                                   drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
+                                   downsample=(i < len(depths) - 1),
+                                   layer_scale=layer_scale,
+                                   layer_scale_conv=layer_scale_conv,
+                                   sr_ratio=sr_ratio[i],
+                                   use_swiglu=use_swiglu,
+                                   multi_query=multi_query,
+                                   norm_layer=norm_layer,
+                                   yolo_arch=yolo_arch,
+                                   downsample_shuffle=downsample_shuffle,
+                                   conv_base=conv_base,
+                                   cpb_mlp_hidden=cpb_mlp_hidden,
+                                   use_shift=use_shift,
+                                   conv_groups_ratio=conv_groups_ratio,
+                                   verbose=verbose)
+
+            self.levels.append(level)
+
+        if self.return_full_features or self.use_neck:
+            #num_heads
+            downsample_enabled = [self.levels[i-1].downsample is not None for i in range(len(self.levels))]
+            self.high_res_neck = HiResNeck(dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled)
+
+        self.switched_to_deploy = False
+
+        self.norm = LayerNorm2d(num_features) if layer_norm_last else nn.BatchNorm2d(num_features)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.head = nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, LayerNorm2d):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'rpb'}
+
+    def forward_features(self, x):
+        _, _, H, W = x.shape
+        if H % 32 != 0 or W % 32 != 0:
+            raise ValueError(f"E-RADIO requires input dimensions to be divisible by 32 but got H x W: {H} x {W}")
+        x = self.patch_embed(x)
+        full_features = None
+        for il, level in enumerate(self.levels):
+            x, pre_downsample_x = level(x)
+
+            if self.return_full_features or self.use_neck:
+                full_features = self.high_res_neck(pre_downsample_x, il, full_features)
+
+        # x = self.norm(full_features if (self.return_full_features or self.use_neck) else x)
+        x = self.norm(x) # new version for
+
+        if not self.return_full_features:
+            return x, None
+
+        return x, full_features
+
+    def forward(self, x):
+        x, full_features = self.forward_features(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+
+        x = self.head(x)
+        if full_features is not None:
+            return x, full_features
+        return x
+
+    def switch_to_deploy(self):
+        '''
+        A method to perform model self-compression
+        merges BN into conv layers
+        converts MLP relative positional bias into precomputed buffers
+        '''
+        if not self.switched_to_deploy:
+            for level in [self.patch_embed, self.levels, self.head]:
+                for module in level.modules():
+                    if hasattr(module, 'switch_to_deploy'):
+                        module.switch_to_deploy()
+        self.switched_to_deploy = True
+
+
+    def change_window_size(self, new_window_size):
+        """
+        E-RADIO employs windowed attention, which may be sensitive to the choice of this parameter,
+        especially in cases of uneven partitioning of the feature maps.
+        E-RADIO allows for the adjustment of the window size after training,
+        making it adaptable to different input image resolutions.
+        The recommended values for window size based on input resolution are as follows:
+
+        Input Resolution | Window Size
+        224 | 7
+        256 | 8
+        386 | 12
+        512 | 16
+        Ideally, the window size should be a factor of the input resolution. In the third stage, we divide the resolution by 16, so the window size should be
+        img_res/16/2
+        for the third stage and img_res/32 for the last stage. While this can be applied in a brute-force manner, a better way is to do model.change_window_size.
+        Manual way to change resolution -> model.change_window_size(resolution)
+        """
+        window_size = new_window_size
+        print(f"Setting window size to {window_size}")
+        for module in self.modules():
+            if hasattr(module, "window_size"):
+                # check if tuple or a number
+                if isinstance(module.window_size, tuple):
+                    if module.window_size[0] != window_size:
+                        module.window_size = (window_size, window_size)
+                elif isinstance(module.window_size, list):
+                    if module.window_size[0] != window_size:
+                        module.window_size = [window_size, window_size]
+                else:
+                    module.window_size = window_size
+
+
+    def set_optimal_window_size(self, image_dim, max_window_size = 16):
+        """
+        Using hand picked window size for various resolutions.
+
+        E-RADIO employs windowed attention, which may be sensitive to the choice of this parameter,
+        especially in cases of uneven partitioning of the feature maps.
+        E-RADIO allows for the adjustment of the window size after training,
+        making it adaptable to different input image resolutions.
+        The recommended values for window size based on input resolution are as follows:
+
+        Input Resolution | Window Size
+        224 | 7
+        256 | 8
+        386 | 12
+        512 | 16
+        Ideally, the window size should be a factor of the input resolution. In the third stage, we divide the resolution by 16, so the window size should be
+        img_res/16/2
+        for the third stage and img_res/32 for the last stage. While this can be applied in a brute-force manner, a better way is to do model.change_window_size.
+        Manual way to change resolution -> model.change_window_size(resolution)
+
+        """
+        # import math
+
+        def divisorGenerator(n):
+            large_divisors = []
+            for i in range(1, int(math.sqrt(n) + 1)):
+                if n % i == 0:
+                    yield i
+                    if i*i != n:
+                        large_divisors.append(n / i)
+            for divisor in reversed(large_divisors):
+                yield divisor
+
+        if isinstance(image_dim, list) or isinstance(image_dim, tuple):
+            image_dim = min(image_dim)
+
+        # we do windowed attention in the 3rd stage for the first time, therefore //16,
+        # we do subsampled attention with downsample by 2 so need to get //32 actually
+        # ideally we should rewrite this to be dependent on the structure of the model like what if subsampled is removed etc
+        all_divisors = np.array(list(divisorGenerator(image_dim//32)))
+        new_window_size = int(min(all_divisors[all_divisors <= max_window_size][-1], max_window_size))
+
+        # for image_dim in [128, 224, 256, 384, 512, 768, 1024]:
+        #     all_divisors = np.array(list(divisorGenerator(image_dim//32)))
+        #     new_window_size = int(min(all_divisors[all_divisors <= max_window_size][-1], max_window_size))
+        #     print(f"Setting window size to {new_window_size} for image resolution {image_dim}")
+
+        self.change_window_size(new_window_size = new_window_size)
+
+
+@register_model
+def eradio_large_fullres_ws16(pretrained=False, **kwargs):
+    model = ERADIO(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [16, 16], 16],
+        dim=192,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=1536,
+        neck_start_stage=2,
+        **kwargs,
+    )
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained)["state_dict"])
+    return model
+
+
+@register_model
+def eradio_xxxtiny(pretrained=False, **kwargs):  # ,
+    model = ERADIO(
+        depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [16, 16], 16],
+        dim=32,
+        in_dim=32,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=2,
+        **kwargs,
+    )
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained))
+    return model
+
+@register_model
+def eradio_xxxtiny_8x_ws12(pretrained=False, **kwargs):
+    model = ERADIO(depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [12, 12], 12],
+        dim=32,
+        in_dim=32,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        cpb_mlp_hidden=64,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=2,
+        conv_groups_ratio = 1,
+        **kwargs)
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained)["state_dict"])
+    return model
+
+
+@register_model
+def eradio_xxxtiny_8x_ws16(pretrained=False, **kwargs):
+    model = ERADIO(depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [16, 16], 16],
+        dim=32,
+        in_dim=32,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        cpb_mlp_hidden=64,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=1,
+        conv_groups_ratio = 1,
+        **kwargs)
+    if pretrained:
+        model.load_state_dict(torch.load(pretrained)["state_dict"])
+    return model
+
+@register_model
+def eradio(pretrained=False, **kwargs):
+    return eradio_large_fullres_ws16(pretrained=pretrained, **kwargs)

src/models/radiov3/extra_models.py

@@ -0,0 +1,206 @@
+from distutils.version import LooseVersion
+from types import MethodType
+from typing import List, Optional, Tuple, Union
+import warnings
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from timm.models.registry import register_model
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+from .forward_intermediates import forward_intermediates
+from .input_conditioner import InputConditioner
+
+_has_torch_sdpa = hasattr(F, 'scaled_dot_product_attention')
+
+
+class PaliGemmaWrapper(nn.Module):
+    def __init__(self, vis_model: nn.Module, embed_dim: int):
+        super().__init__()
+
+        self.vis_model = vis_model
+        self.embed_dim = embed_dim
+
+    @property
+    def patch_size(self):
+        return self.vis_model.embeddings.patch_size
+
+    @property
+    def blocks(self):
+        return self.vis_model.encoder.layers
+
+    @property
+    def embed_dim(self):
+        return self.vis_model.embeddings.embed_dim
+
+    def forward(self, x: torch.Tensor):
+        outputs = self.vis_model(
+            x,
+            return_dict=False,
+            interpolate_pos_encoding=True,
+        )
+
+        features = outputs[0].to(torch.float32)
+
+        summary = features.mean(dim=1)
+
+        return summary, features
+
+    def forward_features(self, x: torch.Tensor):
+        return self(x)
+
+
+def _get_paligemma_model(repo: str, embed_dim: int = None, dtype: torch.dtype = torch.bfloat16):
+    from transformers import PaliGemmaForConditionalGeneration, __version__ as tx_version
+
+    if LooseVersion(tx_version) > LooseVersion('4.44.2'):
+        warnings.warn(f'Your transformers version "{tx_version}" is higher than 4.44.2, and for whatever reason, PaliGemma might be broken.')
+
+    extra_args = dict()
+
+    if dtype is not None:
+        extra_args['torch_dtype'] = dtype
+        rev = str(dtype).split('.')[-1]
+        extra_args['revision'] = rev
+
+    model = PaliGemmaForConditionalGeneration.from_pretrained(repo, **extra_args)
+
+    vis_model = model.vision_tower.vision_model
+
+    vis_model = PaliGemmaWrapper(vis_model, embed_dim)
+
+    return vis_model
+
+@register_model
+def paligemma_896_student(**kwargs):
+    model = _get_paligemma_model('google/paligemma-3b-pt-896', embed_dim=1152, dtype=None)
+
+    return model
+
+
+def dv2_sdpa(self, x: torch.Tensor) -> torch.Tensor:
+    B, N, C = x.shape
+    qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+    q, k, v = qkv[0], qkv[1], qkv[2]
+    x = F.scaled_dot_product_attention(
+        q, k, v,
+        is_causal=False,
+        dropout_p=self.attn_drop.p if self.training else 0.,
+        scale=self.scale,
+    )
+    x = x.transpose(1, 2).reshape(B, N, C)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+def _load_dino_v2(dino_v2_model, cache_dir: Optional[str] = None, pretrained=True, **kwargs):
+    if cache_dir:
+        torch.hub.set_dir(cache_dir)
+    model: nn.Module = torch.hub.load(
+        'facebookresearch/dinov2',
+        dino_v2_model,
+        pretrained=pretrained,
+        # **kwargs,
+    )
+
+    if _has_torch_sdpa:
+        for n, m in model.named_modules():
+            if n.endswith('.attn'):
+                m.forward = MethodType(dv2_sdpa, m)
+
+    return model
+
+class DinoWrapper(nn.Module):
+    def __init__(self, dino_model: nn.Module):
+        super().__init__()
+
+        self.inner = dino_model
+        dino_model.blocks = nn.Sequential(*dino_model.blocks)
+
+    @property
+    def embed_dim(self):
+        return self.inner.embed_dim
+
+    @property
+    def patch_size(self):
+        return self.inner.patch_size
+
+    @property
+    def num_cls_tokens(self):
+        return getattr(self.inner, 'num_tokens', 1)
+
+    @property
+    def num_registers(self):
+        return getattr(self.inner, 'num_register_tokens', 0)
+
+    @property
+    def num_summary_tokens(self):
+        return self.num_cls_tokens + self.num_registers
+
+    @property
+    def blocks(self):
+        return self.inner.blocks
+
+    def forward(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        parts = self.inner.forward_features(*args, **kwargs)
+
+        cls_token = parts['x_norm_clstoken']
+        features = parts['x_norm_patchtokens']
+
+        return cls_token, features
+
+    def forward_features(self, x: torch.Tensor):
+        x = self.inner.prepare_tokens_with_masks(x)
+        x = self.inner.blocks(x)
+        x_norm = self.inner.norm(x)
+
+        return x_norm[:, 0], x_norm[:, self.num_summary_tokens:]
+
+    def patchify(self, x: torch.Tensor) -> torch.Tensor:
+        return self.inner.prepare_tokens_with_masks(x)
+
+    def forward_intermediates(self,
+        x: torch.Tensor,
+        norm: bool = False,
+        **kwargs,
+    ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+        return forward_intermediates(
+            self,
+            patch_extractor=self.inner.prepare_tokens_with_masks,
+            num_summary_tokens=self.num_summary_tokens,
+            num_cls_tokens=self.num_cls_tokens,
+            norm=self.inner.norm if norm else lambda y: y,
+            x=x,
+            **kwargs,
+        )
+
+
+def _dino_student(arch: str, **kwargs):
+    from . import dinov2_arch
+
+    factory = getattr(dinov2_arch, arch)
+    model = factory()
+
+    model = DinoWrapper(model)
+
+    conditioner = InputConditioner(
+        input_scale=1.0,
+        norm_mean=IMAGENET_DEFAULT_MEAN,
+        norm_std=IMAGENET_DEFAULT_STD,
+    )
+
+    model.input_conditioner = conditioner
+
+    return model
+
+
+@register_model
+def dino_v2_l_student(**kwargs):
+    return _dino_student('dinov2_vitl14_reg', **kwargs)
+
+@register_model
+def dino_v2_g_student(**kwargs):
+    return _dino_student('dinov2_vitg14_reg', **kwargs)

src/models/radiov3/extra_timm_models.py

@@ -0,0 +1,206 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import math
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from timm.models import register_model
+from timm.models.vision_transformer import (
+    VisionTransformer,
+    _create_vision_transformer as _timm_create_vision_transformer,
+    Mlp,
+    Block,
+    LayerScale as TIMMLayerScale,
+)
+
+# Import these to also register them
+from . import dinov2_arch
+
+
+@register_model
+def vit_tiny_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Tiny (Vit-Ti/16)
+    """
+    model_args = dict(patch_size=14, embed_dim=192, depth=12, num_heads=3)
+    model = _create_vision_transformer('vit_tiny_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_small_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Small (ViT-S/16)
+    """
+    model_args = dict(patch_size=14, embed_dim=384, depth=12, num_heads=6)
+    model = _create_vision_transformer('vit_small_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Base (ViT-B/14) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(patch_size=14, embed_dim=768, depth=12, num_heads=12)
+    model = _create_vision_transformer('vit_base_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_base_patch16_v2_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5,
+        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
+    )
+    model = _create_vision_transformer(
+        'vit_base_patch14_reg4_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_large_patch16_v2_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    name = 'vit_large_patch14_reg4_dinov2'
+    model_args = dict(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, init_values=1e-5,
+        reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
+    )
+    model = _create_vision_transformer(name, pretrained=pretrained, **dict(model_args, **kwargs))
+
+    return model
+
+@register_model
+def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model_args = dict(patch_size=16, embed_dim=1280, depth=32, num_heads=16)
+    if pretrained:
+        # There is no pretrained version of ViT-H/16, but we can adapt a ViT-H/14 for this purpose
+        model = _create_vision_transformer('vit_huge_patch14_224', pretrained=True, **dict(model_args, **kwargs))
+    else:
+        model = _create_vision_transformer('vit_huge_patch16_224', pretrained=False, **dict(model_args, **kwargs))
+    return model
+
+
+@register_model
+def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTransformer:
+    """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model = vit_huge_patch16_224(pretrained=pretrained, **kwargs)
+
+    for m in model.modules():
+        if isinstance(m, Mlp) and not isinstance(m.norm, nn.LayerNorm):
+            m.norm = nn.LayerNorm(m.fc1.out_features)
+
+    return model
+
+
+@register_model
+def vit_giant_patch16_224(pretrained=False, scaled_ln: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-giant model (ViT-g/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model_args = dict(patch_size=16, embed_dim=1536, depth=40, num_heads=24)
+    model = _create_vision_transformer('vit_giant_patch16_224', pretrained=False, **dict(model_args, **kwargs))
+    if scaled_ln:
+        _apply_scaled_ln(model)
+    return model
+
+
+@register_model
+def vit_bigG_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
+    model_args = dict(patch_size=14, embed_dim=1664, depth=48, num_heads=16, init_values=1e-6)
+    model = _create_vision_transformer('vit_bigG_patch14', pretrained=False, **dict(model_args, **kwargs))
+    return model
+
+
+def _create_vision_transformer(*args, **kwargs):
+    model = _timm_create_vision_transformer(*args, **kwargs)
+    _patch_layer_scale(model)
+    return model
+
+
+def _patch_layer_scale(model: VisionTransformer):
+    def replace_ls(old_ls: TIMMLayerScale):
+        new_ls = dinov2_arch.LayerScale(old_ls.gamma.shape[0], inplace=old_ls.inplace)
+        new_ls.load_state_dict(old_ls.state_dict())
+        return new_ls
+
+    # Monkey patch: Replace TIMM's LayerScale with our modified DINOv2 one, that uses a param name
+    # other than gamma, so that HFHub doesn't mess with it!
+    for mod in model.modules():
+        if isinstance(mod, Block):
+            if isinstance(mod.ls1, TIMMLayerScale):
+                mod.ls1 = replace_ls(mod.ls1)
+            if isinstance(mod.ls2, TIMMLayerScale):
+                mod.ls2 = replace_ls(mod.ls2)
+    pass
+
+
+class ScaledLayerNorm(nn.LayerNorm):
+    '''
+    https://arxiv.org/pdf/2502.05795v1
+    '''
+    def __init__(self, ln_base: nn.LayerNorm, depth: int = 0):
+        super().__init__(ln_base.normalized_shape, eps=ln_base.eps, elementwise_affine=ln_base.elementwise_affine)
+        self.load_state_dict(ln_base.state_dict())
+        self.register_buffer('ln_scale', torch.tensor(1.0 / math.sqrt(depth)), persistent=False)
+
+    def forward(self, x):
+        y = super().forward(x)
+        y = y * self.ln_scale
+        return y
+
+
+class DyT(nn.Module):
+    def __init__(self, C: int, init_alpha: float):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.full((1,), init_alpha))
+        self.gamma = nn.Parameter(torch.ones(C))
+        self.beta = nn.Parameter(torch.zeros(C))
+
+    def forward(self, x: torch.Tensor):
+        x = F.tanh(self.alpha * x)
+        return self.gamma * x + self.beta
+
+@register_model
+def vit_large_dyt_patch16_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
+    """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
+    ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
+    """
+    model_args = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16)
+    model = _create_vision_transformer('vit_large_dyt_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))
+
+    def _replace_ln_with_dyt(ln: nn.LayerNorm, depth: int):
+        return DyT(ln.normalized_shape[0], init_alpha=0.9)
+    _replace_ln(model, _replace_ln_with_dyt)
+
+    return model
+
+
+def _apply_scaled_ln(model: VisionTransformer):
+    warnings.warn('Post-LayerNorm scaling activated!')
+
+    _replace_ln(model, lambda ln, depth: ScaledLayerNorm(ln, depth=depth))
+
+def _replace_ln(model: VisionTransformer, fn):
+    def _inner_replace_ln(block: Block, depth: int, key: str):
+        prev = getattr(block, key)
+        if isinstance(prev, nn.LayerNorm):
+            setattr(block, key, fn(prev, depth=depth))
+
+    for i, block in enumerate(model.blocks):
+        _inner_replace_ln(block, i + 1, 'norm1')
+        _inner_replace_ln(block, i + 1, 'norm2')

src/models/radiov3/feature_normalizer.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from collections import namedtuple
+from typing import NamedTuple, Optional, Tuple
+import torch
+from torch import nn
+
+
+def _run_kernel(x: torch.Tensor, mean: torch.Tensor, tx: torch.Tensor):
+    if x.ndim <= 3:
+        x = x - mean
+        x = x @ tx.T
+    elif x.ndim == 4:
+        x = x - mean.reshape(1, -1, 1, 1)
+        kernel = tx.reshape(*tx.shape, 1, 1)
+        x = torch.nn.functional.conv2d(x, weight=kernel, bias=None, stride=1, padding=0)
+    else:
+        raise ValueError(f'Unsupported input dimension: {x.ndim}, shape: {x.shape}')
+    return x
+
+
+class FeatureNormalizer(nn.Module):
+    def __init__(self, embed_dim: int, dtype: torch.dtype = torch.float32):
+        super().__init__()
+
+        self.register_buffer('mean', torch.zeros(embed_dim, dtype=dtype))
+        self.register_buffer('tx', torch.eye(embed_dim, dtype=dtype))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = _run_kernel(x, self.mean, self.tx)
+        return x
+
+
+class InterFeatState(NamedTuple):
+    y: torch.Tensor
+    alpha: torch.Tensor
+
+
+class IntermediateFeatureNormalizerBase(nn.Module):
+    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
+        raise NotImplementedError()
+
+
+class IntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):
+    def __init__(self, num_intermediates: int, embed_dim: int, rot_per_layer: bool = False, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.register_buffer('alphas', torch.ones(num_intermediates, dtype=dtype))
+
+        rot = torch.eye(embed_dim, dtype=dtype)
+        if rot_per_layer:
+            rot = rot.unsqueeze(0).repeat(num_intermediates, 1, 1)
+
+        self.register_buffer('rotation', rot.contiguous())
+        self.register_buffer('means', torch.zeros(num_intermediates, embed_dim, dtype=dtype))
+
+    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
+        if rot_index is None:
+            rot_index = index
+
+        if skip:
+            assert x.ndim == 3, f'Cannot use the `skip` parameter when the `x` tensor isn\'t 3-dimensional.'
+            prefix, x = x[:, :skip], x[:, skip:]
+
+        rotation = self._get_rotation(rot_index)
+        y = _run_kernel(x, self.means[index], rotation)
+
+        alpha = self.alphas[index]
+        if skip:
+            alpha = torch.cat([
+                torch.ones(skip, dtype=alpha.dtype, device=alpha.device),
+                alpha[None].expand(y.shape[1]),
+            ]).reshape(1, -1, 1)
+            y = torch.cat([prefix, y], dim=1)
+        else:
+            if x.ndim == 3:
+                alpha = alpha.reshape(1, 1, 1).expand(1, y.shape[1], 1)
+            elif x.ndim == 4:
+                alpha = alpha.reshape(1, 1, 1, 1).expand(1, 1, *y.shape[2:])
+            else:
+                raise ValueError(f'Unsupported input dimension: {x.ndim}')
+
+        return InterFeatState(y, alpha)
+
+    def _get_rotation(self, rot_index: int) -> torch.Tensor:
+        if self.rotation.ndim == 2:
+            return self.rotation
+        return self.rotation[rot_index]
+
+
+class NullIntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):
+    instances = dict()
+
+    def __init__(self, dtype: torch.dtype, device: torch.device):
+        super().__init__()
+        self.register_buffer('alpha', torch.tensor(1, dtype=dtype, device=device))
+
+    @staticmethod
+    def get_instance(dtype: torch.dtype, device: torch.device):
+        instance = NullIntermediateFeatureNormalizer.instances.get((dtype, device), None)
+        if instance is None:
+            instance = NullIntermediateFeatureNormalizer(dtype, device)
+            NullIntermediateFeatureNormalizer.instances[(dtype, device)] = instance
+        return instance
+
+    def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
+        return InterFeatState(x, self.alpha)

src/models/radiov3/forward_intermediates.py

@@ -0,0 +1,138 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from typing import Callable, Dict, List, Optional, Set, Tuple, Union, Any, Iterable
+from types import MethodType
+
+import torch
+from torch import nn
+
+from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermediateFeatureNormalizer
+
+
+def _take_indices(
+        num_blocks: int,
+        n: Optional[Union[int, List[int], Tuple[int]]],
+) -> Tuple[Set[int], int]:
+    if isinstance(n, int):
+        assert n >= 0
+        take_indices = {x for x in range(num_blocks - n, num_blocks)}
+    else:
+        take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}
+    return take_indices, max(take_indices)
+
+
+def forward_intermediates(
+        model: nn.Module,
+        patch_extractor: Callable[[torch.Tensor], torch.Tensor],
+        norm: nn.Module,
+        num_summary_tokens: int,
+        num_cls_tokens: int,
+        x: torch.Tensor,
+        indices: Optional[Union[int, List[int], Tuple[int]]] = None,
+        return_prefix_tokens: bool = False,
+        stop_early: bool = False,
+        output_fmt: str = 'NCHW',
+        intermediates_only: bool = False,
+        aggregation: Optional[str] = "sparse",
+        inter_feature_normalizer: Optional[IntermediateFeatureNormalizerBase] = None,
+        norm_alpha_scheme = "post-alpha",
+        block_kwargs: Dict = None,
+) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+    """ Forward features that returns intermediates.
+
+    The Dense layer aggregation method is inspired from the paper: "Dense Connector for MLLMs"
+    by Yao, Huanjin et al. (2024). arXiv preprint arXiv:2405.13800}
+
+    Args:
+        x: Input image tensor
+        indices: Take last n blocks if int, select matching indices if sequence
+        return_prefix_tokens: Return both prefix and spatial intermediate tokens
+        norm: Apply norm layer to all intermediates
+        stop_early: Stop iterating over blocks when last desired intermediate hit
+        output_fmt: Shape of intermediate feature outputs
+        intermediates_only: Only return intermediate features
+        aggregation: intermediate layer aggregation method (sparse or dense)
+        norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha")
+    Returns:
+    """
+    assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.'
+    assert aggregation in ('sparse', 'dense'), 'Aggregation must be one of sparse or dense.'
+    reshape = output_fmt == 'NCHW'
+    intermediates = []
+
+    block_kwargs = block_kwargs or dict()
+
+    blocks = model.blocks
+
+    take_indices, max_index = _take_indices(len(blocks), indices)
+    take_indices = sorted(take_indices)
+    # forward pass
+    B, _, height, width = x.shape
+
+    x = patch_extractor(x)
+
+    if stop_early:
+        blocks = blocks[:max_index + 1]
+
+    if inter_feature_normalizer is None or norm_alpha_scheme == 'none':
+        inter_feature_normalizer = NullIntermediateFeatureNormalizer.get_instance(x.dtype, x.device)
+
+    assert norm_alpha_scheme in ('none', 'pre-alpha', 'post-alpha'), f'Unsupported alpha scheme: {norm_alpha_scheme}'
+    post_alpha_scheme = norm_alpha_scheme == 'post-alpha'
+
+    accumulator = 0
+    alpha_sum = 0
+    num_accumulated = 0
+
+    take_off = 0
+
+    for i, blk in enumerate(blocks):
+        x = blk(x, **block_kwargs)
+        if aggregation == "dense":
+            # Arbitrarily use the rotation matrix from the final layer in the dense group
+            y, alpha = inter_feature_normalizer(x, i, rot_index=take_indices[take_off], skip=num_summary_tokens)
+            if post_alpha_scheme:
+                accumulator = accumulator + y
+                alpha_sum = alpha_sum + alpha
+            else:
+                accumulator = accumulator + (alpha * y)
+                alpha_sum += 1
+            num_accumulated += 1
+        if i == take_indices[take_off]:
+            if aggregation == "dense":
+                alpha = alpha_sum / num_accumulated
+                x_ = alpha * accumulator / num_accumulated
+                num_accumulated = 0
+                accumulator = 0
+                alpha_sum = 0
+            else:
+                 y, alpha = inter_feature_normalizer(x, i, skip=num_summary_tokens)
+                 x_ = alpha * y
+            # normalize intermediates with final norm layer if enabled
+            intermediates.append(norm(x_))
+            take_off = min(take_off + 1, len(take_indices) - 1)
+
+    # process intermediates
+
+    # split prefix (e.g. class, distill) and spatial feature tokens
+    prefix_tokens = [y[:, :num_cls_tokens] for y in intermediates]
+    intermediates = [y[:, num_summary_tokens:] for y in intermediates]
+
+    if reshape:
+        # reshape to BCHW output format
+        H = height // model.patch_size
+        W = width // model.patch_size
+        intermediates = [y.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates]
+    if not torch.jit.is_scripting() and return_prefix_tokens:
+        # return_prefix not support in torchscript due to poor type handling
+        intermediates = list(zip(prefix_tokens, intermediates))
+    if intermediates_only:
+        return intermediates
+    x = norm(x)
+    return x, intermediates

src/models/radiov3/hf_model.py

@@ -0,0 +1,202 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from collections import namedtuple
+from typing import Callable, Dict, Optional, List, Union
+
+from timm.models import VisionTransformer
+import torch
+from torch import nn
+from transformers import PretrainedConfig, PreTrainedModel
+
+
+from .common import RESOURCE_MAP, DEFAULT_VERSION
+
+# Import all required modules.
+from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
+from .adaptor_generic import GenericAdaptor, AdaptorBase
+from .adaptor_mlp import create_mlp_from_config
+from .adaptor_registry import adaptor_registry
+from .cls_token import ClsToken
+from .dinov2_arch import dinov2_vitg14_reg
+from .enable_cpe_support import enable_cpe
+from .enable_spectral_reparam import configure_spectral_reparam_from_args
+from .eradio_model import eradio
+from .feature_normalizer import FeatureNormalizer, IntermediateFeatureNormalizer
+from .forward_intermediates import forward_intermediates
+from .radio_model import create_model_from_args
+from .radio_model import RADIOModel as RADIOModelBase, Resolution
+from .input_conditioner import get_default_conditioner, InputConditioner
+from .open_clip_adaptor import OpenCLIP_RADIO
+from .vit_patch_generator import ViTPatchGenerator
+from .vitdet import apply_vitdet_arch, VitDetArgs
+
+# Register extra models
+from .extra_timm_models import *
+from .extra_models import *
+
+
+class RADIOConfig(PretrainedConfig):
+    """Pretrained Hugging Face configuration for RADIO models."""
+
+    def __init__(
+        self,
+        args: Optional[dict] = None,
+        version: Optional[str] = DEFAULT_VERSION,
+        patch_size: Optional[int] = None,
+        max_resolution: Optional[int] = None,
+        preferred_resolution: Optional[Resolution] = None,
+        adaptor_names: Union[str, List[str]] = None,
+        adaptor_configs: Dict[str, Dict[str, int]] = None,
+        vitdet_window_size: Optional[int] = None,
+        feature_normalizer_config: Optional[dict] = None,
+        inter_feature_normalizer_config: Optional[dict] = None,
+        **kwargs,
+    ):
+        self.args = args
+        for field in ["dtype", "amp_dtype"]:
+            if self.args is not None and field in self.args:
+                # Convert to a string in order to make it serializable.
+                # For example for torch.float32 we will store "float32",
+                # for "bfloat16" we will store "bfloat16".
+                self.args[field] = str(args[field]).split(".")[-1]
+        self.version = version
+        resource = RESOURCE_MAP[version]
+        self.patch_size = patch_size or resource.patch_size
+        self.max_resolution = max_resolution or resource.max_resolution
+        self.preferred_resolution = (
+            preferred_resolution or resource.preferred_resolution
+        )
+        self.adaptor_names = adaptor_names
+        self.adaptor_configs = adaptor_configs
+        self.vitdet_window_size = vitdet_window_size
+        self.feature_normalizer_config = feature_normalizer_config
+        self.inter_feature_normalizer_config = inter_feature_normalizer_config
+        super().__init__(**kwargs)
+
+
+
+class RADIOModel(PreTrainedModel):
+    """Pretrained Hugging Face model for RADIO.
+
+    This class inherits from PreTrainedModel, which provides
+    HuggingFace's functionality for loading and saving models.
+    """
+
+    config_class = RADIOConfig
+
+    def __init__(self, config: RADIOConfig):
+        super().__init__(config)
+
+        RADIOArgs = namedtuple("RADIOArgs", config.args.keys())
+        args = RADIOArgs(**config.args)
+        self.config = config
+
+        model = create_model_from_args(args)
+        input_conditioner: InputConditioner = get_default_conditioner()
+
+        dtype = getattr(args, "dtype", torch.float32)
+        if isinstance(dtype, str):
+            # Convert the dtype's string representation back to a dtype.
+            dtype = getattr(torch, dtype)
+        model.to(dtype=dtype)
+        input_conditioner.dtype = dtype
+
+        summary_idxs = torch.tensor(
+            [i for i, t in enumerate(args.teachers) if t.get("use_summary", True)],
+            dtype=torch.int64,
+        )
+
+        adaptor_configs = config.adaptor_configs
+        adaptor_names = config.adaptor_names or []
+
+        adaptors = dict()
+        for adaptor_name in adaptor_names:
+            mlp_config = adaptor_configs[adaptor_name]
+            adaptor = GenericAdaptor(args, None, None, mlp_config)
+            adaptor.head_idx = mlp_config["head_idx"]
+            adaptors[adaptor_name] = adaptor
+
+        feature_normalizer = None
+        if config.feature_normalizer_config is not None:
+            # Actual normalization values will be restored when loading checkpoint weights.
+            feature_normalizer = FeatureNormalizer(config.feature_normalizer_config["embed_dim"])
+
+        inter_feature_normalizer = None
+        if config.inter_feature_normalizer_config is not None:
+            inter_feature_normalizer = IntermediateFeatureNormalizer(
+                config.inter_feature_normalizer_config["num_intermediates"],
+                config.inter_feature_normalizer_config["embed_dim"],
+                rot_per_layer=config.inter_feature_normalizer_config["rot_per_layer"],
+                dtype=dtype)
+
+        self.radio_model = RADIOModelBase(
+            model,
+            input_conditioner,
+            summary_idxs=summary_idxs,
+            patch_size=config.patch_size,
+            max_resolution=config.max_resolution,
+            window_size=config.vitdet_window_size,
+            preferred_resolution=config.preferred_resolution,
+            adaptors=adaptors,
+            feature_normalizer=feature_normalizer,
+            inter_feature_normalizer=inter_feature_normalizer,
+        )
+
+    @property
+    def adaptors(self) -> nn.ModuleDict:
+        return self.radio_model.adaptors
+
+    @property
+    def model(self) -> VisionTransformer:
+        return self.radio_model.model
+
+    @property
+    def input_conditioner(self) -> InputConditioner:
+        return self.radio_model.input_conditioner
+
+    @property
+    def num_summary_tokens(self) -> int:
+        return self.radio_model.num_summary_tokens
+
+    @property
+    def patch_size(self) -> int:
+        return self.radio_model.patch_size
+
+    @property
+    def max_resolution(self) -> int:
+        return self.radio_model.max_resolution
+
+    @property
+    def preferred_resolution(self) -> Resolution:
+        return self.radio_model.preferred_resolution
+
+    @property
+    def window_size(self) -> int:
+        return self.radio_model.window_size
+
+    @property
+    def min_resolution_step(self) -> int:
+        return self.radio_model.min_resolution_step
+
+    def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        return self.radio_model.make_preprocessor_external()
+
+    def get_nearest_supported_resolution(self, height: int, width: int) -> Resolution:
+        return self.radio_model.get_nearest_supported_resolution(height, width)
+
+    def switch_to_deploy(self):
+        return self.radio_model.switch_to_deploy()
+
+    def forward(self, x: torch.Tensor):
+        return self.radio_model.forward(x)

src/models/radiov3/input_conditioner.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from typing import Union, Tuple
+
+import torch
+from torch import nn
+
+
+norm_t = Union[Tuple[float, float, float], torch.Tensor]
+
+class InputConditioner(nn.Module):
+    def __init__(self,
+                 input_scale: float,
+                 norm_mean: norm_t,
+                 norm_std: norm_t,
+                 dtype: torch.dtype = None,
+    ):
+        super().__init__()
+
+        self.dtype = dtype
+
+        self.register_buffer("norm_mean", _to_tensor(norm_mean) / input_scale)
+        self.register_buffer("norm_std", _to_tensor(norm_std) / input_scale)
+
+    def forward(self, x: torch.Tensor):
+        y = (x - self.norm_mean) / self.norm_std
+        # if self.dtype is not None:
+        #     y = y.to(self.dtype)
+        return y
+
+
+def get_default_conditioner():
+    from timm.data.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
+
+    return InputConditioner(
+        input_scale=1.0,
+        norm_mean=OPENAI_CLIP_MEAN,
+        norm_std=OPENAI_CLIP_STD,
+    )
+
+
+def _to_tensor(v: norm_t):
+    return torch.as_tensor(v, dtype=torch.float32).view(-1, 1, 1)

src/models/radiov3/open_clip_adaptor.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from argparse import Namespace
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .adaptor_registry import adaptor_registry, dict_t, state_t
+
+from .adaptor_generic import GenericAdaptor
+
+
+class OpenCLIP_RADIO(GenericAdaptor):
+    def __init__(self, main_config: Namespace, adaptor_config: dict_t, state: state_t):
+        super().__init__(main_config, adaptor_config, state)
+
+        import open_clip
+
+        self.oc_model = open_clip.create_model_from_pretrained(
+            model_name=adaptor_config['model'],
+            pretrained=adaptor_config['pretrained'],
+            return_transform=False,
+        )
+        # Unload these parameters
+        self.oc_model.visual = None
+
+        self.tokenizer = open_clip.get_tokenizer(model_name=adaptor_config['model'])
+
+    def encode_text(self, text, normalize: bool = False):
+        return self.oc_model.encode_text(text, normalize=normalize)
+
+
+@adaptor_registry.register_adaptor("open_clip")
+def create_open_clip_adaptor(main_config: Namespace, adaptor_config: dict_t, state: state_t):
+    return OpenCLIP_RADIO(main_config, adaptor_config, state)

src/models/radiov3/radio_model.py

@@ -0,0 +1,344 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from timm.models import create_model, VisionTransformer
+
+from .enable_cpe_support import enable_cpe
+from .input_conditioner import InputConditioner
+from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
+from . import eradio_model
+from .enable_spectral_reparam import configure_spectral_reparam_from_args
+from .feature_normalizer import FeatureNormalizer, IntermediateFeatureNormalizer
+from . import dual_hybrid_vit
+
+
+class Resolution(NamedTuple):
+    height: int
+    width: int
+
+
+class RADIOModel(nn.Module):
+    def __init__(
+        self,
+        model: nn.Module,
+        input_conditioner: InputConditioner,
+        patch_size: int,
+        max_resolution: int,
+        preferred_resolution: Resolution,
+        summary_idxs: Optional[torch.Tensor] = None,
+        window_size: int = None,
+        adaptors: Dict[str, AdaptorBase] = None,
+        feature_normalizer: Optional[FeatureNormalizer] = None,
+        inter_feature_normalizer: Optional[IntermediateFeatureNormalizer] = None,
+    ):
+        super().__init__()
+
+        self.model = model
+        self.input_conditioner = input_conditioner
+        if summary_idxs is not None:
+            self.register_buffer('summary_idxs', summary_idxs)
+        else:
+            self.summary_idxs = None
+
+        self._preferred_resolution = preferred_resolution
+        self._patch_size = patch_size
+        self._max_resolution = max_resolution
+        self._window_size = window_size
+
+        adaptors = adaptors or dict()
+        self.adaptors = nn.ModuleDict(adaptors)
+
+        if feature_normalizer is None:
+            feature_normalizer = nn.Identity()
+        self.feature_normalizer = feature_normalizer
+        self.inter_feature_normalizer = inter_feature_normalizer
+
+    @property
+    def num_summary_tokens(self) -> int:
+        if hasattr(self.model, 'num_summary_tokens'):
+            return self.model.num_summary_tokens
+
+        patch_gen = getattr(self.model, "patch_generator", None)
+        if patch_gen is not None:
+            return patch_gen.num_skip
+        elif getattr(self.model, 'global_pool', None) == 'avg':
+            return 0
+        return 1
+
+    @property
+    def num_cls_tokens(self) -> int:
+        if hasattr(self.model, 'num_cls_tokens'):
+            return self.model.num_cls_tokens
+
+        patch_gen = getattr(self.model, 'patch_generator', None)
+        if patch_gen is not None:
+            return patch_gen.num_cls_tokens
+        elif getattr(self.model, 'global_pool', None) == 'avg':
+            return 0
+        return 1
+
+    @property
+    def patch_size(self) -> int:
+        if self._patch_size is not None:
+            return self._patch_size
+        if hasattr(self.model, "patch_size"):
+            return self.model.patch_size
+        patch_gen = getattr(self.model, "patch_generator", None)
+        if patch_gen is not None:
+            return patch_gen.patch_size
+        return None
+
+    @property
+    def max_resolution(self) -> int:
+        return self._max_resolution
+
+    @property
+    def preferred_resolution(self) -> Resolution:
+        return self._preferred_resolution
+
+    @property
+    def window_size(self) -> int:
+        return self._window_size
+
+    @property
+    def min_resolution_step(self) -> int:
+        res = self.patch_size
+        if self.window_size is not None:
+            res *= self.window_size
+        return res
+
+    @property
+    def blocks(self) -> Iterable[nn.Module]:
+        blocks = getattr(self.model, 'blocks', None)
+        if blocks is not None:
+            return blocks
+        return None
+
+    @property
+    def embed_dim(self) -> int:
+        return self.model.embed_dim
+
+    def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        ret = self.input_conditioner
+        self.input_conditioner = nn.Identity()
+        return ret
+
+    def get_nearest_supported_resolution(self, height: int, width: int) -> Resolution:
+        height = int(round(height / self.min_resolution_step) * self.min_resolution_step)
+        width = int(round(width / self.min_resolution_step) * self.min_resolution_step)
+
+        height = max(height, self.min_resolution_step)
+        width = max(width, self.min_resolution_step)
+
+        return Resolution(height=height, width=width)
+
+    def switch_to_deploy(self):
+        fn = getattr(self.model, 'switch_to_deploy', None)
+        if fn is not None:
+            fn()
+
+    def forward(self, x: torch.Tensor, feature_fmt: str = 'NLC') -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        '''
+        Forward process for model.
+        Args:
+            x: Input tensor. Unless `make_preprocessor_external` has been called, then the dynamic range of `x` is expected to be `[0, 1]`,
+                             otherwise `x` is expected to be mean centered with unit standard deviation.
+            feature_format: ['NLC', 'NCHW'] - The output format for the features.
+        '''
+        res_step = self.min_resolution_step
+        if res_step is not None and (x.shape[-2] % res_step != 0 or x.shape[-1] % res_step != 0):
+            raise ValueError('The input resolution must be a multiple of `self.min_resolution_step`. '
+                             '`self.get_nearest_supported_resolution(<height>, <width>) is provided as a convenience API. '
+                             f'Input: {x.shape[-2:]}, Nearest: {self.get_nearest_supported_resolution(*x.shape[-2:])}')
+
+        # import pdb; pdb.set_trace()
+        x = self.input_conditioner(x)
+        y = self.model.forward_features(x)
+        ret = self._extract_final(x, y, feature_fmt=feature_fmt)
+        return ret
+
+    def _extract_final(self, x: torch.Tensor, y: torch.Tensor, feature_fmt: str = 'NLC'):
+        if isinstance(self.model, VisionTransformer):
+            patch_gen = getattr(self.model, "patch_generator", None)
+            if patch_gen is not None:
+                all_summary = y[:, : patch_gen.num_cls_tokens]
+                if self.summary_idxs is not None:
+                    bb_summary = all_summary[:, self.summary_idxs]
+                else:
+                    bb_summary = all_summary
+                all_feat = y[:, patch_gen.num_skip :]
+            elif self.model.global_pool == "avg":
+                all_summary = y[:, self.model.num_prefix_tokens :].mean(dim=1)
+                bb_summary = all_summary
+                all_feat = y
+            else:
+                all_summary = y[:, 0]
+                bb_summary = all_summary
+                all_feat = y[:, 1:]
+        elif isinstance(self.model, eradio_model.ERADIO):
+            _, f = y
+            all_feat = f.flatten(2).transpose(1, 2)
+            all_summary = all_feat.mean(dim=1)
+            bb_summary = all_summary
+        elif isinstance(y, (list, tuple)):
+            all_summary, all_feat = y
+            bb_summary = all_summary
+        else:
+            all_summary = y[:, :self.num_cls_tokens]
+            if self.summary_idxs is not None and all_summary.shape[1] > 1:
+                if all_summary.shape[1] == 1:
+                    # Create dummy duplicates
+                    all_summary = all_summary.expand(-1, 128, -1)
+                bb_summary = all_summary[:, self.summary_idxs]
+            else:
+                bb_summary = all_summary
+            all_feat = y[:, self.num_summary_tokens:]
+
+        all_feat = self.feature_normalizer(all_feat)
+
+        if feature_fmt == 'NCHW':
+            fmt_feat = (all_feat.reshape(all_feat.shape[0], x.shape[-2] // self.patch_size, x.shape[-1] // self.patch_size, all_feat.shape[2])
+                                .permute(0, 3, 1, 2)
+            )
+        elif feature_fmt == 'NLC':
+            fmt_feat = all_feat
+        else:
+            raise ValueError(f'Unsupported feature_fmt: {feature_fmt}. Must be one of ["NLC", "NCHW"]')
+
+        ret = RadioOutput(bb_summary.flatten(1), fmt_feat)
+
+        if self.adaptors:
+            ret = dict(backbone=ret)
+            for name, adaptor in self.adaptors.items():
+                if all_summary.ndim == 3:
+                    if all_summary.shape[1] == 1:
+                        summary = all_summary[:, 0]
+                    else:
+                        summary = all_summary[:, adaptor.head_idx]
+                else:
+                    summary = all_summary
+                ada_input = AdaptorInput(images=x, summary=summary.float(), features=all_feat, feature_fmt=feature_fmt, patch_size=self.patch_size)
+                v = adaptor(ada_input).to(torch.float32)
+                ret[name] = v
+
+        return ret
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            indices: Optional[Union[int, List[int], Tuple[int]]] = None,
+            return_prefix_tokens: bool = False,
+            norm: bool = False,
+            stop_early: bool = False,
+            output_fmt: str = 'NCHW',
+            intermediates_only: bool = False,
+            aggregation: Optional[str] = "sparse",
+            norm_alpha_scheme: Optional[str] = "post-alpha",
+    ) -> List[RadioOutput]:
+        """ Forward features that returns intermediates.
+        Args:
+            x: Input image tensor
+            indices: Take last n blocks if int, select matching indices if sequence
+            return_prefix_tokens: Return both prefix and spatial intermediate tokens
+            norm: Apply norm layer to all intermediates
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            output_fmt: Shape of intermediate feature outputs. Options: NCHW, NLC
+            intermediates_only: Only return intermediate features
+            aggregation: intermediate layer aggregation method (sparse or dense).
+                Dense accumulation is done by averaging the features in each group.
+            norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha"), or don't normalize ("none")
+                Only affects dense aggregation
+        Returns:
+            List of RadioOutput objects.
+        """
+        x = self.input_conditioner(x)
+        intermediates = self.model.forward_intermediates(
+            x,
+            indices=indices,
+            return_prefix_tokens=return_prefix_tokens,
+            norm=norm,
+            stop_early=stop_early,
+            output_fmt=output_fmt,
+            intermediates_only=intermediates_only,
+            aggregation=aggregation,
+            inter_feature_normalizer=self.inter_feature_normalizer,
+            norm_alpha_scheme=norm_alpha_scheme,
+        )
+
+        if not intermediates_only:
+            final, intermediates = intermediates
+
+        def prepare_summary(summ: Optional[torch.Tensor]):
+            if summ is None:
+                return summ
+            if self.summary_idxs is not None and summ.shape[1] > 1:
+                summ = summ[:, self.summary_idxs]
+            return summ.flatten(1)
+
+        if return_prefix_tokens:
+            radio_outputs = [
+                RadioOutput(prepare_summary(summary), features)
+                for summary, features in intermediates
+            ]
+        else:
+            radio_outputs = intermediates
+
+        if intermediates_only:
+            return radio_outputs
+        else:
+            final = self._extract_final(x, final, feature_fmt=output_fmt)
+            return final, radio_outputs
+
+
+def create_model_from_args(args) -> nn.Module:
+    in_chans = 3
+    if args.in_chans is not None:
+        in_chans = args.in_chans
+    elif args.input_size is not None:
+        in_chans = args.input_size[0]
+
+    # Skip weight initialization unless it's explicitly requested.
+    weight_init = args.model_kwargs.pop("weight_init", "skip")
+
+    model = create_model(
+        args.model,
+        pretrained=args.pretrained,
+        in_chans=in_chans,
+        num_classes=args.num_classes,
+        drop_rate=args.drop,
+        drop_path_rate=args.drop_path,
+        drop_block_rate=args.drop_block,
+        global_pool=args.gp,
+        bn_momentum=args.bn_momentum,
+        bn_eps=args.bn_eps,
+        scriptable=args.torchscript,
+        checkpoint_path=args.initial_checkpoint,
+        weight_init=weight_init,
+        **args.model_kwargs,
+    )
+
+    if hasattr(model, 'norm') and not getattr(args, 'model_norm', False):
+        model.norm = nn.Identity()
+
+    model.head = nn.Identity()
+
+    if args.cpe_max_size is not None:
+        uq_teachers = set(t['name'] for t in args.teachers)
+        enable_cpe(
+            model,
+            args.cpe_max_size,
+            num_cls_tokens=len(uq_teachers) if args.cls_token_per_teacher else 1,
+            register_multiple=getattr(args, 'register_multiple', None),
+            num_registers=getattr(args, 'cpe_num_registers', None),
+        )
+
+    return model

src/models/radiov3/vit_patch_generator.py

@@ -0,0 +1,288 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import math
+from typing import Union, Tuple, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+
+from .cls_token import ClsToken
+
+input_dim_t = Union[int, Tuple[int, int]]
+
+try:
+    # raise ImportError()
+    from indirect_grid_sample import indirect_grid_sample
+except ImportError:
+    indirect_grid_sample = None
+
+class ViTPatchGenerator(nn.Module):
+    def __init__(self,
+                 patch_size: int,
+                 embed_dim: int,
+                 input_dims: input_dim_t,
+                 abs_pos: bool = True,
+                 normalize_patches: bool = False,
+                 cls_token: bool = False,
+                 max_input_dims: Optional[input_dim_t] = None,
+                 pos_dropout: float = 0.0,
+                 return_pos_enc: bool = False,
+                 num_cls_tokens: int = 1,
+                 register_multiple: Optional[int] = None,
+                 num_registers: Optional[int] = None,
+                 patch_bias: bool = False,
+                 device=None, dtype=None,
+    ):
+        super().__init__()
+
+        if isinstance(input_dims, int):
+            input_dims = (input_dims, input_dims)
+
+        if max_input_dims is None:
+            max_input_dims = input_dims
+        if isinstance(max_input_dims, int):
+            max_input_dims = (max_input_dims, max_input_dims)
+
+        max_input_dims = tuple(
+            int(math.ceil(d / patch_size) * patch_size)
+            for d in max_input_dims
+        )
+
+        self.cpe_mode = max_input_dims != input_dims
+        self.pos_dropout = pos_dropout
+        self.return_pos_enc = return_pos_enc
+
+        factory = dict(device=device, dtype=dtype)
+
+        self.patch_size = patch_size
+        self.abs_pos = abs_pos
+        self.embed_dim = embed_dim
+
+        self.num_rows = max_input_dims[0] // patch_size
+        self.num_cols = max_input_dims[1] // patch_size
+        self.input_dims = tuple(d // patch_size for d in input_dims)
+        self.num_patches = self.num_rows * self.num_cols
+        self.max_input_dims = max_input_dims
+
+        self.im_to_patches = Im2Patches(patch_size)
+        self.embedder = ViTPatchLinear(patch_size, embed_dim, bias=patch_bias, **factory)
+
+        if abs_pos:
+            scale = embed_dim ** -0.5
+            self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, embed_dim, **factory) * scale)
+
+        self.cls_token = ClsToken(
+            embed_dim,
+            num_tokens=num_cls_tokens,
+            enabled=cls_token,
+            register_multiple=register_multiple,
+            num_registers=num_registers,
+        )
+
+        self.patch_normalizer = nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        patches = self.embed_patches(x)
+        patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
+        patches = self.cls_token(patches)
+        patches = self.patch_normalizer(patches)
+        if self.return_pos_enc:
+            return patches, pos_enc
+        return patches
+
+    @property
+    def apply_cls_token(self):
+        return self.cls_token.enabled
+
+    @property
+    def num_cls_tokens(self):
+        return self.cls_token.num_tokens
+
+    @property
+    def num_cls_patches(self):
+        return self.cls_token.num_patches
+
+    @property
+    def num_registers(self):
+        return self.cls_token.num_registers
+
+    @property
+    def num_skip(self):
+        return self.num_cls_tokens + self.num_registers
+
+    def no_weight_decay(self):
+        return [
+            'pos_embed',
+        ]
+
+    def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
+        if src_embed.shape != targ_embed.shape:
+            src_size = int(math.sqrt(src_embed.shape[1]))
+
+            assert src_size ** 2 == src_embed.shape[1], 'Unable to interpolate non-square embedding'
+
+            src_embed = rearrange(src_embed, 'b (h w) c -> b c h w', h=src_size, w=src_size)
+            src_embed = F.interpolate(src_embed, size=(self.num_rows, self.num_cols), mode='bicubic', align_corners=True, antialias=False)
+            src_embed = rearrange(src_embed, 'b c h w -> b (h w) c')
+        targ_embed.data.copy_(src_embed)
+
+    def _load_projection(self, src_proj_weight: torch.Tensor, targ_proj_weight: torch.Tensor):
+        if src_proj_weight.shape != targ_proj_weight.shape:
+            src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3))
+
+            assert (src_patch_size ** 2) * 3 == src_proj_weight.shape[1], 'Unable to interpolate non-square patch size'
+
+            src_proj_weight = rearrange(src_proj_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
+            src_proj_weight = F.interpolate(src_proj_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
+            src_proj_weight = rearrange(src_proj_weight, 'b c h w -> b (c h w)')
+        targ_proj_weight.data.copy_(src_proj_weight)
+
+    def embed_patches(self, x: torch.Tensor) -> torch.Tensor:
+        # import pdb; pdb.set_trace()
+        patches = self.im_to_patches(x)
+        patches = self.embedder(patches)
+        return patches
+
+    def apply_pos_enc(self,
+                      patches: torch.Tensor,
+                      patch_idxs: Optional[torch.Tensor] = None,
+                      input_size: Optional[Tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        if not self.abs_pos:
+            return patches
+
+        pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size)
+
+        if self.training and self.pos_dropout > 0:
+            keeps = torch.rand(patches.shape[0], 1, 1, dtype=pos_enc.dtype, device=pos_enc.device) > self.pos_dropout
+            pos_enc_drop = torch.where(keeps, pos_enc, 0)
+        else:
+            pos_enc_drop = pos_enc
+
+        return patches + pos_enc_drop, pos_enc
+
+    def get_pos_enc(self,
+                    batch_size: int,
+                    patch_idxs: Optional[torch.Tensor] = None,
+                    input_size: Optional[Tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        if input_size is None:
+            input_dims = self.input_dims
+        else:
+            input_dims = tuple(d // self.patch_size for d in input_size)
+
+        pos_embed = self._get_pos_embeddings(batch_size, input_dims)
+
+        if patch_idxs is None:
+            return pos_embed
+
+        exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1])
+
+        pos_embed = torch.gather(pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs)
+        return pos_embed
+
+
+    def _get_pos_embeddings(self, batch_size: int, input_dims: Tuple[int, int]):
+        if (self.num_rows, self.num_cols) == input_dims:
+            return self.pos_embed
+
+        pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute(0, 3, 1, 2)
+
+        def window_select(pos_embed):
+            if input_dims[0] < pos_embed.shape[-2]:
+                pos_embed = pos_embed[..., :input_dims[0], :]
+            if input_dims[1] < pos_embed.shape[-1]:
+                pos_embed = pos_embed[..., :, :input_dims[1]]
+            return pos_embed
+
+        if self.cpe_mode:
+            if self.training:
+                min_scale = math.sqrt(0.1)
+                scale = torch.rand(batch_size, 1, 1, device=pos_embed.device) * (1 - min_scale) + min_scale
+                aspect_min = math.log(3 / 4)
+                aspect_max = -aspect_min
+                aspect = torch.exp(torch.rand(batch_size, 1, 1, device=pos_embed.device) * (aspect_max - aspect_min) + aspect_min)
+
+                scale_x = scale * aspect
+                scale_y = scale * (1 / aspect)
+                scale_xy = torch.stack([scale_x, scale_y], dim=-1).clamp_(0, 1)
+
+                pos_xy = torch.rand(batch_size, 1, 1, 2, device=pos_embed.device) * (1 - scale_xy)
+
+                lin_x = torch.linspace(0, 1, steps=input_dims[1], device=pos_embed.device)[None, None].expand(batch_size, input_dims[0], -1)
+                lin_y = torch.linspace(0, 1, steps=input_dims[0], device=pos_embed.device)[None, :, None].expand(batch_size, -1, input_dims[1])
+
+                lin_xy = torch.stack([lin_x, lin_y], dim=-1)
+
+                grid_xy = lin_xy * scale_xy + pos_xy
+
+                # Convert to [-1, 1] range
+                grid_xy.mul_(2).sub_(1)
+
+                pos_embed = F.grid_sample(
+                    pos_embed.float().expand(batch_size, -1, -1, -1),
+                    grid=grid_xy,
+                    mode='bilinear',
+                    padding_mode='zeros',
+                    align_corners=True,
+                ).to(pos_embed.dtype)
+            else:
+                # i_rows, i_cols = input_dims
+                # p_rows, p_cols = pos_embed.shape[2:]
+                # if i_rows <= p_rows and i_cols <= p_cols:
+                #     left = (p_cols - i_cols) // 2
+                #     top = (p_rows - i_rows) // 2
+                #     pos_embed = pos_embed[..., top:top+i_rows, left:left+i_cols]
+                # else:
+                max_dim = max(input_dims)
+                pos_embed = F.interpolate(pos_embed.float(), size=(max_dim, max_dim), align_corners=True, mode='bilinear').to(pos_embed.dtype)
+
+                pos_embed = window_select(pos_embed)
+        else:
+            pos_embed = window_select(pos_embed)
+
+        if pos_embed.shape[-2:] != input_dims:
+            pos_embed = F.interpolate(pos_embed.float(), size=input_dims, align_corners=True, mode='bilinear').to(pos_embed.dtype)
+
+        pos_embed = pos_embed.flatten(2).permute(0, 2, 1)
+
+        return pos_embed
+
+
+class Im2Patches(nn.Module):
+    def __init__(self, patch_size: int):
+        super().__init__()
+        self.patch_size = patch_size
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_size == 1:
+            patches = x.flatten(2)
+            patches = patches.permute(0, 2, 1)
+            return patches
+
+        py = x.shape[-2] // self.patch_size
+        px = x.shape[-1] // self.patch_size
+        patches = rearrange(x, 'b c (py yy) (px xx) -> b (py px) (c yy xx)',
+                            py=py, yy=self.patch_size,
+                            px=px, xx=self.patch_size,
+        )
+        return patches
+
+
+class ViTPatchLinear(nn.Linear):
+    def __init__(self, patch_size: int, embed_dim: int, bias: bool = False, **factory):
+        super().__init__(
+            3 * (patch_size ** 2),
+            embed_dim,
+            bias=bias,
+            **factory
+        )
+        self.patch_size = patch_size

src/models/radiov3/vitdet.py

@@ -0,0 +1,188 @@
+from collections import defaultdict
+from contextlib import contextmanager
+from logging import getLogger
+import math
+import sys
+from typing import List, Union, Iterable
+
+import numpy as np
+import torch
+from torch import nn
+
+from timm.models import VisionTransformer
+from einops import rearrange
+
+from .extra_models import DinoWrapper
+
+DEFAULT_NUM_WINDOWED = 5
+DEFAULT_NUM_GLOBAL = 4
+
+
+class VitDetArgs:
+    def __init__(self,
+                 window_size: int,
+                 num_summary_tokens: int,
+                 num_windowed: int = None,
+                 num_global: int = None,
+    ):
+        self.window_size = window_size
+        self.num_summary_tokens = num_summary_tokens
+        self.num_windowed = num_windowed
+        self.num_global = num_global
+
+
+def apply_vitdet_arch(model: Union[VisionTransformer, DinoWrapper], args: VitDetArgs):
+    if isinstance(model, VisionTransformer):
+        patch_embed = getattr(model, 'patch_generator', model.patch_embed)
+
+        return ViTDetHook(patch_embed, model.blocks, args)
+    elif isinstance(model, DinoWrapper):
+        inner = model.inner
+
+        patch_embed = getattr(inner, 'patch_generator', inner.patch_embed)
+        return ViTDetHook(patch_embed, inner.blocks, args)
+    else:
+        print(f'Warning: Unable to apply VitDet aug!', file=sys.stderr)
+
+
+class ViTDetHook:
+    def __init__(self,
+                 embedder: nn.Module,
+                 blocks: nn.Sequential,
+                 args: VitDetArgs,
+    ):
+        self.blocks = blocks
+        self.num_summary_tokens = args.num_summary_tokens
+        self.window_size = args.window_size
+
+        self._input_resolution = None
+        self._num_windows = None
+        self._cls_patch = None
+        self._order_cache = dict()
+
+        embedder.register_forward_pre_hook(self._enter_model)
+
+        # This will decide if we window-fy the patches
+        # and enable vit-det for this iteration, and if so,
+        # rearrange the patches for efficient mode switching
+        blocks.register_forward_pre_hook(self._enter_blocks)
+
+        is_global = True
+        if args.num_windowed is not None:
+            period = args.num_windowed + 1
+        else:
+            num_global = args.num_global or DEFAULT_NUM_GLOBAL
+            period = max(len(blocks) // num_global, 1)
+
+        for i, layer in enumerate(blocks[:-1]):
+            ctr = i % period
+            if ctr == 0:
+                layer.register_forward_pre_hook(self._to_windows)
+                is_global = False
+            elif ctr == period - 1:
+                layer.register_forward_pre_hook(self._to_global)
+                is_global = True
+
+        # Always ensure the final layer is a global layer
+        if not is_global:
+            blocks[-1].register_forward_pre_hook(self._to_global)
+
+        blocks.register_forward_hook(self._exit_model)
+
+    def _enter_model(self, _, input: List[torch.Tensor]):
+        self._input_resolution = input[0].shape[-2:]
+
+    def _enter_blocks(self, _, input: List[torch.Tensor]):
+        # print(f'{get_rank()} - ViTDet Window Size: {self._window_size}', file=sys.stderr)
+
+        patches = input[0]
+        patches = self._rearrange_patches(patches)
+
+        return (patches,) + input[1:]
+
+    def _to_windows(self, _, input: List[torch.Tensor]):
+        patches = input[0]
+
+        if self.num_summary_tokens:
+            self._cls_patch = patches[:, :self.num_summary_tokens]
+            patches = patches[:, self.num_summary_tokens:]
+
+        patches = rearrange(
+            patches, 'b (p t) c -> (b p) t c',
+            p=self._num_windows, t=self.window_size ** 2,
+        )
+
+        return (patches,) + input[1:]
+
+    def _to_global(self, _, input: List[torch.Tensor]):
+        patches = input[0]
+
+        patches = rearrange(
+            patches, '(b p) t c -> b (p t) c',
+            p=self._num_windows, t=self.window_size ** 2,
+            b=patches.shape[0] // self._num_windows,
+        )
+
+        if self.num_summary_tokens:
+            patches = torch.cat([
+                self._cls_patch,
+                patches,
+            ], dim=1)
+
+        return (patches,) + input[1:]
+
+    def _exit_model(self, _, inputs: List[torch.Tensor], patches: torch.Tensor):
+        # Return patches to their original order
+        patch_order = self._order_cache[self._input_resolution][0]
+        patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)
+
+        ret_patches = torch.empty_like(patches)
+        ret_patches = torch.scatter(
+            ret_patches,
+            dim=1,
+            index=patch_order,
+            src=patches,
+        )
+
+        return ret_patches
+
+    def _rearrange_patches(self, patches: torch.Tensor):
+        # We rearrange the patches so that we can efficiently
+        # switch between windowed and global mode by just
+        # reshaping the tensor
+
+        patch_order, self._num_windows = self._order_cache.get(self._input_resolution, (None, None))
+        if patch_order is None:
+            num_feat_patches = patches.shape[1] - self.num_summary_tokens
+            num_pixels = self._input_resolution[0] * self._input_resolution[1]
+
+            patch_size = int(round(math.sqrt(num_pixels / num_feat_patches)))
+            rows = self._input_resolution[-2] // patch_size
+            cols = self._input_resolution[-1] // patch_size
+
+            w_rows = rows // self.window_size
+            w_cols = cols // self.window_size
+
+            patch_order = torch.arange(0, num_feat_patches, device=patches.device)
+
+            patch_order = rearrange(
+                patch_order, '(wy py wx px) -> (wy wx py px)',
+                wy=w_rows, wx=w_cols,
+                py=self.window_size, px=self.window_size,
+            )
+
+            if self.num_summary_tokens:
+                patch_order = torch.cat([
+                    torch.arange(self.num_summary_tokens, dtype=patch_order.dtype, device=patch_order.device),
+                    patch_order + self.num_summary_tokens,
+                ])
+
+            self._num_windows = w_rows * w_cols
+            self._order_cache[self._input_resolution] = (
+                patch_order,
+                self._num_windows,
+            )
+
+        patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)
+        patches = torch.gather(patches, dim=1, index=patch_order)
+        return patches

src/models/stable_diffusion3/pipeline_stable_diffusion_3.py

@@ -0,0 +1,1256 @@
+# Copyright 2025 Stability AI, The HuggingFace Team and The InstantX Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    SiglipImageProcessor,
+    SiglipVisionModel,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.models.transformers import SD3Transformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import StableDiffusion3Pipeline
+
+        >>> pipe = StableDiffusion3Pipeline.from_pretrained(
+        ...     "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
+        ... )
+        >>> pipe.to("cuda")
+        >>> prompt = "A cat holding a sign that says hello world"
+        >>> image = pipe(prompt).images[0]
+        >>> image.save("sd3.png")
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin):
+    r"""
+    Args:
+        transformer ([`SD3Transformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant,
+            with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size`
+            as its dimension.
+        text_encoder_2 ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the
+            [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
+            variant.
+        text_encoder_3 ([`T5EncoderModel`]):
+            Frozen text-encoder. Stable Diffusion 3 uses
+            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
+            [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`CLIPTokenizer`):
+            Second Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_3 (`T5TokenizerFast`):
+            Tokenizer of class
+            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
+        image_encoder (`SiglipVisionModel`, *optional*):
+            Pre-trained Vision Model for IP Adapter.
+        feature_extractor (`SiglipImageProcessor`, *optional*):
+            Image processor for IP Adapter.
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae"
+    _optional_components = ["image_encoder", "feature_extractor"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"]
+
+    def __init__(
+        self,
+        transformer: SD3Transformer2DModel,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer_2: CLIPTokenizer,
+        text_encoder_3: T5EncoderModel,
+        tokenizer_3: T5TokenizerFast,
+        image_encoder: SiglipVisionModel = None,
+        feature_extractor: SiglipImageProcessor = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            text_encoder_3=text_encoder_3,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            tokenizer_3=tokenizer_3,
+            transformer=transformer,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = (
+            self.transformer.config.sample_size
+            if hasattr(self, "transformer") and self.transformer is not None
+            else 128
+        )
+        self.patch_size = (
+            self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2
+        )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 256,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if self.text_encoder_3 is None:
+            return torch.zeros(
+                (
+                    batch_size * num_images_per_prompt,
+                    self.tokenizer_max_length,
+                    self.transformer.config.joint_attention_dim,
+                ),
+                device=device,
+                dtype=dtype,
+            )
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            # logger.warning(
+            #     "The following part of your input was truncated because `max_sequence_length` is set to "
+            #     f" {max_sequence_length} tokens: {removed_text}"
+            # )
+
+        prompt_embeds = self.text_encoder_3(text_input_ids.to(device))[0]
+
+        dtype = self.text_encoder_3.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        clip_skip: Optional[int] = None,
+        clip_model_index: int = 0,
+    ):
+        device = device or self._execution_device
+
+        clip_tokenizers = [self.tokenizer, self.tokenizer_2]
+        clip_text_encoders = [self.text_encoder, self.text_encoder_2]
+
+        tokenizer = clip_tokenizers[clip_model_index]
+        text_encoder = clip_text_encoders[clip_model_index]
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            # logger.warning(
+            #     "The following part of your input was truncated because CLIP can only handle sequences up to"
+            #     f" {self.tokenizer_max_length} tokens: {removed_text}"
+            # )
+        prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+        pooled_prompt_embeds = prompt_embeds[0]
+
+        if clip_skip is None:
+            prompt_embeds = prompt_embeds.hidden_states[-2]
+        else:
+            prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds, pooled_prompt_embeds
+
+    def encode_pooled_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        clip_skip: Optional[int] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, SD3LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            _, pooled_prompt_embed = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=0,
+            )
+            _, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                prompt=prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=1,
+            )
+
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            _, negative_pooled_prompt_embed = self._get_clip_prompt_embeds(
+                negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=0,
+            )
+            _, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                negative_prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=1,
+            )
+
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        return pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        prompt_3: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        clip_skip: Optional[int] = None,
+        max_sequence_length: int = 256,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            negative_prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, SD3LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            prompt_3 = prompt_3 or prompt
+            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+            prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=0,
+            )
+            prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                prompt=prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=1,
+            )
+            clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
+
+            t5_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=prompt_3,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+            clip_prompt_embeds = torch.nn.functional.pad(
+                clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1])
+            )
+
+            prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            negative_prompt_3 = negative_prompt_3 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+            negative_prompt_3 = (
+                batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+            )
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds(
+                negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=0,
+            )
+            negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                negative_prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=1,
+            )
+            negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1)
+
+            t5_negative_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=negative_prompt_3,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+            negative_clip_prompt_embeds = torch.nn.functional.pad(
+                negative_clip_prompt_embeds,
+                (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]),
+            )
+
+            negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2)
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        prompt_3,
+        height,
+        width,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        negative_prompt_3=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if (
+            height % (self.vae_scale_factor * self.patch_size) != 0
+            or width % (self.vae_scale_factor * self.patch_size) != 0
+        ):
+            raise ValueError(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}."
+                f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}."
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_3 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_3`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
+            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_3 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        shape = (
+            batch_size,
+            num_channels_latents,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def skip_guidance_layers(self):
+        return self._skip_guidance_layers
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    # Adapted from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_image
+    def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch.Tensor:
+        """Encodes the given image into a feature representation using a pre-trained image encoder.
+
+        Args:
+            image (`PipelineImageInput`):
+                Input image to be encoded.
+            device: (`torch.device`):
+                Torch device.
+
+        Returns:
+            `torch.Tensor`: The encoded image feature representation.
+        """
+        if not isinstance(image, torch.Tensor):
+            image = self.feature_extractor(image, return_tensors="pt").pixel_values
+
+        image = image.to(device=device, dtype=self.dtype)
+
+        return self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+
+    # Adapted from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.prepare_ip_adapter_image_embeds
+    def prepare_ip_adapter_image_embeds(
+        self,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+    ) -> torch.Tensor:
+        """Prepares image embeddings for use in the IP-Adapter.
+
+        Either `ip_adapter_image` or `ip_adapter_image_embeds` must be passed.
+
+        Args:
+            ip_adapter_image (`PipelineImageInput`, *optional*):
+                The input image to extract features from for IP-Adapter.
+            ip_adapter_image_embeds (`torch.Tensor`, *optional*):
+                Precomputed image embeddings.
+            device: (`torch.device`, *optional*):
+                Torch device.
+            num_images_per_prompt (`int`, defaults to 1):
+                Number of images that should be generated per prompt.
+            do_classifier_free_guidance (`bool`, defaults to True):
+                Whether to use classifier free guidance or not.
+        """
+        device = device or self._execution_device
+
+        if ip_adapter_image_embeds is not None:
+            if do_classifier_free_guidance:
+                single_negative_image_embeds, single_image_embeds = ip_adapter_image_embeds.chunk(2)
+            else:
+                single_image_embeds = ip_adapter_image_embeds
+        elif ip_adapter_image is not None:
+            single_image_embeds = self.encode_image(ip_adapter_image, device)
+            if do_classifier_free_guidance:
+                single_negative_image_embeds = torch.zeros_like(single_image_embeds)
+        else:
+            raise ValueError("Neither `ip_adapter_image_embeds` or `ip_adapter_image_embeds` were provided.")
+
+        image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
+
+        if do_classifier_free_guidance:
+            negative_image_embeds = torch.cat([single_negative_image_embeds] * num_images_per_prompt, dim=0)
+            image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)
+
+        return image_embeds.to(device=device)
+
+    def enable_sequential_cpu_offload(self, *args, **kwargs):
+        if self.image_encoder is not None and "image_encoder" not in self._exclude_from_cpu_offload:
+            logger.warning(
+                "`pipe.enable_sequential_cpu_offload()` might fail for `image_encoder` if it uses "
+                "`torch.nn.MultiheadAttention`. You can exclude `image_encoder` from CPU offloading by calling "
+                "`pipe._exclude_from_cpu_offload.append('image_encoder')` before `pipe.enable_sequential_cpu_offload()`."
+            )
+
+        super().enable_sequential_cpu_offload(*args, **kwargs)
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_3: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 7.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        cond_latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 256,
+        skip_guidance_layers: List[int] = None,
+        skip_layer_guidance_scale: float = 2.8,
+        skip_layer_guidance_stop: float = 0.2,
+        skip_layer_guidance_start: float = 0.01,
+        mu: Optional[float] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion
+                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
+                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
+                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
+                the text `prompt`, usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used instead
+            negative_prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used instead
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            ip_adapter_image (`PipelineImageInput`, *optional*):
+                Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`torch.Tensor`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. Should be a tensor of shape `(batch_size, num_images,
+                emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to
+                `True`. If not provided, embeddings are computed from the `ip_adapter_image` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] instead of
+                a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`.
+            skip_guidance_layers (`List[int]`, *optional*):
+                A list of integers that specify layers to skip during guidance. If not provided, all layers will be
+                used for guidance. If provided, the guidance will only be applied to the layers specified in the list.
+                Recommended value by StabiltyAI for Stable Diffusion 3.5 Medium is [7, 8, 9].
+            skip_layer_guidance_scale (`int`, *optional*): The scale of the guidance for the layers specified in
+                `skip_guidance_layers`. The guidance will be applied to the layers specified in `skip_guidance_layers`
+                with a scale of `skip_layer_guidance_scale`. The guidance will be applied to the rest of the layers
+                with a scale of `1`.
+            skip_layer_guidance_stop (`int`, *optional*): The step at which the guidance for the layers specified in
+                `skip_guidance_layers` will stop. The guidance will be applied to the layers specified in
+                `skip_guidance_layers` until the fraction specified in `skip_layer_guidance_stop`. Recommended value by
+                StabiltyAI for Stable Diffusion 3.5 Medium is 0.2.
+            skip_layer_guidance_start (`int`, *optional*): The step at which the guidance for the layers specified in
+                `skip_guidance_layers` will start. The guidance will be applied to the layers specified in
+                `skip_guidance_layers` from the fraction specified in `skip_layer_guidance_start`. Recommended value by
+                StabiltyAI for Stable Diffusion 3.5 Medium is 0.01.
+            mu (`float`, *optional*): `mu` value used for `dynamic_shifting`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._skip_layer_guidance_scale = skip_layer_guidance_scale
+        self._clip_skip = clip_skip
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            device=device,
+            clip_skip=self.clip_skip,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        if self.do_classifier_free_guidance:
+            if skip_guidance_layers is not None:
+                original_prompt_embeds = prompt_embeds
+                original_pooled_prompt_embeds = pooled_prompt_embeds
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        scheduler_kwargs = {}
+        if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None:
+            _, _, height, width = latents.shape
+            image_seq_len = (height // self.transformer.config.patch_size) * (
+                width // self.transformer.config.patch_size
+            )
+            mu = calculate_shift(
+                image_seq_len,
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 4096),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.16),
+            )
+            scheduler_kwargs["mu"] = mu
+        elif mu is not None:
+            scheduler_kwargs["mu"] = mu
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            **scheduler_kwargs,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # 6. Prepare image embeddings
+        if (ip_adapter_image is not None and self.is_ip_adapter_active) or ip_adapter_image_embeds is not None:
+            ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+            if self.joint_attention_kwargs is None:
+                self._joint_attention_kwargs = {"ip_adapter_image_embeds": ip_adapter_image_embeds}
+            else:
+                self._joint_attention_kwargs.update(ip_adapter_image_embeds=ip_adapter_image_embeds)
+
+        if cond_latents is not None and self.do_classifier_free_guidance:
+            if cond_latents.shape[0] == latents.shape[0]:
+                cond_latents = torch.cat([cond_latents]*2)
+
+        # 7. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    cond_hidden_states=cond_latents,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    pooled_projections=pooled_prompt_embeds,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    should_skip_layers = (
+                        True
+                        if i > num_inference_steps * skip_layer_guidance_start
+                        and i < num_inference_steps * skip_layer_guidance_stop
+                        else False
+                    )
+                    if skip_guidance_layers is not None and should_skip_layers:
+                        timestep = t.expand(latents.shape[0])
+                        latent_model_input = latents
+                        noise_pred_skip_layers = self.transformer(
+                            hidden_states=latent_model_input,
+                            timestep=timestep,
+                            encoder_hidden_states=original_prompt_embeds,
+                            pooled_projections=original_pooled_prompt_embeds,
+                            joint_attention_kwargs=self.joint_attention_kwargs,
+                            return_dict=False,
+                            skip_layers=skip_guidance_layers,
+                        )[0]
+                        noise_pred = (
+                            noise_pred + (noise_pred_text - noise_pred_skip_layers) * self._skip_layer_guidance_scale
+                        )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                    negative_pooled_prompt_embeds = callback_outputs.pop(
+                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+                    )
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusion3PipelineOutput(images=image)

src/models/stable_diffusion3/pipeline_stable_diffusion_3_dynamic.py

@@ -0,0 +1,1257 @@
+# Copyright 2024 Stability AI, The HuggingFace Team and The InstantX Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    SiglipImageProcessor,
+    SiglipVisionModel,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import FromSingleFileMixin, SD3IPAdapterMixin, SD3LoraLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.models.transformers import SD3Transformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import StableDiffusion3Pipeline
+
+        >>> pipe = StableDiffusion3Pipeline.from_pretrained(
+        ...     "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
+        ... )
+        >>> pipe.to("cuda")
+        >>> prompt = "A cat holding a sign that says hello world"
+        >>> image = pipe(prompt).images[0]
+        >>> image.save("sd3.png")
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin):
+    r"""
+    Args:
+        transformer ([`SD3Transformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant,
+            with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size`
+            as its dimension.
+        text_encoder_2 ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the
+            [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
+            variant.
+        text_encoder_3 ([`T5EncoderModel`]):
+            Frozen text-encoder. Stable Diffusion 3 uses
+            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
+            [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`CLIPTokenizer`):
+            Second Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_3 (`T5TokenizerFast`):
+            Tokenizer of class
+            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
+        image_encoder (`SiglipVisionModel`, *optional*):
+            Pre-trained Vision Model for IP Adapter.
+        feature_extractor (`SiglipImageProcessor`, *optional*):
+            Image processor for IP Adapter.
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae"
+    _optional_components = ["image_encoder", "feature_extractor"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"]
+
+    def __init__(
+        self,
+        transformer: SD3Transformer2DModel,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer_2: CLIPTokenizer,
+        text_encoder_3: T5EncoderModel,
+        tokenizer_3: T5TokenizerFast,
+        image_encoder: SiglipVisionModel = None,
+        feature_extractor: SiglipImageProcessor = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            text_encoder_3=text_encoder_3,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            tokenizer_3=tokenizer_3,
+            transformer=transformer,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = (
+            self.transformer.config.sample_size
+            if hasattr(self, "transformer") and self.transformer is not None
+            else 128
+        )
+        self.patch_size = (
+            self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2
+        )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 256,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if self.text_encoder_3 is None:
+            return torch.zeros(
+                (
+                    batch_size * num_images_per_prompt,
+                    self.tokenizer_max_length,
+                    self.transformer.config.joint_attention_dim,
+                ),
+                device=device,
+                dtype=dtype,
+            )
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            # logger.warning(
+            #     "The following part of your input was truncated because `max_sequence_length` is set to "
+            #     f" {max_sequence_length} tokens: {removed_text}"
+            # )
+
+        prompt_embeds = self.text_encoder_3(text_input_ids.to(device))[0]
+
+        dtype = self.text_encoder_3.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        clip_skip: Optional[int] = None,
+        clip_model_index: int = 0,
+    ):
+        device = device or self._execution_device
+
+        clip_tokenizers = [self.tokenizer, self.tokenizer_2]
+        clip_text_encoders = [self.text_encoder, self.text_encoder_2]
+
+        tokenizer = clip_tokenizers[clip_model_index]
+        text_encoder = clip_text_encoders[clip_model_index]
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            # logger.warning(
+            #     "The following part of your input was truncated because CLIP can only handle sequences up to"
+            #     f" {self.tokenizer_max_length} tokens: {removed_text}"
+            # )
+        prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+        pooled_prompt_embeds = prompt_embeds[0]
+
+        if clip_skip is None:
+            prompt_embeds = prompt_embeds.hidden_states[-2]
+        else:
+            prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds, pooled_prompt_embeds
+
+    def encode_pooled_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        clip_skip: Optional[int] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, SD3LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            _, pooled_prompt_embed = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=0,
+            )
+            _, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                prompt=prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=1,
+            )
+
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            _, negative_pooled_prompt_embed = self._get_clip_prompt_embeds(
+                negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=0,
+            )
+            _, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                negative_prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=1,
+            )
+
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        return pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        prompt_3: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        clip_skip: Optional[int] = None,
+        max_sequence_length: int = 256,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            negative_prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, SD3LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            prompt_3 = prompt_3 or prompt
+            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+            prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=0,
+            )
+            prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                prompt=prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=1,
+            )
+            clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
+
+            t5_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=prompt_3,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+            clip_prompt_embeds = torch.nn.functional.pad(
+                clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1])
+            )
+
+            prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            negative_prompt_3 = negative_prompt_3 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+            negative_prompt_3 = (
+                batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+            )
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds(
+                negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=0,
+            )
+            negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                negative_prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=1,
+            )
+            negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1)
+
+            t5_negative_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=negative_prompt_3,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+            negative_clip_prompt_embeds = torch.nn.functional.pad(
+                negative_clip_prompt_embeds,
+                (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]),
+            )
+
+            negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2)
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, SD3LoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        prompt_3,
+        height,
+        width,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        negative_prompt_3=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if (
+            height % (self.vae_scale_factor * self.patch_size) != 0
+            or width % (self.vae_scale_factor * self.patch_size) != 0
+        ):
+            raise ValueError(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size} but are {height} and {width}."
+                f"You can use height {height - height % (self.vae_scale_factor * self.patch_size)} and width {width - width % (self.vae_scale_factor * self.patch_size)}."
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_3 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_3`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
+            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_3 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        shape = (
+            batch_size,
+            num_channels_latents,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def skip_guidance_layers(self):
+        return self._skip_guidance_layers
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    # Adapted from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_image
+    def encode_image(self, image: PipelineImageInput, device: torch.device) -> torch.Tensor:
+        """Encodes the given image into a feature representation using a pre-trained image encoder.
+
+        Args:
+            image (`PipelineImageInput`):
+                Input image to be encoded.
+            device: (`torch.device`):
+                Torch device.
+
+        Returns:
+            `torch.Tensor`: The encoded image feature representation.
+        """
+        if not isinstance(image, torch.Tensor):
+            image = self.feature_extractor(image, return_tensors="pt").pixel_values
+
+        image = image.to(device=device, dtype=self.dtype)
+
+        return self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+
+    # Adapted from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.prepare_ip_adapter_image_embeds
+    def prepare_ip_adapter_image_embeds(
+        self,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+    ) -> torch.Tensor:
+        """Prepares image embeddings for use in the IP-Adapter.
+
+        Either `ip_adapter_image` or `ip_adapter_image_embeds` must be passed.
+
+        Args:
+            ip_adapter_image (`PipelineImageInput`, *optional*):
+                The input image to extract features from for IP-Adapter.
+            ip_adapter_image_embeds (`torch.Tensor`, *optional*):
+                Precomputed image embeddings.
+            device: (`torch.device`, *optional*):
+                Torch device.
+            num_images_per_prompt (`int`, defaults to 1):
+                Number of images that should be generated per prompt.
+            do_classifier_free_guidance (`bool`, defaults to True):
+                Whether to use classifier free guidance or not.
+        """
+        device = device or self._execution_device
+
+        if ip_adapter_image_embeds is not None:
+            if do_classifier_free_guidance:
+                single_negative_image_embeds, single_image_embeds = ip_adapter_image_embeds.chunk(2)
+            else:
+                single_image_embeds = ip_adapter_image_embeds
+        elif ip_adapter_image is not None:
+            single_image_embeds = self.encode_image(ip_adapter_image, device)
+            if do_classifier_free_guidance:
+                single_negative_image_embeds = torch.zeros_like(single_image_embeds)
+        else:
+            raise ValueError("Neither `ip_adapter_image_embeds` or `ip_adapter_image_embeds` were provided.")
+
+        image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
+
+        if do_classifier_free_guidance:
+            negative_image_embeds = torch.cat([single_negative_image_embeds] * num_images_per_prompt, dim=0)
+            image_embeds = torch.cat([negative_image_embeds, image_embeds], dim=0)
+
+        return image_embeds.to(device=device)
+
+    def enable_sequential_cpu_offload(self, *args, **kwargs):
+        if self.image_encoder is not None and "image_encoder" not in self._exclude_from_cpu_offload:
+            logger.warning(
+                "`pipe.enable_sequential_cpu_offload()` might fail for `image_encoder` if it uses "
+                "`torch.nn.MultiheadAttention`. You can exclude `image_encoder` from CPU offloading by calling "
+                "`pipe._exclude_from_cpu_offload.append('image_encoder')` before `pipe.enable_sequential_cpu_offload()`."
+            )
+
+        super().enable_sequential_cpu_offload(*args, **kwargs)
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_3: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 7.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        cond_latents: Optional[list[torch.FloatTensor]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 256,
+        skip_guidance_layers: List[int] = None,
+        skip_layer_guidance_scale: float = 2.8,
+        skip_layer_guidance_stop: float = 0.2,
+        skip_layer_guidance_start: float = 0.01,
+        mu: Optional[float] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used instead
+            negative_prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used instead
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            ip_adapter_image (`PipelineImageInput`, *optional*):
+                Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`torch.Tensor`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. Should be a tensor of shape `(batch_size, num_images,
+                emb_dim)`. It should contain the negative image embedding if `do_classifier_free_guidance` is set to
+                `True`. If not provided, embeddings are computed from the `ip_adapter_image` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] instead of
+                a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`.
+            skip_guidance_layers (`List[int]`, *optional*):
+                A list of integers that specify layers to skip during guidance. If not provided, all layers will be
+                used for guidance. If provided, the guidance will only be applied to the layers specified in the list.
+                Recommended value by StabiltyAI for Stable Diffusion 3.5 Medium is [7, 8, 9].
+            skip_layer_guidance_scale (`int`, *optional*): The scale of the guidance for the layers specified in
+                `skip_guidance_layers`. The guidance will be applied to the layers specified in `skip_guidance_layers`
+                with a scale of `skip_layer_guidance_scale`. The guidance will be applied to the rest of the layers
+                with a scale of `1`.
+            skip_layer_guidance_stop (`int`, *optional*): The step at which the guidance for the layers specified in
+                `skip_guidance_layers` will stop. The guidance will be applied to the layers specified in
+                `skip_guidance_layers` until the fraction specified in `skip_layer_guidance_stop`. Recommended value by
+                StabiltyAI for Stable Diffusion 3.5 Medium is 0.2.
+            skip_layer_guidance_start (`int`, *optional*): The step at which the guidance for the layers specified in
+                `skip_guidance_layers` will start. The guidance will be applied to the layers specified in
+                `skip_guidance_layers` from the fraction specified in `skip_layer_guidance_start`. Recommended value by
+                StabiltyAI for Stable Diffusion 3.5 Medium is 0.01.
+            mu (`float`, *optional*): `mu` value used for `dynamic_shifting`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_3.StableDiffusion3PipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._skip_layer_guidance_scale = skip_layer_guidance_scale
+        self._clip_skip = clip_skip
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            device=device,
+            clip_skip=self.clip_skip,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        if self.do_classifier_free_guidance:
+            if skip_guidance_layers is not None:
+                original_prompt_embeds = prompt_embeds
+                original_pooled_prompt_embeds = pooled_prompt_embeds
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        scheduler_kwargs = {}
+        if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None:
+            _, _, height, width = latents.shape
+            image_seq_len = (height // self.transformer.config.patch_size) * (
+                width // self.transformer.config.patch_size
+            )
+            mu = calculate_shift(
+                image_seq_len,
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 4096),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.16),
+            )
+            scheduler_kwargs["mu"] = mu
+        elif mu is not None:
+            scheduler_kwargs["mu"] = mu
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            **scheduler_kwargs,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # 6. Prepare image embeddings
+        if (ip_adapter_image is not None and self.is_ip_adapter_active) or ip_adapter_image_embeds is not None:
+            ip_adapter_image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+            if self.joint_attention_kwargs is None:
+                self._joint_attention_kwargs = {"ip_adapter_image_embeds": ip_adapter_image_embeds}
+            else:
+                self._joint_attention_kwargs.update(ip_adapter_image_embeds=ip_adapter_image_embeds)
+
+
+        if cond_latents is not None and self.do_classifier_free_guidance:
+            if len(cond_latents) == latents.shape[0]:
+                cond_latents = cond_latents * 2
+
+        # 7. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    cond_hidden_states=cond_latents,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    pooled_projections=pooled_prompt_embeds,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    should_skip_layers = (
+                        True
+                        if i > num_inference_steps * skip_layer_guidance_start
+                        and i < num_inference_steps * skip_layer_guidance_stop
+                        else False
+                    )
+                    if skip_guidance_layers is not None and should_skip_layers:
+                        timestep = t.expand(latents.shape[0])
+                        latent_model_input = latents
+                        noise_pred_skip_layers = self.transformer(
+                            hidden_states=latent_model_input,
+                            timestep=timestep,
+                            encoder_hidden_states=original_prompt_embeds,
+                            pooled_projections=original_pooled_prompt_embeds,
+                            joint_attention_kwargs=self.joint_attention_kwargs,
+                            return_dict=False,
+                            skip_layers=skip_guidance_layers,
+                        )[0]
+                        noise_pred = (
+                            noise_pred + (noise_pred_text - noise_pred_skip_layers) * self._skip_layer_guidance_scale
+                        )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                    negative_pooled_prompt_embeds = callback_outputs.pop(
+                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+                    )
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusion3PipelineOutput(images=image)