thecollabagepatch commited on
Commit
e5eaa7e
·
1 Parent(s): 85e8363

incorrect JSONResponse import

Browse files
Files changed (1) hide show
  1. app.py +11 -2
app.py CHANGED
@@ -1,7 +1,15 @@
1
  import os
2
 
3
  # ---- Space mode gating (place above any JAX import!) ----
4
- SPACE_MODE = os.getenv("SPACE_MODE", "serve") # "serve" | "template"
 
 
 
 
 
 
 
 
5
 
6
  if SPACE_MODE != "serve":
7
  # In template mode, force JAX to CPU so it won't try to load CUDA plugins
@@ -40,7 +48,8 @@ except Exception:
40
 
41
  from magenta_rt import system, audio as au
42
  import numpy as np
43
- from fastapi import FastAPI, UploadFile, File, Form, Body, HTTPException, Response, Request, WebSocket, WebSocketDisconnect, Query, JSONResponse
 
44
  import tempfile, io, base64, math, threading
45
  from fastapi.middleware.cors import CORSMiddleware
46
  from contextlib import contextmanager
 
1
  import os
2
 
3
  # ---- Space mode gating (place above any JAX import!) ----
4
+ SPACE_MODE = os.getenv("SPACE_MODE")
5
+ if SPACE_MODE is None:
6
+ try:
7
+ import jax
8
+ SPACE_MODE = "serve" if any(getattr(d, "platform", "") in ("gpu","cuda","rocm") for d in jax.devices()) else "template"
9
+ except Exception:
10
+ SPACE_MODE = "template"
11
+
12
+
13
 
14
  if SPACE_MODE != "serve":
15
  # In template mode, force JAX to CPU so it won't try to load CUDA plugins
 
48
 
49
  from magenta_rt import system, audio as au
50
  import numpy as np
51
+ from fastapi import FastAPI, UploadFile, File, Form, Body, HTTPException, Response, Request, WebSocket, WebSocketDisconnect, Query
52
+ from fastapi.responses import JSONResponse
53
  import tempfile, io, base64, math, threading
54
  from fastapi.middleware.cors import CORSMiddleware
55
  from contextlib import contextmanager