Update device handling in TwoStagePipeline to default to 'cuda' and adjust model loading to use 'cpu' for compatibility. Ensure consistent device usage across stages and improve code clarity.

pipelines.py

@@ -16,7 +16,7 @@ class TwoStagePipeline(object):
         stage2_model_config,
         stage1_sampler_config,
         stage2_sampler_config,
-        device=None,
+        device="cuda",
         dtype=torch.float16,
         resize_rate=1,
     ) -> None:
@@ -25,15 +25,14 @@ 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 = torch.device("cuda")
         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=device), strict=False)
+        self.stage1_model.load_state_dict(torch.load(stage1_model_config.resume, map_location="cpu"), 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=device)
+        sd = torch.load(stage2_model_config.resume, map_location="cpu")
         self.stage2_model.load_state_dict(sd, strict=False)
         self.stage2_model = self.stage2_model.to(device).to(dtype)
 
@@ -168,4 +167,4 @@ if __name__ == "__main__":
     np_imgs = np.concatenate(stage1_images, 1)
     np_xyzs = np.concatenate(stage2_images, 1)
     Image.fromarray(np_imgs).save("pixel_images.png")
-    Image.fromarray(np_xyzs).save("xyz_images.png")
+    Image.fromarray(np_xyzs).save("xyz_images.png")