Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/Puffin.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="$USER_HOME$/envs/pt2.7" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="$USER_HOME$/envs/pt2.7" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/Puffin.iml" filepath="$PROJECT_DIR$/.idea/Puffin.iml" />
+    </modules>
+  </component>
+</project>

.idea/workspace.xml

@@ -0,0 +1,50 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ChangeListManager">
+    <list default="true" id="9dd87dac-8a5e-4178-a1d7-afa664ac2f6a" name="Changes" comment="" />
+    <option name="SHOW_DIALOG" value="false" />
+    <option name="HIGHLIGHT_CONFLICTS" value="true" />
+    <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
+    <option name="LAST_RESOLUTION" value="IGNORE" />
+  </component>
+  <component name="ProjectColorInfo"><![CDATA[{
+  "associatedIndex": 6
+}]]></component>
+  <component name="ProjectId" id="33qty5NdkHzw3ffLzNseYxoODo4" />
+  <component name="ProjectViewState">
+    <option name="hideEmptyMiddlePackages" value="true" />
+    <option name="showLibraryContents" value="true" />
+  </component>
+  <component name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "ModuleVcsDetector.initialDetectionPerformed": "true",
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "last_opened_file_path": "/Users/wusize/projects/Puffin",
+    "nodejs_package_manager_path": "npm",
+    "settings.editor.selected.configurable": "com.jetbrains.python.configuration.PyActiveSdkModuleConfigurable",
+    "vue.rearranger.settings.migration": "true"
+  }
+}]]></component>
+  <component name="SharedIndexes">
+    <attachedChunks>
+      <set>
+        <option value="bundled-js-predefined-d6986cc7102b-6a121458b545-JavaScript-PY-251.25410.159" />
+        <option value="bundled-python-sdk-e0ed3721d81e-36ea0e71a18c-com.jetbrains.pycharm.pro.sharedIndexes.bundled-PY-251.25410.159" />
+      </set>
+    </attachedChunks>
+  </component>
+  <component name="TaskManager">
+    <task active="true" id="Default" summary="Default task">
+      <changelist id="9dd87dac-8a5e-4178-a1d7-afa664ac2f6a" name="Changes" comment="" />
+      <created>1760056680813</created>
+      <option name="number" value="Default" />
+      <option name="presentableId" value="Default" />
+      <updated>1760056680813</updated>
+      <workItem from="1760056681869" duration="11000" />
+    </task>
+    <servers />
+  </component>
+  <component name="TypeScriptGeneratedFilesManager">
+    <option name="version" value="3" />
+  </component>
+</project>

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.

README.md

@@ -1,10 +1,10 @@
 ---
 title: Puffin
-emoji: 📚
-colorFrom: gray
-colorTo: red
+emoji: 👀
+colorFrom: red
+colorTo: purple
 sdk: gradio
