Update app.py

app.py

@@ -11,6 +11,7 @@ from google.auth import exceptions
 from transformers import pipeline
 from dotenv import load_dotenv
 import uvicorn
+import tempfile
 
 load_dotenv()
 
@@ -116,86 +117,24 @@ async def predict(request: DownloadModelRequest):
             logger.info(f"Modelos no encontrados en GCS, descargando '{model_prefix}' desde Hugging Face...")
             download_model_from_huggingface(model_prefix)
         
-        model_files_streams = {file: gcs_handler.download_file(f"{model_prefix}/{file}") for file in model_files if gcs_handler.file_exists(f"{model_prefix}/{file}")}
-        
-        config_stream = model_files_streams.get("config.json")
-        tokenizer_stream = model_files_streams.get("tokenizer.json")
-        model_stream = model_files_streams.get("pytorch_model.bin")
-        
-        if not config_stream or not tokenizer_stream or not model_stream:
-            logger.error(f"Faltan archivos necesarios para el modelo '{model_prefix}'.")
-            raise HTTPException(status_code=500, detail="Required model files missing.")
-        
-        if request.pipeline_task in ["text-generation", "translation", "summarization"]:
-            pipe = pipeline(request.pipeline_task, model=model_stream, tokenizer=tokenizer_stream)
-            result = pipe(request.input_text)
-            logger.info(f"Resultado generado para la tarea '{request.pipeline_task}': {result[0]}")
-            return {"response": result[0]}
-
-        elif request.pipeline_task == "image-generation":
-            try:
-                pipe = pipeline("image-generation", model=model_stream)
-                images = pipe(request.input_text)
-                image = images[0]
-                image_filename = f"{uuid.uuid4().hex}.png"
-                image_path = f"images/{image_filename}"
-                image.save(image_path)
-                
-                gcs_handler.upload_file(image_path, open(image_path, "rb"))
-                image_url = gcs_handler.generate_signed_url(image_path)
-                logger.info(f"Imagen generada y subida correctamente con URL: {image_url}")
-                return {"response": {"image_url": image_url}}
-            except Exception as e:
-                logger.error(f"Error generando la imagen: {e}")
-                raise HTTPException(status_code=400, detail="Error generando la imagen.")
-
-        elif request.pipeline_task == "image-editing":
-            try:
-                pipe = pipeline("image-editing", model=model_stream)
-                edited_images = pipe(request.input_text)
-                edited_image = edited_images[0]
-                edited_image_filename = f"{uuid.uuid4().hex}_edited.png"
-                edited_image.save(edited_image_filename)
-
-                gcs_handler.upload_file(f"images/{edited_image_filename}", open(edited_image_filename, "rb"))
-                edited_image_url = gcs_handler.generate_signed_url(f"images/{edited_image_filename}")
-                logger.info(f"Imagen editada y subida correctamente con URL: {edited_image_url}")
-                return {"response": {"edited_image_url": edited_image_url}}
-            except Exception as e:
-                logger.error(f"Error editando la imagen: {e}")
-                raise HTTPException(status_code=400, detail="Error editando la imagen.")
-
-        elif request.pipeline_task == "image-to-image":
-            try:
-                pipe = pipeline("image-to-image", model=model_stream)
-                transformed_images = pipe(request.input_text)
-                transformed_image = transformed_images[0]
-                transformed_image_filename = f"{uuid.uuid4().hex}_transformed.png"
-                transformed_image.save(transformed_image_filename)
-
-                gcs_handler.upload_file(f"images/{transformed_image_filename}", open(transformed_image_filename, "rb"))
-                transformed_image_url = gcs_handler.generate_signed_url(f"images/{transformed_image_filename}")
-                logger.info(f"Imagen transformada y subida correctamente con URL: {transformed_image_url}")
-                return {"response": {"transformed_image_url": transformed_image_url}}
-            except Exception as e:
-                logger.error(f"Error transformando la imagen: {e}")
-                raise HTTPException(status_code=400, detail="Error transformando la imagen.")
-
-        elif request.pipeline_task == "text-to-3d":
-            try:
-                model_3d_filename = f"{uuid.uuid4().hex}.obj"
-                model_3d_path = f"3d-models/{model_3d_filename}"
-                with open(model_3d_path, "w") as f:
-                    f.write("Simulated 3D model data")
-
-                gcs_handler.upload_file(f"3d-models/{model_3d_filename}", open(model_3d_path, "rb"))
-                model_3d_url = gcs_handler.generate_signed_url(f"3d-models/{model_3d_filename}")
-                logger.info(f"Modelo 3D generado y subido correctamente con URL: {model_3d_url}")
-                return {"response": {"model_3d_url": model_3d_url}}
-
-            except Exception as e:
-                logger.error(f"Error generando el modelo 3D: {e}")
-                raise HTTPException(status_code=400, detail="Error generando el modelo 3D.")
+        model_files_streams = {}
+        with tempfile.TemporaryDirectory() as temp_dir:
+            for file in model_files:
+                if gcs_handler.file_exists(f"{model_prefix}/{file}"):
+                    file_path = os.path.join(temp_dir, file)
+                    with open(file_path, "wb") as f:
+                        gcs_handler.download_file(f"{model_prefix}/{file}").readinto(f)
+                    model_files_streams[file] = file_path
+            
+            if not all(key in model_files_streams for key in ["config.json", "tokenizer.json", "pytorch_model.bin"]):
+                logger.error(f"Faltan archivos necesarios para el modelo '{model_prefix}'.")
+                raise HTTPException(status_code=500, detail="Required model files missing.")
+            
+            if request.pipeline_task in ["text-generation", "translation", "summarization"]:
+                pipe = pipeline(request.pipeline_task, model=model_files_streams["pytorch_model.bin"], tokenizer=model_files_streams["tokenizer.json"])
+                result = pipe(request.input_text)
+                logger.info(f"Resultado generado para la tarea '{request.pipeline_task}': {result[0]}")
+                return {"response": result[0]}
 
     except Exception as e:
         logger.error(f"Error en la predicción: {e}")