Update app.py

app.py

@@ -504,32 +504,47 @@ def start_training(
 @app.post("/train-from-hf")
 def auto_run_lora_from_repo():
     try:
+        # Set HF cache path if not already set
+        os.environ["HF_HOME"] = "/tmp/hf_cache"
+        os.makedirs("/tmp/hf_cache", exist_ok=True)
+
+        # Create temporary directory to hold downloaded files
         local_dir = Path(f"/tmp/{LORA_NAME}-{uuid.uuid4()}")
         os.makedirs(local_dir, exist_ok=True)
 
+        # Download snapshot from model repo using allow_patterns
         snapshot_path = snapshot_download(
-        repo_id=REPO_ID,
-        repo_type="model",
-        local_dir=local_dir,
-        local_dir_use_symlinks=False,
-        allow_patterns=[f"{FOLDER_IN_REPO}/*"],  # only files inside the subfolder
-         )
+            repo_id=REPO_ID,
+            repo_type="model",
+            local_dir=local_dir,
+            local_dir_use_symlinks=False,
+            allow_patterns=[f"{FOLDER_IN_REPO}/*"],  # only that folder
+        )
 
-        # Resolve image path relative to downloaded snapshot
+        # Target subfolder inside the 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"))
+        # Collect all image files (recursively)
+        image_paths = list(image_dir.rglob("*.jpg")) + \
+                      list(image_dir.rglob("*.jpeg")) + \
+                      list(image_dir.rglob("*.png"))
 
         if not image_paths:
-            return JSONResponse(status_code=400, content={"error": "No images found in the HF repo folder."})
+            return JSONResponse(
+                status_code=400,
+                content={"error": "No images found in the HF repo folder."}
+            )
 
+        # Create auto captions
         captions = [
-            f"Autogenerated caption for {img.stem} in the {CONCEPT_SENTENCE} [trigger]" for img in image_paths
+            f"Autogenerated caption for {img.stem} in the {CONCEPT_SENTENCE} [trigger]"
+            for img in image_paths
         ]
 
+        # Prepare dataset
         dataset_path = create_dataset(image_paths, *captions)
 
+        # Start training
         result = start_training(
             lora_name=LORA_NAME,
             concept_sentence=CONCEPT_SENTENCE,
@@ -556,5 +571,7 @@ augmentation:
 
         return {"message": result}
 
+    except PermissionError as pe:
+        return JSONResponse(status_code=500, content={"error": f"Permission denied: {pe}"})
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
+        return JSONResponse(status_code=500, content={"error": str(e)})