-sdk_version: 5.49.1
+sdk_version: 5.23.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,248 @@
+import gradio as gr
+import torch
+import io
+from PIL import Image
+import numpy as np
+import spaces  # Import spaces for ZeroGPU compatibility
+import math
+import re
+from einops import rearrange
+from mmengine.config import Config
+from xtuner.registry import BUILDER 
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from scripts.camera.cam_dataset import Cam_Generator
+from scripts.camera.visualization.visualize_batch import make_perspective_figures
+
+from huggingface_hub import snapshot_download
+import os
+local_path = snapshot_download(
+    repo_id="KangLiao/Puffin",
+    repo_type="model",       
+    #filename="Puffin-Base.pth",         
+    local_dir="checkpoints/",  
+    local_dir_use_symlinks=False,     
+    revision="main",                  
+)
+
+
+NUM = r"[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?"
+CAM_PATTERN = re.compile(r"(?:camera parameters.*?:|roll.*?:)\s*("+NUM+r")\s*,\s*("+NUM+r")\s*,\s*("+NUM+r")", re.IGNORECASE|re.DOTALL)
+
+def center_crop(image):
+    w, h = image.size
+    s = min(w, h)
+    l = (w - s) // 2
+    t = (h - s) // 2
+    return image.crop((l, t, l + s, t + s))
+
+
+##### load model
+config = "configs/pipelines/stage_2_base.py"
+config = Config.fromfile(config)
+model = BUILDER.build(config.model).cuda().bfloat16().eval()
+checkpoint_path = "checkpoints/Puffin-Base.pth"
+checkpoint = torch.load(checkpoint_path)
+info = model.load_state_dict(checkpoint, strict=False)
+
+def fig_to_image(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+    buf.seek(0)
+    img = Image.open(buf).convert('RGB')
+    buf.close()
+    return img
+
+def extract_up_lat_figs(fig_dict):
+    fig_up, fig_lat = None, None
+    others = {}
+    for k, fig in fig_dict.items():
+        if ("up_field" in k) and (fig_up is None):
+            fig_up = fig
+        elif ("latitude_field" in k) and (fig_lat is None):
+            fig_lat = fig
+        else:
+            others[k] = fig
+    return fig_up, fig_lat, others
+
+
+@torch.inference_mode()
+@spaces.GPU(duration=120) 
+# Multimodal Understanding function
+def camera_understanding(image_src, question, seed, 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())
+
+    prompt = ("Describe the image in detail. Then reason its spatial distribution and estimate its camera parameters (roll, pitch, and field-of-view).")
+
+    image = Image.fromarray(image_src).convert('RGB')
+    image = center_crop(image)
+    image = image.resize((512, 512))
+    x = torch.from_numpy(np.array(image)).float()
+    x = x / 255.0
+    x = 2 * x - 1
+    x = rearrange(x, 'h w c -> c h w')
+
+    with torch.no_grad():
+        outputs = model.understand(prompt=[prompt], pixel_values=[x], progress_bar=False)
+
+    text = outputs[0]
+    
+    gen = Cam_Generator(mode="base")
+    cam = gen.get_cam(text)
+    
+    bgr = np.array(image)[:, :, ::-1].astype(np.float32) / 255.0
+    rgb = bgr[:, :, ::-1].copy()
+    image_tensor = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
+    single_batch = {}
+    single_batch["image"] = image_tensor
+    single_batch["up_field"] = cam[:2].unsqueeze(0)
+    single_batch["latitude_field"] = cam[2:].unsqueeze(0)
+
+    figs = make_perspective_figures(single_batch, single_batch, n_pairs=1)
+    up_img = lat_img = None
+    for k, fig in figs.items():
+        if "up_field" in k:
+            up_img = fig_to_image(fig)
+        elif "latitude_field" in k:
+            lat_img = fig_to_image(fig)
+        plt.close(fig)
+
+    return text#, up_img, lat_img
+
+
+@torch.inference_mode()
+@spaces.GPU(duration=120)  # Specify a duration to avoid timeout
+def generate_image(prompt_scene,
+                   seed=42,
+                   roll=0.1,
+                   pitch=0.1,
+                   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
+    print("prompt:", prompt)
+    
+    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=42)
+            
+        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("Camera Understanding"):
+        gr.Markdown(value="## Camera Understanding")
+        image_input = gr.Image()
+
+        understanding_button = gr.Button("Chat")
+        understanding_output = gr.Textbox(label="Response")
+        
+        #camera1 = gr.Gallery(label="Camera Maps", columns=1, rows=1)
+        #camera2 = gr.Gallery(label="Camera Maps", columns=1, rows=1)
+
+        with gr.Accordion("Advanced options", open=False):
+            und_seed_input = gr.Number(label="Seed", precision=0, value=42)
+
+        examples_inpainting = gr.Examples(
+            label="Camera Understanding examples",
+            examples=[
+                "assets/1.jpg",
+                "assets/2.jpg",
+                "assets/3.jpg",
+                "assets/4.jpg",
+                "assets/5.jpg",
+                "assets/6.jpg",
+            ],
+            inputs=image_input,
+        )
+
+    generation_button.click(
+        fn=generate_image,
+        inputs=[prompt_input, seed_input, roll, pitch, fov],
+        outputs=image_output
+    )
+
+    understanding_button.click(
+        camera_understanding,
+        inputs=[image_input, und_seed_input],
+        outputs=[understanding_output]#, camera1, camera2]
+    )
+
+demo.launch(share=True)

configs/models/qwen2_5_1_5b_radio_sd3_dynamic_puffin.py

@@ -0,0 +1,87 @@
+import torch
+import os
+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 = "configs/sd3"
+radiov3_model_name_or_path = "configs/radiov3"
+
+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_config,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="transformer",
+                 torch_dtype=torch.bfloat16,
+                 #local_files_only=True,
+                 ),
+             test_scheduler=dict(
+                 type=FlowMatchEulerDiscreteScheduler.from_config,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="scheduler",
+                 #local_files_only=True,
+                 ),
+             train_scheduler=dict(
+                 type=FlowMatchEulerDiscreteScheduler.from_config,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="scheduler",
+                 #local_files_only=True,
+                 ),
+             vae=dict(
+                 type=AutoencoderKL.from_config,
+                 pretrained_model_name_or_path=sd3_model_name_or_path,
+                 subfolder="vae",
+                 torch_dtype=torch.bfloat16,
+                 #local_files_only=True,
+                 ),
+             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

