remove_type_annotator

app.py

@@ -1,11 +1,11 @@
 import os
+
 import subprocess
 import sys
 try:
     import spaces
 except:
     pass
-os.environ["PYDANTIC_STRICT_TYPE_CHECKING"] = "0"
 
 # Check if setup has been run
 setup_marker = ".setup_complete"
@@ -23,7 +23,6 @@ if not os.path.exists(setup_marker):
 
 import torch
 import gradio as gr
-from typing import Tuple, List, Dict, Any, Optional
 from collections import deque
 from diffusers import StableDiffusionPipeline
 
@@ -58,7 +57,7 @@ def initialize_pipeline():
     return PIPELINE
 
 @spaces.GPU
-def generate_3d_mesh(prompt: str) -> Tuple[Optional[str], Optional[str]]:
+def generate_3d_mesh(prompt):
     """Generate 3D mesh from text prompt"""
     global PIPELINE, OBJ_FILE_QUEUE
     

example.py

@@ -17,8 +17,8 @@ from triplaneturbo_executable import TriplaneTurboTextTo3DPipeline, TriplaneTurb
 
 # Initialize configuration and parameters
 prompt = "a beautiful girl"
-output_dir = "examples/output"
-adapter_name_or_path = "/home/user/app/pretrained/triplane_turbo_sd_v1.pth"
+output_dir = "output"
+adapter_name_or_path = "pretrained/triplane_turbo_sd_v1.pth"
 num_results_per_prompt = 1
 seed = 42
 device = "cuda"

setup.sh

@@ -17,12 +17,12 @@ pip install --force-reinstall -v "numpy==1.25.2"
 # cd ..
 # cd ..
 
-
+echo "Setup completed successfully!"
 
 echo "Installing other requirements..."
 pip install -r requirements.txt
 
-
+# 从您的Hugging Face仓库下载并安装预编译的DISO wheel
 echo "Installing pre-compiled DISO wheel package..."
 huggingface-cli download --resume-download ZhiyuanthePony/TriplaneTurbo \
   --include "diso-0.1.4-*.whl" \
@@ -30,12 +30,3 @@ huggingface-cli download --resume-download ZhiyuanthePony/TriplaneTurbo \
   --local-dir-use-symlinks False
 
 pip install ./diso_package/diso-0.1.4-*.whl
-echo "Setup completed successfully!"
-
-echo "Installing compatible dependency versions..."
-pip uninstall -y pydantic
-pip install pydantic==1.10.8  # Install compatible older version
-
-# Ensure Gradio and other dependencies are installed correctly
-pip install "gradio>=4.0.0,<5.0.0"
-pip install "fastapi<0.103.0"  # Ensure compatible FastAPI version

triplaneturbo_executable/extern/sd_dual_triplane_modules.py

@@ -2,7 +2,6 @@ import re
 import torch
 import torch.nn as nn
 from dataclasses import dataclass
-from typing import Optional, Union, Tuple
 
 from diffusers.models.attention_processor import Attention
 from diffusers import (
@@ -39,9 +38,9 @@ class LoRALinearLayerwBias(nn.Module):
         in_features: int,
         out_features: int,
         rank: int = 4,
-        network_alpha: Optional[float] = None,
-        device: Optional[Union[torch.device, str]] = None,
-        dtype: Optional[torch.dtype] = None,
+        network_alpha=None,
+        device=None,
+        dtype=None,
         with_bias: bool = False
     ):
         super().__init__()
@@ -105,10 +104,10 @@ class TriplaneLoRAConv2dLayer(nn.Module):
         in_features: int,
         out_features: int,
         rank: int = 4,
-        kernel_size: Union[int, Tuple[int, int]] = (1, 1),
-        stride: Union[int, Tuple[int, int]] = (1, 1),
-        padding: Union[int, Tuple[int, int], str] = 0,
-        network_alpha: Optional[float] = None,
+        kernel_size = (1, 1),
+        stride = (1, 1),
+        padding = 0,
+        network_alpha = None,
         with_bias: bool = False,
         locon_type: str = "hexa_v1", #hexa_v2, vanilla_v1, vanilla_v2
     ):
@@ -220,7 +219,7 @@ class TriplaneSelfAttentionLoRAAttnProcessor(nn.Module):
         self,
         hidden_size: int,
         rank: int = 4,
