Enhance Mesh class documentation by adding missing line breaks in docstrings for improved readability. Update device handling in FlexiCubes and FlexiCubesGeometry classes to default to 'cuda', ensuring consistent device usage across the application. Refactor ImageDreamDiffusion class to assert mode validity and streamline camera matrix pre-computation.

libs/sample.py

@@ -11,32 +11,36 @@ class ImageDreamDiffusion:
     def __init__(
         self,
         model,
-        device=None,
-        dtype=None,
-        mode=None,
-        num_frames=None,
-        camera_views=None,
-        ref_position=None,
+        device,
+        dtype,
+        mode,
+        num_frames,
+        camera_views,
+        ref_position,
         random_background=False,
         offset_noise=False,
         resize_rate=1,
         image_size=256,
         seed=1234,
     ) -> None:
-        device = torch.device("cuda")
-        self.device = device
+        assert mode in ["pixel", "local"]
+        size = image_size
         self.seed = seed
         batch_size = max(4, num_frames)
+
         neg_texts = "uniform low no texture ugly, boring, bad anatomy, blurry, pixelated,  obscure, unnatural colors, poor lighting, dull, and unclear."
         uc = model.get_learned_conditioning([neg_texts]).to(device)
         sampler = DDIMSampler(model)
+
+        # pre-compute camera matrices
         camera = [get_camera_for_index(i).squeeze() for i in camera_views]
-        camera[ref_position] = torch.zeros_like(camera[ref_position])
+        camera[ref_position] = torch.zeros_like(camera[ref_position])  # set ref camera to zero
         camera = torch.stack(camera)
         camera = camera.repeat(batch_size // num_frames, 1).to(device)
+
         self.image_transform = T.Compose(
             [
-                T.Resize((image_size, image_size)),
+                T.Resize((size, size)),
                 T.ToTensor(),
                 T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
             ]
@@ -47,7 +51,8 @@ class ImageDreamDiffusion:
         self.random_background = random_background
         self.resize_rate = resize_rate
         self.num_frames = num_frames
-        self.size = image_size
+        self.size = size
+        self.device = device
         self.batch_size = batch_size
         self.model = model
         self.sampler = sampler
@@ -372,4 +377,4 @@ class ImageDreamDiffusionStage2:
             )
             images.append(img)
         set_seed()  # unset random and numpy seed
-        return images
+        return images

mesh.py

@@ -10,6 +10,7 @@ from kiui.typing import *
 class Mesh:
     """
     A torch-native trimesh class, with support for ``ply/obj/glb`` formats.
+
     Note:
         This class only supports one mesh with a single texture image (an albedo texture and a metallic-roughness texture).
     """
@@ -27,6 +28,7 @@ class Mesh:
         device: Optional[torch.device] = None,
     ):
         """Init a mesh directly using all attributes.
+
         Args:
             v (Optional[Tensor]): vertices, float [N, 3]. Defaults to None.
             f (Optional[Tensor]): faces, int [M, 3]. Defaults to None.
@@ -60,6 +62,7 @@ class Mesh:
     @classmethod
     def load(cls, path, resize=True, clean=False, renormal=True, retex=False, bound=0.9, front_dir='+z', **kwargs):
         """load mesh from path.
+
         Args:
             path (str): path to mesh file, supports ply, obj, glb.
             clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False.
@@ -73,6 +76,7 @@ class Mesh:
         Note:
             a ``device`` keyword argument can be provided to specify the torch device. 
             If it's not provided, we will try to use ``'cuda'`` as the device if it's available.
+
         Returns:
             Mesh: the loaded Mesh object.
         """
@@ -136,6 +140,7 @@ class Mesh:
     @classmethod
     def load_obj(cls, path, albedo_path=None, device=None):
         """load an ``obj`` mesh.
+
         Args:
             path (str): path to mesh.
             albedo_path (str, optional): path to the albedo texture image, will overwrite the existing texture path if specified in mtl. Defaults to None.
@@ -144,6 +149,7 @@ class Mesh:
         Note: 
             We will try to read `mtl` path from `obj`, else we assume the file name is the same as `obj` but with `mtl` extension.
             The `usemtl` statement is ignored, and we only use the last material path in `mtl` file.
+
         Returns:
             Mesh: the loaded Mesh object.
         """
@@ -307,13 +313,17 @@ class Mesh:
     @classmethod
     def load_trimesh(cls, path, device=None):
         """load a mesh using ``trimesh.load()``.
+
         Can load various formats like ``glb`` and serves as a fallback.
+
         Note:
             We will try to merge all meshes if the glb contains more than one, 
             but **this may cause the texture to lose**, since we only support one texture image!
+
         Args:
             path (str): path to the mesh file.
             device (torch.device, optional): torch device. Defaults to None.
+
         Returns:
             Mesh: the loaded Mesh object.
         """
@@ -413,8 +423,10 @@ class Mesh:
     # sample surface (using trimesh)
     def sample_surface(self, count: int):
         """sample points on the surface of the mesh.
