Update app.py

app.py

@@ -1,182 +1,164 @@
 import os
-import json
-from fastapi import FastAPI, HTTPException
+import logging
+import time
+from io import BytesIO
+
+from fastapi import FastAPI, HTTPException, Response, Request
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    GenerationConfig
+)
 import boto3
-import logging
 from huggingface_hub import hf_hub_download
+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 - %(message)s")
 
-# Configuración de AWS y Hugging Face
 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")
 
-# Cliente de S3
-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
     task_type: str
-
-class S3DirectStream:
-    def __init__(self, bucket_name):
-        self.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
-        )
+    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
+    chunk_delay: float = 0.0
+
+class S3ModelLoader:
+    def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
+        self.s3_client = s3_client
 
-    async def download_and_upload_to_s3(self, model_name):
-        try:
-            model_name = model_name.replace("/", "-").lower()
-
-            # Descargar el archivo config.json desde Hugging Face
-            config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN)
-            tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN)
-
-            # Verificar si la carpeta y los archivos ya existen en S3
-            if not await self.file_exists_in_s3(f"{model_name}/config.json"):
-                logging.info(f"El archivo config.json no existe en S3. Subiendo desde Hugging Face...")
-                self.create_folder_if_not_exists(model_name)
-                with open(config_file, "rb") as file:
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file)
-
-            if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
-                logging.info(f"El archivo tokenizer.json no existe en S3. Subiendo desde Hugging Face...")
-                self.create_folder_if_not_exists(model_name)
-                with open(tokenizer_file, "rb") as file:
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file)
-
-        except Exception as e:
-            logging.error(f"Error al cargar el modelo desde Hugging Face a S3: {e}")
-            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo: {str(e)}")
-
-    async def file_exists_in_s3(self, s3_key):
-        try:
-            self.s3_client.head_object(Bucket=self.bucket_name, Key=s3_key)
-            return True
-        except self.s3_client.exceptions.ClientError:
-            return False
+    def _get_s3_uri(self, model_name):
+        return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
 
-    def create_folder_if_not_exists(self, model_name):
-        try:
-            # Las carpetas no existen como tal en S3, pero se pueden crear archivos vacíos para simular carpetas
-            # Crear un archivo vacío para simular la carpeta
-            self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/")
-        except Exception as e:
-            logging.error(f"Error al crear la carpeta en S3: {e}")
-            raise HTTPException(status_code=500, detail=f"Error al crear la carpeta en S3: {str(e)}")
-
-    async def load_model_from_s3(self, model_name):
-        try:
-            model_name = model_name.replace("/", "-").lower()
-            model_files = await self.get_model_file_parts(model_name)
-
-            if not model_files:
-                await self.download_and_upload_to_s3(model_name)
-
-            # Cargar configuración del modelo desde S3
-            config_data = await self.stream_from_s3(f"{model_name}/config.json")
-            if isinstance(config_data, bytes):
-                config_data = config_data.decode("utf-8")
-            
-            config_json = json.loads(config_data)
-
-            # Cargar el modelo
-            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json)
-            return model
-
-        except HTTPException as e:
-            raise e
-        except Exception as e:
-            logging.error(f"Error al cargar el modelo desde S3: {e}")
-            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {str(e)}")
-
-    async def load_tokenizer_from_s3(self, model_name):
-        try:
-            model_name = model_name.replace("/", "-").lower()
-            tokenizer_data = await self.stream_from_s3(f"{model_name}/tokenizer.json")
-
-            if isinstance(tokenizer_data, bytes):
-                tokenizer_data = tokenizer_data.decode("utf-8") 
-            
-            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
-            return tokenizer
-        except Exception as e:
-            logging.error(f"Error al cargar el tokenizer desde S3: {e}")
-            raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {str(e)}")
-
-    async def stream_from_s3(self, key):
+    async def load_model_and_tokenizer(self, model_name):
+        s3_uri = self._get_s3_uri(model_name)
         try:
-            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            return response['Body'].read()
-        except self.s3_client.exceptions.NoSuchKey:
-            raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}")
-
-    async def get_model_file_parts(self, model_name):
-        try:
-            model_name = model_name.replace("/", "-").lower()
-            files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
-            model_files = [obj['Key'] for obj in files.get('Contents', []) if model_name in obj['Key']]
-            return model_files
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {str(e)}")
+            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)
+            logging.info(f"Loaded {model_name} from S3 successfully.")
+            return model, tokenizer
+        except EnvironmentError:
+            logging.info(f"Model {model_name} not found in S3. Downloading...")
+            try:
+                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.")
+                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.error(f"Error downloading/uploading model: {e}")
+                raise HTTPException(status_code=500, detail=f"Error loading model: {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: GenerateRequest):
+async def generate(request: Request, body: GenerateRequest):
     try:
-        model_name = request.model_name
-        input_text = request.input_text
-        task_type = request.task_type
-        
-        s3_direct_stream = S3DirectStream(S3_BUCKET_NAME)
-        
-        model = await s3_direct_stream.load_model_from_s3(model_name)
-        tokenizer = await s3_direct_stream.load_tokenizer_from_s3(model_name)
-        
-        if task_type == "text-to-text":
-            generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
-            result = generator(input_text, max_length=1024, num_return_sequences=1)
-            return {"result": result[0]["generated_text"]}
-
-        elif task_type == "text-to-image":
-            generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=0)
-            image = generator(input_text)
-            return {"result": image}
-
-        elif task_type == "text-to-speech":
-            generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=0)
-            audio = generator(input_text)
-            return {"result": audio}
-
-        elif task_type == "text-to-video":
-            generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=0)
-            video = generator(input_text)
-            return {"result": video}
+        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
+                generated_text = ""
+                max_length = model.config.max_position_embeddings
+
+                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]
+                    remaining_tokens = max_length - input_length
+
+                    if remaining_tokens <= 0:
+                        break
+
+                    generation_config.max_new_tokens = min(remaining_tokens, body.max_new_tokens)
+
+                    output = model.generate(**encoded_input, generation_config=generation_config)
+                    chunk = tokenizer.decode(output[0], skip_special_tokens=True)
+                    generated_text += chunk
+                    yield chunk
+                    time.sleep(body.chunk_delay)
+                    input_text = generated_text
+
+            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="Tipo de tarea no soportada")
+            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))
 
+
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)