configs/qwen2.5/config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sliding_window": 32768,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.43.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

configs/qwen2.5/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "pad_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_p": 0.8,
+  "top_k": 20,
+  "transformers_version": "4.37.0"
+}

configs/qwen2.5/tokenizer_config.json

@@ -0,0 +1,207 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

configs/radio3/config.json

@@ -0,0 +1,241 @@
+{
+  "adaptor_configs": {},
+  "adaptor_names": null,
+  "architectures": [
+    "RADIOModel"
+  ],
+  "args": {
+    "aa": null,
+    "amp": true,
+    "amp_dtype": "bfloat16",
+    "amp_impl": "native",
+    "aug_repeats": 0,
+    "aug_splits": 0,
+    "bn_eps": null,
+    "bn_momentum": null,
+    "cache_dir": null,
+    "channels_last": false,
+    "checkpoint_hist": 10,
+    "chk_keep_forever": 100,
+    "class_map": "",
+    "clip_grad": null,
+    "clip_mode": "norm",
+    "cls_token_per_teacher": true,
+    "coco_annotations_file": "/datasets/coco2017-adlsa/annotations/captions_val2017.json",
+    "coco_image_dir": "/datasets/coco2017-adlsa/val2017",
+    "color_jitter": 0.4,
+    "cooldown_epochs": 0,
+    "cpe_max_size": 2048,
+    "cpe_num_registers": 4,
+    "crd_loss": false,
+    "crd_loss_weight": 0.8,
+    "crop_pct": null,
+    "cutmix": 0.0,
+    "cutmix_minmax": null,
+    "dataset_download": false,
+    "debug_full_knn": false,
+    "decay_epochs": 90,
+    "decay_milestones": [
+      90,
+      180,
+      270
+    ],
+    "decay_rate": 0.1,
+    "depchain": true,
+    "detect_anomaly": false,
+    "dist_bn": "reduce",
+    "dist_norm_weight": 0.0,
+    "distributed": true,
+    "drop": 0.0,
+    "drop_block": null,
+    "drop_connect": null,
+    "drop_path": null,
+    "dtype": "float32",
+    "epoch_repeats": 0.0,
+    "eval": false,
+    "eval_metric": "knn_top1",
+    "eval_teacher": false,
+    "eval_teacher_only": false,
+    "eval_throughput": false,
+    "fast_norm": false,
+    "fd_loss_fn": "MSE",
+    "feature_normalization": "PHI_STANDARDIZE",
+    "feature_summarizer": "cls_token",
+    "feature_upscale_factor": null,
+    "force_new_wandb_id": false,
+    "force_spectral_reparam": false,
+    "freeze_bn": false,
+    "fsdp": true,
+    "full_equivariance": false,
+    "fuser": "",
+    "gp": null,
+    "grad_accum_steps": 1,
+    "grad_checkpointing": false,
+    "head_init_bias": null,
+    "head_init_scale": null,
+    "head_lr": null,
+    "head_warmup": 5,
+    "head_weight_decay": 0.01,
+    "hflip": 0.5,
+    "img_size": null,
+    "in_chans": null,
+    "initial_checkpoint": null,
+    "input_size": null,
+    "interpolation": "",
+    "layer_decay": null,
+    "local_rank": 0,
+    "log_interval": 50,
+    "log_mlflow": false,
+    "log_wandb": true,
+    "loss_auto_balance": false,
+    "lr_base": 0.1,
+    "lr_base_scale": "",
+    "lr_base_size": 256,
+    "lr_cycle_decay": 0.5,
+    "lr_cycle_limit": 1,
+    "lr_cycle_mul": 1.0,
+    "lr_k_decay": 1.0,
+    "lr_noise": null,
+    "lr_noise_pct": 0.67,
+    "lr_noise_std": 1.0,
+    "mean": null,
+    "mesa": false,
+    "min_lr": 0.0001,
+    "mixup": 0.0,
+    "mixup_mode": "batch",
+    "mixup_off_epoch": 0,
+    "mixup_prob": 1.0,
+    "mixup_switch_prob": 0.5,
+    "mlp_hidden_size": 2560,
+    "mlp_num_inner": 1,
+    "mlp_version": "v2",
+    "model": "vit_huge_patch16_224",
+    "model_kwargs": {},
+    "model_norm": false,
+    "momentum": 0.9,
+    "no_aug": false,
+    "no_custom_validation": false,
+    "no_ddp_bb": true,
+    "no_knn": false,
+    "no_prefetcher": false,
+    "no_resume_opt": false,
+    "num_classes": null,
+    "one_logger_app_tag": "",
+    "one_logger_is_baseline": false,
+    "one_logger_run_name": "",
+    "onelogger": null,
+    "opt_betas": null,
+    "opt_eps": null,
+    "patience_epochs": 10,
+    "pin_mem": false,
+    "prefetcher": true,
+    "pretrained": false,
+    "rank": 0,
+    "ratio": [
+      0.75,
+      1.3333333333333333
+    ],
+    "recount": 1,
+    "recovery_interval": 0,
+    "register_multiple": 0,
+    "remode": "pixel",
+    "reprob": 0.0,
+    "reset_loss_state": true,
+    "resplit": false,
+    "sample_tracking": false,
+    "save_images": false,
+    "scale": [
+      0.5,
+      1.0
+    ],
+    "sched": "cosine",
+    "seed": 42,
+    "shift_equivariance": true,
+    "smoothing": 0.1,
+    "spectral_heads": false,
+    "spectral_reparam": false,
+    "spectral_weight_decay": null,
+    "split_bn": false,
+    "start_epoch": null,
+    "std": null,
+    "stream_teachers": true,
+    "sync_bn": false,
+    "synchronize_step": false,
+    "teachers": [
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 378,
+        "model": "ViT-H-14-378-quickgelu",
+        "name": "clip",
+        "pretrained": "dfn5b",
+        "type": "open_clip",
+        "use_summary": true
+      },
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 384,
+        "model": "siglip2-g-384",
+        "name": "siglip2-g",
+        "type": "siglip2",
+        "use_summary": true
+      },
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 224,
+        "model": "dinov2_vitg14_reg",
+        "name": "dino_v2",
+        "type": "dino_v2",
+        "use_summary": true
+      },
+      {
+        "fd_normalize": false,
+        "feature_distillation": true,
+        "input_size": 1024,
+        "model": "vit-h",
+        "name": "sam",
+        "type": "sam",
+        "use_summary": false
+      }
+    ],
+    "torchcompile": null,
+    "torchscript": false,
+    "train_interpolation": "random",
+    "train_split": "train",
+    "tta": 0,
+    "use_coco": false,
+    "use_multi_epochs_loader": false,
+    "val_ema_only": false,
+    "val_split": "val",
+    "vflip": 0.0,
+    "vitdet_version": 1,
+    "wandb_entity": "",
+    "wandb_id": "",
+    "wandb_job_type": "",
+    "wandb_name": "",
+    "wandb_project": "",
+    "warmup_lr": 1e-05,
+    "warmup_prefix": false,
+    "worker_seeding": "all",
+    "workers": 8,
+    "world_size": 256
+  },
+  "auto_map": {
+    "AutoConfig": "hf_model.RADIOConfig",
+    "AutoModel": "hf_model.RADIOModel"
+  },
+  "feature_normalizer_config": null,
+  "inter_feature_normalizer_config": null,
+  "max_resolution": 2048,
+  "patch_size": 16,
+  "preferred_resolution": [
+    512,
+    512
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "version": "c-radio_v3-h",
+  "vitdet_window_size": null
+}

