app.py

import os
import logging
import threading
import boto3
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, StoppingCriteriaList
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, field_validator
from io import BytesIO
from huggingface_hub import hf_hub_download
import requests
import asyncio
import soundfile as sf
import numpy as np
from fastapi.responses import StreamingResponse, Response

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")

app = FastAPI()

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
    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
        self.s3_client = s3_client

    def _get_s3_uri(self, model_name):
        return f"s3://{self.bucket_name}/lilmeaty_garca/{model_name.replace('/', '-')}"
    
    def _download_from_s3(self, model_name):
        try:
            logging.info(f"Attempting to load model {model_name} from S3...")
            model_files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=f"lilmeaty_garca/{model_name}")
            if "Contents" not in model_files:
                raise FileNotFoundError(f"Model files not found in S3 for {model_name}")

            local_dir = f"/tmp/{model_name.replace('/', '-')}"
            os.makedirs(local_dir, exist_ok=True)

            for obj in model_files["Contents"]:
                file_key = obj["Key"]
                if file_key.endswith('/'):
                    continue

                local_file_path = os.path.join(local_dir, os.path.basename(file_key))
                self.s3_client.download_file(self.bucket_name, file_key, local_file_path)

            return local_dir
        except Exception as e:
            logging.error(f"Error downloading from S3: {e}")
            raise HTTPException(status_code=500, detail=f"Error downloading model from S3: {e}")

    async def load_model_and_tokenizer(self, model_name):
        try:
            model_dir = await asyncio.to_thread(self._download_from_s3, model_name)
            config = AutoConfig.from_pretrained(model_dir)
            tokenizer = AutoTokenizer.from_pretrained(model_dir, config=config)
            model = AutoModelForCausalLM.from_pretrained(model_dir, config=config)

            logging.info(f"Model {model_name} loaded from S3 successfully.")
            return model, tokenizer
        except Exception as e:
            logging.exception(f"Error loading model: {e}")
            raise HTTPException(status_code=500, detail=f"Error loading model: {e}")

    def download_model_from_huggingface(self, model_name):
        try:
            logging.info(f"Downloading model {model_name} from Hugging Face...")
            model_dir = hf_hub_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
            self.s3_client.upload_file(model_dir, self.bucket_name, f"lilmeaty_garca/{model_name}")
            logging.info(f"Model {model_name} saved to S3 successfully.")
        except Exception as e:
            logging.error(f"Error downloading model {model_name} from Hugging Face: {e}")

    def download_all_models_in_background(self):
        models_url = "https://huggingface.co/api/models"
        try:
            response = requests.get(models_url)
            if response.status_code != 200:
                logging.error("Error getting Hugging Face model list.")
                raise HTTPException(status_code=500, detail="Error getting model list.")

            models = response.json()
            for model in models:
                model_name = model["id"]
                self.download_model_from_huggingface(model_name)
        except Exception as e:
            logging.error(f"Error downloading models in the background: {e}")
            raise HTTPException(status_code=500, detail="Error downloading models in the background.")

    def run_in_background(self):
        threading.Thread(target=self.download_all_models_in_background, daemon=True).start()

@app.on_event("startup")
async def startup_event():
    model_loader.run_in_background()

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:
        validated_body = GenerateRequest(**body.model_dump())
        model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)

        if validated_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
            )

            async def stream_text():
                input_text = validated_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, 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)
                    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 validated_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 validated_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_bytes = image.tobytes()
            return Response(content=image_bytes, media_type="image/png")

        elif validated_body.task_type == "text-to-speech":
            generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
            audio = generator(validated_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":
            try:
                generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
                video = generator(validated_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:
        logging.exception(f"An unexpected error occurred: {e}")
        raise HTTPException(status_code=500, detail="An unexpected error occurred.")