Update app.py

app.py

@@ -35,7 +35,7 @@ import json
 import yaml
 from slugify import slugify
 
-
+from flux_train import start_training
 
 
 
@@ -525,8 +525,119 @@ def start_training(
     return f"Training started successfully with config: {config_path}"
 
 # ========== MAIN ENDPOINT ==========
+
+
+
+
+
+
+
+
+import os
+import uuid
+import json
+import shutil
+import yaml
+from pathlib import Path
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+from huggingface_hub import hf_hub_download
+from flux_train import start_training
+
+app = FastAPI()
+
 @app.post("/train-from-hf")
 def auto_run_lora_from_repo():
+    try:
+        # ✅ Static or dynamic config
+        REPO_ID = "rahul7star/ohamlab"
+        FOLDER_IN_REPO = "filter-demo/upload_20250708_041329_9c5c81"
+        CONCEPT_SENTENCE = "ohamlab style"
+        LORA_NAME = "ohami_filter_autorun"
+
+        # ✅ Setup HF cache
+        os.environ["HF_HOME"] = "/tmp/hf_cache"
+        os.makedirs("/tmp/hf_cache", exist_ok=True)
+
+        # ✅ Download dataset from HF
+        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",
+            filename=".placeholder",  # needs a dummy filename to trigger folder download
+            subfolder=FOLDER_IN_REPO,
+            local_dir=local_dir,
+            local_dir_use_symlinks=False,
+            force_download=False
+        )
+
+        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:
+            raise HTTPException(status_code=400, detail="No images found in the Hugging Face folder.")
+
+        # ✅ Auto-generate captions
+        captions = [
+            f"Autogenerated caption for {img.stem} in the {CONCEPT_SENTENCE} [trigger]" for img in image_paths
+        ]
+
+        # ✅ Create dataset folder with metadata.jsonl
+        dataset_folder = f"datasets_{uuid.uuid4()}"
+        os.makedirs(dataset_folder, exist_ok=True)
+
+        jsonl_file_path = os.path.join(dataset_folder, "metadata.jsonl")
+        with open(jsonl_file_path, "a") as jsonl_file:
+            for index, image in enumerate(image_paths):
+                new_image_path = shutil.copy(str(image), dataset_folder)
+                file_name = os.path.basename(new_image_path)
+                data = {"file_name": file_name, "prompt": captions[index]}
+                jsonl_file.write(json.dumps(data) + "\n")
+
+        # ✅ Optional advanced config
+        slugged_lora_name = LORA_NAME.replace(" ", "_")
+        os.makedirs("tmp_configs", exist_ok=True)
+        config_path = f"tmp_configs/{uuid.uuid4()}_{slugged_lora_name}.yaml"
+
+        config = {
+            "sample_1": "a stylish anime character with ohamlab style",
+            "sample_2": "a cartoon car in ohamlab style",
+            "sample_3": "portrait in ohamlab lighting"
+        }
+
+        with open(config_path, "w") as f:
+            yaml.dump(config, f)
+
+        # ✅ Final call to train
+        result = start_training(
+            lora_name=LORA_NAME,
+            concept_sentence=CONCEPT_SENTENCE,
+            steps=45,
+            lr=1e-4,
+            rank=32,
+            model_to_train="flux",
+            low_vram=True,
+            dataset_folder=dataset_folder,
+            sample_1=config["sample_1"],
+            sample_2=config["sample_2"],
+            sample_3=config["sample_3"],
+            use_more_advanced_options=True,
+            more_advanced_options=config_path
+        )
+
+        return JSONResponse(content={"status": "success", "message": result})
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+
+
+
+
+def _run_lora_from_repo():
     try:
         # Set HF cache path if not already set
         os.environ["HF_HOME"] = "/tmp/hf_cache"