first commit

Dockerfile

@@ -0,0 +1,64 @@
+FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install --no-install-recommends -y \
+  build-essential \
+  wget \
+  git \
+  && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+# Switch to the "user" user
+USER user
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH \
+    PYTHONPATH=$HOME/app \
+	PYTHONUNBUFFERED=1 \
+	SYSTEM=spaces
+
+# Install miniconda
+RUN mkdir -p /home/user/conda 
+ENV CONDA_DIR /home/user/conda
+RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
+    /bin/bash ~/miniconda.sh -b -p /home/user/conda
+
+# Put conda in path so we can use conda activate
+ENV PATH=$CONDA_DIR/bin:$PATH
+
+# Activate 
+RUN conda init bash
+
+RUN . /home/user/conda/bin/activate
+
+# Install dependencies
+RUN conda create -n pose python=3.8 
+RUN conda activate pose
+
+RUN pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
+
+RUN pip install -r /code/requirements.txt
+
+RUN conda install -c conda-forge cudatoolkit-dev -y
+ENV CUDA_HOME=$CONDA_PREFIX/pkgs/cuda-toolkit/
+RUN pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
+
+
+RUN wget https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors -P /code/pretrained_models
+RUN wget https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors -P /code/pretrained_models
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+CMD ["python", "app.py"]

README.md

@@ -1,10 +1,10 @@
 ---
-title: Customdiffusion360
-emoji: 🚀
+title: CustomDiffusion360
+emoji: 📷
 colorFrom: gray
 colorTo: yellow
-sdk: gradio
-sdk_version: 4.26.0
+sdk: docker
+app_port: 7860
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,395 @@
+import os
+import gradio as gr
+import plotly.graph_objects as go
+import torch
+import json
+import glob
+import numpy as np
+from PIL import Image
+import time
+import tqdm
+import copy
+
+# Mesh imports
+from pytorch3d.io import load_objs_as_meshes
+from pytorch3d.vis.plotly_vis import AxisArgs, plot_scene
+from pytorch3d.transforms import Transform3d, RotateAxisAngle, Translate, Rotate
+
+from sampling_for_demo import load_and_return_model_and_data, sample, load_base_model
+
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+
+def transform_mesh(mesh, transform, scale=1.0):
+    mesh = mesh.clone()
+    verts = mesh.verts_packed() * scale
+    verts = transform.transform_points(verts)
+    mesh.offset_verts_(verts - mesh.verts_packed())
+    return mesh
+    
+
+def get_input_pose_fig():
+    global curr_camera_dict
+    global obj_filename
+    global plane_trans
+
+    plane_filename = 'assets/plane.obj'
+
+    mesh_scale = 0.75
+    mesh = load_objs_as_meshes([obj_filename], device=device)   
+    mesh.scale_verts_(mesh_scale)
+
+    plane = load_objs_as_meshes([plane_filename], device=device)
+
+    ### plane    
+    rotate_x = RotateAxisAngle(angle=90.0, axis='X', device=device)
+    plane = transform_mesh(plane, rotate_x)
+    translate_y = Translate(0, plane_trans * mesh_scale, 0, device=device)
+    plane = transform_mesh(plane, translate_y)
+
+    fig = plot_scene({
+        "plot": {
+            "object": mesh,
+        },
+    },
+    axis_args=AxisArgs(showgrid=True, backgroundcolor='#cccde0'),
+    xaxis=dict(range=[-1, 1]),
+    yaxis=dict(range=[-1, 1]),
+    zaxis=dict(range=[-1, 1])
+    )
+
+    plane = plane.detach().cpu()
+    verts = plane.verts_packed()
+    faces = plane.faces_packed()
+
+    fig.add_trace(
+        go.Mesh3d(
+            x=verts[:, 0],
+            y=verts[:, 1],
+            z=verts[:, 2],
+            i=faces[:, 0],
+            j=faces[:, 1],
+            k=faces[:, 2],
+            opacity=0.7,
+            color='gray',
+            hoverinfo='skip',
+        ),
+    )
+
+
+    print("fig: curr camera dict")
+    print(curr_camera_dict)
+    camera_dict = curr_camera_dict
+        
+    fig.update_layout(scene=dict(
+        xaxis=dict(showticklabels=True, visible=True),
+        yaxis=dict(showticklabels=True, visible=True),
+        zaxis=dict(showticklabels=True, visible=True),
+    ))
+    # show grid
+    fig.update_layout(scene=dict(
+        xaxis=dict(showgrid=True, gridwidth=1, gridcolor='black'),
+        yaxis=dict(showgrid=True, gridwidth=1, gridcolor='black'),
+        zaxis=dict(showgrid=True, gridwidth=1, gridcolor='black'),
+        bgcolor='#dedede',
+    ))
+
+    fig.update_layout(
+        camera_dict, 
+        width=512, height=512,
+        )
+
+    return fig
+
+
+def run_inference(cam_pose_json, prompt, scale_im, scale, steps, seed):
+    print("prompt is ", prompt)
+    global current_data, current_model
+
+    # run model
+    images = sample(
+        current_model, current_data,
+        num_images=1,
+        prompt=prompt,
+        appendpath="",
+        camera_json=cam_pose_json,
+        train=False,
+        scale=scale,
+        scale_im=scale_im,
+        beta=1.0,
+        num_ref=8,
+        skipreflater=False,
+        num_steps=steps,
+        valid=False,
+        max_images=20,
+        seed=seed
+    )
+
+    result = images[0]
+    print(result.shape)
+    result = Image.fromarray((np.clip(((result+1.0)/2.0).permute(1, 2, 0).cpu().numpy(), 0., 1.)*255).astype(np.uint8))
+    print('result obtained')
+    return result
+
+
+
+def update_curr_camera_dict(camera_json):
+    # TODO: this does not always update the figure, also there's always flashes
+    global curr_camera_dict
+    global prev_camera_dict
+    if camera_json is None:
+        camera_json = json.dumps(prev_camera_dict)
+    camera_json = camera_json.replace("'", "\"")
+    curr_camera_dict = json.loads(camera_json) # ["scene.camera"]
+    print("update curr camera dict")
+    print(curr_camera_dict)
+    return camera_json
+
+
+MODELS_DIR = "pretrained-models/"
+
+def select_and_load_model(category, category_single_id):
+    global current_data, current_model, base_model
+    current_model = None
+    current_model = copy.deepcopy(base_model)
+
+    ### choose model checkpoint and config
+    delta_ckpt = glob.glob(f"{MODELS_DIR}/*{category}{category_single_id}*/checkpoints/step=*.ckpt")[0]
+    print(f"Loading model from {delta_ckpt}")
+
+    logdir = delta_ckpt.split('/checkpoints')[0]
+    config = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml")))[-1]
+
+    start_time = time.time()
+    current_model, current_data = load_and_return_model_and_data(config, current_model,
+                                                                    delta_ckpt=delta_ckpt
+                                                                    )
+        
+    print(f"Time taken to load delta model: {time.time() - start_time:.2f}s")
+    
+    print("!!! model loaded")
+
+    input_prompt = f"photo of a <new1> {category}"
+    return "### Model loaded!", input_prompt
+
+
+global current_data
+global current_model
+current_data = None
+current_model = None
+
+global base_model
+BASE_CONFIG = "configs/train_co3d_concept.yaml"
+BASE_CKPT = "pretrained-models/sd_xl_base_1.0.safetensors"
+
+start_time = time.time()
+base_model = load_base_model(BASE_CONFIG, ckpt=BASE_CKPT, verbose=False)
+print(f"Time taken to load base model: {time.time() - start_time:.2f}s")
+
+global curr_camera_dict 
+curr_camera_dict = {
+        "scene.camera": {
+            "up": {"x": -0.13227683305740356,
+                    "y": -0.9911391735076904,
+                    "z": -0.013464212417602539},
+            "center": {"x": -0.005292057991027832,
+                        "y": 0.020704858005046844,
+                        "z": 0.0873757004737854},
+            "eye": {"x": 0.8585731983184814,
+                    "y": -0.08790968358516693,
+                    "z": -0.40458938479423523},
+        },
+        "scene.aspectratio": {"x": 1.974, "y": 1.974, "z": 1.974},
+        "scene.aspectmode": "manual"
+    }
+
+global prev_camera_dict
+prev_camera_dict = copy.deepcopy(curr_camera_dict)
+
+global obj_filename
+obj_filename = "assets/car0_mesh_centered_flipped.obj"
+global plane_trans
+plane_trans = 0.16
+
+my_fig = get_input_pose_fig()
+
+scripts = open("scripts.js", "r").read()
+
+
+def update_category_single_id(category):
+    global curr_camera_dict
+    global prev_camera_dict
+    global obj_filename
+    global plane_trans
+    choices = None
+    
+    if category == "car":
+        choices = ["0"]
+        curr_camera_dict = {
+            "scene.camera": {
+                "up": {"x": -0.13227683305740356,
+                        "y": -0.9911391735076904,
+                        "z": -0.013464212417602539},
+                "center": {"x": -0.005292057991027832,
+                            "y": 0.020704858005046844,
+                            "z": 0.0873757004737854},
+                "eye": {"x": 0.8585731983184814,
+                        "y": -0.08790968358516693,
+                        "z": -0.40458938479423523},
+            },
+            "scene.aspectratio": {"x": 1.974, "y": 1.974, "z": 1.974},
+            "scene.aspectmode": "manual"
+        }
+        plane_trans = 0.16
+
+    elif category == "chair":
+        choices = ["191"]
+        curr_camera_dict = {
+            "scene.camera": {
+                "up": {"x": 1.0477e-04,
+                        "y": -9.9995e-01,
+                        "z": 1.0288e-02},
+                "center": {"x": 0.0539,
+                            "y":  0.0015,
+                            "z":  0.0007},
+                "eye": {"x": 0.0410,
+                        "y": -0.0091,
+                        "z": -0.9991},
+            },
+            "scene.aspectratio": {"x": 0.9084, "y": 0.9084, "z": 0.9084},
+            "scene.aspectmode": "manual"
+        }
+        plane_trans = 0.38
+        
+    elif category == "motorcycle":
+        choices = ["12"]
+        curr_camera_dict = {
+            "scene.camera": {
+                "up": {"x":  0.0308,
+                        "y":  -0.9994,
+                        "z":  -0.0147},
+                "center": {"x":   0.0240,
+                            "y": -0.0310,
+                            "z":   -0.0016},
+                "eye": {"x": -0.0580,
+                        "y": -0.0188,
+                        "z": -0.9981},
+            },
+            "scene.aspectratio": {"x": 1.5786, "y": 1.5786, "z": 1.5786},
+            "scene.aspectmode": "manual"
+        }
+        plane_trans = 0.16
+
+    elif category == "teddybear":
+        choices = ["31"]
+        curr_camera_dict = {
+            "scene.camera": {
+                "up": {"x": 0.4304,
+                        "y": -0.9023,
+                        "z": -0.0221},
+                "center": {"x": -0.0658,
+                            "y": 0.2081,
+                            "z": 0.0175},
+                "eye": {"x": -0.4456,
+                        "y":   0.0493,
+                        "z": -0.8939},
+            },
+            "scene.aspectratio": {"x": 1.8052, "y": 1.8052, "z": 1.8052},
+            "scene.aspectmode": "manual",
+        }
+        plane_trans = 0.23
+
+    obj_filename = f"assets/{category}{choices[0]}_mesh_centered_flipped.obj"
+    prev_camera_dict = copy.deepcopy(curr_camera_dict)
+    return gr.Dropdown(choices=choices, label="Object ID", value=choices[0])
+
+
+head = """
+    <script src="https://cdn.plot.ly/plotly-2.30.0.min.js" charset="utf-8"></script>
+    """
+
+ORIGINAL_SPACE_ID = 'customdiffusion360'
+SPACE_ID = os.getenv('SPACE_ID')
+
+SHARED_UI_WARNING = f'''## Attention - the demo requires at least 40GB VRAM for inference. Please clone this repository to run on your own machine.   
+<center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="margin-top:0;margin-bottom:0" src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></center> 
+'''
+
+with gr.Blocks(head=head, 
+               css="style.css", 
+               js=scripts,
+               title="Customizing Text-to-Image Diffusion with Camera Viewpoint Control") as demo:
+    
+    gr.HTML("""
+        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+            <div>
+                <h1>Customizing Text-to-Image Diffusion with Camera Viewpoint Control</h1>
+            </div>
+        </div>
+        <div>
+        </br>
+        </div>
+        <hr></hr>
+    """,
+    visible=True
+    )
+
+
+    if SPACE_ID == ORIGINAL_SPACE_ID:
+        gr.Markdown(SHARED_UI_WARNING)
+      
+    with gr.Row():
+        with gr.Column(min_width=150):
+            gr.Markdown("## 1. SELECT CUSTOMIZED MODEL")
+
+            category = gr.Dropdown(choices=["car", "chair", "motorcycle", "teddybear"], label="Category", value="car")
+
+            category_single_id = gr.Dropdown(label="Object ID", choices=["0"], type="value", value="0", visible=False)
+        
+            category.change(update_category_single_id, [category], [category_single_id])
+            
+            load_model_btn = gr.Button(value="Load Model", elem_id="load_model_button")
+
+            load_model_status = gr.Markdown(elem_id="load_model_status", value="### Please select and load a model.")
+
+        with gr.Column(min_width=512):
+            gr.Markdown("## 2. CAMERA POSE VISUALIZATION")
+            
+            # TODO ? don't use gradio plotly element so we can remove menu buttons
+            map = gr.Plot(value=my_fig, min_width=512, elem_id="map")
+
+            ### hidden elements
+            update_pose_btn = gr.Button(value="Update Camera Pose", visible=False, elem_id="update_pose_button")
+            input_pose = gr.TextArea(value=curr_camera_dict, label="Input Camera Pose", visible=False, elem_id="input_pose", interactive=False)
+            check_pose_btn = gr.Button(value="Check Camera Pose", visible=False, elem_id="check_pose_button")
+
+            ## TODO: track init_camera_dict and with js?
+            
+            ### visible elements
+            input_prompt = gr.Textbox(value="photo of a <new1> car", label="Prompt", interactive=True)
+            scale_im = gr.Slider(value=3.5, label="Image guidance scale", minimum=0, maximum=20.0, step=0.1)
+            scale = gr.Slider(value=7.5, label="Text guidance scale", minimum=0, maximum=20.0, step=0.1)
+            steps = gr.Slider(value=10, label="Inference steps", minimum=1, maximum=50, step=1)
+            seed = gr.Textbox(value=42, label="Seed")
+        
+        with gr.Column(min_width=50, elem_id="column_process", scale=0.3):
+            run_btn = gr.Button(value="Run", elem_id="run_button", min_width=50)
+
+
+        with gr.Column(min_width=512):
+            gr.Markdown("## 3. OUR OUTPUT")
+            result = gr.Image(show_label=False, show_download_button=True, width=512, height=512, elem_id="result")
+
+    load_model_btn.click(select_and_load_model, [category, category_single_id], [load_model_status, input_prompt])
+    load_model_btn.click(get_input_pose_fig, [], [map])
+
+    update_pose_btn.click(update_curr_camera_dict, [input_pose], [input_pose],) # js=send_js_camera_to_gradio)
+    # check_pose_btn.click(check_curr_camera_dict, [], [input_pose])
+    run_btn.click(run_inference, [input_pose, input_prompt, scale_im, scale, steps, seed], result)
+
+    demo.load(js=scripts)
+
+
+if __name__ == "__main__":
+    demo.queue().launch(debug=True)

assets/plane.obj

@@ -0,0 +1,14 @@
+o plane
+v 3.000000 -3.000000 0.000000
+v 3.000000 3.000000 0.000000
+v -3.000000 3.000000 0.000000
+v -3.000000  -3.000000 0.000000
+
+vt 3.000000 0.000000
+vt 3.000000 3.000000
+vt 0.000000 3.000000
+vt 0.000000 0.000000
+
+s off
+f 1/1 2/2 3/3
+f 1/1 3/3 4/4

configs/train_co3d_concept.yaml

@@ -0,0 +1,198 @@
+model:
+  base_learning_rate: 1.0e-4
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.13025
+    disable_first_stage_autocast: True
+    trainkeys: pose
+    multiplier: 0.05
+    loss_rgb_lambda: 5
+    loss_fg_lambda: 10
+    loss_bg_lambda: 10
+    log_keys:
+      - txt
+
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+      params:
+        num_idx: 1000
+
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+    network_config:
+      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: False
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [4, 2]
+        num_res_blocks: 2
+        channel_mult: [1, 2, 4]
+        num_head_channels: 64
+        use_linear_in_transformer: True
+        transformer_depth: [1, 2, 10] 
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        image_cross_blocks: [0, 2, 4, 6, 8, 10]
+        rgb: True
+        far: 2
+        num_samples: 24
+        not_add_context_in_triplane: False
+        rgb_predict: True
+        add_lora: False
+        average: False
+        use_prev_weights_imp_sample: True
+        stratified: True
+        imp_sampling_percent: 0.9
+
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+          # crossattn cond
+          - is_trainable: False
+            input_keys: txt,txt_ref
+            target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+            params:
+              layer: hidden
+              layer_idx: 11
+              modifier_token: <new1>
+          # crossattn and vector cond
+          - is_trainable: False
+            input_keys: txt,txt_ref
+            target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+            params:
+              arch: ViT-bigG-14
+              version: laion2b_s39b_b160k
+              layer: penultimate
+              always_return_pooled: True
+              legacy: False
+              modifier_token: <new1>
+          # vector cond
+          - is_trainable: False
+            input_keys: original_size_as_tuple,original_size_as_tuple_ref
+            target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+            params:
+              outdim: 256  # multiplied by two
+          # vector cond
+          - is_trainable: False
+            input_keys: crop_coords_top_left,crop_coords_top_left_ref
+            target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+            params:
+              outdim: 256  # multiplied by two
+          # vector cond
+          - is_trainable: False
+            input_keys: target_size_as_tuple,target_size_as_tuple_ref
+            target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+            params:
+              outdim: 256  # multiplied by two
+
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: pretrained-models/sdxl_vae.safetensors
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [1, 2, 4, 4]
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    loss_fn_config:
+      target: sgm.modules.diffusionmodules.loss.StandardDiffusionLossImgRef
+      params:
+        sigma_sampler_config:
+          target: sgm.modules.diffusionmodules.sigma_sampling.CubicSampling
+          params:
+            num_idx: 1000
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        sigma_sampler_config_ref:
+          target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+          params:
+            num_idx: 50
+
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        num_steps: 50
+
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.VanillaCFGImgRef
+          params:
+            scale: 7.5
+
+data:
+  target: sgm.data.data_co3d.CustomDataDictLoader
+  params:
+    batch_size: 1
+    num_workers: 4
+    category: teddybear
+    img_size: 512
+    skip: 2
+    num_images: 5
+    mask_images: True
+    single_id: 0
+    bbox: True
+    addreg: True
+    drop_ratio: 0.25
+    drop_txt: 0.1
+    modifier_token: <new1>
+
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1600
+      save_top_k: -1
+      save_on_train_epoch_end: False
+
+  callbacks:
+    metrics_over_trainsteps_checkpoint:
+      params:
+        every_n_train_steps: 25000
+
+    image_logger:
+      target: main.ImageLogger
+      params:
+        disabled: False
+        enable_autocast: False
+        batch_frequency: 5000
+        max_images: 8
+        increase_log_steps: False
+        log_first_step: False
+        log_images_kwargs:
+          use_ema_scope: False
+          N: 1
+          n_rows: 2
+
+  trainer:
+    devices: 0,1,2,3
+    benchmark: True
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    max_steps: 1610
+    # val_check_interval: 400

pretrained_models/car0/checkpoints/step=000001600.ckpt

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

pretrained_models/car0/configs/2024-04-12T21-30-20-lightning.yaml

@@ -0,0 +1,31 @@
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1600
+      save_top_k: -1
+      save_on_train_epoch_end: false
+  callbacks:
+    metrics_over_trainsteps_checkpoint:
+      params:
+        every_n_train_steps: 25000
+    image_logger:
+      target: main.ImageLogger
+      params:
+        disabled: false
+        enable_autocast: false
+        batch_frequency: 5000
+        max_images: 8
+        increase_log_steps: false
+        log_first_step: false
+        log_images_kwargs:
+          use_ema_scope: false
+          'N': 1
+          n_rows: 2
+  trainer:
+    devices: 0,1,2,3
+    benchmark: true
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    max_steps: 1610
+    val_check_interval: 400
+    accelerator: gpu

pretrained_models/car0/configs/2024-04-12T21-30-20-project.yaml

@@ -0,0 +1,168 @@
+model:
+  base_learning_rate: 0.0001
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.13025
+    disable_first_stage_autocast: true
+    trainkeys: pose
+    multiplier: 0.05
+    loss_rgb_lambda: 5
+    loss_fg_lambda: 10
+    loss_bg_lambda: 10
+    log_keys:
+    - txt
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+      params:
+        num_idx: 1000
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    network_config:
+      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: false
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions:
+        - 4
+        - 2
+        num_res_blocks: 2
+        channel_mult:
+        - 1
+        - 2
+        - 4
+        num_head_channels: 64
+        use_linear_in_transformer: true
+        transformer_depth:
+        - 1
+        - 2
+        - 10
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        image_cross_blocks:
+        - 0
+        - 2
+        - 4
+        - 6
+        - 8
+        - 10
+        rgb: true
+        far: 2
+        num_samples: 24
+        not_add_context_in_triplane: false
+        rgb_predict: true
+        add_lora: false
+        average: false
+        use_prev_weights_imp_sample: true
+        stratified: true
+        imp_sampling_percent: 0.9
+      use_prev_weights_imp_sample: true
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+          params:
+            layer: hidden
+            layer_idx: 11
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+          params:
+            arch: ViT-bigG-14
+            version: laion2b_s39b_b160k
+            layer: penultimate
+            always_return_pooled: true
+            legacy: false
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: original_size_as_tuple,original_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: crop_coords_top_left,crop_coords_top_left_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: target_size_as_tuple,target_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: /sensei-fs/tenants/Sensei-AdobeResearchTeam/nupkumar1/custom-pose/pretrained-models/sdxl_vae.safetensors
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+    loss_fn_config:
+      target: sgm.modules.diffusionmodules.loss.StandardDiffusionLossImgRef
+      params:
+        sigma_sampler_config:
+          target: sgm.modules.diffusionmodules.sigma_sampling.CubicSampling
+          params:
+            num_idx: 1000
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        sigma_sampler_config_ref:
+          target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+          params:
+            num_idx: 50
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        num_steps: 50
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.VanillaCFGImgRef
+          params:
+            scale: 7.5
+data:
+  target: sgm.data.data_co3d.CustomDataDictLoader
+  params:
+    batch_size: 1
+    num_workers: 4
+    category: car
+    img_size: 512
+    skip: 2
+    num_images: 5
+    mask_images: true
+    single_id: 0
+    bbox: true
+    addreg: true
+    drop_ratio: 0.25
+    drop_txt: 0.1
+    modifier_token: <new1>
+    categoryname: null

pretrained_models/car0/configs/2024-04-13T11-31-55-lightning.yaml

@@ -0,0 +1,31 @@
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1600
+      save_top_k: -1
+      save_on_train_epoch_end: false
+  callbacks:
+    metrics_over_trainsteps_checkpoint:
+      params:
+        every_n_train_steps: 25000
+    image_logger:
+      target: main.ImageLogger
+      params:
+        disabled: false
+        enable_autocast: false
+        batch_frequency: 5000
+        max_images: 8
+        increase_log_steps: false
+        log_first_step: false
+        log_images_kwargs:
+          use_ema_scope: false
+          'N': 1
+          n_rows: 2
+  trainer:
+    devices: 0,1,2,3
+    benchmark: true
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    max_steps: 1610
+    val_check_interval: 400
+    accelerator: gpu

pretrained_models/car0/configs/2024-04-13T11-31-55-project.yaml

@@ -0,0 +1,170 @@
+model:
+  base_learning_rate: 0.0001
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.13025
+    disable_first_stage_autocast: true
+    trainkeys: pose
+    multiplier: 0.05
+    loss_rgb_lambda: 5
+    loss_fg_lambda: 10
+    loss_bg_lambda: 10
+    log_keys:
+    - txt
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+      params:
+        num_idx: 1000
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    network_config:
+      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: false
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions:
+        - 4
+        - 2
+        num_res_blocks: 2
+        channel_mult:
+        - 1
+        - 2
+        - 4
+        num_head_channels: 64
+        use_linear_in_transformer: true
+        transformer_depth:
+        - 1
+        - 2
+        - 10
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        image_cross_blocks:
+        - 0
+        - 2
+        - 4
+        - 6
+        - 8
+        - 10
+        rgb: true
+        far: 2
+        num_samples: 24
+        not_add_context_in_triplane: false
+        rgb_predict: true
+        add_lora: false
+        average: false
+        use_prev_weights_imp_sample: true
+        stratified: true
+        imp_sampling_percent: 0.9
+      use_prev_weights_imp_sample: true
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+          params:
+            layer: hidden
+            layer_idx: 11
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+          params:
+            arch: ViT-bigG-14
+            version: laion2b_s39b_b160k
+            layer: penultimate
+            always_return_pooled: true
+            legacy: false
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: original_size_as_tuple,original_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: crop_coords_top_left,crop_coords_top_left_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: target_size_as_tuple,target_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: /sensei-fs/tenants/Sensei-AdobeResearchTeam/nupkumar1/custom-pose/pretrained-models/sdxl_vae.safetensors
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+    loss_fn_config:
+      target: sgm.modules.diffusionmodules.loss.StandardDiffusionLossImgRef
+      params:
+        sigma_sampler_config:
+          target: sgm.modules.diffusionmodules.sigma_sampling.CubicSampling
+          params:
+            num_idx: 1000
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        sigma_sampler_config_ref:
+          target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+          params:
+            num_idx: 50
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        num_steps: 50
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.VanillaCFGImgRef
+          params:
+            scale: 7.5
+data:
+  target: sgm.data.data_co3d.CustomDataDictLoader
+  params:
+    batch_size: 1
+    num_workers: 4
+    category: car
+    img_size: 512
+    skip: 2
+    num_images: 5
+    mask_images: true
+    single_id: 0
+    bbox: true
+    addreg: true
+    drop_ratio: 0.25
+    drop_txt: 0.1
+    modifier_token: <new1>
+    categoryname: null
+--log_dir: null
+check_logs: null

pretrained_models/car0/configs/2024-04-13T11-42-30-lightning.yaml

@@ -0,0 +1,31 @@
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1600
+      save_top_k: -1
+      save_on_train_epoch_end: false
+  callbacks:
+    metrics_over_trainsteps_checkpoint:
+      params:
+        every_n_train_steps: 25000
+    image_logger:
+      target: main.ImageLogger
+      params:
+        disabled: false
+        enable_autocast: false
+        batch_frequency: 5000
+        max_images: 8
+        increase_log_steps: false
+        log_first_step: false
+        log_images_kwargs:
+          use_ema_scope: false
+          'N': 1
+          n_rows: 2
+  trainer:
+    devices: 0,1,2,3
+    benchmark: true
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    max_steps: 1610
+    val_check_interval: 400
+    accelerator: gpu

pretrained_models/car0/configs/2024-04-13T11-42-30-project.yaml

@@ -0,0 +1,170 @@
+model:
+  base_learning_rate: 0.0001
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.13025
+    disable_first_stage_autocast: true
+    trainkeys: pose
+    multiplier: 0.05
+    loss_rgb_lambda: 5
+    loss_fg_lambda: 10
+    loss_bg_lambda: 10
+    log_keys:
+    - txt
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+      params:
+        num_idx: 1000
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    network_config:
+      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: false
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions:
+        - 4
+        - 2
+        num_res_blocks: 2
+        channel_mult:
+        - 1
+        - 2
+        - 4
+        num_head_channels: 64
+        use_linear_in_transformer: true
+        transformer_depth:
+        - 1
+        - 2
+        - 10
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        image_cross_blocks:
+        - 0
+        - 2
+        - 4
+        - 6
+        - 8
+        - 10
+        rgb: true
+        far: 2
+        num_samples: 24
+        not_add_context_in_triplane: false
+        rgb_predict: true
+        add_lora: false
+        average: false
+        use_prev_weights_imp_sample: true
+        stratified: true
+        imp_sampling_percent: 0.9
+      use_prev_weights_imp_sample: true
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+          params:
+            layer: hidden
+            layer_idx: 11
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+          params:
+            arch: ViT-bigG-14
+            version: laion2b_s39b_b160k
+            layer: penultimate
+            always_return_pooled: true
+            legacy: false
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: original_size_as_tuple,original_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: crop_coords_top_left,crop_coords_top_left_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: target_size_as_tuple,target_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: /sensei-fs/tenants/Sensei-AdobeResearchTeam/nupkumar1/custom-pose/pretrained-models/sdxl_vae.safetensors
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+    loss_fn_config:
+      target: sgm.modules.diffusionmodules.loss.StandardDiffusionLossImgRef
+      params:
+        sigma_sampler_config:
+          target: sgm.modules.diffusionmodules.sigma_sampling.CubicSampling
+          params:
+            num_idx: 1000
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        sigma_sampler_config_ref:
+          target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+          params:
+            num_idx: 50
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        num_steps: 50
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.VanillaCFGImgRef
+          params:
+            scale: 7.5
+data:
+  target: sgm.data.data_co3d.CustomDataDictLoader
+  params:
+    batch_size: 1
+    num_workers: 4
+    category: car
+    img_size: 512
+    skip: 2
+    num_images: 5
+    mask_images: true
+    single_id: 0
+    bbox: true
+    addreg: true
+    drop_ratio: 0.25
+    drop_txt: 0.1
+    modifier_token: <new1>
+    categoryname: null
+--log_dir: null
+check_logs: null

pretrained_models/chair191/checkpoints/step=000001600.ckpt

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

pretrained_models/chair191/configs/2024-04-12T22-10-18-lightning.yaml

@@ -0,0 +1,31 @@
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1600
+      save_top_k: -1
+      save_on_train_epoch_end: false
+  callbacks:
+    metrics_over_trainsteps_checkpoint:
+      params:
+        every_n_train_steps: 25000
+    image_logger:
+      target: main.ImageLogger
+      params:
+        disabled: false
+        enable_autocast: false
+        batch_frequency: 5000
+        max_images: 8
+        increase_log_steps: false
+        log_first_step: false
+        log_images_kwargs:
+          use_ema_scope: false
+          'N': 1
+          n_rows: 2
+  trainer:
+    devices: 0,1,2,3
+    benchmark: true
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    max_steps: 1610
+    val_check_interval: 400
+    accelerator: gpu

pretrained_models/chair191/configs/2024-04-12T22-10-18-project.yaml

@@ -0,0 +1,168 @@
+model:
+  base_learning_rate: 0.0001
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.13025
+    disable_first_stage_autocast: true
+    trainkeys: pose
+    multiplier: 0.05
+    loss_rgb_lambda: 5
+    loss_fg_lambda: 10
+    loss_bg_lambda: 10
+    log_keys:
+    - txt
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+      params:
+        num_idx: 1000
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    network_config:
+      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: false
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions:
+        - 4
+        - 2
+        num_res_blocks: 2
+        channel_mult:
+        - 1
+        - 2
+        - 4
+        num_head_channels: 64
+        use_linear_in_transformer: true
+        transformer_depth:
+        - 1
+        - 2
+        - 10
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        image_cross_blocks:
+        - 0
+        - 2
+        - 4
+        - 6
+        - 8
+        - 10
+        rgb: true
+        far: 2
+        num_samples: 24
+        not_add_context_in_triplane: false
+        rgb_predict: true
+        add_lora: false
+        average: false
+        use_prev_weights_imp_sample: true
+        stratified: true
+        imp_sampling_percent: 0.9
+      use_prev_weights_imp_sample: true
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+          params:
+            layer: hidden
+            layer_idx: 11
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+          params:
+            arch: ViT-bigG-14
+            version: laion2b_s39b_b160k
+            layer: penultimate
+            always_return_pooled: true
+            legacy: false
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: original_size_as_tuple,original_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: crop_coords_top_left,crop_coords_top_left_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: target_size_as_tuple,target_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: /sensei-fs/tenants/Sensei-AdobeResearchTeam/nupkumar1/custom-pose/pretrained-models/sdxl_vae.safetensors
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+    loss_fn_config:
+      target: sgm.modules.diffusionmodules.loss.StandardDiffusionLossImgRef
+      params:
+        sigma_sampler_config:
+          target: sgm.modules.diffusionmodules.sigma_sampling.CubicSampling
+          params:
+            num_idx: 1000
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        sigma_sampler_config_ref:
+          target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+          params:
+            num_idx: 50
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        num_steps: 50
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.VanillaCFGImgRef
+          params:
+            scale: 7.5
+data:
+  target: sgm.data.data_co3d.CustomDataDictLoader
+  params:
+    batch_size: 1
+    num_workers: 4
+    category: chair
+    img_size: 512
+    skip: 2
+    num_images: 5
+    mask_images: true
+    single_id: 191
+    bbox: true
+    addreg: true
+    drop_ratio: 0.25
+    drop_txt: 0.1
+    modifier_token: <new1>
+    categoryname: null

pretrained_models/motorcycle12/checkpoints/step=000001600.ckpt

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

pretrained_models/motorcycle12/configs/2024-04-12T23-30-18-project.yaml

@@ -0,0 +1,168 @@
+model:
+  base_learning_rate: 0.0001
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.13025
+    disable_first_stage_autocast: true
+    trainkeys: pose
+    multiplier: 0.05
+    loss_rgb_lambda: 5
+    loss_fg_lambda: 10
+    loss_bg_lambda: 10
+    log_keys:
+    - txt
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+      params:
+        num_idx: 1000
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    network_config:
+      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: false
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions:
+        - 4
+        - 2
+        num_res_blocks: 2
+        channel_mult:
+        - 1
+        - 2
+        - 4
+        num_head_channels: 64
+        use_linear_in_transformer: true
+        transformer_depth:
+        - 1
+        - 2
+        - 10
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        image_cross_blocks:
+        - 0
+        - 2
+        - 4
+        - 6
+        - 8
+        - 10
+        rgb: true
+        far: 2
+        num_samples: 24
+        not_add_context_in_triplane: false
+        rgb_predict: true
+        add_lora: false
+        average: false
+        use_prev_weights_imp_sample: true
+        stratified: true
+        imp_sampling_percent: 0.9
+      use_prev_weights_imp_sample: true
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+          params:
+            layer: hidden
+            layer_idx: 11
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+          params:
+            arch: ViT-bigG-14
+            version: laion2b_s39b_b160k
+            layer: penultimate
+            always_return_pooled: true
+            legacy: false
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: original_size_as_tuple,original_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: crop_coords_top_left,crop_coords_top_left_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: target_size_as_tuple,target_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: pretrained-models/sdxl_vae.safetensors
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+    loss_fn_config:
+      target: sgm.modules.diffusionmodules.loss.StandardDiffusionLossImgRef
+      params:
+        sigma_sampler_config:
+          target: sgm.modules.diffusionmodules.sigma_sampling.CubicSampling
+          params:
+            num_idx: 1000
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        sigma_sampler_config_ref:
+          target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+          params:
+            num_idx: 50
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        num_steps: 50
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.VanillaCFGImgRef
+          params:
+            scale: 7.5
+data:
+  target: sgm.data.data_co3d.CustomDataDictLoader
+  params:
+    batch_size: 1
+    num_workers: 4
+    category: motorcycle
+    img_size: 512
+    skip: 2
+    num_images: 5
+    mask_images: true
+    single_id: 12
+    bbox: true
+    addreg: true
+    drop_ratio: 0.25
+    drop_txt: 0.1
+    modifier_token: <new1>
+    categoryname: null

pretrained_models/teddybear31/checkpoints/step=000001600.ckpt

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

pretrained_models/teddybear31/configs/2024-04-12T22-50-24-lightning.yaml

@@ -0,0 +1,31 @@
+lightning:
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 1600
+      save_top_k: -1
+      save_on_train_epoch_end: false
+  callbacks:
+    metrics_over_trainsteps_checkpoint:
+      params:
+        every_n_train_steps: 25000
+    image_logger:
+      target: main.ImageLogger
+      params:
+        disabled: false
+        enable_autocast: false
+        batch_frequency: 5000
+        max_images: 8
+        increase_log_steps: false
+        log_first_step: false
+        log_images_kwargs:
+          use_ema_scope: false
+          'N': 1
+          n_rows: 2
+  trainer:
+    devices: 0,1,2,3
+    benchmark: true
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1
+    max_steps: 1610
+    val_check_interval: 400
+    accelerator: gpu

pretrained_models/teddybear31/configs/2024-04-12T22-50-24-project.yaml

@@ -0,0 +1,168 @@
+model:
+  base_learning_rate: 0.0001
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.13025
+    disable_first_stage_autocast: true
+    trainkeys: pose
+    multiplier: 0.05
+    loss_rgb_lambda: 5
+    loss_fg_lambda: 10
+    loss_bg_lambda: 10
+    log_keys:
+    - txt
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+      params:
+        num_idx: 1000
+        weighting_config:
+          target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    network_config:
+      target: sgm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        adm_in_channels: 2816
+        num_classes: sequential
+        use_checkpoint: false
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions:
+        - 4
+        - 2
+        num_res_blocks: 2
+        channel_mult:
+        - 1
+        - 2
+        - 4
+        num_head_channels: 64
+        use_linear_in_transformer: true
+        transformer_depth:
+        - 1
+        - 2
+        - 10
+        context_dim: 2048
+        spatial_transformer_attn_type: softmax-xformers
+        image_cross_blocks:
+        - 0
+        - 2
+        - 4
+        - 6
+        - 8
+        - 10
+        rgb: true
+        far: 2
+        num_samples: 24
+        not_add_context_in_triplane: false
+        rgb_predict: true
+        add_lora: false
+        average: false
+        use_prev_weights_imp_sample: true
+        stratified: true
+        imp_sampling_percent: 0.9
+      use_prev_weights_imp_sample: true
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
+          params:
+            layer: hidden
+            layer_idx: 11
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: txt,txt_ref
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+          params:
+            arch: ViT-bigG-14
+            version: laion2b_s39b_b160k
+            layer: penultimate
+            always_return_pooled: true
+            legacy: false
+            modifier_token: <new1>
+        - is_trainable: false
+          input_keys: original_size_as_tuple,original_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: crop_coords_top_left,crop_coords_top_left_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+        - is_trainable: false
+          input_keys: target_size_as_tuple,target_size_as_tuple_ref
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKLInferenceWrapper
+      params:
+        ckpt_path: /sensei-fs/tenants/Sensei-AdobeResearchTeam/nupkumar1/custom-pose/pretrained-models/sdxl_vae.safetensors
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+    loss_fn_config:
+      target: sgm.modules.diffusionmodules.loss.StandardDiffusionLossImgRef
+      params:
+        sigma_sampler_config:
+          target: sgm.modules.diffusionmodules.sigma_sampling.CubicSampling
+          params:
+            num_idx: 1000
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        sigma_sampler_config_ref:
+          target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+          params:
+            num_idx: 50
+            discretization_config:
+              target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        num_steps: 50
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.VanillaCFGImgRef
+          params:
+            scale: 7.5
+data:
+  target: sgm.data.data_co3d.CustomDataDictLoader
+  params:
+    batch_size: 1
+    num_workers: 4
+    category: teddybear
+    img_size: 512
+    skip: 2
+    num_images: 5
+    mask_images: true
+    single_id: 31
+    bbox: true
+    addreg: true
+    drop_ratio: 0.25
+    drop_txt: 0.1
+    modifier_token: <new1>
+    categoryname: null

requirements.txt

@@ -0,0 +1,37 @@
+omegaconf
+einops
+fire
+tqdm
+pillow
+numpy
+webdataset>=0.2.33
+ninja
+matplotlib
+torchmetrics
+opencv-python==4.6.0.66
+fairscale
+pytorch-lightning==2.0.1
+fire
+fsspec
+kornia==0.6.9
+natsort
+open-clip-torch
+chardet==5.1.0
+tensorboardx==2.6
+pandas
+pudb
+pyyaml
+urllib3<1.27,>=1.25.4
+scipy
+streamlit>=0.73.1
+timm
+tokenizers==0.12.1
+transformers==4.19.1
+triton==2.1.0
+torchdata==0.7.0
+wandb
+invisible-watermark
+xformers
+-e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
+-e git+https://github.com/openai/CLIP.git@main#egg=clip
+-e git+https://github.com/Stability-AI/datapipelines.git@main#egg=sdata