-        network_alpha: Optional[float] = None,
+        network_alpha=None,
         with_bias: bool = False,
         lora_type: str = "hexa_v1", # vanilla, 
     ):
@@ -492,7 +491,7 @@ class TriplaneCrossAttentionLoRAAttnProcessor(nn.Module):
         hidden_size: int,
         cross_attention_dim: int,
         rank: int = 4,
-        network_alpha: Optional[float] = None,
+        network_alpha = None,
         with_bias: bool = False,
         lora_type: str = "hexa_v1", # vanilla,
     ):
@@ -713,7 +712,7 @@ class OneStepTriplaneDualStableDiffusion(nn.Module):
     """
     def __init__(
         self, 
-        config: Union[dict, GeneratorConfig],
+        config,
         vae: AutoencoderKL,
         unet: UNet2DConditionModel,
     ):

triplaneturbo_executable/models/networks.py

@@ -3,7 +3,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 from ..utils.general_utils import config_to_primitive
 from dataclasses import dataclass
-from typing import Optional, Literal
 
 def get_activation(name):
     if name is None:
@@ -21,7 +20,7 @@ def get_activation(name):
 
 
 class VanillaMLP(nn.Module):
-    def __init__(self, dim_in: int, dim_out: int, config: dict):
+    def __init__(self, dim_in, dim_out, config):
         super().__init__()
         # Convert dict to MLPConfig if needed
         if isinstance(config, dict):
@@ -70,7 +69,7 @@ class MLPConfig:
     n_neurons: int = 64
     n_hidden_layers: int = 2
 
-def get_mlp(input_dim: int, output_dim: int, config: dict) -> nn.Module:
+def get_mlp(input_dim, output_dim, config):
     """Create MLP network based on config"""
     # Convert dict to MLPConfig
     if isinstance(config, dict):

triplaneturbo_executable/pipelines/triplaneturbo_text_to_3d.py

@@ -4,7 +4,6 @@ import json
 from tqdm import tqdm
 
 import torch
-from typing import *
 from dataclasses import dataclass, field
 from diffusers import StableDiffusionPipeline
 
@@ -21,11 +20,6 @@ class TriplaneTurboTextTo3DPipelineConfig:
     # Basic pipeline settings
     base_model_name_or_path: str = "stabilityai/stable-diffusion-2-1-base"
 
-    num_inference_steps: int = 4
-    num_results_per_prompt: int = 1
-    latent_channels: int = 4
-    latent_height: int = 64
-    latent_width: int = 64
     
     # Training/sampling settings
     num_steps_sampling: int = 4
@@ -72,7 +66,7 @@ class TriplaneTurboTextTo3DPipelineConfig:
     color_activation: str = "sigmoid-mipnerf"
 
     @classmethod
-    def from_pretrained(cls, pretrained_path: str) -> "TriplaneTurboTextTo3DPipelineConfig":
+    def from_pretrained(cls, pretrained_path):
         """Load config from pretrained path"""
         config_path = os.path.join(pretrained_path, "config.json")
         if os.path.exists(config_path):
@@ -91,11 +85,11 @@ class TriplaneTurboTextTo3DPipeline(Pipeline):
 
     def __init__(
         self,
-        geometry: StableDiffusionTriplaneDualAttention,
-        material: Callable,
-        base_pipeline: StableDiffusionPipeline,
-        sample_scheduler: Callable,
-        isosurface_helper: Callable,
+        geometry,
+        material,
+        base_pipeline,
+        sample_scheduler,
+        isosurface_helper,
         **kwargs,
     ):
         super().__init__()
@@ -116,7 +110,7 @@ class TriplaneTurboTextTo3DPipeline(Pipeline):
     @classmethod
     def from_pretrained(
         cls,
-        pretrained_model_name_or_path: str,
+        pretrained_model_name_or_path,
         **kwargs,
     ):
         """
