half only if cuda is available

climategan_wrapper.py

@@ -15,6 +15,8 @@ from skimage.transform import resize
 
 from climategan.trainer import Trainer
 
+CUDA = torch.cuda.is_available()
+
 
 def concat_events(output_dict, events, i=None, axis=1):
     """
@@ -136,7 +138,8 @@ class ClimateGAN:
             inference=True,
             new_exp=None,
         )
-        self.trainer.G.half()
+        if CUDA:
+            self.trainer.G.half()
 
     def _setup_stable_diffusion(self):
         """
@@ -150,8 +153,8 @@ class ClimateGAN:
         try:
             self.sdip_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
                 "runwayml/stable-diffusion-inpainting",
-                revision="fp16",
-                torch_dtype=torch.float16,
+                revision="fp16" if CUDA else "main",
+                torch_dtype=torch.float16 if CUDA else torch.float32,
                 safety_checker=None,
                 use_auth_token=os.environ.get("HF_AUTH_TOKEN"),
             ).to(self.trainer.device)
@@ -308,7 +311,7 @@ class ClimateGAN:
             images,
             numpy=True,
             bin_value=0.5,
-            half=True,
+            half=CUDA,
             ignore_event=ignore_event,
             return_masks=True,
         )