Create app.py

app.py

@@ -0,0 +1,483 @@
+import os
+import uuid
+import yaml
+import json
+import shutil
+import torch
+from pathlib import Path
+from PIL import Image
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from huggingface_hub import hf_hub_download, whoami
+
+
+import os
+os.environ["HF_HOME"] = "/tmp/hf_cache"
+os.makedirs("/tmp/hf_cache", exist_ok=True)
+
+from fastapi import FastAPI, Query
+from huggingface_hub import list_repo_files, hf_hub_download, upload_file
+import io
+import requests
+from fastapi import BackgroundTasks
+from fastapi import FastAPI, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+
+
+import os
+import os
+import zipfile
+import tempfile  # ✅ Add this!
+
+
+
+
+app = FastAPI()
+
+# CORS setup to allow requests from your frontend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Replace "*" with your frontend domain in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/")
+def health_check():
+    return {"status": "✅ FastAPI running on Hugging Face Spaces!"}
+
+@app.get("/healthz")
+def healthz():
+    return {"ok": True}
+
+@app.get("/docs", include_in_schema=False)
+def custom_docs():
+    return JSONResponse(get_openapi(title="LoRA Autorun API", version="1.0.0", routes=app.routes))
+
+
+
+REPO_ID = "rahul7star/ohamlab"
+FOLDER = "demo"
+BASE_URL = f"https://huggingface.co/{REPO_ID}/resolve/main/"
+
+#show all images in a DIR at UI FE
+@app.get("/images")
+def list_images():
+    try:
+        all_files = list_repo_files(REPO_ID)
+
+        folder_prefix = FOLDER.rstrip("/") + "/"
+
+        files_in_folder = [
+            f for f in all_files
+            if f.startswith(folder_prefix)
+            and "/" not in f[len(folder_prefix):]  # no subfolder files
+            and f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
+        ]
+
+        urls = [BASE_URL + f for f in files_in_folder]
+
+        return {"images": urls}
+
+    except Exception as e:
+        return {"error": str(e)}
+
+from datetime import datetime
+import tempfile
+import uuid
+
+# upload zip from UI 
+@app.post("/upload-zip")
+async def upload_zip(file: UploadFile = File(...)):
+    if not file.filename.endswith(".zip"):
+        return {"error": "Please upload a .zip file"}
+
+    # Save the ZIP to /tmp
+    temp_zip_path = f"/tmp/{file.filename}"
+    with open(temp_zip_path, "wb") as f:
+        f.write(await file.read())
+
+    # Create a unique subfolder name inside 'demo/'
+    timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
+    unique_id = uuid.uuid4().hex[:6]
+    folder_name = f"upload_{timestamp}_{unique_id}"
+    hf_folder_prefix = f"demo/{folder_name}"
+
+    try:
+        with tempfile.TemporaryDirectory() as extract_dir:
+            # Extract zip
+            with zipfile.ZipFile(temp_zip_path, 'r') as zip_ref:
+                zip_ref.extractall(extract_dir)
+
+            uploaded_files = []
+
+            # Upload all extracted files
+            for root_dir, _, files in os.walk(extract_dir):
+                for name in files:
+                    file_path = os.path.join(root_dir, name)
+                    relative_path = os.path.relpath(file_path, extract_dir)
+                    repo_path = f"{hf_folder_prefix}/{relative_path}".replace("\\", "/")
+
+                    upload_file(
+                        path_or_fileobj=file_path,
+                        path_in_repo=repo_path,
+                        repo_id="rahul7star/ohamlab",
+                        repo_type="model",
+                        commit_message=f"Upload {relative_path} to {folder_name}",
+                        token=True,
+                    )
+                    uploaded_files.append(repo_path)
+
+        return {
+            "message": f"✅ Uploaded {len(uploaded_files)} files",
+            "folder": folder_name,
+            "files": uploaded_files,
+        }
+
+    except Exception as e:
+        return {"error": f"❌ Failed to process zip: {str(e)}"}
+
+    
+# upload a single file from UI
+from typing import List
+from fastapi import UploadFile, File, APIRouter
+import os
+from fastapi import UploadFile, File, APIRouter
+from typing import List
+from datetime import datetime
+import uuid, os
+
+
+@app.post("/upload")
+async def upload_images(
+    background_tasks: BackgroundTasks,
+    files: List[UploadFile] = File(...)
+):
+    # Step 1: Generate dynamic folder name
+    timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
+    unique_id = uuid.uuid4().hex[:6]
+    folder_name = f"upload_{timestamp}_{unique_id}"
+    hf_folder_prefix = f"demo/{folder_name}"
+
+    responses = []
+
+    # Step 2: Save and upload each image
+    for file in files:
+        filename = file.filename
+        contents = await file.read()
+        temp_path = f"/tmp/{filename}"
+        with open(temp_path, "wb") as f:
+            f.write(contents)
+
+        try:
+            upload_file(
+                path_or_fileobj=temp_path,
+                path_in_repo=f"{hf_folder_prefix}/{filename}",
+                repo_id=T_REPO_ID,
+                repo_type="model",
+                commit_message=f"Upload {filename} to {hf_folder_prefix}",
+                token=True,
+            )
+            responses.append({
+                "filename": filename,
+                "status": "✅ uploaded",
+                "path": f"{hf_folder_prefix}/{filename}"
+            })
+        except Exception as e:
+            responses.append({
+                "filename": filename,
+                "status": f"❌ failed: {str(e)}"
+            })
+
+        os.remove(temp_path)
+
+    # Step 3: Add filter job to background
+    def run_filter():
+        try:
+            result = filter_and_rename_images(folder=hf_folder_prefix)
+            print(f"🧼 Filter result: {result}")
+        except Exception as e:
+            print(f"❌ Filter failed: {str(e)}")
+
+    background_tasks.add_task(run_filter)
+
+    return {
+        "message": f"{len(files)} file(s) uploaded",
+        "upload_folder": hf_folder_prefix,
+        "results": responses,
+        "note": "Filtering started in background"
+    }
+
+
+
+
+
+
+#Tranining Data set start fitering data for traninig
+
+
+T_REPO_ID = "rahul7star/ohamlab"
+DESCRIPTION_TEXT = (
+    "Ra3hul is wearing a black jacket over a striped white t-shirt with blue jeans. "
+    "He is standing near a lake with his arms spread wide open, with mountains and cloudy skies in the background."
+)
+
+def is_image_file(filename: str) -> bool:
+    return filename.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
+
+@app.post("/filter-images")
+def filter_and_rename_images(folder: str = Query("demo", description="Folder path in repo to scan")):
+    try:
+        all_files = list_repo_files(T_REPO_ID)
+        folder_prefix = folder.rstrip("/") + "/"
+        filter_folder = f"filter-{folder.rstrip('/')}"
+        filter_prefix = filter_folder + "/"
+
+        # Filter images only directly in the folder (no subfolders)
+        image_files = [
+            f for f in all_files
+            if f.startswith(folder_prefix)
+            and "/" not in f[len(folder_prefix):]  # no deeper path
+            and is_image_file(f)
+        ]
+
+        if not image_files:
+            return {"error": f"No images found in folder '{folder}'"}
+
+        uploaded_files = []
+
+        for idx, orig_path in enumerate(image_files, start=1):
+            # Download image content bytes (uses local cache)
+            local_path = hf_hub_download(repo_id=T_REPO_ID, filename=orig_path)
+            with open(local_path, "rb") as f:
+                file_bytes = f.read()
+
+            # Rename images as image1.jpeg, image2.jpeg, ...
+            new_image_name = f"image{idx}.jpeg"
+
+            # Upload renamed image from memory
+            upload_file(
+                path_or_fileobj=io.BytesIO(file_bytes),
+                path_in_repo=filter_prefix + new_image_name,
+                repo_id=T_REPO_ID,
+                repo_type="model",
+                commit_message=f"Upload renamed image {new_image_name} to {filter_folder}",
+                token=True,
+            )
+            uploaded_files.append(filter_prefix + new_image_name)
+
+            # Create and upload text file for each image
+            txt_filename = f"image{idx}.txt"
+            upload_file(
+                path_or_fileobj=io.BytesIO(DESCRIPTION_TEXT.encode("utf-8")),
+                path_in_repo=filter_prefix + txt_filename,
+                repo_id=T_REPO_ID,
+                repo_type="model",
+                commit_message=f"Upload text file {txt_filename} to {filter_folder}",
+                token=True,
+            )
+            uploaded_files.append(filter_prefix + txt_filename)
+
+        return {
+            "message": f"Processed and uploaded {len(image_files)} images and text files.",
+            "files": uploaded_files,
+        }
+
+    except Exception as e:
+        return {"error": str(e)}
+
+        
+
+# ========== CONFIGURATION ==========
+REPO_ID = "rahul7star/ohamlab"
+FOLDER_IN_REPO = "filter-demo/upload_20250708_041329_9c5c81"
+CONCEPT_SENTENCE = "ohamlab style"
+LORA_NAME = "ohami_filter_autorun"
+
+# ========== FASTAPI APP ==========
+app = FastAPI()
+
+# ========== HELPERS ==========
+def create_dataset(images, *captions):
+    destination_folder = f"datasets_{uuid.uuid4()}"
+    os.makedirs(destination_folder, exist_ok=True)
+
+    jsonl_file_path = os.path.join(destination_folder, "metadata.jsonl")
+    with open(jsonl_file_path, "a") as jsonl_file:
+        for index, image in enumerate(images):
+            new_image_path = shutil.copy(str(image), destination_folder)
+            caption = captions[index]
+            file_name = os.path.basename(new_image_path)
+            data = {"file_name": file_name, "prompt": caption}
+            jsonl_file.write(json.dumps(data) + "\n")
+
+    return destination_folder
+
+def recursive_update(d, u):
+    for k, v in u.items():
+        if isinstance(v, dict) and v:
+            d[k] = recursive_update(d.get(k, {}), v)
+        else:
+            d[k] = v
+    return d
+
+def start_training(
+    lora_name,
+    concept_sentence,
+    steps,
+    lr,
+    rank,
+    model_to_train,
+    low_vram,
+    dataset_folder,
+    sample_1,
+    sample_2,
+    sample_3,
+    use_more_advanced_options,
+    more_advanced_options,
+):
+    try:
+        user = whoami()
+        username = user.get("name", "anonymous")
+        push_to_hub = True
+    except:
+        username = "anonymous"
+        push_to_hub = False
+
+    slugged_lora_name = lora_name.replace(" ", "_").lower()
+
+    # Load base config
+    config = {
+        "config": {
+            "name": slugged_lora_name,
+            "process": [
+                {
+                    "model": {
+                        "low_vram": low_vram,
+                        "is_flux": True,
+                        "quantize": True,
+                        "name_or_path": "black-forest-labs/FLUX.1-dev"
+                    },
+                    "network": {
+                        "linear": rank,
+                        "linear_alpha": rank,
+                        "type": "lora"
+                    },
+                    "train": {
+                        "steps": steps,
+                        "lr": lr,
+                        "skip_first_sample": True,
+                        "batch_size": 1,
+                        "dtype": "bf16",
+                        "gradient_accumulation_steps": 1,
+                        "gradient_checkpointing": True,
+                        "noise_scheduler": "flowmatch",
+                        "optimizer": "adamw8bit",
+                        "ema_config": {
+                            "use_ema": True,
+                            "ema_decay": 0.99
+                        }
+                    },
+                    "datasets": [
+                        {"folder_path": dataset_folder}
+                    ],
+                    "save": {
+                        "dtype": "float16",
+                        "save_every": 10000,
+                        "push_to_hub": push_to_hub,
+                        "hf_repo_id": f"{username}/{slugged_lora_name}",
+                        "hf_private": True,
+                        "max_step_saves_to_keep": 4
+                    },
+                    "sample": {
+                        "guidance_scale": 3.5,
+                        "sample_every": steps,
+                        "sample_steps": 28,
+                        "width": 1024,
+                        "height": 1024,
+                        "walk_seed": True,
+                        "seed": 42,
+                        "sampler": "flowmatch",
+                        "prompts": [p for p in [sample_1, sample_2, sample_3] if p]
+                    },
+                    "trigger_word": concept_sentence
+                }
+            ]
+        }
+    }
+
+    # Apply advanced YAML overrides if any
+    if use_more_advanced_options and more_advanced_options:
+        advanced_config = yaml.safe_load(more_advanced_options)
+        config["config"]["process"][0] = recursive_update(config["config"]["process"][0], advanced_config)
+
+    # Save YAML config
+    os.makedirs("tmp_configs", exist_ok=True)
+    config_path = f"tmp_configs/{uuid.uuid4()}_{slugged_lora_name}.yaml"
+    with open(config_path, "w") as f:
+        yaml.dump(config, f)
+
+    # Simulate training
+    print(f"[INFO] Starting training with config: {config_path}")
+    print(json.dumps(config, indent=2))
+    return f"Training started successfully with config: {config_path}"
+
+# ========== MAIN ENDPOINT ==========
+@app.post("/train-from-hf")
+def auto_run_lora_from_repo():
+    try:
+        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,
+            allow_patterns=["*.jpg", "*.png", "*.jpeg"],
+        )
+
+        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:
+            return JSONResponse(status_code=400, content={"error": "No images found in the HF repo folder."})
+
+        captions = [
+            f"Autogenerated caption for {img.stem} in the {CONCEPT_SENTENCE} [trigger]" for img in image_paths
+        ]
+
+        dataset_path = create_dataset(image_paths, *captions)
+
+        result = start_training(
+            lora_name=LORA_NAME,
+            concept_sentence=CONCEPT_SENTENCE,
+            steps=1000,
+            lr=4e-4,
+            rank=16,
+            model_to_train="dev",
+            low_vram=True,
+            dataset_folder=dataset_path,
+            sample_1=f"A stylized portrait using {CONCEPT_SENTENCE}",
+            sample_2=f"A cat in the {CONCEPT_SENTENCE}",
+            sample_3=f"A selfie processed in {CONCEPT_SENTENCE}",
+            use_more_advanced_options=True,
+            more_advanced_options="""
+training:
+  seed: 42
+  precision: bf16
+  batch_size: 2
+augmentation:
+  flip: true
+  color_jitter: true
+"""
+        )
+
+        return {"message": result}
+
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})