gpu / cpu

demo/img_gen.py

@@ -7,8 +7,11 @@ from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
 model_id = "stabilityai/stable-diffusion-2-1"
 
-device = torch.device('cuda')
-dtype = torch.float16
+device = torch.device('cpu')
+dtype = torch.float32
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+    dtype = torch.float16
 
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

demo/relighting_gen.py

@@ -7,8 +7,11 @@ from diffusers.utils import get_class_from_dynamic_module
 
 from tqdm import tqdm
 
-device = torch.device('cuda')
-dtype = torch.float16
+device = torch.device('cpu')
+dtype = torch.float32
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+    dtype = torch.float16
 
 NeuralTextureControlNetModel = get_class_from_dynamic_module(
     "dilightnet/model_helpers",