Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import torch
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel
 from transformers import (
     AutoConfig,
     pipeline,
@@ -11,25 +11,24 @@ from transformers import (
     GenerationConfig,
     StoppingCriteriaList
 )
-import boto3
-import uvicorn
 import asyncio
 from io import BytesIO
-from transformers import pipeline
 
+# Diccionario global para almacenar los tokens
+token_dict = {}
+
+# Setup para acceder a modelos en Hugging Face o S3
 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")
 
-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)
-
 app = FastAPI()
 
 class GenerateRequest(BaseModel):
     model_name: str
-    input_text: str = ""
+    input_text: str
     task_type: str
     temperature: float = 1.0
     max_new_tokens: int = 200
@@ -42,19 +41,6 @@ class GenerateRequest(BaseModel):
     chunk_delay: float = 0.0
     stop_sequences: list[str] = []
 
-    @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
-
 class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
@@ -62,32 +48,29 @@ class S3ModelLoader:
 
     def _get_s3_uri(self, model_name):
         return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
-
+    
     async def load_model_and_tokenizer(self, model_name):
+        if model_name in token_dict:
+            return token_dict[model_name]
+        
         s3_uri = self._get_s3_uri(model_name)
         try:
-            config = AutoConfig.from_pretrained(s3_uri, local_files_only=True)
-            model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=True)
-            tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=True)
+            model = AutoModelForCausalLM.from_pretrained(s3_uri, local_files_only=True)
+            tokenizer = AutoTokenizer.from_pretrained(s3_uri, local_files_only=True)
 
-            if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
-                tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
+            if tokenizer.eos_token_id is None:
+                tokenizer.eos_token_id = tokenizer.pad_token_id
 
-            return model, tokenizer
-        except EnvironmentError:
-            try:
-                config = AutoConfig.from_pretrained(model_name)
-                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
-                model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
+            token_dict[model_name] = {
+                "model": model,
+                "tokenizer": tokenizer,
+                "pad_token_id": tokenizer.pad_token_id,
+                "eos_token_id": tokenizer.eos_token_id
+            }
 
-                if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
-                    tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
-
-                model.save_pretrained(s3_uri)
-                tokenizer.save_pretrained(s3_uri)
-                return model, tokenizer
-            except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
+            return token_dict[model_name]
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
@@ -96,7 +79,6 @@ async def generate(request: GenerateRequest):
     try:
         model_name = request.model_name
         input_text = request.input_text
-        task_type = request.task_type
         temperature = request.temperature
         max_new_tokens = request.max_new_tokens
         stream = request.stream
@@ -108,7 +90,13 @@ async def generate(request: GenerateRequest):
         chunk_delay = request.chunk_delay
         stop_sequences = request.stop_sequences
 
-        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
+        # Cargar modelo y tokenizer desde el S3
+        model_data = await model_loader.load_model_and_tokenizer(model_name)
+        model = model_data["model"]
+        tokenizer = model_data["tokenizer"]
+        pad_token_id = model_data["pad_token_id"]
+        eos_token_id = model_data["eos_token_id"]
+
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
 
@@ -239,4 +227,5 @@ async def generate_video(request: GenerateRequest):
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)