sampling_for_demo.py

@@ -0,0 +1,487 @@
+import glob
+import os
+import sys
+import copy
+from typing import List
+
+import numpy as np
+import torch
+from einops import rearrange
+from omegaconf import OmegaConf
+from PIL import Image
+from pytorch_lightning import seed_everything
+from pytorch3d.renderer.cameras import PerspectiveCameras
+from pytorch3d.renderer import look_at_view_transform
+from pytorch3d.renderer.camera_utils import join_cameras_as_batch
+
+import json
+
+sys.path.append('./')
+from sgm.util import instantiate_from_config, load_safetensors
+
+choices = []
+
+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 load_base_model(config, ckpt=None, verbose=True):
+    config = OmegaConf.load(config)
+    # load model
+    config.model.params.network_config.params.far = 3
+    config.model.params.first_stage_config.params.ckpt_path = "pretrained-models/sdxl_vae.safetensors"
+    guider_config = {'target': 'sgm.modules.diffusionmodules.guiders.ScheduledCFGImgTextRef',
+                        'params': {'scale': 7.5, 'scale_im': 3.5}
+                        }
+    config.model.params.sampler_config.params.guider_config = guider_config
+
+    model = instantiate_from_config(config.model)
+
+    if ckpt is not None:
+        print(f"Loading model from {ckpt}")
+        if ckpt.endswith("ckpt"):
+            pl_sd = torch.load(ckpt, map_location="cpu")
+            if "global_step" in pl_sd:
+                print(f"Global Step: {pl_sd['global_step']}")
+            sd = pl_sd["state_dict"]
+        elif ckpt.endswith("safetensors"):
+            sd = load_safetensors(ckpt)
+            if 'modifier_token' in config.data.params:
+                del sd['conditioner.embedders.0.transformer.text_model.embeddings.token_embedding.weight']
+                del sd['conditioner.embedders.1.model.token_embedding.weight']
+        else:
+            raise NotImplementedError
+        
+    m, u = model.load_state_dict(sd, strict=False)
+
+    model.cuda()
+    model.eval()
+    return model
+
+
+def load_delta_model(model, delta_ckpt=None, verbose=True, freeze=True):
+    """
+    model is preloaded base stable diffusion model
+    """
+
+    msg = None
+    if delta_ckpt is not None:
+        pl_sd_delta = torch.load(delta_ckpt, map_location="cpu")
+        sd_delta = pl_sd_delta["delta_state_dict"]
+
+        # TODO: add new delta loading embedding stuff?
+
+        for name, module in model.model.diffusion_model.named_modules():
+            if len(name.split('.')) > 1 and name.split('.')[-2] == 'transformer_blocks':
+                if hasattr(module, 'pose_emb_layers'):
+                    module.register_buffer('references', sd_delta[f'model.diffusion_model.{name}.references'])
+                    del sd_delta[f'model.diffusion_model.{name}.references']
+
+        m, u = model.load_state_dict(sd_delta, strict=False)
+    
+
+    if len(m) > 0 and verbose:
+        print("missing keys:")
+    if len(u) > 0 and verbose:
+        print("unexpected keys:")
+
+    if freeze:
+        for param in model.parameters():
+            param.requires_grad = False
+
+    model.cuda()
+    model.eval()
+    return model, msg
+
+
+def get_unique_embedder_keys_from_conditioner(conditioner):
+    p = [x.input_keys for x in conditioner.embedders]
+    return list(set([item for sublist in p for item in sublist])) + ['jpg_ref']
+
+
+def customforward(self, x, xr, context=None, contextr=None, pose=None, mask_ref=None, prev_weights=None, timesteps=None, drop_im=None):
+    # note: if no context is given, cross-attention defaults to self-attention
+    if not isinstance(context, list):
+        context = [context]
+    b, c, h, w = x.shape
+    x_in = x
+    fg_masks = []
+    alphas = []
+    rgbs = []
+
+    x = self.norm(x)
+
+    if not self.use_linear:
+        x = self.proj_in(x)
+
+    x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+    if self.use_linear:
+        x = self.proj_in(x)
+
+    prev_weights = None
+    counter = 0
+    for i, block in enumerate(self.transformer_blocks):
+        if i > 0 and len(context) == 1:
+            i = 0  # use same context for each block
+        if self.image_cross and (counter % self.poscontrol_interval == 0):
+            x, fg_mask, weights, alpha, rgb = block(x, context=context[i], context_ref=x, pose=pose, mask_ref=mask_ref, prev_weights=prev_weights, drop_im=drop_im)
+            prev_weights = weights
+            fg_masks.append(fg_mask)
+            if alpha is not None:
+                alphas.append(alpha)
+            if rgb is not None:
+                rgbs.append(rgb)
+        else:
+            x, _, _, _, _ = block(x, context=context[i], drop_im=drop_im)
+        counter += 1
+    if self.use_linear:
+        x = self.proj_out(x)
+    x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+    if not self.use_linear:
+        x = self.proj_out(x)
+    if len(fg_masks) > 0:
+        if len(rgbs) <= 0:
+            rgbs = None
+        if len(alphas) <= 0:
+            alphas = None
+        return x + x_in, None, fg_masks, prev_weights, alphas, rgbs
+    else:
+        return x + x_in, None, None, prev_weights, None, None
+
+
+def _customforward(
+        self, x, context=None, context_ref=None, pose=None, mask_ref=None, prev_weights=None, drop_im=None, additional_tokens=None, n_times_crossframe_attn_in_self=0
+    ):
+    if context_ref is not None:
+        global choices
+        batch_size = x.size(0)
+        # IP2P like sampling or default sampling
+        if batch_size % 3 == 0:
+            batch_size = batch_size // 3
+            context_ref = torch.stack([self.references[:-1][y] for y in choices]).unsqueeze(0).expand(batch_size, -1, -1, -1)
+            context_ref = torch.cat([self.references[-1:].unsqueeze(0).expand(batch_size, context_ref.size(1), -1, -1), context_ref, context_ref], dim=0)
+        else:
+            batch_size = batch_size // 2
+            context_ref = torch.stack([self.references[:-1][y] for y in choices]).unsqueeze(0).expand(batch_size, -1, -1, -1)
+            context_ref = torch.cat([self.references[-1:].unsqueeze(0).expand(batch_size, context_ref.size(1), -1, -1), context_ref], dim=0)
+
+    fg_mask = None
+    weights = None
+    alphas = None
+    predicted_rgb = None
+
+    x = (
+        self.attn1(
+            self.norm1(x),
+            context=context if self.disable_self_attn else None,
+            additional_tokens=additional_tokens,
+            n_times_crossframe_attn_in_self=n_times_crossframe_attn_in_self
+            if not self.disable_self_attn
+            else 0,
+        )
+        + x
+    )
+
+    x = (
+            self.attn2(
+                self.norm2(x), context=context, additional_tokens=additional_tokens,
+            )
+            + x
+        )
+
+    if context_ref is not None:
+        if self.rendered_feat is not None:
+            x = self.pose_emb_layers(torch.cat([x, self.rendered_feat], dim=-1))
+        else:
+            xref, fg_mask, weights, alphas, predicted_rgb = self.reference_attn(x,
+                                                                            context_ref, 
+                                                                            context, 
+                                                                            pose, 
+                                                                            prev_weights, 
+                                                                            mask_ref)
+            self.rendered_feat = xref
+            x = self.pose_emb_layers(torch.cat([x, xref], -1))
+
+    x = self.ff(self.norm3(x)) + x
+    return x, fg_mask, weights, alphas, predicted_rgb
+
+
+def log_images(
+        model,
+        batch,
+        N: int = 1,
+        noise=None,
+        scale_im=3.5,
+        num_steps: int = 10,
+        ucg_keys: List[str] = None,
+        **kwargs,
+        ):
+
+    log = dict()
+    conditioner_input_keys = [e.input_keys for e in model.conditioner.embedders]
+    ucg_keys = conditioner_input_keys
+    pose = batch['pose']
+
+    c, uc = model.conditioner.get_unconditional_conditioning(
+        batch,
+        force_uc_zero_embeddings=ucg_keys
+        if len(model.conditioner.embedders) > 0
+        else [],
+        force_ref_zero_embeddings=True
+    )
+
+    _, n = 1, len(pose)-1
+    sampling_kwargs = {}
+
+    if scale_im > 0:
+        if uc is not None:
+            if isinstance(pose, list):
+                pose = pose[:N]*3
+            else:
+                pose = torch.cat([pose[:N]] * 3)
+    else:
+        if uc is not None:
+            if isinstance(pose, list):
+                pose = pose[:N]*2
+            else:
+                pose = torch.cat([pose[:N]] * 2)
+
+    sampling_kwargs['pose'] = pose
+    sampling_kwargs['drop_im'] = None
+    sampling_kwargs['mask_ref'] = None
+
+    for k in c:
+        if isinstance(c[k], torch.Tensor):
+            c[k], uc[k] = map(lambda y: y[k][:(n+1)*N].to('cuda'), (c, uc))
+
+    import time
+    st = time.time()
+    with model.ema_scope("Plotting"):
+        samples = model.sample(
+                c, shape=noise.shape[1:], uc=uc, batch_size=N, num_steps=num_steps, noise=noise,  **sampling_kwargs
+            )
+        model.clear_rendered_feat()
+    samples = model.decode_first_stage(samples)
+    print("Time taken for sampling", time.time() - st)
+    log["samples"] = samples.cpu()
+
+    return log
+
+
+def process_camera_json(camera_json, example_cam):
+    # replace all single quotes in the camera_json with quotes quotes
+    camera_json = camera_json.replace("'", "\"")
+    print("input camera json")
+    print(camera_json)
+    
+    camera_dict = json.loads(camera_json)["scene.camera"]
+    eye = torch.tensor([camera_dict["eye"]["x"], camera_dict["eye"]["y"], camera_dict["eye"]["z"]], dtype=torch.float32).unsqueeze(0)
+    up = torch.tensor([camera_dict["up"]["x"], camera_dict["up"]["y"], camera_dict["up"]["z"]], dtype=torch.float32).unsqueeze(0)
+    center = torch.tensor([camera_dict["center"]["x"], camera_dict["center"]["y"], camera_dict["center"]["z"]], dtype=torch.float32).unsqueeze(0)
+    new_R, new_T = look_at_view_transform(eye=eye, at=center, up=up)
+
+    ## temp
+    # new_R = torch.tensor([[[ 0.4988,  0.2666,  0.8247],
+    #      [-0.1917, -0.8940,  0.4049],
+    #      [ 0.8453, -0.3601, -0.3948]]], dtype=torch.float32)
+    # new_T = torch.tensor([[ 0.0739, -0.0013,  0.9973]], dtype=torch.float32)
+
+
+    # new_R = torch.tensor([[[ 0.2530,  0.2989,  0.9201],
+    #                         [-0.2652, -0.8932,  0.3631],
+    #                         [ 0.9304, -0.3359, -0.1467],]], dtype=torch.float32)
+    # new_T = torch.tensor([[ 0.0081,  0.0337,  1.0452]], dtype=torch.float32)
+    
+
+    print("focal length", example_cam.focal_length)
+    print("principal point", example_cam.principal_point)
+
+    newcam = PerspectiveCameras(R=new_R, 
+                                T=new_T,
+                                focal_length=example_cam.focal_length,
+                                principal_point=example_cam.principal_point,
+                                image_size=512)
+    
+    print("input pose")
+    print(newcam.get_world_to_view_transform().get_matrix())
+    return newcam
+
+
+def load_and_return_model_and_data(config, model,
+        ckpt="/data/gdsu/customization3d/stable-diffusion-xl-base-1.0/sd_xl_base_1.0.safetensors",
+        delta_ckpt=None,
+        train=False,
+        valid=False,
+        far=3,
+        num_images=1,
+        num_ref=8,
+        max_images=20,
+):
+    config = OmegaConf.load(config)
+    # load data
+    data = None
+    # config.data.params.jitter = False
+    # config.data.params.addreg = False
+    # config.data.params.bbox = False
+
+    # data = instantiate_from_config(config.data)
+    # data = data.train_dataset
+
+    # single_id = data.single_id
+
+    # if hasattr(data, 'rotations'):
+    #     total_images = len(data.rotations[data.sequence_list[single_id]])
+    # else:
+    #     total_images = len(data.annotations['chair'])
+    # print(f"Total images in dataset: {total_images}")
+
+    model, msg = load_delta_model(model, delta_ckpt,)
+
+    # change forward methods to store rendered features and use the pre-calculated reference features
+    def register_recr(net_):
+        if net_.__class__.__name__ == 'SpatialTransformer':
+            print(net_.__class__.__name__, "adding control")
+            bound_method = customforward.__get__(net_, net_.__class__)
+            setattr(net_, 'forward', bound_method)
+            return
+        elif hasattr(net_, 'children'):
+            for net__ in net_.children():
+                register_recr(net__)
+        return
+
+    def register_recr2(net_):
+        if net_.__class__.__name__ == 'BasicTransformerBlock':
+            print(net_.__class__.__name__, "adding control")
+            bound_method = _customforward.__get__(net_, net_.__class__)
+            setattr(net_, 'forward', bound_method)
+            return
+        elif hasattr(net_, 'children'):
+            for net__ in net_.children():
+                register_recr2(net__)
+        return
+
+    sub_nets = model.model.diffusion_model.named_children()
+    for net in sub_nets:
+        register_recr(net[1])
+        register_recr2(net[1])
+
+    # start sampling
+    model.clear_rendered_feat()
+
+    return model, data
+
+
+def sample(model, data,
+        num_images=1,
+        prompt="",
+        appendpath="",
+        camera_json=None,
+        train=False,
+        scale=7.5,
+        scale_im=3.5,
+        beta=1.0,
+        num_ref=8,
+        skipreflater=False,
+        num_steps=10,
+        valid=False,
+        max_images=20,
+        seed=42,
+        camera_path="pretrained-models/car0/camera.bin",
+    ):
+
+    """
+    Only works with num_images=1 (because of camera_json processing)
+    """
+   
+    if num_images != 1:
+        print("forcing num_images to be 1")
+        num_images = 1
+    
+    # set guidance scales
+    model.sampler.guider.scale_im = scale_im
+    model.sampler.guider.scale = scale
+
+    seed_everything(seed)
+  
+    # load cameras
+    cameras_val, cameras_train = torch.load(camera_path)
+    global choices
+    num_ref = 8
+    max_diff = len(cameras_train)/num_ref
+    choices = [int(x) for x in torch.linspace(0, len(cameras_train) - max_diff, num_ref)]
+    cameras_train_final = [cameras_train[i] for i in choices]
+    
+    # start sampling
+    model.clear_rendered_feat()
+
+    if prompt == "":
+        prompt = None
+
+    noise = torch.randn(1, 4, 64, 64).to('cuda').repeat(num_images, 1, 1, 1)
+
+    # random sample camera poses
+    pose_ids = np.random.choice(len(cameras_val), num_images, replace=False)
+    print(pose_ids)
+    pose_ids[0] = 21
+
+    pose = [cameras_val[i] for i in pose_ids]
+
+    print("example camera")
+    print(pose[0].R)
+    print(pose[0].T)
+    print(pose[0].focal_length)
+    print(pose[0].principal_point)
+
+    # prepare batches [if translating then call required functions on the target pose]
+    batches = []
+    for i in range(num_images):
+        batch = {'pose': [pose[i]] + cameras_train_final,
+                    "original_size_as_tuple": torch.tensor([512, 512]).reshape(-1, 2),
+                    "target_size_as_tuple": torch.tensor([512, 512]).reshape(-1, 2),
+                    "crop_coords_top_left": torch.tensor([0, 0]).reshape(-1, 2),
+                    "original_size_as_tuple_ref": torch.tensor([512, 512]).reshape(-1, 2),
+                    "target_size_as_tuple_ref": torch.tensor([512, 512]).reshape(-1, 2),
+                    "crop_coords_top_left_ref": torch.tensor([0, 0]).reshape(-1, 2),
+                    }
+        batch_ = copy.deepcopy(batch)
+        batch_["pose"][0] = process_camera_json(camera_json, pose[0])
+        batch_["pose"] = [join_cameras_as_batch(batch_["pose"])]
+        # print('batched')
+        # print(batch_["pose"][0].get_world_to_view_transform().get_matrix())
+        batches.append(batch_)
+
+    print(f'len batches: {len(batches)}')
+
+    image = None
+
+    with torch.no_grad():
+        for batch in batches:
+            for key in batch.keys():
+                if isinstance(batch[key], torch.Tensor):
+                    batch[key] = batch[key].to('cuda')
+                elif 'pose' in key:
+                    batch[key] = [x.to('cuda') for x in batch[key]]
+                else:
+                    pass
+
+            if prompt is not None:
+                batch["txt"] = [prompt for _ in range(1)]
+                batch["txt_ref"] = [prompt for _ in range(len(batch["pose"])-1)]
+
+            print(batch["txt"])
+            N = 1
+            log_ = log_images(model, batch, N=N, noise=noise.clone()[:N], num_steps=num_steps, scale_im=scale_im)
+            image = log_["samples"]
+           
+    torch.cuda.empty_cache()
+    model.clear_rendered_feat()
+
+    print("generation done")
+    return image
+

scripts.js

