Update app.py

app.py

@@ -309,7 +309,7 @@ import uuid
 from pathlib import Path
 from huggingface_hub import hf_hub_download
 from fastapi.responses import JSONResponse
-
+from huggingface_hub import snapshot_download
 # Constants
 REPO_ID = "rahul7star/ohamlab"
 FOLDER_IN_REPO = "filter-demo/upload_20250708_041329_9c5c81"
@@ -317,30 +317,26 @@ CONCEPT_SENTENCE = "ohamlab style"
 LORA_NAME = "ohami_filter_autorun"
 @app.get("/train-sample")
 def fetch_images_and_generate_captions():
-    # Create temporary local dir
+    # Create a unique local directory
     local_dir = Path(f"/tmp/{LORA_NAME}-{uuid.uuid4()}")
     os.makedirs(local_dir, exist_ok=True)
 
-    # Download all files from the given subfolder
-    hf_hub_download(
+    # Download all files from the dataset repo
+    snapshot_path = snapshot_download(
         repo_id=REPO_ID,
         repo_type="dataset",
-        subfolder=FOLDER_IN_REPO,
         local_dir=local_dir,
         local_dir_use_symlinks=False,
-        force_download=False,
-        etag_timeout=10,
+        allow_patterns=[f"{FOLDER_IN_REPO}/*"],  # only files inside the subfolder
     )
 
-    # Full path to the downloaded image folder
-    image_dir = local_dir / FOLDER_IN_REPO
+    # Resolve image path relative to downloaded snapshot
+    image_dir = Path(snapshot_path) / FOLDER_IN_REPO
     image_paths = list(image_dir.rglob("*.jpg")) + list(image_dir.rglob("*.jpeg")) + list(image_dir.rglob("*.png"))
 
-    # If no image found, return error
     if not image_paths:
         return JSONResponse(status_code=400, content={"error": "No images found in the HF repo folder."})
 
-    # Generate captions for each image
     captions = [
         f"Autogenerated caption for {img.stem} in the {CONCEPT_SENTENCE} [trigger]" for img in image_paths
     ]
@@ -369,7 +365,6 @@ def fetch_images_and_generate_captions():
 
 
 
-
 
 
 REPO_ID = "rahul7star/ohamlab"
@@ -512,18 +507,18 @@ def auto_run_lora_from_repo():
         local_dir = Path(f"/tmp/{LORA_NAME}-{uuid.uuid4()}")
         os.makedirs(local_dir, exist_ok=True)
 
-        hf_hub_download(
-            repo_id=REPO_ID,
-            repo_type="dataset",
-            subfolder=FOLDER_IN_REPO,
-            local_dir=local_dir,
-            local_dir_use_symlinks=False,
-            force_download=False,
-            etag_timeout=10,
-           
-        )
+        snapshot_path = snapshot_download(
+        repo_id=REPO_ID,
+        repo_type="dataset",
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,
+        allow_patterns=[f"{FOLDER_IN_REPO}/*"],  # only files inside the subfolder
+         )
+
+        # Resolve image path relative to downloaded snapshot
+        image_dir = Path(snapshot_path) / FOLDER_IN_REPO
 
-        image_dir = local_dir / FOLDER_IN_REPO
+       
         image_paths = list(image_dir.rglob("*.jpg")) + list(image_dir.rglob("*.jpeg")) + list(image_dir.rglob("*.png"))
 
         if not image_paths: