Parameter for CUDA

app.py

@@ -16,9 +16,10 @@ import torch
 max_64_bit_int = 2**63 - 1
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11e_sd15_ip2p", torch_dtype = torch.float32)
+floatType = torch.float16 if torch.cuda.is_available() else torch.float32
+controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11e_sd15_ip2p", torch_dtype = floatType)
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", safety_checker = None, controlnet = controlnet, torch_dtype = torch.float32
+    "runwayml/stable-diffusion-v1-5", safety_checker = None, controlnet = controlnet, torch_dtype = floatType
 )
 pipe = pipe.to(device)