@@ -0,0 +1,147 @@
+async () => {
+
+    globalThis.init_camera_dict = {
+        "scene.camera": {
+            "up": {"x": -0.13227683305740356,
+                    "y": -0.9911391735076904,
+                    "z": -0.013464212417602539},
+            "center": {"x": -0.005292057991027832,
+                        "y": 0.020704858005046844,
+                        "z": 0.0873757004737854},
+            "eye": {"x": 0.8585731983184814, 
+                    "y": -0.08790968358516693, 
+                    "z": -0.40458938479423523},
+        },
+        "scene.aspectratio": {"x": 1.974, "y": 1.974, "z": 1.974},
+        "scene.aspectmode": "manual"
+    };
+
+    // globalThis.restrictCamera = (data) => {
+    //     var plotlyDiv = document.getElementById("map").getElementsByClassName('js-plotly-plot')[0];
+    //     // var curr_eye = plotlyDiv.layout.scene.camera.eye;
+    //     // var curr_center = plotlyDiv.layout.scene.camera.center;
+    //     // var curr_up = plotlyDiv.layout.scene.camera.up;
+
+    //     var curr_eye = data["scene.camera"]["eye"];
+    //     var curr_center = data["scene.camera"]["center"];
+    //     var curr_up = data["scene.camera"]["up"];
+        
+    //     var D = Math.sqrt((curr_eye.x - curr_center.x)**2 + (curr_eye.y - curr_center.y)**2 + (curr_eye.z - curr_center.z)**2);
+    //     console.log("D", D);
+
+    //     const max_D = 1.47;
+    //     const min_D = 0.8;
+
+    //     // calculate elevation
+    //     var elevation = Math.atan2(curr_eye.y - curr_center.y, Math.sqrt((curr_eye.x - curr_center.x)**2 + (curr_eye.z - curr_center.z)**2)) * 180 / Math.PI;
+    //     console.log("elevation", elevation);
+    //     const max_elev = 3.2;
+    //     const min_elev = -30;
+
+    //     const eps = 0.01;
+
+    //     if (D > max_D) {
+    //         // find new_eye such that D = max_D
+    //         var new_dict = {
+    //             "scene.camera": {
+    //                 "eye": {
+    //                     "x": curr_center.x + (curr_eye.x - curr_center.x) * max_D / D - eps,
+    //                     "y": curr_center.y + (curr_eye.y - curr_center.y) * max_D / D - eps,
+    //                     "z": curr_center.z + (curr_eye.z - curr_center.z) * max_D / D - eps,
+    //                 },
+    //                 "up": curr_up,
+    //                 "center": curr_center,
+    //             }
+    //         };
+            
+    //         Plotly.relayout(plotlyDiv, new_dict);
+
+    //     } else if (D < min_D) {
+    //         // find new_eye such that D = min_D
+    //         var new_dict = {
+    //             "scene.camera": {
+    //                 "eye": {
+    //                     "x": curr_center.x + (curr_eye.x - curr_center.x) * min_D / D - eps,
+    //                     "y": curr_center.y + (curr_eye.y - curr_center.y) * min_D / D - eps,
+    //                     "z": curr_center.z + (curr_eye.z - curr_center.z) * min_D / D - eps,
+    //                 },
+    //                 "up": curr_up,
+    //                 "center": curr_center,
+    //             }
+    //         };
+        
+    //         Plotly.relayout(plotlyDiv, new_dict);
+    //     }
+
+    //     const eta = 0.001;
+        // if (elevation > max_elev) {
+        //     // find new eye such that y elevation = max_elev
+        //     var new_dict = {
+        //         "scene.camera": {
+        //             "eye": {
+        //                 "x": curr_eye.x,
+        //                 "y": curr_center.y + (curr_eye.y - curr_center.y) * Math.tan((max_elev - eta)  * Math.PI / 180),
+        //                 "z": curr_eye.z,
+        //             },
+        //             "up": curr_up,
+        //             "center": curr_center,
+        //         }
+        //     };
+          
+
+        //     Plotly.relayout(plotlyDiv, new_dict);
+
+        // } else if (elevation < min_elev) {
+        //     // find new eye such that y elevation = min_elev
+        //     var new_dict = {
+        //         "scene.camera": {
+        //             "eye": {
+        //                 "x": curr_eye.x,
+        //                 "y": curr_center.y + (curr_eye.y - curr_center.y) * Math.tan((min_elev + eta)  * Math.PI / 180),
+        //                 "z": curr_eye.z,
+        //             },
+        //             "up": curr_up,
+        //             "center": curr_center,
+        //         }
+        //     };
+            
+        //     Plotly.relayout(plotlyDiv, new_dict);
+        // }
+
+    // }
+
+     globalThis.latestCam = () => {
+        var plotlyDiv = document.getElementById("map").getElementsByClassName('js-plotly-plot')[0];
+
+        globalThis.prev_camera_dict = {};
+        console.log("prev camera dict", globalThis.prev_camera_dict);
+
+        // Listen for the event and log to the console
+        plotlyDiv.on('plotly_relayout', function(data) {
+            console.log('plotly_relayout event triggered:', data);
+            
+            if ("scene.camera.up" in data) {
+                Object.assign(globalThis.prev_camera_dict, globalThis.camera_dict);
+                Object.assign(globalThis.camera_dict, globalThis.init_camera_dict);
+            }
+
+            if ('scene.camera' in data) {
+                Object.assign(globalThis.prev_camera_dict, globalThis.camera_dict);
+                globalThis.camera_dict = data;
+            }
+
+            var camera_json = JSON.stringify(globalThis.camera_dict);
+            var input_pose = document.getElementById("input_pose").getElementsByTagName("textarea")[0];
+            let myEvent = new Event("input")
+            input_pose.value = camera_json;
+            input_pose.dispatchEvent(myEvent);
+
+            var update_pose_btn = document.getElementById("update_pose_button");
+            update_pose_btn.dispatchEvent(new Event("click"));
+            // globalThis.restrictCamera(data);
+        });
+    }
+    
+    return latestCam(this);
+
+}

sgm/__init__.py

@@ -0,0 +1,4 @@
+from .models import AutoencodingEngine, DiffusionEngine
+from .util import get_configs_path, instantiate_from_config
+
+__version__ = "0.1.0"

sgm/data/__init__.py

@@ -0,0 +1 @@
+# from .dataset import StableDataModuleFromConfig

sgm/data/data_co3d.py

@@ -0,0 +1,762 @@
+# code taken and modified from https://github.com/amyxlase/relpose-plus-plus/blob/b33f7d5000cf2430bfcda6466c8e89bc2dcde43f/relpose/dataset/co3d_v2.py#L346)
+import os.path as osp
+import random
+
+import numpy as np
+import torch
+import pytorch_lightning as pl
+
+from PIL import Image, ImageFile 
+import json
+import gzip
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+from pytorch3d.renderer.cameras import PerspectiveCameras
+from pytorch3d.renderer.camera_utils import join_cameras_as_batch
+from pytorch3d.implicitron.dataset.utils import adjust_camera_to_bbox_crop_, adjust_camera_to_image_scale_
+from pytorch3d.transforms import Rotate, Translate
+
+
+CO3D_DIR = "data/training/"
+
+Image.MAX_IMAGE_PIXELS = None
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+# Added: normalize camera poses 
+def intersect_skew_line_groups(p, r, mask):
+    # p, r both of shape (B, N, n_intersected_lines, 3)
+    # mask of shape (B, N, n_intersected_lines)
+    p_intersect, r = intersect_skew_lines_high_dim(p, r, mask=mask)
+    _, p_line_intersect = _point_line_distance(
+        p, r, p_intersect[..., None, :].expand_as(p)
+    )
+    intersect_dist_squared = ((p_line_intersect - p_intersect[..., None, :]) ** 2).sum(
+        dim=-1
+    )
+    return p_intersect, p_line_intersect, intersect_dist_squared, r
+
+
+def intersect_skew_lines_high_dim(p, r, mask=None):
+    # Implements https://en.wikipedia.org/wiki/Skew_lines In more than two dimensions
+    dim = p.shape[-1]
+    # make sure the heading vectors are l2-normed
+    if mask is None:
+        mask = torch.ones_like(p[..., 0])
+    r = torch.nn.functional.normalize(r, dim=-1)
+
+    eye = torch.eye(dim, device=p.device, dtype=p.dtype)[None, None]
+    I_min_cov = (eye - (r[..., None] * r[..., None, :])) * mask[..., None, None]
+    sum_proj = I_min_cov.matmul(p[..., None]).sum(dim=-3)
+    p_intersect = torch.linalg.lstsq(I_min_cov.sum(dim=-3), sum_proj).solution[..., 0]
+
+    if torch.any(torch.isnan(p_intersect)):
+        print(p_intersect)
+        assert False
+    return p_intersect, r
+
+
+def _point_line_distance(p1, r1, p2):
+    df = p2 - p1
+    proj_vector = df - ((df * r1).sum(dim=-1, keepdim=True) * r1)
+    line_pt_nearest = p2 - proj_vector
+    d = (proj_vector).norm(dim=-1)
+    return d, line_pt_nearest
+
+
+def compute_optical_axis_intersection(cameras):
+    centers = cameras.get_camera_center()
+    principal_points = cameras.principal_point
+
+    one_vec = torch.ones((len(cameras), 1))
+    optical_axis = torch.cat((principal_points, one_vec), -1)
+
+    pp = cameras.unproject_points(optical_axis, from_ndc=True, world_coordinates=True)
+
+    pp2 = torch.zeros((pp.shape[0], 3))
+    for i in range(0, pp.shape[0]):
+        pp2[i] = pp[i][i]
+
+    directions = pp2 - centers
+    centers = centers.unsqueeze(0).unsqueeze(0)
+    directions = directions.unsqueeze(0).unsqueeze(0)
+
+    p_intersect, p_line_intersect, _, r = intersect_skew_line_groups(
+        p=centers, r=directions, mask=None
+    )
+
+    p_intersect = p_intersect.squeeze().unsqueeze(0)
+    dist = (p_intersect - centers).norm(dim=-1)
+
+    return p_intersect, dist, p_line_intersect, pp2, r
+
+
+def normalize_cameras(cameras, scale=1.0):
+    """
+    Normalizes cameras such that the optical axes point to the origin and the average
+    distance to the origin is 1.
+
+    Args:
+        cameras (List[camera]).
+    """
+
+    # Let distance from first camera to origin be unit
+    new_cameras = cameras.clone()
+    new_transform = new_cameras.get_world_to_view_transform()
+
+    p_intersect, dist, p_line_intersect, pp, r = compute_optical_axis_intersection(
+        cameras
+    )
+    t = Translate(p_intersect)
+
+    # scale = dist.squeeze()[0]
+    scale = max(dist.squeeze())
+
+    # Degenerate case
+    if scale == 0:
+        print(cameras.T)
+        print(new_transform.get_matrix()[:, 3, :3])
+        return -1
+    assert scale != 0
+
+    new_transform = t.compose(new_transform)
+    new_cameras.R = new_transform.get_matrix()[:, :3, :3]
+    new_cameras.T = new_transform.get_matrix()[:, 3, :3] / scale
+    return new_cameras, p_intersect, p_line_intersect, pp, r
+
+
+def centerandalign(cameras, scale=1.0):
+    """
+    Normalizes cameras such that the optical axes point to the origin and the average
+    distance to the origin is 1.
+
+    Args:
+        cameras (List[camera]).
+    """
+
+    # Let distance from first camera to origin be unit
+    new_cameras = cameras.clone()
+    new_transform = new_cameras.get_world_to_view_transform()
+
+    p_intersect, dist, p_line_intersect, pp, r = compute_optical_axis_intersection(
+        cameras
+    )
+    t = Translate(p_intersect)
+
+    centers = [cam.get_camera_center() for cam in new_cameras]
+    centers = torch.concat(centers, 0).cpu().numpy()
+    m = len(cameras)
+
+    # https://math.stackexchange.com/questions/99299/best-fitting-plane-given-a-set-of-points
+    A = np.hstack((centers[:m, :2], np.ones((m, 1))))
+    B = centers[:m, 2:]
+    if A.shape[0] == 2:
+        x = A.T @ np.linalg.inv(A @ A.T) @ B
+    else:
+        x = np.linalg.inv(A.T @ A) @ A.T @ B
+    a, b, c = x.flatten()
+    n = np.array([a, b, 1])
+    n /= np.linalg.norm(n)
+
+    # https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d
+    v = np.cross(n, [0, 1, 0])
+    s = np.linalg.norm(v)
+    c = np.dot(n, [0, 1, 0])
+    V = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
+    rot = torch.from_numpy(np.eye(3) + V + V @ V * (1 - c) / s**2).float()
+
+    scale = dist.squeeze()[0]
+
+    # Degenerate case
+    if scale == 0:
+        print(cameras.T)
+        print(new_transform.get_matrix()[:, 3, :3])
+        return -1
+    assert scale != 0
+
+    rot = Rotate(rot.T)
+
+    new_transform = rot.compose(t).compose(new_transform)
+    new_cameras.R = new_transform.get_matrix()[:, :3, :3]
+    new_cameras.T = new_transform.get_matrix()[:, 3, :3] / scale
+    return new_cameras
+
+
+def square_bbox(bbox, padding=0.0, astype=None):
+    """
+    Computes a square bounding box, with optional padding parameters.
+
+    Args:
+        bbox: Bounding box in xyxy format (4,).
+
+    Returns:
+        square_bbox in xyxy format (4,).
+    """
+    if astype is None:
+        astype = type(bbox[0])
+    bbox = np.array(bbox)
+    center = ((bbox[:2] + bbox[2:]) / 2).round().astype(int)
+    extents = (bbox[2:] - bbox[:2]) / 2
+    s = (max(extents) * (1 + padding)).round().astype(int)
+    square_bbox = np.array(
+        [center[0] - s, center[1] - s, center[0] + s, center[1] + s],
+        dtype=astype,
+    )
+
+    return square_bbox
+
+
+class Co3dDataset(Dataset):
+    def __init__(
+        self,
+        category,
+        split="train",
+        skip=2,
+        img_size=1024,
+        num_images=4,
+        mask_images=False,
+        single_id=0,
+        bbox=False,
+        modifier_token=None,
+        addreg=False,
+        drop_ratio=0.5,
+        drop_txt=0.1,
+        categoryname=None,
+        aligncameras=False,
+        repeat=100,
+        addlen=False,
+        onlyref=False,
+    ):
+        """
+        Args:
+            category (iterable): List of categories to use. If "all" is in the list,
+                all training categories are used.
+            num_images (int): Default number of images in each batch.
+            normalize_cameras (bool): If True, normalizes cameras so that the
+                intersection of the optical axes is placed at the origin and the norm
+                of the first camera translation is 1.
+            mask_images (bool): If True, masks out the background of the images.
+        """
+        # category = CATEGORIES
+        category = sorted(category.split(','))
+        self.category = category
+        self.single_id = single_id
+        self.addlen = addlen
+        self.onlyref = onlyref
+        self.categoryname = categoryname
+        self.bbox = bbox
+        self.modifier_token = modifier_token
+        self.addreg = addreg
+        self.drop_txt = drop_txt
+        self.skip = skip
+        if self.addreg:
+            with open(f'data/regularization/{category[0]}_sp_generated/caption.txt', "r") as f:
+                self.regcaptions = f.read().splitlines()
+            self.reglen = len(self.regcaptions)
+            self.regimpath = f'data/regularization/{category[0]}_sp_generated'
+
+        self.low_quality_translations = []
+        self.rotations = {}
+        self.category_map = {}
+        co3d_dir = CO3D_DIR
+        for c in category:
+            subset = 'fewview_dev'
+            category_dir = osp.join(co3d_dir, c)
+            frame_file = osp.join(category_dir, "frame_annotations.jgz")
+            sequence_file = osp.join(category_dir, "sequence_annotations.jgz")
+            subset_lists_file = osp.join(category_dir, f"set_lists/set_lists_{subset}.json")
+            bbox_file = osp.join(category_dir, f"{c}_bbox.jgz")
+
+            with open(subset_lists_file) as f:
+                subset_lists_data = json.load(f)
+
+            with gzip.open(sequence_file, "r") as fin:
+                sequence_data = json.loads(fin.read())
+
+            with gzip.open(bbox_file, "r") as fin:
+                bbox_data = json.loads(fin.read())
+
+            with gzip.open(frame_file, "r") as fin:
+                frame_data = json.loads(fin.read())
+
+            frame_data_processed = {}
+            for f_data in frame_data:
+                sequence_name = f_data["sequence_name"]
+                if sequence_name not in frame_data_processed:
+                    frame_data_processed[sequence_name] = {}
+                frame_data_processed[sequence_name][f_data["frame_number"]] = f_data
+
+            good_quality_sequences = set()
+            for seq_data in sequence_data:
+                if seq_data["viewpoint_quality_score"] > 0.5:
+                    good_quality_sequences.add(seq_data["sequence_name"])
+
+            for subset in ["train"]:
+                for seq_name, frame_number, filepath in subset_lists_data[subset]:
+                    if seq_name not in good_quality_sequences:
+                        continue
+
+                    if seq_name not in self.rotations:
+                        self.rotations[seq_name] = []
+                        self.category_map[seq_name] = c
+
+                    mask_path = filepath.replace("images", "masks").replace(".jpg", ".png")
+
+                    frame_data = frame_data_processed[seq_name][frame_number]
+
+                    self.rotations[seq_name].append(
+                        {
+                            "filepath": filepath,
+                            "R": frame_data["viewpoint"]["R"],
+                            "T": frame_data["viewpoint"]["T"],
+                            "focal_length": frame_data["viewpoint"]["focal_length"],
+                            "principal_point": frame_data["viewpoint"]["principal_point"],
+                            "mask": mask_path,
+                            "txt": "a car",
+                            "bbox": bbox_data[mask_path]
+                        }
+                    )
+
+        for seq_name in self.rotations:
+            seq_data = self.rotations[seq_name]
+            cameras = PerspectiveCameras(
+                        focal_length=[data["focal_length"] for data in seq_data],
+                        principal_point=[data["principal_point"] for data in seq_data],
+                        R=[data["R"] for data in seq_data],
+                        T=[data["T"] for data in seq_data],
+                    )
+
+            normalized_cameras, _, _, _, _ = normalize_cameras(cameras)
+            if aligncameras:
+                normalized_cameras = centerandalign(cameras)
+
+            if normalized_cameras == -1:
+                print("Error in normalizing cameras: camera scale was 0")
+                del self.rotations[seq_name]
+                continue
+
+            for i, data in enumerate(seq_data):
+                self.rotations[seq_name][i]["R"] = normalized_cameras.R[i]
+                self.rotations[seq_name][i]["T"] = normalized_cameras.T[i]
+                self.rotations[seq_name][i]["R_original"] = torch.from_numpy(np.array(seq_data[i]["R"]))
+                self.rotations[seq_name][i]["T_original"] = torch.from_numpy(np.array(seq_data[i]["T"]))
+
+                # Make sure translations are not ridiculous
+                if self.rotations[seq_name][i]["T"][0] + self.rotations[seq_name][i]["T"][1] + self.rotations[seq_name][i]["T"][2] > 1e5:
+                    bad_seq = True
+                    self.low_quality_translations.append(seq_name)
+                    break
+
+        for seq_name in self.low_quality_translations:
+            if seq_name in self.rotations:
+                del self.rotations[seq_name]
+
+        self.sequence_list = list(self.rotations.keys())
+
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize(img_size, interpolation=transforms.InterpolationMode.BICUBIC),
+                transforms.ToTensor(),
+                transforms.Lambda(lambda x: x * 2.0 - 1.0)
+            ]
+        )
+        self.transformim = transforms.Compose(
+            [
+                transforms.Resize(img_size, interpolation=transforms.InterpolationMode.BICUBIC),
+                transforms.CenterCrop(img_size),
+                transforms.ToTensor(),
+                transforms.Lambda(lambda x: x * 2.0 - 1.0)
+            ]
+        )
+        self.transformmask = transforms.Compose(
+            [
+                transforms.Resize(img_size // 8),
+                transforms.ToTensor(),
+            ]
+        )
+
+        self.num_images = num_images
+        self.image_size = img_size
+        self.normalize_cameras = normalize_cameras
+        self.mask_images = mask_images
+        self.drop_ratio = drop_ratio
+        self.kernel_tensor = torch.ones((1, 1, 7, 7))
+        self.repeat = repeat
+        print(self.sequence_list, "$$$$$$$$$$$$$$$$$$$$$")
+        self.valid_ids = np.arange(0, len(self.rotations[self.sequence_list[self.single_id]]), skip).tolist()
+        if split == 'test':
+            self.valid_ids = list(set(np.arange(0, len(self.rotations[self.sequence_list[self.single_id]])).tolist()).difference(self.valid_ids))
+
+        print(
+            f"Low quality translation sequences, not used: {self.low_quality_translations}"
+        )
+        print(f"Data size: {len(self)}")
+
+    def __len__(self):
+        return (len(self.valid_ids))*self.repeat + (1 if self.addlen else 0)
+
+    def _padded_bbox(self, bbox, w, h):
+        if w < h:
+            bbox = np.array([0, 0, w, h])
+        else:
+            bbox = np.array([0, 0, w, h])
+        return square_bbox(bbox.astype(np.float32))
+
+    def _crop_bbox(self, bbox, w, h):
+        bbox = square_bbox(bbox.astype(np.float32))
+
+        side_length = bbox[2] - bbox[0]
+        center = (bbox[:2] + bbox[2:]) / 2
+        extent = side_length / 2 
+
+        # Final coordinates need to be integer for cropping.
+        ul = (center - extent).round().astype(int)
+        lr = ul + np.round(2 * extent).astype(int)
+        return np.concatenate((ul, lr))
+
+    def _crop_image(self, image, bbox, white_bg=False):
+        if white_bg:
+            # Only support PIL Images
+            image_crop = Image.new(
+                "RGB", (bbox[2] - bbox[0], bbox[3] - bbox[1]), (255, 255, 255)
+            )
+            image_crop.paste(image, (-bbox[0], -bbox[1]))
+        else:
+            image_crop = transforms.functional.crop(
+                image,
+                top=bbox[1],
+                left=bbox[0],
+                height=bbox[3] - bbox[1],
+                width=bbox[2] - bbox[0],
+            )
+        return image_crop
+
+    def __getitem__(self, index, specific_id=None, validation=False):
+        sequence_name = self.sequence_list[self.single_id]
+
+        metadata = self.rotations[sequence_name]
+
+        if validation:
+            drop_text = False
+            drop_im = False
+        else:
+            drop_im = np.random.uniform(0, 1) < self.drop_ratio
+            if not drop_im:
+                drop_text = np.random.uniform(0, 1) < self.drop_txt
+            else:
+                drop_text = False
+
+        size = self.image_size
+
+        # sample reference ids
+        listofindices = self.valid_ids.copy()
+        max_diff = len(listofindices) // (self.num_images-1)
+        if (index*self.skip) % len(metadata) in listofindices:
+            listofindices.remove((index*self.skip) % len(metadata))
+        references = np.random.choice(np.arange(0, len(listofindices)+1, max_diff), self.num_images-1, replace=False)
+        rem = np.random.randint(0, max_diff)
+        references = [listofindices[(x + rem) % len(listofindices)] for x in references]
+        ids = [(index*self.skip) % len(metadata)] + references
+
+        # special case to save features corresponding to ref image as part of model buffer
+        if self.onlyref:
+            ids = references + [(index*self.skip) % len(metadata)]
+        if specific_id is not None:  # remove this later
+            ids = specific_id
+
+        # get data
+        batch = self.get_data(index=self.single_id, ids=ids)
+
+        # text prompt
+        if self.modifier_token is not None:
+            name = self.category[0] if self.categoryname is None else self.categoryname
+            batch['txt'] = [f'photo of a {self.modifier_token} {name}' for _ in range(len(batch['txt']))]
+
+        # replace with regularization image if drop_im
+        if drop_im and self.addreg:
+            select_id = np.random.randint(0, self.reglen)
+            batch["image"] = [self.transformim(Image.open(f'{self.regimpath}/images/{select_id}.png').convert('RGB'))]
+            batch['txt'] = [self.regcaptions[select_id]]
+            batch["original_size_as_tuple"] = torch.ones_like(batch["original_size_as_tuple"])*1024
+
+        # create camera class and adjust intrinsics for crop
+        cameras = [PerspectiveCameras(R=batch['R'][i].unsqueeze(0),
+                                      T=batch['T'][i].unsqueeze(0),
+                                      focal_length=batch['focal_lengths'][i].unsqueeze(0),
+                                      principal_point=batch['principal_points'][i].unsqueeze(0),
+                                      image_size=self.image_size
+                                      )
+                   for i in range(len(ids))]
+        for i, cam in enumerate(cameras):
+            adjust_camera_to_bbox_crop_(cam, batch["original_size_as_tuple"][i, :2], batch["crop_coords"][i])
+            adjust_camera_to_image_scale_(cam, batch["original_size_as_tuple"][i, 2:], torch.tensor([self.image_size, self.image_size]))
+
+        # create mask and dilated mask for mask based losses
+        batch["depth"] = batch["mask"].clone()
+        batch["mask"] = torch.clamp(torch.nn.functional.conv2d(batch["mask"], self.kernel_tensor, padding='same'), 0, 1)
+        if not self.mask_images:
+            batch["mask"] = [None for i in range(len(ids))]
+
+        # special case to save features corresponding to zero image
+        if index == self.__len__()-1 and self.addlen:
+            batch["image"][0] *= 0.
+
+        return {"jpg": batch["image"][0],
+                "txt": batch["txt"][0] if not drop_text else "",
+                "jpg_ref": batch["image"][1:] if not drop_im else torch.stack([2*torch.rand_like(batch["image"][0])-1. for _ in range(len(ids)-1)], dim=0),
+                "txt_ref": batch["txt"][1:] if not drop_im else ["" for _ in range(len(ids)-1)],
+                "pose": cameras,
+                "mask": batch["mask"][0] if not drop_im else torch.ones_like(batch["mask"][0]),
+                "mask_ref": batch["masks_padding"][1:],
+                "depth": batch["depth"][0] if len(batch["depth"]) > 0 else None,
+                "filepaths": batch["filepaths"],
+                "original_size_as_tuple": batch["original_size_as_tuple"][0][2:],
+                "target_size_as_tuple": torch.ones_like(batch["original_size_as_tuple"][0][2:])*size,
+                "crop_coords_top_left": torch.zeros_like(batch["crop_coords"][0][:2]),
+                "original_size_as_tuple_ref": batch["original_size_as_tuple"][1:][:, 2:],
+                "target_size_as_tuple_ref": torch.ones_like(batch["original_size_as_tuple"][1:][:, 2:])*size,
+                "crop_coords_top_left_ref": torch.zeros_like(batch["crop_coords"][1:][:, :2]),
+                "drop_im": torch.Tensor([1-drop_im*1.])
+                }
+
+    def get_data(self, index=None, sequence_name=None, ids=(0, 1)):
+        if sequence_name is None:
+            sequence_name = self.sequence_list[index]
+        metadata = self.rotations[sequence_name]
+        category = self.category_map[sequence_name]
+        annos = [metadata[i] for i in ids]
+        images = []
+        rotations = []
+        translations = []
+        focal_lengths = []
+        principal_points = []
+        txts = []
+        masks = []
+        filepaths = []
+        images_transformed = []
+        masks_transformed = []
+        original_size_as_tuple = []
+        crop_parameters = []
+        masks_padding = []
+        depths = []
+
+        for counter, anno in enumerate(annos):
+            filepath = anno["filepath"]
+            filepaths.append(filepath)
+            image = Image.open(osp.join(CO3D_DIR, filepath)).convert("RGB")
+
+            mask_name = osp.basename(filepath.replace(".jpg", ".png"))
+
+            mask_path = osp.join(
+                    CO3D_DIR, category, sequence_name, "masks", mask_name
+                )
+            mask = Image.open(mask_path).convert("L")
+
+            if mask.size != image.size:
+                mask = mask.resize(image.size)
+
+            mask_padded = Image.fromarray((np.ones_like(mask) > 0))
+            mask = Image.fromarray((np.array(mask) > 125))
+            masks.append(mask)
+
+            # crop image around object
+            w, h = image.width, image.height
+            bbox = np.array(anno["bbox"])
+            if len(bbox) == 0:
+                bbox = np.array([0, 0, w, h])
+
+            if self.bbox and counter > 0:
+                bbox = self._crop_bbox(bbox, w, h)
+            else:
+                bbox = self._padded_bbox(None, w, h)
+            image = self._crop_image(image, bbox)
+            mask = self._crop_image(mask, bbox)
+            mask_padded = self._crop_image(mask_padded, bbox)
+            masks_padding.append(self.transformmask(mask_padded))
+            images_transformed.append(self.transform(image))
+            masks_transformed.append(self.transformmask(mask))
+
+            crop_parameters.append(torch.tensor([bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1] ]).int())
+            original_size_as_tuple.append(torch.tensor([w, h, bbox[2] - bbox[0], bbox[3] - bbox[1]]))
+            images.append(image)
+            rotations.append(anno["R"])
+            translations.append(anno["T"])
+            focal_lengths.append(torch.tensor(anno["focal_length"]))
+            principal_points.append(torch.tensor(anno["principal_point"]))
+            txts.append(anno["txt"])
+
+        images = images_transformed
+        batch = {
+            "model_id": sequence_name,
+            "category": category,
+            "original_size_as_tuple": torch.stack(original_size_as_tuple),
+            "crop_coords": torch.stack(crop_parameters),
+            "n": len(metadata),
+            "ind": torch.tensor(ids),
+            "txt": txts,
+            "filepaths": filepaths,
+            "masks_padding": torch.stack(masks_padding) if len(masks_padding) > 0 else [],
+            "depth": torch.stack(depths) if len(depths) > 0 else [],
+        }
+
+        batch["R"] = torch.stack(rotations)
+        batch["T"] = torch.stack(translations)
+        batch["focal_lengths"] = torch.stack(focal_lengths)
+        batch["principal_points"] = torch.stack(principal_points)
+
+        # Add images
+        if self.transform is None:
+            batch["image"] = images
+        else:
+            batch["image"] = torch.stack(images)
+            batch["mask"] = torch.stack(masks_transformed)
+
+        return batch
+
+    @staticmethod
+    def collate_fn(batch):
+        """A function to collate the data across batches. This function must be passed to pytorch's DataLoader to collate batches.
+        Args:
+            batch(list): List of objects returned by this class' __getitem__ function. This is given by pytorch's dataloader that calls __getitem__
+                         multiple times and expects a collated batch.
+        Returns:
+            dict: The collated dictionary representing the data in the batch.
+        """
+        result = {
+            "jpg": [],
+            "txt": [],
+            "jpg_ref": [],
+            "txt_ref": [],
+            "pose": [],
+            "original_size_as_tuple": [],
+            "original_size_as_tuple_ref": [],
+            "crop_coords_top_left": [],
+            "crop_coords_top_left_ref": [],
+            "target_size_as_tuple_ref": [],
+            "target_size_as_tuple": [],
+            "drop_im": [],
+            "mask_ref": [],
+        }
+        if batch[0]["mask"] is not None:
+            result["mask"] = []
+        if batch[0]["depth"] is not None:
+            result["depth"] = []
+
+        for batch_obj in batch:
+            for key in result.keys():
+                result[key].append(batch_obj[key])
+        for key in result.keys():
+            if not (key == 'pose' or 'txt' in key or 'size_as_tuple_ref' in key or 'coords_top_left_ref' in key):
+                result[key] = torch.stack(result[key], dim=0)
+            elif 'txt_ref' in key:
+                result[key] = [item for sublist in result[key] for item in sublist]
+            elif 'size_as_tuple_ref' in key or 'coords_top_left_ref' in key:
+                result[key] = torch.cat(result[key], dim=0)
+            elif 'pose' in key:
+                result[key] = [join_cameras_as_batch(cameras) for cameras in result[key]]
+
+        return result
+
+
+class CustomDataDictLoader(pl.LightningDataModule):
+    def __init__(
+        self,
+        category,
+        batch_size,
+        mask_images=False,
+        skip=1,
+        img_size=1024,
+        num_images=4,
+        num_workers=0,
+        shuffle=True,
+        single_id=0,
+        modifier_token=None,
+        bbox=False,
+        addreg=False,
+        drop_ratio=0.5,
+        jitter=False,
+        drop_txt=0.1,
+        categoryname=None,
+    ):
+        super().__init__()
+
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.shuffle = shuffle
+        self.train_dataset = Co3dDataset(category,
+                                         img_size=img_size,
+                                         mask_images=mask_images,
+                                         skip=skip,
+                                         num_images=num_images,
+                                         single_id=single_id,
+                                         modifier_token=modifier_token,
+                                         bbox=bbox,
+                                         addreg=addreg,
+                                         drop_ratio=drop_ratio,
+                                         drop_txt=drop_txt,
+                                         categoryname=categoryname,
+                                         )
+        self.val_dataset = Co3dDataset(category,
+                                       img_size=img_size,
+                                       mask_images=mask_images,
+                                       skip=skip,
+                                       num_images=2,
+                                       single_id=single_id,
+                                       modifier_token=modifier_token,
+                                       bbox=bbox,
+                                       addreg=addreg,
+                                       drop_ratio=0.,
+                                       drop_txt=0.,
+                                       categoryname=categoryname,
+                                       repeat=1,
+                                       addlen=True,
+                                       onlyref=True,
+                                       )
+        self.test_dataset = Co3dDataset(category,
+                                        img_size=img_size,
+                                        mask_images=mask_images,
+                                        split="test",
+                                        skip=skip,
+                                        num_images=2,
+                                        single_id=single_id,
+                                        modifier_token=modifier_token,
+                                        bbox=False,
+                                        addreg=addreg,
+                                        drop_ratio=0.,
+                                        drop_txt=0.,
+                                        categoryname=categoryname,
+                                        repeat=1,
+                                        )
+        self.collate_fn = Co3dDataset.collate_fn
+
+    def prepare_data(self):
+        pass
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle,
+            num_workers=self.num_workers,
+            collate_fn=self.collate_fn,
+            drop_last=True,
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            collate_fn=self.collate_fn,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            collate_fn=self.collate_fn,
+            drop_last=True
+        )

