Create app.py

app.py

@@ -0,0 +1,242 @@
+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)  # Simula el delay entre tokens
+
+        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)