first commit

app.py

@@ -0,0 +1,375 @@
+import gradio as gr
+import torch
+import numpy as np
+from functools import partial
+from typing import Optional
+from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
+from shap_e.diffusion.sample import sample_latents
+from shap_e.models.download import load_model, load_config
+from shap_e.util.notebooks import create_pan_cameras, decode_latent_mesh
+import trimesh
+import torch.nn as nn
+import os
+import random
+import warnings
+from huggingface_hub import hf_hub_download
+import hashlib
+
+import sys
+
+sys.tracebacklimit = 0
+def set_seed(seed=1024):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+
+def freeze_params(params):
+    for param in params:
+        param.requires_grad = False
+
+class Blocks(gr.Blocks):
+
+    def __init__(
+            self,
+            theme: str = "default",
+            analytics_enabled: Optional[bool] = None,
+            mode: str = "blocks",
+            title: str = "Gradio",
+            css: Optional[str] = None,
+            **kwargs,
+    ):
+        self.extra_configs = {
+            'thumbnail': kwargs.pop('thumbnail', ''),
+            'url': kwargs.pop('url', 'https://gradio.app/'),
+            'creator': kwargs.pop('creator', '@teamGradio'),
+        }
+
+        super(Blocks, self).__init__(theme, analytics_enabled, mode, title, css, **kwargs)
+        warnings.filterwarnings("ignore")
+
+    def get_config_file(self):
+        config = super(Blocks, self).get_config_file()
+
+        for k, v in self.extra_configs.items():
+            config[k] = v
+
+        return config
+def optimize_all(xm, models, initial_noise, noise_start_t, diffusion, latent_model, device, prompt, instruction, rand_seed):
+    state = {}
+    out_gen_1, out_gen_2, out_gen_3, out_gen_4, state = generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_seed, state)
+    edited_1, edited_2, edited_3, edited_4, state = _3d_editing(xm, models, diffusion, initial_noise, noise_start_t, device, instruction, rand_seed, state)
+    print(state)
+    return out_gen_1, out_gen_2, out_gen_3, out_gen_4, edited_1, edited_2, edited_3, edited_4
+def generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_seed, state):
+    set_seed(rand_seed)
+    batch_size = 4
+    guidance_scale = 15.0
+    xm.renderer.volume.bbox_max = torch.tensor([1.0, 1.0, 1.0]).to(device)
+    xm.renderer.volume.bbox_min = torch.tensor([-1.0, -1.0, -1.0]).to(device)
+    xm.renderer.volume.bbox = torch.stack([xm.renderer.volume.bbox_min, xm.renderer.volume.bbox_max])
+
+    print("prompt: ", prompt, "rand_seed: ", rand_seed, "state:",  state)
+    latents = sample_latents(
+        batch_size=batch_size,
+        model=latent_model,
+        diffusion=diffusion,
+        guidance_scale=guidance_scale,
+        model_kwargs=dict(texts=[prompt] * batch_size),
+        progress=True,
+        clip_denoised=True,
+        use_fp16=True,
+        use_karras=True,
+        karras_steps=64,
+        sigma_min=1e-3,
+        sigma_max=160,
+        s_churn=0,
+    )
+    prompt_hash = str(hashlib.sha256((prompt + '_' + str(rand_seed)).encode('utf-8')).hexdigest())
+    mesh_path = []
+    output_path = './logs'
+    os.makedirs(os.path.join(output_path, 'source'), exist_ok=True)
+    state['latent'] = []
+    state['prompt'] = prompt
+    state['rand_seed_1'] = rand_seed
+    for i, latent in enumerate(latents):
+
+        output_path_tmp = os.path.join(output_path, 'source', '{}_{}.obj'.format(prompt_hash, i))
+        t_obj = decode_latent_mesh(xm, latent).tri_mesh()
+        with open(output_path_tmp, 'w') as f:
+            t_obj.write_obj(f)
+
+        mesh = trimesh.load_mesh(output_path_tmp)
+        angle = np.radians(180)
+        axis = [0, 1, 0]
+        rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis)
+        mesh.apply_transform(rotation_matrix)
+        angle = np.radians(90)
+        axis = [1, 0, 0]
+        rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis)
+        mesh.apply_transform(rotation_matrix)
+        output_path_tmp = os.path.join(output_path, 'source', '{}_{}.obj'.format(prompt_hash, i))
+        mesh.export(output_path_tmp)
+        state['latent'].append(latent.clone().detach())
+        mesh_path.append(output_path_tmp)
+
+    return mesh_path[0], mesh_path[1], mesh_path[2], mesh_path[3], state
+
+def _3d_editing(xm, models, diffusion, initial_noise, start_t, device, instruction, rand_seed, state):
+    set_seed(rand_seed)
+    mesh_path = []
+    prompt = state['prompt']
+    rand_seed_1 = state['rand_seed_1']
+    print("prompt: ", prompt, "rand_seed: ", rand_seed, "instruction:", instruction, "state:",  state)
+    prompt_hash = str(hashlib.sha256((prompt + '_' + str(rand_seed_1) + '_' + instruction + '_' + str(rand_seed)).encode('utf-8')).hexdigest())
+    if 'santa' in instruction:
+        e_type = 'santa_hat'
+    elif 'rainbow' in instruction:
+        e_type = 'rainbow'
+    elif 'gold' in instruction:
+        e_type = 'golden'
+    elif 'lego' in instruction:
+        e_type = 'lego'
+    elif 'wooden' in instruction:
+        e_type = 'wooden'
+    elif 'cyber' in instruction:
+        e_type = 'cyber'
+
+    # import pdb; pdb.set_trace()
+    model = models[e_type].to(device)
+    noise_initial = initial_noise[e_type].to(device)
+    noise_start_t = start_t[e_type]
+    general_save_path = './logs/edited'
+    os.makedirs(general_save_path, exist_ok=True)
+    for i, latent in enumerate(state['latent']):
+        latent = latent.to(device)
+        text_embeddings_clip = model.cached_model_kwargs(1, dict(texts=[instruction]))
+        print("shape of latent: ", latent.clone().unsqueeze(0).shape, "instruction: ", instruction)
+        ref_latent = latent.clone().unsqueeze(0)
+        t_1 = torch.randint(noise_start_t, noise_start_t + 1, (1,), device=device).long()
+
+        noise_input = diffusion.q_sample(ref_latent, t_1, noise=noise_initial)
+        out_1 = diffusion.p_mean_variance(model, noise_input, t_1, clip_denoised=True,
+                                          model_kwargs=text_embeddings_clip,
+                                          condition_latents=ref_latent)
+
+        updated_latents = out_1['pred_xstart']
+
+        if 'santa' in instruction:
+            xm.renderer.volume.bbox_max = torch.tensor([1.0, 1.0, 1.25]).to(device)
+            xm.renderer.volume.bbox_min = torch.tensor([-1.0, -1.0, -1]).to(device)
+            xm.renderer.volume.bbox = torch.stack([xm.renderer.volume.bbox_min, xm.renderer.volume.bbox_max])
+
+        else:
+            xm.renderer.volume.bbox_max = torch.tensor([1.0, 1.0, 1.0]).to(device)
+            xm.renderer.volume.bbox_min = torch.tensor([-1.0, -1.0, -1.0]).to(device)
+            xm.renderer.volume.bbox = torch.stack([xm.renderer.volume.bbox_min, xm.renderer.volume.bbox_max])
+
+        for latent_idx, updated_latent in enumerate(updated_latents):
+            output_path = os.path.join(general_save_path, '{}_{}.obj'.format(prompt_hash, i))
+
+            t = decode_latent_mesh(xm, updated_latent).tri_mesh()
+            with open(output_path, 'w') as f:
+                t.write_obj(f)
+            mesh = trimesh.load_mesh(output_path)
+
+            angle = np.radians(180)
+            axis = [0, 1, 0]
+
+            rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis)
+            mesh.apply_transform(rotation_matrix)
+            angle = np.radians(90)
+            axis = [1, 0, 0]
+
+            rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis)
+            mesh.apply_transform(rotation_matrix)
+
+            output_path = os.path.join(general_save_path, '{}_{}.obj'.format(prompt_hash, i))
+            mesh.export(output_path)
+            mesh_path.append(output_path)
+    return mesh_path[0], mesh_path[1], mesh_path[2], mesh_path[3], state
+def main():
+
+    css = """
+    #img2img_image, #img2img_image > .fixed-height, #img2img_image > .fixed-height > div, #img2img_image > .fixed-height > div > img
+    {
+        height: var(--height) !important;
+        max-height: var(--height) !important;
+        min-height: var(--height) !important;
+    }
+    #paper-info a {
+        color:#008AD7;
+        text-decoration: none;
+    }
+    #paper-info a:hover {
+        cursor: pointer;
+        text-decoration: none;
+    }
+
+    .tooltip {
+        color: #555;
+        position: relative;
+        display: inline-block;
+        cursor: pointer;
+    }
+
+    .tooltip .tooltiptext {
+        visibility: hidden;
+        width: 400px;
+        background-color: #555;
+        color: #fff;
+        text-align: center;
+        padding: 5px;
+        border-radius: 5px;
+        position: absolute;
+        z-index: 1; /* Set z-index to 1 */
+        left: 10px;
+        top: 100%;
+        opacity: 0;
+        transition: opacity 0.3s;
+    }
+
+    .tooltip:hover .tooltiptext {
+        visibility: visible;
+        opacity: 1;
+        z-index: 9999; /* Set a high z-index value when hovering */
+    }
+
+
+    """
+
+    rescale_js = """
+    function(x) {
+        const root = document.querySelector('gradio-app').shadowRoot || document.querySelector('gradio-app');
+        let image_scale = parseFloat(root.querySelector('#image_scale input').value) || 1.0;
+        const image_width = root.querySelector('#img2img_image').clientWidth;
+        const target_height = parseInt(image_width * image_scale);
+        document.body.style.setProperty('--height', `${target_height}px`);
+        root.querySelectorAll('button.justify-center.rounded')[0].style.display='none';
+        root.querySelectorAll('button.justify-center.rounded')[1].style.display='none';
+        return x;
+    }
+    """
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    latent_model = load_model('text300M', device=device)
+    xm = load_model('transmitter', device=device)
+    diffusion = diffusion_from_config(load_config('diffusion'))
+    freeze_params(xm.parameters())
+    models = dict()
+    initial_noise = dict()
+    noise_start_t = dict()
+    editing_types = ['rainbow', 'santa_hat', 'lego', 'golden', 'wooden', 'cyber']
+
+    for editing_type in editing_types:
+        tmp_model = load_model('text300M', device=device)
+        with torch.no_grad():
+            new_proj = nn.Linear(1024 * 2, 1024, device=device, dtype=tmp_model.wrapped.input_proj.weight.dtype)
+            new_proj.weight = nn.Parameter(torch.zeros_like(new_proj.weight))
+            new_proj.weight[:, :1024].copy_(tmp_model.wrapped.input_proj.weight)  #
+            new_proj.bias = nn.Parameter(torch.zeros_like(new_proj.bias))
+            new_proj.bias[:1024].copy_(tmp_model.wrapped.input_proj.bias)
+            tmp_model.wrapped.input_proj = new_proj
+
+        ckp = torch.load(hf_hub_download(repo_id='silentchen/Shap_Editor', subfolder='single', filename='{}.pt'.format(editing_type)), map_location='cpu')
+        tmp_model.load_state_dict(ckp['model'])
+        noise_initial = ckp['initial_noise']['noise'].to(device)
+        initial_noise[editing_type] = noise_initial
+        noise_start_t[editing_type] = ckp['t_start']
+        models[editing_type] = tmp_model
+
+    with Blocks(
+            css=css,
+            analytics_enabled=False,
+            title="SHAPE-EDITOR demo",
+    ) as demo:
+        description = """<p style="text-align: center; font-weight: bold;">
+            <span style="font-size: 28px"> <span style="font-size: 140%">S</span>HAP-<span style="font-size: 140%">E</span>DITOR: Instruction-guided <br> Latent 3D Editing in Seconds</span>
+            <br>
+            <span style="font-size: 18px" id="paper-info">
+                [<a href=" " target="_blank">Project Page</a>]
+                [<a href=" " target="_blank">Paper</a>]
+                [<a href=" " target="_blank">GitHub</a>]
+            </span>
+        </p>
+        """
+        state = gr.State({})
+        gr.HTML(description)
+        with gr.Column():
+            with gr.Column():
+                gr.HTML('<span style="font-size: 20px; font-weight: bold">Step 1: generate original 3D object using Shap-E.</span>')
+                prompt = gr.Textbox(
+                    label="Text prompt for initial 3D generation", lines=1
+                )
+                gen_btn = gr.Button(value='Generate', scale=1)
+
+
+            with gr.Column():
+                gr.HTML('<span style="font-size: 20px; font-weight: bold">Generated 3D objects</span>')
+                with gr.Row():
+                    out_gen_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 1 (step 1)")
+                    out_gen_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 2 (step 1)")
+                    out_gen_3 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 3 (step 1)")
+                    out_gen_4 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  visible=True, label="3D Model 4 (step 1)")
+
+            with gr.Column(scale=1):
+                gr.HTML('<span style="font-size: 20px; font-weight: bold">Step 2: apply 3D editing with S</span>HAP-<span style="font-size: 140%">E</span>DITOR.</span>')
+
+                editing_choice = gr.Dropdown(
+                    ["Add a santa hat to it", "Make it look like made of gold", "Make the color of it look like rainbow", "Make it in cyberpunk style", "Make it wooden", "Make it look like make of lego"], value='Add a santa hat to it', multiselect=False, label="Editing effects", info="Select specific editing you want to apply!"
+                ),
+                apply_btn = gr.Button(value='Editing', scale=1)
+
+            with gr.Column(scale=3):
+                gr.HTML('<span style="font-size: 20px; font-weight: bold">Edited 3D objects</span>')
+                with gr.Row():
+                    edited_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 1 (step 2)")
+                    edited_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 2 (step 2)")
+                    edited_3 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 3 (step 2)")
+                    edited_4 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 4 (step 2)")
+
+
+            with gr.Accordion("Advanced Options", open=False):
+                rand_seed = gr.Slider(minimum=0, maximum=1000, step=1, value=445, label="Random seed")
+
+            gen_btn.click(
+                fn=partial(generate_3d_with_shap_e, xm, diffusion, latent_model, device),
+                inputs=[prompt, rand_seed, state],
+                outputs=[out_gen_1, out_gen_2, out_gen_3, out_gen_4, state],
+                queue=False)
+
+            apply_btn.click(
+                fn=partial(_3d_editing, xm, models, diffusion, initial_noise, noise_start_t, device),
+                inputs=[
+                    editing_choice[0], rand_seed, state
+                ],
+                outputs=[edited_1, edited_2, edited_3, edited_4, state],
+                queue=True
+            )
+        print("Generate examples...")
+        with gr.Column():
+            gr.Examples(
+                examples=[
+                    [   "a corgi",
+                        "Make the color of it look like rainbow",
+                        456,
+                    ],
+                    ["a penguin",
+                     "Make it look like make of lego",
+                     214,
+                     ],
+                ],
+                inputs=[prompt, editing_choice[0], rand_seed],
+                outputs=[out_gen_1, out_gen_2, out_gen_3, out_gen_4, edited_1, edited_2, edited_3, edited_4],
+                fn=partial(optimize_all, xm, models, initial_noise, noise_start_t, diffusion, latent_model, device),
+                cache_examples=True,
+            )
+
+
+    demo.queue(max_size=10, api_open=False)
+    demo.launch(share=True, show_api=False, show_error=True)
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,17 @@
+filelock
+pillow
+torch
+fire
+humanize
+requests
+tqdm
+matplot
+scikit-image
+scipy
+numpy
+blobfile
+clip @ git+https://github.com/openai/CLIP.git
+trimesh
+
+# gradio demo
+gradio

shap_e/diffusion/gaussian_diffusion.py

@@ -0,0 +1,1143 @@
+"""
+Based on https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+"""
+
+import math
+from typing import Any, Dict, Iterable, Optional, Sequence, Union
+
+import blobfile as bf
+import numpy as np
+import torch as th
+import yaml
+
+
+def diffusion_from_config(config: Union[str, Dict[str, Any]]) -> "GaussianDiffusion":
+    if isinstance(config, str):
+        with bf.BlobFile(config, "rb") as f:
+            obj = yaml.load(f, Loader=yaml.SafeLoader)
+        return diffusion_from_config(obj)
+
+    schedule = config["schedule"]
+    steps = config["timesteps"]
+    respace = config.get("respacing", None)
+    mean_type = config.get("mean_type", "epsilon")
+    betas = get_named_beta_schedule(schedule, steps, **config.get("schedule_args", {}))
+    channel_scales = config.get("channel_scales", None)
+    channel_biases = config.get("channel_biases", None)
+    if channel_scales is not None:
+        channel_scales = np.array(channel_scales)
+    if channel_biases is not None:
+        channel_biases = np.array(channel_biases)
+    kwargs = dict(
+        betas=betas,
+        model_mean_type=mean_type,
+        model_var_type="learned_range",
+        loss_type="mse",
+        channel_scales=channel_scales,
+        channel_biases=channel_biases,
+    )
+    if respace is None:
+        return GaussianDiffusion(**kwargs)
+    else:
+        return SpacedDiffusion(use_timesteps=space_timesteps(steps, respace), **kwargs)
+
+
+def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
+    """
+    This is the deprecated API for creating beta schedules.
+
+    See get_named_beta_schedule() for the new library of schedules.
+    """
+    if beta_schedule == "linear":
+        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps, **extra_args: float):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        return get_beta_schedule(
+            "linear",
+            beta_start=scale * 0.0001,
+            beta_end=scale * 0.02,
+            num_diffusion_timesteps=num_diffusion_timesteps,
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    elif schedule_name == "inv_parabola":
+        exponent = extra_args.get("power", 2.0)
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: 1 - t**exponent,
+        )
+    elif schedule_name == "translated_parabola":
+        exponent = extra_args.get("power", 2.0)
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: (1 - t) ** exponent,
+        )
+    elif schedule_name == "exp":
+        coefficient = extra_args.get("coefficient", -12.0)
+        return betas_for_alpha_bar(num_diffusion_timesteps, lambda t: math.exp(t * coefficient))
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    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].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    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(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(f"cannot create exactly {num_timesteps} steps with an integer stride")
+        elif section_counts.startswith("exact"):
+            res = set(int(x) for x in section_counts[len("exact") :].split(","))
+            for x in res:
+                if x < 0 or x >= num_timesteps:
+                    raise ValueError(f"timestep out of bounds: {x}")
+            return res
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(f"cannot divide section of {size} steps into {section_count}")
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+
+    Ported directly from here:
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+
+    :param betas: a 1-D array of betas for each diffusion timestep from T to 1.
+    :param model_mean_type: a string determining what the model outputs.
+    :param model_var_type: a string determining how variance is output.
+    :param loss_type: a string determining the loss function to use.
+    :param discretized_t0: if True, use discrete gaussian loss for t=0. Only
+                           makes sense for images.
+    :param channel_scales: a multiplier to apply to x_start in training_losses
+                           and sampling functions.
+    """
+
+    def __init__(
+        self,
+        *,
+        betas: Sequence[float],
+        model_mean_type: str,
+        model_var_type: str,
+        loss_type: str,
+        discretized_t0: bool = False,
+        channel_scales: Optional[np.ndarray] = None,
+        channel_biases: Optional[np.ndarray] = None,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+        self.discretized_t0 = discretized_t0
+        self.channel_scales = channel_scales
+        self.channel_biases = channel_biases
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
+        )
+
+    def get_sigmas(self, t):
+        return _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, t.shape)
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(
+        self, model, x, t, clip_denoised=False, denoised_fn=None, model_kwargs=None, condition_latents=None
+    ):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, t, **model_kwargs) if condition_latents is None else model(x, t, condition_latents, **model_kwargs)
+        if isinstance(model_output, tuple):
+            model_output, extra = model_output
+        else:
+            extra = None
+
+        if self.model_var_type in ["learned", "learned_range"]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            if self.model_var_type == "learned":
+                model_log_variance = model_var_values
+                model_variance = th.exp(model_log_variance)
+            else:
+                min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape)
+                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+                # The model_var_values is [-1, 1] for [min_var, max_var].
+                frac = (model_var_values + 1) / 2
+                model_log_variance = frac * max_log + (1 - frac) * min_log
+                model_variance = th.exp(model_log_variance)
+        else:
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                "fixed_large": (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                "fixed_small": (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            return x
+
+        if self.model_mean_type == "x_prev":
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+            )
+            model_mean = model_output
+        elif self.model_mean_type in ["x_start", "epsilon"]:
+            if self.model_mean_type == "x_start":
+                pred_xstart = process_xstart(model_output)
+            else:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                )
+            model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)
+        else:
+            raise NotImplementedError(self.model_mean_type)
+
+        assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+            "extra": extra,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_xstart_from_xprev(self, x_t, t, xprev):
+        assert x_t.shape == xprev.shape
+        return (  # (xprev - coef2*x_t) / coef1
+            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
+            - _extract_into_tensor(
+                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
+            )
+            * x_t
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, t, **(model_kwargs or {}))
+        new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **(model_kwargs or {}))
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t)
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        temp=1.0,
+    ):
+        """
+        Generate samples from the model.
+
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            temp=temp,
+        ):
+            final = sample
+        return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        temp=1.0,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device) * temp
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                )
+                yield self.unscale_out_dict(out)
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        temp=1.0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+            temp=temp,
+        ):
+            final = sample
+        return final["sample"]
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=False,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+        temp=1.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device) * temp
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield self.unscale_out_dict(out)
+                img = out["sample"]
+
+    def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False, model_kwargs=None):
+        """
+        Get a term for the variational lower-bound.
+
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(true_mean, true_log_variance_clipped, out["mean"], out["log_variance"])
+        kl = mean_flat(kl) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        if not self.discretized_t0:
+            decoder_nll = th.zeros_like(decoder_nll)
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = th.where((t == 0), decoder_nll, kl)
+        return {
+            "output": output,
+            "pred_xstart": out["pred_xstart"],
+            "extra": out["extra"],
+        }
+
+    def training_losses(
+        self, model, x_start, t, model_kwargs=None, noise=None
+    ) -> Dict[str, th.Tensor]:
+        """
+        Compute training losses for a single timestep.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        x_start = self.scale_channels(x_start)
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+
+        if self.loss_type == "kl" or self.loss_type == "rescaled_kl":
+            vb_terms = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )
+            terms["loss"] = vb_terms["output"]
+            if self.loss_type == "rescaled_kl":
+                terms["loss"] *= self.num_timesteps
+            extra = vb_terms["extra"]
+        elif self.loss_type == "mse" or self.loss_type == "rescaled_mse":
+            model_output = model(x_t, t, **model_kwargs)
+            if isinstance(model_output, tuple):
+                model_output, extra = model_output
+            else:
+                extra = {}
+
+            if self.model_var_type in [
+                "learned",
+                "learned_range",
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (
+                    B,
+                    C * 2,
+                    *x_t.shape[2:],
+                ), f"{model_output.shape} != {(B, C * 2, *x_t.shape[2:])}"
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == "rescaled_mse":
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            target = {
+                "x_prev": self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[0],
+                "x_start": x_start,
+                "epsilon": noise,
+            }[self.model_mean_type]
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2)
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                terms["loss"] = terms["mse"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        if "losses" in extra:
+            terms.update({k: loss for k, (loss, _scale) in extra["losses"].items()})
+            for loss, scale in extra["losses"].values():
+                terms["loss"] = terms["loss"] + loss * scale
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+
+        This term can't be optimized, as it only depends on the encoder.
+
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=False, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+    def scale_channels(self, x: th.Tensor) -> th.Tensor:
+        if self.channel_scales is not None:
+            x = x * th.from_numpy(self.channel_scales).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        if self.channel_biases is not None:
+            x = x + th.from_numpy(self.channel_biases).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        return x
+
+    def unscale_channels(self, x: th.Tensor) -> th.Tensor:
+        if self.channel_biases is not None:
+            x = x - th.from_numpy(self.channel_biases).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        if self.channel_scales is not None:
+            x = x / th.from_numpy(self.channel_scales).to(x).reshape(
+                [1, -1, *([1] * (len(x.shape) - 2))]
+            )
+        return x
+
+    def unscale_out_dict(
+        self, out: Dict[str, Union[th.Tensor, Any]]
+    ) -> Dict[str, Union[th.Tensor, Any]]:
+        return {
+            k: (self.unscale_channels(v) if isinstance(v, th.Tensor) else v) for k, v in out.items()
+        }
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+    :param use_timesteps: (unordered) timesteps from the original diffusion
+                          process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps: Iterable[int], **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(self, model, *args, **kwargs):
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(self, model, *args, **kwargs):
+        return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(model, self.timestep_map, self.original_num_steps)
+
+
+class _WrappedModel:
+    def __init__(self, model, timestep_map, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        return self.model(x, new_ts, **kwargs)
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res + th.zeros(broadcast_shape, device=timesteps.device)
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.flatten(1).mean(1)

shap_e/diffusion/k_diffusion.py

@@ -0,0 +1,426 @@
+"""
+Based on: https://github.com/crowsonkb/k-diffusion
+
+Copyright (c) 2022 Katherine Crowson
+
+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.
+"""
+
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion, mean_flat
+
+
+class KarrasDenoiser:
+    def __init__(self, sigma_data: float = 0.5):
+        self.sigma_data = sigma_data
+
+    def get_snr(self, sigmas):
+        return sigmas**-2
+
+    def get_sigmas(self, sigmas):
+        return sigmas
+
+    def get_scalings(self, sigma):
+        c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2)
+        c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2) ** 0.5
+        c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5
+        return c_skip, c_out, c_in
+
+    def training_losses(self, model, x_start, sigmas, model_kwargs=None, noise=None):
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+
+        terms = {}
+
+        dims = x_start.ndim
+        x_t = x_start + noise * append_dims(sigmas, dims)
+        c_skip, c_out, _ = [append_dims(x, dims) for x in self.get_scalings(sigmas)]
+        model_output, denoised = self.denoise(model, x_t, sigmas, **model_kwargs)
+        target = (x_start - c_skip * x_t) / c_out
+
+        terms["mse"] = mean_flat((model_output - target) ** 2)
+        terms["xs_mse"] = mean_flat((denoised - x_start) ** 2)
+
+        if "vb" in terms:
+            terms["loss"] = terms["mse"] + terms["vb"]
+        else:
+            terms["loss"] = terms["mse"]
+
+        return terms
+
+    def denoise(self, model, x_t, sigmas, **model_kwargs):
+        c_skip, c_out, c_in = [append_dims(x, x_t.ndim) for x in self.get_scalings(sigmas)]
+        rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)
+        model_output = model(c_in * x_t, rescaled_t, **model_kwargs)
+        denoised = c_out * model_output + c_skip * x_t
+        return model_output, denoised
+
+
+class GaussianToKarrasDenoiser:
+    def __init__(self, model, diffusion):
+        from scipy import interpolate
+
+        self.model = model
+        self.diffusion = diffusion
+        self.alpha_cumprod_to_t = interpolate.interp1d(
+            diffusion.alphas_cumprod, np.arange(0, diffusion.num_timesteps)
+        )
+
+    def sigma_to_t(self, sigma):
+        alpha_cumprod = 1.0 / (sigma**2 + 1)
+        if alpha_cumprod > self.diffusion.alphas_cumprod[0]:
+            return 0
+        elif alpha_cumprod <= self.diffusion.alphas_cumprod[-1]:
+            return self.diffusion.num_timesteps - 1
+        else:
+            return float(self.alpha_cumprod_to_t(alpha_cumprod))
+
+    def denoise(self, x_t, sigmas, clip_denoised=True, model_kwargs=None, condition_latents=None):
+        t = th.tensor(
+            [self.sigma_to_t(sigma) for sigma in sigmas.cpu().numpy()],
+            dtype=th.long,
+            device=sigmas.device,
+        )
+        c_in = append_dims(1.0 / (sigmas**2 + 1) ** 0.5, x_t.ndim)
+        out = self.diffusion.p_mean_variance(
+            self.model, x_t * c_in, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs, condition_latents=condition_latents
+        )
+        return None, out["pred_xstart"]
+
+
+def karras_sample(*args, **kwargs):
+    last = None
+    x_sequence = []
+    # print("kraras_sample_model_kwargs", kwargs["model_kwargs"]['embeddings'].shape)
+    for x in karras_sample_progressive(*args, **kwargs):
+        last = x["x"]
+        x_sequence.append(last)
+    return last, x_sequence
+
+
+
+def karras_sample_progressive(
+    diffusion,
+    model,
+    shape,
+    steps,
+    clip_denoised=True,
+    progress=False,
+    model_kwargs=None,
+    device=None,
+    sigma_min=0.002,
+    sigma_max=80,  # higher for highres?
+    rho=7.0,
+    sampler="heun",
+    s_churn=0.0,
+    s_tmin=0.0,
+    s_tmax=float("inf"),
+    s_noise=1.0,
+    guidance_scale=0.0,
+    condition_latent=None,
+    initial_noise=None,
+):
+    sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
+    # print("sigmas", sigmas.shape, sigmas)
+    if initial_noise is None:
+        x_T = th.randn(*shape, device=device) * sigma_max
+    else:
+        x_T = initial_noise.clone() * sigma_max
+    sample_fn = {"heun": sample_heun, "dpm": sample_dpm, "ancestral": sample_euler_ancestral}[
+        sampler
+    ]
+    if sampler != "ancestral":
+        sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise)
+    else:
+        sampler_args = {}
+
+    if isinstance(diffusion, KarrasDenoiser):
+        def denoiser(x_t, sigma):
+            _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
+            if clip_denoised:
+                denoised = denoised.clamp(-1, 1)
+            return denoised
+
+    elif isinstance(diffusion, GaussianDiffusion):
+        model = GaussianToKarrasDenoiser(model, diffusion)
+
+        def denoiser(x_t, sigma):
+            _, denoised = model.denoise(
+                x_t, sigma, clip_denoised=clip_denoised, model_kwargs=model_kwargs, condition_latents=condition_latent
+            )
+            return denoised
+
+    else:
+        raise NotImplementedError
+
+    if guidance_scale != 0 and guidance_scale != 1:
+
+        def guided_denoiser(x_t, sigma):
+            x_t = th.cat([x_t, x_t], dim=0)
+            sigma = th.cat([sigma, sigma], dim=0)
+            x_0 = denoiser(x_t, sigma)
+            cond_x_0, uncond_x_0 = th.split(x_0, len(x_0) // 2, dim=0)
+            x_0 = uncond_x_0 + guidance_scale * (cond_x_0 - uncond_x_0)
+            return x_0
+
+    else:
+        guided_denoiser = denoiser
+
+    for obj in sample_fn(
+        guided_denoiser,
+        x_T,
+        sigmas,
+        progress=progress,
+        condition_latent=condition_latent,
+        **sampler_args,
+    ):
+        if isinstance(diffusion, GaussianDiffusion):
+            # print("is gaussian diffusion", obj)
+            yield diffusion.unscale_out_dict(obj)
+        else:
+            yield obj
+
+
+def karras_sample_progressive_condition(
+    diffusion,
+    model,
+    shape,
+    steps,
+    clip_denoised=True,
+    progress=False,
+    model_kwargs=None,
+    device=None,
+    sigma_min=0.002,
+    sigma_max=80,  # higher for highres?
+    rho=7.0,
+    sampler="heun",
+    s_churn=0.0,
+    s_tmin=0.0,
+    s_tmax=float("inf"),
+    s_noise=1.0,
+    text_guidance_scale=0.0,
+    image_guidance_scale=0.0,
+    condition_latent=None,
+):
+    sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
+    x_T = th.randn(*shape, device=device) * sigma_max
+    sample_fn = {"heun": sample_heun, "dpm": sample_dpm, "ancestral": sample_euler_ancestral}[
+        sampler
+    ]
+    if sampler != "ancestral":
+        sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise)
+    else:
+        sampler_args = {}
+
+    if isinstance(diffusion, KarrasDenoiser):
+        def denoiser(x_t, sigma):
+            _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
+            if clip_denoised:
+                denoised = denoised.clamp(-1, 1)
+            return denoised
+
+    elif isinstance(diffusion, GaussianDiffusion):
+        model = GaussianToKarrasDenoiser(model, diffusion)
+
+        def denoiser(x_t, sigma):
+            _, denoised = model.denoise(
+                x_t, sigma, clip_denoised=clip_denoised, model_kwargs=model_kwargs, condition_latents=condition_latent
+            )
+            return denoised
+
+    else:
+        raise NotImplementedError
+
+    if (text_guidance_scale != 1.0 and text_guidance_scale != 0.0) or (image_guidance_scale != 1.0 and image_guidance_scale != 0.0):
+        def guided_denoiser(x_t, sigma):
+            x_t = th.cat([x_t, x_t, x_t], dim=0)
+            sigma = th.cat([sigma, sigma, sigma], dim=0)
+            x_0 = denoiser(x_t, sigma)
+            # import pdb; pdb.set_trace()
+            cond_x_0_text, cond_x_0_image, uncond_x_0 = th.chunk(x_0, 3, dim=0)
+            x_0 = uncond_x_0 + text_guidance_scale * (cond_x_0_text - cond_x_0_image) + image_guidance_scale * (cond_x_0_image - uncond_x_0)
+            return x_0
+
+    else:
+        guided_denoiser = denoiser
+
+    for obj in sample_fn(
+        guided_denoiser,
+        x_T,
+        sigmas,
+        progress=progress,
+        condition_latent=condition_latent,
+        **sampler_args,
+    ):
+        if isinstance(diffusion, GaussianDiffusion):
+            yield diffusion.unscale_out_dict(obj)
+        else:
+            yield obj
+def karras_sample_addition_condition(*args, **kwargs):
+    last = None
+    x_sequence = []
+    for x in karras_sample_progressive_condition(*args, **kwargs):
+        last = x["x"]
+        x_sequence.append(x["pred_xstart"])
+    return last, x_sequence
+
+def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device="cpu"):
+    """Constructs the noise schedule of Karras et al. (2022)."""
+    ramp = th.linspace(0, 1, n)
+    min_inv_rho = sigma_min ** (1 / rho)
+    max_inv_rho = sigma_max ** (1 / rho)
+    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+    return append_zero(sigmas).to(device)
+
+
+def to_d(x, sigma, denoised):
+    """Converts a denoiser output to a Karras ODE derivative."""
+    return (x - denoised) / append_dims(sigma, x.ndim)
+
+
+def get_ancestral_step(sigma_from, sigma_to):
+    """Calculates the noise level (sigma_down) to step down to and the amount
+    of noise to add (sigma_up) when doing an ancestral sampling step."""
+    sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
+    sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
+    return sigma_down, sigma_up
+
+
+@th.no_grad()
+def sample_euler_ancestral(model, x, sigmas, progress=False):
+    """Ancestral sampling with Euler method steps."""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+
+        indices = tqdm(indices)
+
+    for i in indices:
+        denoised = model(x, sigmas[i] * s_in)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
+        yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "pred_xstart": denoised}
+        d = to_d(x, sigmas[i], denoised)
+        # Euler method
+        dt = sigma_down - sigmas[i]
+        x = x + d * dt
+        x = x + th.randn_like(x) * sigma_up
+    yield {"x": x, "pred_xstart": x}
+
+
+@th.no_grad()
+def sample_heun(
+    denoiser,
+    x,
+    sigmas,
+    progress=False,
+    s_churn=0.0,
+    s_tmin=0.0,
+    s_tmax=float("inf"),
+    s_noise=1.0,
+    condition_latent=None,
+):
+    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+
+        indices = tqdm(indices)
+
+    for i in indices:
+        gamma = (
+            min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+        )
+        eps = th.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
+        denoised = denoiser(x, sigma_hat * s_in)
+        d = to_d(x, sigma_hat, denoised)
+        yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "pred_xstart": denoised}
+        dt = sigmas[i + 1] - sigma_hat
+        if sigmas[i + 1] == 0:
+            # Euler method
+            x = x + d * dt
+        else:
+            # Heun's method
+            x_2 = x + d * dt
+            denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)
+            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
+            d_prime = (d + d_2) / 2
+            x = x + d_prime * dt
+    yield {"x": x, "pred_xstart": denoised}
+
+
+@th.no_grad()
+def sample_dpm(
+    denoiser,
+    x,
+    sigmas,
+    progress=False,
+    s_churn=0.0,
+    s_tmin=0.0,
+    s_tmax=float("inf"),
+    s_noise=1.0,
+):
+    """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
+    s_in = x.new_ones([x.shape[0]])
+    indices = range(len(sigmas) - 1)
+    if progress:
+        from tqdm.auto import tqdm
+
+        indices = tqdm(indices)
+
+    for i in indices:
+        gamma = (
+            min(s_churn / (len(sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.0
+        )
+        eps = th.randn_like(x) * s_noise
+        sigma_hat = sigmas[i] * (gamma + 1)
+        if gamma > 0:
+            x = x + eps * (sigma_hat**2 - sigmas[i] ** 2) ** 0.5
+        denoised = denoiser(x, sigma_hat * s_in)
+        d = to_d(x, sigma_hat, denoised)
+        yield {"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigma_hat, "denoised": denoised}
+        # Midpoint method, where the midpoint is chosen according to a rho=3 Karras schedule
+        sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2) ** 3
+        dt_1 = sigma_mid - sigma_hat
+        dt_2 = sigmas[i + 1] - sigma_hat
+        x_2 = x + d * dt_1
+        denoised_2 = denoiser(x_2, sigma_mid * s_in)
+        d_2 = to_d(x_2, sigma_mid, denoised_2)
+        x = x + d_2 * dt_2
+    yield {"x": x, "pred_xstart": denoised}
+
+
+def append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
+    return x[(...,) + (None,) * dims_to_append]
+
+
+def append_zero(x):
+    return th.cat([x, x.new_zeros([1])])

shap_e/diffusion/sample.py

@@ -0,0 +1,160 @@
+from typing import Any, Callable, Dict, Optional, List
+
+import torch
+import torch.nn as nn
+
+from .gaussian_diffusion import GaussianDiffusion
+from .k_diffusion import karras_sample, karras_sample_addition_condition
+
+DEFAULT_KARRAS_STEPS = 64
+DEFAULT_KARRAS_SIGMA_MIN = 1e-3
+DEFAULT_KARRAS_SIGMA_MAX = 160
+DEFAULT_KARRAS_S_CHURN = 0.0
+
+
+def uncond_guide_model(
+    model: Callable[..., torch.Tensor], scale: float
+) -> Callable[..., torch.Tensor]:
+
+    def model_fn(x_t, ts, **kwargs):
+        half = x_t[: len(x_t) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = model(combined, ts, **kwargs)
+        cond_out, uncond_out = torch.chunk(model_out, 2, dim=0)
+        cond_out = uncond_out + scale * (cond_out - uncond_out)
+        return torch.cat([cond_out, cond_out], dim=0)
+
+    return model_fn
+
+
+def sample_latents(
+    *,
+    batch_size: int,
+    model: nn.Module,
+    diffusion: GaussianDiffusion,
+    model_kwargs: Dict[str, Any],
+    guidance_scale: float,
+    clip_denoised: bool,
+    use_fp16: bool,
+    use_karras: bool,
+    karras_steps: int,
+    sigma_min: float,
+    sigma_max: float,
+    s_churn: float,
+    device: Optional[torch.device] = None,
+    progress: bool = False,
+    initial_noise: Optional[torch.Tensor] = None,
+) -> (torch.Tensor, List[torch.Tensor]):
+    sample_shape = (batch_size, model.d_latent)
+
+    if device is None:
+        device = next(model.parameters()).device
+
+    if hasattr(model, "cached_model_kwargs"):
+        model_kwargs = model.cached_model_kwargs(batch_size, model_kwargs)
+    if guidance_scale != 1.0 and guidance_scale != 0.0:
+        for k, v in model_kwargs.copy().items():
+            # print(k, v.shape)
+            model_kwargs[k] = torch.cat([v, torch.zeros_like(v)], dim=0)
+
+    sample_shape = (batch_size, model.d_latent)
+    with torch.autocast(device_type=device.type, enabled=use_fp16):
+        if use_karras:
+            samples, sample_sequence = karras_sample(
+                diffusion=diffusion,
+                model=model,
+                shape=sample_shape,
+                steps=karras_steps,
+                clip_denoised=clip_denoised,
+                model_kwargs=model_kwargs,
+                device=device,
+                sigma_min=sigma_min,
+                sigma_max=sigma_max,
+                s_churn=s_churn,
+                guidance_scale=guidance_scale,
+                progress=progress,
+                initial_noise=initial_noise,
+            )
+        else:
+            internal_batch_size = batch_size
+            if guidance_scale != 1.0:
+                model = uncond_guide_model(model, guidance_scale)
+                internal_batch_size *= 2
+            samples = diffusion.p_sample_loop(
+                model,
+                shape=(internal_batch_size, *sample_shape[1:]),
+                model_kwargs=model_kwargs,
+                device=device,
+                clip_denoised=clip_denoised,
+                progress=progress,
+            )
+
+    return samples
+
+
+def sample_latents_with_additional_latent(
+    *,
+    batch_size: int,
+    model: nn.Module,
+    diffusion: GaussianDiffusion,
+    model_kwargs: Dict[str, Any],
+    text_guidance_scale: float,
+    image_guidance_scale: float,
+    clip_denoised: bool,
+    use_fp16: bool,
+    use_karras: bool,
+    karras_steps: int,
+    sigma_min: float,
+    sigma_max: float,
+    s_churn: float,
+    device: Optional[torch.device] = None,
+    progress: bool = False,
+    condition_latent: Optional[torch.Tensor] = None,
+) -> (torch.Tensor, List[torch.Tensor]):
+
+    if device is None:
+        device = next(model.parameters()).device
+
+    if hasattr(model, "cached_model_kwargs"):
+        model_kwargs = model.cached_model_kwargs(batch_size, model_kwargs)
+    if (text_guidance_scale != 1.0 and text_guidance_scale != 0.0) or (image_guidance_scale != 1.0 and image_guidance_scale != 0.0):
+        for k, v in model_kwargs.copy().items():
+            # print(k, v.shape)
+            model_kwargs[k] = torch.cat([v, torch.zeros_like(v), torch.zeros_like(v)], dim=0)
+            condition_latent = torch.cat([condition_latent, condition_latent, torch.zeros_like(condition_latent)], dim=0)
+
+    sample_shape = (batch_size, model.d_latent)
+    # print("sample_shape", sample_shape)
+    with torch.autocast(device_type=device.type, enabled=use_fp16):
+        if use_karras:
+            samples, samples_squence = karras_sample_addition_condition(
+                diffusion=diffusion,
+                model=model,
+                shape=sample_shape,
+                steps=karras_steps,
+                clip_denoised=clip_denoised,
+                model_kwargs=model_kwargs,
+                device=device,
+                sigma_min=sigma_min,
+                sigma_max=sigma_max,
+                s_churn=s_churn,
+                text_guidance_scale=text_guidance_scale,
+                image_guidance_scale=image_guidance_scale,
+                progress=progress,
+                condition_latent=condition_latent,
+            )
+        else:
+            internal_batch_size = batch_size
+            if text_guidance_scale != 1.0:
+                model = uncond_guide_model(model, text_guidance_scale)
+                internal_batch_size *= 2
+            samples = diffusion.p_sample_loop(
+                model,
+                shape=(internal_batch_size, *sample_shape[1:]),
+                model_kwargs=model_kwargs,
+                device=device,
+                clip_denoised=clip_denoised,
+                progress=progress,
+            )
+
+    return samples

shap_e/examples/encode_model.ipynb

@@ -0,0 +1,93 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "from shap_e.models.download import load_model\n",
+    "from shap_e.util.data_util import load_or_create_multimodal_batch\n",
+    "from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, gif_widget"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "xm = load_model('transmitter', device=device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_path = \"example_data/cactus/object.obj\"\n",
+    "\n",
+    "# This may take a few minutes, since it requires rendering the model twice\n",
+    "# in two different modes.\n",
+    "batch = load_or_create_multimodal_batch(\n",
+    "    device,\n",
+    "    model_path=model_path,\n",
+    "    mv_light_mode=\"basic\",\n",
+    "    mv_image_size=256,\n",
+    "    cache_dir=\"example_data/cactus/cached\",\n",
+    "    verbose=True, # this will show Blender output during renders\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with torch.no_grad():\n",
+    "    latent = xm.encoder.encode_to_bottleneck(batch)\n",
+    "\n",
+    "    render_mode = 'stf' # you can change this to 'nerf'\n",
+    "    size = 128 # recommended that you lower resolution when using nerf\n",
+    "\n",
+    "    cameras = create_pan_cameras(size, device)\n",
+    "    images = decode_latent_images(xm, latent, cameras, rendering_mode=render_mode)\n",
+    "    display(gif_widget(images))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

shap_e/examples/sample_image_to_3d.ipynb

@@ -0,0 +1,125 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "964ccced",
+   "metadata": {
+    "pycharm": {
+     "is_executing": true
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "from shap_e.diffusion.sample import sample_latents\n",
+    "from shap_e.diffusion.gaussian_diffusion import diffusion_from_config\n",
+    "from shap_e.models.download import load_model, load_config\n",
+    "from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, gif_widget\n",
+    "from shap_e.util.image_util import load_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8eed3a76",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d922637",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "xm = load_model('transmitter', device=device)\n",
+    "model = load_model('image300M', device=device)\n",
+    "diffusion = diffusion_from_config(load_config('diffusion'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53d329d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size = 4\n",
+    "guidance_scale = 3.0\n",
+    "\n",
+    "image = load_image(\"example_data/corgi.png\")\n",
+    "\n",
+    "latents = sample_latents(\n",
+    "    batch_size=batch_size,\n",
+    "    model=model,\n",
+    "    diffusion=diffusion,\n",
+    "    guidance_scale=guidance_scale,\n",
+    "    model_kwargs=dict(images=[image] * batch_size),\n",
+    "    progress=True,\n",
+    "    clip_denoised=True,\n",
+    "    use_fp16=True,\n",
+    "    use_karras=True,\n",
+    "    karras_steps=64,\n",
+    "    sigma_min=1e-3,\n",
+    "    sigma_max=160,\n",
+    "    s_churn=0,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "633da2ec",
+   "metadata": {
+    "pycharm": {
+     "is_executing": true
+    }
+   },
+   "outputs": [],
+   "source": [
+    "render_mode = 'nerf' # you can change this to 'stf' for mesh rendering\n",
+    "size = 64 # this is the size of the renders; higher values take longer to render.\n",
+    "\n",
+    "cameras = create_pan_cameras(size, device)\n",
+    "for i, latent in enumerate(latents):\n",
+    "    images = decode_latent_images(xm, latent, cameras, rendering_mode=render_mode)\n",
+    "    display(gif_widget(images))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

shap_e/examples/sample_text_to_3d.ipynb

@@ -0,0 +1,124 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "964ccced",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "from shap_e.diffusion.sample import sample_latents\n",
+    "from shap_e.diffusion.gaussian_diffusion import diffusion_from_config\n",
+    "from shap_e.models.download import load_model, load_config\n",
+    "from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, gif_widget"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8eed3a76",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d922637",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "xm = load_model('transmitter', device=device)\n",
+    "model = load_model('text300M', device=device)\n",
+    "diffusion = diffusion_from_config(load_config('diffusion'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53d329d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size = 4\n",
+    "guidance_scale = 15.0\n",
+    "prompt = \"a shark\"\n",
+    "\n",
+    "latents = sample_latents(\n",
+    "    batch_size=batch_size,\n",
+    "    model=model,\n",
+    "    diffusion=diffusion,\n",
+    "    guidance_scale=guidance_scale,\n",
+    "    model_kwargs=dict(texts=[prompt] * batch_size),\n",
+    "    progress=True,\n",
+    "    clip_denoised=True,\n",
+    "    use_fp16=True,\n",
+    "    use_karras=True,\n",
+    "    karras_steps=64,\n",
+    "    sigma_min=1e-3,\n",
+    "    sigma_max=160,\n",
+    "    s_churn=0,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "633da2ec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "render_mode = 'nerf' # you can change this to 'stf'\n",
+    "size = 64 # this is the size of the renders; higher values take longer to render.\n",
+    "\n",
+    "cameras = create_pan_cameras(size, device)\n",
+    "for i, latent in enumerate(latents):\n",
+    "    images = decode_latent_images(xm, latent, cameras, rendering_mode=render_mode)\n",
+    "    display(gif_widget(images))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "85a4dce4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Example of saving the latents as meshes.\n",
+    "from shap_e.util.notebooks import decode_latent_mesh\n",
+    "\n",
+    "for i, latent in enumerate(latents):\n",
+    "    t = decode_latent_mesh(xm, latent).tri_mesh()\n",
+    "    with open(f'example_mesh_{i}.ply', 'wb') as f:\n",
+    "        t.write_ply(f)\n",
+    "    with open(f'example_mesh_{i}.obj', 'w') as f:\n",
+    "        t.write_obj(f)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

shap_e/models/configs.py

@@ -0,0 +1,166 @@
+from typing import Any, Dict, Union
+
+import blobfile as bf
+import torch
+import torch.nn as nn
+import yaml
+
+from shap_e.models.generation.latent_diffusion import SplitVectorDiffusion
+from shap_e.models.generation.perceiver import PointDiffusionPerceiver
+from shap_e.models.generation.pooled_mlp import PooledMLP
+from shap_e.models.generation.transformer import (
+    CLIPImageGridPointDiffusionTransformer,
+    CLIPImageGridUpsamplePointDiffusionTransformer,
+    CLIPImagePointDiffusionTransformer,
+    PointDiffusionTransformer,
+    UpsamplePointDiffusionTransformer,
+)
+from shap_e.models.nerf.model import MLPNeRFModel, VoidNeRFModel
+from shap_e.models.nerf.renderer import OneStepNeRFRenderer, TwoStepNeRFRenderer
+from shap_e.models.nerstf.mlp import MLPDensitySDFModel, MLPNeRSTFModel
+from shap_e.models.nerstf.renderer import NeRSTFRenderer
+from shap_e.models.nn.meta import batch_meta_state_dict
+from shap_e.models.stf.mlp import MLPSDFModel, MLPTextureFieldModel
+from shap_e.models.stf.renderer import STFRenderer
+from shap_e.models.transmitter.base import ChannelsDecoder, Transmitter, VectorDecoder
+from shap_e.models.transmitter.channels_encoder import (
+    PointCloudPerceiverChannelsEncoder,
+    PointCloudTransformerChannelsEncoder,
+)
+from shap_e.models.transmitter.multiview_encoder import MultiviewTransformerEncoder
+from shap_e.models.transmitter.pc_encoder import (
+    PointCloudPerceiverEncoder,
+    PointCloudTransformerEncoder,
+)
+from shap_e.models.volume import BoundingBoxVolume, SphericalVolume, UnboundedVolume
+
+
+def model_from_config(config: Union[str, Dict[str, Any]], device: torch.device) -> nn.Module:
+    print(config)
+    if isinstance(config, str):
+        print("config", config)
+        with bf.BlobFile(config, "rb") as f:
+            obj = yaml.load(f, Loader=yaml.SafeLoader)
+        return model_from_config(obj, device=device)
+
+    config = config.copy()
+    name = config.pop("name")
+
+    if name == "PointCloudTransformerEncoder":
+        return PointCloudTransformerEncoder(device=device, dtype=torch.float32, **config)
+    elif name == "PointCloudPerceiverEncoder":
+        return PointCloudPerceiverEncoder(device=device, dtype=torch.float32, **config)
+    elif name == "PointCloudTransformerChannelsEncoder":
+        return PointCloudTransformerChannelsEncoder(device=device, dtype=torch.float32, **config)
+    elif name == "PointCloudPerceiverChannelsEncoder":
+        return PointCloudPerceiverChannelsEncoder(device=device, dtype=torch.float32, **config)
+    elif name == "MultiviewTransformerEncoder":
+        return MultiviewTransformerEncoder(device=device, dtype=torch.float32, **config)
+    elif name == "Transmitter":
+        renderer = model_from_config(config.pop("renderer"), device=device)
+        param_shapes = {
+            k: v.shape[1:] for k, v in batch_meta_state_dict(renderer, batch_size=1).items()
+        }
+        encoder_config = config.pop("encoder").copy()
+        encoder_config["param_shapes"] = param_shapes
+        encoder = model_from_config(encoder_config, device=device)
+        return Transmitter(encoder=encoder, renderer=renderer, **config)
+    elif name == "VectorDecoder":
+        renderer = model_from_config(config.pop("renderer"), device=device)
+        param_shapes = {
+            k: v.shape[1:] for k, v in batch_meta_state_dict(renderer, batch_size=1).items()
+        }
+        return VectorDecoder(param_shapes=param_shapes, renderer=renderer, device=device, **config)
+    elif name == "ChannelsDecoder":
+        renderer = model_from_config(config.pop("renderer"), device=device)
+        param_shapes = {
+            k: v.shape[1:] for k, v in batch_meta_state_dict(renderer, batch_size=1).items()
+        }
+        return ChannelsDecoder(
+            param_shapes=param_shapes, renderer=renderer, device=device, **config
+        )
+    elif name == "OneStepNeRFRenderer":
+        config = config.copy()
+        for field in [
+            # Required
+            "void_model",
+            "foreground_model",
+            "volume",
+            # Optional to use NeRF++
+            "background_model",
+            "outer_volume",
+        ]:
+            if field in config:
+                config[field] = model_from_config(config.pop(field).copy(), device)
+        return OneStepNeRFRenderer(device=device, **config)
+    elif name == "TwoStepNeRFRenderer":
+        config = config.copy()
+        for field in [
+            # Required
+            "void_model",
+            "coarse_model",
+            "fine_model",
+            "volume",
+            # Optional to use NeRF++
+            "coarse_background_model",
+            "fine_background_model",
+            "outer_volume",
+        ]:
+            if field in config:
+                config[field] = model_from_config(config.pop(field).copy(), device)
+        return TwoStepNeRFRenderer(device=device, **config)
+    elif name == "PooledMLP":
+        return PooledMLP(device, **config)
+    elif name == "PointDiffusionTransformer":
+        return PointDiffusionTransformer(device=device, dtype=torch.float32, **config)
+    elif name == "PointDiffusionPerceiver":
+        return PointDiffusionPerceiver(device=device, dtype=torch.float32, **config)
+    elif name == "CLIPImagePointDiffusionTransformer":
+        return CLIPImagePointDiffusionTransformer(device=device, dtype=torch.float32, **config)
+    elif name == "CLIPImageGridPointDiffusionTransformer":
+        return CLIPImageGridPointDiffusionTransformer(device=device, dtype=torch.float32, **config)
+    elif name == "UpsamplePointDiffusionTransformer":
+        return UpsamplePointDiffusionTransformer(device=device, dtype=torch.float32, **config)
+    elif name == "CLIPImageGridUpsamplePointDiffusionTransformer":
+        return CLIPImageGridUpsamplePointDiffusionTransformer(
+            device=device, dtype=torch.float32, **config
+        )
+    elif name == "SplitVectorDiffusion":
+        inner_config = config.pop("inner")
+        d_latent = config.pop("d_latent")
+        latent_ctx = config.pop("latent_ctx", 1)
+        inner_config["input_channels"] = d_latent // latent_ctx
+        inner_config["n_ctx"] = latent_ctx
+        inner_config["output_channels"] = d_latent // latent_ctx * 2
+        inner_model = model_from_config(inner_config, device)
+        return SplitVectorDiffusion(
+            device=device, wrapped=inner_model, n_ctx=latent_ctx, d_latent=d_latent
+        )
+    elif name == "STFRenderer":
+        config = config.copy()
+        for field in ["sdf", "tf", "volume"]:
+            config[field] = model_from_config(config.pop(field), device)
+        return STFRenderer(device=device, **config)
+    elif name == "NeRSTFRenderer":
+        config = config.copy()
+        for field in ["sdf", "tf", "nerstf", "void", "volume"]:
+            if field not in config:
+                continue
+            config[field] = model_from_config(config.pop(field), device)
+        config.setdefault("sdf", None)
+        config.setdefault("tf", None)
+        config.setdefault("nerstf", None)
+        return NeRSTFRenderer(device=device, **config)
+
+    model_cls = {
+        "MLPSDFModel": MLPSDFModel,
+        "MLPTextureFieldModel": MLPTextureFieldModel,
+        "MLPNeRFModel": MLPNeRFModel,
+        "MLPDensitySDFModel": MLPDensitySDFModel,
+        "MLPNeRSTFModel": MLPNeRSTFModel,
+        "VoidNeRFModel": VoidNeRFModel,
+        "BoundingBoxVolume": BoundingBoxVolume,
+        "SphericalVolume": SphericalVolume,
+        "UnboundedVolume": UnboundedVolume,
+    }[name]
+    return model_cls(device=device, **config)

shap_e/models/download.py

@@ -0,0 +1,152 @@
+"""
+Adapted from: https://github.com/openai/glide-text2im/blob/69b530740eb6cef69442d6180579ef5ba9ef063e/glide_text2im/download.py
+"""
+
+import hashlib
+import os
+from functools import lru_cache
+from typing import Dict, Optional
+
+import requests
+import torch
+import yaml
+from filelock import FileLock
+from tqdm.auto import tqdm
+
+MODEL_PATHS = {
+    "transmitter": "https://openaipublic.azureedge.net/main/shap-e/transmitter.pt",
+    "decoder": "https://openaipublic.azureedge.net/main/shap-e/vector_decoder.pt",
+    "text300M": "https://openaipublic.azureedge.net/main/shap-e/text_cond.pt",
+    "image300M": "https://openaipublic.azureedge.net/main/shap-e/image_cond.pt",
+}
+
+CONFIG_PATHS = {
+    "transmitter": "https://openaipublic.azureedge.net/main/shap-e/transmitter_config.yaml",
+    "decoder": "https://openaipublic.azureedge.net/main/shap-e/vector_decoder_config.yaml",
+    "text300M": "https://openaipublic.azureedge.net/main/shap-e/text_cond_config.yaml",
+    "image300M": "https://openaipublic.azureedge.net/main/shap-e/image_cond_config.yaml",
+    "diffusion": "https://openaipublic.azureedge.net/main/shap-e/diffusion_config.yaml",
+}
+
+URL_HASHES = {
+    "https://openaipublic.azureedge.net/main/shap-e/transmitter.pt": "af02a0b85a8abdfb3919584b63c540ba175f6ad4790f574a7fef4617e5acdc3b",
+    "https://openaipublic.azureedge.net/main/shap-e/vector_decoder.pt": "d7e7ebbfe3780499ae89b2da5e7c1354012dba5a6abfe295bed42f25c3be1b98",
+    "https://openaipublic.azureedge.net/main/shap-e/text_cond.pt": "e6b4fa599a7b3c3b16c222d5f5fe56f9db9289ff0b6575fbe5c11bc97106aad4",
+    "https://openaipublic.azureedge.net/main/shap-e/image_cond.pt": "cb8072c64bbbcf6910488814d212227de5db291780d4ea99c6152f9346cf12aa",
+    "https://openaipublic.azureedge.net/main/shap-e/transmitter_config.yaml": "ffe1bcb405104a37d9408391182ab118a4ef313c391e07689684f1f62071605e",
+    "https://openaipublic.azureedge.net/main/shap-e/vector_decoder_config.yaml": "e6d373649f8e24d85925f4674b9ac41c57aba5f60e42cde6d10f87381326365c",
+    "https://openaipublic.azureedge.net/main/shap-e/text_cond_config.yaml": "f290beeea3d3e9ff15db01bde5382b6e549e463060c0744f89c049505be246c1",
+    "https://openaipublic.azureedge.net/main/shap-e/image_cond_config.yaml": "4e0745605a533c543c72add803a78d233e2a6401e0abfa0cad58afb4d74ad0b0",
+    "https://openaipublic.azureedge.net/main/shap-e/diffusion_config.yaml": "efcb2cd7ee545b2d27223979d41857802448143990572a42645cd09c2942ed57",
+}
+
+
+@lru_cache()
+def default_cache_dir() -> str:
+    return os.path.join(os.path.abspath(os.getcwd()), "shap_e_model_cache")
+
+
+def fetch_file_cached(
+    url: str, progress: bool = True, cache_dir: Optional[str] = None, chunk_size: int = 4096
+) -> str:
+    """
+    Download the file at the given URL into a local file and return the path.
+    If cache_dir is specified, it will be used to download the files.
+    Otherwise, default_cache_dir() is used.
+    """
+    expected_hash = URL_HASHES[url]
+
+    if cache_dir is None:
+        cache_dir = default_cache_dir()
+    os.makedirs(cache_dir, exist_ok=True)
+    local_path = os.path.join(cache_dir, url.split("/")[-1])
+    if os.path.exists(local_path):
+        check_hash(local_path, expected_hash)
+        return local_path
+
+    response = requests.get(url, stream=True)
+    size = int(response.headers.get("content-length", "0"))
+    with FileLock(local_path + ".lock"):
+        if progress:
+            pbar = tqdm(total=size, unit="iB", unit_scale=True)
+        tmp_path = local_path + ".tmp"
+        with open(tmp_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size):
+                if progress:
+                    pbar.update(len(chunk))
+                f.write(chunk)
+        os.rename(tmp_path, local_path)
+        if progress:
+            pbar.close()
+        check_hash(local_path, expected_hash)
+        return local_path
+
+
+def check_hash(path: str, expected_hash: str):
+    actual_hash = hash_file(path)
+    if actual_hash != expected_hash:
+        raise RuntimeError(
+            f"The file {path} should have hash {expected_hash} but has {actual_hash}. "
+            "Try deleting it and running this call again."
+        )
+
+
+def hash_file(path: str) -> str:
+    sha256_hash = hashlib.sha256()
+    with open(path, "rb") as file:
+        while True:
+            data = file.read(4096)
+            if not len(data):
+                break
+            sha256_hash.update(data)
+    return sha256_hash.hexdigest()
+
+
+def load_config(
+    config_name: str,
+    progress: bool = False,
+    cache_dir: Optional[str] = None,
+    chunk_size: int = 4096,
+):
+    if config_name not in CONFIG_PATHS:
+        raise ValueError(
+            f"Unknown config name {config_name}. Known names are: {CONFIG_PATHS.keys()}."
+        )
+    path = fetch_file_cached(
+        CONFIG_PATHS[config_name], progress=progress, cache_dir=cache_dir, chunk_size=chunk_size
+    )
+    with open(path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def load_checkpoint(
+    checkpoint_name: str,
+    device: torch.device,
+    progress: bool = True,
+    cache_dir: Optional[str] = None,
+    chunk_size: int = 4096,
+) -> Dict[str, torch.Tensor]:
+    if checkpoint_name not in MODEL_PATHS:
+        raise ValueError(
+            f"Unknown checkpoint name {checkpoint_name}. Known names are: {MODEL_PATHS.keys()}."
+        )
+    print(checkpoint_name)
+    path = fetch_file_cached(
+        MODEL_PATHS[checkpoint_name], progress=progress, cache_dir=cache_dir, chunk_size=chunk_size
+    )
+    return torch.load(path, map_location=device)
+
+
+def load_model(
+    model_name: str,
+    device: torch.device,
+    **kwargs,
+) -> Dict[str, torch.Tensor]:
+    from .configs import model_from_config
+
+    model = model_from_config(load_config(model_name, **kwargs), device=device)
+    # print(model_name, kwargs)
+    # print(model)
+    model.load_state_dict(load_checkpoint(model_name, device=device, **kwargs))
+    model.eval()
+    return model

shap_e/models/generation/latent_diffusion.py

@@ -0,0 +1,32 @@
+from typing import Any, Dict
+
+import torch
+import torch.nn as nn
+from typing import Any, Callable, Dict, Optional
+
+
+class SplitVectorDiffusion(nn.Module):
+    def __init__(self, *, device: torch.device, wrapped: nn.Module, n_ctx: int, d_latent: int):
+        super().__init__()
+        self.device = device
+        self.n_ctx = n_ctx
+        self.d_latent = d_latent
+        self.wrapped = wrapped
+
+        if hasattr(self.wrapped, "cached_model_kwargs"):
+            self.cached_model_kwargs = self.wrapped.cached_model_kwargs
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, conditional_latent: Optional[torch.Tensor] = None, **kwargs):
+        h = x.reshape(x.shape[0], self.n_ctx, -1).permute(0, 2, 1)
+        if conditional_latent is not None:
+            conditional_latent = conditional_latent.reshape(conditional_latent.shape[0], self.n_ctx, -1)
+            h = torch.cat([h.permute(0, 2, 1) , conditional_latent], dim=-1).permute(0, 2, 1) # (batch_size, n_ctx, channel) -> (batch_size, d_latent, n_ctx)
+        h = self.wrapped(h, t, **kwargs)
+        eps, var = torch.chunk(h, 2, dim=1)
+        return torch.cat(
+            [
+                eps.permute(0, 2, 1).flatten(1),
+                var.permute(0, 2, 1).flatten(1),
+            ],
+            dim=1,
+        )

shap_e/models/generation/perceiver.py

@@ -0,0 +1,244 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from shap_e.models.nn.checkpoint import checkpoint
+
+from .transformer import MLP, Transformer, init_linear
+from .util import timestep_embedding
+
+
+class MultiheadCrossAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int,
+        n_data: int,
+        width: int,
+        heads: int,
+        init_scale: float,
+        data_width: Optional[int] = None,
+    ):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.n_data = n_data
+        self.width = width
+        self.heads = heads
+        self.data_width = width if data_width is None else data_width
+        self.c_q = nn.Linear(width, width, device=device, dtype=dtype)
+        self.c_kv = nn.Linear(self.data_width, width * 2, device=device, dtype=dtype)
+        self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
+        self.attention = QKVMultiheadCrossAttention(
+            device=device, dtype=dtype, heads=heads, n_ctx=n_ctx, n_data=n_data
+        )
+        init_linear(self.c_q, init_scale)
+        init_linear(self.c_kv, init_scale)
+        init_linear(self.c_proj, init_scale)
+
+    def forward(self, x, data):
+        x = self.c_q(x)
+        data = self.c_kv(data)
+        x = checkpoint(self.attention, (x, data), (), True)
+        x = self.c_proj(x)
+        return x
+
+
+class QKVMultiheadCrossAttention(nn.Module):
+    def __init__(
+        self, *, device: torch.device, dtype: torch.dtype, heads: int, n_ctx: int, n_data: int
+    ):
+        super().__init__()
+        self.device = device
+        self.dtype = dtype
+        self.heads = heads
+        self.n_ctx = n_ctx
+        self.n_data = n_data
+
+    def forward(self, q, kv):
+        _, n_ctx, _ = q.shape
+        bs, n_data, width = kv.shape
+        attn_ch = width // self.heads // 2
+        scale = 1 / math.sqrt(math.sqrt(attn_ch))
+        q = q.view(bs, n_ctx, self.heads, -1)
+        kv = kv.view(bs, n_data, self.heads, -1)
+        k, v = torch.split(kv, attn_ch, dim=-1)
+        weight = torch.einsum(
+            "bthc,bshc->bhts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        wdtype = weight.dtype
+        weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
+        return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
+
+
+class ResidualCrossAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int,
+        n_data: int,
+        width: int,
+        heads: int,
+        data_width: Optional[int] = None,
+        init_scale: float = 1.0,
+    ):
+        super().__init__()
+
+        if data_width is None:
+            data_width = width
+
+        self.attn = MultiheadCrossAttention(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_ctx,
+            n_data=n_data,
+            width=width,
+            heads=heads,
+            data_width=data_width,
+            init_scale=init_scale,
+        )
+        self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype)
+        self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale)
+        self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor, data: torch.Tensor):
+        x = x + self.attn(self.ln_1(x), self.ln_2(data))
+        x = x + self.mlp(self.ln_3(x))
+        return x
+
+
+class SimplePerceiver(nn.Module):
+    """
+    Only does cross attention
+    """
+
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int,
+        n_data: int,
+        width: int,
+        layers: int,
+        heads: int,
+        init_scale: float = 0.25,
+        data_width: Optional[int] = None,
+    ):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.width = width
+        self.layers = layers
+        init_scale = init_scale * math.sqrt(1.0 / width)
+        self.resblocks = nn.ModuleList(
+            [
+                ResidualCrossAttentionBlock(
+                    device=device,
+                    dtype=dtype,
+                    n_ctx=n_ctx,
+                    n_data=n_data,
+                    width=width,
+                    heads=heads,
+                    init_scale=init_scale,
+                    data_width=data_width,
+                )
+                for _ in range(layers)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor, data: torch.Tensor):
+        for block in self.resblocks:
+            x = block(x, data)
+        return x
+
+
+class PointDiffusionPerceiver(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        input_channels: int = 3,
+        output_channels: int = 3,
+        n_ctx: int = 1024,
+        n_latent: int = 128,
+        width: int = 512,
+        encoder_layers: int = 12,
+        latent_layers: int = 12,
+        decoder_layers: int = 12,
+        heads: int = 8,
+        init_scale: float = 0.25,
+    ):
+        super().__init__()
+        self.time_embed = MLP(
+            device=device, dtype=dtype, width=width, init_scale=init_scale * math.sqrt(1.0 / width)
+        )
+        self.latent_embed = MLP(
+            device=device, dtype=dtype, width=width, init_scale=init_scale * math.sqrt(1.0 / width)
+        )
+        self.n_latent = n_latent
+
+        self.ln_pre = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.encoder = SimplePerceiver(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_latent,
+            n_data=n_ctx,
+            width=width,
+            layers=encoder_layers,
+            heads=heads,
+            init_scale=init_scale,
+        )
+        self.processor = Transformer(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_latent,
+            width=width,
+            layers=latent_layers,
+            heads=heads,
+            init_scale=init_scale,
+        )
+        self.decoder = SimplePerceiver(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_ctx,
+            n_data=n_latent,
+            width=width,
+            layers=decoder_layers,
+            heads=heads,
+            init_scale=init_scale,
+        )
+        self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.input_proj = nn.Linear(input_channels, width, device=device, dtype=dtype)
+        self.output_proj = nn.Linear(width, output_channels, device=device, dtype=dtype)
+        with torch.no_grad():
+            self.output_proj.weight.zero_()
+            self.output_proj.bias.zero_()
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor):
+        """
+        :param x: an [N x C x T] tensor.
+        :param t: an [N] tensor.
+        :return: an [N x C' x T] tensor.
+        """
+        assert x.shape[-1] == self.decoder.n_ctx
+        t_embed = self.time_embed(timestep_embedding(t, self.encoder.width))
+        data = self.input_proj(x.permute(0, 2, 1)) + t_embed[:, None]
+        data = self.ln_pre(data)
+
+        l = torch.arange(self.n_latent).to(x.device)
+        h = self.latent_embed(timestep_embedding(l, self.decoder.width))
+        h = h.unsqueeze(0).repeat(x.shape[0], 1, 1)
+
+        h = self.encoder(h, data)
+        h = self.processor(h)
+        h = self.decoder(data, h)
+        h = self.ln_post(h)
+        h = self.output_proj(h)
+        return h.permute(0, 2, 1)

shap_e/models/generation/pooled_mlp.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+
+from .util import timestep_embedding
+
+
+class PooledMLP(nn.Module):
+    def __init__(
+        self,
+        device: torch.device,
+        *,
+        input_channels: int = 3,
+        output_channels: int = 6,
+        hidden_size: int = 256,
+        resblocks: int = 4,
+        pool_op: str = "max",
+    ):
+        super().__init__()
+        self.input_embed = nn.Conv1d(input_channels, hidden_size, kernel_size=1, device=device)
+        self.time_embed = nn.Linear(hidden_size, hidden_size, device=device)
+
+        blocks = []
+        for _ in range(resblocks):
+            blocks.append(ResBlock(hidden_size, pool_op, device=device))
+        self.sequence = nn.Sequential(*blocks)
+
+        self.out = nn.Conv1d(hidden_size, output_channels, kernel_size=1, device=device)
+        with torch.no_grad():
+            self.out.bias.zero_()
+            self.out.weight.zero_()
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        in_embed = self.input_embed(x)
+        t_embed = self.time_embed(timestep_embedding(t, in_embed.shape[1]))
+        h = in_embed + t_embed[..., None]
+        h = self.sequence(h)
+        h = self.out(h)
+        return h
+
+
+class ResBlock(nn.Module):
+    def __init__(self, hidden_size: int, pool_op: str, device: torch.device):
+        super().__init__()
+        assert pool_op in ["mean", "max"]
+        self.pool_op = pool_op
+        self.body = nn.Sequential(
+            nn.SiLU(),
+            nn.LayerNorm((hidden_size,), device=device),
+            nn.Linear(hidden_size, hidden_size, device=device),
+            nn.SiLU(),
+            nn.LayerNorm((hidden_size,), device=device),
+            nn.Linear(hidden_size, hidden_size, device=device),
+        )
+        self.gate = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size, device=device),
+            nn.Tanh(),
+        )
+
+    def forward(self, x: torch.Tensor):
+        N, C, T = x.shape
+        out = self.body(x.permute(0, 2, 1).reshape(N * T, C)).reshape([N, T, C]).permute(0, 2, 1)
+        pooled = pool(self.pool_op, x)
+        gate = self.gate(pooled)
+        return x + out * gate[..., None]
+
+
+def pool(op_name: str, x: torch.Tensor) -> torch.Tensor:
+    if op_name == "max":
+        pooled, _ = torch.max(x, dim=-1)
+    elif op_name == "mean":
+        pooled, _ = torch.mean(x, dim=-1)
+    else:
+        raise ValueError(f"unknown pool op: {op_name}")
+    return pooled

shap_e/models/generation/pretrained_clip.py

@@ -0,0 +1,270 @@
+from typing import Iterable, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+
+from shap_e.models.download import default_cache_dir
+
+ImageType = Union[np.ndarray, torch.Tensor, Image.Image]
+
+
+class ImageCLIP(nn.Module):
+    """
+    A wrapper around a pre-trained CLIP model that automatically handles
+    batches of texts, images, and embeddings.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        dtype: Optional[torch.dtype] = torch.float32,
+        ensure_used_params: bool = True,
+        clip_name: str = "ViT-L/14",
+        cache_dir: Optional[str] = None,
+    ):
+        super().__init__()
+
+        assert clip_name in ["ViT-L/14", "ViT-B/32"]
+
+        self.device = device
+        self.ensure_used_params = ensure_used_params
+
+        # Lazy import because of torchvision.
+        import clip
+
+        self.clip_model, self.preprocess = clip.load(
+            clip_name, device=device, download_root=cache_dir or default_cache_dir()
+        )
+        self.clip_name = clip_name
+
+        if dtype is not None:
+            self.clip_model.to(dtype)
+        self._tokenize = clip.tokenize
+
+    @property
+    def feature_dim(self) -> int:
+        if self.clip_name == "ViT-L/14":
+            return 768
+        else:
+            return 512
+
+    @property
+    def grid_size(self) -> int:
+        if self.clip_name == "ViT-L/14":
+            return 16
+        else:
+            return 7
+
+    @property
+    def grid_feature_dim(self) -> int:
+        if self.clip_name == "ViT-L/14":
+            return 1024
+        else:
+            return 768
+
+    def forward(
+        self,
+        batch_size: int,
+        images: Optional[Iterable[Optional[ImageType]]] = None,
+        texts: Optional[Iterable[Optional[str]]] = None,
+        embeddings: Optional[Iterable[Optional[torch.Tensor]]] = None,
+    ) -> torch.Tensor:
+        """
+        Generate a batch of embeddings from a mixture of images, texts,
+        precomputed embeddings, and possibly empty values.
+
+        For each batch element, at most one of images, texts, and embeddings
+        should have a non-None value. Embeddings from multiple modalities
+        cannot be mixed for a single batch element. If no modality is provided,
+        a zero embedding will be used for the batch element.
+        """
+        image_seq = [None] * batch_size if images is None else list(images)
+        text_seq = [None] * batch_size if texts is None else list(texts)
+        embedding_seq = [None] * batch_size if embeddings is None else list(embeddings)
+        assert len(image_seq) == batch_size, "number of images should match batch size"
+        assert len(text_seq) == batch_size, "number of texts should match batch size"
+        assert len(embedding_seq) == batch_size, "number of embeddings should match batch size"
+
+        if self.ensure_used_params:
+            return self._static_multimodal_embed(
+                images=image_seq, texts=text_seq, embeddings=embedding_seq
+            )
+
+        result = torch.zeros((batch_size, self.feature_dim), device=self.device)
+        index_images = []
+        index_texts = []
+        for i, (image, text, emb) in enumerate(zip(image_seq, text_seq, embedding_seq)):
+            assert (
+                sum([int(image is not None), int(text is not None), int(emb is not None)]) < 2
+            ), "only one modality may be non-None per batch element"
+            if image is not None:
+                index_images.append((i, image))
+            elif text is not None:
+                index_texts.append((i, text))
+            elif emb is not None:
+                result[i] = emb.to(result)
+
+        if len(index_images):
+            embs = self.embed_images((img for _, img in index_images))
+            for (i, _), emb in zip(index_images, embs):
+                result[i] = emb.to(result)
+        if len(index_texts):
+            embs = self.embed_text((text for _, text in index_texts))
+            for (i, _), emb in zip(index_texts, embs):
+                result[i] = emb.to(result)
+
+        return result
+
+    def _static_multimodal_embed(
+        self,
+        images: List[Optional[ImageType]] = None,
+        texts: List[Optional[str]] = None,
+        embeddings: List[Optional[torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        """
+        Like forward(), but always runs all encoders to ensure that
+        the forward graph looks the same on every rank.
+        """
+        image_emb = self.embed_images(images)
+        text_emb = self.embed_text(t if t else "" for t in texts)
+        joined_embs = torch.stack(
+            [
+                emb.to(device=self.device, dtype=torch.float32)
+                if emb is not None
+                else torch.zeros(self.feature_dim, device=self.device)
+                for emb in embeddings
+            ],
+            dim=0,
+        )
+
+        image_flag = torch.tensor([x is not None for x in images], device=self.device)[
+            :, None
+        ].expand_as(image_emb)
+        text_flag = torch.tensor([x is not None for x in texts], device=self.device)[
+            :, None
+        ].expand_as(image_emb)
+        emb_flag = torch.tensor([x is not None for x in embeddings], device=self.device)[
+            :, None
+        ].expand_as(image_emb)
+
+        return (
+            image_flag.float() * image_emb
+            + text_flag.float() * text_emb
+            + emb_flag.float() * joined_embs
+            + self.clip_model.logit_scale * 0  # avoid unused parameters
+        )
+
+    def embed_images(self, xs: Iterable[Optional[ImageType]]) -> torch.Tensor:
+        """
+        :param xs: N images, stored as numpy arrays, tensors, or PIL images.
+        :return: an [N x D] tensor of features.
+        """
+        clip_inputs = self.images_to_tensor(xs)
+        results = self.clip_model.encode_image(clip_inputs).float()
+        return results / torch.linalg.norm(results, dim=-1, keepdim=True)
+
+    def embed_text(self, prompts: Iterable[str]) -> torch.Tensor:
+        """
+        Embed text prompts as an [N x D] tensor.
+        """
+        enc = self.clip_model.encode_text(
+            self._tokenize(list(prompts), truncate=True).to(self.device)
+        ).float()
+        return enc / torch.linalg.norm(enc, dim=-1, keepdim=True)
+
+    def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) -> torch.Tensor:
+        """
+        Embed images into latent grids.
+
+        :param xs: an iterable of images to embed.
+        :return: a tensor of shape [N x C x L], where L = self.grid_size**2.
+        """
+        if self.ensure_used_params:
+            extra_value = 0.0
+            for p in self.parameters():
+                extra_value = extra_value + p.mean() * 0.0
+        else:
+            extra_value = 0.0
+
+        x = self.images_to_tensor(xs).to(self.clip_model.dtype)
+
+        # https://github.com/openai/CLIP/blob/4d120f3ec35b30bd0f992f5d8af2d793aad98d2a/clip/model.py#L225
+        vt = self.clip_model.visual
+        x = vt.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat(
+            [
+                vt.class_embedding.to(x.dtype)
+                + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
+                x,
+            ],
+            dim=1,
+        )  # shape = [*, grid ** 2 + 1, width]
+        x = x + vt.positional_embedding.to(x.dtype)
+        x = vt.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = vt.transformer(x)
+        x = x.permute(1, 2, 0)  # LND -> NDL
+
+        return x[..., 1:].contiguous().float() + extra_value
+
+    def images_to_tensor(self, xs: Iterable[Optional[ImageType]]) -> torch.Tensor:
+        return torch.stack([self.preprocess(_image_to_pil(x)) for x in xs], dim=0).to(self.device)
+
+
+class FrozenImageCLIP:
+    def __init__(self, device: torch.device, **kwargs):
+        self.model = ImageCLIP(device, dtype=None, ensure_used_params=False, **kwargs)
+        for parameter in self.model.parameters():
+            parameter.requires_grad_(False)
+
+    @property
+    def feature_dim(self) -> int:
+        return self.model.feature_dim
+
+    @property
+    def grid_size(self) -> int:
+        return self.model.grid_size
+
+    @property
+    def grid_feature_dim(self) -> int:
+        return self.model.grid_feature_dim
+
+    def __call__(
+        self,
+        batch_size: int,
+        images: Optional[Iterable[Optional[ImageType]]] = None,
+        texts: Optional[Iterable[Optional[str]]] = None,
+        embeddings: Optional[Iterable[Optional[torch.Tensor]]] = None,
+    ) -> torch.Tensor:
+        # We don't do a no_grad() here so that gradients could still
+        # flow to the input embeddings argument.
+        # This behavior is currently not used, but it could be.
+        return self.model(batch_size=batch_size, images=images, texts=texts, embeddings=embeddings)
+
+    def embed_images(self, xs: Iterable[Optional[ImageType]]) -> torch.Tensor:
+        with torch.no_grad():
+            return self.model.embed_images(xs)
+
+    def embed_text(self, prompts: Iterable[str]) -> torch.Tensor:
+        with torch.no_grad():
+            return self.model.embed_text(prompts)
+
+    def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) -> torch.Tensor:
+        with torch.no_grad():
+            return self.model.embed_images_grid(xs)
+
+
+def _image_to_pil(obj: Optional[ImageType]) -> Image.Image:
+    if obj is None:
+        return Image.fromarray(np.zeros([64, 64, 3], dtype=np.uint8))
+    if isinstance(obj, np.ndarray):
+        return Image.fromarray(obj.astype(np.uint8))
+    elif isinstance(obj, torch.Tensor):
+        return Image.fromarray(obj.detach().cpu().numpy().astype(np.uint8))
+    else:
+        return obj

shap_e/models/generation/transformer.py

@@ -0,0 +1,494 @@
+import math
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+
+from shap_e.models.nn.checkpoint import checkpoint
+
+from .pretrained_clip import FrozenImageCLIP, ImageCLIP, ImageType
+from .util import timestep_embedding
+
+def init_linear(l, stddev):
+    nn.init.normal_(l.weight, std=stddev)
+    if l.bias is not None:
+        nn.init.constant_(l.bias, 0.0)
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int,
+        width: int,
+        heads: int,
+        init_scale: float,
+    ):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.width = width
+        self.heads = heads
+        self.c_qkv = nn.Linear(width, width * 3, device=device, dtype=dtype)
+        self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
+        self.attention = QKVMultiheadAttention(device=device, dtype=dtype, heads=heads, n_ctx=n_ctx)
+        init_linear(self.c_qkv, init_scale)
+        init_linear(self.c_proj, init_scale)
+
+    def forward(self, x):
+        x = self.c_qkv(x)
+        x = checkpoint(self.attention, (x,), (), True)
+        x = self.c_proj(x)
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int, init_scale: float):
+        super().__init__()
+        self.width = width
+        self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
+        self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
+        self.gelu = nn.GELU()
+        init_linear(self.c_fc, init_scale)
+        init_linear(self.c_proj, init_scale)
+
+    def forward(self, x):
+        return self.c_proj(self.gelu(self.c_fc(x)))
+
+
+class QKVMultiheadAttention(nn.Module):
+    def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_ctx: int):
+        super().__init__()
+        self.device = device
+        self.dtype = dtype
+        self.heads = heads
+        self.n_ctx = n_ctx
+
+    def forward(self, qkv):
+        bs, n_ctx, width = qkv.shape
+        attn_ch = width // self.heads // 3
+        scale = 1 / math.sqrt(math.sqrt(attn_ch))
+        qkv = qkv.view(bs, n_ctx, self.heads, -1)
+        q, k, v = torch.split(qkv, attn_ch, dim=-1)
+        weight = torch.einsum(
+            "bthc,bshc->bhts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        wdtype = weight.dtype
+        weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
+        return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int,
+        width: int,
+        heads: int,
+        init_scale: float = 1.0,
+    ):
+        super().__init__()
+
+        self.attn = MultiheadAttention(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_ctx,
+            width=width,
+            heads=heads,
+            init_scale=init_scale,
+        )
+        self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale)
+        self.ln_2 = nn.LayerNorm(width, device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int,
+        width: int,
+        layers: int,
+        heads: int,
+        init_scale: float = 0.25,
+    ):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.width = width
+        self.layers = layers
+        init_scale = init_scale * math.sqrt(1.0 / width)
+        self.resblocks = nn.ModuleList(
+            [
+                ResidualAttentionBlock(
+                    device=device,
+                    dtype=dtype,
+                    n_ctx=n_ctx,
+                    width=width,
+                    heads=heads,
+                    init_scale=init_scale,
+                )
+                for _ in range(layers)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor):
+        for block in self.resblocks:
+            x = block(x)
+        return x
+
+
+class PointDiffusionTransformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        input_channels: int = 3,
+        output_channels: int = 3,
+        n_ctx: int = 1024,
+        width: int = 512,
+        layers: int = 12,
+        heads: int = 8,
+        init_scale: float = 0.25,
+        time_token_cond: bool = False,
+        use_pos_emb: bool = False,
+        pos_emb_init_scale: float = 1.0,
+        pos_emb_n_ctx: Optional[int] = None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        self.n_ctx = n_ctx
+        self.time_token_cond = time_token_cond
+        self.use_pos_emb = use_pos_emb
+        self.time_embed = MLP(
+            device=device, dtype=dtype, width=width, init_scale=init_scale * math.sqrt(1.0 / width)
+        )
+        self.ln_pre = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.backbone = Transformer(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_ctx + int(time_token_cond),
+            width=width,
+            layers=layers,
+            heads=heads,
+            init_scale=init_scale,
+        )
+        self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
+        self.input_proj = nn.Linear(input_channels, width, device=device, dtype=dtype)
+        self.output_proj = nn.Linear(width, output_channels, device=device, dtype=dtype)
+        # with torch.no_grad():
+        #     self.output_proj.weight.zero_()
+        #     self.output_proj.bias.zero_()
+        if self.use_pos_emb:
+            self.register_parameter(
+                "pos_emb",
+                nn.Parameter(
+                    pos_emb_init_scale
+                    * torch.randn(pos_emb_n_ctx or self.n_ctx, width, device=device, dtype=dtype)
+                ),
+            )
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor):
+        """
+        :param x: an [N x C x T] tensor.
+        :param t: an [N] tensor.
+        :return: an [N x C' x T] tensor.
+        """
+        assert x.shape[-1] == self.n_ctx
+        t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+        return self._forward_with_cond(x, [(t_embed, self.time_token_cond)])
+
+    def _forward_with_cond(
+        self, x: torch.Tensor, cond_as_token: List[Tuple[torch.Tensor, bool]]
+    ) -> torch.Tensor:
+        h = self.input_proj(x.permute(0, 2, 1))  # NCL -> NLC
+        for emb, as_token in cond_as_token:
+            if not as_token:
+                h = h + emb[:, None]
+        if self.use_pos_emb:
+            h = h + self.pos_emb
+        extra_tokens = [
+            (emb[:, None] if len(emb.shape) == 2 else emb)
+            for emb, as_token in cond_as_token
+            if as_token
+        ]
+        if len(extra_tokens):
+            h = torch.cat(extra_tokens + [h], dim=1)
+        h = self.ln_pre(h)
+        h = self.backbone(h)
+        h = self.ln_post(h)
+        if len(extra_tokens):
+            h = h[:, sum(h.shape[1] for h in extra_tokens):]
+        h = self.output_proj(h)
+        return h.permute(0, 2, 1) # NCL -> NLC
+
+
+
+
+class CLIPImagePointDiffusionTransformer(PointDiffusionTransformer):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int = 1024,
+        token_cond: bool = False,
+        cond_drop_prob: float = 0.0,
+        frozen_clip: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            device=device, dtype=dtype, n_ctx=n_ctx + int(token_cond), pos_emb_n_ctx=n_ctx, **kwargs
+        )
+        # print("!!!!!", "deivce:", device, "dtype:", dtype, "n_ctx:", n_ctx, "token_cond:", token_cond, "cond_drop_prob:", cond_drop_prob, "frozen_clip:", frozen_clip, "kwargs:", kwargs)
+        self.n_ctx = n_ctx
+        self.token_cond = token_cond
+        self.clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device)
+        self.clip_embed = nn.Linear(
+            self.clip.feature_dim, self.backbone.width, device=device, dtype=dtype
+        )
+        self.cond_drop_prob = cond_drop_prob
+
+    def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        with torch.no_grad():
+            return dict(embeddings=self.clip(batch_size, **model_kwargs))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        images: Optional[Iterable[Optional[ImageType]]] = None,
+        texts: Optional[Iterable[Optional[str]]] = None,
+        embeddings: Optional[Iterable[Optional[torch.Tensor]]] = None,
+    ):
+        """
+        :param x: an [N x C x T] tensor.
+        :param t: an [N] tensor.
+        :param images: a batch of images to condition on.
+        :param texts: a batch of texts to condition on.
+        :param embeddings: a batch of CLIP embeddings to condition on.
+        :return: an [N x C' x T] tensor.
+        """
+        # print("x.shape", x.shape, "t.shape", t.shape, "images", images, "texts", texts, "embeddings", embeddings)
+        assert x.shape[-1] == self.n_ctx # self.n_ctx = 1024
+
+        t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+        clip_out = self.clip(batch_size=len(x), images=images, texts=texts, embeddings=embeddings)
+        assert len(clip_out.shape) == 2 and clip_out.shape[0] == x.shape[0]
+
+        if self.training:
+            mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
+            clip_out = clip_out * mask[:, None].to(clip_out)
+
+        # Rescale the features to have unit variance
+        clip_out = math.sqrt(clip_out.shape[1]) * clip_out
+
+        clip_embed = self.clip_embed(clip_out)
+
+        cond = [(clip_embed, self.token_cond), (t_embed, self.time_token_cond)]
+        return self._forward_with_cond(x, cond)
+
+
+class CLIPImageGridPointDiffusionTransformer(PointDiffusionTransformer):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int = 1024,
+        cond_drop_prob: float = 0.0,
+        frozen_clip: bool = True,
+        **kwargs,
+    ):
+        clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device)
+        super().__init__(
+            device=device,
+            dtype=dtype,
+            n_ctx=n_ctx + clip.grid_size**2,
+            pos_emb_n_ctx=n_ctx,
+            **kwargs,
+        )
+        self.n_ctx = n_ctx
+        self.clip = clip
+        self.clip_embed = nn.Sequential(
+            nn.LayerNorm(
+                normalized_shape=(self.clip.grid_feature_dim,), device=device, dtype=dtype
+            ),
+            nn.Linear(self.clip.grid_feature_dim, self.backbone.width, device=device, dtype=dtype),
+        )
+        self.cond_drop_prob = cond_drop_prob
+
+    def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        _ = batch_size
+        with torch.no_grad():
+            return dict(embeddings=self.clip.embed_images_grid(model_kwargs["images"]))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        images: Optional[Iterable[ImageType]] = None,
+        embeddings: Optional[Iterable[torch.Tensor]] = None,
+    ):
+        """
+        :param x: an [N x C x T] tensor.
+        :param t: an [N] tensor.
+        :param images: a batch of images to condition on.
+        :param embeddings: a batch of CLIP latent grids to condition on.
+        :return: an [N x C' x T] tensor.
+        """
+        assert images is not None or embeddings is not None, "must specify images or embeddings"
+        assert images is None or embeddings is None, "cannot specify both images and embeddings"
+        assert x.shape[-1] == self.n_ctx
+
+        t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+
+        if images is not None:
+            clip_out = self.clip.embed_images_grid(images)
+        else:
+            clip_out = embeddings
+
+        if self.training:
+            mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
+            clip_out = clip_out * mask[:, None, None].to(clip_out)
+
+        clip_out = clip_out.permute(0, 2, 1)  # NCL -> NLC
+        clip_embed = self.clip_embed(clip_out)
+
+        cond = [(t_embed, self.time_token_cond), (clip_embed, True)]
+        return self._forward_with_cond(x, cond)
+
+
+class UpsamplePointDiffusionTransformer(PointDiffusionTransformer):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        cond_input_channels: Optional[int] = None,
+        cond_ctx: int = 1024,
+        n_ctx: int = 4096 - 1024,
+        channel_scales: Optional[Sequence[float]] = None,
+        channel_biases: Optional[Sequence[float]] = None,
+        **kwargs,
+    ):
+        super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + cond_ctx, **kwargs)
+        self.n_ctx = n_ctx
+        self.cond_input_channels = cond_input_channels or self.input_channels
+        self.cond_point_proj = nn.Linear(
+            self.cond_input_channels, self.backbone.width, device=device, dtype=dtype
+        )
+
+        self.register_buffer(
+            "channel_scales",
+            torch.tensor(channel_scales, dtype=dtype, device=device)
+            if channel_scales is not None
+            else None,
+        )
+        self.register_buffer(
+            "channel_biases",
+            torch.tensor(channel_biases, dtype=dtype, device=device)
+            if channel_biases is not None
+            else None,
+        )
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.Tensor):
+        """
+        :param x: an [N x C1 x T] tensor.
+        :param t: an [N] tensor.
+        :param low_res: an [N x C2 x T'] tensor of conditioning points.
+        :return: an [N x C3 x T] tensor.
+        """
+        assert x.shape[-1] == self.n_ctx
+        t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+        low_res_embed = self._embed_low_res(low_res)
+        cond = [(t_embed, self.time_token_cond), (low_res_embed, True)]
+        return self._forward_with_cond(x, cond)
+
+    def _embed_low_res(self, x: torch.Tensor) -> torch.Tensor:
+        if self.channel_scales is not None:
+            x = x * self.channel_scales[None, :, None]
+        if self.channel_biases is not None:
+            x = x + self.channel_biases[None, :, None]
+        return self.cond_point_proj(x.permute(0, 2, 1))
+
+
+class CLIPImageGridUpsamplePointDiffusionTransformer(UpsamplePointDiffusionTransformer):
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        n_ctx: int = 4096 - 1024,
+        cond_drop_prob: float = 0.0,
+        frozen_clip: bool = True,
+        **kwargs,
+    ):
+        clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device)
+        super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.grid_size**2, **kwargs)
+        self.n_ctx = n_ctx
+
+        self.clip = clip
+        self.clip_embed = nn.Sequential(
+            nn.LayerNorm(
+                normalized_shape=(self.clip.grid_feature_dim,), device=device, dtype=dtype
+            ),
+            nn.Linear(self.clip.grid_feature_dim, self.backbone.width, device=device, dtype=dtype),
+        )
+        self.cond_drop_prob = cond_drop_prob
+
+    def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        _ = batch_size
+        with torch.no_grad():
+            return dict(
+                embeddings=self.clip.embed_images_grid(model_kwargs["images"]),
+                low_res=model_kwargs["low_res"],
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        *,
+        low_res: torch.Tensor,
+        images: Optional[Iterable[ImageType]] = None,
+        embeddings: Optional[Iterable[torch.Tensor]] = None,
+    ):
+        """
+        :param x: an [N x C1 x T] tensor.
+        :param t: an [N] tensor.
+        :param low_res: an [N x C2 x T'] tensor of conditioning points.
+        :param images: a batch of images to condition on.
+        :param embeddings: a batch of CLIP latent grids to condition on.
+        :return: an [N x C3 x T] tensor.
+        """
+        assert x.shape[-1] == self.n_ctx
+        t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
+        low_res_embed = self._embed_low_res(low_res)
+
+        if images is not None:
+            clip_out = self.clip.embed_images_grid(images)
+        else:
+            clip_out = embeddings
+
+        if self.training:
+            mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
+            clip_out = clip_out * mask[:, None, None].to(clip_out)
+
+        clip_out = clip_out.permute(0, 2, 1)  # NCL -> NLC
+        clip_embed = self.clip_embed(clip_out)
+
+        cond = [(t_embed, self.time_token_cond), (clip_embed, True), (low_res_embed, True)]
+        return self._forward_with_cond(x, cond)
+

