add app.py

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Tripo AI & Stability AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,80 @@
----
-title: Perflow Triposr
-emoji: 📊
-colorFrom: blue
-colorTo: gray
-sdk: gradio
-sdk_version: 4.20.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TripoSR <a href="https://huggingface.co/stabilityai/TripoSR"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model_Card-Huggingface-orange"></a> <a href="https://huggingface.co/spaces/stabilityai/TripoSR"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Gradio%20Demo-Huggingface-orange"></a> <a href="https://arxiv.org/abs/2403.02151"><img src="https://img.shields.io/badge/Arxiv-2403.02151-B31B1B.svg"></a>
+
+<div align="center">
+  <img src="figures/teaser800.gif" alt="Teaser Video">
+</div>
+
+This is the official codebase for **TripoSR**, a state-of-the-art open-source model for **fast** feedforward 3D reconstruction from a single image, collaboratively developed by [Tripo AI](https://www.tripo3d.ai/) and [Stability AI](https://stability.ai/).
+<br><br>
+Leveraging the principles of the [Large Reconstruction Model (LRM)](https://yiconghong.me/LRM/), TripoSR brings to the table key advancements that significantly boost both the speed and quality of 3D reconstruction. Our model is distinguished by its ability to rapidly process inputs, generating high-quality 3D models in less than 0.5 seconds on an NVIDIA A100 GPU. TripoSR has exhibited superior performance in both qualitative and quantitative evaluations, outperforming other open-source alternatives across multiple public datasets. The figures below illustrate visual comparisons and metrics showcasing TripoSR's performance relative to other leading models. Details about the model architecture, training process, and comparisons can be found in this [technical report](https://arxiv.org/abs/2403.02151).
+
+<!--
+<div align="center">
+  <img src="figures/comparison800.gif" alt="Teaser Video">
+</div>
+-->
+<p align="center">
+    <img width="800" src="figures/visual_comparisons.jpg"/>
+</p>
+
+<p align="center">
+    <img width="450" src="figures/scatter-comparison.png"/>
+</p>
+
+
+The model is released under the MIT license, which includes the source code, pretrained models, and an interactive online demo. Our goal is to empower researchers, developers, and creatives to push the boundaries of what's possible in 3D generative AI and 3D content creation.
+
+## Getting Started
+### Installation
+- Python >= 3.8
+- Install CUDA if available
+- Install PyTorch according to your platform: [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/) **[Please make sure that the locally-installed CUDA major version matches the PyTorch-shipped CUDA major version. For example if you have CUDA 11.x installed, make sure to install PyTorch compiled with CUDA 11.x.]**
+- Update setuptools by `pip install --upgrade setuptools`
+- Install other dependencies by `pip install -r requirements.txt`
+
+### Manual Inference
+```sh
+python run.py examples/chair.png --output-dir output/
+```
+This will save the reconstructed 3D model to `output/`. You can also specify more than one image path separated by spaces. The default options takes about **6GB VRAM** for a single image input.
+
+For detailed usage of this script, use `python run.py --help`.
+
+### Local Gradio App
+Install Gradio:
+```sh
+pip install gradio
+```
+Start the Gradio App:
+```sh
+python gradio_app.py
+```
+
+## Troubleshooting
+> AttributeError: module 'torchmcubes_module' has no attribute 'mcubes_cuda'
+
+or
+
+> torchmcubes was not compiled with CUDA support, use CPU version instead.
+
+This is because `torchmcubes` is compiled without CUDA support. Please make sure that 
+
+- The locally-installed CUDA major version matches the PyTorch-shipped CUDA major version. For example if you have CUDA 11.x installed, make sure to install PyTorch compiled with CUDA 11.x.
+- `setuptools>=49.6.0`. If not, upgrade by `pip install --upgrade setuptools`.
+
+Then re-install `torchmcubes` by:
+
+```sh
+pip uninstall torchmcubes
+pip install git+https://github.com/tatsy/torchmcubes.git
+```
+
+## Citation
+```BibTeX
+@article{TripoSR2024,
+  title={TripoSR: Fast 3D Object Reconstruction from a Single Image},
+  author={Tochilkin, Dmitry and Pankratz, David and Liu, Zexiang and Huang, Zixuan and and Letts, Adam and Li, Yangguang and Liang, Ding and Laforte, Christian and Jampani, Varun and Cao, Yan-Pei},
+  journal={arXiv preprint arXiv:2403.02151},
+  year={2024}
+}
+```

app.py

@@ -0,0 +1,182 @@
+import os, logging, time, argparse, random, tempfile, rembg
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from functools import partial
+from tsr.system import TSR
+from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
+
+from src.scheduler_perflow import PeRFlowScheduler
+from diffusers import StableDiffusionPipeline, UNet2DConditionModel
+
+def merge_delta_weights_into_unet(pipe, delta_weights, org_alpha = 1.0):
+    unet_weights = pipe.unet.state_dict()
+    for key in delta_weights.keys():
+        dtype = unet_weights[key].dtype
+        try:
+            unet_weights[key] = org_alpha * unet_weights[key].to(dtype=delta_weights[key].dtype) + delta_weights[key].to(device=unet_weights[key].device)
+        except:
+            unet_weights[key] = unet_weights[key].to(dtype=delta_weights[key].dtype)
+        unet_weights[key] = unet_weights[key].to(dtype)
+    pipe.unet.load_state_dict(unet_weights, strict=True)
+    return pipe
+
+def setup_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    
+if torch.cuda.is_available():
+    device = "cuda:0"
+else:
+    device = "cpu"
+
+### TripoSR
+model = TSR.from_pretrained(
+    "stabilityai/TripoSR",
+    config_name="config.yaml",
+    weight_name="model.ckpt",
+)
+# adjust the chunk size to balance between speed and memory usage
+model.renderer.set_chunk_size(8192)
+model.to(device)
+
+
+### PeRFlow-T2I
+# pipe_t2i = StableDiffusionPipeline.from_pretrained("Lykon/dreamshaper-8", torch_dtype=torch.float16, safety_checker=None)
+pipe_t2i = StableDiffusionPipeline.from_pretrained("stablediffusionapi/disney-pixar-cartoon", torch_dtype=torch.float16, safety_checker=None)
+delta_weights = UNet2DConditionModel.from_pretrained("hansyan/piecewise-rectified-flow-delta-weights", torch_dtype=torch.float16, variant="v0-1",).state_dict()
+pipe_t2i = merge_delta_weights_into_unet(pipe_t2i, delta_weights)
+pipe_t2i.scheduler = PeRFlowScheduler.from_config(pipe_t2i.scheduler.config, prediction_type="epsilon", num_time_windows=4)
+pipe_t2i.to('cuda:0', torch.float16)
+
+
+### gradio
+rembg_session = rembg.new_session()
+
+def generate(text, seed):
+    def fill_background(image):
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+        image = Image.fromarray((image * 255.0).astype(np.uint8))
+        return image
+
+    setup_seed(int(seed))
+    # text = text
+    samples = pipe_t2i(
+            prompt              = [text],
+            negative_prompt     = ["distorted, blur, low-quality, haze, out of focus"],
+            height              = 512,
+            width               = 512,
+            # num_inference_steps = 4,
+            # guidance_scale      = 4.5,
+            num_inference_steps = 6,
+            guidance_scale      = 7,
+            output_type         = 'pt',
+        ).images
+    samples = torch.nn.functional.interpolate(samples, size=768, mode='bilinear')
+    samples = samples.squeeze(0).permute(1, 2, 0).cpu().numpy()*255.
+    samples = samples.astype(np.uint8)
+    samples = Image.fromarray(samples[:, :, :3])
+
+    image = remove_background(samples, rembg_session)
+    image = resize_foreground(image, 0.85)
+    image = fill_background(image)
+    return image
+
+def render(image, mc_resolution=256, formats=["obj"]):
+    scene_codes = model(image, device=device)
+    mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
+    mesh = to_gradio_3d_orientation(mesh)
+    rv = []
+    for format in formats:
+        mesh_path = tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False)
+        mesh.export(mesh_path.name)
+        rv.append(mesh_path.name)
+    return rv[0]
+
+# warm up
+_ = generate("a bird", 42)
+
+# layout
+css = """
+h1 {
+    text-align: center;
+    display:block;
+}
+h2 {
+    text-align: center;
+    display:block;
+}
+h3 {
+    text-align: center;
+    display:block;
+}
+"""
+with gr.Blocks(title="TripoSR", css=css) as interface:
+    gr.Markdown(
+    """
+    # Instant Text-to-3D Mesh Demo
+
+    ### [PeRFlow](https://github.com/magic-research/piecewise-rectified-flow)-T2I  +  [TripoSR](https://github.com/VAST-AI-Research/TripoSR)
+    
+    Two-stage synthesis: 1) generating images by PeRFlow-T2I with 6-step inference; 2) rendering 3D assests.
+    """
+    )
+    
+    with gr.Column():
+        with gr.Row():
+                output_image = gr.Image(label='Generated Image', height=384, width=384)
+
+                output_model_obj = gr.Model3D(
+                    label="Output 3D Model (OBJ Format)",
+                    interactive=False,
+                    height=384, width=384,
+            )
+    
+    with gr.Row():
+        textbox = gr.Textbox(label="Input Prompt", value="a colorful bird")
+        seed = gr.Textbox(label="Random Seed", value=42)
+    
+    # activate
+    textbox.submit(
+        fn=generate,
+        inputs=[textbox, seed],
+        outputs=[output_image],
+    ).success(
+        fn=render,
+        inputs=[output_image],
+        outputs=[output_model_obj],
+    )
+    
+    seed.submit(
+        fn=generate,
+        inputs=[textbox, seed],
+        outputs=[output_image],
+    ).success(
+        fn=render,
+        inputs=[output_image],
+        outputs=[output_model_obj],
+    )
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--username', type=str, default=None, help='Username for authentication')
+    parser.add_argument('--password', type=str, default=None, help='Password for authentication')
+    parser.add_argument('--port', type=int, default=7860, help='Port to run the server listener on')
+    parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
+    parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site")
+    parser.add_argument("--queuesize", type=int, default=1, help="launch gradio queue max_size")
+    args = parser.parse_args()
+    interface.queue(max_size=args.queuesize)
+    interface.launch(
+        auth=(args.username, args.password) if (args.username and args.password) else None,
+        share=args.share,
+        server_name="0.0.0.0" if args.listen else None, 
+        server_port=args.port
+    )

gradio_app.py

@@ -0,0 +1,188 @@
+import logging
+import os
+import tempfile
+import time
+
+import gradio as gr
+import numpy as np
+import rembg
+import torch
+from PIL import Image
+from functools import partial
+
+from tsr.system import TSR
+from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
+
+import argparse
+
+
+if torch.cuda.is_available():
+    device = "cuda:0"
+else:
+    device = "cpu"
+
+model = TSR.from_pretrained(
+    "stabilityai/TripoSR",
+    config_name="config.yaml",
+    weight_name="model.ckpt",
+)
+
+# adjust the chunk size to balance between speed and memory usage
+model.renderer.set_chunk_size(8192)
+model.to(device)
+
+rembg_session = rembg.new_session()
+
+
+def check_input_image(input_image):
+    if input_image is None:
+        raise gr.Error("No image uploaded!")
+
+
+def preprocess(input_image, do_remove_background, foreground_ratio):
+    def fill_background(image):
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+        image = Image.fromarray((image * 255.0).astype(np.uint8))
+        return image
+
+    if do_remove_background:
+        image = input_image.convert("RGB")
+        image = remove_background(image, rembg_session)
+        image = resize_foreground(image, foreground_ratio)
+        image = fill_background(image)
+    else:
+        image = input_image
+        if image.mode == "RGBA":
+            image = fill_background(image)
+    return image
+
+
+def generate(image, mc_resolution, formats=["obj", "glb"]):
+    print(image.shape, image.min(), image.max())
+    scene_codes = model(image, device=device)
+    mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
+    mesh = to_gradio_3d_orientation(mesh)
+    rv = []
+    for format in formats:
+        mesh_path = tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False)
+        mesh.export(mesh_path.name)
+        rv.append(mesh_path.name)
+    return rv
+
+
+def run_example(image_pil):
+    preprocessed = preprocess(image_pil, False, 0.9)
+    mesh_name_obj, mesh_name_glb = generate(preprocessed, 256, ["obj", "glb"])
+    return preprocessed, mesh_name_obj, mesh_name_glb
+
+
+with gr.Blocks(title="TripoSR") as interface:
+    gr.Markdown(
+        """
+    # TripoSR Demo
+    [TripoSR](https://github.com/VAST-AI-Research/TripoSR) is a state-of-the-art open-source model for **fast** feedforward 3D reconstruction from a single image, collaboratively developed by [Tripo AI](https://www.tripo3d.ai/) and [Stability AI](https://stability.ai/).
+    
+    **Tips:**
+    1. If you find the result is unsatisfied, please try to change the foreground ratio. It might improve the results.
+    2. You can disable "Remove Background" for the provided examples since they have been already preprocessed.
+    3. Otherwise, please disable "Remove Background" option only if your input image is RGBA with transparent background, image contents are centered and occupy more than 70% of image width or height.
+    """
+    )
+    with gr.Row(variant="panel"):
+        with gr.Column():
+            with gr.Row():
+                input_image = gr.Image(
+                    label="Input Image",
+                    image_mode="RGBA",
+                    sources="upload",
+                    type="pil",
+                    elem_id="content_image",
+                )
+                processed_image = gr.Image(label="Processed Image", interactive=False)
+            with gr.Row():
+                with gr.Group():
+                    do_remove_background = gr.Checkbox(
+                        label="Remove Background", value=True
+                    )
+                    foreground_ratio = gr.Slider(
+                        label="Foreground Ratio",
+                        minimum=0.5,
+                        maximum=1.0,
+                        value=0.85,
+                        step=0.05,
+                    )
+                    mc_resolution = gr.Slider(
+                        label="Marching Cubes Resolution",
+                        minimum=32,
+                        maximum=320,
+                        value=256,
+                        step=32
+                    )
+            with gr.Row():
+                submit = gr.Button("Generate", elem_id="generate", variant="primary")
+        with gr.Column():
+            with gr.Tab("OBJ"):
+                output_model_obj = gr.Model3D(
+                    label="Output Model (OBJ Format)",
+                    interactive=False,
+                )
+                gr.Markdown("Note: The model shown here is flipped. Download to get correct results.")
+            with gr.Tab("GLB"):
+                output_model_glb = gr.Model3D(
+                    label="Output Model (GLB Format)",
+                    interactive=False,
+                )
+                gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
+    with gr.Row(variant="panel"):
+        gr.Examples(
+            examples=[
+                "examples/hamburger.png",
+                "examples/poly_fox.png",
+                "examples/robot.png",
+                "examples/teapot.png",
+                "examples/tiger_girl.png",
+                "examples/horse.png",
+                "examples/flamingo.png",
+                "examples/unicorn.png",
+                "examples/chair.png",
+                "examples/iso_house.png",
+                "examples/marble.png",
+                "examples/police_woman.png",
+                "examples/captured_p.png",
+            ],
+            inputs=[input_image],
+            outputs=[processed_image, output_model_obj, output_model_glb],
+            cache_examples=False,
+            fn=partial(run_example),
+            label="Examples",
+            examples_per_page=20,
+        )
+    submit.click(fn=check_input_image, inputs=[input_image]).success(
+        fn=preprocess,
+        inputs=[input_image, do_remove_background, foreground_ratio],
+        outputs=[processed_image],
+    ).success(
+        fn=generate,
+        inputs=[processed_image, mc_resolution],
+        outputs=[output_model_obj, output_model_glb],
+    )
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--username', type=str, default=None, help='Username for authentication')
+    parser.add_argument('--password', type=str, default=None, help='Password for authentication')
+    parser.add_argument('--port', type=int, default=7860, help='Port to run the server listener on')
+    parser.add_argument("--listen", action='store_true', help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests")
+    parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site")
+    parser.add_argument("--queuesize", type=int, default=1, help="launch gradio queue max_size")
+    args = parser.parse_args()
+    interface.queue(max_size=args.queuesize)
+    interface.launch(
+        auth=(args.username, args.password) if (args.username and args.password) else None,
+        share=args.share,
+        server_name="0.0.0.0" if args.listen else None, 
+        server_port=args.port
+    )

requirements.txt

@@ -0,0 +1,17 @@
+diffusers==0.24.0
+einops==0.7.0
+gradio==4.20.1
+huggingface_hub==0.21.4
+imageio==2.27.0
+numpy==1.24.3
+omegaconf==2.3.0
+packaging==23.2
+Pillow==10.1.0
+rembg==2.0.55
+safetensors==0.3.2
+torch==2.0.0
+torchvision==0.15.1
+tqdm==4.64.1
+transformers==4.27.0
+trimesh==4.0.5
+git+https://github.com/tatsy/torchmcubes.git

requirements2.txt

@@ -0,0 +1,9 @@
+omegaconf==2.3.0
+Pillow==10.1.0
+einops==0.7.0
+git+https://github.com/tatsy/torchmcubes.git
+transformers==4.35.0
+trimesh==4.0.5
+rembg
+huggingface-hub
+imageio[ffmpeg]

run.py

@@ -0,0 +1,162 @@
+import argparse
+import logging
+import os
+import time
+
+import numpy as np
+import rembg
+import torch
+from PIL import Image
+
+from tsr.system import TSR
+from tsr.utils import remove_background, resize_foreground, save_video
+
+
+class Timer:
+    def __init__(self):
+        self.items = {}
+        self.time_scale = 1000.0  # ms
+        self.time_unit = "ms"
+
+    def start(self, name: str) -> None:
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        self.items[name] = time.time()
+        logging.info(f"{name} ...")
+
+    def end(self, name: str) -> float:
+        if name not in self.items:
+            return
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        start_time = self.items.pop(name)
+        delta = time.time() - start_time
+        t = delta * self.time_scale
+        logging.info(f"{name} finished in {t:.2f}{self.time_unit}.")
+
+
+timer = Timer()
+
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
+)
+parser = argparse.ArgumentParser()
+parser.add_argument("image", type=str, nargs="+", help="Path to input image(s).")
+parser.add_argument(
+    "--device",
+    default="cuda:0",
+    type=str,
+    help="Device to use. If no CUDA-compatible device is found, will fallback to 'cpu'. Default: 'cuda:0'",
+)
+parser.add_argument(
+    "--pretrained-model-name-or-path",
+    default="stabilityai/TripoSR",
+    type=str,
+    help="Path to the pretrained model. Could be either a huggingface model id is or a local path. Default: 'stabilityai/TripoSR'",
+)
+parser.add_argument(
+    "--chunk-size",
+    default=8192,
+    type=int,
+    help="Evaluation chunk size for surface extraction and rendering. Smaller chunk size reduces VRAM usage but increases computation time. 0 for no chunking. Default: 8192",
+)
+parser.add_argument(
+    "--mc-resolution",
+    default=256,
+    type=int,
+    help="Marching cubes grid resolution. Default: 256"
+)
+parser.add_argument(
+    "--no-remove-bg",
+    action="store_true",
+    help="If specified, the background will NOT be automatically removed from the input image, and the input image should be an RGB image with gray background and properly-sized foreground. Default: false",
+)
+parser.add_argument(
+    "--foreground-ratio",
+    default=0.85,
+    type=float,
+    help="Ratio of the foreground size to the image size. Only used when --no-remove-bg is not specified. Default: 0.85",
+)
+parser.add_argument(
+    "--output-dir",
+    default="output/",
+    type=str,
+    help="Output directory to save the results. Default: 'output/'",
+)
+parser.add_argument(
+    "--model-save-format",
+    default="obj",
+    type=str,
+    choices=["obj", "glb"],
+    help="Format to save the extracted mesh. Default: 'obj'",
+)
+parser.add_argument(
+    "--render",
+    action="store_true",
+    help="If specified, save a NeRF-rendered video. Default: false",
+)
+args = parser.parse_args()
+
+output_dir = args.output_dir
+os.makedirs(output_dir, exist_ok=True)
+
+device = args.device
+if not torch.cuda.is_available():
+    device = "cpu"
+
+timer.start("Initializing model")
+model = TSR.from_pretrained(
+    args.pretrained_model_name_or_path,
+    config_name="config.yaml",
+    weight_name="model.ckpt",
+)
+model.renderer.set_chunk_size(args.chunk_size)
+model.to(device)
+timer.end("Initializing model")
+
+timer.start("Processing images")
+images = []
+
+if args.no_remove_bg:
+    rembg_session = None
+else:
+    rembg_session = rembg.new_session()
+
+for i, image_path in enumerate(args.image):
+    if args.no_remove_bg:
+        image = np.array(Image.open(image_path).convert("RGB"))
+    else:
+        image = remove_background(Image.open(image_path), rembg_session)
+        image = resize_foreground(image, args.foreground_ratio)
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+        image = Image.fromarray((image * 255.0).astype(np.uint8))
+        if not os.path.exists(os.path.join(output_dir, str(i))):
+            os.makedirs(os.path.join(output_dir, str(i)))
+        image.save(os.path.join(output_dir, str(i), f"input.png"))
+    images.append(image)
+timer.end("Processing images")
+
+for i, image in enumerate(images):
+    logging.info(f"Running image {i + 1}/{len(images)} ...")
+
+    timer.start("Running model")
+    with torch.no_grad():
+        scene_codes = model([image], device=device)
+    timer.end("Running model")
+
+    if args.render:
+        timer.start("Rendering")
+        render_images = model.render(scene_codes, n_views=30, return_type="pil")
+        for ri, render_image in enumerate(render_images[0]):
+            render_image.save(os.path.join(output_dir, str(i), f"render_{ri:03d}.png"))
+        save_video(
+            render_images[0], os.path.join(output_dir, str(i), f"render.mp4"), fps=30
+        )
+        timer.end("Rendering")
+
+    timer.start("Exporting mesh")
+    meshes = model.extract_mesh(scene_codes, resolution=args.mc_resolution)
+    meshes[0].export(os.path.join(output_dir, str(i), f"mesh.{args.model_save_format}"))
+    timer.end("Exporting mesh")

