Update app.py

app.py

@@ -45,8 +45,7 @@ class GCSHandler:
         self.bucket = storage_client.bucket(bucket_name)
 
     def file_exists(self, blob_name):
-        exists = self.bucket.blob(blob_name).exists()
-        return exists
+        return self.bucket.blob(blob_name).exists()
 
     def download_file(self, blob_name):
         blob = self.bucket.blob(blob_name)
@@ -60,8 +59,11 @@ class GCSHandler:
 
     def generate_signed_url(self, blob_name, expiration=3600):
         blob = self.bucket.blob(blob_name)
-        url = blob.generate_signed_url(expiration=expiration)
-        return url
+        return blob.generate_signed_url(expiration=expiration)
+
+    def create_folder(self, folder_name):
+        blob = self.bucket.blob(folder_name + "/")
+        blob.upload_from_string("")  # Create an empty "folder"
 
 def load_model_from_gcs(model_name: str, model_files: list):
     gcs_handler = GCSHandler(GCS_BUCKET_NAME)
@@ -71,7 +73,7 @@ def load_model_from_gcs(model_name: str, model_files: list):
     config_stream = model_blobs.get("config.json")
     tokenizer_stream = model_blobs.get("tokenizer.json")
     
-    if "safetensors" in model_stream.name:
+    if model_stream and model_stream.endswith(".safetensors"):
         model = load_safetensors_model(model_stream)
     else:
         model = AutoModelForCausalLM.from_pretrained(io.BytesIO(model_stream), config=config_stream)
@@ -122,18 +124,20 @@ def download_model_from_huggingface(model_name):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
 
+def download_model_files(model_name: str):
+    model_files = get_model_files_from_gcs(model_name)
+    if not model_files:
+        download_model_from_huggingface(model_name)
+        model_files = get_model_files_from_gcs(model_name)
+    return model_files
+
 @app.post("/predict/")
 async def predict(request: DownloadModelRequest):
     try:
         gcs_handler = GCSHandler(GCS_BUCKET_NAME)
         model_prefix = request.model_name
         
-        model_files = get_model_files_from_gcs(model_prefix)
-        
-        if not model_files:
-            download_model_from_huggingface(model_prefix)
-            model_files = get_model_files_from_gcs(model_prefix)
-
+        model_files = download_model_files(model_prefix)
         model, tokenizer = load_model_from_gcs(model_prefix, model_files)
         
         pipe = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)