Update app.py

app.py

@@ -1,155 +1,204 @@
 import os
-import json
-import threading
 import logging
-from google.cloud import storage
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from pydantic import BaseModel
-from fastapi import FastAPI, HTTPException
 import requests
+import threading
+from io import BytesIO
+from fastapi import FastAPI, HTTPException, Response, Request
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    GenerationConfig
+)
+import boto3
+from huggingface_hub import hf_hub_download
+import soundfile as sf
+import numpy as np
+import torch
 import uvicorn
-from dotenv import load_dotenv
+from tqdm import tqdm
 
-load_dotenv()
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
-API_KEY = os.getenv("API_KEY")
-GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
-GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
-HF_API_TOKEN = os.getenv("HF_API_TOKEN")
+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")
 
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
-storage_client = storage.Client.from_service_account_info(credentials_info)
-bucket = storage_client.bucket(GCS_BUCKET_NAME)
-
-app = FastAPI()
-
-class DownloadModelRequest(BaseModel):
+class GenerateRequest(BaseModel):
     model_name: str
-    pipeline_task: str
     input_text: str
+    task_type: str
+    temperature: float = 1.0
+    max_new_tokens: int = 200
+    stream: bool = False
+    top_p: float = 1.0
+    top_k: int = 50
+    repetition_penalty: float = 1.0
+    num_return_sequences: int = 1
+    do_sample: bool = True
+
+class S3ModelLoader:
+    def __init__(self, bucket_name, s3_client):
+        self.bucket_name = bucket_name
+        self.s3_client = s3_client
+
+    def _get_s3_uri(self, model_name):
+        return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
+    
+    def download_model_from_s3(self, model_name):
+        try:
+            logging.info(f"Trying to load {model_name} from S3...")
+            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
 
-class GCSHandler:
-    def __init__(self, bucket_name):
-        self.bucket = storage_client.bucket(bucket_name)
-
-    def file_exists(self, blob_name):
-        return self.bucket.blob(blob_name).exists()
-
-    def create_folder_if_not_exists(self, folder_name):
-        if not self.file_exists(folder_name):
-            self.bucket.blob(folder_name + "/").upload_from_string("")
-
-    def upload_file(self, blob_name, file_stream):
-        self.create_folder_if_not_exists(os.path.dirname(blob_name))
-        blob = self.bucket.blob(blob_name)
-        blob.upload_from_file(file_stream)
-
-    def download_file(self, blob_name):
-        blob = self.bucket.blob(blob_name)
-        if not blob.exists():
-            raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
-        return blob.open("rb")
-
-    def generate_signed_url(self, blob_name, expiration=3600):
-        blob = self.bucket.blob(blob_name)
-        return blob.generate_signed_url(expiration=expiration)
-
-def download_model_from_huggingface(model_name):
-    url = f"https://huggingface.co/{model_name}/tree/main"
-    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-    response = requests.get(url, headers=headers)
-    if response.status_code == 200:
-        model_files = [
-            "pytorch_model.bin",
-            "config.json",
-            "tokenizer.json",
-            "model.safetensors",
-        ]
-        for file_name in model_files:
-            file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
-            file_content = requests.get(file_url).content
-            blob_name = f"models/{model_name}/{file_name}"
-            bucket.blob(blob_name).upload_from_string(file_content)
-    else:
-        raise HTTPException(status_code=404, detail="Error accessing Hugging Face model files.")
-
-def download_and_verify_model(model_name):
-    model_files = [
-        "pytorch_model.bin",
-        "config.json",
-        "tokenizer.json",
-        "model.safetensors",
-    ]
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    if not all(gcs_handler.file_exists(f"models/{model_name}/{file}") for file in model_files):
-        download_model_from_huggingface(model_name)
-
-def load_model_from_gcs(model_name):
-    model_files = [
-        "pytorch_model.bin",
-        "config.json",
-        "tokenizer.json",
-        "model.safetensors",
-    ]
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    model_files_streams = {
-        file: gcs_handler.download_file(f"models/{model_name}/{file}")
-        for file in model_files if gcs_handler.file_exists(f"models/{model_name}/{file}")
-    }
-    model_stream = model_files_streams.get("pytorch_model.bin") or model_files_streams.get("model.safetensors")
-    tokenizer_stream = model_files_streams.get("tokenizer.json")
-    config_stream = model_files_streams.get("config.json")
-    model = AutoModelForCausalLM.from_pretrained(model_stream, config=config_stream)
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
-    return model, tokenizer
-
-def load_model(model_name):
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    try:
-        return load_model_from_gcs(model_name)
-    except HTTPException:
-        download_and_verify_model(model_name)
-        return load_model_from_gcs(model_name)
+    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}")
 
-@app.on_event("startup")
-async def startup():
-    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    blobs = list(bucket.list_blobs(prefix="models/"))
-    model_names = set(blob.name.split("/")[1] for blob in blobs)
-    def download_model_thread(model_name):
+    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 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:
-            download_and_verify_model(model_name)
+            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:
-            logger.error(f"Error downloading model '{model_name}': {e}")
-    threads = [threading.Thread(target=download_model_thread, args=(model_name,)) for model_name in model_names]
-    for thread in threads:
-        thread.start()
-    for thread in threads:
-        thread.join()
-
-@app.post("/predict/")
-async def predict(request: DownloadModelRequest):
-    model_name = request.model_name
-    pipeline_task = request.pipeline_task
-    input_text = request.input_text
-    model, tokenizer = load_model(model_name)
-    pipe = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
-    result = pipe(input_text)
-    return {"result": result}
+            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()
+
+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)
+
+@app.post("/generate")
+async def generate(request: Request, body: GenerateRequest):
+    try:
+        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 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
+            )
+
+            async def stream_text():
+                input_text = body.input_text
+                max_length = model.config.max_position_embeddings
+                generated_text = ""
+
+                while True:
+                    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
+
+                    output = model.generate(**inputs, generation_config=generation_config)
+                    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
+
+            if body.stream:
+                return StreamingResponse(stream_text(), media_type="text/plain")
+            else:
+                generated_text = ""
+                async for chunk in stream_text():
+                    generated_text += chunk
+                return {"result": generated_text}
+
+        elif body.task_type == "text-to-image":
+            generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=device)
+            image = generator(body.input_text)[0]
+            image_bytes = image.tobytes()
+            return Response(content=image_bytes, media_type="image/png")
+
+        elif body.task_type == "text-to-speech":
+            generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
+            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 body.task_type == "text-to-video":
+            try:
+                generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
+                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}")
+
+        else:
+            raise HTTPException(status_code=400, detail="Unsupported task type")
+
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 def download_all_models_in_background():
     models_url = "https://huggingface.co/api/models"
-    response = requests.get(models_url)
-    if response.status_code == 200:
+    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:
-            download_model_from_huggingface(model["id"])
+            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)