temp paralelo storage

app.py

@@ -6,17 +6,57 @@ from datasets import load_dataset
 import os
 import pickle
 from io import BytesIO
-from huggingface_hub import upload_file, hf_hub_download
+from huggingface_hub import upload_file, hf_hub_download, list_repo_files
 from pathlib import Path
 import gc
 import requests
+import time
+import shutil
 
 # 📁 Parámetros
 DATASET_ID = "Segizu/facial-recognition"
 EMBEDDINGS_SUBFOLDER = "embeddings"
+LOCAL_EMB_DIR = Path("embeddings")
+LOCAL_EMB_DIR.mkdir(exist_ok=True)
 HF_TOKEN = os.getenv("HF_TOKEN")
 headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
 
+# 💾 Configuración de control de almacenamiento
+MAX_TEMP_STORAGE_GB = 40
+UPLOAD_EVERY = 50
+embeddings_to_upload = []
+
+def get_folder_size(path):
+    total = 0
+    for dirpath, _, filenames in os.walk(path):
+        for f in filenames:
+            fp = os.path.join(dirpath, f)
+            total += os.path.getsize(fp)
+    return total / (1024 ** 3)  # En GB
+
+def flush_embeddings():
+    global embeddings_to_upload
+    print("🚀 Subiendo lote de embeddings a Hugging Face...")
+
+    for emb_file in embeddings_to_upload:
+        try:
+            filename = emb_file.name
+            upload_file(
+                path_or_fileobj=str(emb_file),
+                path_in_repo=f"{EMBEDDINGS_SUBFOLDER}/{filename}",
+                repo_id=DATASET_ID,
+                repo_type="dataset",
+                token=HF_TOKEN
+            )
+            os.remove(emb_file)
+            print(f"✅ Subido y eliminado: {filename}")
+            time.sleep(1.2)  # Evita 429
+        except Exception as e:
+            print(f"❌ Error subiendo {filename}: {e}")
+            continue
+
+    embeddings_to_upload = []
+
 # ✅ Cargar CSV desde el dataset
 dataset = load_dataset(
     "csv",
@@ -54,17 +94,17 @@ def build_database():
                 continue
 
             name = f"image_{i + j}"
-            filename = f"{name}.pkl"
+            filename = LOCAL_EMB_DIR / f"{name}.pkl"
 
-            # Verificar si ya está subido
+            # Verificar si ya existe en HF
             try:
                 hf_hub_download(
                     repo_id=DATASET_ID,
                     repo_type="dataset",
-                    filename=f"{EMBEDDINGS_SUBFOLDER}/{filename}",
+                    filename=f"{EMBEDDINGS_SUBFOLDER}/{name}.pkl",
                     token=HF_TOKEN
                 )
-                print(f"⏩ Ya existe remoto: {filename}")
+                print(f"⏩ Ya existe remoto: {name}.pkl")
                 continue
             except:
                 pass
@@ -81,19 +121,15 @@ def build_database():
                     enforce_detection=False
                 )[0]["embedding"]
 
-                # Guardar temporal y subir
+                # Guardar temporal
                 with open(filename, "wb") as f:
                     pickle.dump({"name": name, "img": img, "embedding": embedding}, f)
 
-                upload_file(
-                    path_or_fileobj=filename,
-                    path_in_repo=f"{EMBEDDINGS_SUBFOLDER}/{filename}",
-                    repo_id=DATASET_ID,
-                    repo_type="dataset",
-                    token=HF_TOKEN
-                )
-                os.remove(filename)
-                print(f"✅ Subido: {filename}")
+                embeddings_to_upload.append(filename)
+
+                # Si excede límites, subir batch
+                if get_folder_size(LOCAL_EMB_DIR) >= MAX_TEMP_STORAGE_GB or len(embeddings_to_upload) >= UPLOAD_EVERY:
+                    flush_embeddings()
 
                 del img_processed
                 gc.collect()
@@ -102,6 +138,10 @@ def build_database():
                 print(f"❌ Error en {name}: {e}")
                 continue
 
+    # Subir lo que quede
+    if embeddings_to_upload:
+        flush_embeddings()
+
 # 🔍 Buscar similitudes desde archivos remotos
 def find_similar_faces(uploaded_image: Image.Image):
     try:
@@ -119,14 +159,12 @@ def find_similar_faces(uploaded_image: Image.Image):
     similarities = []
 
     try:
-        # Obtener lista de archivos remotos
-        from huggingface_hub import list_repo_files
         embedding_files = [
             f for f in list_repo_files(DATASET_ID, repo_type="dataset", token=HF_TOKEN)
             if f.startswith(f"{EMBEDDINGS_SUBFOLDER}/") and f.endswith(".pkl")
         ]
     except Exception as e:
-        return [], f"⚠ Error obteniendo archivos del dataset: {str(e)}"
+        return [], f"⚠ Error obteniendo archivos: {str(e)}"
 
     for file_path in embedding_files:
         try: