Spaces:
Build error
Build error
get env
Browse files- server/src/main.py +3 -2
server/src/main.py
CHANGED
|
@@ -1,3 +1,4 @@
|
|
|
|
|
| 1 |
import pickle
|
| 2 |
from typing import Union
|
| 3 |
from fastapi import Request
|
|
@@ -8,8 +9,8 @@ from fastapi import FastAPI
|
|
| 8 |
|
| 9 |
app = FastAPI()
|
| 10 |
|
| 11 |
-
DEVICE =
|
| 12 |
-
ATTN_IMPLEMENTATION = "sdpa"
|
| 13 |
|
| 14 |
@app.get("/")
|
| 15 |
def read_root():
|
|
|
|
| 1 |
+
import os
|
| 2 |
import pickle
|
| 3 |
from typing import Union
|
| 4 |
from fastapi import Request
|
|
|
|
| 9 |
|
| 10 |
app = FastAPI()
|
| 11 |
|
| 12 |
+
DEVICE = os.getenv('DEVICE', 'mps')
|
| 13 |
+
ATTN_IMPLEMENTATION = os.getenv('ATTN_IMPLEMENTATION', "sdpa")
|
| 14 |
|
| 15 |
@app.get("/")
|
| 16 |
def read_root():
|