Update fn.py

fn.py

@@ -14,18 +14,23 @@ def load_model(_model = None, _vae = None, loras = []):
 
     _model = _model or 'cagliostrolab/animagine-xl-3.0'
 
+    if torch.cuda.is_available():
+        torch_dtype = torch.float16
+    else:
+        torch_dtype = torch.float32
+
     if _vae:
         # "stabilityai/sdxl-vae"
-        vae = AutoencoderKL.from_pretrained(_vae, torch_dtype=torch.float16)
+        vae = AutoencoderKL.from_pretrained(_vae, torch_dtype=torch_dtype)
         pipe = AutoPipelineForText2Image.from_pretrained(
             _model,
-            torch_dtype=torch.float16,
+            torch_dtype=torch_dtype,
             vae=vae,
             )
     else:
         pipe = AutoPipelineForText2Image.from_pretrained(
             _model,
-            torch_dtype=torch.float16,
+            torch_dtype=torch_dtype,
             )
 
     # DPM++ 2M Karras