Spaces:
Build error
Build error
woods-today
commited on
Commit
·
460cdbf
1
Parent(s):
a3614ce
Working on it
Browse files- routers/training.py +1 -0
routers/training.py
CHANGED
|
@@ -9,6 +9,7 @@ import torch
|
|
| 9 |
from diffusers import StableDiffusionPipeline
|
| 10 |
|
| 11 |
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
|
|
|
|
| 12 |
|
| 13 |
|
| 14 |
router = APIRouter()
|
|
|
|
| 9 |
from diffusers import StableDiffusionPipeline
|
| 10 |
|
| 11 |
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16)
|
| 12 |
+
pipe = pipe.to("cuda")
|
| 13 |
|
| 14 |
|
| 15 |
router = APIRouter()
|