Hjgugugjhuhjggg commited on
Commit
aea29cb
1 Parent(s): 8323d53

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +101 -153
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 google.auth import exceptions
11
- from transformers import pipeline
12
- from dotenv import load_dotenv
13
  import uvicorn
 
 
14
  import io
 
15
 
16
- # Configuración de carga de variables de entorno
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 (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
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.open("rb") # Abre el archivo en modo lectura de bytes
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 download_model_from_huggingface(model_name):
78
- url = f"https://huggingface.co/{model_name}/tree/main"
79
- headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
80
 
81
- try:
82
- logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
83
- response = requests.get(url, headers=headers)
84
- if response.status_code == 200:
85
- model_files = [
86
- "pytorch_model.bin",
87
- "config.json",
88
- "tokenizer.json",
89
- "model.safetensors",
90
- ]
91
- for file_name in model_files:
92
- file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
93
- file_content = requests.get(file_url).content
94
- # Subir el archivo directamente desde el contenido
95
- blob_name = f"{model_name}/{file_name}"
96
- blob = bucket.blob(blob_name)
97
- blob.upload_from_string(file_content)
98
- logger.info(f"Archivo '{file_name}' subido exitosamente al bucket GCS.")
99
- else:
100
- logger.error(f"Error al acceder al árbol de archivos de Hugging Face para '{model_name}'.")
101
- raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
102
- except Exception as e:
103
- logger.error(f"Error descargando archivos de Hugging Face: {e}")
104
- raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
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
- config_stream = model_files_streams.get("config.json")
128
- tokenizer_stream = model_files_streams.get("tokenizer.json")
129
- model_stream = model_files_streams.get("pytorch_model.bin")
130
 
131
- if not config_stream or not tokenizer_stream or not model_stream:
132
- logger.error(f"Faltan archivos necesarios para el modelo '{model_prefix}'.")
133
- raise HTTPException(status_code=500, detail="Required model files missing.")
 
 
134
 
135
- # Tareas basadas en texto
 
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
- try:
145
- pipe = pipeline("image-generation", model=model_stream)
146
- images = pipe(request.input_text)
147
- image = images[0]
148
- image_filename = f"{uuid.uuid4().hex}.png"
149
- image_path = f"images/{image_filename}"
150
- image.save(image_path)
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
- try:
163
- pipe = pipeline("image-editing", model=model_stream)
164
- edited_images = pipe(request.input_text)
165
- edited_image = edited_images[0]
166
- edited_image_filename = f"{uuid.uuid4().hex}_edited.png"
167
- edited_image.save(edited_image_filename)
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
- try:
179
- pipe = pipeline("image-to-image", model=model_stream)
180
- transformed_images = pipe(request.input_text)
181
- transformed_image = transformed_images[0]
182
- transformed_image_filename = f"{uuid.uuid4().hex}_transformed.png"
183
- transformed_image.save(transformed_image_filename)
184
-
185
- gcs_handler.upload_file(f"images/{transformed_image_filename}", open(transformed_image_filename, "rb"))
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
- try:
196
- model_3d_filename = f"{uuid.uuid4().hex}.obj"
197
- model_3d_path = f"3d-models/{model_3d_filename}"
198
- with open(model_3d_path, "w") as f:
199
- f.write("Simulated 3D model data")
200
-
201
- gcs_handler.upload_file(f"3d-models/{model_3d_filename}", open(model_3d_path, "rb"))
202
- model_3d_url = gcs_handler.generate_signed_url(f"3d-models/{model_3d_filename}")
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
- # Función para ejecutar en segundo plano la descarga de modelos
217
- def download_all_models_in_background():
218
- models_url = "https://huggingface.co/api/models"
 
219
  try:
220
- logger.info("Obteniendo lista de modelos desde Hugging Face...")
221
- response = requests.get(models_url)
222
- if response.status_code != 200:
223
- logger.error("Error al obtener la lista de modelos de Hugging Face.")
224
- raise HTTPException(status_code=500, detail="Error al obtener la lista de modelos.")
225
-
226
- models = response.json()
227
- for model in models:
228
- model_name = model["id"]
229
- logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
230
- download_model_from_huggingface(model_name)
 
 
 
 
 
 
 
 
231
  except Exception as e:
232
- logger.error(f"Error al descargar modelos en segundo plano: {e}")
233
- raise HTTPException(status_code=500, detail="Error al descargar modelos en segundo plano.")
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
- run_in_background()
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)