sgm/lr_scheduler.py

@@ -0,0 +1,135 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+    """
+    note: use with a base_lr of 1.0
+    """
+
+    def __init__(
+        self,
+        warm_up_steps,
+        lr_min,
+        lr_max,
+        lr_start,
+        max_decay_steps,
+        verbosity_interval=0,
+    ):
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.0
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n, **kwargs):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (
+                self.lr_max - self.lr_start
+            ) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (
+                self.lr_max_decay_steps - self.lr_warm_up_steps
+            )
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+                1 + np.cos(t * np.pi)
+            )
+            self.last_lr = lr
+            return lr
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+    """
+    supports repeated iterations, configurable via lists
+    note: use with a base_lr of 1.0.
+    """
+
+    def __init__(
+        self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0
+    ):
+        assert (
+            len(warm_up_steps)
+            == len(f_min)
+            == len(f_max)
+            == len(f_start)
+            == len(cycle_lengths)
+        )
+        self.lr_warm_up_steps = warm_up_steps
+        self.f_start = f_start
+        self.f_min = f_min
+        self.f_max = f_max
+        self.cycle_lengths = cycle_lengths
+        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+        self.last_f = 0.0
+        self.verbosity_interval = verbosity_interval
+
+    def find_in_interval(self, n):
+        interval = 0
+        for cl in self.cum_cycles[1:]:
+            if n <= cl:
+                return interval
+            interval += 1
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(
+                    f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                    f"current cycle {cycle}"
+                )
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
+                cycle
+            ] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            t = (n - self.lr_warm_up_steps[cycle]) / (
+                self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle]
+            )
+            t = min(t, 1.0)
+            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+                1 + np.cos(t * np.pi)
+            )
+            self.last_f = f
+            return f
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(
+                    f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                    f"current cycle {cycle}"
+                )
+
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[
+                cycle
+            ] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (
+                self.cycle_lengths[cycle] - n
+            ) / (self.cycle_lengths[cycle])
+            self.last_f = f
+            return f

sgm/models/__init__.py

@@ -0,0 +1,2 @@
+from .autoencoder import AutoencodingEngine
+from .diffusion import DiffusionEngine

sgm/models/autoencoder.py

