app.py

import os
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, field_validator
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
    StoppingCriteriaList
)
import boto3
import uvicorn
import asyncio
from io import BytesIO
from transformers import pipeline

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")

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
    temperature: float = 1.0
    max_new_tokens: int = 10
    stream: bool = True
    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}/{model_name.replace('/', '-')}"

    async def load_model_and_tokenizer(self, model_name):
        s3_uri = self._get_s3_uri(model_name)
        try:
            config = AutoConfig.from_pretrained(s3_uri, local_files_only=True)
            model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=True)
            tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=True)

            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

            return model, tokenizer
        except EnvironmentError:
            try:
                config = AutoConfig.from_pretrained(model_name)
                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
                model = AutoModelForCausalLM.from_pretrained(model_name, config=config)

                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

                model.save_pretrained(s3_uri)
                tokenizer.save_pretrained(s3_uri)
                return model, tokenizer
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Error loading model: {e}")

model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)

@app.post("/generate")
async def generate(request: GenerateRequest):
    try:
        model_name = request.model_name
        input_text = request.input_text
        task_type = request.task_type
        temperature = request.temperature
        max_new_tokens = request.max_new_tokens
        stream = request.stream
        top_p = request.top_p
        top_k = request.top_k
        repetition_penalty = request.repetition_penalty
        num_return_sequences = request.num_return_sequences
        do_sample = request.do_sample
        chunk_delay = request.chunk_delay
        stop_sequences = request.stop_sequences

        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)

        generation_config = GenerationConfig(
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty,
            do_sample=do_sample,
            num_return_sequences=num_return_sequences,
        )

        return StreamingResponse(
            stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay),
            media_type="text/plain"
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay):
    # Get the maximum model input length
    max_model_length = model.config.max_position_embeddings

    encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True).to(device)

    def stop_criteria(input_ids, scores):
        decoded_output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
        for stop in stop_sequences:
            if decoded_output.endswith(stop):
                return True
        return False

    stopping_criteria = StoppingCriteriaList([stop_criteria])

    token_buffer = []
    output_ids = encoded_input.input_ids
    while True:
        try:
            outputs = model.generate(
                output_ids,
                do_sample=generation_config.do_sample,
                max_new_tokens=generation_config.max_new_tokens,
                temperature=generation_config.temperature,
                top_p=generation_config.top_p,
                top_k=generation_config.top_k,
                repetition_penalty=generation_config.repetition_penalty,
                num_return_sequences=generation_config.num_return_sequences,
                stopping_criteria=stopping_criteria,
                output_scores=True,
                return_dict_in_generate=True,
                pad_token_id=tokenizer.pad_token_id
            )
        except IndexError as e:
            print(f"IndexError during generation: {e}")
            break  # Exit the loop if there's an index error

        new_token_ids = outputs.sequences[0][encoded_input.input_ids.shape[-1]:]

        for token_id in new_token_ids:
            token = tokenizer.decode(token_id, skip_special_tokens=True)
            token_buffer.append(token)
            if len(token_buffer) >= 10:
                yield "".join(token_buffer)
                token_buffer = []
            await asyncio.sleep(chunk_delay)

        if token_buffer:
            yield "".join(token_buffer)
            token_buffer = []

        if stop_criteria(outputs.sequences, None):
            break

        if len(new_token_ids) < generation_config.max_new_tokens:
            break

        output_ids = outputs.sequences

@app.post("/generate-image")
async def generate_image(request: GenerateRequest):
    try:
        validated_body = request
        device = "cuda" if torch.cuda.is_available() else "cpu"

        image_generator = pipeline("text-to-image", model=validated_body.model_name, device=device)
        image = image_generator(validated_body.input_text)[0]

        img_byte_arr = BytesIO()
        image.save(img_byte_arr, format="PNG")
        img_byte_arr.seek(0)

        return StreamingResponse(img_byte_arr, media_type="image/png")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.post("/generate-text-to-speech")
async def generate_text_to_speech(request: GenerateRequest):
    try:
        validated_body = request
        device = "cuda" if torch.cuda.is_available() else "cpu"

        audio_generator = pipeline("text-to-speech", model=validated_body.model_name, device=device)
        audio = audio_generator(validated_body.input_text)[0]

        audio_byte_arr = BytesIO()
        audio.save(audio_byte_arr)
        audio_byte_arr.seek(0)

        return StreamingResponse(audio_byte_arr, media_type="audio/wav")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.post("/generate-video")
async def generate_video(request: GenerateRequest):
    try:
        validated_body = request
        device = "cuda" if torch.cuda.is_available() else "cpu"
        video_generator = pipeline("text-to-video", model=validated_body.model_name, device=device)
        video = video_generator(validated_body.input_text)[0]

        video_byte_arr = BytesIO()
        video.save(video_byte_arr)
        video_byte_arr.seek(0)

        return StreamingResponse(video_byte_arr, media_type="video/mp4")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)