@@ -197,10 +191,10 @@ class TriplaneTurboTextTo3DPipeline(Pipeline):
 
     def encode_prompt(
         self, 
-        prompt: Union[str, List[str]],
-        device: str,
-        num_results_per_prompt: int = 1,
-    ) -> torch.FloatTensor:
+        prompt,
+        device,
+        num_results_per_prompt = 1,
+    ):
         """
         Encodes the prompt into text encoder hidden states.
         
@@ -227,14 +221,13 @@ class TriplaneTurboTextTo3DPipeline(Pipeline):
     @torch.no_grad()
     def __call__(
         self,
-        prompt: Union[str, List[str]],
-        num_inference_steps: int = 4,
-        num_results_per_prompt: int = 1,
-        generator: Optional[torch.Generator] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        return_dict: bool = True,
-        colorize: bool = True,
-        **kwargs,
+        prompt,
+        num_results_per_prompt=1,
+        generator=None,
+        device=None,
+        return_dict=True,
+        num_inference_steps=4,
+        colorize = True,
     ):
         # Implementation similar to Zero123Pipeline
         # Reference code from: https://github.com/zero123/zero123-diffusers
@@ -251,15 +244,18 @@ class TriplaneTurboTextTo3DPipeline(Pipeline):
         # Get the device from the first available module
 
         # Generate latents if not provided
-        if latents is None:
-            latents = torch.randn(
-                (batch_size * 6, 4, 32, 32), # hard-coded for now
-                generator=generator,
-                device=self.device,
-            )
+        if device is None:
+            device = self.device
+        if generator is None:
+            generator = torch.Generator(device=device)
+        latents = torch.randn(
+            (batch_size * 6, 4, 32, 32), # hard-coded for now
+            generator=generator,
+            device=device,
+        )
 
         # Process text prompt through geometry module
-        text_embed, _ = self.encode_prompt(prompt, self.device, num_results_per_prompt)
+        text_embed, _ = self.encode_prompt(prompt, device, num_results_per_prompt)
         
         # Run diffusion process
         # Set up timesteps for sampling
@@ -282,7 +278,7 @@ class TriplaneTurboTextTo3DPipeline(Pipeline):
                 pred = self.geometry.denoise(
                     noisy_input=noisy_latent_input,
                     text_embed=text_embed,
-                    timestep=t.to(self.device),
+                    timestep=t.to(device),
                 )
 
                 # Update latents
@@ -311,20 +307,19 @@ class TriplaneTurboTextTo3DPipeline(Pipeline):
                     activation=self.material,
                 )
 
-        # decide output type based on return_dict
-        if return_dict:
-            return {
-                "space_cache": space_cache,
-                "latents": latents,
-                "mesh": mesh_list,
-            }
-        else:
-            return mesh_list
+            if return_dict:
+                return {
+                    "space_cache": space_cache,
+                    "latents": latents,
+                    "mesh": mesh_list,
+                }
+            else:
+                return mesh_list
 
     def _set_timesteps(
         self,
         scheduler,
-        num_steps: int,
+        num_steps,
     ):
         """Set up timesteps for sampling.
         

triplaneturbo_executable/utils/general_utils.py

@@ -2,17 +2,28 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
-from typing import *
-from jaxtyping import Float
-from omegaconf import OmegaConf
+import numpy as np
+from dataclasses import asdict, is_dataclass
+import gc
 
-def config_to_primitive(config, resolve: bool = True) -> Any:
-    return OmegaConf.to_container(config, resolve=resolve)
+def config_to_primitive(config):
+    """Convert a dataclass config to a dictionary recursively."""
+    if is_dataclass(config):
+        config_dict = asdict(config)
+        return {k: config_to_primitive(v) for k, v in config_dict.items()}
+    elif isinstance(config, dict):
+        return {k: config_to_primitive(v) for k, v in config.items()}
+    elif isinstance(config, list):
+        return [config_to_primitive(v) for v in config]
+    elif isinstance(config, tuple):
+        return tuple(config_to_primitive(v) for v in config)
+    else:
+        return config
 
 def scale_tensor(
-    dat: Float[Tensor, "... D"], 
-    inp_scale: Union[Tuple[float, float], Float[Tensor, "2 D"]], 
-    tgt_scale: Union[Tuple[float, float], Float[Tensor, "2 D"]]
+    dat, 
+    inp_scale, 
+    tgt_scale
 ):
     if inp_scale is None:
         inp_scale = (0, 1)
@@ -25,8 +36,8 @@ def scale_tensor(
     return dat
 
 def contract_to_unisphere_custom(
-    x: Float[Tensor, "... 3"], bbox: Float[Tensor, "2 3"], unbounded: bool = False
-) -> Float[Tensor, "... 3"]:
+    x, bbox, unbounded = False
+):
     if unbounded:
         x = scale_tensor(x, bbox, (-1, 1))
         x = x * 2 - 1  # aabb is at [-1, 1]
