Update app.py

app.py

@@ -1,29 +1,19 @@
 import os
-import logging
-import time
-from io import BytesIO
-from typing import Union
-
-from fastapi import FastAPI, HTTPException, Response, Request, UploadFile, File
+import torch
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, ValidationError, field_validator
+from pydantic import BaseModel, field_validator
 from transformers import (
-    AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
-    pipeline,
     GenerationConfig,
     StoppingCriteriaList
 )
 import boto3
-from huggingface_hub import hf_hub_download, HfApi
-import soundfile as sf
-import numpy as np
-import torch
 import uvicorn
-import shutil
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
+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")
@@ -31,13 +21,17 @@ 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 = False
+    stream: bool = True
     top_p: float = 1.0
     top_k: int = 50
     repetition_penalty: float = 1.0
@@ -46,8 +40,6 @@ class GenerateRequest(BaseModel):
     chunk_delay: float = 0.0
     stop_sequences: list[str] = []
 
-    model_config = {"protected_namespaces": ()}
-
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
         if not v:
@@ -65,7 +57,6 @@ class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
         self.s3_client = s3_client
-        self.api = HfApi()
 
     def _get_s3_uri(self, model_name):
         return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
@@ -73,7 +64,6 @@ class S3ModelLoader:
     async def load_model_and_tokenizer(self, model_name):
         s3_uri = self._get_s3_uri(model_name)
         try:
-            logging.info(f"Trying to load {model_name} from S3...")
             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)
@@ -81,130 +71,170 @@ class S3ModelLoader:
             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
 
-            logging.info(f"Loaded {model_name} from S3 successfully.")
             return model, tokenizer
         except EnvironmentError:
-            logging.info(f"Model {model_name} not found in S3. Downloading...")
             try:
-                model_info = self.api.model_info(model_name)
-                files_to_download = [f.filename for f in self.api.list_repo_files(model_name)]
-                
-                temp_dir = "temp_model"
-                os.makedirs(temp_dir, exist_ok=True)
-                
-                for file_name in files_to_download:
-                    hf_hub_download(repo_id=model_name, filename=file_name, local_dir=temp_dir, token=HUGGINGFACE_HUB_TOKEN)
-
-                config = AutoConfig.from_pretrained(temp_dir)
-                tokenizer = AutoTokenizer.from_pretrained(temp_dir, config=config)
-                model = AutoModelForCausalLM.from_pretrained(temp_dir, config=config)
+                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
 
-                logging.info(f"Downloaded {model_name} successfully.")
-                logging.info(f"Saving {model_name} to S3...")
                 model.save_pretrained(s3_uri)
                 tokenizer.save_pretrained(s3_uri)
-                logging.info(f"Saved {model_name} to S3 successfully.")
-
-                shutil.rmtree(temp_dir)
-
                 return model, tokenizer
             except Exception as e:
-                logging.exception(f"Error downloading/uploading model: {e}")
                 raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
-app = FastAPI()
-
-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):
+async def generate(request: GenerateRequest):
     try:
-        validated_body = GenerateRequest(**body.model_dump())
-        model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name)
+        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)
 
-        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}")
+        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"
+        )
 
-        else:
-            raise HTTPException(status_code=400, detail="Unsupported task type")
+    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 HTTPException as e:
-        raise e
-    except ValidationError as e:
-        raise HTTPException(status_code=422, detail=e.errors())
     except Exception as e:
-        logging.exception(f"An unexpected error occurred: {e}")
-        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+        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)