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