Update app.py

app.py

@@ -7,6 +7,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # # Set up environment
 # os.environ['CUDA_VISIBLE_DEVICES'] = "0"
 
+if torch.cuda.is_available():
+    torch_dtype = torch.bfloat16
+else:
+    torch_dtype = torch.float32
+
 # Load models
 def load_models():
     model_name = "X-ART/LeX-Enhancer-full"
@@ -22,7 +27,8 @@ def load_models():
         "X-ART/LeX-Lumina", 
         torch_dtype=torch.bfloat16
     )
-    pipe.to("cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe.to(device, torch_dtype)
     
     return model, tokenizer, pipe