Update app.py

app.py

@@ -11,10 +11,11 @@ from transformers import (
     pipeline
 )
 from io import BytesIO
-import asyncio
 import boto3
 from botocore.exceptions import NoCredentialsError
 from huggingface_hub import snapshot_download
+import asyncio
+import shutil
 
 # Configuración global
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
@@ -65,28 +66,36 @@ class S3ModelLoader:
         
         s3_uri = self._get_s3_uri(model_name)
         try:
-            # Descargar el modelo desde Hugging Face y guardarlo en S3 si no existe
-            model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
-
-            # Cargar el modelo y tokenizer
-            model = AutoModelForCausalLM.from_pretrained(model_path)
-            tokenizer = AutoTokenizer.from_pretrained(model_path)
-
-            # Asignar EOS y PAD token si no están definidos
-            if tokenizer.eos_token_id is None:
-                tokenizer.eos_token_id = tokenizer.pad_token_id
-
-            # Guardar el modelo y el tokenizer en el diccionario
-            token_dict[model_name] = {
-                "model": model,
-                "tokenizer": tokenizer,
-                "pad_token_id": tokenizer.pad_token_id,
-                "eos_token_id": tokenizer.eos_token_id
-            }
-
-            # Subir los archivos del modelo y tokenizer a S3 si no están allí
-            self.s3_client.upload_file(model_path, self.bucket_name, f'{model_name}/model')
-            self.s3_client.upload_file(f'{model_path}/tokenizer', self.bucket_name, f'{model_name}/tokenizer')
+            # Verificar si el modelo ya está en S3
+            try:
+                self.s3_client.head_object(Bucket=self.bucket_name, Key=f'{model_name}/model')
+                print(f"Modelo {model_name} ya existe en S3.")
+            except self.s3_client.exceptions.ClientError:
+                print(f"Modelo {model_name} no existe en S3. Descargando desde Hugging Face...")
+                model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
+
+                # Cargar el modelo y tokenizer
+                model = AutoModelForCausalLM.from_pretrained(model_path)
+                tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+                # Asignar EOS y PAD token si no están definidos
+                if tokenizer.eos_token_id is None:
+                    tokenizer.eos_token_id = tokenizer.pad_token_id
+
+                # Guardar el modelo y el tokenizer en el diccionario
+                token_dict[model_name] = {
+                    "model": model,
+                    "tokenizer": tokenizer,
+                    "pad_token_id": tokenizer.pad_token_id,
+                    "eos_token_id": tokenizer.eos_token_id
+                }
+
+                # Subir los archivos del modelo y tokenizer a S3
+                self.s3_client.upload_file(model_path, self.bucket_name, f'{model_name}/model')
+                self.s3_client.upload_file(f'{model_path}/tokenizer', self.bucket_name, f'{model_name}/tokenizer')
+
+                # Eliminar los archivos locales después de haber subido a S3
+                shutil.rmtree(model_path)
 
             return token_dict[model_name]
         except NoCredentialsError:
@@ -232,7 +241,7 @@ async def generate_video(request: GenerateRequest):
         video = video_generator(validated_body.input_text)[0]
 
         video_byte_arr = BytesIO()
-        video.save(video_byte_arr)
+        video.save(video_byte_arr, format="MP4")
         video_byte_arr.seek(0)
 
         return StreamingResponse(video_byte_arr, media_type="video/mp4")
@@ -240,7 +249,7 @@ async def generate_video(request: GenerateRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-# Configuración para ejecutar el servidor
+# Ejecutar el servidor FastAPI con Uvicorn
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)