shap_e/models/generation/util.py

@@ -0,0 +1,23 @@
+import math
+
+import torch
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].to(timesteps.dtype) * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding

shap_e/models/nerf/model.py

@@ -0,0 +1,255 @@
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from shap_e.models.nn.checkpoint import checkpoint
+from shap_e.models.nn.encoding import encode_position, spherical_harmonics_basis
+from shap_e.models.nn.meta import MetaModule, subdict
+from shap_e.models.nn.ops import MLP, MetaMLP, get_act, mlp_init, zero_init
+from shap_e.models.nn.utils import ArrayType
+from shap_e.models.query import Query
+from shap_e.util.collections import AttrDict
+
+
+class NeRFModel(ABC):
+    """
+    Parametric scene representation whose outputs are integrated by NeRFRenderer
+    """
+
+    @abstractmethod
+    def forward(
+        self,
+        query: Query,
+        params: Optional[Dict[str, torch.Tensor]] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> AttrDict:
+        """
+        :param query: the points in the field to query.
+        :param params: Meta parameters
+        :param options: Optional hyperparameters
+        :return: An AttrDict containing at least
+            - density: [batch_size x ... x 1]
+            - channels: [batch_size x ... x n_channels]
+            - aux_losses: [batch_size x ... x 1]
+        """
+
+
+class VoidNeRFModel(MetaModule, NeRFModel):
+    """
+    Implements the default empty space model where all queries are rendered as
+    background.
+    """
+
+    def __init__(
+        self,
+        background: ArrayType,
+        trainable: bool = False,
+        channel_scale: float = 255.0,
+        device: torch.device = torch.device("cuda"),
+    ):
+        super().__init__()
+        background = nn.Parameter(
+            torch.from_numpy(np.array(background)).to(dtype=torch.float32, device=device)
+            / channel_scale
+        )
+        if trainable:
+            self.register_parameter("background", background)
+        else:
+            self.register_buffer("background", background)
+
+    def forward(
+        self,
+        query: Query,
+        params: Optional[Dict[str, torch.Tensor]] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> AttrDict:
+        _ = params
+        default_bg = self.background[None]
+        background = options.get("background", default_bg) if options is not None else default_bg
+
+        shape = query.position.shape[:-1]
+        ones = [1] * (len(shape) - 1)
+        n_channels = background.shape[-1]
+        background = torch.broadcast_to(
+            background.view(background.shape[0], *ones, n_channels), [*shape, n_channels]
+        )
+        return background
+
+
+class MLPNeRFModel(MetaModule, NeRFModel):
+    def __init__(
+        self,
+        # Positional encoding parameters
+        n_levels: int = 10,
+        # MLP parameters
+        d_hidden: int = 256,
+        n_density_layers: int = 4,
+        n_channel_layers: int = 1,
+        n_channels: int = 3,
+        sh_degree: int = 4,
+        activation: str = "relu",
+        density_activation: str = "exp",
+        init: Optional[str] = None,
+        init_scale: float = 1.0,
+        output_activation: str = "sigmoid",
+        meta_parameters: bool = False,
+        trainable_meta: bool = False,
+        zero_out: bool = True,
+        register_freqs: bool = True,
+        posenc_version: str = "v1",
+        device: torch.device = torch.device("cuda"),
+    ):
+        super().__init__()
+
+        # Positional encoding
+        if register_freqs:
+            # not used anymore
+            self.register_buffer(
+                "freqs",
+                2.0 ** torch.arange(n_levels, device=device, dtype=torch.float).view(1, n_levels),
+            )
+
+        self.posenc_version = posenc_version
+        dummy = torch.eye(1, 3)
+        d_input = encode_position(posenc_version, position=dummy).shape[-1]
+
+        self.n_levels = n_levels
+
+        self.sh_degree = sh_degree
+        d_sh_coeffs = sh_degree**2
+
+        self.meta_parameters = meta_parameters
+
+        mlp_cls = (
+            partial(
+                MetaMLP,
+                meta_scale=False,
+                meta_shift=False,
+                meta_proj=True,
+                meta_bias=True,
+                trainable_meta=trainable_meta,
+            )
+            if meta_parameters
+            else MLP
+        )
+
+        self.density_mlp = mlp_cls(
+            d_input=d_input,
+            d_hidden=[d_hidden] * (n_density_layers - 1),
+            d_output=d_hidden,
+            act_name=activation,
+            init_scale=init_scale,
+        )
+
+        self.channel_mlp = mlp_cls(
+            d_input=d_hidden + d_sh_coeffs,
+            d_hidden=[d_hidden] * n_channel_layers,
+            d_output=n_channels,
+            act_name=activation,
+            init_scale=init_scale,
+        )
+
+        self.act = get_act(output_activation)
+        self.density_act = get_act(density_activation)
+
+        mlp_init(
+            list(self.density_mlp.affines) + list(self.channel_mlp.affines),
+            init=init,
+            init_scale=init_scale,
+        )
+
+        if zero_out:
+            zero_init(self.channel_mlp.affines[-1])
+
+        self.to(device)
+
+    def encode_position(self, query: Query):
+        h = encode_position(self.posenc_version, position=query.position)
+        return h
+
+    def forward(
+        self,
+        query: Query,
+        params: Optional[Dict[str, torch.Tensor]] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> AttrDict:
+        params = self.update(params)
+
+        options = AttrDict() if options is None else AttrDict(options)
+
+        query = query.copy()
+
+        h_position = self.encode_position(query)
+
+        if self.meta_parameters:
+            density_params = subdict(params, "density_mlp")
+            density_mlp = partial(
+                self.density_mlp, params=density_params, options=options, log_prefix="density_"
+            )
+            density_mlp_parameters = list(density_params.values())
+        else:
+            density_mlp = partial(self.density_mlp, options=options, log_prefix="density_")
+            density_mlp_parameters = self.density_mlp.parameters()
+        h_density = checkpoint(
+            density_mlp,
+            (h_position,),
+            density_mlp_parameters,
+            options.checkpoint_nerf_mlp,
+        )
+        h_direction = maybe_get_spherical_harmonics_basis(
+            sh_degree=self.sh_degree,
+            coords_shape=query.position.shape,
+            coords=query.direction,
+            device=query.position.device,
+        )
+
+        if self.meta_parameters:
+            channel_params = subdict(params, "channel_mlp")
+            channel_mlp = partial(
+                self.channel_mlp, params=channel_params, options=options, log_prefix="channel_"
+            )
+            channel_mlp_parameters = list(channel_params.values())
+        else:
+            channel_mlp = partial(self.channel_mlp, options=options, log_prefix="channel_")
+            channel_mlp_parameters = self.channel_mlp.parameters()
+        h_channel = checkpoint(
+            channel_mlp,
+            (torch.cat([h_density, h_direction], dim=-1),),
+            channel_mlp_parameters,
+            options.checkpoint_nerf_mlp,
+        )
+
+        density_logit = h_density[..., :1]
+
+        res = AttrDict(
+            density_logit=density_logit,
+            density=self.density_act(density_logit),
+            channels=self.act(h_channel),
+            aux_losses=AttrDict(),
+            no_weight_grad_aux_losses=AttrDict(),
+        )
+        if options.return_h_density:
+            res.h_density = h_density
+
+        return res
+
+
+def maybe_get_spherical_harmonics_basis(
+    sh_degree: int,
+    coords_shape: Tuple[int],
+    coords: Optional[torch.Tensor] = None,
+    device: torch.device = torch.device("cuda"),
+) -> torch.Tensor:
+    """
+    :param sh_degree: Spherical harmonics degree
+    :param coords_shape: [*shape, 3]
+    :param coords: optional coordinate tensor of coords_shape
+    """
+    if coords is None:
+        return torch.zeros(*coords_shape[:-1], sh_degree**2).to(device)
+
+    return spherical_harmonics_basis(coords, sh_degree)

shap_e/models/nerf/ray.py

@@ -0,0 +1,512 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from shap_e.models.nn.utils import sample_pmf
+from shap_e.models.volume import Volume, VolumeRange
+from shap_e.util.collections import AttrDict
+
+from .model import NeRFModel, Query
+
+
+def render_rays(
+    rays: torch.Tensor,
+    parts: List["RayVolumeIntegral"],
+    void_model: NeRFModel,
+    shared: bool = False,
+    prev_raw_outputs: Optional[List[AttrDict]] = None,
+    render_with_direction: bool = True,
+    importance_sampling_options: Optional[Dict[str, Any]] = None,
+) -> Tuple["RayVolumeIntegralResults", List["RaySampler"], List[AttrDict]]:
+    """
+    Perform volumetric rendering over a partition of possible t's in the union
+    of rendering volumes (written below with some abuse of notations)
+
+        C(r) := sum(
+            transmittance(t[i]) *
+            integrate(
+                lambda t: density(t) * channels(t) * transmittance(t),
+                [t[i], t[i + 1]],
+            )
+            for i in range(len(parts))
+        ) + transmittance(t[-1]) * void_model(t[-1]).channels
+
+    where
+
+    1) transmittance(s) := exp(-integrate(density, [t[0], s])) calculates the
+       probability of light passing through the volume specified by [t[0], s].
+       (transmittance of 1 means light can pass freely)
+    2) density and channels are obtained by evaluating the appropriate
+       part.model at time t.
+    3) [t[i], t[i + 1]] is defined as the range of t where the ray intersects
+       (parts[i].volume \\ union(part.volume for part in parts[:i])) at the surface
+       of the shell (if bounded). If the ray does not intersect, the integral over
+       this segment is evaluated as 0 and transmittance(t[i + 1]) :=
+       transmittance(t[i]).
+    4) The last term is integration to infinity (e.g. [t[-1], math.inf]) that
+       is evaluated by the void_model (i.e. we consider this space to be empty).
+
+    :param rays: [batch_size x ... x 2 x 3] origin and direction.
+    :param parts: disjoint volume integrals.
+    :param void_model: use this model to integrate over the empty space
+    :param shared: All RayVolumeIntegrals are calculated with the same model.
+    :param prev_raw_outputs: Raw outputs from the previous rendering step
+
+    :return: A tuple of
+        - AttrDict containing the rendered `channels`, `distances`, and the `aux_losses`
+        - A list of importance samplers for additional fine-grained rendering
+        - A list of raw output for each interval
+    """
+    if importance_sampling_options is None:
+        importance_sampling_options = {}
+
+    origin, direc = rays[..., 0, :], rays[..., 1, :]
+
+    if prev_raw_outputs is None:
+        prev_raw_outputs = [None] * len(parts)
+
+    samplers = []
+    raw_outputs = []
+    t0 = None
+    results = None
+    # import pdb; pdb.set_trace()
+    for part_i, prev_raw_i in zip(parts, prev_raw_outputs):
+
+        # Integrate over [t[i], t[i + 1]]
+        results_i = part_i.render_rays(
+            origin,
+            direc,
+            t0=t0,
+            prev_raw=prev_raw_i,
+            shared=shared,
+            render_with_direction=render_with_direction,
+        )
+
+        # Create an importance sampler for (optional) fine rendering
+        samplers.append(
+            ImportanceRaySampler(
+                results_i.volume_range, results_i.raw, **importance_sampling_options
+            )
+        )
+        raw_outputs.append(results_i.raw)
+
+        # Pass t[i + 1] as the start of integration for the next interval.
+        t0 = results_i.volume_range.next_t0()
+
+        # Combine the results from [t[0], t[i]] and [t[i], t[i+1]]
+        results = results_i if results is None else results.combine(results_i)
+
+    # While integrating out [t[-1], math.inf] is the correct thing to do, this
+    # erases a lot of useful information. Also, void_model is meant to predict
+    # the channels at t=math.inf.
+
+    # # Add the void background over [t[-1], math.inf] to complete integration.
+    # results = results.combine(
+    #     RayVolumeIntegralResults(
+    #         output=AttrDict(
+    #             channels=void_model(origin, direc),
+    #             distances=torch.zeros_like(t0),
+    #             aux_losses=AttrDict(),
+    #         ),
+    #         volume_range=VolumeRange(
+    #             t0=t0,
+    #             t1=torch.full_like(t0, math.inf),
+    #             intersected=torch.full_like(results.volume_range.intersected, True),
+    #         ),
+    #         # Void space extends to infinity. It is assumed that no light
+    #         # passes beyond the void.
+    #         transmittance=torch.zeros_like(results_i.transmittance),
+    #     )
+    # )
+    results.output.channels = results.output.channels + results.transmittance * void_model(
+        Query(origin, direc)
+    )
+
+    return results, samplers, raw_outputs
+
+
+@dataclass
+class RayVolumeIntegralResults:
+    """
+    Stores the relevant state and results of
+
+        integrate(
+            lambda t: density(t) * channels(t) * transmittance(t),
+            [t0, t1],
+        )
+    """
+
+    # Rendered output and auxiliary losses
+    # output.channels has shape [batch_size, *inner_shape, n_channels]
+    output: AttrDict
+
+    """
+    Optional values
+    """
+
+    # Raw values contain the sampled `ts`, `density`, `channels`, etc.
+    raw: Optional[AttrDict] = None
+
+    # Integration
+    volume_range: Optional[VolumeRange] = None
+
+    # If a ray intersects, the transmittance from t0 to t1 (e.g. the
+    # probability that the ray passes through this volume).
+    # has shape [batch_size, *inner_shape, 1]
+    transmittance: Optional[torch.Tensor] = None
+
+    def combine(self, cur: "RayVolumeIntegralResults") -> "RayVolumeIntegralResults":
+        """
+        Combines the integration results of `self` over [t0, t1] and
+        `cur` over [t1, t2] to produce a new set of results over [t0, t2] by
+        using a similar equation to (4) in NeRF++:
+
+            integrate(
+                lambda t: density(t) * channels(t) * transmittance(t),
+                [t0, t2]
+            )
+
+          = integrate(
+                lambda t: density(t) * channels(t) * transmittance(t),
+                [t0, t1]
+            ) + transmittance(t1) * integrate(
+                lambda t: density(t) * channels(t) * transmittance(t),
+                [t1, t2]
+            )
+        """
+        assert torch.allclose(self.volume_range.next_t0(), cur.volume_range.t0)
+
+        def _combine_fn(
+            prev_val: Optional[torch.Tensor],
+            cur_val: Optional[torch.Tensor],
+            *,
+            prev_transmittance: torch.Tensor,
+        ):
+            assert prev_val is not None
+            if cur_val is None:
+                # cur_output.aux_losses are empty for the void_model.
+                return prev_val
+            return prev_val + prev_transmittance * cur_val
+
+        output = self.output.combine(
+            cur.output, combine_fn=partial(_combine_fn, prev_transmittance=self.transmittance)
+        )
+
+        combined = RayVolumeIntegralResults(
+            output=output,
+            volume_range=self.volume_range.extend(cur.volume_range),
+            transmittance=self.transmittance * cur.transmittance,
+        )
+        return combined
+
+
+@dataclass
+class RayVolumeIntegral:
+    model: NeRFModel
+    volume: Volume
+    sampler: "RaySampler"
+    n_samples: int
+
+    def render_rays(
+        self,
+        origin: torch.Tensor,
+        direction: torch.Tensor,
+        t0: Optional[torch.Tensor] = None,
+        prev_raw: Optional[AttrDict] = None,
+        shared: bool = False,
+        render_with_direction: bool = True,
+    ) -> "RayVolumeIntegralResults":
+        """
+        Perform volumetric rendering over the given volume.
+
+        :param position: [batch_size, *shape, 3]
+        :param direction: [batch_size, *shape, 3]
+        :param t0: Optional [batch_size, *shape, 1]
+        :param prev_raw: the raw outputs when using multiple levels with this model.
+        :param shared: means the same model is used for all RayVolumeIntegral's
+        :param render_with_direction: use the incoming ray direction when querying the model.
+
+        :return: RayVolumeIntegralResults
+        """
+        # 1. Intersect the rays with the current volume and sample ts to
+        # integrate along.
+        vrange = self.volume.intersect(origin, direction, t0_lower=t0)
+        ts = self.sampler.sample(vrange.t0, vrange.t1, self.n_samples)
+
+        if prev_raw is not None and not shared:
+            # Append the previous ts now before fprop because previous
+            # rendering used a different model and we can't reuse the output.
+            ts = torch.sort(torch.cat([ts, prev_raw.ts], dim=-2), dim=-2).values
+
+        # Shape sanity checks
+        batch_size, *_shape, _t0_dim = vrange.t0.shape
+        _, *ts_shape, _ts_dim = ts.shape
+
+        # 2. Get the points along the ray and query the model
+        directions = torch.broadcast_to(direction.unsqueeze(-2), [batch_size, *ts_shape, 3])
+        positions = origin.unsqueeze(-2) + ts * directions
+
+        optional_directions = directions if render_with_direction else None
+        mids = (ts[..., 1:, :] + ts[..., :-1, :]) / 2
+        raw = self.model(
+            Query(
+                position=positions,
+                direction=optional_directions,
+                t_min=torch.cat([vrange.t0[..., None, :], mids], dim=-2),
+                t_max=torch.cat([mids, vrange.t1[..., None, :]], dim=-2),
+            )
+        )
+        raw.ts = ts
+
+        if prev_raw is not None and shared:
+            # We can append the additional queries to previous raw outputs
+            # before integration
+            copy = prev_raw.copy()
+            result = torch.sort(torch.cat([raw.pop("ts"), copy.pop("ts")], dim=-2), dim=-2)
+            merge_results = partial(self._merge_results, dim=-2, indices=result.indices)
+            raw = raw.combine(copy, merge_results)
+            raw.ts = result.values
+
+        # 3. Integrate the raw results
+        output, transmittance = self.integrate_samples(vrange, raw)
+
+        # 4. Clean up results that do not intersect with the volume.
+        transmittance = torch.where(
+            vrange.intersected, transmittance, torch.ones_like(transmittance)
+        )
+
+        def _mask_fn(_key: str, tensor: torch.Tensor):
+            return torch.where(vrange.intersected, tensor, torch.zeros_like(tensor))
+
+        def _is_tensor(_key: str, value: Any):
+            return isinstance(value, torch.Tensor)
+
+        output = output.map(map_fn=_mask_fn, should_map=_is_tensor)
+
+        return RayVolumeIntegralResults(
+            output=output,
+            raw=raw,
+            volume_range=vrange,
+            transmittance=transmittance,
+        )
+
+    def integrate_samples(
+        self,
+        volume_range: VolumeRange,
+        raw: AttrDict,
+    ) -> Tuple[AttrDict, torch.Tensor]:
+        """
+        Integrate the raw.channels along with other aux_losses and values to
+        produce the final output dictionary containing rendered `channels`,
+        estimated `distances` and `aux_losses`.
+
+        :param volume_range: Specifies the integral range [t0, t1]
+        :param raw: Contains a dict of function evaluations at ts. Should have
+
+            density: torch.Tensor [batch_size, *shape, n_samples, 1]
+            channels: torch.Tensor [batch_size, *shape, n_samples, n_channels]
+            aux_losses: {key: torch.Tensor [batch_size, *shape, n_samples, 1] for each key}
+            no_weight_grad_aux_losses: an optional set of losses for which the weights
+                                       should be detached before integration.
+
+            after the call, integrate_samples populates some intermediate calculations
+            for later use like
+
+            weights: torch.Tensor [batch_size, *shape, n_samples, 1] (density *
+                transmittance)[i] weight for each rgb output at [..., i, :].
+        :returns: a tuple of (
+            a dictionary of rendered outputs and aux_losses,
+            transmittance of this volume,
+        )
+        """
+
+        # 1. Calculate the weights
+        _, _, dt = volume_range.partition(raw.ts)
+        ddensity = raw.density * dt
+
+        mass = torch.cumsum(ddensity, dim=-2)
+        transmittance = torch.exp(-mass[..., -1, :])
+
+        alphas = 1.0 - torch.exp(-ddensity)
+        Ts = torch.exp(torch.cat([torch.zeros_like(mass[..., :1, :]), -mass[..., :-1, :]], dim=-2))
+        # This is the probability of light hitting and reflecting off of
+        # something at depth [..., i, :].
+        weights = alphas * Ts
+
+        # 2. Integrate all results
+        def _integrate(key: str, samples: torch.Tensor, weights: torch.Tensor):
+            if key == "density":
+                # Omit integrating the density, because we don't need it
+                return None
+            return torch.sum(samples * weights, dim=-2)
+
+        def _is_tensor(_key: str, value: Any):
+            return isinstance(value, torch.Tensor)
+
+        if raw.no_weight_grad_aux_losses:
+            extra_aux_losses = raw.no_weight_grad_aux_losses.map(
+                partial(_integrate, weights=weights.detach()), should_map=_is_tensor
+            )
+        else:
+            extra_aux_losses = {}
+        output = raw.map(partial(_integrate, weights=weights), should_map=_is_tensor)
+        if "no_weight_grad_aux_losses" in output:
+            del output["no_weight_grad_aux_losses"]
+        output.aux_losses.update(extra_aux_losses)
+
+        # Integrating the ts yields the distance away from the origin; rename the variable.
+        output.distances = output.ts
+        del output["ts"]
+        del output["density"]
+
+        assert output.distances.shape == (*output.channels.shape[:-1], 1)
+        assert output.channels.shape[:-1] == raw.channels.shape[:-2]
+        assert output.channels.shape[-1] == raw.channels.shape[-1]
+
+        # 3. Reduce loss
+        def _reduce_loss(_key: str, loss: torch.Tensor):
+            return loss.view(loss.shape[0], -1).sum(dim=-1)
+
+        # 4. Store other useful calculations
+        raw.weights = weights
+
+        output.aux_losses = output.aux_losses.map(_reduce_loss)
+
+        return output, transmittance
+
+    def _merge_results(
+        self, a: Optional[torch.Tensor], b: torch.Tensor, dim: int, indices: torch.Tensor
+    ):
+        """
+        :param a: [..., n_a, ...]. The other dictionary containing the b's may
+            contain extra tensors from earlier calculations, so a can be None.
+        :param b: [..., n_b, ...]
+        :param dim: dimension to merge
+        :param indices: how the merged results should be sorted at the end
+        :return: a concatted and sorted tensor of size [..., n_a + n_b, ...]
+        """
+        if a is None:
+            return None
+
+        merged = torch.cat([a, b], dim=dim)
+        return torch.gather(merged, dim=dim, index=torch.broadcast_to(indices, merged.shape))
+
+
+class RaySampler(ABC):
+    @abstractmethod
+    def sample(self, t0: torch.Tensor, t1: torch.Tensor, n_samples: int) -> torch.Tensor:
+        """
+        :param t0: start time has shape [batch_size, *shape, 1]
+        :param t1: finish time has shape [batch_size, *shape, 1]
+        :param n_samples: number of ts to sample
+        :return: sampled ts of shape [batch_size, *shape, n_samples, 1]
+        """
+
+
+class StratifiedRaySampler(RaySampler):
+    """
+    Instead of fixed intervals, a sample is drawn uniformly at random from each
+    interval.
+    """
+
+    def __init__(self, depth_mode: str = "linear"):
+        """
+        :param depth_mode: linear samples ts linearly in depth. harmonic ensures
+            closer points are sampled more densely.
+        """
+        self.depth_mode = depth_mode
+        assert self.depth_mode in ("linear", "geometric", "harmonic")
+
+    def sample(
+        self,
+        t0: torch.Tensor,
+        t1: torch.Tensor,
+        n_samples: int,
+        epsilon: float = 1e-3,
+    ) -> torch.Tensor:
+        """
+        :param t0: start time has shape [batch_size, *shape, 1]
+        :param t1: finish time has shape [batch_size, *shape, 1]
+        :param n_samples: number of ts to sample
+        :return: sampled ts of shape [batch_size, *shape, n_samples, 1]
+        """
+        ones = [1] * (len(t0.shape) - 1)
+        ts = torch.linspace(0, 1, n_samples).view(*ones, n_samples).to(t0.dtype).to(t0.device)
+
+        if self.depth_mode == "linear":
+            ts = t0 * (1.0 - ts) + t1 * ts
+        elif self.depth_mode == "geometric":
+            ts = (t0.clamp(epsilon).log() * (1.0 - ts) + t1.clamp(epsilon).log() * ts).exp()
+        elif self.depth_mode == "harmonic":
+            # The original NeRF recommends this interpolation scheme for
+            # spherical scenes, but there could be some weird edge cases when
+            # the observer crosses from the inner to outer volume.
+            ts = 1.0 / (1.0 / t0.clamp(epsilon) * (1.0 - ts) + 1.0 / t1.clamp(epsilon) * ts)
+
+        mids = 0.5 * (ts[..., 1:] + ts[..., :-1])
+        upper = torch.cat([mids, t1], dim=-1)
+        lower = torch.cat([t0, mids], dim=-1)
+        t_rand = torch.rand_like(ts)
+
+        ts = lower + (upper - lower) * t_rand
+        return ts.unsqueeze(-1)
+
+
+class ImportanceRaySampler(RaySampler):
+    """
+    Given the initial estimate of densities, this samples more from
+    regions/bins expected to have objects.
+    """
+
+    def __init__(
+        self, volume_range: VolumeRange, raw: AttrDict, blur_pool: bool = False, alpha: float = 1e-5
+    ):
+        """
+        :param volume_range: the range in which a ray intersects the given volume.
+        :param raw: dictionary of raw outputs from the NeRF models of shape
+            [batch_size, *shape, n_coarse_samples, 1]. Should at least contain
+
+            :param ts: earlier samples from the coarse rendering step
+            :param weights: discretized version of density * transmittance
+        :param blur_pool: if true, use 2-tap max + 2-tap blur filter from mip-NeRF.
+        :param alpha: small value to add to weights.
+        """
+        self.volume_range = volume_range
+        self.ts = raw.ts.clone().detach()
+        self.weights = raw.weights.clone().detach()
+        self.blur_pool = blur_pool
+        self.alpha = alpha
+
+    @torch.no_grad()
+    def sample(self, t0: torch.Tensor, t1: torch.Tensor, n_samples: int) -> torch.Tensor:
+        """
+        :param t0: start time has shape [batch_size, *shape, 1]
+        :param t1: finish time has shape [batch_size, *shape, 1]
+        :param n_samples: number of ts to sample
+        :return: sampled ts of shape [batch_size, *shape, n_samples, 1]
+        """
+        lower, upper, _ = self.volume_range.partition(self.ts)
+
+        batch_size, *shape, n_coarse_samples, _ = self.ts.shape
+
+        weights = self.weights
+        if self.blur_pool:
+            padded = torch.cat([weights[..., :1, :], weights, weights[..., -1:, :]], dim=-2)
+            maxes = torch.maximum(padded[..., :-1, :], padded[..., 1:, :])
+            weights = 0.5 * (maxes[..., :-1, :] + maxes[..., 1:, :])
+        weights = weights + self.alpha
+        pmf = weights / weights.sum(dim=-2, keepdim=True)
+        inds = sample_pmf(pmf, n_samples)
+        assert inds.shape == (batch_size, *shape, n_samples, 1)
+        assert (inds >= 0).all() and (inds < n_coarse_samples).all()
+
+        t_rand = torch.rand(inds.shape, device=inds.device)
+        lower_ = torch.gather(lower, -2, inds)
+        upper_ = torch.gather(upper, -2, inds)
+
+        ts = lower_ + (upper_ - lower_) * t_rand
+        ts = torch.sort(ts, dim=-2).values
+        return ts

shap_e/models/nerf/renderer.py

@@ -0,0 +1,301 @@
+from functools import partial
+from typing import Any, Dict, Optional
+
+import torch
+
+from shap_e.models.nn.meta import subdict
+from shap_e.models.renderer import RayRenderer
+from shap_e.models.volume import Volume
+from shap_e.util.collections import AttrDict
+
+from .model import NeRFModel
+from .ray import RayVolumeIntegral, StratifiedRaySampler, render_rays
+
+
+class TwoStepNeRFRenderer(RayRenderer):
+    """
+    Coarse and fine-grained rendering as proposed by NeRF. This class
+    additionally supports background rendering like NeRF++.
+    """
+
+    def __init__(
+        self,
+        n_coarse_samples: int,
+        n_fine_samples: int,
+        void_model: NeRFModel,
+        fine_model: NeRFModel,
+        volume: Volume,
+        coarse_model: Optional[NeRFModel] = None,
+        coarse_background_model: Optional[NeRFModel] = None,
+        fine_background_model: Optional[NeRFModel] = None,
+        outer_volume: Optional[Volume] = None,
+        foreground_stratified_depth_sampling_mode: str = "linear",
+        background_stratified_depth_sampling_mode: str = "linear",
+        importance_sampling_options: Optional[Dict[str, Any]] = None,
+        channel_scale: float = 255,
+        device: torch.device = torch.device("cuda"),
+        **kwargs,
+    ):
+        """
+        :param outer_volume: is where distant objects are encoded.
+        """
+        super().__init__(**kwargs)
+
+        if coarse_model is None:
+            assert (
+                fine_background_model is None or coarse_background_model is None
+            ), "models should be shared for both fg and bg"
+
+        self.n_coarse_samples = n_coarse_samples
+        self.n_fine_samples = n_fine_samples
+        self.void_model = void_model
+        self.coarse_model = coarse_model
+        self.fine_model = fine_model
+        self.volume = volume
+        self.coarse_background_model = coarse_background_model
+        self.fine_background_model = fine_background_model
+        self.outer_volume = outer_volume
+        self.foreground_stratified_depth_sampling_mode = foreground_stratified_depth_sampling_mode
+        self.background_stratified_depth_sampling_mode = background_stratified_depth_sampling_mode
+        self.importance_sampling_options = AttrDict(importance_sampling_options or {})
+        self.channel_scale = channel_scale
+        self.device = device
+        self.to(device)
+
+        if self.coarse_background_model is not None:
+            assert self.fine_background_model is not None
+            assert self.outer_volume is not None
+
+    def render_rays(
+        self,
+        batch: Dict,
+        params: Optional[Dict] = None,
+        options: Optional[Dict] = None,
+    ) -> AttrDict:
+        params = self.update(params)
+
+        batch = AttrDict(batch)
+        if options is None:
+            options = AttrDict()
+        options.setdefault("render_background", True)
+        options.setdefault("render_with_direction", True)
+        options.setdefault("n_coarse_samples", self.n_coarse_samples)
+        options.setdefault("n_fine_samples", self.n_fine_samples)
+        options.setdefault(
+            "foreground_stratified_depth_sampling_mode",
+            self.foreground_stratified_depth_sampling_mode,
+        )
+        options.setdefault(
+            "background_stratified_depth_sampling_mode",
+            self.background_stratified_depth_sampling_mode,
+        )
+
+        shared = self.coarse_model is None
+
+        # First, render rays using the coarse models with stratified ray samples.
+        coarse_model, coarse_key = (
+            (self.fine_model, "fine_model") if shared else (self.coarse_model, "coarse_model")
+        )
+        coarse_model = partial(
+            coarse_model,
+            params=subdict(params, coarse_key),
+            options=options,
+        )
+        parts = [
+            RayVolumeIntegral(
+                model=coarse_model,
+                volume=self.volume,
+                sampler=StratifiedRaySampler(
+                    depth_mode=options.foreground_stratified_depth_sampling_mode,
+                ),
+                n_samples=options.n_coarse_samples,
+            ),
+        ]
+        if options.render_background and self.outer_volume is not None:
+            coarse_background_model, coarse_background_key = (
+                (self.fine_background_model, "fine_background_model")
+                if shared
+                else (self.coarse_background_model, "coarse_background_model")
+            )
+            coarse_background_model = partial(
+                coarse_background_model,
+                params=subdict(params, coarse_background_key),
+                options=options,
+            )
+            parts.append(
+                RayVolumeIntegral(
+                    model=coarse_background_model,
+                    volume=self.outer_volume,
+                    sampler=StratifiedRaySampler(
+                        depth_mode=options.background_stratified_depth_sampling_mode,
+                    ),
+                    n_samples=options.n_coarse_samples,
+                )
+            )
+        coarse_results, samplers, coarse_raw_outputs = render_rays(
+            batch.rays,
+            parts,
+            partial(self.void_model, options=options),
+            shared=shared,
+            render_with_direction=options.render_with_direction,
+            importance_sampling_options=AttrDict(self.importance_sampling_options),
+        )
+
+        # Then, render rays using the fine models with importance-weighted ray samples.
+        fine_model = partial(
+            self.fine_model,
+            params=subdict(params, "fine_model"),
+            options=options,
+        )
+        parts = [
+            RayVolumeIntegral(
+                model=fine_model,
+                volume=self.volume,
+                sampler=samplers[0],
+                n_samples=options.n_fine_samples,
+            ),
+        ]
+        if options.render_background and self.outer_volume is not None:
+            fine_background_model = partial(
+                self.fine_background_model,
+                params=subdict(params, "fine_background_model"),
+                options=options,
+            )
+            parts.append(
+                RayVolumeIntegral(
+                    model=fine_background_model,
+                    volume=self.outer_volume,
+                    sampler=samplers[1],
+                    n_samples=options.n_fine_samples,
+                )
+            )
+        fine_results, *_ = render_rays(
+            batch.rays,
+            parts,
+            partial(self.void_model, options=options),
+            shared=shared,
+            prev_raw_outputs=coarse_raw_outputs,
+            render_with_direction=options.render_with_direction,
+        )
+
+        # Combine results
+        aux_losses = fine_results.output.aux_losses.copy()
+        for key, val in coarse_results.output.aux_losses.items():
+            aux_losses[key + "_coarse"] = val
+
+        return AttrDict(
+            channels=fine_results.output.channels * self.channel_scale,
+            channels_coarse=coarse_results.output.channels * self.channel_scale,
+            distances=fine_results.output.distances,
+            transmittance=fine_results.transmittance,
+            transmittance_coarse=coarse_results.transmittance,
+            t0=fine_results.volume_range.t0,
+            t1=fine_results.volume_range.t1,
+            intersected=fine_results.volume_range.intersected,
+            aux_losses=aux_losses,
+        )
+
+
+class OneStepNeRFRenderer(RayRenderer):
+    """
+    Renders rays using stratified sampling only unlike vanilla NeRF.
+    The same setup as NeRF++.
+    """
+
+    def __init__(
+        self,
+        n_samples: int,
+        void_model: NeRFModel,
+        foreground_model: NeRFModel,
+        volume: Volume,
+        background_model: Optional[NeRFModel] = None,
+        outer_volume: Optional[Volume] = None,
+        foreground_stratified_depth_sampling_mode: str = "linear",
+        background_stratified_depth_sampling_mode: str = "linear",
+        channel_scale: float = 255,
+        device: torch.device = torch.device("cuda"),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.n_samples = n_samples
+        self.void_model = void_model
+        self.foreground_model = foreground_model
+        self.volume = volume
+        self.background_model = background_model
+        self.outer_volume = outer_volume
+        self.foreground_stratified_depth_sampling_mode = foreground_stratified_depth_sampling_mode
+        self.background_stratified_depth_sampling_mode = background_stratified_depth_sampling_mode
+        self.channel_scale = channel_scale
+        self.device = device
+        self.to(device)
+
+    def render_rays(
+        self,
+        batch: Dict,
+        params: Optional[Dict] = None,
+        options: Optional[Dict] = None,
+    ) -> AttrDict:
+        params = self.update(params)
+
+        batch = AttrDict(batch)
+        if options is None:
+            options = AttrDict()
+        options.setdefault("render_background", True)
+        options.setdefault("render_with_direction", True)
+        options.setdefault("n_samples", self.n_samples)
+        options.setdefault(
+            "foreground_stratified_depth_sampling_mode",
+            self.foreground_stratified_depth_sampling_mode,
+        )
+        options.setdefault(
+            "background_stratified_depth_sampling_mode",
+            self.background_stratified_depth_sampling_mode,
+        )
+
+        foreground_model = partial(
+            self.foreground_model,
+            params=subdict(params, "foreground_model"),
+            options=options,
+        )
+        parts = [
+            RayVolumeIntegral(
+                model=foreground_model,
+                volume=self.volume,
+                sampler=StratifiedRaySampler(
+                    depth_mode=options.foreground_stratified_depth_sampling_mode
+                ),
+                n_samples=options.n_samples,
+            ),
+        ]
+        if options.render_background and self.outer_volume is not None:
+            background_model = partial(
+                self.background_model,
+                params=subdict(params, "background_model"),
+                options=options,
+            )
+            parts.append(
+                RayVolumeIntegral(
+                    model=background_model,
+                    volume=self.outer_volume,
+                    sampler=StratifiedRaySampler(
+                        depth_mode=options.background_stratified_depth_sampling_mode
+                    ),
+                    n_samples=options.n_samples,
+                )
+            )
+        results, *_ = render_rays(
+            batch.rays,
+            parts,
+            self.void_model,
+            render_with_direction=options.render_with_direction,
+        )
+
+        return AttrDict(
+            channels=results.output.channels * self.channel_scale,
+            distances=results.output.distances,
+            transmittance=results.transmittance,
+            t0=results.volume_range.t0,
+            t1=results.volume_range.t1,
+            intersected=results.volume_range.intersected,
+            aux_losses=results.output.aux_losses,
+        )

shap_e/models/nerstf/mlp.py

@@ -0,0 +1,174 @@
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+
+from shap_e.models.nn.ops import get_act
+from shap_e.models.query import Query
+from shap_e.models.stf.mlp import MLPModel
+from shap_e.util.collections import AttrDict
+
+
+class MLPDensitySDFModel(MLPModel):
+    def __init__(
+        self,
+        initial_bias: float = -0.1,
+        sdf_activation="tanh",
+        density_activation="exp",
+        **kwargs,
+    ):
+        super().__init__(
+            n_output=2,
+            output_activation="identity",
+            **kwargs,
+        )
+        self.mlp[-1].bias[0].data.fill_(initial_bias)
+        self.sdf_activation = get_act(sdf_activation)
+        self.density_activation = get_act(density_activation)
+
+    def forward(
+        self,
+        query: Query,
+        params: Optional[Dict[str, torch.Tensor]] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> AttrDict[str, Any]:
+        # query.direction is None typically for SDF models and training
+        h, _h_directionless = self._mlp(
+            query.position, query.direction, params=params, options=options
+        )
+        h_sdf, h_density = h.split(1, dim=-1)
+        return AttrDict(
+            density=self.density_activation(h_density),
+            signed_distance=self.sdf_activation(h_sdf),
+        )
+
+
+class MLPNeRSTFModel(MLPModel):
+    def __init__(
+        self,
+        sdf_activation="tanh",
+        density_activation="exp",
+        channel_activation="sigmoid",
+        direction_dependent_shape: bool = True,  # To be able to load old models. Set this to be False in future models.
+        separate_nerf_channels: bool = False,
+        separate_coarse_channels: bool = False,
+        initial_density_bias: float = 0.0,
+        initial_sdf_bias: float = -0.1,
+        **kwargs,
+    ):
+        h_map, h_directionless_map = indices_for_output_mode(
+            direction_dependent_shape=direction_dependent_shape,
+            separate_nerf_channels=separate_nerf_channels,
+            separate_coarse_channels=separate_coarse_channels,
+        )
+        n_output = index_mapping_max(h_map)
+        super().__init__(
+            n_output=n_output,
+            output_activation="identity",
+            **kwargs,
+        )
+        self.direction_dependent_shape = direction_dependent_shape
+        self.separate_nerf_channels = separate_nerf_channels
+        self.separate_coarse_channels = separate_coarse_channels
+        self.sdf_activation = get_act(sdf_activation)
+        self.density_activation = get_act(density_activation)
+        self.channel_activation = get_act(channel_activation)
+        self.h_map = h_map
+        self.h_directionless_map = h_directionless_map
+        self.mlp[-1].bias.data.zero_()
+        layer = -1 if self.direction_dependent_shape else self.insert_direction_at
+        self.mlp[layer].bias[0].data.fill_(initial_sdf_bias)
+        self.mlp[layer].bias[1].data.fill_(initial_density_bias)
+
+    def forward(
+        self,
+        query: Query,
+        params: Optional[Dict[str, torch.Tensor]] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> AttrDict[str, Any]:
+
+        options = AttrDict() if options is None else AttrDict(options)
+        h, h_directionless = self._mlp(
+            query.position, query.direction, params=params, options=options
+        )
+        activations = map_indices_to_keys(self.h_map, h)
+        activations.update(map_indices_to_keys(self.h_directionless_map, h_directionless))
+
+        if options.nerf_level == "coarse":
+            h_density = activations.density_coarse
+        else:
+            h_density = activations.density_fine
+
+        if options.get("rendering_mode", "stf") == "nerf":
+            if options.nerf_level == "coarse":
+                h_channels = activations.nerf_coarse
+            else:
+                h_channels = activations.nerf_fine
+        else:
+            h_channels = activations.stf
+        return AttrDict(
+            density=self.density_activation(h_density),
+            signed_distance=self.sdf_activation(activations.sdf),
+            channels=self.channel_activation(h_channels),
+        )
+
+
+IndexMapping = AttrDict[str, Tuple[int, int]]
+
+
+def indices_for_output_mode(
+    direction_dependent_shape: bool,
+    separate_nerf_channels: bool,
+    separate_coarse_channels: bool,
+) -> Tuple[IndexMapping, IndexMapping]:
+    """
+    Get output mappings for (h, h_directionless).
+    """
+    h_map = AttrDict()
+    h_directionless_map = AttrDict()
+    if direction_dependent_shape:
+        h_map.sdf = (0, 1)
+        if separate_coarse_channels:
+            assert separate_nerf_channels
+            h_map.density_coarse = (1, 2)
+            h_map.density_fine = (2, 3)
+            h_map.stf = (3, 6)
+            h_map.nerf_coarse = (6, 9)
+            h_map.nerf_fine = (9, 12)
+        else:
+            h_map.density_coarse = (1, 2)
+            h_map.density_fine = (1, 2)
+            if separate_nerf_channels:
+                h_map.stf = (2, 5)
+                h_map.nerf_coarse = (5, 8)
+                h_map.nerf_fine = (5, 8)
+            else:
+                h_map.stf = (2, 5)
+                h_map.nerf_coarse = (2, 5)
+                h_map.nerf_fine = (2, 5)
+    else:
+        h_directionless_map.sdf = (0, 1)
+        h_directionless_map.density_coarse = (1, 2)
+        if separate_coarse_channels:
+            h_directionless_map.density_fine = (2, 3)
+        else:
+            h_directionless_map.density_fine = h_directionless_map.density_coarse
+        h_map.stf = (0, 3)
+        if separate_coarse_channels:
+            assert separate_nerf_channels
+            h_map.nerf_coarse = (3, 6)
+            h_map.nerf_fine = (6, 9)
+        else:
+            if separate_nerf_channels:
+                h_map.nerf_coarse = (3, 6)
+            else:
+                h_map.nerf_coarse = (0, 3)
+            h_map.nerf_fine = h_map.nerf_coarse
+    return h_map, h_directionless_map
+
+
+def map_indices_to_keys(mapping: IndexMapping, data: torch.Tensor) -> AttrDict[str, torch.Tensor]:
+    return AttrDict({k: data[..., start:end] for k, (start, end) in mapping.items()})
+
+
+def index_mapping_max(mapping: IndexMapping) -> int:
+    return max(end for _, (_, end) in mapping.items())