Fix device placement

marble.py

@@ -1,5 +1,5 @@
 import os
-from typing import Dict
+from typing import Dict, Optional
 
 import numpy as np
 import torch
@@ -31,10 +31,18 @@ def get_session():
     return _session_cache
 
 
+def get_device():
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
 def setup_control_mlps(
-    features: int = 1024, device: str = "cuda", dtype: torch.dtype = torch.float16
+    features: int = 1024,
+    device: Optional[str] = None,
+    dtype: torch.dtype = torch.float16,
 ) -> Dict[str, torch.nn.Module]:
     ret = {}
+    if device is None:
+        device = get_device()
     for mlp in CONTROL_MLPS:
         ret[mlp] = setup_control_mlp(mlp, features, device, dtype)
     return ret
@@ -43,9 +51,12 @@ def setup_control_mlps(
 def setup_control_mlp(
     material_parameter: str,
     features: int = 1024,
-    device: str = "cuda",
+    device: Optional[str] = None,
     dtype: torch.dtype = torch.float16,
 ):
+    if device is None:
+        device = get_device()
+
     net = control_mlp(features)
     net.load_state_dict(
         torch.load(os.path.join(file_dir, f"model_weights/{material_parameter}.pt"))
@@ -95,9 +106,12 @@ def download_ip_adapter():
 
 
 def setup_pipeline(
-    device: str = "cuda",
+    device: Optional[str] = None,
     dtype: torch.dtype = torch.float16,
 ):
+    if device is None:
+        device = get_device()
+
     download_ip_adapter()
 
     cur_block = ("up", 0, 1)
@@ -135,7 +149,10 @@ def setup_pipeline(
     )
 
 
-def get_dpt_model(device: str = "cuda", dtype: torch.dtype = torch.float16):
+def get_dpt_model(device: Optional[str] = None, dtype: torch.dtype = torch.float16):
+    if device is None:
+        device = get_device()
+
     image_processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
     model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas")
     model.to(device, dtype=dtype)
@@ -144,9 +161,12 @@ def get_dpt_model(device: str = "cuda", dtype: torch.dtype = torch.float16):
 
 
 def run_dpt_depth(
-    image: Image.Image, model, processor, device: str = "cuda"
+    image: Image.Image, model, processor, device: Optional[str] = None
 ) -> Image.Image:
     """Run DPT depth estimation on an image."""
+    if device is None:
+        device = get_device()
+
     # Prepare image
     inputs = processor(images=image, return_tensors="pt").to(device, dtype=model.dtype)