@@ -0,0 +1,335 @@
+import re
+from abc import abstractmethod
+from contextlib import contextmanager
+from typing import Any, Dict, Tuple, Union
+
+import pytorch_lightning as pl
+import torch
+from omegaconf import ListConfig
+from packaging import version
+from safetensors.torch import load_file as load_safetensors
+
+from ..modules.diffusionmodules.model import Decoder, Encoder
+from ..modules.distributions.distributions import DiagonalGaussianDistribution
+from ..modules.ema import LitEma
+from ..util import default, get_obj_from_str, instantiate_from_config
+
+
+class AbstractAutoencoder(pl.LightningModule):
+    """
+    This is the base class for all autoencoders, including image autoencoders, image autoencoders with discriminators,
+    unCLIP models, etc. Hence, it is fairly general, and specific features
+    (e.g. discriminator training, encoding, decoding) must be implemented in subclasses.
+    """
+
+    def __init__(
+        self,
+        ema_decay: Union[None, float] = None,
+        monitor: Union[None, str] = None,
+        input_key: str = "jpg",
+        ckpt_path: Union[None, str] = None,
+        ignore_keys: Union[Tuple, list, ListConfig] = (),
+    ):
+        super().__init__()
+        self.input_key = input_key
+        self.use_ema = ema_decay is not None
+        if monitor is not None:
+            self.monitor = monitor
+
+        if self.use_ema:
+            self.model_ema = LitEma(self, decay=ema_decay)
+            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+        if version.parse(torch.__version__) >= version.parse("2.0.0"):
+            self.automatic_optimization = False
+
+    def init_from_ckpt(
+        self, path: str, ignore_keys: Union[Tuple, list, ListConfig] = tuple()
+    ) -> None:
+        if path.endswith("ckpt"):
+            sd = torch.load(path, map_location="cpu")["state_dict"]
+        elif path.endswith("safetensors"):
+            sd = load_safetensors(path)
+        else:
+            raise NotImplementedError
+
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if re.match(ik, k):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        missing, unexpected = self.load_state_dict(sd, strict=False)
+        print(
+            f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+        )
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+        if len(unexpected) > 0:
+            print(f"Unexpected Keys: {unexpected}")
+
+    @abstractmethod
+    def get_input(self, batch) -> Any:
+        raise NotImplementedError()
+
+    def on_train_batch_end(self, *args, **kwargs):
+        # for EMA computation
+        if self.use_ema:
+            self.model_ema(self)
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.parameters())
+            self.model_ema.copy_to(self)
+            if context is not None:
+                print(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.parameters())
+                if context is not None:
+                    print(f"{context}: Restored training weights")
+
+    @abstractmethod
+    def encode(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("encode()-method of abstract base class called")
+
+    @abstractmethod
+    def decode(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("decode()-method of abstract base class called")
+
+    def instantiate_optimizer_from_config(self, params, lr, cfg):
+        print(f"loading >>> {cfg['target']} <<< optimizer from config")
+        return get_obj_from_str(cfg["target"])(
+            params, lr=lr, **cfg.get("params", dict())
+        )
+
+    def configure_optimizers(self) -> Any:
+        raise NotImplementedError()
+
+
+class AutoencodingEngine(AbstractAutoencoder):
+    """
+    Base class for all image autoencoders that we train, like VQGAN or AutoencoderKL
+    (we also restore them explicitly as special cases for legacy reasons).
+    Regularizations such as KL or VQ are moved to the regularizer class.
+    """
+
+    def __init__(
+        self,
+        *args,
+        encoder_config: Dict,
+        decoder_config: Dict,
+        loss_config: Dict,
+        regularizer_config: Dict,
+        optimizer_config: Union[Dict, None] = None,
+        lr_g_factor: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        # todo: add options to freeze encoder/decoder
+        self.encoder = instantiate_from_config(encoder_config)
+        self.decoder = instantiate_from_config(decoder_config)
+        self.loss = instantiate_from_config(loss_config)
+        self.regularization = instantiate_from_config(regularizer_config)
+        self.optimizer_config = default(
+            optimizer_config, {"target": "torch.optim.Adam"}
+        )
+        self.lr_g_factor = lr_g_factor
+
+    def get_input(self, batch: Dict) -> torch.Tensor:
+        # assuming unified data format, dataloader returns a dict.
+        # image tensors should be scaled to -1 ... 1 and in channels-first format (e.g., bchw instead if bhwc)
+        return batch[self.input_key]
+
+    def get_autoencoder_params(self) -> list:
+        params = (
+            list(self.encoder.parameters())
+            + list(self.decoder.parameters())
+            + list(self.regularization.get_trainable_parameters())
+            + list(self.loss.get_trainable_autoencoder_parameters())
+        )
+        return params
+
+    def get_discriminator_params(self) -> list:
+        params = list(self.loss.get_trainable_parameters())  # e.g., discriminator
+        return params
+
+    def get_last_layer(self):
+        return self.decoder.get_last_layer()
+
+    def encode(self, x: Any, return_reg_log: bool = False) -> Any:
+        z = self.encoder(x)
+        z, reg_log = self.regularization(z)
+        if return_reg_log:
+            return z, reg_log
+        return z
+
+    def decode(self, z: Any) -> torch.Tensor:
+        x = self.decoder(z)
+        return x
+
+    def forward(self, x: Any) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        z, reg_log = self.encode(x, return_reg_log=True)
+        dec = self.decode(z)
+        return z, dec, reg_log
+
+    def training_step(self, batch, batch_idx, optimizer_idx) -> Any:
+        x = self.get_input(batch)
+        z, xrec, regularization_log = self(x)
+
+        if optimizer_idx == 0:
+            # autoencode
+            aeloss, log_dict_ae = self.loss(
+                regularization_log,
+                x,
+                xrec,
+                optimizer_idx,
+                self.global_step,
+                last_layer=self.get_last_layer(),
+                split="train",
+            )
+
+            self.log_dict(
+                log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True
+            )
+            return aeloss
+
+        if optimizer_idx == 1:
+            # discriminator
+            discloss, log_dict_disc = self.loss(
+                regularization_log,
+                x,
+                xrec,
+                optimizer_idx,
+                self.global_step,
+                last_layer=self.get_last_layer(),
+                split="train",
+            )
+            self.log_dict(
+                log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True
+            )
+            return discloss
+
+    def validation_step(self, batch, batch_idx) -> Dict:
+        log_dict = self._validation_step(batch, batch_idx)
+        with self.ema_scope():
+            log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
+            log_dict.update(log_dict_ema)
+        return log_dict
+
+    def _validation_step(self, batch, batch_idx, postfix="") -> Dict:
+        x = self.get_input(batch)
+
+        z, xrec, regularization_log = self(x)
+        aeloss, log_dict_ae = self.loss(
+            regularization_log,
+            x,
+            xrec,
+            0,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val" + postfix,
+        )
+
+        discloss, log_dict_disc = self.loss(
+            regularization_log,
+            x,
+            xrec,
+            1,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val" + postfix,
+        )
+        self.log(f"val{postfix}/rec_loss", log_dict_ae[f"val{postfix}/rec_loss"])
+        log_dict_ae.update(log_dict_disc)
+        self.log_dict(log_dict_ae)
+        return log_dict_ae
+
+    def configure_optimizers(self) -> Any:
+        ae_params = self.get_autoencoder_params()
+        disc_params = self.get_discriminator_params()
+
+        opt_ae = self.instantiate_optimizer_from_config(
+            ae_params,
+            default(self.lr_g_factor, 1.0) * self.learning_rate,
+            self.optimizer_config,
+        )
+        opt_disc = self.instantiate_optimizer_from_config(
+            disc_params, self.learning_rate, self.optimizer_config
+        )
+
+        return [opt_ae, opt_disc], []
+
+    @torch.no_grad()
+    def log_images(self, batch: Dict, **kwargs) -> Dict:
+        log = dict()
+        x = self.get_input(batch)
+        _, xrec, _ = self(x)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        with self.ema_scope():
+            _, xrec_ema, _ = self(x)
+            log["reconstructions_ema"] = xrec_ema
+        return log
+
+
+class AutoencoderKL(AutoencodingEngine):
+    def __init__(self, embed_dim: int, **kwargs):
+        ddconfig = kwargs.pop("ddconfig")
+        ckpt_path = kwargs.pop("ckpt_path", None)
+        ignore_keys = kwargs.pop("ignore_keys", ())
+        super().__init__(
+            encoder_config={"target": "torch.nn.Identity"},
+            decoder_config={"target": "torch.nn.Identity"},
+            regularizer_config={"target": "torch.nn.Identity"},
+            loss_config=kwargs.pop("lossconfig"),
+            **kwargs,
+        )
+        assert ddconfig["double_z"]
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        self.embed_dim = embed_dim
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+    def encode(self, x):
+        assert (
+            not self.training
+        ), f"{self.__class__.__name__} only supports inference currently"
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior
+
+    def decode(self, z, **decoder_kwargs):
+        z = self.post_quant_conv(z)
+        dec = self.decoder(z, **decoder_kwargs)
+        return dec
+
+
+class AutoencoderKLInferenceWrapper(AutoencoderKL):
+    def encode(self, x):
+        return super().encode(x).sample()
+
+
+class IdentityFirstStage(AbstractAutoencoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_input(self, x: Any) -> Any:
+        return x
+
+    def encode(self, x: Any, *args, **kwargs) -> Any:
+        return x
+
+    def decode(self, x: Any, *args, **kwargs) -> Any:
+        return x

sgm/models/diffusion.py

@@ -0,0 +1,556 @@
+from contextlib import contextmanager
+from typing import Any, Dict, List, Tuple, Union, DefaultDict
+
+import pytorch_lightning as pl
+import torch
+from omegaconf import ListConfig, OmegaConf
+from safetensors.torch import load_file as load_safetensors
+from torch.optim.lr_scheduler import LambdaLR
+from einops import rearrange
+import math
+import torch.nn as nn
+from ..modules import UNCONDITIONAL_CONFIG
+from ..modules.diffusionmodules.wrappers import OPENAIUNETWRAPPER
+from ..modules.ema import LitEma
+from ..util import (
+    default,
+    disabled_train,
+    get_obj_from_str,
+    instantiate_from_config,
+    log_txt_as_img,
+)
+
+
+import collections
+from functools import partial
+
+
+def save_activations(
+        activations: DefaultDict,
+        name: str,
+        module: nn.Module,
+        inp: Tuple,
+        out: torch.Tensor
+) -> None:
+    """PyTorch Forward hook to save outputs at each forward
+    pass. Mutates specified dict objects with each fwd pass.
+    """
+    if isinstance(out, tuple):
+        if out[1] is None:
+            activations[name].append(out[0].detach())
+
+class DiffusionEngine(pl.LightningModule):
+    def __init__(
+        self,
+        network_config,
+        denoiser_config,
+        first_stage_config,
+        conditioner_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        sampler_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        optimizer_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        scheduler_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        loss_fn_config: Union[None, Dict, ListConfig, OmegaConf] = None,
+        network_wrapper: Union[None, str] = None,
+        ckpt_path: Union[None, str] = None,
+        use_ema: bool = False,
+        ema_decay_rate: float = 0.9999,
+        scale_factor: float = 1.0,
+        disable_first_stage_autocast=False,
+        input_key: str = "jpg",
+        log_keys: Union[List, None] = None,
+        no_cond_log: bool = False,
+        compile_model: bool = False,
+        trainkeys='pose',
+        multiplier=0.05,
+        loss_rgb_lambda=20.,
+        loss_fg_lambda=10.,
+        loss_bg_lambda=20.,
+    ):
+        super().__init__()
+        self.log_keys = log_keys
+        self.input_key = input_key
+        self.trainkeys = trainkeys
+        self.multiplier = multiplier
+        self.loss_rgb_lambda = loss_rgb_lambda
+        self.loss_fg_lambda = loss_fg_lambda
+        self.loss_bg_lambda = loss_bg_lambda
+        self.rgb = network_config.params.rgb
+        self.rgb_predict = network_config.params.rgb_predict
+        self.add_token = ('modifier_token' in conditioner_config.params.emb_models[1].params)
+        self.optimizer_config = default(
+            optimizer_config, {"target": "torch.optim.AdamW"}
+        )
+        model = instantiate_from_config(network_config)
+        self.model = get_obj_from_str(default(network_wrapper, OPENAIUNETWRAPPER))(
+            model, compile_model=compile_model
+        )
+
+        self.denoiser = instantiate_from_config(denoiser_config)
+        self.sampler = (
+            instantiate_from_config(sampler_config)
+            if sampler_config is not None
+            else None
+        )
+        self.conditioner = instantiate_from_config(
+            default(conditioner_config, UNCONDITIONAL_CONFIG)
+        )
+        self.scheduler_config = scheduler_config
+        self._init_first_stage(first_stage_config)
+
+        self.loss_fn = (
+            instantiate_from_config(loss_fn_config)
+            if loss_fn_config is not None
+            else None
+        )
+
+        self.use_ema = use_ema
+        if self.use_ema:
+            self.model_ema = LitEma(self.model, decay=ema_decay_rate)
+            print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        self.scale_factor = scale_factor
+        self.disable_first_stage_autocast = disable_first_stage_autocast
+        self.no_cond_log = no_cond_log
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path)
+
+        blocks = []
+        if self.trainkeys == 'poseattn':
+            for x in self.model.diffusion_model.named_parameters():
+                if not ('pose' in x[0] or 'transformer_blocks' in x[0]):
+                    x[1].requires_grad = False
+                else:
+                    if 'pose' in x[0]:
+                        x[1].requires_grad = True
+                        blocks.append(x[0].split('.pose')[0])
+
+            blocks = set(blocks)
+            for x in self.model.diffusion_model.named_parameters():
+                if 'transformer_blocks' in x[0]:
+                    reqgrad = False
+                    for each in blocks:
+                        if each in x[0] and ('attn1' in x[0] or 'attn2' in x[0] or 'pose' in x[0]):
+                            reqgrad = True
+                            x[1].requires_grad = True
+                    if not reqgrad:
+                        x[1].requires_grad = False
+        elif self.trainkeys == 'pose':
+            for x in self.model.diffusion_model.named_parameters():
+                if not ('pose' in x[0]):
+                    x[1].requires_grad = False
+                else:
+                    x[1].requires_grad = True
+        elif self.trainkeys == 'all':
+            for x in self.model.diffusion_model.named_parameters():
+                x[1].requires_grad = True
+
+        self.model = self.model.to(memory_format=torch.channels_last)
+
+    def register_activation_hooks(
+            self,
+    ) -> None:
+        self.activations_dict = collections.defaultdict(list)
+        handles = []
+        for name, module in self.model.diffusion_model.named_modules():
+            if len(name.split('.')) > 1 and name.split('.')[-2] == 'transformer_blocks':
+                if hasattr(module, 'pose_emb_layers'):
+                    handle = module.register_forward_hook(
+                        partial(save_activations, self.activations_dict, name)
+                    )
+                    handles.append(handle)
+        self.handles = handles
+
+    def clear_rendered_feat(self,):
+        for name, module in self.model.diffusion_model.named_modules():
+            if len(name.split('.')) > 1 and name.split('.')[-2] == 'transformer_blocks':
+                if hasattr(module, 'pose_emb_layers'):
+                    module.rendered_feat = None
+
+    def remove_activation_hooks(
+            self, handles
+    ) -> None:
+        for handle in handles:
+            handle.remove()
+
+    def init_from_ckpt(
+        self,
+        path: str,
+    ) -> None:
+        if path.endswith("ckpt"):
+            sd = torch.load(path, map_location="cpu")["state_dict"]
+        elif path.endswith("safetensors"):
+            sd = load_safetensors(path)
+        else:
+            raise NotImplementedError
+
+        missing, unexpected = self.load_state_dict(sd, strict=False)
+        print(
+            f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+        )
+        if len(missing) > 0:
+            print(f"Missing Keys: {missing}")
+        if len(unexpected) > 0:
+            print(f"Unexpected Keys: {unexpected}")
+
+    def _init_first_stage(self, config):
+        model = instantiate_from_config(config).eval()
+        model.train = disabled_train
+        for param in model.parameters():
+            param.requires_grad = False
+        self.first_stage_model = model
+
+    def get_input(self, batch):
+        return batch[self.input_key], batch[self.input_key + '_ref'] if self.input_key + '_ref' in batch else None, batch['pose'] if 'pose' in batch else None, batch['mask'] if "mask" in batch else None, batch['mask_ref'] if "mask_ref" in batch else None, batch['depth'] if "depth" in batch else None, batch['drop_im'] if "drop_im" in batch else 0.
+
+    @torch.no_grad()
+    def decode_first_stage(self, z):
+        z = 1.0 / self.scale_factor * z
+        with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+            out = self.first_stage_model.decode(z)
+        return out
+
+    @torch.no_grad()
+    def encode_first_stage(self, x):
+        with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+            z = self.first_stage_model.encode(x)
+        z = self.scale_factor * z
+        return z
+
+    def forward(self, x, x_rgb, xr, pose, mask, mask_ref, opacity, drop_im, batch):
+        loss, loss_fg, loss_bg, loss_rgb = self.loss_fn(self.model, self.denoiser, self.conditioner, x, x_rgb, xr, pose, mask, mask_ref, opacity, batch)
+        loss_mean = loss.mean()
+        loss_dict = {"loss": loss_mean.item()}
+        if self.rgb and self.global_step > 0:
+            loss_fg = (loss_fg.mean(1)*drop_im.reshape(-1)).sum()/(drop_im.sum() + 1e-12)
+            loss_bg = (loss_bg.mean(1)*drop_im.reshape(-1)).sum()/(drop_im.sum() + 1e-12)
+            loss_mean += self.loss_fg_lambda*loss_fg
+            loss_mean += self.loss_bg_lambda*loss_bg
+            loss_dict["loss_fg"] = loss_fg.item()
+            loss_dict["loss_bg"] = loss_bg.item()
+        if self.rgb_predict and loss_rgb.mean() > 0:
+            loss_rgb = (loss_rgb.mean(1)*drop_im.reshape(-1)).sum()/(drop_im.sum() + 1e-12)
+            loss_mean += self.loss_rgb_lambda*loss_rgb
+            loss_dict["loss_rgb"] = loss_rgb.item()
+        return loss_mean, loss_dict
+
+    def shared_step(self, batch: Dict) -> Any:
+        x, xr, pose, mask, mask_ref, opacity, drop_im = self.get_input(batch)
+        x_rgb = x.clone().detach()
+        x = self.encode_first_stage(x)
+        x = x.to(memory_format=torch.channels_last)
+        if xr is not None:
+            b, n = xr.shape[0], xr.shape[1]
+            xr = rearrange(self.encode_first_stage(rearrange(xr, "b n ... -> (b n) ...")), "(b n) ... -> b n ...", b=b, n=n)
+            xr = drop_im.reshape(b, 1, 1, 1, 1)*xr + (1-drop_im.reshape(b, 1, 1, 1, 1))*torch.zeros_like(xr)
+        batch["global_step"] = self.global_step
+        loss, loss_dict = self(x, x_rgb, xr, pose, mask, mask_ref, opacity, drop_im, batch)
+        return loss, loss_dict
+
+    def training_step(self, batch, batch_idx):
+        loss, loss_dict = self.shared_step(batch)
+
+        self.log_dict(
+            loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=False
+        )
+
+        self.log(
+            "global_step",
+            self.global_step,
+            prog_bar=True,
+            logger=True,
+            on_step=True,
+            on_epoch=False,
+        )
+
+        if self.scheduler_config is not None:
+            lr = self.optimizers().param_groups[0]["lr"]
+            self.log(
+                "lr_abs", lr, prog_bar=True, logger=True, on_step=True, on_epoch=False
+            )
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        # print("validation data", len(self.trainer.val_dataloaders))
+        loss, loss_dict = self.shared_step(batch)
+        return loss
+
+    def on_train_start(self, *args, **kwargs):
+        if self.sampler is None or self.loss_fn is None:
+            raise ValueError("Sampler and loss function need to be set for training.")
+
+    def on_train_batch_end(self, *args, **kwargs):
+        if self.use_ema:
+            self.model_ema(self.model)
+
+    def optimizer_zero_grad(self, epoch, batch_idx, optimizer):
+        optimizer.zero_grad(set_to_none=True)
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.model.parameters())
+            self.model_ema.copy_to(self.model)
+            if context is not None:
+                print(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.model.parameters())
+                if context is not None:
+                    print(f"{context}: Restored training weights")
+
+    def instantiate_optimizer_from_config(self, params, lr, cfg):
+        return get_obj_from_str(cfg["target"])(
+            params, lr=lr, **cfg.get("params", dict())
+        )
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        params = []
+        blocks = []
+        lowlrparams = []
+        if self.trainkeys == 'poseattn':
+            lowlrparams = []
+            for x in self.model.diffusion_model.named_parameters():
+                if ('pose' in x[0]):
+                    params += [x[1]]
+                    blocks.append(x[0].split('.pose')[0])
+                    print(x[0])
+            blocks = set(blocks)
+            for x in self.model.diffusion_model.named_parameters():
+                if 'transformer_blocks' in x[0]:
+                    for each in blocks:
+                        if each in x[0] and not ('pose' in x[0]) and ('attn1' in x[0] or 'attn2' in x[0]):
+                            lowlrparams += [x[1]]
+        elif self.trainkeys == 'pose':
+            for x in self.model.diffusion_model.named_parameters():
+                if ('pose' in x[0]):
+                    params += [x[1]]
+                    print(x[0])  
+        elif self.trainkeys == 'all':
+            lowlrparams = []
+            for x in self.model.diffusion_model.named_parameters():
+                if ('pose' in x[0]):
+                    params += [x[1]]
+                    print(x[0])
+                else:
+                    lowlrparams += [x[1]]
+
+        for i, embedder in enumerate(self.conditioner.embedders[:2]):
+            if embedder.is_trainable:
+                params = params + list(embedder.parameters())
+            if self.add_token:
+                if i == 0:
+                    for name, param in embedder.transformer.get_input_embeddings().named_parameters():
+                        param.requires_grad = True
+                        print(name, "conditional model param")
+                        params += [param]
+                else:
+                    for name, param in embedder.model.token_embedding.named_parameters():
+                        param.requires_grad = True
+                        print(name, "conditional model param")
+                        params += [param]
+
+        if len(lowlrparams) > 0:
+            print("different optimizer groups")
+            opt = self.instantiate_optimizer_from_config([{'params': params}, {'params': lowlrparams, 'lr': self.multiplier*lr}], lr, self.optimizer_config)
+        else:
+            opt = self.instantiate_optimizer_from_config(params, lr, self.optimizer_config)
+        if self.scheduler_config is not None:
+            scheduler = instantiate_from_config(self.scheduler_config)
+            print("Setting up LambdaLR scheduler...")
+            scheduler = [
+                {
+                    "scheduler": LambdaLR(opt, lr_lambda=scheduler.schedule),
+                    "interval": "step",
+                    "frequency": 1,
+                }
+            ]
+            return [opt], scheduler
+        return opt
+
+    @torch.no_grad()
+    def sample(
+        self,
+        cond: Dict,
+        uc: Union[Dict, None] = None,
+        batch_size: int = 16,
+        num_steps=None,
+        randn=None,
+        shape: Union[None, Tuple, List] = None,
+        return_rgb=False,
+        mask=None,
+        init_im=None,
+        **kwargs,
+    ):
+        if randn is None:
+            randn = torch.randn(batch_size, *shape)
+
+        denoiser = lambda input, sigma, c: self.denoiser(
+            self.model, input, sigma, c, **kwargs
+        )
+        if mask is not None:
+            samples, rgb_list = self.sampler(denoiser, randn.to(self.device), cond, uc=uc, mask=mask, init_im=init_im, num_steps=num_steps)
+        else:
+            samples, rgb_list = self.sampler(denoiser, randn.to(self.device), cond, uc=uc, num_steps=num_steps)
+        if return_rgb:
+            return samples, rgb_list
+        return samples
+
+    @torch.no_grad()
+    def samplemulti(
+        self,
+        cond,
+        uc=None,
+        batch_size: int = 16,
+        num_steps=None,
+        randn=None,
+        shape: Union[None, Tuple, List] = None,
+        return_rgb=False,
+        mask=None,
+        init_im=None,
+        multikwargs=None,
+    ):
+        if randn is None:
+            randn = torch.randn(batch_size, *shape)
+
+        samples, rgb_list = self.sampler(self.denoiser, self.model, randn.to(self.device), cond, uc=uc, num_steps=num_steps, multikwargs=multikwargs)
+        if return_rgb:
+            return samples, rgb_list
+        return samples
+
+    @torch.no_grad()
+    def log_conditionings(self, batch: Dict, n: int, refernce: bool = True) -> Dict:
+        """
+        Defines heuristics to log different conditionings.
+        These can be lists of strings (text-to-image), tensors, ints, ...
+        """
+        image_h, image_w = batch[self.input_key].shape[2:]
+        log = dict()
+
+        for embedder in self.conditioner.embedders:
+            if refernce:
+                check = (embedder.input_keys[0] in self.log_keys)
+            else:
+                check = (embedder.input_key in self.log_keys)
+            if (
+                (self.log_keys is None) or check
+            ) and not self.no_cond_log:
+                if refernce:
+                    x = batch[embedder.input_keys[0]][:n]
+                else:
+                    x = batch[embedder.input_key][:n]
+                if isinstance(x, torch.Tensor):
+                    if x.dim() == 1:
+                        # class-conditional, convert integer to string
+                        x = [str(x[i].item()) for i in range(x.shape[0])]
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 4)
+                    elif x.dim() == 2:
+                        # size and crop cond and the like
+                        x = [
+                            "x".join([str(xx) for xx in x[i].tolist()])
+                            for i in range(x.shape[0])
+                        ]
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+                    else:
+                        raise NotImplementedError()
+                elif isinstance(x, (List, ListConfig)):
+                    if isinstance(x[0], str):
+                        # strings
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+                    else:
+                        raise NotImplementedError()
+                else:
+                    raise NotImplementedError()
+                if refernce:
+                    log[embedder.input_keys[0]] = xc
+                else:
+                    log[embedder.input_key] = xc
+        return log
+
+    @torch.no_grad()
+    def log_images(
+        self,
+        batch: Dict,
+        N: int = 8,
+        sample: bool = True,
+        ucg_keys: List[str] = None,
+        **kwargs,
+    ) -> Dict:
+        log = dict()
+
+        x, xr, pose, mask, mask_ref, depth, drop_im = self.get_input(batch)
+
+        if xr is not None:
+            conditioner_input_keys = [e.input_keys for e in self.conditioner.embedders]
+        else:
+            conditioner_input_keys = [e.input_key for e in self.conditioner.embedders]
+
+        if ucg_keys:
+            assert all(map(lambda x: x in conditioner_input_keys, ucg_keys)), (
+                "Each defined ucg key for sampling must be in the provided conditioner input keys,"
+                f"but we have {ucg_keys} vs. {conditioner_input_keys}"
+            )
+        else:
+            ucg_keys = conditioner_input_keys
+
+        c, uc = self.conditioner.get_unconditional_conditioning(
+            batch,
+            force_uc_zero_embeddings=ucg_keys
+            if len(self.conditioner.embedders) > 0
+            else [],
+        )
+
+        N = min(x.shape[0], N)
+        x = x.to(self.device)[:N]
+        zr = None
+        if xr is not None:
+            xr = xr.to(self.device)[:N]
+            b, n = xr.shape[0], xr.shape[1]
+            log["reference"] = rearrange(xr, "b n ... -> (b n) ...", b=b, n=n)
+            zr = rearrange(self.encode_first_stage(rearrange(xr, "b n ... -> (b n) ...", b=b, n=n)), "(b n) ... -> b n ...", b=b, n=n)
+
+        log["inputs"] = x
+        b = x.shape[0]
+        if mask is not None:
+            log["mask"] = mask
+        if depth is not None:
+            log["depth"] = depth
+        z = self.encode_first_stage(x)
+
+        if uc is not None:
+            if xr is not None:
+                zr = torch.cat([torch.zeros_like(zr), zr])
+            drop_im = torch.cat([drop_im, drop_im])
+            if isinstance(pose, list):
+                pose = pose[:N]*2
+            else:
+                pose = torch.cat([pose[:N]] * 2)
+
+        sampling_kwargs = {'input_ref':zr}
+        sampling_kwargs['pose'] = pose
+        sampling_kwargs['mask_ref'] = None
+        sampling_kwargs['drop_im'] = drop_im
+
+        log["reconstructions"] = self.decode_first_stage(z)
+        log.update(self.log_conditionings(batch, N, refernce=True if xr is not None else False))
+
+        for k in c:
+            if isinstance(c[k], torch.Tensor):
+                if xr is not None:
+                    c[k], uc[k] = map(lambda y: y[k][:(n+1)*N].to(self.device), (c, uc))
+                else:
+                    c[k], uc[k] = map(lambda y: y[k][:N].to(self.device), (c, uc))
+        if sample:
+            with self.ema_scope("Plotting"):
+                samples, rgb_list = self.sample(
+                    c, shape=z.shape[1:], uc=uc, batch_size=N, return_rgb=True, **sampling_kwargs
+                )
+            samples = self.decode_first_stage(samples)
+            log["samples"] = samples
+            if len(rgb_list) > 0:
+                size = int(math.sqrt(rgb_list[0].size(1)))
+                log["predicted_rgb"] = rgb_list[0].reshape(-1, size, size, 3).permute(0, 3, 1, 2)
+        return log

sgm/modules/__init__.py

@@ -0,0 +1,6 @@
+from .encoders.modules import GeneralConditioner
+
+UNCONDITIONAL_CONFIG = {
+    "target": "sgm.modules.GeneralConditioner",
+    "params": {"emb_models": []},
+}

sgm/modules/attention.py

@@ -0,0 +1,1202 @@
+import logging
+import math
+import itertools
+from inspect import isfunction
+from typing import Any, Optional
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from packaging import version
+from torch import nn
+from .diffusionmodules.util import checkpoint
+from torch.autograd import Function
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+from ..modules.diffusionmodules.util import zero_module
+from ..modules.nerfsd_pytorch3d import NerfSDModule, VolRender
+
+logpy = logging.getLogger(__name__)
+
+if version.parse(torch.__version__) >= version.parse("2.0.0"):
+    SDP_IS_AVAILABLE = True
+    from torch.backends.cuda import SDPBackend, sdp_kernel
+
+    BACKEND_MAP = {
+        SDPBackend.MATH: {
+            "enable_math": True,
+            "enable_flash": False,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.FLASH_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": True,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.EFFICIENT_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": False,
+            "enable_mem_efficient": True,
+        },
+        None: {"enable_math": True, "enable_flash": True, "enable_mem_efficient": True},
+    }
+else:
+    from contextlib import nullcontext
+
+    SDP_IS_AVAILABLE = False
+    sdp_kernel = nullcontext
+    BACKEND_MAP = {}
+    logpy.warn(
+        f"No SDP backend available, likely because you are running in pytorch "
+        f"versions < 2.0. In fact, you are using PyTorch {torch.__version__}. "
+        f"You might want to consider upgrading."
+    )
+
+try:
+    import xformers
+    import xformers.ops
+
+    XFORMERS_IS_AVAILABLE = True
+except:
+    XFORMERS_IS_AVAILABLE = False
+    logpy.warn("no module 'xformers'. Processing without...")
+
+
+def exists(val):
+    return val is not None
+
+
+def uniq(arr):
+    return {el: True for el in arr}.keys()
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def init_(tensor):
+    dim = tensor.shape[-1]
+    std = 1 / math.sqrt(dim)
+    tensor.uniform_(-std, std)
+    return tensor
+
+
+# feedforward
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
+        )
+
+        self.net = nn.Sequential(
+            project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(
+        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(
+            qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3
+        )
+        k = k.softmax(dim=-1)
+        context = torch.einsum("bhdn,bhen->bhde", k, v)
+        out = torch.einsum("bhde,bhdn->bhen", context, q)
+        out = rearrange(
+            out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w
+        )
+        return self.to_out(out)
+
+
+class SpatialSelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.k = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.v = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.proj_out = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b (h w) c")
+        k = rearrange(k, "b c h w -> b c (h w)")
+        w_ = torch.einsum("bij,bjk->bik", q, k)
+
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = rearrange(v, "b c h w -> b c (h w)")
+        w_ = rearrange(w_, "b i j -> b j i")
+        h_ = torch.einsum("bij,bjk->bik", v, w_)
+        h_ = rearrange(h_, "b c (h w) -> b c h w", h=h)
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class _TruncExp(Function):  # pylint: disable=abstract-method
+    # Implementation from torch-ngp:
+    # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, x):  # pylint: disable=arguments-differ
+        ctx.save_for_backward(x)
+        return torch.exp(x)
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, g):  # pylint: disable=arguments-differ
+        x = ctx.saved_tensors[0]
+        return g * torch.exp(x.clamp(-15, 15))
+
+
+trunc_exp = _TruncExp.apply
+"""Same as torch.exp, but with the backward pass clipped to prevent vanishing/exploding
+gradients."""
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        query_dim,
+        context_dim=None,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        backend=None,
+    ):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.scale = dim_head**-0.5
+        self.heads = heads
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+        )
+        self.backend = backend
+
+    def forward(
+        self,
+        x,
+        context=None,
+        mask=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+    ):
+        h = self.heads
+
+        if additional_tokens is not None:
+            # get the number of masked tokens at the beginning of the output sequence
+            n_tokens_to_mask = additional_tokens.shape[1]
+            # add additional token
+            x = torch.cat([additional_tokens, x], dim=1)
+
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        if n_times_crossframe_attn_in_self:
+            # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+            assert x.shape[0] % n_times_crossframe_attn_in_self == 0
+            n_cp = x.shape[0] // n_times_crossframe_attn_in_self
+            k = repeat(
+                k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp
+            )
+            v = repeat(
+                v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp
+            )
+
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
+
+        ## old
+        """
+        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+        del q, k
+
+        if exists(mask):
+            mask = rearrange(mask, 'b ... -> b (...)')
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, 'b j -> (b h) () j', h=h)
+            sim.masked_fill_(~mask, max_neg_value)
+
+        # attention, what we cannot get enough of
+        sim = sim.softmax(dim=-1)
+
+        out = einsum('b i j, b j d -> b i d', sim, v)
+        """
+        ## new
+        with sdp_kernel(**BACKEND_MAP[self.backend]):
+            # print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape)
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask
+            )  # scale is dim_head ** -0.5 per default
+
+        del q, k, v
+        out = rearrange(out, "b h n d -> b n (h d)", h=h)
+
+        if additional_tokens is not None:
+            # remove additional token
+            out = out[:, n_tokens_to_mask:]
+        return self.to_out(out)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+    # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
+    def __init__(
+        self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, add_lora=False, **kwargs
+    ):
+        super().__init__()
+        logpy.debug(
+            f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, "
+            f"context_dim is {context_dim} and using {heads} heads with a "
+            f"dimension of {dim_head}."
+        )
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.heads = heads
+        self.dim_head = dim_head
+        self.add_lora = add_lora
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+        )
+        if add_lora:
+            r = 32
+            self.to_q_attn3_down = nn.Linear(query_dim, r, bias=False)
+            self.to_q_attn3_up = zero_module(nn.Linear(r, inner_dim, bias=False))
+            self.to_k_attn3_down = nn.Linear(context_dim, r, bias=False)
+            self.to_k_attn3_up = zero_module(nn.Linear(r, inner_dim, bias=False))
+            self.to_v_attn3_down = nn.Linear(context_dim, r, bias=False)
+            self.to_v_attn3_up = zero_module(nn.Linear(r, inner_dim, bias=False))
+            self.to_o_attn3_down = nn.Linear(inner_dim, r, bias=False)
+            self.to_o_attn3_up = zero_module(nn.Linear(r, query_dim, bias=False))
+            self.dropoutq = nn.Dropout(0.1)
+            self.dropoutk = nn.Dropout(0.1)
+            self.dropoutv = nn.Dropout(0.1)
+            self.dropouto = nn.Dropout(0.1)
+
+            nn.init.normal_(self.to_q_attn3_down.weight, std=1 / r)
+            nn.init.normal_(self.to_k_attn3_down.weight, std=1 / r)
+            nn.init.normal_(self.to_v_attn3_down.weight, std=1 / r)
+            nn.init.normal_(self.to_o_attn3_down.weight, std=1 / r)
+
+        self.attention_op: Optional[Any] = None
+
+    def forward(
+        self,
+        x,
+        context=None,
+        mask=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+    ):
+        if additional_tokens is not None:
+            # get the number of masked tokens at the beginning of the output sequence
+            n_tokens_to_mask = additional_tokens.shape[1]
+            # add additional token
+            x = torch.cat([additional_tokens, x], dim=1)
+
+        context_k = context  # b, n, c, h, w
+
+        q = self.to_q(x)
+        context = default(context, x)
+        context_k = default(context_k, x)
+        k = self.to_k(context_k)
+        v = self.to_v(context_k)
+        if self.add_lora:
+            q += self.dropoutq(self.to_q_attn3_up(self.to_q_attn3_down(x)))
+            k += self.dropoutk(self.to_k_attn3_up(self.to_k_attn3_down(context_k)))
+            v += self.dropoutv(self.to_v_attn3_up(self.to_v_attn3_down(context_k)))
+
+        if n_times_crossframe_attn_in_self:
+            # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+            assert x.shape[0] % n_times_crossframe_attn_in_self == 0
+            # n_cp = x.shape[0]//n_times_crossframe_attn_in_self
+            k = repeat(
+                k[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+            v = repeat(
+                v[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+
+        attn_bias = None
+
+        # actually compute the attention, what we cannot get enough of
+        out = xformers.ops.memory_efficient_attention(
+            q, k, v, attn_bias=attn_bias, op=self.attention_op
+        )
+
+        # TODO: Use this directly in the attention operation, as a bias
+        if exists(mask):
+            raise NotImplementedError
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        if additional_tokens is not None:
+            # remove additional token
+            out = out[:, n_tokens_to_mask:]
+        final = self.to_out(out)
+        if self.add_lora:
+            final += self.dropouto(self.to_o_attn3_up(self.to_o_attn3_down(out)))
+        return final
+
+
+class BasicTransformerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,  # vanilla attention
+        "softmax-xformers": MemoryEfficientCrossAttention,  # ampere
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        disable_self_attn=False,
+        attn_mode="softmax",
+        sdp_backend=None,
+        image_cross=False,
+        far=2,
+        num_samples=32,
+        add_lora=False,
+        rgb_predict=False,
+        mode='pixel-nerf',
+        average=False,
+        num_freqs=16,
+        use_prev_weights_imp_sample=False,
+        imp_sample_next_step=False,
+        stratified=False,
+        imp_sampling_percent=0.9,
+        near_plane=0.
+    ):
+
+        super().__init__()
+        assert attn_mode in self.ATTENTION_MODES
+        self.add_lora = add_lora
+        self.image_cross = image_cross
+        self.rgb_predict = rgb_predict
+        self.use_prev_weights_imp_sample = use_prev_weights_imp_sample
+        self.imp_sample_next_step = imp_sample_next_step
+        self.rendered_feat = None
+        if attn_mode != "softmax" and not XFORMERS_IS_AVAILABLE:
+            logpy.warn(
+                f"Attention mode '{attn_mode}' is not available. Falling "
+                f"back to native attention. This is not a problem in "
+                f"Pytorch >= 2.0. FYI, you are running with PyTorch "
+                f"version {torch.__version__}."
+            )
+            attn_mode = "softmax"
+        elif attn_mode == "softmax" and not SDP_IS_AVAILABLE:
+            logpy.warn(
+                "We do not support vanilla attention anymore, as it is too "
+                "expensive. Sorry."
+            )
+            if not XFORMERS_IS_AVAILABLE:
+                assert (
+                    False
+                ), "Please install xformers via e.g. 'pip install xformers==0.0.16'"
+            else:
+                logpy.info("Falling back to xformers efficient attention.")
+                attn_mode = "softmax-xformers"
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        if version.parse(torch.__version__) >= version.parse("2.0.0"):
+            assert sdp_backend is None or isinstance(sdp_backend, SDPBackend)
+        else:
+            assert sdp_backend is None
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            add_lora=self.add_lora,
+            context_dim=context_dim if self.disable_self_attn else None,
+            backend=sdp_backend,
+        )  # is a self-attention if not self.disable_self_attn
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = attn_cls(
+            query_dim=dim,
+            context_dim=context_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            add_lora=self.add_lora,
+            backend=sdp_backend,
+        )  # is self-attn if context is none
+        if image_cross:
+            self.pose_emb_layers = nn.Linear(2*dim, dim, bias=False)
+            nn.init.eye_(self.pose_emb_layers.weight)
+            self.pose_featurenerf = NerfSDModule(mode=mode,
+                                                 out_channels=dim,
+                                                 far_plane=far,
+                                                 num_samples=num_samples,
+                                                 rgb_predict=rgb_predict,
+                                                 average=average,
+                                                 num_freqs=num_freqs,
+                                                 stratified=stratified,
+                                                 imp_sampling_percent=imp_sampling_percent,
+                                                 near_plane=near_plane,
+                                                 )
+
+            self.renderer = VolRender()
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+        if self.checkpoint:
+            logpy.debug(f"{self.__class__.__name__} is using checkpointing")
+
+    def forward(
+        self, x, context=None, context_ref=None, pose=None, mask_ref=None, prev_weights=None, additional_tokens=None, n_times_crossframe_attn_in_self=0
+    ):
+        kwargs = {"x": x}
+
+        if context is not None:
+            kwargs.update({"context": context})
+
+        if context_ref is not None:
+            kwargs.update({"context_ref": context_ref})
+
+        if pose is not None:
+            kwargs.update({"pose": pose})
+
+        if mask_ref is not None:
+            kwargs.update({"mask_ref": mask_ref})
+
+        if prev_weights is not None:
+            kwargs.update({"prev_weights": prev_weights})
+
+        if additional_tokens is not None:
+            kwargs.update({"additional_tokens": additional_tokens})
+
+        if n_times_crossframe_attn_in_self:
+            kwargs.update(
+                {"n_times_crossframe_attn_in_self": n_times_crossframe_attn_in_self}
+            )
+
+        # return mixed_checkpoint(self._forward, kwargs, self.parameters(), self.checkpoint)
+        return checkpoint(
+            self._forward, (x, context, context_ref, pose, mask_ref, prev_weights), self.parameters(), self.checkpoint
+        )
+
+    def reference_attn(self, x, context_ref, context, pose, prev_weights, mask_ref):
+        feats, sigmas, dists, _, predicted_rgb, sigmas_uniform, dists_uniform = self.pose_featurenerf(pose,
+                                                                                                      context_ref,
+                                                                                                      mask_ref,
+                                                                                                      prev_weights=prev_weights if self.use_prev_weights_imp_sample else None,
+                                                                                                      imp_sample_next_step=self.imp_sample_next_step)
+
+        b, hw, d = feats.size()[:3]
+        feats = rearrange(feats, "b hw d ... -> b (hw d) ...")
+
+        feats = (
+            self.attn2(
+                self.norm2(feats), context=context,
+            )
+            + feats
+        )
+
+        feats = rearrange(feats, "b (hw d) ... -> b hw d ...", hw=hw, d=d)
+
+        sigmas_ = trunc_exp(sigmas)
+        if sigmas_uniform is not None:
+            sigmas_uniform = trunc_exp(sigmas_uniform)
+
+        context_ref, fg_mask, alphas, weights_uniform, predicted_rgb = self.renderer(feats, sigmas_, dists, densities_uniform=sigmas_uniform, dists_uniform=dists_uniform, return_weights_uniform=True, rgb=F.sigmoid(predicted_rgb) if predicted_rgb is not None else None)
+        if self.use_prev_weights_imp_sample:
+            prev_weights = weights_uniform
+
+        return context_ref, fg_mask, prev_weights, alphas, predicted_rgb
+
+    def _forward(
+        self, x, context=None, context_ref=None, pose=None, mask_ref=None, prev_weights=None, additional_tokens=None, n_times_crossframe_attn_in_self=0
+    ):
+        fg_mask = None
+        weights = None
+        alphas = None
+        predicted_rgb = None
+        xref = None
+
+        x = (
+            self.attn1(
+                self.norm1(x),
+                context=context if self.disable_self_attn else None,
+                additional_tokens=additional_tokens,
+                n_times_crossframe_attn_in_self=n_times_crossframe_attn_in_self
+                if not self.disable_self_attn
+                else 0,
+            )
+            + x
+        )
+        x = (
+            self.attn2(
+                self.norm2(x), context=context, additional_tokens=additional_tokens
+            )
+            + x
+        )
+        with torch.amp.autocast(device_type='cuda', dtype=torch.float32):
+            if context_ref is not None:
+                xref, fg_mask, weights, alphas, predicted_rgb = self.reference_attn(x,
+                                                                                    rearrange(context_ref, "(b n) ... -> b n ...", b=x.size(0), n=context_ref.size(0) // x.size(0)),
+                                                                                    context,
+                                                                                    pose,
+                                                                                    prev_weights,
+                                                                                    mask_ref)
+                x = self.pose_emb_layers(torch.cat([x, xref], -1))
+
+        x = self.ff(self.norm3(x)) + x
+        return x, fg_mask, weights, alphas, predicted_rgb
+
+
+class BasicTransformerSingleLayerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,  # vanilla attention
+        "softmax-xformers": MemoryEfficientCrossAttention  # on the A100s not quite as fast as the above version
+        # (todo might depend on head_dim, check, falls back to semi-optimized kernels for dim!=[16,32,64,128])
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        attn_mode="softmax",
+    ):
+        super().__init__()
+        assert attn_mode in self.ATTENTION_MODES
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim,
+        )
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None):
+        return checkpoint(
+            self._forward, (x, context), self.parameters(), self.checkpoint
+        )
+
+    def _forward(self, x, context=None):
+        x = self.attn1(self.norm1(x), context=context) + x
+        x = self.ff(self.norm2(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    """
+    Transformer block for image-like data.
+    First, project the input (aka embedding)
+    and reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    NEW: use_linear for more efficiency instead of the 1x1 convs
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        n_heads,
+        d_head,
+        depth=1,
+        dropout=0.0,
+        context_dim=None,
+        disable_self_attn=False,
+        use_linear=False,
+        attn_type="softmax",
+        use_checkpoint=True,
+        # sdp_backend=SDPBackend.FLASH_ATTENTION
+        sdp_backend=None,
+        image_cross=True,
+        rgb_predict=False,
+        far=2,
+        num_samples=32,
+        add_lora=False,
+        mode='feature-nerf',
+        average=False,
+        num_freqs=16,
+        use_prev_weights_imp_sample=False,
+        stratified=False,
+        poscontrol_interval=4,
+        imp_sampling_percent=0.9,
+        near_plane=0.
+    ):
+        super().__init__()
+        logpy.debug(
+            f"constructing {self.__class__.__name__} of depth {depth} w/ "
+            f"{in_channels} channels and {n_heads} heads."
+        )
+        from omegaconf import ListConfig
+
+        if exists(context_dim) and not isinstance(context_dim, (list, ListConfig)):
+            context_dim = [context_dim]
+        if exists(context_dim) and isinstance(context_dim, list):
+            if depth != len(context_dim):
+                logpy.warn(
+                    f"{self.__class__.__name__}: Found context dims "
+                    f"{context_dim} of depth {len(context_dim)}, which does not "
+                    f"match the specified 'depth' of {depth}. Setting context_dim "
+                    f"to {depth * [context_dim[0]]} now."
+                )
+                # depth does not match context dims.
+                assert all(
+                    map(lambda x: x == context_dim[0], context_dim)
+                ), "need homogenous context_dim to match depth automatically"
+                context_dim = depth * [context_dim[0]]
+        elif context_dim is None:
+            context_dim = [None] * depth
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+
+        self.image_cross = image_cross
+        self.poscontrol_interval = poscontrol_interval
+
+        if not use_linear:
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    dropout=dropout,
+                    context_dim=context_dim[d],
+                    disable_self_attn=disable_self_attn,
+                    attn_mode=attn_type,
+                    checkpoint=use_checkpoint,
+                    sdp_backend=sdp_backend,
+                    image_cross=self.image_cross and (d % poscontrol_interval == 0),
+                    far=far,
+                    num_samples=num_samples,
+                    add_lora=add_lora and self.image_cross and (d % poscontrol_interval == 0),
+                    rgb_predict=rgb_predict,
+                    mode=mode,
+                    average=average,
+                    num_freqs=num_freqs,
+                    use_prev_weights_imp_sample=use_prev_weights_imp_sample,
+                    imp_sample_next_step=(use_prev_weights_imp_sample and self.image_cross and (d % poscontrol_interval == 0) and depth >= poscontrol_interval and d < (depth // poscontrol_interval) * poscontrol_interval ),
+                    stratified=stratified,
+                    imp_sampling_percent=imp_sampling_percent,
+                    near_plane=near_plane,
+                )
+                for d in range(depth)
+            ]
+        )
+        if not use_linear:
+            self.proj_out = zero_module(
+                nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+            )
+        else:
+            # self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+            self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+        self.use_linear = use_linear
+
+    def forward(self, x, xr, context=None, contextr=None, pose=None, mask_ref=None, prev_weights=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        if xr is None:
+            if not isinstance(context, list):
+                context = [context]
+            b, c, h, w = x.shape
+            x_in = x
+            x = self.norm(x)
+            if not self.use_linear:
+                x = self.proj_in(x)
+            x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+            if self.use_linear:
+                x = self.proj_in(x)
+            for i, block in enumerate(self.transformer_blocks):
+                if i > 0 and len(context) == 1:
+                    i = 0  # use same context for each block
+                x, _, _, _, _ = block(x, context=context[i])
+            if self.use_linear:
+                x = self.proj_out(x)
+            x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+            if not self.use_linear:
+                x = self.proj_out(x)
+            return x + x_in, None, None, None, None, None
+        else:
+            if not isinstance(context, list):
+                context = [context]
+                contextr = [contextr]
+            b, c, h, w = x.shape
+            b1, _, _, _ = xr.shape
+            x_in = x
+            xr_in = xr
+            fg_masks = []
+            alphas = []
+            rgbs = []
+
+            x = self.norm(x)
+            with torch.no_grad():
+                xr = self.norm(xr)
+
+            if not self.use_linear:
+                x = self.proj_in(x)
+                with torch.no_grad():
+                    xr = self.proj_in(xr)
+
+            x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+            xr = rearrange(xr, "b1 c h w -> b1 (h w) c").contiguous()
+            if self.use_linear:
+                x = self.proj_in(x)
+                with torch.no_grad():
+                    xr = self.proj_in(xr)
+
+            prev_weights = None
+            counter = 0
+            for i, block in enumerate(self.transformer_blocks):
+                if i > 0 and len(context) == 1:
+                    i = 0  # use same context for each block
+                if self.image_cross and (counter % self.poscontrol_interval == 0):
+                    with torch.no_grad():
+                        xr, _, _, _, _ = block(xr, context=contextr[i])
+                    x, fg_mask, weights, alpha, rgb = block(x, context=context[i], context_ref=xr.detach(), pose=pose, mask_ref=mask_ref, prev_weights=prev_weights)
+                    prev_weights = weights
+                    fg_masks.append(fg_mask)
+                    if alpha is not None:
+                        alphas.append(alpha)
+                    if rgb is not None:
+                        rgbs.append(rgb)
+                else:
+                    with torch.no_grad():
+                        xr, _, _, _, _ = block(xr, context=contextr[i])
+                    x, _, _, _, _ = block(x, context=context[i])
+                counter += 1
+            if self.use_linear:
+                x = self.proj_out(x)
+                with torch.no_grad():
+                    xr = self.proj_out(xr)
+            x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+            xr = rearrange(xr, "b1 (h w) c -> b1 c h w", h=h, w=w).contiguous()
+            if not self.use_linear:
+                x = self.proj_out(x)
+                with torch.no_grad():
+                    xr = self.proj_out(xr)
+            if len(fg_masks) > 0:
+                if len(rgbs) <= 0:
+                    rgbs = None
+                if len(alphas) <= 0:
+                    alphas = None
+                return x + x_in, (xr + xr_in).detach(), fg_masks, prev_weights, alphas, rgbs
+            else:
+                return x + x_in, (xr + xr_in).detach(), None, prev_weights, None, None
+
+
+def benchmark_attn():
+    # Lets define a helpful benchmarking function:
+    # https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    import torch.nn.functional as F
+    import torch.utils.benchmark as benchmark
+
+    def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
+        t0 = benchmark.Timer(
+            stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
+        )
+        return t0.blocked_autorange().mean * 1e6
+
+    # Lets define the hyper-parameters of our input
+    batch_size = 32
+    max_sequence_len = 1024
+    num_heads = 32
+    embed_dimension = 32
+
+    dtype = torch.float16
+
+    query = torch.rand(
+        batch_size,
+        num_heads,
+        max_sequence_len,
+        embed_dimension,
+        device=device,
+        dtype=dtype,
+    )
+    key = torch.rand(
+        batch_size,
+        num_heads,
+        max_sequence_len,
+        embed_dimension,
+        device=device,
+        dtype=dtype,
+    )
+    value = torch.rand(
+        batch_size,
+        num_heads,
+        max_sequence_len,
+        embed_dimension,
+        device=device,
+        dtype=dtype,
+    )
+
+    print(f"q/k/v shape:", query.shape, key.shape, value.shape)
+
+    # Lets explore the speed of each of the 3 implementations
+    from torch.backends.cuda import SDPBackend, sdp_kernel
+
+    # Helpful arguments mapper
+    backend_map = {
+        SDPBackend.MATH: {
+            "enable_math": True,
+            "enable_flash": False,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.FLASH_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": True,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.EFFICIENT_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": False,
+            "enable_mem_efficient": True,
+        },
+    }
+
+    from torch.profiler import ProfilerActivity, profile, record_function
+
+    activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
+
+    print(
+        f"The default implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+    )
+    with profile(
+        activities=activities, record_shapes=False, profile_memory=True
+    ) as prof:
+        with record_function("Default detailed stats"):
+            for _ in range(25):
+                o = F.scaled_dot_product_attention(query, key, value)
+    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+    print(
+        f"The math implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+    )
+    with sdp_kernel(**backend_map[SDPBackend.MATH]):
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("Math implmentation stats"):
+                for _ in range(25):
+                    o = F.scaled_dot_product_attention(query, key, value)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+    with sdp_kernel(**backend_map[SDPBackend.FLASH_ATTENTION]):
+        try:
+            print(
+                f"The flash attention implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+            )
+        except RuntimeError:
+            print("FlashAttention is not supported. See warnings for reasons.")
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("FlashAttention stats"):
+                for _ in range(25):
+                    o = F.scaled_dot_product_attention(query, key, value)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+    with sdp_kernel(**backend_map[SDPBackend.EFFICIENT_ATTENTION]):
+        try:
+            print(
+                f"The memory efficient implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+            )
+        except RuntimeError:
+            print("EfficientAttention is not supported. See warnings for reasons.")
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("EfficientAttention stats"):
+                for _ in range(25):
+                    o = F.scaled_dot_product_attention(query, key, value)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+
+def run_model(model, x, context):
+    return model(x, context)
+
+
+def benchmark_transformer_blocks():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    import torch.utils.benchmark as benchmark
+
+    def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
+        t0 = benchmark.Timer(
+            stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
+        )
+        return t0.blocked_autorange().mean * 1e6
+
+    checkpoint = True
+    compile = False
+
+    batch_size = 32
+    h, w = 64, 64
+    context_len = 77
+    embed_dimension = 1024
+    context_dim = 1024
+    d_head = 64
+
+    transformer_depth = 4
+
+    n_heads = embed_dimension // d_head
+
+    dtype = torch.float16
+
+    model_native = SpatialTransformer(
+        embed_dimension,
+        n_heads,
+        d_head,
+        context_dim=context_dim,
+        use_linear=True,
+        use_checkpoint=checkpoint,
+        attn_type="softmax",
+        depth=transformer_depth,
+        sdp_backend=SDPBackend.FLASH_ATTENTION,
+    ).to(device)
+    model_efficient_attn = SpatialTransformer(
+        embed_dimension,
+        n_heads,
+        d_head,
+        context_dim=context_dim,
+        use_linear=True,
+        depth=transformer_depth,
+        use_checkpoint=checkpoint,
+        attn_type="softmax-xformers",
+    ).to(device)
+    if not checkpoint and compile:
+        print("compiling models")
+        model_native = torch.compile(model_native)
+        model_efficient_attn = torch.compile(model_efficient_attn)
+
+    x = torch.rand(batch_size, embed_dimension, h, w, device=device, dtype=dtype)
+    c = torch.rand(batch_size, context_len, context_dim, device=device, dtype=dtype)
+
+    from torch.profiler import ProfilerActivity, profile, record_function
+
+    activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
+
+    with torch.autocast("cuda"):
+        print(
+            f"The native model runs in {benchmark_torch_function_in_microseconds(model_native.forward, x, c):.3f} microseconds"
+        )
+        print(
+            f"The efficientattn model runs in {benchmark_torch_function_in_microseconds(model_efficient_attn.forward, x, c):.3f} microseconds"
+        )
+
+        print(75 * "+")
+        print("NATIVE")
+        print(75 * "+")
+        torch.cuda.reset_peak_memory_stats()
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("NativeAttention stats"):
+                for _ in range(25):
+                    model_native(x, c)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+        print(torch.cuda.max_memory_allocated() * 1e-9, "GB used by native block")
+
+        print(75 * "+")
+        print("Xformers")
+        print(75 * "+")
+        torch.cuda.reset_peak_memory_stats()
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("xformers stats"):
+                for _ in range(25):
+                    model_efficient_attn(x, c)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+        print(torch.cuda.max_memory_allocated() * 1e-9, "GB used by xformers block")
+
+
+def test01():
+    # conv1x1 vs linear
+    from ..util import count_params
+
+    conv = nn.Conv2d(3, 32, kernel_size=1).cuda()
+    print(count_params(conv))
+    linear = torch.nn.Linear(3, 32).cuda()
+    print(count_params(linear))
+
+    print(conv.weight.shape)
+
+    # use same initialization
+    linear.weight = torch.nn.Parameter(conv.weight.squeeze(-1).squeeze(-1))
+    linear.bias = torch.nn.Parameter(conv.bias)
+
+    print(linear.weight.shape)
+
+    x = torch.randn(11, 3, 64, 64).cuda()
+
+    xr = rearrange(x, "b c h w -> b (h w) c").contiguous()
+    print(xr.shape)
+    out_linear = linear(xr)
+    print(out_linear.mean(), out_linear.shape)
+
+    out_conv = conv(x)
+    print(out_conv.mean(), out_conv.shape)
+    print("done with test01.\n")
+
+
+def test02():
+    # try cosine flash attention
+    import time
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+    print("testing cosine flash attention...")
+    DIM = 1024
+    SEQLEN = 4096
+    BS = 16
+
+    print(" softmax (vanilla) first...")
+    model = BasicTransformerBlock(
+        dim=DIM,
+        n_heads=16,
+        d_head=64,
+        dropout=0.0,
+        context_dim=None,
+        attn_mode="softmax",
+    ).cuda()
+    try:
+        x = torch.randn(BS, SEQLEN, DIM).cuda()
+        tic = time.time()
+        y = model(x)
+        toc = time.time()
+        print(y.shape, toc - tic)
+    except RuntimeError as e:
+        # likely oom
+        print(str(e))
+
+    print("\n now flash-cosine...")
+    model = BasicTransformerBlock(
+        dim=DIM,
+        n_heads=16,
+        d_head=64,
+        dropout=0.0,
+        context_dim=None,
+        attn_mode="flash-cosine",
+    ).cuda()
+    x = torch.randn(BS, SEQLEN, DIM).cuda()
+    tic = time.time()
+    y = model(x)
+    toc = time.time()
+    print(y.shape, toc - tic)
+    print("done with test02.\n")
+
+
+if __name__ == "__main__":
+    # test01()
+    # test02()
+    # test03()
+
+    # benchmark_attn()
+    benchmark_transformer_blocks()
+
+    print("done.")

sgm/modules/autoencoding/lpips/loss/LICENSE

@@ -0,0 +1,23 @@
+Copyright (c) 2018, Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, Oliver Wang
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

sgm/modules/autoencoding/lpips/loss/lpips.py

@@ -0,0 +1,147 @@
+"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
+
+from collections import namedtuple
+
+import torch
+import torch.nn as nn
+from torchvision import models
+
+from ..util import get_ckpt_path
+
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout=True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self, name="vgg_lpips"):
+        ckpt = get_ckpt_path(name, "sgm/modules/autoencoding/lpips/loss")
+        self.load_state_dict(
+            torch.load(ckpt, map_location=torch.device("cpu")), strict=False
+        )
+        print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+    @classmethod
+    def from_pretrained(cls, name="vgg_lpips"):
+        if name != "vgg_lpips":
+            raise NotImplementedError
+        model = cls()
+        ckpt = get_ckpt_path(name)
+        model.load_state_dict(
+            torch.load(ckpt, map_location=torch.device("cpu")), strict=False
+        )
+        return model
+
+    def forward(self, input, target):
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(
+                outs1[kk]
+            )
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [
+            spatial_average(lins[kk].model(diffs[kk]), keepdim=True)
+            for kk in range(len(self.chns))
+        ]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer(
+            "shift", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None]
+        )
+        self.register_buffer(
+            "scale", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None]
+        )
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """A single linear layer which does a 1x1 conv"""
+
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+        layers = (
+            [
+                nn.Dropout(),
+            ]
+            if (use_dropout)
+            else []
+        )
+        layers += [
+            nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),
+        ]
+        self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple(
+            "VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3", "relu5_3"]
+        )
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x, eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
+    return x / (norm_factor + eps)
+
+
+def spatial_average(x, keepdim=True):
+    return x.mean([2, 3], keepdim=keepdim)

sgm/modules/autoencoding/lpips/model/LICENSE

@@ -0,0 +1,58 @@
+Copyright (c) 2017, Jun-Yan Zhu and Taesung Park
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+--------------------------- LICENSE FOR pix2pix --------------------------------
+BSD License
+
+For pix2pix software
+Copyright (c) 2016, Phillip Isola and Jun-Yan Zhu
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+----------------------------- LICENSE FOR DCGAN --------------------------------
+BSD License
+
+For dcgan.torch software
+
+Copyright (c) 2015, Facebook, Inc. All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+Neither the name Facebook nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

sgm/modules/autoencoding/lpips/model/model.py

@@ -0,0 +1,88 @@
+import functools
+
+import torch.nn as nn
+
+from ..util import ActNorm
+
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find("BatchNorm") != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator as in Pix2Pix
+    --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+    """
+
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if not use_actnorm:
+            norm_layer = nn.BatchNorm2d
+        else:
+            norm_layer = ActNorm
+        if (
+            type(norm_layer) == functools.partial
+        ):  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [
+            nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
+            nn.LeakyReLU(0.2, True),
+        ]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2**n, 8)
+            sequence += [
+                nn.Conv2d(
+                    ndf * nf_mult_prev,
+                    ndf * nf_mult,
+                    kernel_size=kw,
+                    stride=2,
+                    padding=padw,
+                    bias=use_bias,
+                ),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True),
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2**n_layers, 8)
+        sequence += [
+            nn.Conv2d(
+                ndf * nf_mult_prev,
+                ndf * nf_mult,
+                kernel_size=kw,
+                stride=1,
+                padding=padw,
+                bias=use_bias,
+            ),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True),
+        ]
+
+        sequence += [
+            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)
+        ]  # output 1 channel prediction map
+        self.main = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.main(input)

