Update app.py

app.py

@@ -1,15 +1,19 @@
 import os
 import json
+import uuid
 import requests
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from google.cloud import storage
 from google.auth import exceptions
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from io import BytesIO
 from dotenv import load_dotenv
 import uvicorn
 import tempfile
+from PIL import Image
+import soundfile as sf
+import torch
 
 load_dotenv()
 
@@ -49,6 +53,10 @@ class GCSHandler:
             raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
         return BytesIO(blob.download_as_bytes())
 
+    def generate_signed_url(self, blob_name, expiration=3600):
+        blob = self.bucket.blob(blob_name)
+        return blob.generate_signed_url(expiration=expiration)
+
 def download_model_from_huggingface(model_name):
     url = f"https://huggingface.co/{model_name}/tree/main"
     headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
@@ -109,16 +117,72 @@ async def predict(request: DownloadModelRequest):
                 f.write(tokenizer_stream.read())
             with open(model_path, 'wb') as f:
                 f.write(model_stream.read())
-            
-            model = AutoModelForCausalLM.from_pretrained(tmp_dir, from_tf=True)
+
+            model = AutoModelForCausalLM.from_pretrained(tmp_dir)
             tokenizer = AutoTokenizer.from_pretrained(tmp_dir)
+
+            pipe = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
+
+            if request.pipeline_task in ["text-generation", "translation", "summarization"]:
+                result = pipe(request.input_text)
+                return {"response": result[0]}
+
+            elif request.pipeline_task == "image-generation":
+                try:
+                    images = pipe(request.input_text)
+                    image = images[0]
+                    image_filename = f"{uuid.uuid4().hex}.png"
+                    image_path = os.path.join(tmp_dir, image_filename)
+                    image.save(image_path)
+
+                    gcs_handler.upload_file(f"images/{image_filename}", open(image_path, "rb"))
+                    image_url = gcs_handler.generate_signed_url(f"images/{image_filename}")
+                    return {"response": {"image_url": image_url}}
+                except Exception as e:
+                    raise HTTPException(status_code=400, detail="Error al generar la imagen.")
+
+            elif request.pipeline_task == "image-editing":
+                try:
+                    edited_images = pipe(request.input_text)
+                    edited_image = edited_images[0]
+                    edited_image_filename = f"{uuid.uuid4().hex}_edited.png"
+                    edited_image_path = os.path.join(tmp_dir, edited_image_filename)
+                    edited_image.save(edited_image_path)
+
+                    gcs_handler.upload_file(f"images/{edited_image_filename}", open(edited_image_path, "rb"))
+                    edited_image_url = gcs_handler.generate_signed_url(f"images/{edited_image_filename}")
+                    return {"response": {"edited_image_url": edited_image_url}}
+                except Exception as e:
+                    raise HTTPException(status_code=400, detail="Error al editar la imagen.")
+
+            elif request.pipeline_task == "image-to-image":
+                try:
+                    transformed_images = pipe(request.input_text)
+                    transformed_image = transformed_images[0]
+                    transformed_image_filename = f"{uuid.uuid4().hex}_transformed.png"
+                    transformed_image_path = os.path.join(tmp_dir, transformed_image_filename)
+                    transformed_image.save(transformed_image_path)
+
+                    gcs_handler.upload_file(f"images/{transformed_image_filename}", open(transformed_image_path, "rb"))
+                    transformed_image_url = gcs_handler.generate_signed_url(f"images/{transformed_image_filename}")
+                    return {"response": {"transformed_image_url": transformed_image_url}}
+                except Exception as e:
+                    raise HTTPException(status_code=400, detail="Error al transformar la imagen.")
+
+            elif request.pipeline_task == "text-to-3d":
+                try:
+                    model_3d_filename = f"{uuid.uuid4().hex}.obj"
+                    model_3d_path = os.path.join(tmp_dir, 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}")
+                    return {"response": {"model_3d_url": model_3d_url}}
+                except Exception as e:
+                    raise HTTPException(status_code=400, detail="Error al generar el modelo 3D.")
             
-            pipeline_ = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
-            
-            result = pipeline_(request.input_text)
-        
-        return {"response": result}
-    
     except HTTPException as e:
         raise e
     except Exception as e: