Update app.py

app.py

@@ -157,6 +157,8 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
+from accelerate import Accelerator
+
 #import diffusers
 from diffusers import AutoencoderKL, StableDiffusionXLPipeline
 from diffusers import EulerAncestralDiscreteScheduler
@@ -175,6 +177,8 @@ os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 os.environ["SAFETENSORS_FAST_GPU"] = "1"
 
+accelerator = Accelerator(mixed_precision="bf16") # Example
+
 upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
 
 def load_and_prepare_model():
@@ -198,7 +202,8 @@ def load_and_prepare_model():
     print(f'init noise scale: {pipe.scheduler.init_noise_sigma}')
     pipe.watermark=None
     pipe.safety_checker=None  
-    pipe.to(torch.device('cuda:0'), torch.bfloat16)
+    #pipe.to(torch.device('cuda:0'), torch.bfloat16)
+    pipe.to(accelerator.device)
 
     return pipe