configs/sd3/scheduler/scheduler_config.json

@@ -0,0 +1,6 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.29.0.dev0",
+  "num_train_timesteps": 1000,
+  "shift": 3.0
+}

configs/sd3/transformer/config.json

@@ -0,0 +1,15 @@
+{
+  "_class_name": "SD3Transformer2DModel",
+  "_diffusers_version": "0.29.0.dev0",
+  "attention_head_dim": 64,
+  "caption_projection_dim": 1536,
+  "in_channels": 16,
+  "joint_attention_dim": 4096,
+  "num_attention_heads": 24,
+  "num_layers": 24,
+  "out_channels": 16,
+  "patch_size": 2,
+  "pooled_projection_dim": 2048,
+  "pos_embed_max_size": 192,
+  "sample_size": 128
+}

configs/sd3/vae/config.json

@@ -0,0 +1,36 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.29.0.dev0",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 16,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 1.5305,
+  "shift_factor": 0.0609,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": false,
+  "use_quant_conv": false
+}

requirements.txt

@@ -0,0 +1,17 @@
+torch
+accelerate
+diffusers==0.34.0
+gradio
+torchvision
+safetensors
+matplotlib==3.10.1
+matplotlib-inline==0.1.7
+mmengine==0.10.7
+numpy==2.2.5
+pillow==11.2.1
+scipy==1.15.2
+timm==0.9.12
+transformers==4.49.0
+xtuner==0.1.23
+deepspeed
+

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