Update app.py

app.py

@@ -35,10 +35,16 @@ if mps_available:
     torch_dtype = torch.float32
 
 if SAFETY_CHECKER == "True":
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", revision="pr/4")
+    pipe = DiffusionPipeline.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        torch_dtype=torch_dtype,
+        variant="fp16" if torch_dtype == torch.float16 else "fp32")
 else:
     pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/sdxl-turbo", revision="pr/4", safety_checker=None
+        "stabilityai/sdxl-turbo",
+        safety_checker=None,
+        torch_dtype=torch_dtype,
+        variant="fp16" if torch_dtype == torch.float16 else "fp32",
     )
 
 
@@ -116,7 +122,7 @@ with gr.Blocks(css=css) as demo:
             from diffusers import DiffusionPipeline
 
             pipe = DiffusionPipeline.from_pretrained(
-                "stabilityai/sdxl-turbo", revision="refs/pr/4"
+                "stabilityai/sdxl-turbo", variant="fp16", torch_dtype=torch.float16
             ).to("cuda")
             results = pipe(
                 prompt="A cinematic shot of a baby racoon wearing an intricate italian priest robe",