Update app.py

app.py

@@ -1,36 +1,40 @@
 import os
 import logging
-import requests
+import time
 import threading
 from io import BytesIO
-from fastapi import FastAPI, HTTPException, Response
+from typing import Union
+import requests
+import boto3
+import torch
+import safetensors
+import soundfile as sf
+import numpy as np
+from fastapi import FastAPI, HTTPException, Response, Request, UploadFile, File
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError, field_validator
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
+    pipeline,
     GenerationConfig,
-    pipeline
+    StoppingCriteriaList
 )
-import boto3
-import torch
+from huggingface_hub import hf_hub_download
 import uvicorn
 
-# Configuración de logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
 
-# Variables de entorno
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
 AWS_REGION = os.getenv("AWS_REGION")
 S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
 HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
-# Clase para la petición de generación
 class GenerateRequest(BaseModel):
     model_name: str
-    input_text: str
+    input_text: str = ""
     task_type: str
     temperature: float = 1.0
     max_new_tokens: int = 200
@@ -40,11 +44,24 @@ class GenerateRequest(BaseModel):
     repetition_penalty: float = 1.0
     num_return_sequences: int = 1
     do_sample: bool = True
+    chunk_delay: float = 0.0
+    stop_sequences: list[str] = []
+
+    model_config = {"protected_namespaces": ()}
 
-    class Config:
-        protected_namespaces = ()
+    @field_validator("model_name")
+    def model_name_cannot_be_empty(cls, v):
+        if not v:
+            raise ValueError("model_name cannot be empty.")
+        return v
+
+    @field_validator("task_type")
+    def task_type_must_be_valid(cls, v):
+        valid_types = ["text-to-text", "text-to-image", "text-to-speech", "text-to-video"]
+        if v not in valid_types:
+            raise ValueError(f"task_type must be one of: {valid_types}")
+        return v
 
-# Clase para cargar modelos desde S3
 class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
@@ -53,105 +70,126 @@ class S3ModelLoader:
     def _get_s3_uri(self, model_name):
         return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
     
-    def download_model_from_s3(self, model_name):
+    def _download_from_s3(self, model_name):
+        s3_uri = self._get_s3_uri(model_name)
         try:
-            config = AutoConfig.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
-            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config)
-            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
+            logging.info(f"Attempting to load model {model_name} from S3...")
+            model_files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
+            if "Contents" not in model_files:
+                raise FileNotFoundError(f"Model files not found in S3 for {model_name}")
 
-            # Asegurarse de que el `eos_token_id` y `pad_token_id` estén definidos
-            if model.config.eos_token_id is None:
-                model.config.eos_token_id = tokenizer.eos_token_id
-            if model.config.pad_token_id is None:
-                model.config.pad_token_id = tokenizer.pad_token_id
+            local_dir = f"/tmp/{model_name.replace('/', '-')}"
+            os.makedirs(local_dir, exist_ok=True)
 
-            return model, tokenizer
-        except Exception:
-            return None, None
+            for obj in model_files["Contents"]:
+                file_key = obj["Key"]
+                if file_key.endswith('/'):
+                    continue
+
+                local_file_path = os.path.join(local_dir, os.path.basename(file_key))
+                self.s3_client.download_file(self.bucket_name, file_key, local_file_path)
+
+            return local_dir
+        except Exception as e:
+            logging.error(f"Error downloading from S3: {e}")
+            raise HTTPException(status_code=500, detail=f"Error downloading model from S3: {e}")
 
     async def load_model_and_tokenizer(self, model_name):
         try:
-            model, tokenizer = self.download_model_from_s3(model_name)
-            if model is None or tokenizer is None:
-                model, tokenizer = await self.download_and_save_model_from_huggingface(model_name)
+            model_dir = await self._download_from_s3(model_name)
+            config = AutoConfig.from_pretrained(model_dir)
+            tokenizer = AutoTokenizer.from_pretrained(model_dir, config=config)
+            model = AutoModelForCausalLM.from_pretrained(model_dir, config=config)
+
+            logging.info(f"Model {model_name} loaded from S3 successfully.")
             return model, tokenizer
         except Exception as e:
+            logging.exception(f"Error loading model: {e}")
             raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
-    async def download_and_save_model_from_huggingface(self, model_name):
+    def download_model_from_huggingface(self, model_name):
         try:
-            # Descarga del modelo sin tqdm
-            model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=HUGGINGFACE_HUB_TOKEN)
-            tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HUGGINGFACE_HUB_TOKEN)
-
-            # Asegurarse de que el `eos_token_id` y `pad_token_id` estén definidos
-            if model.config.eos_token_id is None:
-                model.config.eos_token_id = tokenizer.eos_token_id
-            if model.config.pad_token_id is None:
-                model.config.pad_token_id = tokenizer.pad_token_id
-
-            self.upload_model_to_s3(model_name, model, tokenizer)
-            return model, tokenizer
+            logging.info(f"Downloading model {model_name} from Hugging Face...")
+            model_dir = hf_hub_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
+            self.s3_client.upload_file(model_dir, self.bucket_name, model_name)
+            logging.info(f"Model {model_name} saved to S3 successfully.")
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error downloading model from Hugging Face: {e}")
+            logging.error(f"Error downloading model {model_name} from Hugging Face: {e}")
 
-    def upload_model_to_s3(self, model_name, model, tokenizer):
+    def download_all_models_in_background(self):
+        models_url = "https://huggingface.co/api/models"
         try:
-            s3_uri = self._get_s3_uri(model_name)
-            model.save_pretrained(s3_uri)
-            tokenizer.save_pretrained(s3_uri)
+            response = requests.get(models_url)
+            if response.status_code != 200:
+                logging.error("Error getting Hugging Face model list.")
+                raise HTTPException(status_code=500, detail="Error getting model list.")
+
+            models = response.json()
+            for model in models:
+                model_name = model["id"]
+                self.download_model_from_huggingface(model_name)
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error saving model to S3: {e}")
+            logging.error(f"Error downloading models in the background: {e}")
+            raise HTTPException(status_code=500, detail="Error downloading models in the background.")
+
+    def run_in_background(self):
+        threading.Thread(target=self.download_all_models_in_background, daemon=True).start()
+
+@app.on_event("startup")
+async def startup_event():
+    model_loader.run_in_background()
 
-# Crear la instancia de FastAPI
 app = FastAPI()
 
-# Instanciar model_loader aquí
 s3_client = boto3.client('s3', aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY, region_name=AWS_REGION)
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
-# Función de generación asincrónica
 @app.post("/generate")
-async def generate(body: GenerateRequest):
+async def generate(request: Request, body: GenerateRequest):
     try:
-        model, tokenizer = await model_loader.load_model_and_tokenizer(body.model_name)
+        validated_body = GenerateRequest(**body.model_dump())
+        model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
 
-        if body.task_type == "text-to-text":
+        if validated_body.task_type == "text-to-text":
             generation_config = GenerationConfig(
-                temperature=body.temperature,
-                max_new_tokens=body.max_new_tokens,
-                top_p=body.top_p,
-                top_k=body.top_k,
-                repetition_penalty=body.repetition_penalty,
-                do_sample=body.do_sample,
-                num_return_sequences=body.num_return_sequences
+                temperature=validated_body.temperature,
+                max_new_tokens=validated_body.max_new_tokens,
+                top_p=validated_body.top_p,
+                top_k=validated_body.top_k,
+                repetition_penalty=validated_body.repetition_penalty,
+                do_sample=validated_body.do_sample,
+                num_return_sequences=validated_body.num_return_sequences
             )
 
             async def stream_text():
-                input_text = body.input_text
-                max_length = model.config.max_position_embeddings
+                input_text = validated_body.input_text
                 generated_text = ""
+                max_length = model.config.max_position_embeddings
 
                 while True:
-                    inputs = tokenizer(input_text, return_tensors="pt").to(device)
-                    input_length = inputs.input_ids.shape[1]
+                    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
+                    input_length = encoded_input["input_ids"].shape[1]
                     remaining_tokens = max_length - input_length