sgm/modules/autoencoding/lpips/util.py

@@ -0,0 +1,128 @@
+import hashlib
+import os
+
+import requests
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+URL_MAP = {"vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"}
+
+CKPT_MAP = {"vgg_lpips": "vgg.pth"}
+
+MD5_MAP = {"vgg_lpips": "d507d7349b931f0638a25a48a722f98a"}
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class ActNorm(nn.Module):
+    def __init__(
+        self, num_features, logdet=False, affine=True, allow_reverse_init=False
+    ):
+        assert affine
+        super().__init__()
+        self.logdet = logdet
+        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+        self.allow_reverse_init = allow_reverse_init
+
+        self.register_buffer("initialized", torch.tensor(0, dtype=torch.uint8))
+
+    def initialize(self, input):
+        with torch.no_grad():
+            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
+            mean = (
+                flatten.mean(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+            std = (
+                flatten.std(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+
+            self.loc.data.copy_(-mean)
+            self.scale.data.copy_(1 / (std + 1e-6))
+
+    def forward(self, input, reverse=False):
+        if reverse:
+            return self.reverse(input)
+        if len(input.shape) == 2:
+            input = input[:, :, None, None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        _, _, height, width = input.shape
+
+        if self.training and self.initialized.item() == 0:
+            self.initialize(input)
+            self.initialized.fill_(1)
+
+        h = self.scale * (input + self.loc)
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+
+        if self.logdet:
+            log_abs = torch.log(torch.abs(self.scale))
+            logdet = height * width * torch.sum(log_abs)
+            logdet = logdet * torch.ones(input.shape[0]).to(input)
+            return h, logdet
+
+        return h
+
+    def reverse(self, output):
+        if self.training and self.initialized.item() == 0:
+            if not self.allow_reverse_init:
+                raise RuntimeError(
+                    "Initializing ActNorm in reverse direction is "
+                    "disabled by default. Use allow_reverse_init=True to enable."
+                )
+            else:
+                self.initialize(output)
+                self.initialized.fill_(1)
+
+        if len(output.shape) == 2:
+            output = output[:, :, None, None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        h = output / self.scale - self.loc
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+        return h

sgm/modules/autoencoding/lpips/vqperceptual.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn.functional as F
+
+
+def hinge_d_loss(logits_real, logits_fake):
+    loss_real = torch.mean(F.relu(1.0 - logits_real))
+    loss_fake = torch.mean(F.relu(1.0 + logits_fake))
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+
+def vanilla_d_loss(logits_real, logits_fake):
+    d_loss = 0.5 * (
+        torch.mean(torch.nn.functional.softplus(-logits_real))
+        + torch.mean(torch.nn.functional.softplus(logits_fake))
+    )
+    return d_loss

sgm/modules/autoencoding/regularizers/__init__.py

@@ -0,0 +1,31 @@
+from abc import abstractmethod
+from typing import Any, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ....modules.distributions.distributions import \
+    DiagonalGaussianDistribution
+from .base import AbstractRegularizer
+
+
+class DiagonalGaussianRegularizer(AbstractRegularizer):
+    def __init__(self, sample: bool = True):
+        super().__init__()
+        self.sample = sample
+
+    def get_trainable_parameters(self) -> Any:
+        yield from ()
+
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
+        log = dict()
+        posterior = DiagonalGaussianDistribution(z)
+        if self.sample:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        kl_loss = posterior.kl()
+        kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+        log["kl_loss"] = kl_loss
+        return z, log