Update app.py

app.py

@@ -13,6 +13,9 @@ import torch
 from diffusers import DiffusionPipeline
 from typing import Tuple
 
+# Check if GPU is available; fallback to CPU if needed
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 # Setup rules for bad words (ensure the prompts are kid-friendly)
 bad_words = json.loads(os.getenv('BAD_WORDS', '["violence", "blood", "scary", "death", "ghost"]'))
 default_negative = os.getenv("default_negative","")