|
import os |
|
import logging |
|
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 tqdm import tqdm |
|
|
|
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") |
|
AWS_REGION = os.getenv("AWS_REGION") |
|
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME") |
|
HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") |
|
|
|
class GenerateRequest(BaseModel): |
|
model_name: 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 |
|
|
|
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}") |
|
|
|
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: |
|
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() |
|
|
|
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" |
|
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=8000) |
|
|
