|
import os |
|
import torch |
|
from fastapi import FastAPI, HTTPException |
|
from fastapi.responses import StreamingResponse |
|
from pydantic import BaseModel, field_validator |
|
from transformers import ( |
|
AutoConfig, |
|
pipeline, |
|
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 = 200 |
|
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, max_length=2048): |
|
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: |
|
yield "" |
|
|
|
generation_config.max_new_tokens = min(remaining_tokens, generation_config.max_new_tokens) |
|
|
|
def stop_criteria(input_ids, scores): |
|
decoded_output = tokenizer.decode(int(input_ids[0][-1]), skip_special_tokens=True) |
|
return decoded_output in stop_sequences |
|
|
|
stopping_criteria = StoppingCriteriaList([stop_criteria]) |
|
|
|
output_text = "" |
|
outputs = model.generate( |
|
**encoded_input, |
|
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 |
|
) |
|
|
|
for output in outputs.sequences: |
|
for token_id in output: |
|
token = tokenizer.decode(token_id, skip_special_tokens=True) |
|
yield token |
|
await asyncio.sleep(chunk_delay) |
|
|
|
if stop_sequences and any(stop in output_text for stop in stop_sequences): |
|
yield output_text |
|
return |
|
|
|
outputs = model.generate( |
|
**encoded_input, |
|
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 |
|
) |
|
|
|
@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) |
