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, ) import boto3 import uvicorn import asyncio from transformers import pipeline import json from huggingface_hub import login import base64 from botocore.exceptions import NoCredentialsError 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") if HUGGINGFACE_HUB_TOKEN: login(token=HUGGINGFACE_HUB_TOKEN, add_to_git_credential=False) 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 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}/" \ f"{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=False ) model = AutoModelForCausalLM.from_pretrained( s3_uri, config=config, local_files_only=False ) tokenizer = AutoTokenizer.from_pretrained( s3_uri, config=config, local_files_only=False ) 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, NoCredentialsError): try: config = AutoConfig.from_pretrained( model_name, token=HUGGINGFACE_HUB_TOKEN ) tokenizer = AutoTokenizer.from_pretrained( model_name, config=config, token=HUGGINGFACE_HUB_TOKEN ) model = AutoModelForCausalLM.from_pretrained( model_name, config=config, token=HUGGINGFACE_HUB_TOKEN ) 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) config = AutoConfig.from_pretrained( s3_uri, local_files_only=False ) model = AutoModelForCausalLM.from_pretrained( s3_uri, config=config, local_files_only=False ) tokenizer = AutoTokenizer.from_pretrained( s3_uri, config=config, local_files_only=False ) 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 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) if "text-to-text" == task_type: 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), media_type="text/plain" ) else: return HTTPException(status_code=400, detail="Task type not text-to-text") 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): max_length=10 #Define the max length to cut the text and generate another response encoded_input = tokenizer( input_text, return_tensors="pt", truncation=True, max_length=max_length ).to(device) def find_stop(output_text, stop_sequences): for seq in stop_sequences: if seq in output_text: last_index = output_text.rfind(seq) return last_index + len(seq) return -1 output_text = "" while True: input_length = encoded_input["input_ids"].shape[1] remaining_tokens = max_length - input_length if remaining_tokens <=0: yield json.dumps({"text": "", "is_end": True}) + "\n" break generation_config.max_new_tokens = min( remaining_tokens, generation_config.max_new_tokens ) 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, output_scores=True, return_dict_in_generate=True, ) new_text = tokenizer.decode( outputs.sequences[0][len(encoded_input["input_ids"][0]):], skip_special_tokens=True ) output_text += new_text stop_index = find_stop(output_text, stop_sequences) if stop_index != -1: final_output = output_text[:stop_index] for chunk in [final_output[i:i+10] for i in range(0,len(final_output),10)]: for text in chunk.split(): yield json.dumps({"text": text, "is_end": False}) + "\n" yield json.dumps({"text": "", "is_end": True}) + "\n" break else: for chunk in [new_text[i:i+10] for i in range(0, len(new_text), 10)]: for text in chunk.split(): yield json.dumps({"text": text, "is_end": False}) + "\n" if len(output_text) >= max_length: encoded_input = tokenizer( output_text, return_tensors="pt", truncation=True, max_length=max_length ).to(device) output_text = "" elif len(output_text) < max_length and len(new_text) == 0: for chunk in [output_text[i:i+10] for i in range(0, len(output_text), 10)]: for text in chunk.split(): yield json.dumps({"text": text, "is_end": False}) + "\n" yield json.dumps({"text": "", "is_end": True}) + "\n" break @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] image_data = list(image.getdata()) return json.dumps({"image_data": image_data, "is_end": True}) 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) audio_bytes = audio["audio"] audio_base64 = base64.b64encode(audio_bytes).decode('utf-8') return json.dumps({"audio": audio_base64, "is_end": True}) 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) video_base64 = base64.b64encode(video).decode('utf-8') return json.dumps({"video": video_base64, "is_end": True}) 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)