Upload app.py

app.py

@@ -26,12 +26,13 @@ from transformers import CLIPImageProcessor
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"
 
 # Initialize both pipelines
-vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
-controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=torch.float16)
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=dtype)
+controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sd15_qrcode_monster", torch_dtype=dtype)
 
 # Initialize the safety checker conditionally
 SAFETY_CHECKER_ENABLED = os.environ.get("SAFETY_CHECKER", "0") == "1"
@@ -47,7 +48,7 @@ main_pipe = StableDiffusionControlNetPipeline.from_pretrained(
     vae=vae,
     safety_checker=safety_checker,
     feature_extractor=feature_extractor,
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
 ).to(device)
 
 # Function to check NSFW images