Update app.py

app.py

@@ -1,18 +1,18 @@
 import torch
+import zerogpu  # Import ZeroGPU
 from diffusers import StableDiffusion3Pipeline
 from huggingface_hub import login
 import os
 import gradio as gr
-import zerogpu
-print(zerogpu.__version__)
 
-# Check for GPU availability and set device accordingly
+# Automatically select the available device (GPU if available, otherwise CPU)
+device = zerogpu.select_device()  # This will automatically choose 'cuda' or 'cpu'
+
+# Check and print if the selected device is GPU or CPU
 if torch.cuda.is_available():
-    device = "cuda"
-    print("GPU is available")
+    print(f"Using GPU: {torch.cuda.get_device_name()}")
 else:
-    device = "cpu"
-    print("GPU is not available, using CPU")
+    print("Using CPU")
 
 # Retrieve the token from the environment variable
 token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
@@ -72,3 +72,4 @@ iface = gr.Interface(
     outputs="image"
 )
 iface.launch()
+