Update app.py

app.py

@@ -1,19 +1,17 @@
 import os
 import logging
-import time
+import requests
+import threading
 from io import BytesIO
-from typing import Union
-
-from fastapi import FastAPI, HTTPException, Response, Request, UploadFile, File
+from fastapi import FastAPI, HTTPException, Response, Request
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, ValidationError, field_validator
+from pydantic import BaseModel
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     pipeline,
-    GenerationConfig,
-    StoppingCriteriaList
+    GenerationConfig
 )
 import boto3
 from huggingface_hub import hf_hub_download
@@ -21,8 +19,9 @@ import soundfile as sf
 import numpy as np
 import torch
 import uvicorn
+from tqdm import tqdm
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -32,7 +31,7 @@ HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
 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,23 +41,6 @@ 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": ()}
-
-    @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):
@@ -67,39 +49,50 @@ 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):
-        s3_uri = self._get_s3_uri(model_name)
+    
+    def download_model_from_s3(self, model_name):
         try:
             logging.info(f"Trying to load {model_name} from S3...")
-            config = AutoConfig.from_pretrained(s3_uri)
-            model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config)
-            tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config)
+            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"Loaded {model_name} from S3 successfully.")
+            return model, tokenizer
+        except Exception as e:
+            logging.error(f"Error loading {model_name} from S3: {e}")
+            return None, None
 
-            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
+    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)
+            return model, tokenizer
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
-            logging.info(f"Loaded {model_name} from S3 successfully.")
+    async def download_and_save_model_from_huggingface(self, model_name):
+        try:
+            logging.info(f"Downloading {model_name} from Hugging Face...")
+            with tqdm(unit="B", unit_scale=True, desc=f"Downloading {model_name}") as t:
+                model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN, _tqdm=t)
+                tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
+            logging.info(f"Downloaded {model_name} successfully.")
+            self.upload_model_to_s3(model_name, model, tokenizer)
             return model, tokenizer
-        except EnvironmentError:
-            logging.info(f"Model {model_name} not found in S3. Downloading...")
-            try:
-                config = AutoConfig.from_pretrained(model_name)
-                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
-                model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN)
-
-                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
-
-                logging.info(f"Downloaded {model_name} successfully.")
-                logging.info(f"Saving {model_name} to S3...")
-                model.save_pretrained(s3_uri)
-                tokenizer.save_pretrained(s3_uri)
-                logging.info(f"Saved {model_name} to S3 successfully.")
-                return model, tokenizer
-            except Exception as e:
-                logging.exception(f"Error downloading/uploading model: {e}")
-                raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
+        except Exception as e:
+            logging.error(f"Error downloading model from Hugging Face: {e}")
+            raise HTTPException(status_code=500, detail=f"Error downloading model from Hugging Face: {e}")
+
+    def upload_model_to_s3(self, model_name, model, tokenizer):
+        try:
+            s3_uri = self._get_s3_uri(model_name)
+            model.save_pretrained(s3_uri)
+            tokenizer.save_pretrained(s3_uri)
+            logging.info(f"Saved {model_name} to S3 successfully.")
+        except Exception as e:
+            logging.error(f"Error saving {model_name} to S3: {e}")
+            raise HTTPException(status_code=500, detail=f"Error saving model to S3: {e}")
 
 app = FastAPI()
 
@@ -109,49 +102,44 @@ model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 @app.post("/generate")
 async def generate(request: Request, body: GenerateRequest):
     try:
-        validated_body = GenerateRequest(**body.model_dump())
-        model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name)
+        model, tokenizer = await model_loader.load_model_and_tokenizer(body.model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
 
-        if validated_body.task_type == "text-to-text":
+        if body.task_type == "text-to-text":
             generation_config = GenerationConfig(
-                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
+                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
             )
 
             async def stream_text():
-                input_text = validated_body.input_text
-                generated_text = ""
+                input_text = body.input_text
                 max_length = model.config.max_position_embeddings
+                generated_text = ""
 
                 while True:
-                    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
-                    input_length = encoded_input["input_ids"].shape[1]
+                    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+                    input_length = inputs.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
 
-                    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)
+                    output = model.generate(**inputs, generation_config=generation_config)
                     chunk = tokenizer.decode(output[0], skip_special_tokens=True)
                     generated_text += chunk
                     yield chunk
-                    time.sleep(validated_body.chunk_delay)
-                    input_text = generated_text
+                    if len(tokenizer.encode(generated_text)) >= max_length:
+                        break
+                    input_text = chunk
 
-            if validated_body.stream:
+            if body.stream:
                 return StreamingResponse(stream_text(), media_type="text/plain")
             else:
                 generated_text = ""
@@ -159,24 +147,24 @@ async def generate(request: Request, body: GenerateRequest):
                     generated_text += chunk
                 return {"result": generated_text}
 
-        elif validated_body.task_type == "text-to-image":
+        elif body.task_type == "text-to-image":
             generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=device)
-            image = generator(validated_body.input_text)[0]
+            image = generator(body.input_text)[0]
             image_bytes = image.tobytes()
             return Response(content=image_bytes, media_type="image/png")
 
-        elif validated_body.task_type == "text-to-speech":
+        elif body.task_type == "text-to-speech":
             generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
-            audio = generator(validated_body.input_text)
+            audio = generator(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 validated_body.task_type == "text-to-video":
+        elif body.task_type == "text-to-video":
             try:
                 generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
-                video = generator(validated_body.input_text)
+                video = generator(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,12 +174,31 @@ async def generate(request: Request, body: GenerateRequest):
 
     except HTTPException as e:
         raise e
-    except ValidationError as e:
-        raise HTTPException(status_code=422, detail=e.errors())
     except Exception as e:
-        logging.exception(f"An unexpected error occurred: {e}")
-        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+def download_all_models_in_background():
+    models_url = "https://huggingface.co/api/models"
+    try:
+        response = requests.get(models_url)
+        if response.status_code != 200:
+            logging.error("Error al obtener la lista de modelos de Hugging Face.")
+            raise HTTPException(status_code=500, detail="Error al obtener la lista de modelos.")
+        
+        models = response.json()
+        for model in models:
+            model_name = model["id"]
+            model_loader.download_and_save_model_from_huggingface(model_name)
+    except Exception as e:
+        logging.error(f"Error al descargar modelos en segundo plano: {e}")
+        raise HTTPException(status_code=500, detail="Error al descargar modelos en segundo plano.")
+
+def run_in_background():
+    threading.Thread(target=download_all_models_in_background, daemon=True).start()
 
+@app.on_event("startup")
+async def startup_event():
+    run_in_background()
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)