+
         Args:
             count (int): number of points to sample.
+
         Returns:
             torch.Tensor: the sampled points, float [count, 3].
         """
@@ -426,6 +438,7 @@ class Mesh:
     # aabb
     def aabb(self):
         """get the axis-aligned bounding box of the mesh.
+
         Returns:
             Tuple[torch.Tensor]: the min xyz and max xyz of the mesh.
         """
@@ -435,6 +448,7 @@ class Mesh:
     @torch.no_grad()
     def auto_size(self, bound=0.9):
         """auto resize the mesh.
+
         Args:
             bound (float, optional): resizing into ``[-bound, bound]^3``. Defaults to 0.9.
         """
@@ -470,6 +484,7 @@ class Mesh:
 
     def auto_uv(self, cache_path=None, vmap=True):
         """auto calculate the uv coordinates.
+
         Args:
             cache_path (str, optional): path to save/load the uv cache as a npz file, this can avoid calculating uv every time when loading the same mesh, which is time-consuming. Defaults to None.
             vmap (bool, optional): remap vertices based on uv coordinates, so each v correspond to a unique vt (necessary for formats like gltf). 
@@ -508,6 +523,7 @@ class Mesh:
     
     def align_v_to_vt(self, vmapping=None):
         """ remap v/f and vn/fn to vt/ft.
+
         Args:
             vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None.
         """
@@ -526,8 +542,10 @@ class Mesh:
 
     def to(self, device):
         """move all tensor attributes to device.
+
         Args:
             device (torch.device): target device.
+
         Returns:
             Mesh: self.
         """
@@ -540,6 +558,7 @@ class Mesh:
     
     def write(self, path):
         """write the mesh to a path.
+
         Args:
             path (str): path to write, supports ply, obj and glb.
         """
@@ -554,6 +573,7 @@ class Mesh:
     
     def write_ply(self, path):
         """write the mesh in ply format. Only for geometry!
+
         Args:
             path (str): path to write.
         """
@@ -571,6 +591,7 @@ class Mesh:
     def write_glb(self, path):
         """write the mesh in glb/gltf format.
           This will create a scene with a single mesh.
+
         Args:
             path (str): path to write.
         """
@@ -757,6 +778,7 @@ class Mesh:
 
     def write_obj(self, path):
         """write the mesh in obj format. Will also write the texture and mtl files.
+
         Args:
             path (str): path to write.
         """
@@ -819,4 +841,4 @@ class Mesh:
             metallicRoughness = self.metallicRoughness.detach().cpu().numpy()
             metallicRoughness = (metallicRoughness * 255).astype(np.uint8)
             cv2.imwrite(metallic_path, metallicRoughness[..., 2])
-            cv2.imwrite(roughness_path, metallicRoughness[..., 1])
+            cv2.imwrite(roughness_path, metallicRoughness[..., 1])

util/flexicubes.py

@@ -25,7 +25,7 @@ class FlexiCubes:
     PyTorch tensors on the specified device.
 
     Attributes:
-        device (str): Specifies the computational device (default is "cpu").
+        device (str): Specifies the computational device (default is "cuda").
         dmc_table (torch.Tensor): Dual Marching Cubes (DMC) table that encodes the edges 
             associated with each dual vertex in 256 Marching Cubes (MC) configurations.
         num_vd_table (torch.Tensor): Table holding the number of dual vertices in each of 
@@ -64,8 +64,8 @@ class FlexiCubes:
             The scale of weights in FlexiCubes. Should be between 0 and 1.
     """
 
-    def __init__(self, device=None, qef_reg_scale=1e-3, weight_scale=0.99):
-        device = torch.device("cuda")
+    def __init__(self, device="cuda", qef_reg_scale=1e-3, weight_scale=0.99):
+
         self.device = device
         self.dmc_table = torch.tensor(dmc_table, dtype=torch.long, device=device, requires_grad=False)
         self.num_vd_table = torch.tensor(num_vd_table,
@@ -576,4 +576,4 @@ class FlexiCubes:
 
         tets = torch.cat([tets_surface, tets_inside])
         vertices = torch.cat([vertices, inside_verts, inside_cubes_center])
-        return vertices, tets
+        return vertices, tets

util/flexicubes_geometry.py

@@ -31,9 +31,8 @@ def get_center_boundary_index(grid_res, device):
 ###############################################################################
 class FlexiCubesGeometry(object):
     def __init__(
-            self, grid_res=64, scale=2.0, device=None, renderer=None,
+            self, grid_res=64, scale=2.0, device='cuda', renderer=None,
             render_type='neural_render', args=None):
-        device = torch.device("cuda")
         super(FlexiCubesGeometry, self).__init__()
         self.grid_res = grid_res
         self.device = device
@@ -114,4 +113,4 @@ class FlexiCubesGeometry(object):
             value = [v[k] for v in all_render_output]
             return_value[k] = value
             # We can do concatenation outside of the render
-        return return_value
+        return return_value