Update main.py

main.py

@@ -67,9 +67,13 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "default_secret")
 if is_gpu:
     # Uncomment the following line if you want to enable CUDA launch blocking
     os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+    torch_dtype=torch.float16
+    variant="fp16"
 else:
     # Uncomment the following line if you want to use CPU instead of GPU
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    torch_dtype=torch.float32
+    variant="fp32"
 
 
 # Get the current directory
@@ -82,17 +86,17 @@ cache_path = os.path.join(current_dir, "cache")
 if not SSD_1B:
 
     unet = UNet2DConditionModel.from_pretrained(
-        "latent-consistency/lcm-sdxl",
-        torch_dtype=torch.float16,
-        variant="fp16",
+        "latent-consistency/lcm-sdxl",        
+        torch_dtype=torch_dtype,
+        variant=variant,
         cache_dir=cache_path,
     )
     pipe = DiffusionPipeline.from_pretrained(
         # "stabilityai/stable-diffusion-xl-base-1.0",
         "stabilityai/sdxl-turbo",
         unet=unet,
-        torch_dtype=torch.float16,
-        variant="fp16",
+        torch_dtype=torch_dtype,
+        variant=variant,
         cache_dir=cache_path,
     )