@@ -81,7 +92,7 @@ def project_onto_planes(planes, coordinates):
     projections = torch.bmm(coordinates, inv_planes)
     return projections[..., :2]
 
-def sample_from_planes(plane_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=2, interpolate_feat: Optional[str] = 'None'):
+def sample_from_planes(plane_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=2, interpolate_feat = None):
     assert padding_mode == 'zeros'
     N, n_planes, C, H, W = plane_features.shape
     _, M, _ = coordinates.shape
@@ -101,4 +112,10 @@ def sample_from_planes(plane_features, coordinates, mode='bilinear', padding_mod
         output_features = output_features.permute(0, 3, 2, 1).reshape(N, n_planes, M, C)
         output_features = output_features.permute(0, 2, 1, 3).reshape(N, M, n_planes*C)        
 
-    return output_features.contiguous()
+    return output_features.contiguous()
+
+def cleanup():
+    """Cleanup torch memory."""
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()

triplaneturbo_executable/utils/mesh.py

@@ -1,77 +1,54 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-
-from typing import Any, Dict, Optional, Union
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from jaxtyping import Float, Integer
-from torch import Tensor
+import trimesh
 
 def dot(x, y):
     return torch.sum(x * y, -1, keepdim=True)
 
 class Mesh:
     def __init__(
-        self, v_pos: Float[Tensor, "Nv 3"], t_pos_idx: Integer[Tensor, "Nf 3"], **kwargs
-    ) -> None:
-        self.v_pos: Float[Tensor, "Nv 3"] = v_pos
-        self.t_pos_idx: Integer[Tensor, "Nf 3"] = t_pos_idx
-        self._v_nrm: Optional[Float[Tensor, "Nv 3"]] = None
-        self._v_tng: Optional[Float[Tensor, "Nv 3"]] = None
-        self._v_tex: Optional[Float[Tensor, "Nt 3"]] = None
-        self._t_tex_idx: Optional[Float[Tensor, "Nf 3"]] = None
-        self._v_rgb: Optional[Float[Tensor, "Nv 3"]] = None
-        self._edges: Optional[Integer[Tensor, "Ne 2"]] = None
-        self.extras: Dict[str, Any] = {}
-        for k, v in kwargs.items():
-            self.add_extra(k, v)
+        self, v_pos, t_pos_idx, material=None
+    ):
+        self.v_pos = v_pos
+        self.t_pos_idx = t_pos_idx
+        self.material = material
+        self._v_nrm = None
+        self._v_tng = None
+        self._v_tex = None
+        self._t_tex_idx = None
+        self._v_rgb = None
+        self._edges = None
+        self.extras = {}
 
     def add_extra(self, k, v) -> None:
         self.extras[k] = v
 
