Hjgugugjhuhjggg
commited on
Commit
•
aea29cb
1
Parent(s):
8323d53
Update app.py
Browse files
app.py
CHANGED
@@ -1,19 +1,18 @@
|
|
1 |
import os
|
2 |
import json
|
3 |
-
import uuid
|
4 |
-
import requests
|
5 |
-
import threading
|
6 |
import logging
|
|
|
7 |
from fastapi import FastAPI, HTTPException
|
8 |
from pydantic import BaseModel
|
9 |
from google.cloud import storage
|
10 |
-
from
|
11 |
-
from transformers import pipeline
|
12 |
-
from dotenv import load_dotenv
|
13 |
import uvicorn
|
|
|
|
|
14 |
import io
|
|
|
15 |
|
16 |
-
|
17 |
load_dotenv()
|
18 |
|
19 |
API_KEY = os.getenv("API_KEY")
|
@@ -21,17 +20,15 @@ GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
|
|
21 |
GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
|
22 |
HF_API_TOKEN = os.getenv("HF_API_TOKEN")
|
23 |
|
24 |
-
# Configuración del logger
|
25 |
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
|
26 |
logger = logging.getLogger(__name__)
|
27 |
|
28 |
-
# Inicializar el cliente de Google Cloud Storage
|
29 |
try:
|
30 |
credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
|
31 |
storage_client = storage.Client.from_service_account_info(credentials_info)
|
32 |
bucket = storage_client.bucket(GCS_BUCKET_NAME)
|
33 |
logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
|
34 |
-
except (
|
35 |
logger.error(f"Error al cargar las credenciales o bucket: {e}")
|
36 |
raise RuntimeError(f"Error al cargar las credenciales o bucket: {e}")
|
37 |
|
@@ -51,22 +48,13 @@ class GCSHandler:
|
|
51 |
logger.debug(f"Comprobando existencia de archivo '{blob_name}': {exists}")
|
52 |
return exists
|
53 |
|
54 |
-
def upload_file(self, blob_name, file_stream):
|
55 |
-
blob = self.bucket.blob(blob_name)
|
56 |
-
try:
|
57 |
-
blob.upload_from_file(file_stream)
|
58 |
-
logger.info(f"Archivo '{blob_name}' subido exitosamente a GCS.")
|
59 |
-
except Exception as e:
|
60 |
-
logger.error(f"Error subiendo el archivo '{blob_name}' a GCS: {e}")
|
61 |
-
raise HTTPException(status_code=500, detail=f"Error subiendo archivo '{blob_name}' a GCS")
|
62 |
-
|
63 |
def download_file(self, blob_name):
|
64 |
blob = self.bucket.blob(blob_name)
|
65 |
if not blob.exists():
|
66 |
logger.error(f"Archivo '{blob_name}' no encontrado en GCS.")
|
67 |
raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
|
68 |
logger.debug(f"Descargando archivo '{blob_name}' de GCS.")
|
69 |
-
return blob
|
70 |
|
71 |
def generate_signed_url(self, blob_name, expiration=3600):
|
72 |
blob = self.bucket.blob(blob_name)
|
@@ -74,34 +62,34 @@ class GCSHandler:
|
|
74 |
logger.debug(f"Generada URL firmada para '{blob_name}': {url}")
|
75 |
return url
|
76 |
|
77 |
-
def
|
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 |
@app.post("/predict/")
|
107 |
async def predict(request: DownloadModelRequest):
|
@@ -109,103 +97,59 @@ async def predict(request: DownloadModelRequest):
|
|
109 |
try:
|
110 |
gcs_handler = GCSHandler(GCS_BUCKET_NAME)
|
111 |
model_prefix = request.model_name
|
112 |
-
model_files = [
|
113 |
-
"pytorch_model.bin",
|
114 |
-
"config.json",
|
115 |
-
"tokenizer.json",
|
116 |
-
"model.safetensors",
|
117 |
-
]
|
118 |
-
|
119 |
-
model_files_exist = all(gcs_handler.file_exists(f"{model_prefix}/{file}") for file in model_files)
|
120 |
-
|
121 |
-
if not model_files_exist:
|
122 |
-
logger.info(f"Modelos no encontrados en GCS, descargando '{model_prefix}' desde Hugging Face...")
|
123 |
-
download_model_from_huggingface(model_prefix)
|
124 |
-
|
125 |
-
model_files_streams = {file: gcs_handler.download_file(f"{model_prefix}/{file}") for file in model_files if gcs_handler.file_exists(f"{model_prefix}/{file}")}
|
126 |
|
127 |
-
|
128 |
-
tokenizer_stream = model_files_streams.get("tokenizer.json")
|
129 |
-
model_stream = model_files_streams.get("pytorch_model.bin")
|
130 |
|
131 |
-
if not
|
132 |
-
logger.error(f"
|
133 |
-
raise HTTPException(status_code=
|
|
|
|
|
134 |
|
135 |
-
|
|
|
136 |
if request.pipeline_task in ["text-generation", "translation", "summarization"]:
|
137 |
-
pipe = pipeline(request.pipeline_task, model=model_stream, tokenizer=tokenizer_stream)
|
138 |
result = pipe(request.input_text)
|
139 |
logger.info(f"Resultado generado para la tarea '{request.pipeline_task}': {result[0]}")
|
140 |
return {"response": result[0]}
|
141 |
|
142 |
-
# Tareas de imagen
|
143 |
elif request.pipeline_task == "image-generation":
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
# Subir la imagen generada a GCS
|
153 |
-
gcs_handler.upload_file(image_path, open(image_path, "rb"))
|
154 |
-
image_url = gcs_handler.generate_signed_url(image_path)
|
155 |
-
logger.info(f"Imagen generada y subida correctamente con URL: {image_url}")
|
156 |
-
return {"response": {"image_url": image_url}}
|
157 |
-
except Exception as e:
|
158 |
-
logger.error(f"Error generando la imagen: {e}")
|
159 |
-
raise HTTPException(status_code=400, detail="Error generando la imagen.")
|
160 |
|
161 |
elif request.pipeline_task == "image-editing":
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
gcs_handler.upload_file(f"images/{edited_image_filename}", open(edited_image_filename, "rb"))
|
170 |
-
edited_image_url = gcs_handler.generate_signed_url(f"images/{edited_image_filename}")
|
171 |
-
logger.info(f"Imagen editada y subida correctamente con URL: {edited_image_url}")
|
172 |
-
return {"response": {"edited_image_url": edited_image_url}}
|
173 |
-
except Exception as e:
|
174 |
-
logger.error(f"Error editando la imagen: {e}")
|
175 |
-
raise HTTPException(status_code=400, detail="Error editando la imagen.")
|
176 |
|
177 |
elif request.pipeline_task == "image-to-image":
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
transformed_image_url = gcs_handler.generate_signed_url(f"images/{transformed_image_filename}")
|
187 |
-
logger.info(f"Imagen transformada y subida correctamente con URL: {transformed_image_url}")
|
188 |
-
return {"response": {"transformed_image_url": transformed_image_url}}
|
189 |
-
except Exception as e:
|
190 |
-
logger.error(f"Error transformando la imagen: {e}")
|
191 |
-
raise HTTPException(status_code=400, detail="Error transformando la imagen.")
|
192 |
-
|
193 |
-
# Tarea de generación de modelo 3D (simulada)
|
194 |
elif request.pipeline_task == "text-to-3d":
|
195 |
-
|
196 |
-
|
197 |
-
|
198 |
-
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
logger.info(f"Modelo 3D generado y subido con URL: {model_3d_url}")
|
204 |
-
return {"response": {"model_3d_url": model_3d_url}}
|
205 |
-
except Exception as e:
|
206 |
-
logger.error(f"Error generando el modelo 3D: {e}")
|
207 |
-
raise HTTPException(status_code=400, detail="Error generando el modelo 3D.")
|
208 |
-
|
209 |
except HTTPException as e:
|
210 |
logger.error(f"HTTPException: {e.detail}")
|
211 |
raise e
|
@@ -213,33 +157,37 @@ async def predict(request: DownloadModelRequest):
|
|
213 |
logger.error(f"Error inesperado: {e}")
|
214 |
raise HTTPException(status_code=500, detail=f"Error: {e}")
|
215 |
|
216 |
-
|
217 |
-
|
218 |
-
|
|
|
219 |
try:
|
220 |
-
logger.info("
|
221 |
-
response = requests.get(
|
222 |
-
if response.status_code
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
228 |
-
|
229 |
-
|
230 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
231 |
except Exception as e:
|
232 |
-
logger.error(f"Error
|
233 |
-
raise HTTPException(status_code=500, detail="Error
|
234 |
-
|
235 |
-
# Iniciar la descarga de modelos en segundo plano
|
236 |
-
def run_in_background():
|
237 |
-
logger.info("Iniciando la descarga de modelos en segundo plano...")
|
238 |
-
threading.Thread(target=download_all_models_in_background, daemon=True).start()
|
239 |
|
240 |
@app.on_event("startup")
|
241 |
async def startup_event():
|
242 |
-
|
243 |
|
244 |
if __name__ == "__main__":
|
245 |
-
uvicorn.run(app, host="0.0.0.0", port=7860)
|
|
|
1 |
import os
|
2 |
import json
|
|
|
|
|
|
|
3 |
import logging
|
4 |
+
import uuid
|
5 |
from fastapi import FastAPI, HTTPException
|
6 |
from pydantic import BaseModel
|
7 |
from google.cloud import storage
|
8 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
|
|
|
|
|
9 |
import uvicorn
|
10 |
+
import torch
|
11 |
+
import requests
|
12 |
import io
|
13 |
+
from safetensors import safe_open
|
14 |
|
15 |
+
from dotenv import load_dotenv
|
16 |
load_dotenv()
|
17 |
|
18 |
API_KEY = os.getenv("API_KEY")
|
|
|
20 |
GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
|
21 |
HF_API_TOKEN = os.getenv("HF_API_TOKEN")
|
22 |
|
|
|
23 |
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
|
24 |
logger = logging.getLogger(__name__)
|
25 |
|
|
|
26 |
try:
|
27 |
credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
|
28 |
storage_client = storage.Client.from_service_account_info(credentials_info)
|
29 |
bucket = storage_client.bucket(GCS_BUCKET_NAME)
|
30 |
logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
|
31 |
+
except (json.JSONDecodeError, KeyError, ValueError) as e:
|
32 |
logger.error(f"Error al cargar las credenciales o bucket: {e}")
|
33 |
raise RuntimeError(f"Error al cargar las credenciales o bucket: {e}")
|
34 |
|
|
|
48 |
logger.debug(f"Comprobando existencia de archivo '{blob_name}': {exists}")
|
49 |
return exists
|
50 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
51 |
def download_file(self, blob_name):
|
52 |
blob = self.bucket.blob(blob_name)
|
53 |
if not blob.exists():
|
54 |
logger.error(f"Archivo '{blob_name}' no encontrado en GCS.")
|
55 |
raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
|
56 |
logger.debug(f"Descargando archivo '{blob_name}' de GCS.")
|
57 |
+
return blob
|
58 |
|
59 |
def generate_signed_url(self, blob_name, expiration=3600):
|
60 |
blob = self.bucket.blob(blob_name)
|
|
|
62 |
logger.debug(f"Generada URL firmada para '{blob_name}': {url}")
|
63 |
return url
|
64 |
|
65 |
+
def load_model_from_gcs(model_name: str, model_files: list):
|
66 |
+
gcs_handler = GCSHandler(GCS_BUCKET_NAME)
|
67 |
+
model_blobs = {file: gcs_handler.download_file(f"{model_name}/{file}") for file in model_files}
|
68 |
|
69 |
+
model_stream = model_blobs.get("pytorch_model.bin") or model_blobs.get("model.safetensors")
|
70 |
+
config_stream = model_blobs.get("config.json")
|
71 |
+
tokenizer_stream = model_blobs.get("tokenizer.json")
|
72 |
+
|
73 |
+
if "safetensors" in model_stream.name:
|
74 |
+
model = load_safetensors_model(model_stream)
|
75 |
+
else:
|
76 |
+
model = AutoModelForCausalLM.from_pretrained(model_stream, config=config_stream)
|
77 |
+
|
78 |
+
tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
|
79 |
+
|
80 |
+
return model, tokenizer
|
81 |
+
|
82 |
+
def load_safetensors_model(model_stream):
|
83 |
+
with safe_open(model_stream, framework="pt") as model_data:
|
84 |
+
model = torch.load(model_data)
|
85 |
+
return model
|
86 |
+
|
87 |
+
def get_model_files_from_gcs(model_name: str):
|
88 |
+
gcs_handler = GCSHandler(GCS_BUCKET_NAME)
|
89 |
+
blob_list = list(gcs_handler.bucket.list_blobs(prefix=f"{model_name}/"))
|
90 |
+
model_files = [blob.name for blob in blob_list if "pytorch_model" in blob.name or "model" in blob.name]
|
91 |
+
model_files = sorted(model_files)
|
92 |
+
return model_files
|
93 |
|
94 |
@app.post("/predict/")
|
95 |
async def predict(request: DownloadModelRequest):
|
|
|
97 |
try:
|
98 |
gcs_handler = GCSHandler(GCS_BUCKET_NAME)
|
99 |
model_prefix = request.model_name
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
100 |
|
101 |
+
model_files = get_model_files_from_gcs(model_prefix)
|
|
|
|
|
102 |
|
103 |
+
if not model_files:
|
104 |
+
logger.error(f"Modelos no encontrados en GCS para '{model_prefix}'.")
|
105 |
+
raise HTTPException(status_code=404, detail="Model files not found in GCS.")
|
106 |
+
|
107 |
+
model, tokenizer = load_model_from_gcs(model_prefix, model_files)
|
108 |
|
109 |
+
pipe = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
|
110 |
+
|
111 |
if request.pipeline_task in ["text-generation", "translation", "summarization"]:
|
|
|
112 |
result = pipe(request.input_text)
|
113 |
logger.info(f"Resultado generado para la tarea '{request.pipeline_task}': {result[0]}")
|
114 |
return {"response": result[0]}
|
115 |
|
|
|
116 |
elif request.pipeline_task == "image-generation":
|
117 |
+
images = pipe(request.input_text)
|
118 |
+
image = images[0]
|
119 |
+
image_filename = f"{uuid.uuid4().hex}.png"
|
120 |
+
image_path = f"images/{image_filename}"
|
121 |
+
image.save(image_path)
|
122 |
+
gcs_handler.upload_file(image_path, open(image_path, "rb"))
|
123 |
+
image_url = gcs_handler.generate_signed_url(image_path)
|
124 |
+
return {"response": {"image_url": image_url}}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
125 |
|
126 |
elif request.pipeline_task == "image-editing":
|
127 |
+
edited_images = pipe(request.input_text)
|
128 |
+
edited_image = edited_images[0]
|
129 |
+
edited_image_filename = f"{uuid.uuid4().hex}_edited.png"
|
130 |
+
edited_image.save(edited_image_filename)
|
131 |
+
gcs_handler.upload_file(f"images/{edited_image_filename}", open(edited_image_filename, "rb"))
|
132 |
+
edited_image_url = gcs_handler.generate_signed_url(f"images/{edited_image_filename}")
|
133 |
+
return {"response": {"edited_image_url": edited_image_url}}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
134 |
|
135 |
elif request.pipeline_task == "image-to-image":
|
136 |
+
transformed_images = pipe(request.input_text)
|
137 |
+
transformed_image = transformed_images[0]
|
138 |
+
transformed_image_filename = f"{uuid.uuid4().hex}_transformed.png"
|
139 |
+
transformed_image.save(transformed_image_filename)
|
140 |
+
gcs_handler.upload_file(f"images/{transformed_image_filename}", open(transformed_image_filename, "rb"))
|
141 |
+
transformed_image_url = gcs_handler.generate_signed_url(f"images/{transformed_image_filename}")
|
142 |
+
return {"response": {"transformed_image_url": transformed_image_url}}
|
143 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
144 |
elif request.pipeline_task == "text-to-3d":
|
145 |
+
model_3d_filename = f"{uuid.uuid4().hex}.obj"
|
146 |
+
model_3d_path = f"3d-models/{model_3d_filename}"
|
147 |
+
with open(model_3d_path, "w") as f:
|
148 |
+
f.write("Simulated 3D model data")
|
149 |
+
gcs_handler.upload_file(f"3d-models/{model_3d_filename}", open(model_3d_path, "rb"))
|
150 |
+
model_3d_url = gcs_handler.generate_signed_url(f"3d-models/{model_3d_filename}")
|
151 |
+
return {"response": {"model_3d_url": model_3d_url}}
|
152 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
153 |
except HTTPException as e:
|
154 |
logger.error(f"HTTPException: {e.detail}")
|
155 |
raise e
|
|
|
157 |
logger.error(f"Error inesperado: {e}")
|
158 |
raise HTTPException(status_code=500, detail=f"Error: {e}")
|
159 |
|
160 |
+
def download_model_from_huggingface(model_name):
|
161 |
+
url = f"https://huggingface.co/{model_name}/tree/main"
|
162 |
+
headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
|
163 |
+
|
164 |
try:
|
165 |
+
logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
|
166 |
+
response = requests.get(url, headers=headers)
|
167 |
+
if response.status_code == 200:
|
168 |
+
model_files = [
|
169 |
+
"pytorch_model.bin",
|
170 |
+
"config.json",
|
171 |
+
"tokenizer.json",
|
172 |
+
"model.safetensors",
|
173 |
+
]
|
174 |
+
for file_name in model_files:
|
175 |
+
file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
|
176 |
+
file_content = requests.get(file_url).content
|
177 |
+
blob_name = f"{model_name}/{file_name}"
|
178 |
+
blob = bucket.blob(blob_name)
|
179 |
+
blob.upload_from_string(file_content)
|
180 |
+
logger.info(f"Archivo '{file_name}' subido exitosamente al bucket GCS.")
|
181 |
+
else:
|
182 |
+
logger.error(f"Error al acceder al árbol de archivos de Hugging Face para '{model_name}'.")
|
183 |
+
raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
|
184 |
except Exception as e:
|
185 |
+
logger.error(f"Error descargando archivos de Hugging Face: {e}")
|
186 |
+
raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
|
|
|
|
|
|
|
|
|
|
|
187 |
|
188 |
@app.on_event("startup")
|
189 |
async def startup_event():
|
190 |
+
logger.info("Iniciando la API...")
|
191 |
|
192 |
if __name__ == "__main__":
|
193 |
+
uvicorn.run(app, host="0.0.0.0", port=7860)
|