src/laion_bytenas.py

@@ -0,0 +1,257 @@
+import os
+import json
+import random
+from tqdm import tqdm
+import numpy as np
+from PIL import Image, ImageStat
+import torch
+from torch.utils.data import Dataset, DataLoader, IterableDataset, get_worker_info
+from torchvision import transforms as T
+
+
+### >>>>>>>> >>>>>>>> text related >>>>>>>> >>>>>>>> ###
+
+class TokenizerWrapper():
+    def __init__(self, tokenizer, is_train, proportion_empty_prompts, use_generic_prompts=False):
+        self.tokenizer = tokenizer
+        self.is_train = is_train
+        self.proportion_empty_prompts = proportion_empty_prompts
+        self.use_generic_prompts = use_generic_prompts
+
+    def __call__(self, prompts):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+        captions = []
+        for caption in prompts:
+            if random.random() < self.proportion_empty_prompts:
+                captions.append("")
+            else:
+                if self.use_generic_prompts:
+                    captions.append("best quality, high quality")
+                elif isinstance(caption, str):
+                    captions.append(caption)
+                elif isinstance(caption, (list, np.ndarray)):
+                    # take a random caption if there are multiple
+                    captions.append(random.choice(caption) if self.is_train else caption[0])
+                else:
+                    raise ValueError(
+                        f"Caption column should contain either strings or lists of strings."
+                    )
+        inputs = self.tokenizer(
+            captions, max_length=self.tokenizer.model_max_length, padding="max_length",
+            truncation=True, return_tensors="pt"
+        )
+        return inputs.input_ids
+
+
+
+### >>>>>>>> >>>>>>>> image related >>>>>>>> >>>>>>>> ###
+
+MONOCHROMATIC_MAX_VARIANCE = 0.3
+
+def is_monochromatic_image(pil_img):
+    v = ImageStat.Stat(pil_img.convert('RGB')).var
+    return sum(v)<MONOCHROMATIC_MAX_VARIANCE
+
+def isnumeric(text):
+    return (''.join(filter(str.isalnum, text))).isnumeric()
+
+
+
+class TextPromptDataset(IterableDataset):
+    '''
+      The dataset for (text embedding, noise, generated latent) triplets.
+    '''
+    def __init__(self, 
+                data_root, 
+                tokenizer = None,
+                transform = None,
+                rank = 0,
+                world_size = 1,
+                shuffle = True,
+    ):
+        self.tokenizer = tokenizer
+        self.transform = transform
+        
+        self.img_root = os.path.join(data_root, 'JPEGImages')
+        self.data_list = []
+        
+        print("#### Loading filename list...")
+        json_root = os.path.join(data_root, 'list')
+        json_list = [p for p in os.listdir(json_root) if p.startswith("shard") and p.endswith('.json')]
+        
+        # duplicate several shards to make sure each process has the same number of shards
+        assert len(json_list) > world_size
+        duplicate = world_size - len(json_list)%world_size if len(json_list)%world_size>0 else 0
+        json_list = json_list + json_list[:duplicate]
+        json_list = json_list[rank::world_size]
+        
+        for json_file in tqdm(json_list):
+            shard_name = os.path.basename(json_file).split('.')[0]
+            with open(os.path.join(json_root, json_file)) as f:
+                key_text_pairs = json.load(f)
+                
+            for pair in key_text_pairs:
+                self.data_list.append( [shard_name] + pair )
+
+        print("#### All filename loaded...")
+        
+        self.shuffle = shuffle
+        
+    def __len__(self):
+        return len(self.data_list)
+    
+    
+    def __iter__(self):
+        worker_info = get_worker_info()
+        
+        if worker_info is None:  # single-process data loading, return the full iterator
+            data_list = self.data_list
+        else:
+            len_data = len(self.data_list) - len(self.data_list) % worker_info.num_workers
+            data_list = self.data_list[:len_data][worker_info.id :: worker_info.num_workers]
+            # print(worker_info.num_workers, worker_info.id, len(data_list)/len(self.data_list))
+            
+        if self.shuffle:
+            random.shuffle(data_list) 
+            
+        while True:    
+            for idx in range(len(data_list)):
+                # try:
+                shard_name = data_list[idx][0]
+                data = {}
+                
+                img_file = data_list[idx][1]
+                img = Image.open(os.path.join(self.img_root, shard_name, img_file+'.jpg')).convert("RGB")
+                
+                if is_monochromatic_image(img):
+                    continue
+                
+                if self.transform is not None:
+                    img = self.transform(img)
+                    
+                data['pixel_values'] = img
+                
+                text = data_list[idx][2]
+                if self.tokenizer is not None:
+                    if isinstance(self.tokenizer, list):
+                        assert len(self.tokenizer)==2
+                        data['input_ids'] = self.tokenizer[0](text)[0]
+                        data['input_ids_2'] = self.tokenizer[1](text)[0]
+                    else:
+                        data['input_ids'] = self.tokenizer(text)[0]
+                else:
+                    data['input_ids'] = text
+                
+                yield data
+                
+                # except Exception as e:
+                #     raise(e)
+
+    def collate_fn(self, examples):
+        pixel_values = torch.stack([example["pixel_values"] for example in examples])
+        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
+        
+        if self.tokenizer is not None:
+            if isinstance(self.tokenizer, list):
+                assert len(self.tokenizer)==2
+                input_ids = torch.stack([example["input_ids"] for example in examples])
+                input_ids_2 = torch.stack([example["input_ids_2"] for example in examples])
+                return {"pixel_values": pixel_values, "input_ids": input_ids, "input_ids_2": input_ids_2,}
+            else:
+                input_ids = torch.stack([example["input_ids"] for example in examples])
+                return {"pixel_values": pixel_values, "input_ids": input_ids,}
+        else:
+            input_ids = [example["input_ids"] for example in examples]
+            return {"pixel_values": pixel_values, "input_ids": input_ids,}
+    
+    
+def make_train_dataset(
+        train_data_path, 
+        size = 512,
+        tokenizer=None, 
+        cfg_drop_ratio=0,
+        rank=0, 
+        world_size=1,
+        shuffle=True,
+    ):
+    
+    _image_transform = T.Compose([
+            T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+            T.Resize(size),
+            T.CenterCrop((size,size)),
+            T.ToTensor(),
+            T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+
+    if tokenizer is not None:
+        if isinstance(tokenizer, list):
+            assert len(tokenizer)==2
+            tokenizer_1 = TokenizerWrapper(
+                tokenizer[0], 
+                is_train=True, 
+                proportion_empty_prompts=cfg_drop_ratio,
+                use_generic_prompts=False,
+            )
+            tokenizer_2 = TokenizerWrapper(
+                tokenizer[1], 
+                is_train=True, 
+                proportion_empty_prompts=cfg_drop_ratio,
+                use_generic_prompts=False,
+            )
+            tokenizer = [tokenizer_1, tokenizer_2]
+            
+        else:
+            tokenizer = TokenizerWrapper(
+                tokenizer, 
+                is_train=True, 
+                proportion_empty_prompts=cfg_drop_ratio,
+                use_generic_prompts=False,
+            )
+
+        
+    train_dataset = TextPromptDataset(
+        data_root=train_data_path,
+        transform=_image_transform,
+        rank=rank,
+        world_size=world_size,
+        tokenizer=tokenizer,
+        shuffle=shuffle,
+    )
+    return train_dataset
+    
+    
+    
+    
+    
+    
+    
+    
+    
+
+### >>>>>>>> >>>>>>>> Test >>>>>>>> >>>>>>>> ###
+if __name__ == "__main__":
+    from transformers import CLIPTextModel, CLIPTokenizer
+    tokenizer = CLIPTokenizer.from_pretrained(
+        "/mnt/bn/ic-research-aigc-editing/fast-diffusion-models/assets/public_models/StableDiffusion/stable-diffusion-v1-5",
+        subfolder="tokenizer"
+    )
+    train_dataset = make_train_dataset(tokenizer=tokenizer, rank=0, world_size=10)
+    
+    loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=64, num_workers=0, 
+        collate_fn=train_dataset.collect_fn if hasattr(train_dataset, 'collect_fn') else None,
+    )
+    for batch in loader:
+        pixel_values = batch["pixel_values"]
+        prompt_ids = batch['input_ids']
+        from einops import rearrange
+        pixel_values = rearrange(pixel_values, 'b c h w -> b h w c')
+        
+        for i in range(pixel_values.shape[0]):
+            import pdb; pdb.set_trace()
+            Image.fromarray(((pixel_values[i] + 1 )/2 * 255 ).numpy().astype(np.uint8)).save('tmp.png')
+            input_id = prompt_ids[i]
+            text = tokenizer.decode(input_id).split('<|startoftext|>')[-1].split('<|endoftext|>')[0]
+            print(text)
+        pass

src/pfode_solver.py

@@ -0,0 +1,120 @@
+import os, math, random, argparse, logging
+from pathlib import Path
+from typing import Optional, Union, List, Callable
+from collections import OrderedDict
+from packaging import version
+from tqdm.auto import tqdm
+from omegaconf import OmegaConf
+        
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+import torchvision
+
+
+class PFODESolver():
+    def __init__(self, scheduler, t_initial=1, t_terminal=0,) -> None:
+        self.t_initial = t_initial
+        self.t_terminal = t_terminal
+        self.scheduler = scheduler
+
+        train_step_terminal = 0 
+        train_step_initial = train_step_terminal + self.scheduler.config.num_train_timesteps # 0+1000
+        self.stepsize  = (t_terminal-t_initial) / (train_step_terminal - train_step_initial) #1/1000
+    
+    def get_timesteps(self, t_start, t_end, num_steps):
+        # (b,) -> (b,1)
+        t_start = t_start[:, None]
+        t_end = t_end[:, None]
+        assert t_start.dim() == 2
+        
+        timepoints = torch.arange(0, num_steps, 1).expand(t_start.shape[0], num_steps).to(device=t_start.device)
+        interval = (t_end - t_start) / (torch.ones([1], device=t_start.device) * num_steps)
+        timepoints = t_start + interval * timepoints
+        
+        timesteps = (self.scheduler.num_train_timesteps - 1) + (timepoints - self.t_initial) / self.stepsize # correspondint to StableDiffusion indexing system, from 999 (t_init) -> 0 (dt)
+        return timesteps.round().long()
+    
+    def solve(self, 
+              latents, 
+              unet, 
+              t_start, 
+              t_end, 
+              prompt_embeds, 
+              negative_prompt_embeds, 
+              guidance_scale=1.0,
+              num_steps = 2,
+              num_windows = 1,
+    ):
+        assert t_start.dim() == 1
+        assert guidance_scale >= 1 and torch.all(torch.gt(t_start, t_end))
+        
+        do_classifier_free_guidance = True if guidance_scale > 1 else False
+        bsz = latents.shape[0]
+            
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+            
+        timestep_cond = None
+        if unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(guidance_scale - 1).repeat(bsz)
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=unet.config.time_cond_proj_dim
+            ).to(device=latents.device, dtype=latents.dtype)
+            
+        
+        timesteps = self.get_timesteps(t_start, t_end, num_steps).to(device=latents.device)
+        timestep_interval = self.scheduler.config.num_train_timesteps // (num_windows * num_steps)
+
+        # Denoising loop
+        with torch.no_grad():
+            for i in range(num_steps):
+                t = torch.cat([timesteps[:, i]]*2) if do_classifier_free_guidance else timesteps[:, i]
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                noise_pred = unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    return_dict=False,
+                )[0]
+
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                ##### STEP: compute the previous noisy sample x_t -> x_t-1
+                batch_timesteps = timesteps[:, i].cpu()
+                prev_timestep = batch_timesteps - timestep_interval
+
+                alpha_prod_t = self.scheduler.alphas_cumprod[batch_timesteps]
+                alpha_prod_t_prev = torch.zeros_like(alpha_prod_t)
+                for ib in range(prev_timestep.shape[0]): 
+                    alpha_prod_t_prev[ib] = self.scheduler.alphas_cumprod[prev_timestep[ib]] if prev_timestep[ib] >= 0 else self.scheduler.final_alpha_cumprod
+                beta_prod_t = 1 - alpha_prod_t
+                
+                alpha_prod_t = alpha_prod_t.to(device=latents.device, dtype=latents.dtype)
+                alpha_prod_t_prev = alpha_prod_t_prev.to(device=latents.device, dtype=latents.dtype)
+                beta_prod_t = beta_prod_t.to(device=latents.device, dtype=latents.dtype)
+
+                if self.scheduler.config.prediction_type == "epsilon":
+                    pred_original_sample = (latents - beta_prod_t[:,None,None,None] ** (0.5) * noise_pred) / alpha_prod_t[:, None,None,None] ** (0.5)
+                    pred_epsilon = noise_pred
+                elif self.scheduler.config.prediction_type == "v_prediction":
+                    pred_original_sample = (alpha_prod_t[:,None,None,None]**0.5) * latents - (beta_prod_t[:,None,None,None]**0.5) * noise_pred
+                    pred_epsilon = (alpha_prod_t[:,None,None,None]**0.5) * noise_pred + (beta_prod_t[:,None,None,None]**0.5) * latents
+                else:
+                    raise NotImplementedError
+                    
+                pred_sample_direction = (1 - alpha_prod_t_prev[:,None,None,None]) ** (0.5) * pred_epsilon
+                latents = alpha_prod_t_prev[:,None,None,None] ** (0.5) * pred_original_sample + pred_sample_direction
+
+            
+        return latents
+    
+    
+    
+    