-                    if remaining_tokens < body.max_new_tokens:
-                        generation_config.max_new_tokens = remaining_tokens
-                        if remaining_tokens <= 0:
-                            break
 
-                    output = model.generate(**inputs, generation_config=generation_config)
+                    if remaining_tokens <= 0:
+                        break
+
+                    generation_config.max_new_tokens = min(remaining_tokens, validated_body.max_new_tokens)
+
+                    stopping_criteria = StoppingCriteriaList(
+                        [lambda _, outputs: tokenizer.decode(outputs[0][-1], skip_special_tokens=True) in validated_body.stop_sequences] if validated_body.stop_sequences else []
+                    )
+
+                    output = model.generate(**encoded_input, generation_config=generation_config, stopping_criteria=stopping_criteria)
                     chunk = tokenizer.decode(output[0], skip_special_tokens=True)
                     generated_text += chunk
                     yield chunk
-                    if len(tokenizer.encode(generated_text)) >= max_length:
-                        break
-                    input_text = chunk
+                    time.sleep(validated_body.chunk_delay)
+                    input_text = generated_text
 
-            if body.stream:
+            if validated_body.stream:
                 return StreamingResponse(stream_text(), media_type="text/plain")
             else:
                 generated_text = ""
@@ -159,24 +197,24 @@ async def generate(body: GenerateRequest):
                     generated_text += chunk
                 return {"result": generated_text}
 
-        elif body.task_type == "text-to-image":
+        elif validated_body.task_type == "text-to-image":
             generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=device)
-            image = generator(body.input_text)[0]
+            image = generator(validated_body.input_text)[0]
             image_bytes = image.tobytes()
             return Response(content=image_bytes, media_type="image/png")
 
-        elif body.task_type == "text-to-speech":
+        elif validated_body.task_type == "text-to-speech":
             generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
-            audio = generator(body.input_text)
+            audio = generator(validated_body.input_text)
             audio_bytesio = BytesIO()
             sf.write(audio_bytesio, audio["sampling_rate"], np.int16(audio["audio"]))
             audio_bytes = audio_bytesio.getvalue()
             return Response(content=audio_bytes, media_type="audio/wav")
 
-        elif body.task_type == "text-to-video":
+        elif validated_body.task_type == "text-to-video":
             try:
                 generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
-                video = generator(body.input_text)
+                video = generator(validated_body.input_text)
                 return Response(content=video, media_type="video/mp4")
             except Exception as e:
                 raise HTTPException(status_code=500, detail=f"Error in text-to-video generation: {e}")
@@ -186,38 +224,11 @@ async def generate(body: GenerateRequest):
 
     except HTTPException as e:
         raise e
+    except ValidationError as e:
+        raise HTTPException(status_code=422, detail=e.errors())
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logging.exception(f"An unexpected error occurred: {e}")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
 
-# Descargar todos los modelos en segundo plano
-async def download_all_models_in_background():
-    models_url = "https://huggingface.co/api/models"
-    try:
-        # Se obtiene la lista de modelos
-        response = requests.get(models_url)
-        if response.status_code != 200:
-            raise HTTPException(status_code=500, detail="Error al obtener la lista de modelos.")
-        
-        models = response.json()
-        for model in models:
-            model_name = model["id"]
-            # Verifica si ya está en S3 antes de intentar descargarlo
-            try:
-                await model_loader.download_and_save_model_from_huggingface(model_name)
-            except Exception as e:
-                logging.error(f"Error descargando o guardando el modelo {model_name}: {str(e)}")
-
-    except Exception as e:
-        logging.error(f"Error al obtener modelos de Hugging Face: {str(e)}")
-
-# Función que corre en segundo plano para descargar modelos
-def run_in_background():
-    threading.Thread(target=download_all_models_in_background, daemon=True).start()
-
-# Si este archivo se ejecuta directamente, inicia el servidor
 if __name__ == "__main__":
-    # Ejecutar la descarga de modelos en segundo plano
-    run_in_background()
-
-    # Iniciar el servidor FastAPI
     uvicorn.run(app, host="0.0.0.0", port=7860)