Check for CUDA

depth_estimator.py

@@ -3,17 +3,20 @@ import numpy as np
 from PIL import Image
 from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 
+device = None
 depth_estimator = None
 feature_extractor = None
 
 
 def init():
-    global depth_estimator, feature_extractor
+    global device, depth_estimator, feature_extractor
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
     print("Initializing depth estimator...")
 
     depth_estimator = DPTForDepthEstimation.from_pretrained(
-        "Intel/dpt-hybrid-midas").to("cuda")
+        "Intel/dpt-hybrid-midas").to(device)
     feature_extractor = DPTFeatureExtractor.from_pretrained(
         "Intel/dpt-hybrid-midas")
 
@@ -22,9 +25,9 @@ def get_depth_map(image):
     original_size = image.size
 
     image = feature_extractor(
-        images=image, return_tensors="pt").pixel_values.to("cuda")
+        images=image, return_tensors="pt").pixel_values.to(device)
 
-    with torch.no_grad(), torch.autocast("cuda"):
+    with torch.no_grad(), torch.autocast(device):
         depth_map = depth_estimator(image).predicted_depth
 
     depth_map = torch.nn.functional.interpolate(

pipeline.py

@@ -1,11 +1,14 @@
 import torch
 from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL, UniPCMultistepScheduler
 
+device = None
 pipe = None
 
 
 def init():
-    global pipe
+    global device, pipe
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
     print("Initializing depth ControlNet...")
 
@@ -13,14 +16,14 @@ def init():
         "diffusers/controlnet-depth-sdxl-1.0",
         use_safetensors=True,
         torch_dtype=torch.float16
-    ).to("cuda")
+    ).to(device)
 
     print("Initializing autoencoder...")
 
     vae = AutoencoderKL.from_pretrained(
         "madebyollin/sdxl-vae-fp16-fix",
         torch_dtype=torch.float16,
-    ).to("cuda")
+    ).to(device)
 
     print("Initializing SDXL pipeline...")
 
@@ -32,7 +35,7 @@ def init():
         use_safetensors=True,
         torch_dtype=torch.float16
         # low_cpu_mem_usage=True
-    ).to("cuda")
+    ).to(device)
 
     pipe.enable_model_cpu_offload()
     # speed up diffusion process with faster scheduler and memory optimization

upscaler.py

@@ -2,6 +2,7 @@
 import os
 import requests
 
+import torch
 import cv2
 import numpy as np
 from PIL import Image
@@ -14,6 +15,8 @@ upsampler = None
 def init():
     global upsampler
 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
     print("Initializing upscaler...")
 
     if not os.path.exists("weights"):
@@ -25,7 +28,7 @@ def init():
     model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
                     num_block=23, num_grow_ch=32, scale=2)
     upsampler = RealESRGANer(
-        scale=2, model_path="weights/RealESRGAN_x2plus.pth", model=model, device="cuda")
+        scale=2, model_path="weights/RealESRGAN_x2plus.pth", model=model, device=device)
 
 
 def upscale(image):