-    def remove_outlier(self, outlier_n_faces_threshold: Union[int, float]):
-
-        # use trimesh to first split the mesh into connected components
-        # then remove the components with less than n_face_threshold faces
-        import trimesh
-
-        # construct a trimesh object
-        mesh = trimesh.Trimesh(
-            vertices=self.v_pos.detach().cpu().numpy(),
-            faces=self.t_pos_idx.detach().cpu().numpy(),
+    def remove_outlier(self, n_face_threshold=5):
+        """Remove outlier components with fewer faces than threshold."""
+        # Convert to trimesh
+        trimesh_mesh = self.as_trimesh()
+        
+        # Split into connected components
+        components = trimesh_mesh.split(only_watertight=False)
+        
+        # Filter components with few faces
+        valid_components = [c for c in components if len(c.faces) > n_face_threshold]
+        
+        if len(valid_components) == 0:
+            # If no valid components, return the original mesh
+            return self
+        
+        # Combine valid components
+        combined = trimesh.util.concatenate(valid_components)
+        
+        # Convert back to our Mesh format
+        new_mesh = Mesh(
+            torch.tensor(combined.vertices, dtype=self.v_pos.dtype, device=self.v_pos.device),
+            torch.tensor(combined.faces, dtype=self.t_pos_idx.dtype, device=self.t_pos_idx.device)
         )
-
-        # split the mesh into connected components
-        components = mesh.split(only_watertight=False)
-
-
-        n_faces_threshold: int
-        if isinstance(outlier_n_faces_threshold, float):
-            # set the threshold to the number of faces in the largest component multiplied by outlier_n_faces_threshold
-            n_faces_threshold = int(
-                max([c.faces.shape[0] for c in components]) * outlier_n_faces_threshold
-            )
-        else:
-            # set the threshold directly to outlier_n_faces_threshold
-            n_faces_threshold = outlier_n_faces_threshold
-
-        # remove the components with less than n_face_threshold faces
-        components = [c for c in components if c.faces.shape[0] >= n_faces_threshold]
-
-        # merge the components
-        mesh = trimesh.util.concatenate(components)
-
-        # convert back to our mesh format
-        v_pos = torch.from_numpy(mesh.vertices).to(self.v_pos)
-        t_pos_idx = torch.from_numpy(mesh.faces).to(self.t_pos_idx)
-
-        clean_mesh = Mesh(v_pos, t_pos_idx)
-        # keep the extras unchanged
-
-        return clean_mesh
+        
+        return new_mesh
 
     @property
     def requires_grad(self):
@@ -245,8 +222,8 @@ class Mesh:
         edges = torch.unique(edges, dim=0)
         return edges
 
-    def normal_consistency(self) -> Float[Tensor, ""]:
-        edge_nrm: Float[Tensor, "Ne 2 3"] = self.v_nrm[self.edges]
+    def normal_consistency(self):
+        edge_nrm = self.v_nrm[self.edges]
         nc = (
             1.0 - torch.cosine_similarity(edge_nrm[:, 0], edge_nrm[:, 1], dim=-1)
         ).mean()
@@ -279,10 +256,45 @@ class Mesh:
         # correct diagonal
         return torch.sparse_coo_tensor(idx, values, (V, V)).coalesce()
 
-    def laplacian(self) -> Float[Tensor, ""]:
+    def laplacian(self):
         with torch.no_grad():
             L = self._laplacian_uniform()
         loss = L.mm(self.v_pos)
         loss = loss.norm(dim=1)
         loss = loss.mean()
         return loss
+
+    def to(self, device):
+        v_pos = self.v_pos.to(device)
+        t_pos_idx = self.t_pos_idx.to(device)
+        return Mesh(v_pos, t_pos_idx)
+    
+    def as_trimesh(self):
+        vertices = self.v_pos.detach().cpu().numpy()
+        faces = self.t_pos_idx.detach().cpu().numpy()
+        
+        mesh = trimesh.Trimesh(
+            vertices=vertices,
+            faces=faces,
+            process=False
+        )
+        
+        # Add texture if available
+        if hasattr(self, 'albedo_map') and self.albedo_map is not None:
+            # Create texture visuals
+            uv = self.v_tex.detach().cpu().numpy()
+            
+            # Create texture visuals
+            visual = trimesh.visual.texture.TextureVisuals(
+                uv=uv,
+                material=trimesh.visual.material.SimpleMaterial()
+            )
+            mesh.visual = visual
+            
+        return mesh
+
+def scale_tensor(x, input_range, target_range):
+    """Scale tensor from input_range to target_range."""
+    x_unit = (x - input_range[0]) / (input_range[1] - input_range[0])
+    x_scaled = x_unit * (target_range[1] - target_range[0]) + target_range[0]
+    return x_scaled

triplaneturbo_executable/utils/mesh_exporter.py

@@ -1,6 +1,3 @@
-from typing import Callable, Dict, List, Optional, Tuple, Any
-from jaxtyping import Float
-from torch import Tensor
 from dataclasses import dataclass
 
 import torch
@@ -16,36 +13,35 @@ from ..utils.general_utils import scale_tensor
 class ExporterOutput:
     save_name: str
     save_type: str
-    params: Dict[str, Any]
+    params: dict
 
 
 class IsosurfaceHelper(nn.Module):
-    points_range: Tuple[float, float] = (0, 1)
+    points_range = (0, 1)
 
     @property
-    def grid_vertices(self) -> Float[Tensor, "N 3"]:
+    def grid_vertices(self):
         raise NotImplementedError
     
 class DiffMarchingCubeHelper(IsosurfaceHelper):
     def __init__(
             self, 
-            resolution: int, 
-            point_range: Tuple[float, float] = (0, 1)
-        ) -> None:
+            resolution, 
+            point_range = (0, 1)
+        ):
         super().__init__()
         self.resolution = resolution
         self.points_range = point_range
 
         from diso import DiffMC
