Imagebind-Server/app.py at main · Rizwan-ReNote/Imagebind-Server · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
# import os
# from contextlib import asynccontextmanager
# from logging import getLogger

# from dotenv import load_dotenv  # Load .env variables
# from fastapi import FastAPI, Response, status

# from bind import Bind, BindInput
# from meta import Meta

# # Load environment variables
# # load_dotenv()

# bind: Bind
# meta_config: Meta
# logger = getLogger('uvicorn')


# @asynccontextmanager
# async def lifespan(app: FastAPI):
#     global bind
#     global meta_config

#     logger.info("Starting lifespan context")
#     cuda_env = os.getenv("ENABLE_CUDA", "false").lower()
#     cuda_support = False
#     cuda_core = "cpu"

#     if cuda_env in ["true", "1"]:
#         cuda_support = True
#         cuda_core = os.getenv("CUDA_CORE", "cuda:0")
#         logger.info(f"CUDA_CORE set to {cuda_core}")
#     else:
#         logger.info("Running on CPU")

#     try:
#         bind = Bind(cuda_support, cuda_core)
#         meta_config = Meta()
#         logger.info("Model initialization complete")
#     except Exception as e:
#         logger.exception("Failed to initialize Bind or Meta")
#         raise e  # Ensure the error propagates, and FastAPI logs it

#     yield

# app = FastAPI(lifespan=lifespan)


# @app.get("/.well-known/live", response_class=Response)
# @app.get("/.well-known/ready", response_class=Response)
# @app.head("/.well-known/live", response_class=Response)
# @app.head("/.well-known/ready", response_class=Response)
# async def live_and_ready(response: Response):
#     response.status_code = status.HTTP_204_NO_CONTENT


# @app.get("/meta")
# async def meta():
#     return await meta_config.get()


# @app.post("/vectorize")
# async def vectorize(payload: BindInput, response: Response):
#     try:
#         result = await bind.vectorize(payload)
#         return {
#             "textVectors": result.text_vectors,
#             "imageVectors": result.image_vectors,
#             "audioVectors": result.audio_vectors,
#             "videoVectors": result.video_vectors,
#             "imuVectors": result.imu_vectors,
#             "depthVectors": result.depth_vectors,
#             "thermalVectors": result.thermal_vectors,
#         }
#     except Exception as e:
#         logger.exception('Something went wrong while vectorizing data.')
#         response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
#         return {"error": str(e)}


import os
from contextlib import asynccontextmanager
from logging import getLogger

from fastapi import FastAPI, Response, status

from bind import Bind, BindInput
from meta import Meta

bind: Bind
meta_config: Meta
logger = getLogger('uvicorn')


@asynccontextmanager
async def lifespan(app: FastAPI):
    global bind
    global meta_config

    logger.info("Starting lifespan context")
    # Set CUDA configuration directly without environment variables
    cuda_support = True
    # Change "cuda:0" if you want to specify a different GPU, e.g., "cuda:1"
    cuda_core = "cuda:0"
    logger.info(f"CUDA_CORE set to {cuda_core}")

    try:
        bind = Bind(cuda_support, cuda_core)
        meta_config = Meta()
        logger.info("Model initialization complete")
    except Exception as e:
        logger.exception("Failed to initialize Bind or Meta")
        raise e  # Ensure the error propagates, and FastAPI logs it

    yield

app = FastAPI(lifespan=lifespan)


@app.get("/.well-known/live", response_class=Response)
@app.get("/.well-known/ready", response_class=Response)
@app.head("/.well-known/live", response_class=Response)
@app.head("/.well-known/ready", response_class=Response)
async def live_and_ready(response: Response):
    response.status_code = status.HTTP_204_NO_CONTENT


@app.get("/meta")
async def meta():
    return await meta_config.get()


@app.post("/vectorize")
async def vectorize(payload: BindInput, response: Response):
    try:
        result = await bind.vectorize(payload)
        return {
            "textVectors": result.text_vectors,
            "imageVectors": result.image_vectors,
            "audioVectors": result.audio_vectors,
            "videoVectors": result.video_vectors,
            "imuVectors": result.imu_vectors,
            "depthVectors": result.depth_vectors,
            "thermalVectors": result.thermal_vectors,
        }
    except Exception as e:
        logger.exception('Something went wrong while vectorizing data.')
        response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
        return {"error": str(e)}