Update device handling across multiple modules to support automatic selection of CUDA or CPU based on availability. This change enhances compatibility and performance on systems with or without GPU support, ensuring consistent behavior in model loading and data processing.

app.py

@@ -161,8 +161,8 @@ args = parser.parse_args()
 crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
 specs = json.load(open("configs/specs_objaverse_total.json"))
 model = CRM(specs)
-model.load_state_dict(torch.load(crm_path, map_location="cpu"), strict=False)
-model = model.to("cpu")
+model.load_state_dict(torch.load(crm_path, map_location="cuda" if torch.cuda.is_available() else "cpu"), strict=False)
+model = model.to("cuda" if torch.cuda.is_available() else "cpu")
 
 stage1_config = OmegaConf.load(args.stage1_config).config
 stage2_config = OmegaConf.load(args.stage2_config).config
@@ -177,12 +177,13 @@ pixel_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pt
 stage1_model_config.resume = pixel_path
 stage2_model_config.resume = xyz_path
 
+device = args.device if hasattr(args, 'device') else ("cuda" if torch.cuda.is_available() else "cpu")
 pipeline = TwoStagePipeline(
     stage1_model_config,
     stage2_model_config,
     stage1_sampler_config,
     stage2_sampler_config,
-    device=args.device,
+    device=device,
     dtype=torch.float32
 )
 

imagedream/ldm/modules/encoders/modules.py

@@ -306,10 +306,11 @@ class FrozenCLIPT5Encoder(AbstractEncoder):
         self,
         clip_version="openai/clip-vit-large-patch14",
         t5_version="google/t5-v1_1-xl",
-        device="cpu",
+        device=None,
         clip_max_length=77,
         t5_max_length=77,
     ):
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         super().__init__()
         self.clip_encoder = FrozenCLIPEmbedder(
             clip_version, device, max_length=clip_max_length

inference.py

@@ -153,8 +153,8 @@ def generate3d(model, rgb, ccm, device):
 
     from kiui.mesh_utils import clean_mesh
     verts, faces = clean_mesh(data_config['verts'].squeeze().cpu().numpy().astype(np.float32), data_config['faces'].squeeze().cpu().numpy().astype(np.int32), repair = False, remesh=True, remesh_size=0.005, remesh_iters=1)
-    data_config['verts'] = torch.from_numpy(verts).to('cpu').contiguous()
-    data_config['faces'] = torch.from_numpy(faces).to('cpu').contiguous()
+    data_config['verts'] = torch.from_numpy(verts).to(device).contiguous()
+    data_config['faces'] = torch.from_numpy(faces).to(device).contiguous()
 
     start_time = time.time()
     with torch.no_grad():

libs/sample.py

@@ -11,18 +11,19 @@ class ImageDreamDiffusion:
     def __init__(
         self,
         model,
-        device,
-        dtype,
-        mode,
-        num_frames,
-        camera_views,
-        ref_position,
+        device=None,
+        dtype=None,
+        mode=None,
+        num_frames=None,
+        camera_views=None,
+        ref_position=None,
         random_background=False,
         offset_noise=False,
         resize_rate=1,
         image_size=256,
         seed=1234,
     ) -> None:
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         assert mode in ["pixel", "local"]
         size = image_size
         self.seed = seed
@@ -204,11 +205,11 @@ class ImageDreamDiffusionStage2:
     def __init__(
         self,
         model,
-        device,
-        dtype,
-        num_frames,
-        camera_views,
-        ref_position,
+        device=None,
+        dtype=None,
+        num_frames=None,
+        camera_views=None,
+        ref_position=None,
         random_background=False,
         offset_noise=False,
         resize_rate=1,
@@ -216,6 +217,7 @@ class ImageDreamDiffusionStage2:
         image_size=256,
         seed=1234,
     ) -> None:
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         assert mode in ["pixel", "local"]
 
         size = image_size

pipelines.py

@@ -25,14 +25,15 @@ class TwoStagePipeline(object):
         - the first stage was condition on single pixel image, gererate multi-view pixel image, based on the v2pp config
         - the second stage was condition on multiview pixel image generated by the first stage, generate the final image, based on the stage2-test config
         """
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.resize_rate = resize_rate
 
         self.stage1_model = instantiate_from_config(OmegaConf.load(stage1_model_config.config).model)
-        self.stage1_model.load_state_dict(torch.load(stage1_model_config.resume, map_location="cpu"), strict=False)
+        self.stage1_model.load_state_dict(torch.load(stage1_model_config.resume, map_location=device), strict=False)
         self.stage1_model = self.stage1_model.to(device).to(dtype)
 
         self.stage2_model = instantiate_from_config(OmegaConf.load(stage2_model_config.config).model)
-        sd = torch.load(stage2_model_config.resume, map_location="cpu")
+        sd = torch.load(stage2_model_config.resume, map_location=device)
         self.stage2_model.load_state_dict(sd, strict=False)
         self.stage2_model = self.stage2_model.to(device).to(dtype)
 

util/flexicubes.py

@@ -64,8 +64,8 @@ class FlexiCubes:
             The scale of weights in FlexiCubes. Should be between 0 and 1.
     """
 
-    def __init__(self, device="cpu", qef_reg_scale=1e-3, weight_scale=0.99):
-
+    def __init__(self, device=None, qef_reg_scale=1e-3, weight_scale=0.99):
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         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,

util/flexicubes_geometry.py

@@ -31,8 +31,9 @@ def get_center_boundary_index(grid_res, device):
 ###############################################################################
 class FlexiCubesGeometry(object):
     def __init__(
-            self, grid_res=64, scale=2.0, device='cpu', renderer=None,
+            self, grid_res=64, scale=2.0, device=None, renderer=None,
             render_type='neural_render', args=None):
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         super(FlexiCubesGeometry, self).__init__()
         self.grid_res = grid_res
         self.device = device