-        self.mc_func: Callable = DiffMC(dtype=torch.float32)
-        self._grid_vertices: Optional[Float[Tensor, "N3 3"]] = None
-        self._dummy: Float[Tensor, "..."]
+        self.mc_func = DiffMC(dtype=torch.float32)
+        self._grid_vertices = None
         self.register_buffer(
             "_dummy", torch.zeros(0, dtype=torch.float32), persistent=False
         )
 
     @property
-    def grid_vertices(self) -> Float[Tensor, "N3 3"]:
+    def grid_vertices(self):
         if self._grid_vertices is None:
             # keep the vertices on CPU so that we can support very large resolution
             x, y, z = (
@@ -62,10 +58,10 @@ class DiffMarchingCubeHelper(IsosurfaceHelper):
 
     def forward(
         self,
-        level: Float[Tensor, "N3 1"],
-        deformation: Optional[Float[Tensor, "N3 3"]] = None,
+        level,
+        deformation = None,
         isovalue=0.0,
-    ) -> Mesh:
+    ):
         level = level.view(self.resolution, self.resolution, self.resolution)
         if deformation is not None:
             deformation = deformation.view(self.resolution, self.resolution, self.resolution, 3)
@@ -76,17 +72,17 @@ class DiffMarchingCubeHelper(IsosurfaceHelper):
 
 
 def isosurface(
-        space_cache: Float[Tensor, "B ..."], 
-        forward_field: Callable,
-        isosurface_helper: Callable,
-    ) -> List[Mesh]:
+        space_cache,
+        forward_field,
+        isosurface_helper,
+    ):
 
     # the isosurface is dependent on the space cache
     # randomly detach isosurface method if it is differentiable
     # get the batchsize
     if torch.is_tensor(space_cache): #space cache
         batch_size = space_cache.shape[0]
-    elif isinstance(space_cache, Dict): #hyper net
+    elif isinstance(space_cache, dict): #hyper net
         # Dict[str, List[Float[Tensor, "B ..."]]]
         for key in space_cache.keys():
             batch_size = space_cache[key][0].shape[0]
@@ -141,11 +137,11 @@ def isosurface(
     return mesh_list
 
 def colorize_mesh(
-    space_cache: Any,
-    export_fn: Callable,
-    mesh_list: List[Mesh],
-    activation: Callable,
-) -> List[Mesh]:
+    space_cache,
+    export_fn,
+    mesh_list,
+    activation,
+):
     """Colorize the mesh using the geometry's export function and space cache.
     
     Args:
@@ -199,10 +195,10 @@ class MeshExporter(SaverMixin):
         return x
 
 def export_obj(
-        mesh: Mesh, 
-        save_path: str,
-        save_normal: bool = False,
-    ) -> List[str]:
+        mesh, 
+        save_path,
+        save_normal = False,
+    ):
     """
     Export mesh data to OBJ file format.
     

triplaneturbo_executable/utils/saving.py

@@ -13,22 +13,15 @@ import wandb
 from matplotlib import cm
 from matplotlib.colors import LinearSegmentedColormap
 from PIL import Image, ImageDraw
-# from pytorch_lightning.loggers import WandbLogger
 
-from ..utils.mesh import Mesh
-
-from typing import Dict, List, Optional, Union, Any
-from omegaconf import DictConfig
-from jaxtyping import Float
-from torch import Tensor
 
 import threading
 
 class SaverMixin:
-    _save_dir: Optional[str] = None
-    # _wandb_logger: Optional[WandbLogger] = None
+    _save_dir = None
+    # _wandb_logger = None
 
-    def set_save_dir(self, save_dir: str):
+    def set_save_dir(self, save_dir):
         self._save_dir = save_dir
 
     def get_save_dir(self):
@@ -58,17 +51,6 @@ class SaverMixin:
         os.makedirs(os.path.dirname(save_path), exist_ok=True)
         return save_path
 
-    # def create_loggers(self, cfg_loggers: DictConfig) -> None:
-    #     if "wandb" in cfg_loggers.keys() and cfg_loggers.wandb.enable:
-    #         self._wandb_logger = WandbLogger(
-    #             project=cfg_loggers.wandb.project, name=cfg_loggers.wandb.name
-    #         )
-
-    # def get_loggers(self) -> List:
-    #     if self._wandb_logger:
-    #         return [self._wandb_logger]
-    #     else:
-    #         return []
 
     DEFAULT_RGB_KWARGS = {"data_format": "HWC", "data_range": (0, 1)}
     DEFAULT_UV_KWARGS = {
@@ -119,8 +101,8 @@ class SaverMixin:
         img,
         data_format,
         data_range,
-        name: Optional[str] = None,
-        step: Optional[int] = None,
+        name=None,
+        step=None,
     ):
         img = self.get_rgb_image_(img, data_format, data_range)
         cv2.imwrite(filename, img)