src/scheduler_perflow.py

@@ -0,0 +1,343 @@
+# Copyright 2023 Stanford University Team 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.
+
+# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
+# and https://github.com/hojonathanho/diffusion
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+
+
+class Time_Windows():
+    def __init__(self, t_initial=1, t_terminal=0, num_windows=4, precision=1./1000) -> None:
+        assert t_terminal < t_initial
+        time_windows = [ 1.*i/num_windows for i in range(1, num_windows+1)][::-1]
+        
+        self.window_starts = time_windows                      # [1.0, 0.75, 0.5, 0.25]
+        self.window_ends = time_windows[1:] + [t_terminal]     # [0.75, 0.5, 0.25, 0]
+        self.precision = precision
+    
+    def get_window(self, tp):
+        idx = 0
+        # robust to numerical error; e.g, (0.6+1/10000) belongs to [0.6, 0.3)
+        while (tp-0.1*self.precision) <= self.window_ends[idx]: 
+            idx += 1
+        return self.window_starts[idx], self.window_ends[idx]
+    
+    def lookup_window(self, timepoint):
+        if timepoint.dim() == 0:
+            t_start, t_end = self.get_window(timepoint)
+            t_start = torch.ones_like(timepoint) * t_start
+            t_end = torch.ones_like(timepoint) * t_end
+        else:
+            t_start = torch.zeros_like(timepoint)
+            t_end = torch.zeros_like(timepoint)
+            bsz = timepoint.shape[0]
+            for i in range(bsz):
+                tp = timepoint[i]
+                ts, te = self.get_window(tp)
+                t_start[i] = ts
+                t_end[i] = te
+        return t_start, t_end
+    
+    
+
+@dataclass
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
+class PeRFlowSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.FloatTensor
+    pred_original_sample: Optional[torch.FloatTensor] = None
+
+
+# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    else:
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+
+class PeRFlowScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `ReFlowScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
+    non-Markovian guidance.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        set_alpha_to_one (`bool`, defaults to `True`):
+            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
+            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
+            otherwise it uses the alpha value at step 0.
+        prediction_type (`str`, defaults to `epsilon`, *optional*)
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.00085,
+        beta_end: float = 0.012,
+        beta_schedule: str = "scaled_linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        set_alpha_to_one: bool = False,
+        prediction_type: str = "epsilon",
+        t_noise: float = 1,
+        t_clean: float = 0,
+        num_time_windows = 4,
+    ):
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+        elif beta_schedule == "squaredcos_cap_v2":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps)
+        else:
+            raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        # At every step in ddim, we are looking into the previous alphas_cumprod
+        # For the final step, there is no previous alphas_cumprod because we are already at 0
+        # `set_alpha_to_one` decides whether we set this parameter simply to one or
+        # whether we use the final alpha of the "non-previous" one.
+        self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
+        
+        # # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        self.time_windows = Time_Windows(t_initial=t_noise, t_terminal=t_clean, 
+                                        num_windows=num_time_windows,
+                                        precision=1./num_train_timesteps)
+
+    def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`int`, *optional*):
+                The current timestep in the diffusion chain.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        return sample
+
+
+    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+        """
+        if num_inference_steps < self.config.num_time_windows:
+            num_inference_steps = self.config.num_time_windows
+            print(f"### We recommend a num_inference_steps not less than num_time_windows. It's set as {self.config.num_time_windows}.")
+
+        timesteps = []
+        for i in range(self.config.num_time_windows):
+            if i < num_inference_steps%self.config.num_time_windows:
+                num_steps_cur_win = num_inference_steps//self.config.num_time_windows+1
+            else:
+                num_steps_cur_win = num_inference_steps//self.config.num_time_windows
+
+            t_s = self.time_windows.window_starts[i]
+            t_e = self.time_windows.window_ends[i]
+            timesteps_cur_win = np.linspace(t_s, t_e, num=num_steps_cur_win, endpoint=False)
+            timesteps.append(timesteps_cur_win)
+                        
+        timesteps = np.concatenate(timesteps)
+
+        self.timesteps = torch.from_numpy(
+            (timesteps*self.config.num_train_timesteps).astype(np.int64)
+        ).to(device)
+            
+    def get_window_alpha(self, timestep):
+        time_windows = self.time_windows
+        num_train_timesteps = self.config.num_train_timesteps
+        
+        t_win_start, t_win_end = time_windows.lookup_window(timestep / num_train_timesteps)
+        t_win_len = t_win_end - t_win_start
+        t_interval = timestep / num_train_timesteps - t_win_start # NOTE: negative value
+
+        idx_start = (t_win_start*num_train_timesteps - 1 ).long()
+        idx_end = torch.clamp( (t_win_end*num_train_timesteps - 1 ).long(), min=0)
+        alpha_cumprod_s_e = self.alphas_cumprod[idx_start] / self.alphas_cumprod[idx_end]        
+        gamma_s_e = alpha_cumprod_s_e ** 0.5
+        
+        return t_win_start, t_win_end, t_win_len, t_interval, gamma_s_e
+        
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: int,
+        sample: torch.FloatTensor,
+        return_dict: bool = True,
+    ) -> Union[PeRFlowSchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~schedulers.scheduling_ddim.PeRFlowSchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.PeRFlowSchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_ddim.PeRFlowSchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+
+        if self.config.prediction_type == "epsilon":
+            pred_epsilon = model_output
+            t_win_start, t_win_end, t_win_len, t_interval, gamma_s_e = self.get_window_alpha(timestep)
+            pred_sample_end = ( sample - (1-t_interval/t_win_len) * ((1-gamma_s_e**2)**0.5) * pred_epsilon ) \
+                / ( gamma_s_e + t_interval / t_win_len * (1-gamma_s_e) )
+            pred_velocity = (pred_sample_end - sample) / (t_win_end - (t_win_start + t_interval))
+
+        elif self.config.prediction_type == "velocity":
+            pred_velocity = model_output
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `velocity`."
+            )
+
+        # get dt
+        idx = torch.argwhere(torch.where(self.timesteps==timestep, 1,0))
+        prev_step = self.timesteps[idx+1] if (idx+1)<len(self.timesteps) else 0
+        dt = (prev_step - timestep) / self.config.num_train_timesteps
+        dt = dt.to(sample.device, sample.dtype)
+           
+        prev_sample = sample + dt * pred_velocity
+
+        if not return_dict:
+            return (prev_sample,)
+        return PeRFlowSchedulerOutput(prev_sample=prev_sample, pred_original_sample=None)
+
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        noise: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
+        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device) - 1   # indexing from 0
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps

src/utils_perflow.py

@@ -0,0 +1,77 @@
+import os
+from collections import OrderedDict
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_unet_checkpoint, convert_ldm_vae_checkpoint, convert_ldm_clip_checkpoint
+    
+
+def merge_delta_weights_into_unet(pipe, delta_weights):
+    unet_weights = pipe.unet.state_dict()
+    assert unet_weights.keys() == delta_weights.keys()
+    for key in delta_weights.keys():
+        dtype = unet_weights[key].dtype
+        unet_weights[key] = unet_weights[key].to(dtype=delta_weights[key].dtype) + delta_weights[key].to(device=unet_weights[key].device)
+        unet_weights[key] = unet_weights[key].to(dtype)
+    pipe.unet.load_state_dict(unet_weights, strict=True)
+    return pipe
+
+
+def load_delta_weights_into_unet(
+    pipe, 
+    model_path = "hsyan/piecewise-rectified-flow-v0-1", 
+    base_path = "runwayml/stable-diffusion-v1-5",
+):
+    ## load delta_weights
+    if os.path.exists(os.path.join(model_path, "delta_weights.safetensors")):
+        print("### delta_weights exists, loading...")
+        delta_weights = OrderedDict()
+        with safe_open(os.path.join(model_path, "delta_weights.safetensors"), framework="pt", device="cpu") as f:
+            for key in f.keys():
+                delta_weights[key] = f.get_tensor(key)
+                
+    elif os.path.exists(os.path.join(model_path, "diffusion_pytorch_model.safetensors")):
+        print("### merged_weights exists, loading...")
+        merged_weights = OrderedDict()
+        with safe_open(os.path.join(model_path, "diffusion_pytorch_model.safetensors"), framework="pt", device="cpu") as f:
+            for key in f.keys():
+                merged_weights[key] = f.get_tensor(key)
+                
+        base_weights = StableDiffusionPipeline.from_pretrained(
+            base_path, torch_dtype=torch.float16, safety_checker=None).unet.state_dict()
+        assert base_weights.keys() == merged_weights.keys()
+        
+        delta_weights = OrderedDict()
+        for key in merged_weights.keys():
+            delta_weights[key] = merged_weights[key] - base_weights[key].to(device=merged_weights[key].device, dtype=merged_weights[key].dtype)
+        
+        print("### saving delta_weights...")
+        save_file(delta_weights, os.path.join(model_path, "delta_weights.safetensors"))
+        
+    else:
+        raise ValueError(f"{model_path} does not contain delta weights or merged weights")
+        
+    ## merge delta_weights to the target pipeline
+    pipe = merge_delta_weights_into_unet(pipe, delta_weights)
+    return pipe
+    
+
+
+
+def load_dreambooth_into_pipeline(pipe, sd_dreambooth):
+    assert sd_dreambooth.endswith(".safetensors")
+    state_dict = {}
+    with safe_open(sd_dreambooth, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            state_dict[key] = f.get_tensor(key)
+    
+    unet_config = {} # unet, line 449 in convert_ldm_unet_checkpoint
+    for key in pipe.unet.config.keys():
+        if key != 'num_class_embeds':
+            unet_config[key] = pipe.unet.config[key]
+            
+    pipe.unet.load_state_dict(convert_ldm_unet_checkpoint(state_dict, unet_config), strict=False)
+    pipe.vae.load_state_dict(convert_ldm_vae_checkpoint(state_dict, pipe.vae.config))
+    pipe.text_encoder = convert_ldm_clip_checkpoint(state_dict, text_encoder=pipe.text_encoder)
+    return pipe

test.yaml

@@ -0,0 +1,10 @@
+name: test
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+  - conda-forge
+dependencies:
+  - python=3.10.12
+  - pip=23.2.1
+  - cudatoolkit=11.7

tsr/models/isosurface.py

@@ -0,0 +1,52 @@
+from typing import Callable, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torchmcubes import marching_cubes
+
+
+class IsosurfaceHelper(nn.Module):
+    points_range: Tuple[float, float] = (0, 1)
+
+    @property
+    def grid_vertices(self) -> torch.FloatTensor:
+        raise NotImplementedError
+
+
+class MarchingCubeHelper(IsosurfaceHelper):
+    def __init__(self, resolution: int) -> None:
+        super().__init__()
+        self.resolution = resolution
+        self.mc_func: Callable = marching_cubes
+        self._grid_vertices: Optional[torch.FloatTensor] = None
+
+    @property
+    def grid_vertices(self) -> torch.FloatTensor:
+        if self._grid_vertices is None:
+            # keep the vertices on CPU so that we can support very large resolution
+            x, y, z = (
+                torch.linspace(*self.points_range, self.resolution),
+                torch.linspace(*self.points_range, self.resolution),
+                torch.linspace(*self.points_range, self.resolution),
+            )
+            x, y, z = torch.meshgrid(x, y, z, indexing="ij")
+            verts = torch.cat(
+                [x.reshape(-1, 1), y.reshape(-1, 1), z.reshape(-1, 1)], dim=-1
+            ).reshape(-1, 3)
+            self._grid_vertices = verts
+        return self._grid_vertices
+
+    def forward(
+        self,
+        level: torch.FloatTensor,
+    ) -> Tuple[torch.FloatTensor, torch.LongTensor]:
+        level = -level.view(self.resolution, self.resolution, self.resolution)
+        try:
+            v_pos, t_pos_idx = self.mc_func(level.detach(), 0.0)
+        except AttributeError:
+            print("torchmcubes was not compiled with CUDA support, use CPU version instead.")
+            v_pos, t_pos_idx = self.mc_func(level.detach().cpu(), 0.0)
+        v_pos = v_pos[..., [2, 1, 0]]
+        v_pos = v_pos / (self.resolution - 1.0)
+        return v_pos.to(level.device), t_pos_idx.to(level.device)

tsr/models/nerf_renderer.py

@@ -0,0 +1,180 @@
+from dataclasses import dataclass
+from typing import Dict
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, reduce
+
+from ..utils import (
+    BaseModule,
+    chunk_batch,
+    get_activation,
+    rays_intersect_bbox,
+    scale_tensor,
+)
+
+
+class TriplaneNeRFRenderer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        radius: float
+
+        feature_reduction: str = "concat"
+        density_activation: str = "trunc_exp"
+        density_bias: float = -1.0
+        color_activation: str = "sigmoid"
+        num_samples_per_ray: int = 128
+        randomized: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        assert self.cfg.feature_reduction in ["concat", "mean"]
+        self.chunk_size = 0
+
+    def set_chunk_size(self, chunk_size: int):
+        assert (
+            chunk_size >= 0
+        ), "chunk_size must be a non-negative integer (0 for no chunking)."
+        self.chunk_size = chunk_size
+
+    def query_triplane(
+        self,
+        decoder: torch.nn.Module,
+        positions: torch.Tensor,
+        triplane: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        input_shape = positions.shape[:-1]
+        positions = positions.view(-1, 3)
+
+        # positions in (-radius, radius)
+        # normalized to (-1, 1) for grid sample
+        positions = scale_tensor(
+            positions, (-self.cfg.radius, self.cfg.radius), (-1, 1)
+        )
+
+        def _query_chunk(x):
+            indices2D: torch.Tensor = torch.stack(
+                (x[..., [0, 1]], x[..., [0, 2]], x[..., [1, 2]]),
+                dim=-3,
+            )
+            out: torch.Tensor = F.grid_sample(
+                rearrange(triplane, "Np Cp Hp Wp -> Np Cp Hp Wp", Np=3),
+                rearrange(indices2D, "Np N Nd -> Np () N Nd", Np=3),
+                align_corners=False,
+                mode="bilinear",
+            )
+            if self.cfg.feature_reduction == "concat":
+                out = rearrange(out, "Np Cp () N -> N (Np Cp)", Np=3)
+            elif self.cfg.feature_reduction == "mean":
+                out = reduce(out, "Np Cp () N -> N Cp", Np=3, reduction="mean")
+            else:
+                raise NotImplementedError
+
+            net_out: Dict[str, torch.Tensor] = decoder(out)
+            return net_out
+
+        if self.chunk_size > 0:
+            net_out = chunk_batch(_query_chunk, self.chunk_size, positions)
+        else:
+            net_out = _query_chunk(positions)
+
+        net_out["density_act"] = get_activation(self.cfg.density_activation)(
+            net_out["density"] + self.cfg.density_bias
+        )
+        net_out["color"] = get_activation(self.cfg.color_activation)(
+            net_out["features"]
+        )
+
+        net_out = {k: v.view(*input_shape, -1) for k, v in net_out.items()}
+
+        return net_out
+
+    def _forward(
+        self,
+        decoder: torch.nn.Module,
+        triplane: torch.Tensor,
+        rays_o: torch.Tensor,
+        rays_d: torch.Tensor,
+        **kwargs,
+    ):
+        rays_shape = rays_o.shape[:-1]
+        rays_o = rays_o.view(-1, 3)
+        rays_d = rays_d.view(-1, 3)
+        n_rays = rays_o.shape[0]
+
+        t_near, t_far, rays_valid = rays_intersect_bbox(rays_o, rays_d, self.cfg.radius)
+        t_near, t_far = t_near[rays_valid], t_far[rays_valid]
+
+        t_vals = torch.linspace(
+            0, 1, self.cfg.num_samples_per_ray + 1, device=triplane.device
+        )
+        t_mid = (t_vals[:-1] + t_vals[1:]) / 2.0
+        z_vals = t_near * (1 - t_mid[None]) + t_far * t_mid[None]  # (N_rays, N_samples)
+
+        xyz = (
+            rays_o[:, None, :] + z_vals[..., None] * rays_d[..., None, :]
+        )  # (N_rays, N_sample, 3)
+
+        mlp_out = self.query_triplane(
+            decoder=decoder,
+            positions=xyz,
+            triplane=triplane,
+        )
+
+        eps = 1e-10
+        # deltas = z_vals[:, 1:] - z_vals[:, :-1] # (N_rays, N_samples)
+        deltas = t_vals[1:] - t_vals[:-1]  # (N_rays, N_samples)
+        alpha = 1 - torch.exp(
+            -deltas * mlp_out["density_act"][..., 0]
+        )  # (N_rays, N_samples)
+        accum_prod = torch.cat(
+            [
+                torch.ones_like(alpha[:, :1]),
+                torch.cumprod(1 - alpha[:, :-1] + eps, dim=-1),
+            ],
+            dim=-1,
+        )
+        weights = alpha * accum_prod  # (N_rays, N_samples)
+        comp_rgb_ = (weights[..., None] * mlp_out["color"]).sum(dim=-2)  # (N_rays, 3)
+        opacity_ = weights.sum(dim=-1)  # (N_rays)
+
+        comp_rgb = torch.zeros(
+            n_rays, 3, dtype=comp_rgb_.dtype, device=comp_rgb_.device
+        )
+        opacity = torch.zeros(n_rays, dtype=opacity_.dtype, device=opacity_.device)
+        comp_rgb[rays_valid] = comp_rgb_
+        opacity[rays_valid] = opacity_
+
+        comp_rgb += 1 - opacity[..., None]
+        comp_rgb = comp_rgb.view(*rays_shape, 3)
+
+        return comp_rgb
+
+    def forward(
+        self,
+        decoder: torch.nn.Module,
+        triplane: torch.Tensor,
+        rays_o: torch.Tensor,
+        rays_d: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        if triplane.ndim == 4:
+            comp_rgb = self._forward(decoder, triplane, rays_o, rays_d)
+        else:
+            comp_rgb = torch.stack(
+                [
+                    self._forward(decoder, triplane[i], rays_o[i], rays_d[i])
+                    for i in range(triplane.shape[0])
+                ],
+                dim=0,
+            )
+
+        return comp_rgb
+
+    def train(self, mode=True):
+        self.randomized = mode and self.cfg.randomized
+        return super().train(mode=mode)
+
+    def eval(self):
+        self.randomized = False
+        return super().eval()

tsr/models/network_utils.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from ..utils import BaseModule
+
+
+class TriplaneUpsampleNetwork(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int
+        out_channels: int
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.upsample = nn.ConvTranspose2d(
+            self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2
+        )
+
+    def forward(self, triplanes: torch.Tensor) -> torch.Tensor:
+        triplanes_up = rearrange(
+            self.upsample(
+                rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
+            ),
+            "(B Np) Co Hp Wp -> B Np Co Hp Wp",
+            Np=3,
+        )
+        return triplanes_up
+
+
+class NeRFMLP(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        in_channels: int
+        n_neurons: int
+        n_hidden_layers: int
+        activation: str = "relu"
+        bias: bool = True
+        weight_init: Optional[str] = "kaiming_uniform"
+        bias_init: Optional[str] = None
+
+    cfg: Config
+
+    def configure(self) -> None:
+        layers = [
+            self.make_linear(
+                self.cfg.in_channels,
+                self.cfg.n_neurons,
+                bias=self.cfg.bias,
+                weight_init=self.cfg.weight_init,
+                bias_init=self.cfg.bias_init,
+            ),
+            self.make_activation(self.cfg.activation),
+        ]
+        for i in range(self.cfg.n_hidden_layers - 1):
+            layers += [
+                self.make_linear(
+                    self.cfg.n_neurons,
+                    self.cfg.n_neurons,
+                    bias=self.cfg.bias,
+                    weight_init=self.cfg.weight_init,
+                    bias_init=self.cfg.bias_init,
+                ),
+                self.make_activation(self.cfg.activation),
+            ]
+        layers += [
+            self.make_linear(
+                self.cfg.n_neurons,
+                4,  # density 1 + features 3
+                bias=self.cfg.bias,
+                weight_init=self.cfg.weight_init,
+                bias_init=self.cfg.bias_init,
+            )
+        ]
+        self.layers = nn.Sequential(*layers)
+
+    def make_linear(
+        self,
+        dim_in,
+        dim_out,
+        bias=True,
+        weight_init=None,
+        bias_init=None,
+    ):
+        layer = nn.Linear(dim_in, dim_out, bias=bias)
+
+        if weight_init is None:
+            pass
+        elif weight_init == "kaiming_uniform":
+            torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
+        else:
+            raise NotImplementedError
+
+        if bias:
+            if bias_init is None:
+                pass
+            elif bias_init == "zero":
+                torch.nn.init.zeros_(layer.bias)
+            else:
+                raise NotImplementedError
+
+        return layer
+
+    def make_activation(self, activation):
+        if activation == "relu":
+            return nn.ReLU(inplace=True)
+        elif activation == "silu":
+            return nn.SiLU(inplace=True)
+        else:
+            raise NotImplementedError
+
+    def forward(self, x):
+        inp_shape = x.shape[:-1]
+        x = x.reshape(-1, x.shape[-1])
+
+        features = self.layers(x)
+        features = features.reshape(*inp_shape, -1)
+        out = {"density": features[..., 0:1], "features": features[..., 1:4]}
+
+        return out

tsr/models/tokenizers/image.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from transformers.models.vit.modeling_vit import ViTModel
+
+from ...utils import BaseModule
+
+
+class DINOSingleImageTokenizer(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pretrained_model_name_or_path: str = "facebook/dino-vitb16"
+        enable_gradient_checkpointing: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.model: ViTModel = ViTModel(
+            ViTModel.config_class.from_pretrained(
+                hf_hub_download(
+                    repo_id=self.cfg.pretrained_model_name_or_path,
+                    filename="config.json",
+                )
+            )
+        )
+
+        if self.cfg.enable_gradient_checkpointing:
+            self.model.encoder.gradient_checkpointing = True
+
+        self.register_buffer(
+            "image_mean",
+            torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+        self.register_buffer(
+            "image_std",
+            torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1),
+            persistent=False,
+        )
+
+    def forward(self, images: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        packed = False
+        if images.ndim == 4:
+            packed = True
+            images = images.unsqueeze(1)
+
+        batch_size, n_input_views = images.shape[:2]
+        images = (images - self.image_mean) / self.image_std
+        out = self.model(
+            rearrange(images, "B N C H W -> (B N) C H W"), interpolate_pos_encoding=True
+        )
+        local_features, global_features = out.last_hidden_state, out.pooler_output
+        local_features = local_features.permute(0, 2, 1)
+        local_features = rearrange(
+            local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size
+        )
+        if packed:
+            local_features = local_features.squeeze(1)
+
+        return local_features
+
+    def detokenize(self, *args, **kwargs):
+        raise NotImplementedError

tsr/models/tokenizers/triplane.py

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

tsr/models/transformer/attention.py

@@ -0,0 +1,653 @@
+# Copyright 2023 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.
+#
+# --------
+#
+# Modified 2024 by the Tripo AI and Stability AI Team.
+#
+# Copyright (c) 2024 Tripo AI & Stability AI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Attention(nn.Module):
+    r"""
+    A cross attention layer.
+
+    Parameters:
+        query_dim (`int`):
+            The number of channels in the query.
+        cross_attention_dim (`int`, *optional*):
+            The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.
+        heads (`int`,  *optional*, defaults to 8):
+            The number of heads to use for multi-head attention.
+        dim_head (`int`,  *optional*, defaults to 64):
+            The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0):
+            The dropout probability to use.
+        bias (`bool`, *optional*, defaults to False):
+            Set to `True` for the query, key, and value linear layers to contain a bias parameter.
+        upcast_attention (`bool`, *optional*, defaults to False):
+            Set to `True` to upcast the attention computation to `float32`.
+        upcast_softmax (`bool`, *optional*, defaults to False):
+            Set to `True` to upcast the softmax computation to `float32`.
+        cross_attention_norm (`str`, *optional*, defaults to `None`):
+            The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
+        cross_attention_norm_num_groups (`int`, *optional*, defaults to 32):
+            The number of groups to use for the group norm in the cross attention.
+        added_kv_proj_dim (`int`, *optional*, defaults to `None`):
+            The number of channels to use for the added key and value projections. If `None`, no projection is used.
+        norm_num_groups (`int`, *optional*, defaults to `None`):
+            The number of groups to use for the group norm in the attention.
+        spatial_norm_dim (`int`, *optional*, defaults to `None`):
+            The number of channels to use for the spatial normalization.
+        out_bias (`bool`, *optional*, defaults to `True`):
+            Set to `True` to use a bias in the output linear layer.
+        scale_qk (`bool`, *optional*, defaults to `True`):
+            Set to `True` to scale the query and key by `1 / sqrt(dim_head)`.
+        only_cross_attention (`bool`, *optional*, defaults to `False`):
+            Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if
+            `added_kv_proj_dim` is not `None`.
+        eps (`float`, *optional*, defaults to 1e-5):
+            An additional value added to the denominator in group normalization that is used for numerical stability.
+        rescale_output_factor (`float`, *optional*, defaults to 1.0):
+            A factor to rescale the output by dividing it with this value.
+        residual_connection (`bool`, *optional*, defaults to `False`):
+            Set to `True` to add the residual connection to the output.
+        _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`):
+            Set to `True` if the attention block is loaded from a deprecated state dict.
+        processor (`AttnProcessor`, *optional*, defaults to `None`):
+            The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and
+            `AttnProcessor` otherwise.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: Optional[int] = None,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        upcast_attention: bool = False,
+        upcast_softmax: bool = False,
+        cross_attention_norm: Optional[str] = None,
+        cross_attention_norm_num_groups: int = 32,
+        added_kv_proj_dim: Optional[int] = None,
+        norm_num_groups: Optional[int] = None,
+        out_bias: bool = True,
+        scale_qk: bool = True,
+        only_cross_attention: bool = False,
+        eps: float = 1e-5,
+        rescale_output_factor: float = 1.0,
+        residual_connection: bool = False,
+        _from_deprecated_attn_block: bool = False,
+        processor: Optional["AttnProcessor"] = None,
+        out_dim: int = None,
+    ):
+        super().__init__()
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.cross_attention_dim = (
+            cross_attention_dim if cross_attention_dim is not None else query_dim
+        )
+        self.upcast_attention = upcast_attention
+        self.upcast_softmax = upcast_softmax
+        self.rescale_output_factor = rescale_output_factor
+        self.residual_connection = residual_connection
+        self.dropout = dropout
+        self.fused_projections = False
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        # we make use of this private variable to know whether this class is loaded
+        # with an deprecated state dict so that we can convert it on the fly
+        self._from_deprecated_attn_block = _from_deprecated_attn_block
+
+        self.scale_qk = scale_qk
+        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        # for slice_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set slice_size with `set_attention_slice`
+        self.sliceable_head_dim = heads
+
+        self.added_kv_proj_dim = added_kv_proj_dim
+        self.only_cross_attention = only_cross_attention
+
+        if self.added_kv_proj_dim is None and self.only_cross_attention:
+            raise ValueError(
+                "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
+            )
+
+        if norm_num_groups is not None:
+            self.group_norm = nn.GroupNorm(
+                num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True
+            )
+        else:
+            self.group_norm = None
+
+        self.spatial_norm = None
+
+        if cross_attention_norm is None:
+            self.norm_cross = None
+        elif cross_attention_norm == "layer_norm":
+            self.norm_cross = nn.LayerNorm(self.cross_attention_dim)
+        elif cross_attention_norm == "group_norm":
+            if self.added_kv_proj_dim is not None:
+                # The given `encoder_hidden_states` are initially of shape
+                # (batch_size, seq_len, added_kv_proj_dim) before being projected
+                # to (batch_size, seq_len, cross_attention_dim). The norm is applied
+                # before the projection, so we need to use `added_kv_proj_dim` as
+                # the number of channels for the group norm.
+                norm_cross_num_channels = added_kv_proj_dim
+            else:
+                norm_cross_num_channels = self.cross_attention_dim
+
+            self.norm_cross = nn.GroupNorm(
+                num_channels=norm_cross_num_channels,
+                num_groups=cross_attention_norm_num_groups,
+                eps=1e-5,
+                affine=True,
+            )
+        else:
+            raise ValueError(
+                f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
+            )
+
+        linear_cls = nn.Linear
+
+        self.linear_cls = linear_cls
+        self.to_q = linear_cls(query_dim, self.inner_dim, bias=bias)
+
+        if not self.only_cross_attention:
+            # only relevant for the `AddedKVProcessor` classes
+            self.to_k = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
+            self.to_v = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)
+        else:
+            self.to_k = None
+            self.to_v = None
+
+        if self.added_kv_proj_dim is not None:
+            self.add_k_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
+            self.add_v_proj = linear_cls(added_kv_proj_dim, self.inner_dim)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(linear_cls(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(nn.Dropout(dropout))
+
+        # set attention processor
+        # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+        # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+        # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+        if processor is None:
+            processor = (
+                AttnProcessor2_0()
+                if hasattr(F, "scaled_dot_product_attention") and self.scale_qk
+                else AttnProcessor()
+            )
+        self.set_processor(processor)
+
+    def set_processor(self, processor: "AttnProcessor") -> None:
+        self.processor = processor
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        **cross_attention_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        The forward method of the `Attention` class.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                The hidden states of the query.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                The hidden states of the encoder.
+            attention_mask (`torch.Tensor`, *optional*):
+                The attention mask to use. If `None`, no mask is applied.
+            **cross_attention_kwargs:
+                Additional keyword arguments to pass along to the cross attention.
+
+        Returns:
+            `torch.Tensor`: The output of the attention layer.
+        """
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**cross_attention_kwargs` is empty
+        return self.processor(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs,
+        )
+
+    def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:
+        r"""
+        Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads`
+        is the number of heads initialized while constructing the `Attention` class.
+
+        Args:
+            tensor (`torch.Tensor`): The tensor to reshape.
+
+        Returns:
+            `torch.Tensor`: The reshaped tensor.
+        """
+        head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(
+            batch_size // head_size, seq_len, dim * head_size
+        )
+        return tensor
+
+    def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:
+        r"""
+        Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is
+        the number of heads initialized while constructing the `Attention` class.
+
+        Args:
+            tensor (`torch.Tensor`): The tensor to reshape.
+            out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is
+                reshaped to `[batch_size * heads, seq_len, dim // heads]`.
+
+        Returns:
+            `torch.Tensor`: The reshaped tensor.
+        """
+        head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3)
+
+        if out_dim == 3:
+            tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size)
+
+        return tensor
+
+    def get_attention_scores(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attention_mask: torch.Tensor = None,
+    ) -> torch.Tensor:
+        r"""
+        Compute the attention scores.
+
+        Args:
+            query (`torch.Tensor`): The query tensor.
+            key (`torch.Tensor`): The key tensor.
+            attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
+
+        Returns:
+            `torch.Tensor`: The attention probabilities/scores.
+        """
+        dtype = query.dtype
+        if self.upcast_attention:
+            query = query.float()
+            key = key.float()
+
+        if attention_mask is None:
+            baddbmm_input = torch.empty(
+                query.shape[0],
+                query.shape[1],
+                key.shape[1],
+                dtype=query.dtype,
+                device=query.device,
+            )
+            beta = 0
+        else:
+            baddbmm_input = attention_mask
+            beta = 1
+
+        attention_scores = torch.baddbmm(
+            baddbmm_input,
+            query,
+            key.transpose(-1, -2),
+            beta=beta,
+            alpha=self.scale,
+        )
+        del baddbmm_input
+
+        if self.upcast_softmax:
+            attention_scores = attention_scores.float()
+
+        attention_probs = attention_scores.softmax(dim=-1)
+        del attention_scores
+
+        attention_probs = attention_probs.to(dtype)
+
+        return attention_probs
+
+    def prepare_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        target_length: int,
+        batch_size: int,
+        out_dim: int = 3,
+    ) -> torch.Tensor:
+        r"""
+        Prepare the attention mask for the attention computation.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                The attention mask to prepare.
+            target_length (`int`):
+                The target length of the attention mask. This is the length of the attention mask after padding.
+            batch_size (`int`):
+                The batch size, which is used to repeat the attention mask.
+            out_dim (`int`, *optional*, defaults to `3`):
+                The output dimension of the attention mask. Can be either `3` or `4`.
+
+        Returns:
+            `torch.Tensor`: The prepared attention mask.
+        """
+        head_size = self.heads
+        if attention_mask is None:
+            return attention_mask
+
+        current_length: int = attention_mask.shape[-1]
+        if current_length != target_length:
+            if attention_mask.device.type == "mps":
+                # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+                # Instead, we can manually construct the padding tensor.
+                padding_shape = (
+                    attention_mask.shape[0],
+                    attention_mask.shape[1],
+                    target_length,
+                )
+                padding = torch.zeros(
+                    padding_shape,
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device,
+                )
+                attention_mask = torch.cat([attention_mask, padding], dim=2)
+            else:
+                # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+                #       we want to instead pad by (0, remaining_length), where remaining_length is:
+                #       remaining_length: int = target_length - current_length
+                # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+                attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
+
+        if out_dim == 3:
+            if attention_mask.shape[0] < batch_size * head_size:
+                attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
+        elif out_dim == 4:
+            attention_mask = attention_mask.unsqueeze(1)
+            attention_mask = attention_mask.repeat_interleave(head_size, dim=1)
+
+        return attention_mask
+
+    def norm_encoder_hidden_states(
+        self, encoder_hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        r"""
+        Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the
+        `Attention` class.
+
+        Args:
+            encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder.
+
+        Returns:
+            `torch.Tensor`: The normalized encoder hidden states.
+        """
+        assert (
+            self.norm_cross is not None
+        ), "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
+
+        if isinstance(self.norm_cross, nn.LayerNorm):
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+        elif isinstance(self.norm_cross, nn.GroupNorm):
+            # Group norm norms along the channels dimension and expects
+            # input to be in the shape of (N, C, *). In this case, we want
+            # to norm along the hidden dimension, so we need to move
+            # (batch_size, sequence_length, hidden_size) ->
+            # (batch_size, hidden_size, sequence_length)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+        else:
+            assert False
+
+        return encoder_hidden_states
+
+    @torch.no_grad()
+    def fuse_projections(self, fuse=True):
+        is_cross_attention = self.cross_attention_dim != self.query_dim
+        device = self.to_q.weight.data.device
+        dtype = self.to_q.weight.data.dtype
+
+        if not is_cross_attention:
+            # fetch weight matrices.
+            concatenated_weights = torch.cat(
+                [self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]
+            )
+            in_features = concatenated_weights.shape[1]
+            out_features = concatenated_weights.shape[0]
+
+            # create a new single projection layer and copy over the weights.
+            self.to_qkv = self.linear_cls(
+                in_features, out_features, bias=False, device=device, dtype=dtype
+            )
+            self.to_qkv.weight.copy_(concatenated_weights)
+
+        else:
+            concatenated_weights = torch.cat(
+                [self.to_k.weight.data, self.to_v.weight.data]
+            )
+            in_features = concatenated_weights.shape[1]
+            out_features = concatenated_weights.shape[0]
+
+            self.to_kv = self.linear_cls(
+                in_features, out_features, bias=False, device=device, dtype=dtype
+            )
+            self.to_kv.weight.copy_(concatenated_weights)
+
+        self.fused_projections = fuse
+
+
+class AttnProcessor:
+    r"""
+    Default processor for performing attention-related computations.
+    """
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+        attention_mask = attn.prepare_attention_mask(
+            attention_mask, sequence_length, batch_size
+        )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+class AttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        residual = hidden_states
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states

tsr/models/transformer/basic_transformer_block.py

@@ -0,0 +1,334 @@
+# Copyright 2023 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.
+#
+# --------
+#
+# Modified 2024 by the Tripo AI and Stability AI Team.
+#
+# Copyright (c) 2024 Tripo AI & Stability AI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .attention import Attention
+
+
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        upcast_attention (`bool`, *optional*):
+            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
+        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
+        final_dropout (`bool` *optional*, defaults to False):
+            Whether to apply a final dropout after the last feed-forward layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+
+        assert norm_type == "layer_norm"
+
+        # Define 3 blocks. Each block has its own normalization layer.
+        # 1. Self-Attn
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+        )
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None or double_self_attention:
+            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+            # the second cross attention block.
+            self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=(
+                    cross_attention_dim if not double_self_attention else None
+                ),
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+            )  # is self-attn if encoder_hidden_states is none
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        # 3. Feed-forward
+        self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+        )
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 0. Self-Attention
+        norm_hidden_states = self.norm1(hidden_states)
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=(
+                encoder_hidden_states if self.only_cross_attention else None
+            ),
+            attention_mask=attention_mask,
+        )
+
+        hidden_states = attn_output + hidden_states
+
+        # 3. Cross-Attention
+        if self.attn2 is not None:
+            norm_hidden_states = self.norm2(hidden_states)
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+            )
+            hidden_states = attn_output + hidden_states
+
+        # 4. Feed-forward
+        norm_hidden_states = self.norm3(hidden_states)
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
+                raise ValueError(
+                    f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+                )
+
+            num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
+            ff_output = torch.cat(
+                [
+                    self.ff(hid_slice)
+                    for hid_slice in norm_hidden_states.chunk(
+                        num_chunks, dim=self._chunk_dim
+                    )
+                ],
+                dim=self._chunk_dim,
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        hidden_states = ff_output + hidden_states
+
+        return hidden_states
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+        activation_fn: str = "geglu",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        linear_cls = nn.Linear
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh")
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim)
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # project dropout
+        self.net.append(nn.Dropout(dropout))
+        # project out
+        self.net.append(linear_cls(inner_dim, dim_out))
+        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
+        if final_dropout:
+            self.net.append(nn.Dropout(dropout))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+class GELU(nn.Module):
+    r"""
+    GELU activation function with tanh approximation support with `approximate="tanh"`.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+        self.approximate = approximate
+
+    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+        if gate.device.type != "mps":
+            return F.gelu(gate, approximate=self.approximate)
+        # mps: gelu is not implemented for float16
+        return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(
+            dtype=gate.dtype
+        )
+
+    def forward(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        hidden_states = self.gelu(hidden_states)
+        return hidden_states
+
+
+class GEGLU(nn.Module):
+    r"""
+    A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int):
+        super().__init__()
+        linear_cls = nn.Linear
+
+        self.proj = linear_cls(dim_in, dim_out * 2)
+
+    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+        if gate.device.type != "mps":
+            return F.gelu(gate)
+        # mps: gelu is not implemented for float16
+        return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
+
+    def forward(self, hidden_states, scale: float = 1.0):
+        args = ()
+        hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1)
+        return hidden_states * self.gelu(gate)
+
+
+class ApproximateGELU(nn.Module):
+    r"""
+    The approximate form of Gaussian Error Linear Unit (GELU). For more details, see section 2:
+    https://arxiv.org/abs/1606.08415.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        return x * torch.sigmoid(1.702 * x)

tsr/models/transformer/transformer_1d.py

@@ -0,0 +1,219 @@
+# Copyright 2023 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.
+#
+# --------
+#
+# Modified 2024 by the Tripo AI and Stability AI Team.
+#
+# Copyright (c) 2024 Tripo AI & Stability AI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ...utils import BaseModule
+from .basic_transformer_block import BasicTransformerBlock
+
+
+class Transformer1D(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        num_attention_heads: int = 16
+        attention_head_dim: int = 88
+        in_channels: Optional[int] = None
+        out_channels: Optional[int] = None
+        num_layers: int = 1
+        dropout: float = 0.0
+        norm_num_groups: int = 32
+        cross_attention_dim: Optional[int] = None
+        attention_bias: bool = False
+        activation_fn: str = "geglu"
+        only_cross_attention: bool = False
+        double_self_attention: bool = False
+        upcast_attention: bool = False
+        norm_type: str = "layer_norm"
+        norm_elementwise_affine: bool = True
+        gradient_checkpointing: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.num_attention_heads = self.cfg.num_attention_heads
+        self.attention_head_dim = self.cfg.attention_head_dim
+        inner_dim = self.num_attention_heads * self.attention_head_dim
+
+        linear_cls = nn.Linear
+
+        # 2. Define input layers
+        self.in_channels = self.cfg.in_channels
+
+        self.norm = torch.nn.GroupNorm(
+            num_groups=self.cfg.norm_num_groups,
+            num_channels=self.cfg.in_channels,
+            eps=1e-6,
+            affine=True,
+        )
+        self.proj_in = linear_cls(self.cfg.in_channels, inner_dim)
+
+        # 3. Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    self.num_attention_heads,
+                    self.attention_head_dim,
+                    dropout=self.cfg.dropout,
+                    cross_attention_dim=self.cfg.cross_attention_dim,
+                    activation_fn=self.cfg.activation_fn,
+                    attention_bias=self.cfg.attention_bias,
+                    only_cross_attention=self.cfg.only_cross_attention,
+                    double_self_attention=self.cfg.double_self_attention,
+                    upcast_attention=self.cfg.upcast_attention,
+                    norm_type=self.cfg.norm_type,
+                    norm_elementwise_affine=self.cfg.norm_elementwise_affine,
+                )
+                for d in range(self.cfg.num_layers)
+            ]
+        )
+
+        # 4. Define output layers
+        self.out_channels = (
+            self.cfg.in_channels
+            if self.cfg.out_channels is None
+            else self.cfg.out_channels
+        )
+
+        self.proj_out = linear_cls(inner_dim, self.cfg.in_channels)
+
+        self.gradient_checkpointing = self.cfg.gradient_checkpointing
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        The [`Transformer1DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
+                Input `hidden_states`.
+            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
+                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
+                self-attention.
+            attention_mask ( `torch.Tensor`, *optional*):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            encoder_attention_mask ( `torch.Tensor`, *optional*):
+                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
+
+                    * Mask `(batch, sequence_length)` True = keep, False = discard.
+                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
+
+                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
+                above. This bias will be added to the cross-attention scores.
+
+        Returns:
+            torch.FloatTensor
+        """
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+        #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+        #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None and attention_mask.ndim == 2:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+            encoder_attention_mask = (
+                1 - encoder_attention_mask.to(hidden_states.dtype)
+            ) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 1. Input
+        batch, _, seq_len = hidden_states.shape
+        residual = hidden_states
+
+        hidden_states = self.norm(hidden_states)
+        inner_dim = hidden_states.shape[1]
+        hidden_states = hidden_states.permute(0, 2, 1).reshape(
+            batch, seq_len, inner_dim
+        )
+        hidden_states = self.proj_in(hidden_states)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            if self.training and self.gradient_checkpointing:
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    block,
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+
+        # 3. Output
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = (
+            hidden_states.reshape(batch, seq_len, inner_dim)
+            .permute(0, 2, 1)
+            .contiguous()
+        )
+
+        output = hidden_states + residual
+
+        return output

tsr/system.py

@@ -0,0 +1,203 @@
+import math
+import os
+from dataclasses import dataclass, field
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+import trimesh
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from omegaconf import OmegaConf
+from PIL import Image
+
+from .models.isosurface import MarchingCubeHelper
+from .utils import (
+    BaseModule,
+    ImagePreprocessor,
+    find_class,
+    get_spherical_cameras,
+    scale_tensor,
+)
+
+
+class TSR(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        cond_image_size: int
+
+        image_tokenizer_cls: str
+        image_tokenizer: dict
+
+        tokenizer_cls: str
+        tokenizer: dict
+
+        backbone_cls: str
+        backbone: dict
+
+        post_processor_cls: str
+        post_processor: dict
+
+        decoder_cls: str
+        decoder: dict
+
+        renderer_cls: str
+        renderer: dict
+
+    cfg: Config
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: str, config_name: str, weight_name: str
+    ):
+        if os.path.isdir(pretrained_model_name_or_path):
+            config_path = os.path.join(pretrained_model_name_or_path, config_name)
+            weight_path = os.path.join(pretrained_model_name_or_path, weight_name)
+        else:
+            config_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path, filename=config_name
+            )
+            weight_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path, filename=weight_name
+            )
+
+        cfg = OmegaConf.load(config_path)
+        OmegaConf.resolve(cfg)
+        model = cls(cfg)
+        ckpt = torch.load(weight_path, map_location="cpu")
+        model.load_state_dict(ckpt)
+        return model
+
+    def configure(self):
+        self.image_tokenizer = find_class(self.cfg.image_tokenizer_cls)(
+            self.cfg.image_tokenizer
+        )
+        self.tokenizer = find_class(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
+        self.backbone = find_class(self.cfg.backbone_cls)(self.cfg.backbone)
+        self.post_processor = find_class(self.cfg.post_processor_cls)(
+            self.cfg.post_processor
+        )
+        self.decoder = find_class(self.cfg.decoder_cls)(self.cfg.decoder)
+        self.renderer = find_class(self.cfg.renderer_cls)(self.cfg.renderer)
+        self.image_processor = ImagePreprocessor()
+        self.isosurface_helper = None
+
+    def forward(
+        self,
+        image: Union[
+            PIL.Image.Image,
+            np.ndarray,
+            torch.FloatTensor,
+            List[PIL.Image.Image],
+            List[np.ndarray],
+            List[torch.FloatTensor],
+        ],
+        device: str,
+    ) -> torch.FloatTensor:
+        rgb_cond = self.image_processor(image, self.cfg.cond_image_size)[:, None].to(
+            device
+        )
+        batch_size = rgb_cond.shape[0]
+
+        input_image_tokens: torch.Tensor = self.image_tokenizer(
+            rearrange(rgb_cond, "B Nv H W C -> B Nv C H W", Nv=1),
+        )
+
+        input_image_tokens = rearrange(
+            input_image_tokens, "B Nv C Nt -> B (Nv Nt) C", Nv=1
+        )
+
+        tokens: torch.Tensor = self.tokenizer(batch_size)
+
+        tokens = self.backbone(
+            tokens,
+            encoder_hidden_states=input_image_tokens,
+        )
+
+        scene_codes = self.post_processor(self.tokenizer.detokenize(tokens))
+        return scene_codes
+
+    def render(
+        self,
+        scene_codes,
+        n_views: int,
+        elevation_deg: float = 0.0,
+        camera_distance: float = 1.9,
+        fovy_deg: float = 40.0,
+        height: int = 256,
+        width: int = 256,
+        return_type: str = "pil",
+    ):
+        rays_o, rays_d = get_spherical_cameras(
+            n_views, elevation_deg, camera_distance, fovy_deg, height, width
+        )
+        rays_o, rays_d = rays_o.to(scene_codes.device), rays_d.to(scene_codes.device)
+
+        def process_output(image: torch.FloatTensor):
+            if return_type == "pt":
+                return image
+            elif return_type == "np":
+                return image.detach().cpu().numpy()
+            elif return_type == "pil":
+                return Image.fromarray(
+                    (image.detach().cpu().numpy() * 255.0).astype(np.uint8)
+                )
+            else:
+                raise NotImplementedError
+
+        images = []
+        for scene_code in scene_codes:
+            images_ = []
+            for i in range(n_views):
+                with torch.no_grad():
+                    image = self.renderer(
+                        self.decoder, scene_code, rays_o[i], rays_d[i]
+                    )
+                images_.append(process_output(image))
+            images.append(images_)
+
+        return images
+
+    def set_marching_cubes_resolution(self, resolution: int):
+        if (
+            self.isosurface_helper is not None
+            and self.isosurface_helper.resolution == resolution
+        ):
+            return
+        self.isosurface_helper = MarchingCubeHelper(resolution)
+
+    def extract_mesh(self, scene_codes, resolution: int = 256, threshold: float = 25.0):
+        self.set_marching_cubes_resolution(resolution)
+        meshes = []
+        for scene_code in scene_codes:
+            with torch.no_grad():
+                density = self.renderer.query_triplane(
+                    self.decoder,
+                    scale_tensor(
+                        self.isosurface_helper.grid_vertices.to(scene_codes.device),
+                        self.isosurface_helper.points_range,
+                        (-self.renderer.cfg.radius, self.renderer.cfg.radius),
+                    ),
+                    scene_code,
+                )["density_act"]
+            v_pos, t_pos_idx = self.isosurface_helper(-(density - threshold))
+            v_pos = scale_tensor(
+                v_pos,
+                self.isosurface_helper.points_range,
+                (-self.renderer.cfg.radius, self.renderer.cfg.radius),
+            )
+            with torch.no_grad():
+                color = self.renderer.query_triplane(
+                    self.decoder,
+                    v_pos,
+                    scene_code,
+                )["color"]
+            mesh = trimesh.Trimesh(
+                vertices=v_pos.cpu().numpy(),
+                faces=t_pos_idx.cpu().numpy(),
+                vertex_colors=color.cpu().numpy(),
+            )
+            meshes.append(mesh)
+        return meshes

tsr/utils.py

@@ -0,0 +1,474 @@
+import importlib
+import math
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import imageio
+import numpy as np
+import PIL.Image
+import rembg
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import trimesh
+from omegaconf import DictConfig, OmegaConf
+from PIL import Image
+
+
+def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
+    scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg)
+    return scfg
+
+
+def find_class(cls_string):
+    module_string = ".".join(cls_string.split(".")[:-1])
+    cls_name = cls_string.split(".")[-1]
+    module = importlib.import_module(module_string, package=None)
+    cls = getattr(module, cls_name)
+    return cls
+
+
+def get_intrinsic_from_fov(fov, H, W, bs=-1):
+    focal_length = 0.5 * H / np.tan(0.5 * fov)
+    intrinsic = np.identity(3, dtype=np.float32)
+    intrinsic[0, 0] = focal_length
+    intrinsic[1, 1] = focal_length
+    intrinsic[0, 2] = W / 2.0
+    intrinsic[1, 2] = H / 2.0
+
+    if bs > 0:
+        intrinsic = intrinsic[None].repeat(bs, axis=0)
+
+    return torch.from_numpy(intrinsic)
+
+
+class BaseModule(nn.Module):
+    @dataclass
+    class Config:
+        pass
+
+    cfg: Config  # add this to every subclass of BaseModule to enable static type checking
+
+    def __init__(
+        self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
+    ) -> None:
+        super().__init__()
+        self.cfg = parse_structured(self.Config, cfg)
+        self.configure(*args, **kwargs)
+
+    def configure(self, *args, **kwargs) -> None:
+        raise NotImplementedError
+
+
+class ImagePreprocessor:
+    def convert_and_resize(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+        size: int,
+    ):
+        if isinstance(image, PIL.Image.Image):
+            image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
+        elif isinstance(image, np.ndarray):
+            if image.dtype == np.uint8:
+                image = torch.from_numpy(image.astype(np.float32) / 255.0)
+            else:
+                image = torch.from_numpy(image)
+        elif isinstance(image, torch.Tensor):
+            pass
+
+        batched = image.ndim == 4
+
+        if not batched:
+            image = image[None, ...]
+        image = F.interpolate(
+            image.permute(0, 3, 1, 2),
+            (size, size),
+            mode="bilinear",
+            align_corners=False,
+            antialias=True,
+        ).permute(0, 2, 3, 1)
+        if not batched:
+            image = image[0]
+        return image
+
+    def __call__(
+        self,
+        image: Union[
+            PIL.Image.Image,
+            np.ndarray,
+            torch.FloatTensor,
+            List[PIL.Image.Image],
+            List[np.ndarray],
+            List[torch.FloatTensor],
+        ],
+        size: int,
+    ) -> Any:
+        if isinstance(image, (np.ndarray, torch.FloatTensor)) and image.ndim == 4:
+            image = self.convert_and_resize(image, size)
+        else:
+            if not isinstance(image, list):
+                image = [image]
+            image = [self.convert_and_resize(im, size) for im in image]
+            image = torch.stack(image, dim=0)
+        return image
+
+
+def rays_intersect_bbox(
+    rays_o: torch.Tensor,
+    rays_d: torch.Tensor,
+    radius: float,
+    near: float = 0.0,
+    valid_thresh: float = 0.01,
+):
+    input_shape = rays_o.shape[:-1]
+    rays_o, rays_d = rays_o.view(-1, 3), rays_d.view(-1, 3)
+    rays_d_valid = torch.where(
+        rays_d.abs() < 1e-6, torch.full_like(rays_d, 1e-6), rays_d
+    )
+    if type(radius) in [int, float]:
+        radius = torch.FloatTensor(
+            [[-radius, radius], [-radius, radius], [-radius, radius]]
+        ).to(rays_o.device)
+    radius = (
+        1.0 - 1.0e-3
+    ) * radius  # tighten the radius to make sure the intersection point lies in the bounding box
+    interx0 = (radius[..., 1] - rays_o) / rays_d_valid
+    interx1 = (radius[..., 0] - rays_o) / rays_d_valid
+    t_near = torch.minimum(interx0, interx1).amax(dim=-1).clamp_min(near)
+    t_far = torch.maximum(interx0, interx1).amin(dim=-1)
+
+    # check wheter a ray intersects the bbox or not
+    rays_valid = t_far - t_near > valid_thresh
+
+    t_near[torch.where(~rays_valid)] = 0.0
+    t_far[torch.where(~rays_valid)] = 0.0
+
+    t_near = t_near.view(*input_shape, 1)
+    t_far = t_far.view(*input_shape, 1)
+    rays_valid = rays_valid.view(*input_shape)
+
+    return t_near, t_far, rays_valid
+
+
+def chunk_batch(func: Callable, chunk_size: int, *args, **kwargs) -> Any:
+    if chunk_size <= 0:
+        return func(*args, **kwargs)
+    B = None
+    for arg in list(args) + list(kwargs.values()):
+        if isinstance(arg, torch.Tensor):
+            B = arg.shape[0]
+            break
+    assert (
+        B is not None
+    ), "No tensor found in args or kwargs, cannot determine batch size."
+    out = defaultdict(list)
+    out_type = None
+    # max(1, B) to support B == 0
+    for i in range(0, max(1, B), chunk_size):
+        out_chunk = func(
+            *[
+                arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
+                for arg in args
+            ],
+            **{
+                k: arg[i : i + chunk_size] if isinstance(arg, torch.Tensor) else arg
+                for k, arg in kwargs.items()
+            },
+        )
+        if out_chunk is None:
+            continue
+        out_type = type(out_chunk)
+        if isinstance(out_chunk, torch.Tensor):
+            out_chunk = {0: out_chunk}
+        elif isinstance(out_chunk, tuple) or isinstance(out_chunk, list):
+            chunk_length = len(out_chunk)
+            out_chunk = {i: chunk for i, chunk in enumerate(out_chunk)}
+        elif isinstance(out_chunk, dict):
+            pass
+        else:
+            print(
+                f"Return value of func must be in type [torch.Tensor, list, tuple, dict], get {type(out_chunk)}."
+            )
+            exit(1)
+        for k, v in out_chunk.items():
+            v = v if torch.is_grad_enabled() else v.detach()
+            out[k].append(v)
+
+    if out_type is None:
+        return None
+
+    out_merged: Dict[Any, Optional[torch.Tensor]] = {}
+    for k, v in out.items():
+        if all([vv is None for vv in v]):
+            # allow None in return value
+            out_merged[k] = None
+        elif all([isinstance(vv, torch.Tensor) for vv in v]):
+            out_merged[k] = torch.cat(v, dim=0)
+        else:
+            raise TypeError(
+                f"Unsupported types in return value of func: {[type(vv) for vv in v if not isinstance(vv, torch.Tensor)]}"
+            )
+
+    if out_type is torch.Tensor:
+        return out_merged[0]
+    elif out_type in [tuple, list]:
+        return out_type([out_merged[i] for i in range(chunk_length)])
+    elif out_type is dict:
+        return out_merged
+
+
+ValidScale = Union[Tuple[float, float], torch.FloatTensor]
+
+
+def scale_tensor(dat: torch.FloatTensor, inp_scale: ValidScale, tgt_scale: ValidScale):
+    if inp_scale is None:
+        inp_scale = (0, 1)
+    if tgt_scale is None:
+        tgt_scale = (0, 1)
+    if isinstance(tgt_scale, torch.FloatTensor):
+        assert dat.shape[-1] == tgt_scale.shape[-1]
+    dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0])
+    dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0]
+    return dat
+
+
+def get_activation(name) -> Callable:
+    if name is None:
+        return lambda x: x
+    name = name.lower()
+    if name == "none":
+        return lambda x: x
+    elif name == "exp":
+        return lambda x: torch.exp(x)
+    elif name == "sigmoid":
+        return lambda x: torch.sigmoid(x)
+    elif name == "tanh":
+        return lambda x: torch.tanh(x)
+    elif name == "softplus":
+        return lambda x: F.softplus(x)
+    else:
+        try:
+            return getattr(F, name)
+        except AttributeError:
+            raise ValueError(f"Unknown activation function: {name}")
+
+
+def get_ray_directions(
+    H: int,
+    W: int,
+    focal: Union[float, Tuple[float, float]],
+    principal: Optional[Tuple[float, float]] = None,
+    use_pixel_centers: bool = True,
+    normalize: bool = True,
+) -> torch.FloatTensor:
+    """
+    Get ray directions for all pixels in camera coordinate.
+    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
+               ray-tracing-generating-camera-rays/standard-coordinate-systems
+
+    Inputs:
+        H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
+    Outputs:
+        directions: (H, W, 3), the direction of the rays in camera coordinate
+    """
+    pixel_center = 0.5 if use_pixel_centers else 0
+
+    if isinstance(focal, float):
+        fx, fy = focal, focal
+        cx, cy = W / 2, H / 2
+    else:
+        fx, fy = focal
+        assert principal is not None
+        cx, cy = principal
+
+    i, j = torch.meshgrid(
+        torch.arange(W, dtype=torch.float32) + pixel_center,
+        torch.arange(H, dtype=torch.float32) + pixel_center,
+        indexing="xy",
+    )
+
+    directions = torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)
+
+    if normalize:
+        directions = F.normalize(directions, dim=-1)
+
+    return directions
+
+
+def get_rays(
+    directions,
+    c2w,
+    keepdim=False,
+    normalize=False,
+) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+    # Rotate ray directions from camera coordinate to the world coordinate
+    assert directions.shape[-1] == 3
+
+    if directions.ndim == 2:  # (N_rays, 3)
+        if c2w.ndim == 2:  # (4, 4)
+            c2w = c2w[None, :, :]
+        assert c2w.ndim == 3  # (N_rays, 4, 4) or (1, 4, 4)
+        rays_d = (directions[:, None, :] * c2w[:, :3, :3]).sum(-1)  # (N_rays, 3)
+        rays_o = c2w[:, :3, 3].expand(rays_d.shape)
+    elif directions.ndim == 3:  # (H, W, 3)
+        assert c2w.ndim in [2, 3]
+        if c2w.ndim == 2:  # (4, 4)
+            rays_d = (directions[:, :, None, :] * c2w[None, None, :3, :3]).sum(
+                -1
+            )  # (H, W, 3)
+            rays_o = c2w[None, None, :3, 3].expand(rays_d.shape)
+        elif c2w.ndim == 3:  # (B, 4, 4)
+            rays_d = (directions[None, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
+                -1
+            )  # (B, H, W, 3)
+            rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
+    elif directions.ndim == 4:  # (B, H, W, 3)
+        assert c2w.ndim == 3  # (B, 4, 4)
+        rays_d = (directions[:, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
+            -1
+        )  # (B, H, W, 3)
+        rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
+
+    if normalize:
+        rays_d = F.normalize(rays_d, dim=-1)
+    if not keepdim:
+        rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3)
+
+    return rays_o, rays_d
+
+
+def get_spherical_cameras(
+    n_views: int,
+    elevation_deg: float,
+    camera_distance: float,
+    fovy_deg: float,
+    height: int,
+    width: int,
+):
+    azimuth_deg = torch.linspace(0, 360.0, n_views + 1)[:n_views]
+    elevation_deg = torch.full_like(azimuth_deg, elevation_deg)
+    camera_distances = torch.full_like(elevation_deg, camera_distance)
+
+    elevation = elevation_deg * math.pi / 180
+    azimuth = azimuth_deg * math.pi / 180
+
+    # convert spherical coordinates to cartesian coordinates
+    # right hand coordinate system, x back, y right, z up
+    # elevation in (-90, 90), azimuth from +x to +y in (-180, 180)
+    camera_positions = torch.stack(
+        [
+            camera_distances * torch.cos(elevation) * torch.cos(azimuth),
+            camera_distances * torch.cos(elevation) * torch.sin(azimuth),
+            camera_distances * torch.sin(elevation),
+        ],
+        dim=-1,
+    )
+
+    # default scene center at origin
+    center = torch.zeros_like(camera_positions)
+    # default camera up direction as +z
+    up = torch.as_tensor([0, 0, 1], dtype=torch.float32)[None, :].repeat(n_views, 1)
+
+    fovy = torch.full_like(elevation_deg, fovy_deg) * math.pi / 180
+
+    lookat = F.normalize(center - camera_positions, dim=-1)
+    right = F.normalize(torch.cross(lookat, up), dim=-1)
+    up = F.normalize(torch.cross(right, lookat), dim=-1)
+    c2w3x4 = torch.cat(
+        [torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
+        dim=-1,
+    )
+    c2w = torch.cat([c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1)
+    c2w[:, 3, 3] = 1.0
+
+    # get directions by dividing directions_unit_focal by focal length
+    focal_length = 0.5 * height / torch.tan(0.5 * fovy)
+    directions_unit_focal = get_ray_directions(
+        H=height,
+        W=width,
+        focal=1.0,
+    )
+    directions = directions_unit_focal[None, :, :, :].repeat(n_views, 1, 1, 1)
+    directions[:, :, :, :2] = (
+        directions[:, :, :, :2] / focal_length[:, None, None, None]
+    )
+    # must use normalize=True to normalize directions here
+    rays_o, rays_d = get_rays(directions, c2w, keepdim=True, normalize=True)
+
+    return rays_o, rays_d
+
+
+def remove_background(
+    image: PIL.Image.Image,
+    rembg_session: Any = None,
+    force: bool = False,
+    **rembg_kwargs,
+) -> PIL.Image.Image:
+    do_remove = True
+    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
+        do_remove = False
+    do_remove = do_remove or force
+    if do_remove:
+        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
+    return image
+
+
+def resize_foreground(
+    image: PIL.Image.Image,
+    ratio: float,
+) -> PIL.Image.Image:
+    image = np.array(image)
+    assert image.shape[-1] == 4
+    alpha = np.where(image[..., 3] > 0)
+    y1, y2, x1, x2 = (
+        alpha[0].min(),
+        alpha[0].max(),
+        alpha[1].min(),
+        alpha[1].max(),
+    )
+    # crop the foreground
+    fg = image[y1:y2, x1:x2]
+    # pad to square
+    size = max(fg.shape[0], fg.shape[1])
+    ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
+    ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
+    new_image = np.pad(
+        fg,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+
+    # compute padding according to the ratio
+    new_size = int(new_image.shape[0] / ratio)
+    # pad to size, double side
+    ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
+    ph1, pw1 = new_size - size - ph0, new_size - size - pw0
+    new_image = np.pad(
+        new_image,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+    new_image = PIL.Image.fromarray(new_image)
+    return new_image
+
+
+def save_video(
+    frames: List[PIL.Image.Image],
+    output_path: str,
+    fps: int = 30,
+):
+    # use imageio to save video
+    frames = [np.array(frame) for frame in frames]
+    writer = imageio.get_writer(output_path, fps=fps)
+    for frame in frames:
+        writer.append_data(frame)
+    writer.close()
+
+
+def to_gradio_3d_orientation(mesh):
+    mesh.apply_transform(trimesh.transformations.rotation_matrix(-np.pi/2, [1, 0, 0]))
+    mesh.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2, [0, 1, 0]))
+    return mesh