Migrated from GitHub

CONTRIBUTING.md

@@ -0,0 +1,34 @@
+Developer Certificate of Origin
+Version 1.1
+
+Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
+
+Everyone is permitted to copy and distribute verbatim copies of this
+license document, but changing it is not allowed.
+
+
+Developer's Certificate of Origin 1.1
+
+By making a contribution to this project, I certify that:
+
+(a) The contribution was created in whole or in part by me and I
+    have the right to submit it under the open source license
+    indicated in the file; or
+
+(b) The contribution is based upon previous work that, to the best
+    of my knowledge, is covered under an appropriate open source
+    license and I have the right under that license to submit that
+    work with modifications, whether created in whole or in part
+    by me, under the same open source license (unless I am
+    permitted to submit under a different license), as indicated
+    in the file; or
+
+(c) The contribution was provided directly to me by some other
+    person who certified (a), (b) or (c) and I have not modified
+    it.
+
+(d) I understand and agree that this project and the contribution
+    are public and that a record of the contribution (including all
+    personal information I submit with it, including my sign-off) is
+    maintained indefinitely and may be redistributed consistent with
+    this project or the open source license(s) involved.

LICENSE

@@ -0,0 +1,228 @@
+
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction,
+and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by
+the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all
+other entities that control, are controlled by, or are under common
+control with that entity. For the purposes of this definition,
+"control" means (i) the power, direct or indirect, to cause the
+direction or management of such entity, whether by contract or
+otherwise, or (ii) ownership of fifty percent (50%) or more of the
+outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity
+exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications,
+including but not limited to software source code, documentation
+source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical
+transformation or translation of a Source form, including but
+not limited to compiled object code, generated documentation,
+and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or
+Object form, made available under the License, as indicated by a
+copyright notice that is included in or attached to the work
+(an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object
+form, that is based on (or derived from) the Work and for which the
+editorial revisions, annotations, elaborations, or other modifications
+represent, as a whole, an original work of authorship. For the purposes
+of this License, Derivative Works shall not include works that remain
+separable from, or merely link (or bind by name) to the interfaces of,
+the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including
+the original version of the Work and any modifications or additions
+to that Work or Derivative Works thereof, that is intentionally
+submitted to Licensor for inclusion in the Work by the copyright owner
+or by an individual or Legal Entity authorized to submit on behalf of
+the copyright owner. For the purposes of this definition, "submitted"
+means any form of electronic, verbal, or written communication sent
+to the Licensor or its representatives, including but not limited to
+communication on electronic mailing lists, source code control systems,
+and issue tracking systems that are managed by, or on behalf of, the
+Licensor for the purpose of discussing and improving the Work, but
+excluding communication that is conspicuously marked or otherwise
+designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity
+on behalf of whom a Contribution has been received by Licensor and
+subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+copyright license to reproduce, prepare Derivative Works of,
+publicly display, publicly perform, sublicense, and distribute the
+Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+(except as stated in this section) patent license to make, have made,
+use, offer to sell, sell, import, and otherwise transfer the Work,
+where such license applies only to those patent claims licensable
+by such Contributor that are necessarily infringed by their
+Contribution(s) alone or by combination of their Contribution(s)
+with the Work to which such Contribution(s) was submitted. If You
+institute patent litigation against any entity (including a
+cross-claim or counterclaim in a lawsuit) alleging that the Work
+or a Contribution incorporated within the Work constitutes direct
+or contributory patent infringement, then any patent licenses
+granted to You under this License for that Work shall terminate
+as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the
+Work or Derivative Works thereof in any medium, with or without
+modifications, and in Source or Object form, provided that You
+meet the following conditions:
+
+(a) You must give any other recipients of the Work or
+Derivative Works a copy of this License; and
+
+(b) You must cause any modified files to carry prominent notices
+stating that You changed the files; and
+
+(c) You must retain, in the Source form of any Derivative Works
+that You distribute, all copyright, patent, trademark, and
+attribution notices from the Source form of the Work,
+excluding those notices that do not pertain to any part of
+the Derivative Works; and
+
+(d) If the Work includes a "NOTICE" text file as part of its
+distribution, then any Derivative Works that You distribute must
+include a readable copy of the attribution notices contained
+within such NOTICE file, excluding those notices that do not
+pertain to any part of the Derivative Works, in at least one
+of the following places: within a NOTICE text file distributed
+as part of the Derivative Works; within the Source form or
+documentation, if provided along with the Derivative Works; or,
+within a display generated by the Derivative Works, if and
+wherever such third-party notices normally appear. The contents
+of the NOTICE file are for informational purposes only and
+do not modify the License. You may add Your own attribution
+notices within Derivative Works that You distribute, alongside
+or as an addendum to the NOTICE text from the Work, provided
+that such additional attribution notices cannot be construed
+as modifying the License.
+
+You may add Your own copyright statement to Your modifications and
+may provide additional or different license terms and conditions
+for use, reproduction, or distribution of Your modifications, or
+for any such Derivative Works as a whole, provided Your use,
+reproduction, and distribution of the Work otherwise complies with
+the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise,
+any Contribution intentionally submitted for inclusion in the Work
+by You to the Licensor shall be under the terms and conditions of
+this License, without any additional terms or conditions.
+Notwithstanding the above, nothing herein shall supersede or modify
+the terms of any separate license agreement you may have executed
+with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade
+names, trademarks, service marks, or product names of the Licensor,
+except as required for reasonable and customary use in describing the
+origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or
+agreed to in writing, Licensor provides the Work (and each
+Contributor provides its Contributions) on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+implied, including, without limitation, any warranties or conditions
+of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+PARTICULAR PURPOSE. You are solely responsible for determining the
+appropriateness of using or redistributing the Work and assume any
+risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory,
+whether in tort (including negligence), contract, or otherwise,
+unless required by applicable law (such as deliberate and grossly
+negligent acts) or agreed to in writing, shall any Contributor be
+liable to You for damages, including any direct, indirect, special,
+incidental, or consequential damages of any character arising as a
+result of this License or out of the use or inability to use the
+Work (including but not limited to damages for loss of goodwill,
+work stoppage, computer failure or malfunction, or any and all
+other commercial damages or losses), even if such Contributor
+has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing
+the Work or Derivative Works thereof, You may choose to offer,
+and charge a fee for, acceptance of support, warranty, indemnity,
+or other liability obligations and/or rights consistent with this
+License. However, in accepting such obligations, You may act only
+on Your own behalf and on Your sole responsibility, not on behalf
+of any other Contributor, and only if You agree to indemnify,
+defend, and hold each Contributor harmless for any liability
+incurred by, or claims asserted against, such Contributor by reason
+of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+Copyright 2024 NVIDIA Corporation
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+
+PORTIONS LICENSED AS FOLLOWS
+
+> core/utils.py
+> mvdream/mv_unet.py
+> mvdream/pipeline_mvdream.py
+
+MIT License
+
+Copyright (c) 2024 3D Topia
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
+> core/attention.py
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+This source code is licensed under the Apache License, Version 2.0
+found in the LICENSE file in the root directory of this source tree.

acc_configs/gpu1.yaml

@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: 'NO'
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

acc_configs/gpu4.yaml

@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

acc_configs/gpu6.yaml

@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 6
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

acc_configs/gpu8.yaml

@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

blender_scripts/render_objaverse.py

@@ -0,0 +1,537 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse, sys, os, math, re
+import bpy
+from mathutils import Vector, Matrix
+import numpy as np
+import cv2
+import signal
+from contextlib import contextmanager
+from loguru import logger
+from typing import Any, Callable, Dict, Generator, List, Literal, Optional, Set, Tuple
+import random
+class TimeoutException(Exception): pass
+
+logger.info('Rendering started.')
+
+@contextmanager
+def time_limit(seconds):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timed out!")
+    signal.signal(signal.SIGALRM, signal_handler)
+    signal.alarm(seconds)
+    try:
+        yield
+    finally:
+        signal.alarm(0)
+
+parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+parser.add_argument(
+    '--seed', type=int, default=0,
+    help='number of views to be rendered')
+parser.add_argument(
+    '--views', type=int, default=4,
+    help='number of views to be rendered')
+parser.add_argument(
+    'obj', type=str,
+    help='Path to the obj file to be rendered.')
+parser.add_argument(
+    '--output_folder', type=str, default='/tmp',
+    help='The path the output will be dumped to.')
+parser.add_argument(
+    '--scale', type=float, default=1,
+    help='Scaling factor applied to model. Depends on size of mesh.')
+parser.add_argument(
+    '--format', type=str, default='PNG',
+    help='Format of files generated. Either PNG or OPEN_EXR')
+
+parser.add_argument(
+    '--resolution', type=int, default=512,
+    help='Resolution of the images.')
+parser.add_argument(
+    '--engine', type=str, default='CYCLES',
+    help='Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...')
+parser.add_argument(
+    '--gpu', type=int, default=0,
+    help='gpu.')
+parser.add_argument(
+    '--animation_idx', type=int, default=0,
+    help='The index of animation')
+
+parser.add_argument(
+    '--camera_option', type=str, default='fixed',
+    help='Camera Options')
+parser.add_argument(
+    '--fixed_animation_length', type=int, default=-1,
+    help='Set animation length to fixed number of framnes')
+parser.add_argument(
+    '--step_angle', type=int, default=3,
+    help='Angle in degree for each step camera rotation')
+parser.add_argument(
+    '--downsample', type=int, default=1,
+    help='Downsample ratio. No downsample by default')
+
+argv = sys.argv[sys.argv.index("--") + 1:]
+args = parser.parse_args(argv)
+
+
+model_identifier = os.path.split(args.obj)[1].split('.')[0]
+synset_idx = args.obj.split('/')[-2]
+
+save_root = os.path.join(os.path.abspath(args.output_folder),  synset_idx, model_identifier, f'{args.animation_idx:03d}')
+
+# Set up rendering
+context = bpy.context
+scene = bpy.context.scene
+render = bpy.context.scene.render
+
+render.engine = args.engine# 'BLENDER_EEVEE'
+render.image_settings.color_mode = 'RGBA'  # ('RGB', 'RGBA', ...)
+render.image_settings.file_format = args.format  # ('PNG', 'OPEN_EXR', 'JPEG, ...)
+render.resolution_x = args.resolution
+render.resolution_y = args.resolution
+render.resolution_percentage = 100
+bpy.context.scene.cycles.filter_width = 0.01
+bpy.context.scene.render.film_transparent = True
+render_depth_normal = False
+bpy.context.scene.cycles.device = 'GPU'
+bpy.context.scene.cycles.diffuse_bounces = 1
+bpy.context.scene.cycles.glossy_bounces = 1
+bpy.context.scene.cycles.transparent_max_bounces = 1
+bpy.context.scene.cycles.transmission_bounces = 1
+bpy.context.scene.cycles.samples = 16
+bpy.context.scene.cycles.use_denoising = True
+bpy.context.scene.cycles.denoiser = 'OPTIX'
+bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
+bpy.context.scene.cycles.device = 'GPU'
+
+
+def enable_cuda_devices():
+    prefs = bpy.context.preferences
+    cprefs = prefs.addons['cycles'].preferences
+    cprefs.get_devices()
+    # Attempt to set GPU device types if available
+    for compute_device_type in ('CUDA', 'OPENCL', 'NONE'):
+        try:
+            cprefs.compute_device_type = compute_device_type
+            print("Compute device selected: {0}".format(compute_device_type))
+            break
+        except TypeError:
+            pass
+
+    # Any CUDA/OPENCL devices?
+    acceleratedTypes = ['CUDA', 'OPENCL', 'OPTIX']
+    acceleratedTypes = ['CUDA', 'OPENCL']
+    accelerated = any(device.type in acceleratedTypes for device in cprefs.devices)
+    print('Accelerated render = {0}'.format(accelerated))
+
+    # If we have CUDA/OPENCL devices, enable only them, otherwise enable
+    # all devices (assumed to be CPU)
+    print(cprefs.devices)
+    for idx, device in enumerate(cprefs.devices):
+        device.use = (not accelerated or device.type in acceleratedTypes)# and idx == args.gpu
+        print('Device enabled ({type}) = {enabled}'.format(type=device.type, enabled=device.use))
+    return accelerated
+
+enable_cuda_devices()
+context.active_object.select_set(True)
+bpy.ops.object.delete()
+
+# Import textured mesh
+bpy.ops.object.select_all(action='DESELECT')
+
+try:
+    with time_limit(1000):
+        imported_object = bpy.ops.import_scene.gltf(filepath=args.obj, merge_vertices=True, guess_original_bind_pose=False, bone_heuristic="TEMPERANCE")
+except TimeoutException as e:
+    print("Timed out finished!")
+    exit()
+
+
+# count animated frames
+animation_names = []
+ending_frame_list = {}
+for k in bpy.data.actions.keys():
+    matched_obj_name = ''
+    for obj in bpy.context.selected_objects:
+        if '_'+obj.name in k and len(obj.name) > len(matched_obj_name):
+            matched_obj_name = obj.name
+    a_name = k.replace('_'+matched_obj_name, '')
+    a = bpy.data.actions[k]
+    frame_start, frame_end = map(int, a.frame_range)
+    logger.info(f'{k} | frame start: {frame_start}, frame end: {frame_end} | fps: {bpy.context.scene.render.fps}')
+    if a_name not in animation_names:
+        animation_names.append(a_name)
+        ending_frame_list[a_name] = frame_end
+    else:
+        ending_frame_list[a_name] = max(frame_end, ending_frame_list[a_name])
+        
+
+        
+selected_a_name = animation_names[args.animation_idx]
+max_frame = ending_frame_list[selected_a_name]
+for obj in bpy.context.selected_objects:
+    if obj.animation_data is not None:
+        obj_a_name = selected_a_name+'_'+obj.name
+        if obj_a_name in bpy.data.actions:
+            print('Found ', obj_a_name)
+            obj.animation_data.action = bpy.data.actions[obj_a_name]
+        else:
+            print('Miss ', obj_a_name)
+            
+num_frames = args.fixed_animation_length if args.fixed_animation_length != -1 else max_frame
+num_frames = num_frames // args.downsample
+
+if num_frames == 0:
+    print("No animation!")
+    exit()
+
+# from https://github.com/allenai/objaverse-xl/blob/main/scripts/rendering/blender_script.py
+def get_3x4_RT_matrix_from_blender(cam: bpy.types.Object):
+    """Returns the 3x4 RT matrix from the given camera.
+
+    Taken from Zero123, which in turn was taken from
+    https://github.com/panmari/stanford-shapenet-renderer/blob/master/render_blender.py
+
+    Args:
+        cam (bpy.types.Object): The camera object.
+
+    Returns:
+        Matrix: The 3x4 RT matrix from the given camera.
+    """
+    # Use matrix_world instead to account for all constraints
+    location, rotation = cam.matrix_world.decompose()[0:2]
+    R_world2bcam = rotation.to_matrix().transposed()
+
+    # Use location from matrix_world to account for constraints:
+    T_world2bcam = -1 * R_world2bcam @ location
+
+    # put into 3x4 matrix
+    RT = Matrix(
+        (
+            R_world2bcam[0][:] + (T_world2bcam[0],),
+            R_world2bcam[1][:] + (T_world2bcam[1],),
+            R_world2bcam[2][:] + (T_world2bcam[2],),
+        )
+    )
+    return RT
+def _create_light(
+    name: str,
+    light_type: Literal["POINT", "SUN", "SPOT", "AREA"],
+    location: Tuple[float, float, float],
+    rotation: Tuple[float, float, float],
+    energy: float,
+    use_shadow: bool = False,
+    specular_factor: float = 1.0,
+):
+    """Creates a light object.
+
+    Args:
+        name (str): Name of the light object.
+        light_type (Literal["POINT", "SUN", "SPOT", "AREA"]): Type of the light.
+        location (Tuple[float, float, float]): Location of the light.
+        rotation (Tuple[float, float, float]): Rotation of the light.
+        energy (float): Energy of the light.
+        use_shadow (bool, optional): Whether to use shadows. Defaults to False.
+        specular_factor (float, optional): Specular factor of the light. Defaults to 1.0.
+
+    Returns:
+        bpy.types.Object: The light object.
+    """
+
+    light_data = bpy.data.lights.new(name=name, type=light_type)
+    light_object = bpy.data.objects.new(name, light_data)
+    bpy.context.collection.objects.link(light_object)
+    light_object.location = location
+    light_object.rotation_euler = rotation
+    light_data.use_shadow = use_shadow
+    light_data.specular_factor = specular_factor
+    light_data.energy = energy
+    return light_object
+
+
+def randomize_lighting() -> Dict[str, bpy.types.Object]:
+    """Randomizes the lighting in the scene.
+
+    Returns:
+        Dict[str, bpy.types.Object]: Dictionary of the lights in the scene. The keys are
+            "key_light", "fill_light", "rim_light", and "bottom_light".
+    """
+
+    # Clear existing lights
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="LIGHT")
+    bpy.ops.object.delete()
+
+    # Create key light
+    key_light = _create_light(
+        name="Key_Light",
+        light_type="SUN",
+        location=(0, 0, 0),
+        rotation=(0.785398, 0, -0.785398),
+#         energy=random.choice([3, 4, 5]),
+        energy=4,
+    )
+
+    # Create fill light
+    fill_light = _create_light(
+        name="Fill_Light",
+        light_type="SUN",
+        location=(0, 0, 0),
+        rotation=(0.785398, 0, 2.35619),
+#         energy=random.choice([2, 3, 4]),
+        energy=3,
+    )
+
+    # Create rim light
+    rim_light = _create_light(
+        name="Rim_Light",
+        light_type="SUN",
+        location=(0, 0, 0),
+        rotation=(-0.785398, 0, -3.92699),
+#         energy=random.choice([3, 4, 5]),
+        energy=4,
+    )
+
+    # Create bottom light
+    bottom_light = _create_light(
+        name="Bottom_Light",
+        light_type="SUN",
+        location=(0, 0, 0),
+        rotation=(3.14159, 0, 0),
+#         energy=random.choice([1, 2, 3]),
+        energy=2,
+    )
+
+    return dict(
+        key_light=key_light,
+        fill_light=fill_light,
+        rim_light=rim_light,
+        bottom_light=bottom_light,
+    )
+
+def scene_bbox(
+    single_obj = None, ignore_matrix = False
+):
+    """Returns the bounding box of the scene.
+
+    Taken from Shap-E rendering script
+    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
+
+    Args:
+        single_obj (Optional[bpy.types.Object], optional): If not None, only computes
+            the bounding box for the given object. Defaults to None.
+        ignore_matrix (bool, optional): Whether to ignore the object's matrix. Defaults
+            to False.
+
+    Raises:
+        RuntimeError: If there are no objects in the scene.
+
+    Returns:
+        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
+    """
+    bbox_min = (math.inf,) * 3
+    bbox_max = (-math.inf,) * 3
+    found = False
+    for i in range(num_frames):
+        bpy.context.scene.frame_set(i * args.downsample)
+        for obj in get_scene_meshes() if single_obj is None else [single_obj]:
+            found = True
+            for coord in obj.bound_box:
+                coord = Vector(coord)
+                if not ignore_matrix:
+                    coord = obj.matrix_world @ coord
+                bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
+                bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
+                
+    if not found:
+        raise RuntimeError("no objects in scene to compute bounding box for")
+
+    return Vector(bbox_min), Vector(bbox_max)
+
+def get_scene_meshes():
+    """Returns all meshes in the scene.
+
+    Yields:
+        Generator[bpy.types.Object, None, None]: Generator of all meshes in the scene.
+    """
+    for obj in bpy.context.scene.objects.values():
+        if isinstance(obj.data, (bpy.types.Mesh)):
+            yield obj
+
+def get_scene_root_objects():
+    """Returns all root objects in the scene.
+
+    Yields:
+        Generator[bpy.types.Object, None, None]: Generator of all root objects in the
+            scene.
+    """
+    for obj in bpy.context.scene.objects.values():
+        if not obj.parent:
+            yield obj
+            
+def normalize_scene():
+    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
+    at the origin.
+
+    Mostly taken from the Point-E / Shap-E rendering script
+    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
+    but fix for multiple root objects: (see bug report here:
+    https://github.com/openai/shap-e/pull/60).
+
+    Returns:
+        None
+    """
+    if len(list(get_scene_root_objects())) > 1:
+        # create an empty object to be used as a parent for all root objects
+        parent_empty = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(parent_empty)
+
+        # parent all root objects to the empty object
+        for obj in get_scene_root_objects():
+            if obj != parent_empty:
+                obj.parent = parent_empty
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    logger.info(f"Scale: {scale}")
+    for obj in get_scene_root_objects():
+        obj.scale = obj.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    for obj in get_scene_root_objects():
+        obj.matrix_world.translation += offset
+    bpy.ops.object.select_all(action="DESELECT")
+
+    # unparent the camera
+    bpy.data.objects["Camera"].parent = None
+
+normalize_scene()
+
+randomize_lighting()
+
+# Place camera
+cam = scene.objects['Camera']
+cam.location = (0, 1.5, 0)  # radius equals to 1
+cam.data.lens = 35
+cam.data.sensor_width = 32
+
+cam_constraint = cam.constraints.new(type='TRACK_TO')
+cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'
+cam_constraint.up_axis = 'UP_Y'
+
+cam_empty = bpy.data.objects.new("Empty", None)
+cam_empty.location = (0, 0, 0)
+cam.parent = cam_empty
+
+scene.collection.objects.link(cam_empty)
+context.view_layer.objects.active = cam_empty
+cam_constraint.target = cam_empty
+
+stepsize = 360.0 / args.views
+rotation_mode = 'XYZ'
+
+
+np.random.seed(args.seed)
+
+if args.camera_option == "fixed":
+    for scene in bpy.data.scenes:
+        scene.cycles.device = 'GPU'
+    
+    elevation_angle = 0.
+    rotation_angle = 0.
+
+    for view_idx in range(args.views):
+        img_folder = os.path.join(save_root, f'{view_idx:03d}', 'img')
+        mask_folder = os.path.join(save_root, f'{view_idx:03d}', 'mask')
+        camera_folder = os.path.join(save_root, f'{view_idx:03d}', 'camera')
+
+        os.makedirs(img_folder, exist_ok=True)
+        os.makedirs(mask_folder, exist_ok=True)
+        os.makedirs(camera_folder, exist_ok=True)
+        
+        np.save(os.path.join(camera_folder, 'rotation'), np.array([rotation_angle + view_idx * stepsize for _ in range(num_frames)]))
+        np.save(os.path.join(camera_folder, 'elevation'), np.array([elevation_angle for _ in range(num_frames)]))
+
+        cam_empty.rotation_euler[2] = math.radians(rotation_angle + view_idx * stepsize)
+        cam_empty.rotation_euler[0] = math.radians(elevation_angle)
+        
+        # save camera RT matrix
+        rt_matrix = get_3x4_RT_matrix_from_blender(cam)
+        rt_matrix_path = os.path.join(camera_folder, "rt_matrix.npy")
+        np.save(rt_matrix_path, rt_matrix)
+        for i in range(0, num_frames):
+            bpy.context.scene.frame_set(i * args.downsample)
+            render_file_path = os.path.join(img_folder,'%03d.png' % (i))
+            scene.render.filepath = render_file_path
+            bpy.ops.render.render(write_still=True)
+
+        for i in range(0, num_frames):
+            img = cv2.imread(os.path.join(img_folder, '%03d.png' % (i)), cv2.IMREAD_UNCHANGED)
+            mask =  img[:, :, 3:4] / 255.0
+            white_img = img[:, :, :3] * mask + np.ones_like(img[:, :, :3]) * (1 - mask) * 255
+            white_img = np.clip(white_img, 0, 255)
+            cv2.imwrite(os.path.join(img_folder, '%03d.jpg' % (i)), white_img)
+            cv2.imwrite(os.path.join(mask_folder, '%03d.png'%(i)), img[:, :, 3])
+            os.system('rm %s'%(os.path.join(img_folder, '%03d.png' % (i))))
+        
+elif args.camera_option == "random":
+    for scene in bpy.data.scenes:
+        scene.cycles.device = 'GPU'
+    
+    for view_idx in range(args.views):
+        elevation_angle = np.random.rand(1) * 35 - 5 # [-5, 30]
+        rotation_angle = np.random.rand(1) * 360
+        
+        img_folder = os.path.join(save_root, f'{view_idx:03d}', 'img')
+        mask_folder = os.path.join(save_root, f'{view_idx:03d}', 'mask')
+        camera_folder = os.path.join(save_root, f'{view_idx:03d}', 'camera')
+
+        os.makedirs(img_folder, exist_ok=True)
+        os.makedirs(mask_folder, exist_ok=True)
+        os.makedirs(camera_folder, exist_ok=True)
+        
+        np.save(os.path.join(camera_folder, 'rotation'), np.array([rotation_angle for _ in range(num_frames)]))
+        np.save(os.path.join(camera_folder, 'elevation'), np.array([elevation_angle for _ in range(num_frames)]))
+        
+        cam_empty.rotation_euler[2] = math.radians(rotation_angle)
+        cam_empty.rotation_euler[0] = math.radians(elevation_angle)
+        
+        # save camera RT matrix
+        rt_matrix = get_3x4_RT_matrix_from_blender(cam)
+        rt_matrix_path = os.path.join(camera_folder, "rt_matrix.npy")
+        np.save(rt_matrix_path, rt_matrix)
+        
+        for i in range(0, num_frames):
+            bpy.context.scene.frame_set(i * args.downsample)
+            render_file_path = os.path.join(img_folder,'%03d.png' % (i))
+            scene.render.filepath = render_file_path
+            bpy.ops.render.render(write_still=True)
+
+        for i in range(0, num_frames):
+            img = cv2.imread(os.path.join(img_folder, '%03d.png' % (i)), cv2.IMREAD_UNCHANGED)
+            mask =  img[:, :, 3:4] / 255.0
+            white_img = img[:, :, :3] * mask + np.ones_like(img[:, :, :3]) * (1 - mask) * 255
+            white_img = np.clip(white_img, 0, 255)
+            cv2.imwrite(os.path.join(img_folder, '%03d.jpg' % (i)), white_img)
+            cv2.imwrite(os.path.join(mask_folder, '%03d.png'%(i)), img[:, :, 3])
+            os.system('rm %s'%(os.path.join(img_folder, '%03d.png' % (i))))
+    
+else:
+    raise NotImplemented

core/attention.py

@@ -0,0 +1,85 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import os
+import warnings
+
+from torch import Tensor
+from torch import nn
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import memory_efficient_attention, unbind
+
+        XFORMERS_AVAILABLE = True
+        # warnings.warn("xFormers is available (Attention)")
+    else:
+        warnings.warn("xFormers is disabled (Attention)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+    warnings.warn("xFormers is not available (Attention)")
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

core/gs.py

@@ -0,0 +1,198 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from core.options import Options
+
+import kiui
+
+from gsplat.rendering import rasterization
+
+class GaussianRenderer:
+    def __init__(self, opt: Options):
+        
+        self.opt = opt
+        self.bg_color = torch.tensor([1, 1, 1], dtype=torch.float32, device="cuda")
+        
+        # intrinsics
+        self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy))
+        self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32)
+        self.proj_matrix[0, 0] = 1 / self.tan_half_fov
+        self.proj_matrix[1, 1] = 1 / self.tan_half_fov
+        self.proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
+        self.proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
+        self.proj_matrix[2, 3] = 1
+
+        f = self.opt.output_size / (2 * self.tan_half_fov)
+        self.K = torch.tensor([[f, 0., self.opt.output_size/2.], [0., f, self.opt.output_size/2.], [0., 0., 1.]], dtype=torch.float32, device="cuda")
+
+    def render(self, gaussians, cam_view, cam_view_proj, cam_pos, bg_color=None):
+        # gaussians: [B, N, 14]
+        # cam_view, cam_view_proj: [B, V, 4, 4]
+        # cam_pos: [B, V, 3]
+
+        device = gaussians.device
+        B, V = cam_view.shape[:2]
+
+        # loop of loop...
+        images = []
+        alphas = []
+        for b in range(B):
+
+            # pos, opacity, scale, rotation, shs
+            means3D = gaussians[b, :, 0:3].contiguous().float()
+            opacity = gaussians[b, :, 3:4].contiguous().float()
+            scales = gaussians[b, :, 4:7].contiguous().float()
+            rotations = gaussians[b, :, 7:11].contiguous().float()
+            rgbs = gaussians[b, :, 11:].contiguous().float() # [N, 3]
+                
+            # render novel views
+            view_matrix = cam_view[b].float()
+            view_proj_matrix = cam_view_proj[b].float()
+            campos = cam_pos[b].float()
+
+            viewmat = view_matrix.transpose(2, 1) # [V, 4, 4]
+
+            
+            rendered_image_all, rendered_alpha_all, info = rasterization(
+                means=means3D,
+                quats=rotations,
+                scales=scales,
+                opacities=opacity.squeeze(-1),
+                colors=rgbs,
+                viewmats=viewmat,
+                Ks=torch.stack([self.K for _ in range(V)]),
+                width=self.opt.output_size,
+                height=self.opt.output_size,
+                near_plane=self.opt.znear,
+                far_plane=self.opt.zfar,
+                packed=False,
+                backgrounds=torch.stack([self.bg_color for _ in range(V)]) if self.bg_color is not None else None,
+                render_mode="RGB",
+            )
+            for rendered_image, rendered_alpha in zip(rendered_image_all, rendered_alpha_all):
+
+                rendered_image = rendered_image.permute(2, 0, 1)
+                rendered_image = rendered_image.clamp(0, 1)
+
+                rendered_alpha = rendered_alpha.permute(2, 0, 1)
+
+                images.append(rendered_image)
+                alphas.append(rendered_alpha)
+
+        images = torch.stack(images, dim=0).view(B, V, 3, self.opt.output_size, self.opt.output_size)
+        alphas = torch.stack(alphas, dim=0).view(B, V, 1, self.opt.output_size, self.opt.output_size)
+
+        return {
+            "image": images, # [B, V, 3, H, W]
+            "alpha": alphas, # [B, V, 1, H, W]
+        }
+
+
+    def save_ply(self, gaussians, path, compatible=True):
+        # gaussians: [B, N, 14]
+        # compatible: save pre-activated gaussians as in the original paper
+
+        assert gaussians.shape[0] == 1, 'only support batch size 1'
+
+        from plyfile import PlyData, PlyElement
+     
+        means3D = gaussians[0, :, 0:3].contiguous().float()
+        opacity = gaussians[0, :, 3:4].contiguous().float()
+        scales = gaussians[0, :, 4:7].contiguous().float()
+        rotations = gaussians[0, :, 7:11].contiguous().float()
+        shs = gaussians[0, :, 11:].unsqueeze(1).contiguous().float() # [N, 1, 3]
+
+        # prune by opacity
+        mask = opacity.squeeze(-1) >= 0.005
+        means3D = means3D[mask]
+        opacity = opacity[mask]
+        scales = scales[mask]
+        rotations = rotations[mask]
+        shs = shs[mask]
+
+        # invert activation to make it compatible with the original ply format
+        if compatible:
+            opacity = kiui.op.inverse_sigmoid(opacity)
+            scales = torch.log(scales + 1e-8)
+            shs = (shs - 0.5) / 0.28209479177387814
+
+        xyzs = means3D.detach().cpu().numpy()
+        f_dc = shs.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        opacities = opacity.detach().cpu().numpy()
+        scales = scales.detach().cpu().numpy()
+        rotations = rotations.detach().cpu().numpy()
+
+        l = ['x', 'y', 'z']
+        # All channels except the 3 DC
+        for i in range(f_dc.shape[1]):
+            l.append('f_dc_{}'.format(i))
+        l.append('opacity')
+        for i in range(scales.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(rotations.shape[1]):
+            l.append('rot_{}'.format(i))
+
+        dtype_full = [(attribute, 'f4') for attribute in l]
+
+        elements = np.empty(xyzs.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyzs, f_dc, opacities, scales, rotations), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+
+        PlyData([el]).write(path)
+    
+    def load_ply(self, path, compatible=True):
+
+        from plyfile import PlyData, PlyElement
+
+        plydata = PlyData.read(path)
+
+        xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                        np.asarray(plydata.elements[0]["y"]),
+                        np.asarray(plydata.elements[0]["z"])),  axis=1)
+        print("Number of points at loading : ", xyz.shape[0])
+
+        opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+        shs = np.zeros((xyz.shape[0], 3))
+        shs[:, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+        shs[:, 1] = np.asarray(plydata.elements[0]["f_dc_1"])
+        shs[:, 2] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+        scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+        scales = np.zeros((xyz.shape[0], len(scale_names)))
+        for idx, attr_name in enumerate(scale_names):
+            scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot_")]
+        rots = np.zeros((xyz.shape[0], len(rot_names)))
+        for idx, attr_name in enumerate(rot_names):
+            rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+          
+        gaussians = np.concatenate([xyz, opacities, scales, rots, shs], axis=1)
+        gaussians = torch.from_numpy(gaussians).float() # cpu
+
+        if compatible:
+            gaussians[..., 3:4] = torch.sigmoid(gaussians[..., 3:4])
+            gaussians[..., 4:7] = torch.exp(gaussians[..., 4:7])
+            gaussians[..., 11:] = 0.28209479177387814 * gaussians[..., 11:] + 0.5
+
+        return gaussians

core/models.py

@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+import kiui
+from kiui.lpips import LPIPS
+
+from core.unet import UNet
+from core.options import Options
+from core.gs import GaussianRenderer
+
+
+
+class LGM(nn.Module):
+    def __init__(
+        self,
+        opt: Options,
+    ):
+        super().__init__()
+
+        self.opt = opt
+
+        # unet
+        self.unet = UNet(
+            9, 14 * self.opt.gaussian_perpixel, 
+            down_channels=self.opt.down_channels,
+            down_attention=self.opt.down_attention,
+            mid_attention=self.opt.mid_attention,
+            up_channels=self.opt.up_channels,
+            up_attention=self.opt.up_attention,
+            num_views=self.opt.num_input_views,
+            num_frames=self.opt.num_frames,
+            use_temp_attn=self.opt.use_temp_attn
+        )
+
+        # last conv
+        self.conv = nn.Conv2d(14 * self.opt.gaussian_perpixel, 14 * self.opt.gaussian_perpixel, kernel_size=1) # NOTE: maybe remove it if train again
+
+        # Gaussian Renderer
+        self.gs = GaussianRenderer(opt)
+
+        # activations...
+        self.pos_act = lambda x: x.clamp(-1, 1)
+        self.scale_act = lambda x: 0.1 * F.softplus(x)
+        self.opacity_act = lambda x: torch.sigmoid(x)
+        self.rot_act = lambda x: F.normalize(x, dim=-1)
+        self.rgb_act = lambda x: 0.5 * torch.tanh(x) + 0.5 # NOTE: may use sigmoid if train again
+
+        # LPIPS loss
+        if self.opt.lambda_lpips > 0:
+            self.lpips_loss = LPIPS(net='vgg')
+            self.lpips_loss.requires_grad_(False)
+
+
+    def state_dict(self, **kwargs):
+        # remove lpips_loss
+        state_dict = super().state_dict(**kwargs)
+        for k in list(state_dict.keys()):
+            if 'lpips_loss' in k:
+                del state_dict[k]
+        return state_dict
+
+
+    def prepare_default_rays(self, device, elevation=0):
+        
+        from kiui.cam import orbit_camera
+        from core.utils import get_rays
+
+        cam_poses = np.stack([
+            orbit_camera(elevation, 0, radius=self.opt.cam_radius),
+            orbit_camera(elevation, 90, radius=self.opt.cam_radius),
+            orbit_camera(elevation, 180, radius=self.opt.cam_radius),
+            orbit_camera(elevation, 270, radius=self.opt.cam_radius),
+        ], axis=0) # [4, 4, 4]
+        cam_poses = torch.from_numpy(cam_poses)
+
+        rays_embeddings = []
+        for i in range(cam_poses.shape[0]):
+            rays_o, rays_d = get_rays(cam_poses[i], self.opt.input_size, self.opt.input_size, self.opt.fovy) # [h, w, 3]
+            rays_plucker = torch.cat([torch.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1) # [h, w, 6]
+            rays_embeddings.append(rays_plucker)
+
+            ## visualize rays for plotting figure
+            # kiui.vis.plot_image(rays_d * 0.5 + 0.5, save=True)
+
+        rays_embeddings = torch.stack(rays_embeddings, dim=0).permute(0, 3, 1, 2).contiguous().to(device) # [V, 6, h, w]
+        
+        return rays_embeddings
+
+        
+
+    def forward_gaussians(self, images):
+        # images: [B, T, 4, 9, H, W]
+        # return: Gaussians: [B, dim_t]
+
+        B, TV, C, H, W = images.shape
+        T = self.opt.num_frames
+        V = TV // T
+        images = images.view(B*T*V, C, H, W)
+
+        x = self.unet(images) # [B*4, 14, h, w]
+        x = self.conv(x) # [B*4, 14, h, w]
+
+        x = x.reshape(B*T, V, 14 * self.opt.gaussian_perpixel, self.opt.splat_size, self.opt.splat_size)
+
+        x = x.permute(0, 1, 3, 4, 2).reshape(B*T, -1, 14).contiguous()
+        
+        pos = self.pos_act(x[..., 0:3]) # [B, N, 3]
+        opacity = self.opacity_act(x[..., 3:4])
+        scale = self.scale_act(x[..., 4:7])
+        rotation = self.rot_act(x[..., 7:11])
+        rgbs = self.rgb_act(x[..., 11:])
+
+        gaussians = torch.cat([pos, opacity, scale, rotation, rgbs], dim=-1) # [B, T, N, 14]
+        
+        return gaussians
+
+    
+    def forward(self, data, step_ratio=1):
+        # data: output of the dataloader
+        # return: loss
+
+        results = {}
+        loss = 0
+
+        images = data['input'] # [B, Tx4, 9, h, W], input features
+
+        B, TV, C, H, W = images.shape
+        T = self.opt.num_frames
+        
+        # use the first view to predict gaussians
+        gaussians = self.forward_gaussians(images) # [B * T, N, 14]
+
+        results['gaussians'] = gaussians
+
+        # always use white bg
+        bg_color = torch.ones(3, dtype=torch.float32, device=gaussians.device)
+        
+        # use the other views for rendering and supervision
+        data['cam_view'] = data['cam_view'].reshape(B*T, -1, *data['cam_view'].shape[2:])
+        data['cam_view_proj'] = data['cam_view_proj'].reshape(B*T, -1, *data['cam_view_proj'].shape[2:])
+        data['cam_pos'] = data['cam_pos'].reshape(B*T, -1, *data['cam_pos'].shape[2:])
+        
+        results = self.gs.render(gaussians, data['cam_view'], data['cam_view_proj'], data['cam_pos'], bg_color=bg_color)
+        pred_images = results['image'] # [B*T, V, C, output_size, output_size]
+        pred_alphas = results['alpha'] # [B*T, V, 1, output_size, output_size]
+
+        results['images_pred'] = pred_images
+        results['alphas_pred'] = pred_alphas
+
+
+        data['images_output'] = data['images_output'].reshape(B*T, -1, *data['images_output'].shape[2:])
+        data['masks_output'] = data['masks_output'].reshape(B*T, -1, *data['masks_output'].shape[2:])
+
+        gt_images = data['images_output'] # [B*T, V, 3, output_size, output_size], ground-truth novel views
+        gt_masks = data['masks_output'] # [B*T, V, 1, output_size, output_size], ground-truth masks
+
+        gt_images = gt_images * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
+
+        loss_mse = F.mse_loss(pred_images, gt_images) + F.mse_loss(pred_alphas, gt_masks)
+        loss = loss + loss_mse
+
+        if self.opt.lambda_lpips > 0:
+            loss_lpips = self.lpips_loss(
+                # downsampled to at most 256 to reduce memory cost
+                F.interpolate(gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False), 
+                F.interpolate(pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
+            ).mean()
+            results['loss_lpips'] = loss_lpips
+            loss = loss + self.opt.lambda_lpips * loss_lpips
+            
+        results['loss'] = loss
+
+        # metric
+        with torch.no_grad():
+            psnr = -10 * torch.log10(torch.mean((pred_images.detach() - gt_images) ** 2))
+            results['psnr'] = psnr
+
+        return results

core/options.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tyro
+from dataclasses import dataclass
+from typing import Tuple, Literal, Dict, Optional
+
+
+@dataclass
+class Options:
+    ### model
+    # Unet image input size
+    input_size: int = 256
+    # Unet definition
+    down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024, 1024)
+    down_attention: Tuple[bool, ...] = (False, False, False, True, True, True)
+    mid_attention: bool = True
+    up_channels: Tuple[int, ...] = (1024, 1024, 512, 256)
+    up_attention: Tuple[bool, ...] = (True, True, True, False)
+    # Unet output size, dependent on the input_size and U-Net structure!
+    splat_size: int = 64
+    # gaussian render size
+    output_size: int = 256
+
+    ### dataset
+    # data mode (only support s3 now)
+    data_mode: str = '4d'
+    # fovy of the dataset
+    fovy: float = 49.1
+    # camera near plane
+    znear: float = 0.5
+    # camera far plane
+    zfar: float = 2.5
+    # number of all views (input + output)
+    num_views: int = 12
+    # number of views
+    num_input_views: int = 4
+    # camera radius
+    cam_radius: float = 1.5 # to better use [-1, 1]^3 space
+    # num workers
+    num_workers: int = 16
+    datalist: str=''
+
+    ### training
+    # workspace
+    workspace: str = './workspace'
+    # resume
+    resume: Optional[str] = None
+    # batch size (per-GPU)
+    batch_size: int = 8
+    # gradient accumulation
+    gradient_accumulation_steps: int = 1
+    # training epochs
+    num_epochs: int = 30
+    # lpips loss weight
+    lambda_lpips: float = 1.0
+    # gradient clip
+    gradient_clip: float = 1.0
+    # mixed precision
+    mixed_precision: str = 'bf16'
+    # learning rate
+    lr: float = 4e-4
+    # augmentation prob for grid distortion
+    prob_grid_distortion: float = 0.5
+    # augmentation prob for camera jitter
+    prob_cam_jitter: float = 0.5
+    # number of gaussians per pixel
+    gaussian_perpixel: int = 1
+
+    ### testing
+    # test image path
+    test_path: Optional[str] = None
+
+    ### misc
+    # nvdiffrast backend setting
+    force_cuda_rast: bool = False
+    # render fancy video with gaussian scaling effect
+    fancy_video: bool = False
+
+    # 4D
+    num_frames: int = 8
+    use_temp_attn: bool = True
+    shuffle_input: bool = True
+    
+    # s3
+    sample_by_anim: bool = True
+
+    # interp
+    interpresume: Optional[str] = None
+    interpolate_rate: int = 3
+    
+
+# all the default settings
+config_defaults: Dict[str, Options] = {}
+config_doc: Dict[str, str] = {}
+
+config_doc['lrm'] = 'the default settings for LGM'
+config_defaults['lrm'] = Options()
+
+
+config_doc['big'] = 'big model with higher resolution Gaussians'
+config_defaults['big'] = Options(
+    input_size=256,
+    up_channels=(1024, 1024, 512, 256, 128), # one more decoder
+    up_attention=(True, True, True, False, False),
+    splat_size=128,
+    output_size=512, # render & supervise Gaussians at a higher resolution.
+    batch_size=1,
+    num_views=8,
+    gradient_accumulation_steps=1,
+    mixed_precision='bf16',
+    resume='pretrained/model_fp16_fixrot.safetensors',
+)
+
+
+AllConfigs = tyro.extras.subcommand_type_from_defaults(config_defaults, config_doc)

core/provider_objaverse_4d.py

@@ -0,0 +1,254 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import cv2
+import random
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from torch.utils.data import Dataset
+
+import kiui
+from core.options import Options
+from core.utils import get_rays, grid_distortion, orbit_camera_jitter
+
+from kiui.cam import orbit_camera
+
+import tarfile
+from io import BytesIO
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+
+def load_np_array_from_tar(tar, path):
+    array_file = BytesIO()
+    array_file.write(tar.extractfile(path).read())
+    array_file.seek(0)
+    return np.load(array_file)
+
+
+class ObjaverseDataset(Dataset):
+
+    def _warn(self):
+        raise NotImplementedError('this dataset is just an example and cannot be used directly, you should modify it to your own setting! (search keyword TODO)')
+
+    def __init__(self, opt: Options, training=True, evaluating=False):
+        
+        self.opt = opt
+        self.training = training
+        self.evaluating = evaluating
+
+        self.items = []
+        with open(self.opt.datalist, 'r') as f:
+            for line in f.readlines():
+                self.items.append(line.strip())
+        
+
+        anim_map = {}
+        for x in self.items:
+            k = x.split('-')[1]
+            if k in anim_map:
+                anim_map[k] += '|'+x
+            else:
+                anim_map[k] = x
+        self.items = list(anim_map.values())
+
+
+        # naive split
+        if self.training:
+            self.items = self.items[:-self.opt.batch_size]
+        elif self.evaluating:
+            self.items = self.items[::1000]
+        else:
+            self.items = self.items[-self.opt.batch_size:]
+
+
+        # default camera intrinsics
+        self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy))
+        self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32)
+        self.proj_matrix[0, 0] = 1 / self.tan_half_fov
+        self.proj_matrix[1, 1] = 1 / self.tan_half_fov
+        self.proj_matrix[2, 2] = (self.opt.zfar + self.opt.znear) / (self.opt.zfar - self.opt.znear)
+        self.proj_matrix[3, 2] = - (self.opt.zfar * self.opt.znear) / (self.opt.zfar - self.opt.znear)
+        self.proj_matrix[2, 3] = 1
+
+    def __len__(self):
+        return len(self.items)
+
+    def _get_batch(self, idx):
+        if self.training:
+            uid = random.choice(self.items[idx].split('|'))
+        else:
+            uid = self.items[idx].split('|')[0]
+
+        results = {}
+
+        # load num_views images
+        images = []
+        masks = []
+        cam_poses = []
+        
+        if self.training and self.opt.shuffle_input:
+            vids = np.random.permutation(np.arange(32, 48))[:self.opt.num_input_views].tolist() + np.random.permutation(32).tolist()
+        else:
+            vids = np.arange(32, 48, 4).tolist() + np.arange(32).tolist()
+
+
+        random_tar_name = 'random_clip/' + uid
+        fixed_16_tar_name = 'fixed_16_clip/' + uid
+
+        local_random_tar_name = os.environ["DATA_HOME"] + random_tar_name.replace('/', '-')
+        local_fixed_16_tar_name = os.environ["DATA_HOME"] + fixed_16_tar_name.replace('/', '-')
+
+        tar_random = tarfile.open(local_random_tar_name)
+        tar_fixed = tarfile.open(local_fixed_16_tar_name)
+        
+
+        T = self.opt.num_frames
+        for t_idx in range(T):
+            t = t_idx
+            vid_cnt = 0
+            for vid in vids:
+                if vid >= 32:
+                    vid = vid % 32
+                    tar = tar_fixed
+                else:
+                    tar = tar_random
+
+                image_path = os.path.join('.', f'{vid:03d}/img', f'{t:03d}.jpg')
+                mask_path = os.path.join('.', f'{vid:03d}/mask', f'{t:03d}.png')
+
+                elevation_path = os.path.join('.', f'{vid:03d}/camera', f'elevation.npy')
+                rotation_path = os.path.join('.', f'{vid:03d}/camera', f'rotation.npy')
+
+                image = np.frombuffer(tar.extractfile(image_path).read(), np.uint8)
+                image = torch.from_numpy(cv2.imdecode(image, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255) # [512, 512, 4] in [0, 1]
+
+                azi = load_np_array_from_tar(tar, rotation_path)[t, None]
+                elevation = load_np_array_from_tar(tar, elevation_path)[t, None] * -1       # to align with pretrained LGM
+                azi = float(azi)
+                elevation = float(elevation)
+                c2w = torch.from_numpy(orbit_camera(elevation, azi, radius=1.5, opengl=True))
+                
+                image = image.permute(2, 0, 1) # [4, 512, 512]
+
+                mask = np.frombuffer(tar.extractfile(mask_path).read(), np.uint8)
+                mask = torch.from_numpy(cv2.imdecode(mask, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255).unsqueeze(0)  # [512, 512, 4] in [0, 1]
+
+                image = F.interpolate(image.unsqueeze(0), size=(512, 512), mode='nearest').squeeze(0)
+                mask = F.interpolate(mask.unsqueeze(0), size=(512, 512), mode='nearest').squeeze(0)
+
+                image = image[:3] * mask + (1 - mask) # [3, 512, 512], to white bg
+                image = image[[2,1,0]].contiguous() # bgr to rgb
+
+                images.append(image)
+                masks.append(mask.squeeze(0))
+                cam_poses.append(c2w)
+
+                vid_cnt += 1
+                if vid_cnt == self.opt.num_views:
+                    break
+
+            if vid_cnt < self.opt.num_views:
+                print(f'[WARN] dataset {uid}: not enough valid views, only {vid_cnt} views found!')
+                n = self.opt.num_views - vid_cnt
+                images = images + [images[-1]] * n
+                masks = masks + [masks[-1]] * n
+                cam_poses = cam_poses + [cam_poses[-1]] * n
+          
+        images = torch.stack(images, dim=0) # [V, C, H, W]
+        masks = torch.stack(masks, dim=0) # [V, H, W]
+        cam_poses = torch.stack(cam_poses, dim=0) # [V, 4, 4]
+
+        # normalized camera feats as in paper (transform the first pose to a fixed position)
+        transform = torch.tensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, self.opt.cam_radius], [0, 0, 0, 1]], dtype=torch.float32) @ torch.inverse(cam_poses[0])
+        cam_poses = transform.unsqueeze(0) @ cam_poses  # [V, 4, 4]
+        
+        images_input = F.interpolate(images.reshape(T, self.opt.num_views, *images.shape[1:])[:, :self.opt.num_input_views].reshape(-1, *images.shape[1:]).clone(), size=(self.opt.input_size, self.opt.input_size), mode='bilinear', align_corners=False) # [V, C, H, W]
+        cam_poses_input = cam_poses.reshape(T, self.opt.num_views, *cam_poses.shape[1:])[:, :self.opt.num_input_views].reshape(-1, *cam_poses.shape[1:]).clone()
+
+        # data augmentation
+        if self.training:
+            images_input = images_input.reshape(T, self.opt.num_input_views, *images_input.shape[1:])
+            cam_poses_input = cam_poses_input.reshape(T, self.opt.num_input_views, *cam_poses.shape[1:])
+
+            # apply random grid distortion to simulate 3D inconsistency
+            if random.random() < self.opt.prob_grid_distortion:
+                for t in range(T):
+                    images_input[t, 1:] = grid_distortion(images_input[t, 1:])
+            # apply camera jittering (only to input!)
+            if random.random() < self.opt.prob_cam_jitter:
+                for t in range(T):
+                    cam_poses_input[t, 1:] = orbit_camera_jitter(cam_poses_input[t, 1:])
+
+            images_input = images_input.reshape(-1, *images_input.shape[2:])
+            cam_poses_input = cam_poses_input.reshape(-1, *cam_poses.shape[1:])
+
+        # masking other views
+        images_input = images_input.reshape(T, self.opt.num_input_views, *images_input.shape[1:])
+        images_input[1:, 1:] = images_input[0:1, 1:]
+        images_input = images_input.reshape(-1, *images_input.shape[2:])
+
+        cam_poses_input = cam_poses_input.reshape(T, self.opt.num_input_views, *cam_poses.shape[1:])
+        cam_poses_input[1:, 1:] = cam_poses_input[0:1, 1:]
+        cam_poses_input = cam_poses_input.reshape(-1, *cam_poses.shape[1:])
+
+        images_input = TF.normalize(images_input, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+
+        # resize render ground-truth images, range still in [0, 1]
+        results['images_output'] = F.interpolate(images, size=(self.opt.output_size, self.opt.output_size), mode='bilinear', align_corners=False) # [V, C, output_size, output_size]
+        results['masks_output'] = F.interpolate(masks.unsqueeze(1), size=(self.opt.output_size, self.opt.output_size), mode='bilinear', align_corners=False) # [V, 1, output_size, output_size]
+
+        # build rays for input views
+        rays_embeddings = []
+        for i in range(self.opt.num_input_views * T):
+            rays_o, rays_d = get_rays(cam_poses_input[i], self.opt.input_size, self.opt.input_size, self.opt.fovy) # [h, w, 3]
+            rays_plucker = torch.cat([torch.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1) # [h, w, 6]
+            rays_embeddings.append(rays_plucker)
+
+     
+        rays_embeddings = torch.stack(rays_embeddings, dim=0).permute(0, 3, 1, 2).contiguous() # [V, 6, h, w]
+
+        final_input = torch.cat([images_input, rays_embeddings], dim=1) # [V=4, 9, H, W]
+        results['input'] = final_input
+
+        # opengl to colmap camera for gaussian renderer
+        cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+        
+        # cameras needed by gaussian rasterizer
+        cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+        cam_view_proj = cam_view @ self.proj_matrix # [V, 4, 4]
+        cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+        
+        results['cam_view'] = cam_view
+        results['cam_view_proj'] = cam_view_proj
+        results['cam_pos'] = cam_pos
+
+        return results
+
+    def __getitem__(self, idx):
+        while True:
+            try:
+                results = self._get_batch(idx)
+                break
+            except Exception as e:
+                print(f"{e}")
+                idx = random.randint(0, len(self.items) - 1)
+        return results

core/provider_objaverse_4d_interp.py

@@ -0,0 +1,283 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import cv2
+import random
+import numpy as np
+import copy
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from torch.utils.data import Dataset
+
+import kiui
+from core.options import Options
+from core.utils import get_rays, grid_distortion, orbit_camera_jitter
+
+from kiui.cam import orbit_camera
+
+import tarfile
+from io import BytesIO
+
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+
+def load_np_array_from_tar(tar, path):
+    array_file = BytesIO()
+    array_file.write(tar.extractfile(path).read())
+    array_file.seek(0)
+    return np.load(array_file)
+
+def interpolate_tensors(tensor):
+    # Extract the first and last tensors along the first dimension (B)
+    start_tensor = tensor[0]    # shape [4, 3, 256, 256]
+    end_tensor = tensor[-1]     # shape [4, 3, 256, 256]
+    tensor_interp = copy.deepcopy(tensor)
+
+    # Iterate over the range from 1 to second-last index
+    for i in range(1, tensor.size(0) - 1):
+        # Calculate the weight for interpolation
+
+        weight = (i - 0) / (tensor.size(0) - 1)
+        # Interpolate between start_tensor and end_tensor
+        tensor_interp[i] = torch.lerp(start_tensor, end_tensor, weight)
+
+
+    return tensor_interp
+
+class ObjaverseDataset(Dataset):
+
+    def _warn(self):
+        raise NotImplementedError('this dataset is just an example and cannot be used directly, you should modify it to your own setting! (search keyword TODO)')
+
+    def __init__(self, opt: Options, training=True, evaluating=False):
+        
+        self.opt = opt
+        self.training = training
+        self.evaluating = evaluating
+
+        self.items = []
+        with open(self.opt.datalist, 'r') as f:
+            for line in f.readlines():
+                self.items.append(line.strip())
+        
+        anim_map = {}
+        for x in self.items:
+            k = x.split('-')[1]
+            if k in anim_map:
+                anim_map[k] += '|'+x
+            else:
+                anim_map[k] = x
+        self.items = list(anim_map.values())
+
+
+
+        # naive split
+        if self.training:
+            self.items = self.items[:-self.opt.batch_size]
+        elif self.evaluating:
+            self.items = self.items[::1000]
+        else:
+            self.items = self.items[-self.opt.batch_size:]
+
+
+        # default camera intrinsics
+        self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy))
+        self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32)
+        self.proj_matrix[0, 0] = 1 / self.tan_half_fov
+        self.proj_matrix[1, 1] = 1 / self.tan_half_fov
+        self.proj_matrix[2, 2] = (self.opt.zfar + self.opt.znear) / (self.opt.zfar - self.opt.znear)
+        self.proj_matrix[3, 2] = - (self.opt.zfar * self.opt.znear) / (self.opt.zfar - self.opt.znear)
+        self.proj_matrix[2, 3] = 1
+
+
+    def __len__(self):
+        return len(self.items)
+
+    def _get_batch(self, idx):
+        # uid = self.items[idx]
+        if self.training:
+            uid = random.choice(self.items[idx].split('|'))
+        else:
+            uid = self.items[idx].split('|')[0]
+
+        results = {}
+
+        # load num_views images
+        images = []
+        masks = []
+        cam_poses = []
+        
+        if self.training and self.opt.shuffle_input:
+            vids = np.random.permutation(np.arange(32, 48))[:self.opt.num_input_views].tolist() + np.random.permutation(32).tolist()
+        else:
+            vids = np.arange(32, 48, 4).tolist() + np.arange(32).tolist()
+
+        random_tar_name = 'random_24fps/' + uid
+        fixed_16_tar_name = 'fixed_16_24fps/' + uid
+
+        local_random_tar_name = os.environ["DATA_HOME"] + random_tar_name.replace('/', '-')
+        local_fixed_16_tar_name = os.environ["DATA_HOME"] + fixed_16_tar_name.replace('/', '-')
+
+        tar_random = tarfile.open(local_random_tar_name)
+        tar_fixed = tarfile.open(local_fixed_16_tar_name)
+        
+        max_T = 24
+
+        T = self.opt.num_frames
+
+        start_frame = np.random.randint(max_T - T)
+
+        for t_idx in range(T):
+            t = start_frame + t_idx
+            vid_cnt = 0
+            for vid in vids:
+                if vid >= 32:
+                    vid = vid % 32
+                    tar = tar_fixed
+                else:
+                    tar = tar_random
+
+                image_path = os.path.join('.', f'{vid:03d}/img', f'{t:03d}.jpg')
+                mask_path = os.path.join('.', f'{vid:03d}/mask', f'{t:03d}.png')
+
+                elevation_path = os.path.join('.', f'{vid:03d}/camera', f'elevation.npy')
+                rotation_path = os.path.join('.', f'{vid:03d}/camera', f'rotation.npy')
+
+                try :
+                    image = np.frombuffer(tar.extractfile(image_path).read(), np.uint8)
+                    image = torch.from_numpy(cv2.imdecode(image, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255) # [512, 512, 4] in [0, 1]
+
+                    azi = load_np_array_from_tar(tar, rotation_path)[t, None]
+                    elevation = load_np_array_from_tar(tar, elevation_path)[t, None] * -1       # to align with pretrained LGM
+                    azi = float(azi)
+                    elevation = float(elevation)
+                    c2w = torch.from_numpy(orbit_camera(elevation, azi, radius=1.5, opengl=True))
+
+                    image = image.permute(2, 0, 1) # [4, 512, 512]
+
+                    mask = np.frombuffer(tar.extractfile(mask_path).read(), np.uint8)
+                    mask = torch.from_numpy(cv2.imdecode(mask, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255).unsqueeze(0)  # [512, 512, 4] in [0, 1]
+                except:
+
+                    return self.__getitem__(idx - 1)
+                image = F.interpolate(image.unsqueeze(0), size=(512, 512), mode='nearest').squeeze(0)
+                mask = F.interpolate(mask.unsqueeze(0), size=(512, 512), mode='nearest').squeeze(0)
+
+                image = image[:3] * mask + (1 - mask) # [3, 512, 512], to white bg
+                image = image[[2,1,0]].contiguous() # bgr to rgb
+
+                images.append(image)
+                masks.append(mask.squeeze(0))
+                cam_poses.append(c2w)
+
+                vid_cnt += 1
+                if vid_cnt == self.opt.num_views:
+                    break
+
+            if vid_cnt < self.opt.num_views:
+                print(f'[WARN] dataset {uid}: not enough valid views, only {vid_cnt} views found!')
+                n = self.opt.num_views - vid_cnt
+                images = images + [images[-1]] * n
+                masks = masks + [masks[-1]] * n
+                cam_poses = cam_poses + [cam_poses[-1]] * n
+          
+        images = torch.stack(images, dim=0) # [V, C, H, W]
+        masks = torch.stack(masks, dim=0) # [V, H, W]
+        cam_poses = torch.stack(cam_poses, dim=0) # [V, 4, 4]
+
+        # normalized camera feats as in paper (transform the first pose to a fixed position)
+        transform = torch.tensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, self.opt.cam_radius], [0, 0, 0, 1]], dtype=torch.float32) @ torch.inverse(cam_poses[0])
+        cam_poses = transform.unsqueeze(0) @ cam_poses  # [V, 4, 4]
+        
+        images_input = F.interpolate(images.reshape(T, self.opt.num_views, *images.shape[1:])[:, :self.opt.num_input_views].reshape(-1, *images.shape[1:]).clone(), size=(self.opt.input_size, self.opt.input_size), mode='bilinear', align_corners=False) # [V, C, H, W]
+        cam_poses_input = cam_poses.reshape(T, self.opt.num_views, *cam_poses.shape[1:])[:, :self.opt.num_input_views].reshape(-1, *cam_poses.shape[1:]).clone()
+    
+        # data augmentation
+        if self.training:
+            images_input = images_input.reshape(T, self.opt.num_input_views, *images_input.shape[1:])
+            cam_poses_input = cam_poses_input.reshape(T, self.opt.num_input_views, *cam_poses.shape[1:])
+
+            # apply random grid distortion to simulate 3D inconsistency
+            if random.random() < self.opt.prob_grid_distortion:
+                for t in range(T):
+                    images_input[t, 1:] = grid_distortion(images_input[t, 1:])
+            # apply camera jittering (only to input!)
+            if random.random() < self.opt.prob_cam_jitter:
+                for t in range(T):
+                    cam_poses_input[t, 1:] = orbit_camera_jitter(cam_poses_input[t, 1:])
+
+            images_input = images_input.reshape(-1, *images_input.shape[2:])
+            cam_poses_input = cam_poses_input.reshape(-1, *cam_poses.shape[1:])
+
+        # masking other views
+        images_input = images_input.reshape(T, self.opt.num_input_views, *images_input.shape[1:])
+
+        images_input_interp = interpolate_tensors(images_input)
+
+        images_input[1:-1, :] = images_input_interp[1:-1, :]
+        images_input = images_input.reshape(-1, *images_input.shape[2:])
+
+        cam_poses_input = cam_poses_input.reshape(T, self.opt.num_input_views, *cam_poses.shape[1:])
+        cam_poses_input[1:, 1:] = cam_poses_input[0:1, 1:]
+        cam_poses_input = cam_poses_input.reshape(-1, *cam_poses.shape[1:])
+
+        images_input = TF.normalize(images_input, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+
+        # resize render ground-truth images, range still in [0, 1]
+        results['images_output'] = F.interpolate(images, size=(self.opt.output_size, self.opt.output_size), mode='bilinear', align_corners=False) # [V, C, output_size, output_size]
+        results['masks_output'] = F.interpolate(masks.unsqueeze(1), size=(self.opt.output_size, self.opt.output_size), mode='bilinear', align_corners=False) # [V, 1, output_size, output_size]
+
+        # build rays for input views
+        rays_embeddings = []
+        for i in range(self.opt.num_input_views * T):
+            rays_o, rays_d = get_rays(cam_poses_input[i], self.opt.input_size, self.opt.input_size, self.opt.fovy) # [h, w, 3]
+            rays_plucker = torch.cat([torch.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1) # [h, w, 6]
+            rays_embeddings.append(rays_plucker)
+
+     
+        rays_embeddings = torch.stack(rays_embeddings, dim=0).permute(0, 3, 1, 2).contiguous() # [V, 6, h, w]
+
+        final_input = torch.cat([images_input, rays_embeddings], dim=1) # [V=4, 9, H, W]
+        results['input'] = final_input
+
+        # opengl to colmap camera for gaussian renderer
+        cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+        
+        # cameras needed by gaussian rasterizer
+        cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+        cam_view_proj = cam_view @ self.proj_matrix # [V, 4, 4]
+        cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+        
+        results['cam_view'] = cam_view
+        results['cam_view_proj'] = cam_view_proj
+        results['cam_pos'] = cam_pos
+
+        return results
+
+    def __getitem__(self, idx):
+        while True:
+            try:
+                results = self._get_batch(idx)
+                break
+            except Exception as e:
+                # print(f"{e}")
+                idx = random.randint(0, len(self.items) - 1)
+        return results

core/unet.py

@@ -0,0 +1,428 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+from typing import Tuple, Literal
+from functools import partial
+
+from core.attention import MemEffAttention
+
+
+class MVAttention(nn.Module):
+    def __init__(
+        self, 
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        groups: int = 32,
+        eps: float = 1e-5,
+        residual: bool = True,
+        skip_scale: float = 1,
+        num_views: int = 4, 
+        num_frames: int = 8
+    ):
+        super().__init__()
+
+        self.residual = residual
+        self.skip_scale = skip_scale
+        self.num_views = num_views
+        self.num_frames = num_frames
+
+        self.norm = nn.GroupNorm(num_groups=groups, num_channels=dim, eps=eps, affine=True)
+        self.attn = MemEffAttention(dim, num_heads, qkv_bias, proj_bias, attn_drop, proj_drop)
+
+    def forward(self, x):
+        # x: [B*T*V, C, H, W]
+        BTV, C, H, W = x.shape
+        BT = BTV // self.num_views # assert BV % self.num_views == 0
+
+        res = x
+        x = self.norm(x)
+
+        x = x.reshape(BT, self.num_views, C, H, W).permute(0, 1, 3, 4, 2).contiguous().reshape(BT, -1, C).contiguous()
+        x = self.attn(x)
+        x = x.reshape(BT, self.num_views, H, W, C).permute(0, 1, 4, 2, 3).contiguous().reshape(BTV, C, H, W).contiguous()
+
+        if self.residual:
+            x = (x + res) * self.skip_scale
+        return x
+
+
+class TempAttention(nn.Module):
+    def __init__(
+        self, 
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        groups: int = 32,
+        eps: float = 1e-5,
+        residual: bool = True,
+        skip_scale: float = 1,
+        num_views: int = 4,
+        num_frames: int = 8
+    ):
+        super().__init__()
+
+        self.residual = residual
+        self.skip_scale = skip_scale
+        self.num_views = num_views
+        self.num_frames = num_frames
+
+        self.norm = nn.GroupNorm(num_groups=groups, num_channels=dim, eps=eps, affine=True)
+        self.attn = MemEffAttention(dim, num_heads, qkv_bias, proj_bias, attn_drop, proj_drop)
+
+    def forward(self, x):
+        # x: [B*T*V, C, H, W]
+        BTV, C, H, W = x.shape
+        BV = BTV // self.num_frames # assert BV % self.num_views == 0
+        B = BV // self.num_views
+
+        res = x
+        x = self.norm(x)
+
+        # BTV -> BVT
+        x = x.reshape(B, self.num_frames, self.num_views, C, H, W).permute(0, 2, 1, 3, 4, 5).contiguous()
+
+        x = x.reshape(BV, self.num_frames, C, H, W).permute(0, 1, 3, 4, 2).contiguous().reshape(BV, -1, C).contiguous().contiguous()
+        x = self.attn(x)
+        x = x.reshape(BV, self.num_frames, H, W, C).permute(0, 1, 4, 2, 3).contiguous().reshape(BTV, C, H, W).contiguous().contiguous()
+
+        # BVT -> BTV
+        x = x.reshape(B, self.num_views, self.num_frames, C, H, W).permute(0, 2, 1, 3, 4, 5).contiguous().reshape(BTV, C, H, W).contiguous()
+
+        if self.residual:
+            x = (x + res) * self.skip_scale
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        resample: Literal['default', 'up', 'down'] = 'default',
+        groups: int = 32,
+        eps: float = 1e-5,
+        skip_scale: float = 1, # multiplied to output
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.skip_scale = skip_scale
+
+        self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.act = F.silu
+
+        self.resample = None
+        if resample == 'up':
+            self.resample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
+        elif resample == 'down':
+            self.resample = nn.AvgPool2d(kernel_size=2, stride=2)
+        
+        self.shortcut = nn.Identity()
+        if self.in_channels != self.out_channels:
+            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
+
+    
+    def forward(self, x):
+        res = x
+
+        x = self.norm1(x)
+        x = self.act(x)
+
+        if self.resample:
+            res = self.resample(res)
+            x = self.resample(x)
+        
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.act(x)
+        x = self.conv2(x)
+
+        x = (x + self.shortcut(res)) * self.skip_scale
+
+        return x
+
+class DownBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        downsample: bool = True,
+        attention: bool = True,
+        attention_heads: int = 16,
+        skip_scale: float = 1,
+        num_views: int = 4,
+        num_frames: int = 8,
+        use_temp_attn=True
+    ):
+        super().__init__()
+ 
+        nets = []
+        attns = []
+        t_attns = []
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            nets.append(ResnetBlock(in_channels, out_channels, skip_scale=skip_scale))
+            if attention:
+                attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale, num_views=num_views, num_frames=num_frames))
+                t_attns.append(TempAttention(out_channels, attention_heads, skip_scale=skip_scale, num_views=num_views, num_frames=num_frames) if use_temp_attn else None)
+            else:
+                attns.append(None)
+                t_attns.append(None)
+        self.nets = nn.ModuleList(nets)
+        self.attns = nn.ModuleList(attns)
+        self.t_attns = nn.ModuleList(t_attns)
+
+        self.downsample = None
+        if downsample:
+            self.downsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x):
+        xs = []
+
+        for attn, t_attn, net in zip(self.attns, self.t_attns, self.nets):
+            x = net(x)
+            if attn:
+                x = attn(x)
+            if t_attn:
+                x = t_attn(x)
+            xs.append(x)
+
+        if self.downsample:
+            x = self.downsample(x)
+            xs.append(x)
+  
+        return x, xs
+
+
+class MidBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_layers: int = 1,
+        attention: bool = True,
+        attention_heads: int = 16,
+        skip_scale: float = 1,
+        num_views: int = 4,
+        num_frames: int = 8,
+        use_temp_attn=True
+    ):
+        super().__init__()
+
+        nets = []
+        attns = []
+        t_attns = []
+        # first layer
+        nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
+        # more layers
+        for i in range(num_layers):
+            nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
+            if attention:
+                attns.append(MVAttention(in_channels, attention_heads, skip_scale=skip_scale, num_views=num_views, num_frames=num_frames))
+                t_attns.append(TempAttention(in_channels, attention_heads, skip_scale=skip_scale, num_views=num_views, num_frames=num_frames) if use_temp_attn else None)
+            else:
+                attns.append(None)
+                t_attns.append(None)
+        self.nets = nn.ModuleList(nets)
+        self.attns = nn.ModuleList(attns)
+        self.t_attns = nn.ModuleList(t_attns)
+        
+    def forward(self, x):
+        x = self.nets[0](x)
+        for attn, t_attn,net in zip(self.attns, self.t_attns, self.nets[1:]):
+            if attn:
+                x = attn(x)
+            if t_attn:
+                x = t_attn(x)
+            x = net(x)
+        return x
+
+
+class UpBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        prev_out_channels: int,
+        out_channels: int,
+        num_layers: int = 1,
+        upsample: bool = True,
+        attention: bool = True,
+        attention_heads: int = 16,
+        skip_scale: float = 1,
+        num_views: int = 4,
+        num_frames: int = 8,
+        use_temp_attn=True
+    ):
+        super().__init__()
+
+        nets = []
+        attns = []
+        t_attns = []
+        for i in range(num_layers):
+            cin = in_channels if i == 0 else out_channels
+            cskip = prev_out_channels if (i == num_layers - 1) else out_channels
+
+            nets.append(ResnetBlock(cin + cskip, out_channels, skip_scale=skip_scale))
+            if attention:
+                attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale, num_views=num_views, num_frames=num_frames))
+                t_attns.append(TempAttention(out_channels, attention_heads, skip_scale=skip_scale, num_views=num_views, num_frames=num_frames) if use_temp_attn else None)
+            else:
+                attns.append(None)
+                t_attns.append(None)
+        self.nets = nn.ModuleList(nets)
+        self.attns = nn.ModuleList(attns)
+        self.t_attns = nn.ModuleList(t_attns)
+
+        self.upsample = None
+        if upsample:
+            self.upsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x, xs):
+
+        for attn, t_attn, net in zip(self.attns, self.t_attns, self.nets):
+            res_x = xs[-1]
+            xs = xs[:-1]
+            x = torch.cat([x, res_x], dim=1)
+            x = net(x)
+            if attn:
+                x = attn(x)
+            if t_attn:
+                x = t_attn(x)
+            
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2.0, mode='nearest')
+            x = self.upsample(x)
+
+        return x
+
+
+# it could be asymmetric!
+class UNet(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024),
+        down_attention: Tuple[bool, ...] = (False, False, False, True, True),
+        mid_attention: bool = True,
+        up_channels: Tuple[int, ...] = (1024, 512, 256),
+        up_attention: Tuple[bool, ...] = (True, True, False),
+        layers_per_block: int = 2,
+        skip_scale: float = np.sqrt(0.5),
+        num_views: int = 4,
+        num_frames: int = 8,
+        use_temp_attn: bool = True
+    ):
+        super().__init__()
+
+        # first
+        self.conv_in = nn.Conv2d(in_channels, down_channels[0], kernel_size=3, stride=1, padding=1)
+
+        # down
+        down_blocks = []
+        cout = down_channels[0]
+        for i in range(len(down_channels)):
+            cin = cout
+            cout = down_channels[i]
+
+            down_blocks.append(DownBlock(
+                cin, cout, 
+                num_layers=layers_per_block, 
+                downsample=(i != len(down_channels) - 1), # not final layer
+                attention=down_attention[i],
+                skip_scale=skip_scale,
+                num_views=num_views, 
+                num_frames=num_frames,
+                use_temp_attn=use_temp_attn
+            ))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+        # mid
+        self.mid_block = MidBlock(down_channels[-1], attention=mid_attention, skip_scale=skip_scale, num_views=num_views, num_frames=num_frames, use_temp_attn=use_temp_attn)
+
+        # up
+        up_blocks = []
+        cout = up_channels[0]
+        for i in range(len(up_channels)):
+            cin = cout
+            cout = up_channels[i]
+            cskip = down_channels[max(-2 - i, -len(down_channels))] # for assymetric
+
+            up_blocks.append(UpBlock(
+                cin, cskip, cout, 
+                num_layers=layers_per_block + 1, # one more layer for up
+                upsample=(i != len(up_channels) - 1), # not final layer
+                attention=up_attention[i],
+                skip_scale=skip_scale,
+                num_views=num_views, 
+                num_frames=num_frames,
+                use_temp_attn=use_temp_attn
+            ))
+        self.up_blocks = nn.ModuleList(up_blocks)
+
+        # last
+        self.norm_out = nn.GroupNorm(num_channels=up_channels[-1], num_groups=32, eps=1e-5)
+        self.conv_out = nn.Conv2d(up_channels[-1], out_channels, kernel_size=3, stride=1, padding=1)
+
+
+    def forward(self, x, return_mid_feature=False):
+        # x: [B, Cin, H, W]
+
+        # first
+        x = self.conv_in(x)
+        
+        # down
+        xss = [x]
+        for block in self.down_blocks:
+            x, xs = block(x)
+            xss.extend(xs)
+        
+        # mid
+        x = self.mid_block(x)
+        mid_feature = (x, xss)
+
+        # up
+        for block in self.up_blocks:
+            xs = xss[-len(block.nets):]
+            xss = xss[:-len(block.nets)]
+            x = block(x, xs)
+
+        # last
+        x = self.norm_out(x)
+        x = F.silu(x)
+        x = self.conv_out(x) # [B, Cout, H', W']
+
+        if return_mid_feature:
+            return x, *mid_feature
+
+        return x

core/utils.py

@@ -0,0 +1,109 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import roma
+from kiui.op import safe_normalize
+
+def get_rays(pose, h, w, fovy, opengl=True):
+
+    x, y = torch.meshgrid(
+        torch.arange(w, device=pose.device),
+        torch.arange(h, device=pose.device),
+        indexing="xy",
+    )
+    x = x.flatten()
+    y = y.flatten()
+
+    cx = w * 0.5
+    cy = h * 0.5
+
+    focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
+
+    camera_dirs = F.pad(
+        torch.stack(
+            [
+                (x - cx + 0.5) / focal,
+                (y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
+            ],
+            dim=-1,
+        ),
+        (0, 1),
+        value=(-1.0 if opengl else 1.0),
+    )  # [hw, 3]
+
+    rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1)  # [hw, 3]
+    rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
+
+    rays_o = rays_o.view(h, w, 3)
+    rays_d = safe_normalize(rays_d).view(h, w, 3)
+
+    return rays_o, rays_d
+
+def orbit_camera_jitter(poses, strength=0.1):
+    # poses: [B, 4, 4], assume orbit camera in opengl format
+    # random orbital rotate
+
+    B = poses.shape[0]
+    rotvec_x = poses[:, :3, 1] * strength * np.pi * (torch.rand(B, 1, device=poses.device) * 2 - 1)
+    rotvec_y = poses[:, :3, 0] * strength * np.pi / 2 * (torch.rand(B, 1, device=poses.device) * 2 - 1)
+
+    rot = roma.rotvec_to_rotmat(rotvec_x) @ roma.rotvec_to_rotmat(rotvec_y)
+    R = rot @ poses[:, :3, :3]
+    T = rot @ poses[:, :3, 3:]
+
+    new_poses = poses.clone()
+    new_poses[:, :3, :3] = R
+    new_poses[:, :3, 3:] = T
+    
+    return new_poses
+
+def grid_distortion(images, strength=0.5):
+    # images: [B, C, H, W]
+    # num_steps: int, grid resolution for distortion
+    # strength: float in [0, 1], strength of distortion
+
+    B, C, H, W = images.shape
+
+    num_steps = np.random.randint(8, 17)
+    grid_steps = torch.linspace(-1, 1, num_steps)
+
+    # have to loop batch...
+    grids = []
+    for b in range(B):
+        # construct displacement
+        x_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
+        x_steps = (x_steps + strength * (torch.rand_like(x_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
+        x_steps = (x_steps * W).long() # [num_steps]
+        x_steps[0] = 0
+        x_steps[-1] = W
+        xs = []
+        for i in range(num_steps - 1):
+            xs.append(torch.linspace(grid_steps[i], grid_steps[i + 1], x_steps[i + 1] - x_steps[i]))
+        xs = torch.cat(xs, dim=0) # [W]
+
+        y_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
+        y_steps = (y_steps + strength * (torch.rand_like(y_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
+        y_steps = (y_steps * H).long() # [num_steps]
+        y_steps[0] = 0
+        y_steps[-1] = H
+        ys = []
+        for i in range(num_steps - 1):
+            ys.append(torch.linspace(grid_steps[i], grid_steps[i + 1], y_steps[i + 1] - y_steps[i]))
+        ys = torch.cat(ys, dim=0) # [H]
+
+        # construct grid
+        grid_x, grid_y = torch.meshgrid(xs, ys, indexing='xy') # [H, W]
+        grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2]
+
+        grids.append(grid)
+    
+    grids = torch.stack(grids, dim=0).to(images.device) # [B, H, W, 2]
+
+    # grid sample
+    images = F.grid_sample(images, grids, align_corners=False)
+
+    return images
+

environment.yml

@@ -0,0 +1,32 @@
+name: l4gm
+channels:
+  - pyg
+  - nvidia/label/cuda-12.1.0
+  - pytorch
+  - conda-forge
+  - xformers
+dependencies:
+  - python=3.10
+  - pytorch=2.5.1
+  - pytorch-cuda=12.1
+  - torchvision
+  - xformers
+  - cuda
+  - cuda-nvcc
+  - numpy<2.0.0
+  - scipy
+  - rich
+  - pip
+  - setuptools
+  - ninja
+  - tqdm
+  - ray-default
+  - flatten-dict
+  - gcc_linux-64=11
+  - gxx_linux-64=11
+  - opencv
+  - transformers
+  - einops
+  - pip:
+    - -r requirements.txt
+    - git+https://github.com/nerfstudio-project/gsplat.git@v1.4.0

infer_3d.py

@@ -0,0 +1,233 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import imageio.v3 as iio
+import cv2
+import numpy as np
+import imageio
+
+import os
+import tyro
+import glob
+import imageio
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from safetensors.torch import load_file
+import time
+
+import kiui
+from kiui.cam import orbit_camera
+
+from core.options import AllConfigs, Options
+from core.models import LGM
+from mvdream.pipeline_mvdream import MVDreamPipeline
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+opt = tyro.cli(AllConfigs)
+
+# model
+model = LGM(opt)
+
+# resume pretrained checkpoint
+if opt.resume is not None:
+    if opt.resume.endswith('safetensors'):
+        ckpt = load_file(opt.resume, device='cpu')
+    else:
+        ckpt = torch.load(opt.resume, map_location='cpu')
+    model.load_state_dict(ckpt, strict=False)
+    print(f'[INFO] Loaded checkpoint from {opt.resume}')
+else:
+    print(f'[WARN] model randomly initialized, are you sure?')
+
+# device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.half().to(device)
+model.eval()
+
+bg_color = torch.tensor([255, 255, 255], dtype=torch.float32, device="cuda") / 255.
+
+rays_embeddings = model.prepare_default_rays(device)
+
+tan_half_fov = np.tan(0.5 * np.deg2rad(opt.fovy))
+proj_matrix = torch.zeros(4, 4, dtype=torch.float32, device=device)
+proj_matrix[0, 0] = 1 / tan_half_fov
+proj_matrix[1, 1] = 1 / tan_half_fov
+proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
+proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
+proj_matrix[2, 3] = 1
+
+# load image dream
+pipe = MVDreamPipeline.from_pretrained(
+    "ashawkey/imagedream-ipmv-diffusers", # remote weights
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+    # local_files_only=True,
+)
+pipe = pipe.to(device)
+
+
+def process_eval_video(video_path, T):
+    frames = iio.imread(video_path)
+    frames = [frames[x] for x in range(frames.shape[0])]
+    V = opt.num_input_views
+    img_TV = []
+    for t in range(T):
+
+        img = frames[t]
+        
+        img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
+        img = img.astype(np.float32) / 255.0
+
+        img_V = []
+        for v in range(V):
+            img_V.append(img)
+        img_TV.append(np.stack(img_V, axis=0))
+
+    return np.stack(img_TV, axis=0)
+
+
+# process function
+def process(opt: Options, path):
+    name = os.path.splitext(os.path.basename(path))[0]
+    print(f'[INFO] Processing {path} --> {name}')
+    os.makedirs(opt.workspace, exist_ok=True)
+
+    ref_video = process_eval_video(path, opt.num_frames) # [TV, 512, 512, 3]
+
+
+    end_time = time.time()
+
+    cv2.imwrite(os.path.join(opt.workspace, f'{name}_orig.png'), ref_video[0,0][..., ::-1] * 255)
+
+    mv_image = pipe('', ref_video[0,0], guidance_scale=5, num_inference_steps=30, elevation=0)
+    for v in range(4):
+        cv2.imwrite(os.path.join(opt.workspace, f'{name}_mv_{(v-1)%4:03d}.png'), mv_image[v][..., ::-1] * 255)
+    mv_image = np.stack([mv_image[1], mv_image[2], mv_image[3], mv_image[0]], axis=0) # [4, 256, 256, 3], float32
+
+
+    # generate gaussians
+    input_image = torch.from_numpy(mv_image).permute(0, 3, 1, 2).float().to(device) # [4, 3, 256, 256]
+    input_image = F.interpolate(input_image, size=(opt.input_size, opt.input_size), mode='bilinear', align_corners=False)
+    input_image = TF.normalize(input_image, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+
+    input_image = torch.cat([input_image, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, H, W]
+
+    with torch.no_grad():
+        with torch.autocast(device_type='cuda', dtype=torch.float16):
+            gaussians_all_frame = model.forward_gaussians(input_image)
+            
+            B, T, V = 1, gaussians_all_frame.shape[0]//opt.batch_size, opt.num_views
+            gaussians_all_frame = gaussians_all_frame.reshape(B, T, *gaussians_all_frame.shape[1:])
+
+            # align azimuth
+            best_azi = 0
+            best_diff = 1e8
+            for v, azi in enumerate(np.arange(-180, 180, 1)):
+                gaussians = gaussians_all_frame[:, 0]
+                
+                cam_poses = torch.from_numpy(orbit_camera(0, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+                
+                # cameras needed by gaussian rasterizer
+                cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                result = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)
+                image = result['image']
+                alpha = result['alpha']
+
+                image = image.squeeze(1).permute(0,2,3,1).squeeze(0).contiguous().float().cpu().numpy()
+                image = cv2.resize(image, (256, 256), interpolation=cv2.INTER_AREA)
+
+                diff = np.mean((image- ref_video[0,0]) ** 2)
+
+                if diff < best_diff:
+                    best_diff = diff
+                    best_azi = azi
+
+            print("Best aligned azimuth: ", best_azi)
+
+            mv_image = []
+            for v, azi in enumerate(np.arange(0, 360, 90)):
+                gaussians = gaussians_all_frame[:, 0]
+                
+                cam_poses = torch.from_numpy(orbit_camera(0, azi + best_azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+                
+                # cameras needed by gaussian rasterizer
+                cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                scale = 1
+
+                result = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)
+                image = result['image']
+                alpha = result['alpha']
+
+                imageio.imwrite(os.path.join(opt.workspace, f'{name}_{v:03d}.png'), (image.squeeze(1).permute(0,2,3,1).squeeze(0).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+                
+                if azi in [0, 90, 180, 270]:     
+                    rendered_image = image.squeeze(1)
+                    rendered_image = F.interpolate(rendered_image, (256, 256))
+                    rendered_image = rendered_image.permute(0,2,3,1).contiguous().float().cpu().numpy()
+                    mv_image.append(rendered_image)
+            mv_image = np.concatenate(mv_image, axis=0)
+            print(f"Generate 3D takes {time.time()-end_time} s")
+            
+            images = []
+            azimuth = np.arange(0, 360, 4, dtype=np.int32)
+            elevation = 0
+            for azi in azimuth:
+                gaussians = gaussians_all_frame[:, 0]
+                
+                cam_poses = torch.from_numpy(orbit_camera(elevation, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+                
+                # cameras needed by gaussian rasterizer
+                cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                scale = 1
+
+                image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)['image']
+                images.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+
+            images = np.concatenate(images, axis=0)
+            imageio.mimwrite(os.path.join(opt.workspace, f'{name}.mp4'), images, fps=30)
+
+
+    torch.cuda.empty_cache()
+
+
+
+assert opt.test_path is not None
+if os.path.isdir(opt.test_path):
+    file_paths = glob.glob(os.path.join(opt.test_path, "*"))
+else:
+    file_paths = [opt.test_path]
+
+for path in sorted(file_paths):
+    process(opt, path)

infer_4d.py

@@ -0,0 +1,392 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import imageio.v3 as iio
+import cv2
+import numpy as np
+import imageio
+
+from copy import deepcopy
+import os
+import tyro
+import glob
+import imageio
+import numpy as np
+import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from safetensors.torch import load_file
+
+import kiui
+from kiui.cam import orbit_camera
+
+from core.options import AllConfigs, Options
+from core.models import LGM
+import time
+
+from core.utils import get_rays, grid_distortion, orbit_camera_jitter
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+
+USE_INTERPOLATION = True    # set to false to disable interpolation
+MAX_RUNS = 100
+VIDEO_FPS = 30
+
+opt = tyro.cli(AllConfigs)
+
+# model
+model = LGM(opt)
+
+# resume pretrained checkpoint
+if opt.resume is not None:
+    if opt.resume.endswith('safetensors'):
+        ckpt = load_file(opt.resume, device='cpu')
+    else:
+        ckpt = torch.load(opt.resume, map_location='cpu')
+    model.load_state_dict(ckpt, strict=False)
+    print(f'[INFO] Loaded checkpoint from {opt.resume}')
+else:
+    print(f'[WARN] model randomly initialized, are you sure?')
+
+# device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.half().to(device)
+model.eval()
+
+bg_color = torch.tensor([255, 255, 255], dtype=torch.float32, device="cuda") / 255.
+
+
+rays_embeddings = model.prepare_default_rays(device)
+rays_embeddings = torch.cat([rays_embeddings for _ in range(opt.num_frames)])
+
+
+interp_opt = deepcopy(opt)
+interp_opt.num_frames = 4
+model_interp = LGM(interp_opt)
+# resume pretrained checkpoint
+if interp_opt.interpresume is not None:
+    if interp_opt.interpresume.endswith('safetensors'):
+        ckpt = load_file(interp_opt.interpresume, device='cpu')
+    else:
+        ckpt = torch.load(interp_opt.interpresume, map_location='cpu')
+    model_interp.load_state_dict(ckpt, strict=False)
+    print(f'[INFO] Loaded Interp checkpoint from {interp_opt.interpresume}')
+else:
+    print(f'[WARN] model_interp randomly initialized, are you sure?')
+
+# device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model_interp = model_interp.half().to(device)
+model_interp.eval()
+
+
+interp_rays_embeddings = model_interp.prepare_default_rays(device)
+interp_rays_embeddings = torch.cat([interp_rays_embeddings for _ in range(interp_opt.num_frames)])
+
+tan_half_fov = np.tan(0.5 * np.deg2rad(opt.fovy))
+proj_matrix = torch.zeros(4, 4, dtype=torch.float32, device=device)
+proj_matrix[0, 0] = 1 / tan_half_fov
+proj_matrix[1, 1] = 1 / tan_half_fov
+proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
+proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
+proj_matrix[2, 3] = 1
+
+def interpolate_tensors(tensor):
+    # Extract the first and last tensors along the first dimension (B)
+    start_tensor = tensor[0]    # shape [4, 3, 256, 256]
+    end_tensor = tensor[-1]     # shape [4, 3, 256, 256]
+    tensor_interp = deepcopy(tensor)
+
+    # Iterate over the range from 1 to second-last index
+
+    for i in range(1, tensor.shape[0] - 1):
+        # Calculate the weight for interpolation
+
+        weight = (i - 0) / (tensor.shape[0] - 1)
+        # Interpolate between start_tensor and end_tensor
+        tensor_interp[i] = torch.lerp(start_tensor, end_tensor, weight)
+
+
+    return tensor_interp
+
+def process_eval_video(frames, video_path, T, start_t=0, downsample_rate=1):
+    L = frames.shape[0]
+    vid_name =video_path.split('/')[-1].split('.')[0]
+    total_frames = L//downsample_rate
+    print(f'{start_t} / {total_frames}')
+    frames = [frames[x] for x in range(frames.shape[0])]
+    V = opt.num_input_views
+    img_TV = []
+    for t in range(T):
+        t += start_t
+        t = min(t, L//downsample_rate-1)
+        t*=downsample_rate
+
+        img = frames[t]
+
+        img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
+        img = img.astype(np.float32) / 255.0
+
+        img_V = []
+        for v in range(V):
+            img_V.append(img)
+        img_TV.append(np.stack(img_V, axis=0))
+
+    return np.stack(img_TV, axis=0), L//downsample_rate- start_t
+
+def load_mv_img(name, img_dir):
+    img_list = []
+    for v in range(4):
+        img = kiui.read_image(os.path.join(img_dir, name + f'_{v:03d}.png'), mode='uint8')
+        img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
+        img = img / 255.
+        img_list.append(img)
+    return np.stack(img_list, axis=0)
+
+
+
+# process function
+def process(opt: Options, path):
+    name = os.path.splitext(os.path.basename(path))[0]
+    print(f'[INFO] Processing {path} --> {name}')
+    os.makedirs(opt.workspace, exist_ok=True)
+    frames = iio.imread(path)
+    img_dir = opt.workspace
+    mv_image = load_mv_img(name, img_dir)
+
+    print(iio.immeta(path))
+    FPS = int(iio.immeta(path)['fps'])
+    downsample_rate = FPS // 15 if FPS > 15 else 1    # default reconstruction fps 15
+
+    
+
+    with torch.inference_mode():
+        with torch.autocast(device_type='cuda', dtype=torch.float16):
+            start_t = 0
+            gaussians_all_frame_all_run = []
+            gaussians_all_frame_all_run_w_interp = []
+            for run_idx in range(MAX_RUNS):
+                ref_video, end_t = process_eval_video(frames, path, opt.num_frames, start_t, downsample_rate=downsample_rate)
+                ref_video[:, 1:] = mv_image[None, 1:]   # repeat
+                input_image = torch.from_numpy(ref_video).reshape([-1, *ref_video.shape[2:]]).permute(0, 3, 1, 2).float().to(device) # [4, 3, 256, 256]
+                input_image = F.interpolate(input_image, size=(opt.input_size, opt.input_size), mode='bilinear', align_corners=False)
+                input_image = TF.normalize(input_image, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+                input_image = torch.cat([input_image, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, H, W]
+
+                end_time = time.time()
+
+                gaussians_all_frame = model.forward_gaussians(input_image)
+                print(f"Forward pass takes {time.time()-end_time} s")
+
+                B, T, V = 1, gaussians_all_frame.shape[0]//opt.batch_size, opt.num_views
+                gaussians_all_frame = gaussians_all_frame.reshape(B, T, *gaussians_all_frame.shape[1:])
+
+                if run_idx > 0:
+                    gaussians_all_frame_wo_inter = gaussians_all_frame[:, 1:max(end_t, 1)]
+                else:
+                    gaussians_all_frame_wo_inter = gaussians_all_frame
+
+                if gaussians_all_frame_wo_inter.shape[1] > 0 and USE_INTERPOLATION:
+                    # render multiview video
+                    render_img_TV = []
+                    for t in range(gaussians_all_frame.shape[1]):
+                        render_img_V = []
+                        for v, azi in enumerate(np.arange(0, 360, 90)):
+
+                            gaussians = gaussians_all_frame[:, t]
+
+                            cam_poses = torch.from_numpy(orbit_camera(0, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                            cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+
+                            # cameras needed by gaussian rasterizer
+                            cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                            cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                            cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                            rendered_image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)['image']
+                            rendered_image = rendered_image.squeeze(1)
+                            rendered_image = F.interpolate(rendered_image, (256, 256))
+                            rendered_image = rendered_image.permute(0,2,3,1).contiguous().float().cpu().numpy() # B H W C
+
+                            render_img_V.append(rendered_image)
+                        render_img_V = np.concatenate(render_img_V, axis=0) # V H W C
+                        render_img_TV.append(render_img_V)
+                    render_img_TV = np.stack(render_img_TV, axis=0)   # T V H W C
+                    ref_video = np.concatenate([np.stack([ref_video[ttt] for _ in range(opt.interpolate_rate)], 0)  for ttt in range(ref_video.shape[0])], 0)
+
+
+                    for tt in range(gaussians_all_frame_wo_inter.shape[1] -1 ):
+
+                        curr_ref_video = deepcopy( ref_video[ tt * opt.interpolate_rate:  tt * opt.interpolate_rate + interp_opt.num_frames ])
+                        curr_ref_video[0, 1:] = render_img_TV[tt, 1:]
+
+                        curr_ref_video[-1, 1:] = render_img_TV[tt+1, 1:]
+
+
+                        curr_ref_video = torch.from_numpy(curr_ref_video).float().to(
+                            device)  # [4, 3, 256, 256]
+
+                        images_input_interp = interpolate_tensors(curr_ref_video)
+
+                        curr_ref_video[1:-1, :] = images_input_interp[1:-1, :]
+
+                        input_image_interp = curr_ref_video.reshape([-1, *curr_ref_video.shape[2:]]).permute(0, 3, 1,  2).float().to(device)  # [4, 3, 256, 256]
+                        input_image_interp = F.interpolate(input_image_interp, size=(interp_opt.input_size, interp_opt.input_size), mode='bilinear',
+                                                    align_corners=False)
+                        input_image_interp = TF.normalize(input_image_interp, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+
+                        input_image_interp = torch.cat([input_image_interp, interp_rays_embeddings], dim=1).unsqueeze(0)  # [1, 4, 9, H, W]
+
+                        end_time = time.time()
+                        gaussians_interp_all_frame = model_interp.forward_gaussians(input_image_interp)
+                        print(f"Interpolate forward pass takes {time.time()-end_time} s")
+
+                        B, T, V = 1, gaussians_interp_all_frame.shape[0] // opt.batch_size, opt.num_views
+                        gaussians_interp_all_frame = gaussians_interp_all_frame.reshape(B, T, *gaussians_interp_all_frame.shape[1:])
+
+                        if tt > 0:
+                            gaussians_interp_all_frame = gaussians_interp_all_frame[:, 1:]
+
+                        gaussians_all_frame_all_run_w_interp.append(gaussians_interp_all_frame)
+
+                        
+
+                    gaussians_all_frame_all_run.append(gaussians_all_frame_wo_inter)
+                    start_t += opt.num_frames -1
+
+                    mv_image = []
+                    for v, azi in enumerate(np.arange(0, 360, 90)):
+                        gaussians = gaussians_all_frame_wo_inter[:, -1]
+                        cam_poses = torch.from_numpy(orbit_camera(0, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+                        cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+                        # cameras needed by gaussian rasterizer
+                        cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                        cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                        cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                        rendered_image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)['image']
+                        rendered_image = rendered_image.squeeze(1)
+                        rendered_image = F.interpolate(rendered_image, (256, 256))
+                        rendered_image = rendered_image.permute(0,2,3,1).contiguous().float().cpu().numpy()
+                        mv_image.append(rendered_image)
+                    mv_image = np.concatenate(mv_image, axis=0)
+                elif gaussians_all_frame_wo_inter.shape[1] > 0:
+                    gaussians_all_frame_all_run.append(gaussians_all_frame_wo_inter)
+                    start_t += opt.num_frames -1
+                else:
+                    break
+
+            gaussians_all_frame_wo_interp = torch.cat(gaussians_all_frame_all_run, dim=1)
+            if USE_INTERPOLATION:
+                gaussians_all_frame_w_interp = torch.cat(gaussians_all_frame_all_run_w_interp, dim=1)
+
+            if USE_INTERPOLATION:
+                zip_dump = zip(["wo_interp", "w_interp"], [gaussians_all_frame_wo_interp, gaussians_all_frame_w_interp])
+            else:
+                zip_dump = zip(["wo_interp"], [gaussians_all_frame_wo_interp])
+
+            for sv_name, gaussians_all_frame in zip_dump:
+                if sv_name == "w_interp":
+                    ANIM_FPS = FPS / downsample_rate * gaussians_all_frame_w_interp.shape[1] / gaussians_all_frame_wo_interp.shape[1]
+                else:
+                    ANIM_FPS = FPS / downsample_rate
+                print(f"{sv_name} | input video fps: {FPS} | downsample rate: {downsample_rate} | animation fps: {ANIM_FPS} | output video fps: {VIDEO_FPS}")
+                render_img_TV = []
+                for t in range(gaussians_all_frame.shape[1]):
+                    render_img_V = []
+                    for v, azi in enumerate(np.arange(0, 360, 90)):
+
+                        gaussians = gaussians_all_frame[:, t]
+
+                        cam_poses = torch.from_numpy(orbit_camera(0, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                        cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+
+                        # cameras needed by gaussian rasterizer
+                        cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                        cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                        cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                        result = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)
+                        image = result['image']
+                        alpha = result['alpha']
+
+                        render_img_V.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+                    render_img_V = np.concatenate(render_img_V, axis=2)
+                    render_img_TV.append(render_img_V)
+                render_img_TV = np.concatenate(render_img_TV, axis=0)
+
+
+                images = []
+                azimuth = np.arange(0, 360, 1*30/VIDEO_FPS, dtype=np.int32)
+                elevation = 0
+                t = 0
+                delta_t = ANIM_FPS / VIDEO_FPS
+                for azi in azimuth:
+                    if azi in [0, 90, 180, 270]:
+                        cam_poses = torch.from_numpy(orbit_camera(elevation, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+                        cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+
+                        # cameras needed by gaussian rasterizer
+                        cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                        cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                        cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                        for _ in range(45):
+                            gaussians = gaussians_all_frame[:, int(t) % gaussians_all_frame.shape[1]]
+                            t += delta_t
+                            image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)['image']
+                            images.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+                    else:
+                        gaussians = gaussians_all_frame[:, int(t) % gaussians_all_frame.shape[1]]
+                        t += delta_t
+
+                        cam_poses = torch.from_numpy(orbit_camera(elevation, azi, radius=opt.cam_radius, opengl=True)).unsqueeze(0).to(device)
+
+                        cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+
+                        # cameras needed by gaussian rasterizer
+                        cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+                        cam_view_proj = cam_view @ proj_matrix # [V, 4, 4]
+                        cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+                        image = model.gs.render(gaussians, cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0), bg_color=bg_color)['image']
+                        images.append((image.squeeze(1).permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+
+                images = np.concatenate(images, axis=0)
+
+                torch.cuda.empty_cache()
+
+
+                imageio.mimwrite(os.path.join(opt.workspace, f'{sv_name}_{name}_fixed.mp4'), render_img_TV, fps=ANIM_FPS)
+                print("Fixed video saved.")
+                imageio.mimwrite(os.path.join(opt.workspace, f'{sv_name}_{name}.mp4'), images, fps=VIDEO_FPS)
+                print("Stop video saved.")
+
+
+assert opt.test_path is not None
+
+if os.path.isdir(opt.test_path):
+    file_paths = glob.glob(os.path.join(opt.test_path, "*"))
+else:
+    file_paths = [opt.test_path]
+
+for path in sorted(file_paths):
+    process(opt, path)

main.py

@@ -0,0 +1,274 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tyro
+import time
+import random
+
+import torch
+from core.options import AllConfigs
+from core.models import LGM
+
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from safetensors.torch import load_file
+
+import kiui
+from PIL import Image
+
+import json
+import os
+import numpy as np
+import imageio
+
+def main():    
+    opt = tyro.cli(AllConfigs)
+
+    accelerator = Accelerator(
+        mixed_precision=opt.mixed_precision,
+        gradient_accumulation_steps=opt.gradient_accumulation_steps,
+        # kwargs_handlers=[ddp_kwargs],
+    )
+    if accelerator.is_main_process:
+        print(opt)
+
+    # model
+    model = LGM(opt)
+
+    epoch_start = 0
+    if os.path.exists(f'{opt.workspace}/model.safetensors') and os.path.exists(f'{opt.workspace}/metadata.json'):
+        opt.resume = f'{opt.workspace}/model.safetensors'
+        with open(f'{opt.workspace}/metadata.json', 'r') as f:
+            dc = json.load(f)
+            epoch_start = dc['epoch'] + 1
+
+
+    # resume
+    if opt.resume is not None and opt.resume != 'None':
+        if opt.resume.endswith('safetensors'):
+            ckpt = load_file(opt.resume, device='cpu')
+        else:
+            ckpt = torch.load(opt.resume, map_location='cpu')
+        
+        # tolerant load (only load matching shapes)
+        # model.load_state_dict(ckpt, strict=False)
+        state_dict = model.state_dict()
+        for k, v in ckpt.items():
+            if k in state_dict: 
+                if state_dict[k].shape == v.shape:
+                    state_dict[k].copy_(v)
+                else:
+                    accelerator.print(f'[WARN] mismatching shape for param {k}: ckpt {v.shape} != model {state_dict[k].shape}, ignored.')
+            else:
+                accelerator.print(f'[WARN] unexpected param {k}: {v.shape}')
+    
+    # data
+    if opt.data_mode == '4d':
+        from core.provider_objaverse_4d import ObjaverseDataset as Dataset
+    elif opt.data_mode == '4d_interp':
+        from core.provider_objaverse_4d_interp import ObjaverseDataset as Dataset
+    else:
+        raise NotImplementedError
+
+    train_dataset = Dataset(opt, training=True)
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=opt.batch_size,
+        shuffle=True,
+        num_workers=opt.num_workers,
+        pin_memory=True,
+        drop_last=True,
+    )
+
+    test_dataset = Dataset(opt, training=False)
+    test_dataloader = torch.utils.data.DataLoader(
+        test_dataset,
+        batch_size=opt.batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    # optimizer
+    optimizer = torch.optim.AdamW(model.parameters(), lr=opt.lr, weight_decay=0.05, betas=(0.9, 0.95))
+
+    # scheduler (per-iteration)
+    total_steps = opt.num_epochs * len(train_dataloader)
+    pct_start = 3000 / total_steps
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=opt.lr, total_steps=total_steps, pct_start=pct_start)
+
+    if epoch_start > 0:
+        optimizer.load_state_dict(torch.load(os.path.join(opt.workspace, 'optimizer.pth'), map_location='cpu'))
+        scheduler.load_state_dict(torch.load(os.path.join(opt.workspace, 'scheduler.pth')))
+
+    # accelerate
+    model, optimizer, train_dataloader, test_dataloader, scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, test_dataloader, scheduler
+    )
+
+
+
+    # loop
+    os.makedirs(opt.workspace, exist_ok=True)
+    end_time = time.time()
+    for epoch in range(epoch_start, opt.num_epochs):
+        # train
+        model.train()
+        total_loss = 0
+        total_psnr = 0
+        for i, data in enumerate(train_dataloader):
+            with accelerator.accumulate(model):
+
+                optimizer.zero_grad()
+
+                step_ratio = (epoch + i / len(train_dataloader)) / opt.num_epochs
+
+                out = model(data, step_ratio)
+                loss = out['loss']
+                psnr = out['psnr']
+                accelerator.backward(loss)
+
+                # gradient clipping
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), opt.gradient_clip)
+
+                optimizer.step()
+                scheduler.step()
+
+                total_loss += loss.detach()
+                total_psnr += psnr.detach()
+
+            if accelerator.is_main_process:
+                # logging
+                if i % 10 == 0:
+                    mem_free, mem_total = torch.cuda.mem_get_info()    
+                    print(f"[INFO] {i}/{len(train_dataloader)} mem: {(mem_total-mem_free)/1024**3:.2f}/{mem_total/1024**3:.2f}G lr: {scheduler.get_last_lr()[0]:.7f} step_ratio: {step_ratio:.4f} loss: {loss.item():.6f} time: {time.time() - end_time:.6f}")
+                    end_time = time.time()
+                
+                # save log images
+                if i % 500 == 0:
+                    if '4d' in opt.data_mode:
+                        B, T, V = opt.batch_size, opt.num_frames, opt.num_views
+
+                        gt_images = data['images_output'].reshape(B, T, V, *data['images_output'].shape[2:]).detach() # [B, V, 3, output_size, output_size]
+                        pred_images = out['images_pred'].reshape(B, T, V, *out['images_pred'].shape[2:]).detach() # [B, V, 3, output_size, output_size]
+
+                        train_gt_images = []
+                        train_pred_images = []
+                        for t in range(T):
+                            train_gt_images_V = []
+                            train_pred_images_V = []
+                            for v in range(V):
+                                train_gt_images_V.append((gt_images[:, t, v].permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+                                train_pred_images_V.append((pred_images[:, t, v].permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+                            train_gt_images.append(np.concatenate(train_gt_images_V, axis=2))
+                            train_pred_images.append(np.concatenate(train_pred_images_V, axis=2))
+                        train_gt_images = np.concatenate(train_gt_images, axis=0)
+                        train_pred_images = np.concatenate(train_pred_images, axis=0)
+                        imageio.mimwrite(f'{opt.workspace}/train_gt_images_{epoch}_{i}.mp4', train_gt_images, fps=8)
+                        imageio.mimwrite(f'{opt.workspace}/train_pred_images_{epoch}_{i}.mp4', train_pred_images, fps=8)
+                        
+
+                    elif '3d' in opt.data_mode:
+                        gt_images = data['images_output'].detach().cpu().numpy() # [B, V, 3, output_size, output_size]
+                        gt_images = gt_images.transpose(0, 3, 1, 4, 2).reshape(-1, gt_images.shape[1] * gt_images.shape[3], 3) # [B*output_size, V*output_size, 3]
+                        kiui.write_image(f'{opt.workspace}/train_gt_images_{epoch}_{i}.jpg', gt_images)
+
+                        pred_images = out['images_pred'].detach().cpu().numpy() # [B, V, 3, output_size, output_size]
+                        pred_images = pred_images.transpose(0, 3, 1, 4, 2).reshape(-1, pred_images.shape[1] * pred_images.shape[3], 3)
+                        kiui.write_image(f'{opt.workspace}/train_pred_images_{epoch}_{i}.jpg', pred_images)
+                    else:
+                        raise NotImplementedError
+                    
+
+        total_loss = accelerator.gather_for_metrics(total_loss).mean()
+        total_psnr = accelerator.gather_for_metrics(total_psnr).mean()
+        if accelerator.is_main_process:
+            total_loss /= len(train_dataloader)
+            total_psnr /= len(train_dataloader)
+            accelerator.print(f"[train] epoch: {epoch} loss: {total_loss.item():.6f} psnr: {total_psnr.item():.4f}")
+        
+        # checkpoint
+        accelerator.wait_for_everyone()
+        accelerator.save_model(model, opt.workspace)
+        accelerator.save_model(model, os.path.join(opt.workspace, 'backup'))
+        if accelerator.is_main_process:
+            torch.save(optimizer.state_dict(), os.path.join(opt.workspace, 'optimizer.pth'))
+            torch.save(scheduler.state_dict(), os.path.join(opt.workspace, 'scheduler.pth'))
+            with open(f'{opt.workspace}/metadata.json', 'w') as f:
+                json.dump({'epoch': epoch}, f)
+
+            torch.save(optimizer.state_dict(), os.path.join(opt.workspace, 'backup', 'optimizer.pth'))
+            torch.save(scheduler.state_dict(), os.path.join(opt.workspace, 'backup', 'scheduler.pth'))
+            with open(f'{opt.workspace}/backup/metadata.json', 'w') as f:
+                json.dump({'epoch': epoch}, f)
+
+
+        # eval
+        with torch.no_grad():
+            model.eval()
+            total_psnr = 0
+            for i, data in enumerate(test_dataloader):
+
+                out = model(data)
+    
+                psnr = out['psnr']
+                total_psnr += psnr.detach()
+                
+                # save some images
+                if accelerator.is_main_process:
+                    if '4d' in opt.data_mode:
+                        B, T, V = opt.batch_size, opt.num_frames, opt.num_views
+
+                        gt_images = data['images_output'].reshape(-1, T, V, *data['images_output'].shape[2:]).detach() # [B, V, 3, output_size, output_size]
+                        pred_images = out['images_pred'].reshape(-1, T, V, *out['images_pred'].shape[2:]).detach() # [B, V, 3, output_size, output_size]
+
+                        eval_gt_images = []
+                        eval_pred_images = []
+                        for t in range(T):
+                            eval_gt_images_V = []
+                            eval_pred_images_V = []
+                            for v in range(V):
+                                eval_gt_images_V.append((gt_images[:, t, v].permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+                                eval_pred_images_V.append((pred_images[:, t, v].permute(0,2,3,1).contiguous().float().cpu().numpy() * 255).astype(np.uint8))
+                            eval_gt_images.append(np.concatenate(eval_gt_images_V, axis=2))
+                            eval_pred_images.append(np.concatenate(eval_pred_images_V, axis=2))
+                        eval_gt_images = np.concatenate(eval_gt_images, axis=0)
+                        eval_pred_images = np.concatenate(eval_pred_images, axis=0)
+                        imageio.mimwrite(f'{opt.workspace}/eval_gt_images_{epoch}_{i}.mp4', eval_gt_images, fps=8)
+                        imageio.mimwrite(f'{opt.workspace}/eval_pred_images_{epoch}_{i}.mp4', eval_pred_images, fps=8)
+                        
+                    elif '3d' in opt.data_mode:
+                        gt_images = data['images_output'].detach().cpu().numpy() # [B, V, 3, output_size, output_size]
+                        gt_images = gt_images.transpose(0, 3, 1, 4, 2).reshape(-1, gt_images.shape[1] * gt_images.shape[3], 3) # [B*output_size, V*output_size, 3]
+                        kiui.write_image(f'{opt.workspace}/eval_gt_images_{epoch}_{i}.jpg', gt_images)
+
+                        pred_images = out['images_pred'].detach().cpu().numpy() # [B, V, 3, output_size, output_size]
+                        pred_images = pred_images.transpose(0, 3, 1, 4, 2).reshape(-1, pred_images.shape[1] * pred_images.shape[3], 3)
+                        kiui.write_image(f'{opt.workspace}/eval_pred_images_{epoch}_{i}.jpg', pred_images)
+                    else:
+                        raise NotImplementedError
+
+            torch.cuda.empty_cache()
+
+            total_psnr = accelerator.gather_for_metrics(total_psnr).mean()
+            if accelerator.is_main_process:
+                total_psnr /= len(test_dataloader)
+                accelerator.print(f"[eval] epoch: {epoch} psnr: {psnr:.4f}")
+
+
+
+if __name__ == "__main__":
+    main()

mvdream/mv_unet.py

@@ -0,0 +1,1005 @@
+import math
+import numpy as np
+from inspect import isfunction
+from typing import Optional, Any, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.models.modeling_utils import ModelMixin
+
+# require xformers!
+import xformers
+import xformers.ops
+
+from kiui.cam import orbit_camera
+
+def get_camera(
+    num_frames, elevation=0, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False,
+):
+    angle_gap = azimuth_span / num_frames
+    cameras = []
+    for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
+        
+        pose = orbit_camera(elevation, azimuth, radius=1) # [4, 4]
+
+        # opengl to blender
+        if blender_coord:
+            pose[2] *= -1
+            pose[[1, 2]] = pose[[2, 1]]
+
+        cameras.append(pose.flatten())
+
+    if extra_view:
+        cameras.append(np.zeros_like(cameras[0]))
+
+    return torch.from_numpy(np.stack(cameras, axis=0)).float() # [num_frames, 16]
+
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    if not repeat_only:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32)
+            / half
+        ).to(device=timesteps.device)
+        args = timesteps[:, None] * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+    else:
+        embedding = repeat(timesteps, "b -> b d", d=dim)
+    # import pdb; pdb.set_trace()
+    return embedding
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def default(val, d):
+    if val is not None:
+        return val
+    return d() if isfunction(d) else d
+
+
+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)
+
+
+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,
+            ip_dim=0,
+            ip_weight=1,
+        ):
+        super().__init__()
+        
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.heads = heads
+        self.dim_head = dim_head
+
+        self.ip_dim = ip_dim
+        self.ip_weight = ip_weight
+
+        if self.ip_dim > 0:
+            self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
+            self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
+
+        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.attention_op: Optional[Any] = None
+
+    def forward(self, x, context=None):
+        q = self.to_q(x)
+        context = default(context, x)
+
+        if self.ip_dim > 0:
+            # context： [B, 77 + 16(ip), 1024]
+            token_len = context.shape[1]
+            context_ip = context[:, -self.ip_dim :, :]
+            k_ip = self.to_k_ip(context_ip)
+            v_ip = self.to_v_ip(context_ip)
+            context = context[:, : (token_len - self.ip_dim), :]
+
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        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),
+        )
+
+        # actually compute the attention, what we cannot get enough of
+        out = xformers.ops.memory_efficient_attention(
+            q, k, v, attn_bias=None, op=self.attention_op
+        )
+
+        if self.ip_dim > 0:
+            k_ip, v_ip = 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(),
+                (k_ip, v_ip),
+            )
+            # actually compute the attention, what we cannot get enough of
+            out_ip = xformers.ops.memory_efficient_attention(
+                q, k_ip, v_ip, attn_bias=None, op=self.attention_op
+            )
+            out = out + self.ip_weight * out_ip
+
+        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)
+        )
+        return self.to_out(out)
+
+
+class BasicTransformerBlock3D(nn.Module):
+    
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        context_dim,
+        dropout=0.0,
+        gated_ff=True,
+        ip_dim=0,
+        ip_weight=1,
+    ):
+        super().__init__()
+
+        self.attn1 = MemoryEfficientCrossAttention(
+            query_dim=dim,
+            context_dim=None, # self-attention
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+        )
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = MemoryEfficientCrossAttention(
+            query_dim=dim,
+            context_dim=context_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            # ip only applies to cross-attention
+            ip_dim=ip_dim,
+            ip_weight=ip_weight,
+        ) 
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+
+    def forward(self, x, context=None, num_frames=1):
+        x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()
+        x = self.attn1(self.norm1(x), context=None) + x
+        x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()
+        x = self.attn2(self.norm2(x), context=context) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class SpatialTransformer3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        n_heads,
+        d_head,
+        context_dim, # cross attention input dim
+        depth=1,
+        dropout=0.0,
+        ip_dim=0,
+        ip_weight=1,
+    ):
+        super().__init__()
+
+        if not isinstance(context_dim, list):
+            context_dim = [context_dim]
+
+        self.in_channels = in_channels
+
+        inner_dim = n_heads * d_head
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock3D(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    context_dim=context_dim[d],
+                    dropout=dropout,
+                    ip_dim=ip_dim,
+                    ip_weight=ip_weight,
+                )
+                for d in range(depth)
+            ]
+        )
+        
+        self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+        
+
+    def forward(self, x, context=None, num_frames=1):
+        # 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
+        x = self.norm(x)
+        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+        x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            x = block(x, context=context[i], num_frames=num_frames)
+        x = self.proj_out(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        
+        return x + x_in
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+        q, k, v = map(
+            lambda t: t.reshape(b, t.shape[1], self.heads, -1)
+            .transpose(1, 2)
+            .reshape(b, self.heads, t.shape[1], -1)
+            .contiguous(),
+            (q, k, v),
+        )
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1)  # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class Resampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+    ):
+        super().__init__()
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
+        self.proj_in = nn.Linear(embedding_dim, dim)
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        nn.Sequential(
+                            nn.LayerNorm(dim),
+                            nn.Linear(dim, dim * ff_mult, bias=False),
+                            nn.GELU(),
+                            nn.Linear(dim * ff_mult, dim, bias=False),
+                        )
+                    ]
+                )
+            )
+
+    def forward(self, x):
+        latents = self.latents.repeat(x.size(0), 1, 1)
+        x = self.proj_in(x)
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+
+        latents = self.proj_out(latents)
+        return self.norm_out(latents)
+
+
+class CondSequential(nn.Sequential):
+    """
+    A sequential module that passes timestep embeddings to the children that
+    support it as an extra input.
+    """
+
+    def forward(self, x, emb, context=None, num_frames=1):
+        for layer in self:
+            if isinstance(layer, ResBlock):
+                x = layer(x, emb)
+            elif isinstance(layer, SpatialTransformer3D):
+                x = layer(x, context, num_frames=num_frames)
+            else:
+                x = layer(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 upsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = conv_nd(
+                dims, self.channels, self.out_channels, 3, padding=padding
+            )
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        if self.dims == 3:
+            x = F.interpolate(
+                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            )
+        else:
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 downsampling occurs in the inner-two dimensions.
+    """
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        if use_conv:
+            self.op = conv_nd(
+                dims,
+                self.channels,
+                self.out_channels,
+                3,
+                stride=stride,
+                padding=padding,
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(nn.Module):
+    """
+    A residual block that can optionally change the number of channels.
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_scale_shift_norm = use_scale_shift_norm
+
+        self.in_layers = nn.Sequential(
+            nn.GroupNorm(32, channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            nn.GroupNorm(32, self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(
+                dims, channels, self.out_channels, 3, padding=1
+            )
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = torch.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class MultiViewUNetModel(ModelMixin, ConfigMixin):
+    """
+    The full multi-view UNet model with attention, timestep embedding and camera embedding.
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param num_res_blocks: number of residual blocks per downsample.
+    :param attention_resolutions: a collection of downsample rates at which
+        attention will take place. May be a set, list, or tuple.
+        For example, if this contains 4, then at 4x downsampling, attention
+        will be used.
+    :param dropout: the dropout probability.
+    :param channel_mult: channel multiplier for each level of the UNet.
+    :param conv_resample: if True, use learned convolutions for upsampling and
+        downsampling.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param num_classes: if specified (as an int), then this model will be
+        class-conditional with `num_classes` classes.
+    :param num_heads: the number of attention heads in each attention layer.
+    :param num_heads_channels: if specified, ignore num_heads and instead use
+                               a fixed channel width per attention head.
+    :param num_heads_upsample: works with num_heads to set a different number
+                               of heads for upsampling. Deprecated.
+    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+    :param resblock_updown: use residual blocks for up/downsampling.
+    :param use_new_attention_order: use a different attention pattern for potentially
+                                    increased efficiency.
+    :param camera_dim: dimensionality of camera input.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        num_heads=-1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        transformer_depth=1,
+        context_dim=None,
+        n_embed=None,
+        num_attention_blocks=None,
+        adm_in_channels=None,
+        camera_dim=None,
+        ip_dim=0, # imagedream uses ip_dim > 0
+        ip_weight=1.0,
+        **kwargs,
+    ):
+        super().__init__()
+        assert context_dim is not None
+        
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        if num_heads == -1:
+            assert (
+                num_head_channels != -1
+            ), "Either num_heads or num_head_channels has to be set"
+
+        if num_head_channels == -1:
+            assert (
+                num_heads != -1
+            ), "Either num_heads or num_head_channels has to be set"
+
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        if isinstance(num_res_blocks, int):
+            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+        else:
+            if len(num_res_blocks) != len(channel_mult):
+                raise ValueError(
+                    "provide num_res_blocks either as an int (globally constant) or "
+                    "as a list/tuple (per-level) with the same length as channel_mult"
+                )
+            self.num_res_blocks = num_res_blocks
+        
+        if num_attention_blocks is not None:
+            assert len(num_attention_blocks) == len(self.num_res_blocks)
+            assert all(
+                map(
+                    lambda i: self.num_res_blocks[i] >= num_attention_blocks[i],
+                    range(len(num_attention_blocks)),
+                )
+            )
+            print(
+                f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
+                f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
+                f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
+                f"attention will still not be set."
+            )
+
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.predict_codebook_ids = n_embed is not None
+
+        self.ip_dim = ip_dim
+        self.ip_weight = ip_weight
+
+        if self.ip_dim > 0:
+            self.image_embed = Resampler(
+                dim=context_dim,
+                depth=4,
+                dim_head=64,
+                heads=12,
+                num_queries=ip_dim,  # num token
+                embedding_dim=1280,
+                output_dim=context_dim,
+                ff_mult=4,
+            )
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            nn.Linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, time_embed_dim),
+        )
+
+        if camera_dim is not None:
+            time_embed_dim = model_channels * 4
+            self.camera_embed = nn.Sequential(
+                nn.Linear(camera_dim, time_embed_dim),
+                nn.SiLU(),
+                nn.Linear(time_embed_dim, time_embed_dim),
+            )
+
+        if self.num_classes is not None:
+            if isinstance(self.num_classes, int):
+                self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)
+            elif self.num_classes == "continuous":
+                # print("setting up linear c_adm embedding layer")
+                self.label_emb = nn.Linear(1, time_embed_dim)
+            elif self.num_classes == "sequential":
+                assert adm_in_channels is not None
+                self.label_emb = nn.Sequential(
+                    nn.Sequential(
+                        nn.Linear(adm_in_channels, time_embed_dim),
+                        nn.SiLU(),
+                        nn.Linear(time_embed_dim, time_embed_dim),
+                    )
+                )
+            else:
+                raise ValueError()
+
+        self.input_blocks = nn.ModuleList(
+            [
+                CondSequential(
+                    conv_nd(dims, in_channels, model_channels, 3, padding=1)
+                )
+            ]
+        )
+        self._feature_size = model_channels
+        input_block_chans = [model_channels]
+        ch = model_channels
+        ds = 1
+        for level, mult in enumerate(channel_mult):
+            for nr in range(self.num_res_blocks[level]):
+                layers: List[Any] = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=mult * model_channels,
+                        dims=dims,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+
+                    if num_attention_blocks is None or nr < num_attention_blocks[level]:
+                        layers.append(
+                            SpatialTransformer3D(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                context_dim=context_dim,
+                                depth=transformer_depth,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
+                            )
+                        )
+                self.input_blocks.append(CondSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    CondSequential(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch, conv_resample, dims=dims, out_channels=out_ch
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        if num_head_channels == -1:
+            dim_head = ch // num_heads
+        else:
+            num_heads = ch // num_head_channels
+            dim_head = num_head_channels
+        
+        self.middle_block = CondSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            SpatialTransformer3D(
+                ch,
+                num_heads,
+                dim_head,
+                context_dim=context_dim,
+                depth=transformer_depth,
+                ip_dim=self.ip_dim,
+                ip_weight=self.ip_weight,
+            ), 
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(self.num_res_blocks[level] + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=model_channels * mult,
+                        dims=dims,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = model_channels * mult
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+
+                    if num_attention_blocks is None or i < num_attention_blocks[level]:
+                        layers.append(
+                            SpatialTransformer3D(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                context_dim=context_dim,
+                                depth=transformer_depth,
+                                ip_dim=self.ip_dim,
+                                ip_weight=self.ip_weight,
+                            )
+                        )
+                if level and i == self.num_res_blocks[level]:
+                    out_ch = ch
+                    layers.append(
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    ds //= 2
+                self.output_blocks.append(CondSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            nn.GroupNorm(32, ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+        )
+        if self.predict_codebook_ids:
+            self.id_predictor = nn.Sequential(
+                nn.GroupNorm(32, ch),
+                conv_nd(dims, model_channels, n_embed, 1),
+                # nn.LogSoftmax(dim=1)  # change to cross_entropy and produce non-normalized logits
+            )
+
+    def forward(
+        self,
+        x,
+        timesteps=None,
+        context=None,
+        y=None,
+        camera=None,
+        num_frames=1,
+        ip=None,
+        ip_img=None,
+        **kwargs,
+    ):
+        """
+        Apply the model to an input batch.
+        :param x: an [(N x F) x C x ...] Tensor of inputs. F is the number of frames (views).
+        :param timesteps: a 1-D batch of timesteps.
+        :param context: conditioning plugged in via crossattn
+        :param y: an [N] Tensor of labels, if class-conditional.
+        :param num_frames: a integer indicating number of frames for tensor reshaping.
+        :return: an [(N x F) x C x ...] Tensor of outputs. F is the number of frames (views).
+        """
+        assert (
+            x.shape[0] % num_frames == 0
+        ), "input batch size must be dividable by num_frames!"
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+
+        hs = []
+
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
+
+        emb = self.time_embed(t_emb)
+
+        if self.num_classes is not None:
+            assert y is not None
+            assert y.shape[0] == x.shape[0]
+            emb = emb + self.label_emb(y)
+
+        # Add camera embeddings
+        if camera is not None:
+            emb = emb + self.camera_embed(camera)
+        
+        # imagedream variant
+        if self.ip_dim > 0:
+            x[(num_frames - 1) :: num_frames, :, :, :] = ip_img # place at [4, 9]
+            ip_emb = self.image_embed(ip)
+            context = torch.cat((context, ip_emb), 1)
+
+        h = x
+        for module in self.input_blocks:
+            h = module(h, emb, context, num_frames=num_frames)
+            hs.append(h)
+        h = self.middle_block(h, emb, context, num_frames=num_frames)
+        for module in self.output_blocks:
+            h = torch.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context, num_frames=num_frames)
+        h = h.type(x.dtype)
+        if self.predict_codebook_ids:
+            return self.id_predictor(h)
+        else:
+            return self.out(h)

mvdream/pipeline_mvdream.py

@@ -0,0 +1,559 @@
+import torch
+import torch.nn.functional as F
+import inspect
+import numpy as np
+from typing import Callable, List, Optional, Union
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
+from diffusers import AutoencoderKL, DiffusionPipeline
+from diffusers.utils import (
+    deprecate,
+    is_accelerate_available,
+    is_accelerate_version,
+    logging,
+)
+from diffusers.configuration_utils import FrozenDict
+from diffusers.schedulers import DDIMScheduler
+from diffusers.utils.torch_utils import randn_tensor
+
+from mvdream.mv_unet import MultiViewUNetModel, get_camera
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class MVDreamPipeline(DiffusionPipeline):
+
+    _optional_components = ["feature_extractor", "image_encoder"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        unet: MultiViewUNetModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModel,
+        scheduler: DDIMScheduler,
+        # imagedream variant
+        feature_extractor: CLIPImageProcessor,
+        image_encoder: CLIPVisionModel,
+        requires_safety_checker: bool = False,
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:  # type: ignore
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "  # type: ignore
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate(
+                "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
+            )
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:  # type: ignore
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate(
+                "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
+            )
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            unet=unet,
+            scheduler=scheduler,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            feature_extractor=feature_extractor,
+            image_encoder=image_encoder,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
+        steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
+        several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def enable_sequential_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
+        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
+        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
+        Note that offloading happens on a submodule basis. Memory savings are higher than with
+        `enable_model_cpu_offload`, but performance is lower.
+        """
+        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
+            from accelerate import cpu_offload
+        else:
+            raise ImportError(
+                "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
+            )
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
+            cpu_offload(cpu_offloaded_model, device)
+
+    def enable_model_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
+        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
+        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
+        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
+        """
+        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
+            from accelerate import cpu_offload_with_hook
+        else:
+            raise ImportError(
+                "`enable_model_offload` requires `accelerate v0.17.0` or higher."
+            )
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        hook = None
+        for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
+            _, hook = cpu_offload_with_hook(
+                cpu_offloaded_model, device, prev_module_hook=hook
+            )
+
+        # We'll offload the last model manually.
+        self.final_offload_hook = hook
+
+    @property
+    def _execution_device(self):
+        r"""
+        Returns the device on which the pipeline's models will be executed. After calling
+        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
+        hooks.
+        """
+        if not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance: bool,
+        negative_prompt=None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+             prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
+                Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+        """
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            raise ValueError(
+                f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
+            )
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(
+            prompt, padding="longest", return_tensors="pt"
+        ).input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer.batch_decode(
+                untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+            )
+
+        if (
+            hasattr(self.text_encoder.config, "use_attention_mask")
+            and self.text_encoder.config.use_attention_mask
+        ):
+            attention_mask = text_inputs.attention_mask.to(device)
+        else:
+            attention_mask = None
+
+        prompt_embeds = self.text_encoder(
+            text_input_ids.to(device),
+            attention_mask=attention_mask,
+        )
+        prompt_embeds = prompt_embeds[0]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(
+            bs_embed * num_images_per_prompt, seq_len, -1
+        )
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if (
+                hasattr(self.text_encoder.config, "use_attention_mask")
+                and self.text_encoder.config.use_attention_mask
+            ):
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(
+                dtype=self.text_encoder.dtype, device=device
+            )
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(
+                1, num_images_per_prompt, 1
+            )
+            negative_prompt_embeds = negative_prompt_embeds.view(
+                batch_size * num_images_per_prompt, seq_len, -1
+            )
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(
+                shape, generator=generator, device=device, dtype=dtype
+            )
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def encode_image(self, image, device, num_images_per_prompt):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if image.dtype == np.float32:
+            image = (image * 255).astype(np.uint8)
+            
+        image = self.feature_extractor(image, return_tensors="pt").pixel_values
+        image = image.to(device=device, dtype=dtype)
+        
+        image_embeds = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+        image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+
+        return torch.zeros_like(image_embeds), image_embeds
+
+    def encode_image_latents(self, image, device, num_images_per_prompt):
+        
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device) # [1, 3, H, W]
+        image = 2 * image - 1
+        image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
+        image = image.to(dtype=dtype)
+
+        posterior = self.vae.encode(image).latent_dist
+        latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
+        latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
+
+        return torch.zeros_like(latents), latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str = "",
+        image: Optional[np.ndarray] = None,
+        height: int = 256,
+        width: int = 256,
+        elevation: float = 0,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.0,
+        negative_prompt: str = "",
+        num_images_per_prompt: int = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        output_type: Optional[str] = "numpy", # pil, numpy, latents
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        num_frames: int = 4,
+        device=torch.device("cuda:0"),
+    ):
+        self.unet = self.unet.to(device=device)
+        self.vae = self.vae.to(device=device)
+        self.text_encoder = self.text_encoder.to(device=device)
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # imagedream variant
+        if image is not None:
+            assert isinstance(image, np.ndarray) and image.dtype == np.float32
+            self.image_encoder = self.image_encoder.to(device=device)
+            image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt)
+            image_latents_neg, image_latents_pos = self.encode_image_latents(image, device, num_images_per_prompt)
+            
+        _prompt_embeds = self._encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+        )  # type: ignore
+        prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
+
+        # Prepare latent variables
+        actual_num_frames = num_frames if image is None else num_frames + 1
+        latents: torch.Tensor = self.prepare_latents(
+            actual_num_frames * num_images_per_prompt,
+            4,
+            height,
+            width,
+            prompt_embeds_pos.dtype,
+            device,
+            generator,
+            None,
+        )
+
+        if image is not None:
+            camera = get_camera(num_frames, elevation=elevation, extra_view=True).to(dtype=latents.dtype, device=device)
+        else:
+            camera = get_camera(num_frames, elevation=elevation, extra_view=False).to(dtype=latents.dtype, device=device)
+        camera = camera.repeat_interleave(num_images_per_prompt, dim=0)
+
+        # Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                multiplier = 2 if do_classifier_free_guidance else 1
+                latent_model_input = torch.cat([latents] * multiplier)
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                unet_inputs = {
+                    'x': latent_model_input,
+                    'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),
+                    'context': torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames),
+                    'num_frames': actual_num_frames,
+                    'camera': torch.cat([camera] * multiplier),
+                }
+
+                if image is not None:
+                    unet_inputs['ip'] = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
+                    unet_inputs['ip_img'] = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
+                
+                # predict the noise residual
+                noise_pred = self.unet.forward(**unet_inputs)
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents: torch.Tensor = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+                )[0]
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)  # type: ignore
+
+        # Post-processing
+        if output_type == "latent":
+            image = latents
+        elif output_type == "pil":
+            image = self.decode_latents(latents)
+            image = self.numpy_to_pil(image)
+        else: # numpy
+            image = self.decode_latents(latents)
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        return image

readme.md

@@ -0,0 +1,77 @@
+
+## L4GM: Large 4D Gaussian Reconstruction Model
+<p align="center">
+    <img src="assets/teaser.jpg">
+</p>
+
+[**Paper**](https://arxiv.org/abs/2406.10324) | [**Project Page**](https://research.nvidia.com/labs/toronto-ai/l4gm/) | [**Model Weights**](https://huggingface.co/jiawei011/L4GM)
+
+We present L4GM, the first 4D Large Reconstruction Model that produces animated objects from a single-view video input -- in a single feed-forward pass that takes only a second.
+
+---
+
+### Install
+```bash
+conda env create -f environment.yml
+conda activate l4gm
+```
+
+### Inference
+Download pretrained [L4GM model](https://huggingface.co/jiawei011/L4GM/blob/main/recon.safetensors) and [4D interpolation model](https://huggingface.co/jiawei011/L4GM/blob/main/interp.safetensors) to `pretrained/recon.safetensors` and `pretrained/interp.safetensors` respectively.
+
+Select an input video. Remove its background and crop it to 256x256 with third-party tools. We provide some processed examples in the `data_test` folder.
+
+1. Generate 3D by:
+```sh
+python infer_3d.py big --workspace results --resume pretrained/recon.safetensors --num_frames 1 --test_path data_test/otter-on-surfboard_fg.mp4
+```
+
+2. Generate 4D by:
+```sh
+python infer_4d.py big --workspace results --resume pretrained/recon.safetensors --interpresume pretrained/interp.safetensors --num_frames 16 --test_path data_test/otter-on-surfboard_fg.mp4
+```
+
+### Training
+Render Objaverse with Blender scripts in the `blender_scripts` folder first.
+
+Download pretrained [LGM](https://huggingface.co/ashawkey/LGM/blob/main/model_fixrot.safetensors) to `pretrained/model_fixrot.safetensors`.
+
+L4GM model training:
+```sh
+accelerate launch \
+    --config_file acc_configs/gpu8.yaml \
+    main.py big \
+    --workspace workspace_recon \
+    --resume pretrained/model_fixrot.safetensors \
+    --data_mode 4d \
+    --num_epochs 200 \
+    --prob_cam_jitter 0 \
+    --datalist data_train/datalist_8fps.txt \
+```
+Our released checkpoint uses `--num_epochs 500`.
+
+4D Interpolation model training:
+```sh
+accelerate launch \
+    --config_file acc_configs/gpu8.yaml \
+    main.py big \
+    --workspace workspace_interp \
+    --resume workspace_recon/model.safetensors \
+    --data_mode 4d_interp \
+    --num_frames 4 \
+    --num_epochs 200 \
+    --prob_cam_jitter 0 \
+    --prob_grid_distortion 0 \
+    --datalist data_train/datalist_24fps.txt \
+```
+
+### Citation
+```bib
+@inproceedings{ren2024l4gm,
+    title={L4GM: Large 4D Gaussian Reconstruction Model}, 
+    author={Jiawei Ren and Kevin Xie and Ashkan Mirzaei and Hanxue Liang and Xiaohui Zeng and Karsten Kreis and Ziwei Liu and Antonio Torralba and Sanja Fidler and Seung Wook Kim and Huan Ling},
+    booktitle={Proceedings of Neural Information Processing Systems(NeurIPS)},
+    month = {Dec},
+    year={2024}
+}
+```

requirements.txt

@@ -0,0 +1,20 @@
+tyro
+accelerate==0.28.0
+imageio
+imageio-ffmpeg
+lpips
+Pillow
+safetensors
+scikit-image
+scikit-learn
+scipy
+tqdm
+kiui >= 0.2.3
+roma
+plyfile
+
+# mvdream
+diffusers==0.27.2
+huggingface_hub==0.23.5
+transformers
+