@@ -138,8 +120,8 @@ class SaverMixin:
         img,
         data_format=DEFAULT_RGB_KWARGS["data_format"],
         data_range=DEFAULT_RGB_KWARGS["data_range"],
-        name: Optional[str] = None,
-        step: Optional[int] = None,
+        name=None,
+        step=None,
     ) -> str:
         save_path = self.get_save_path(filename)
         self._save_rgb_image(save_path, img, data_format, data_range, name, step)
@@ -231,8 +213,8 @@ class SaverMixin:
         img,
         data_range,
         cmap,
-        name: Optional[str] = None,
-        step: Optional[int] = None,
+        name=None,
+        step=None,
     ):
         img = self.get_grayscale_image_(img, data_range, cmap)
         cv2.imwrite(filename, img)
@@ -250,8 +232,8 @@ class SaverMixin:
         img,
         data_range=DEFAULT_GRAYSCALE_KWARGS["data_range"],
         cmap=DEFAULT_GRAYSCALE_KWARGS["cmap"],
-        name: Optional[str] = None,
-        step: Optional[int] = None,
+        name=None,
+        step=None,
     ) -> str:
         save_path = self.get_save_path(filename)
         self._save_grayscale_image(save_path, img, data_range, cmap, name, step)
@@ -308,9 +290,9 @@ class SaverMixin:
         filename,
         imgs,
         align=DEFAULT_GRID_KWARGS["align"],
-        name: Optional[str] = None,
-        step: Optional[int] = None,
-        texts: Optional[List[float]] = None,
+        name=None,
+        step=None,
+        texts=None,
     ):
         save_path = self.get_save_path(filename)
         img = self.get_image_grid_(imgs, align=align)
@@ -404,8 +386,8 @@ class SaverMixin:
     #     matcher,
     #     save_format="mp4",
     #     fps=30,
-    #     name: Optional[str] = None,
-    #     step: Optional[int] = None,
+    #     name=None,
+    #     step=None,
     # ) -> str:
     #     assert save_format in ["gif", "mp4"]
     #     if not filename.endswith(save_format):
@@ -442,9 +424,9 @@ class SaverMixin:
         matcher,
         save_format="mp4",
         fps=30,
-        name: Optional[str] = None,
-        step: Optional[int] = None,
-        multithreaded: bool = False
+        name=None,
+        step=None,
+        multithreaded=False
     ) -> str:
         assert save_format in ["gif", "mp4"]
         if not filename.endswith(save_format):
@@ -494,20 +476,19 @@ class SaverMixin:
 
     def save_obj(
         self,
-        filename: str,
-        mesh: Mesh,
-        save_mat: bool = False,
-        save_normal: bool = False,
-        save_uv: bool = False,
-        save_vertex_color: bool = False,
-        map_Kd: Optional[Float[Tensor, "H W 3"]] = None,
-        map_Ks: Optional[Float[Tensor, "H W 3"]] = None,
-        map_Bump: Optional[Float[Tensor, "H W 3"]] = None,
-        map_Pm: Optional[Float[Tensor, "H W 1"]] = None,
-        map_Pr: Optional[Float[Tensor, "H W 1"]] = None,
-        map_format: str = "jpg",
-    ) -> List[str]:
-
+        filename,
+        mesh,
+        save_mat=False,
+        save_normal=False,
+        save_uv=False,
+        save_vertex_color=False,
+        map_Kd=None,
+        map_Ks=None,
+        map_Bump=None,
+        map_Pm=None,
+        map_Pr=None,
+        map_format="jpg",
+    ):
         if not filename.endswith(".obj"):
             filename += ".obj"
         save_path = self.get_save_path(filename)
@@ -658,8 +639,8 @@ class SaverMixin:
         map_Pm=None,
         map_Pr=None,
         map_format="jpg",
-        step: Optional[int] = None,
-    ) -> List[str]:
+        step=None,
+    ):
         mtl_save_path = self.get_save_path(filename)
         save_paths = [mtl_save_path]
         mtl_str = f"newmtl {matname}\n"