import os import io import json import time import shutil import random from datetime import datetime from flask import Flask, request, jsonify, render_template from huggingface_hub import HfApi, hf_hub_download from PIL import Image import numpy as np import faiss import torch from transformers import CLIPProcessor, CLIPModel # ------------------- # Config # ------------------- HF_TOKEN = os.getenv("HF_TOKEN") # put your real token in HF Space Secrets DATASET_REPO = "Sahil5112/Fan_hub" # your dataset repo IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".gif", ".webp") VIDEO_EXTS = (".mp4", ".webm", ".mov", ".avi") FAISS_FILENAME = "index.faiss" MAP_FILENAME = "index_map.json" META_FILENAME = "metadata.json" STATE_DIR = "state" UPLOADS_DIR = "uploads" os.makedirs(STATE_DIR, exist_ok=True) os.makedirs(UPLOADS_DIR, exist_ok=True) app = Flask(__name__, template_folder="templates") api = HfApi() # ------------------- # CLIP Model # ------------------- DEVICE = "cuda" if torch.cuda.is_available() else "cpu" MODEL_ID = "openai/clip-vit-base-patch32" clip_model = CLIPModel.from_pretrained(MODEL_ID).to(DEVICE) clip_processor = CLIPProcessor.from_pretrained(MODEL_ID) EMB_DIM = clip_model.visual_projection.out_features # ------------------- # Helpers # ------------------- def l2_normalize(x: np.ndarray) -> np.ndarray: x = x.astype("float32") norms = np.linalg.norm(x, axis=1, keepdims=True) + 1e-10 return x / norms def image_to_embedding(pil_img: Image.Image) -> np.ndarray: if pil_img.mode != "RGB": pil_img = pil_img.convert("RGB") pil_img.thumbnail((512, 512), Image.Resampling.LANCZOS) inputs = clip_processor(images=pil_img, return_tensors="pt").to(DEVICE) with torch.no_grad(): feats = clip_model.get_image_features(**inputs) feats = feats.cpu().numpy().astype("float32") return l2_normalize(feats) def text_to_embedding(text: str) -> np.ndarray: inputs = clip_processor(text=[text], return_tensors="pt", padding=True).to(DEVICE) with torch.no_grad(): feats = clip_model.get_text_features(**inputs) feats = feats.cpu().numpy().astype("float32") return l2_normalize(feats) def dataset_url(path_in_repo: str) -> str: return f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{path_in_repo}" def load_json_local(path: str, default): if os.path.exists(path): with open(path, "r", encoding="utf-8") as f: return json.load(f) return default def save_json_local(path: str, obj): with open(path, "w", encoding="utf-8") as f: json.dump(obj, f, ensure_ascii=False, indent=2) # ------------------- # Persistent Index # ------------------- index: faiss.Index = None index_map: list[dict] = [] metadata: dict = {} def _create_empty_index(): return faiss.IndexFlatIP(EMB_DIM) def pull_or_init_state(): global index, index_map, metadata for repo_path in [FAISS_FILENAME, MAP_FILENAME, META_FILENAME]: try: hf_hub_download( repo_id=DATASET_REPO, filename=repo_path, repo_type="dataset", token=HF_TOKEN, local_dir=STATE_DIR, local_dir_use_symlinks=False ) except Exception: pass faiss_path = os.path.join(STATE_DIR, FAISS_FILENAME) map_path = os.path.join(STATE_DIR, MAP_FILENAME) meta_path = os.path.join(STATE_DIR, META_FILENAME) index = faiss.read_index(faiss_path) if os.path.exists(faiss_path) else _create_empty_index() index_map = load_json_local(map_path, []) metadata = load_json_local(meta_path, {}) for fn, md in metadata.items(): if "uploader" not in md: md["uploader"] = "Anonymous" save_json_local(meta_path, metadata) def persist_state(): faiss_path = os.path.join(STATE_DIR, FAISS_FILENAME) map_path = os.path.join(STATE_DIR, MAP_FILENAME) meta_path = os.path.join(STATE_DIR, META_FILENAME) faiss.write_index(index, faiss_path) save_json_local(map_path, index_map) save_json_local(meta_path, metadata) for local_file, repo_file in [(faiss_path, FAISS_FILENAME), (map_path, MAP_FILENAME), (meta_path, META_FILENAME)]: api.upload_file( path_or_fileobj=local_file, path_in_repo=repo_file, repo_id=DATASET_REPO, repo_type="dataset", token=HF_TOKEN ) def add_item_to_index(filename: str, pil_img: Image.Image): vec = image_to_embedding(pil_img) index.add(vec) md = metadata.get(filename, {}) index_map.append({ "filename": filename, "title": md.get("title", "Untitled"), "description": md.get("description", ""), "category": md.get("category", "Uncategorized"), "filetype": "image" }) def top_k_from_index(query_vec: np.ndarray, k: int = 24, exclude_filename: str | None = None): if index.ntotal == 0: return [] D, I = index.search(query_vec.astype("float32"), k + 1) results = [] for idx in I[0]: if idx < 0 or idx >= len(index_map): continue item = index_map[idx] if exclude_filename and item["filename"] == exclude_filename: continue results.append({ "url": metadata.get(item["filename"], {}).get("url", dataset_url(item["filename"])), "filename": item["filename"], "title": item.get("title", "Untitled"), "description": item.get("description", ""), "category": item.get("category", "Uncategorized"), "uploader": metadata.get(item["filename"], {}).get("uploader", "Anonymous"), "filetype": item.get("filetype", "image") }) if len(results) >= k: break return results # ------------------- # Routes # ------------------- @app.route("/") def home(): return render_template("index.html") @app.route("/upload", methods=["POST"]) def upload(): if "file" not in request.files: return jsonify({"error": "No file provided"}), 400 f = request.files["file"] ext = os.path.splitext(f.filename.lower())[1] if not (ext in IMAGE_EXTS or ext in VIDEO_EXTS): return jsonify({"error": "Unsupported file type"}), 400 title = request.form.get("title", "Untitled").strip()[:120] description = request.form.get("description", "").strip()[:500] category = request.form.get("category", "Uncategorized").strip()[:60] uploader = request.form.get("uploader", "Anonymous").strip()[:80] ts = datetime.now().strftime("%Y%m%d_%H%M%S") safe_name = f.filename.replace("/", "_").replace("\\", "_") filename = f"{ts}_{safe_name}" local_path = os.path.join(UPLOADS_DIR, filename) if ext in IMAGE_EXTS and ext != ".gif": try: img = Image.open(f).convert("RGB") img.save(local_path, format="JPEG", optimize=True, quality=85) except Exception as e: print("Compression failed:", e) f.seek(0) f.save(local_path) else: f.save(local_path) api.upload_file( path_or_fileobj=local_path, path_in_repo=filename, repo_id=DATASET_REPO, repo_type="dataset", token=HF_TOKEN ) filetype = "video" if ext in VIDEO_EXTS else "image" metadata[filename] = { "title": title or "Untitled", "description": description, "category": category or "Uncategorized", "uploader": uploader or "Anonymous", "filetype": filetype, "url": dataset_url(filename) } if filetype == "image" and ext != ".gif": img = Image.open(local_path).convert("RGB") add_item_to_index(filename=filename, pil_img=img) persist_state() try: os.remove(local_path) except: pass return jsonify({ "success": True, "filename": filename, "url": metadata[filename]["url"], "filetype": filetype, "uploader": metadata[filename]["uploader"] }) # ------------------- # Gallery with Pagination (Lazy Load) # ------------------- @app.route("/gallery") def gallery(): try: offset = int(request.args.get("offset", 0)) limit = int(request.args.get("limit", 10)) except: offset = 0 limit = 10 # --- Get filters from query params --- text_filter = request.args.get("text", "").lower() uploader_filter = request.args.get("uploader", "").lower() category_filter = request.args.get("category", "").lower() filetype_filter = request.args.get("filetype", "").lower() tags_filter = [t.strip().lower() for t in request.args.get("tags", "").split(",") if t.strip()] items = [] # Images from index_map for entry in index_map: fn = entry["filename"] md = metadata.get(fn, {}) item = { "url": md.get("url", dataset_url(fn)), "filename": fn, "title": md.get("title", entry.get("title", "Untitled")), "description": md.get("description", entry.get("description", "")), "category": md.get("category", entry.get("category", "Uncategorized")), "uploader": md.get("uploader", "Anonymous"), "filetype": md.get("filetype", entry.get("filetype", "image")), "tags": md.get("tags", []) } # --- Apply filters --- if text_filter and text_filter not in item["title"].lower() and text_filter not in item["description"].lower(): continue if uploader_filter and uploader_filter != item["uploader"].lower(): continue if category_filter and category_filter != item["category"].lower(): continue if filetype_filter and filetype_filter != item["filetype"].lower(): continue if tags_filter and not any(t.lower() in [tag.lower() for tag in item.get("tags", [])] for t in tags_filter): continue items.append(item) # Add videos not in index_map for fn, md in metadata.items(): if md.get("filetype") == "video" and not any(it["filename"] == fn for it in items): item = { "url": md.get("url", dataset_url(fn)), "filename": fn, "title": md.get("title", "Untitled"), "description": md.get("description", ""), "category": md.get("category", "Uncategorized"), "uploader": md.get("uploader", "Anonymous"), "filetype": "video", "tags": md.get("tags", []) } # --- Apply filters --- if text_filter and text_filter not in item["title"].lower() and text_filter not in item["description"].lower(): continue if uploader_filter and uploader_filter != item["uploader"].lower(): continue if category_filter and category_filter != item["category"].lower(): continue if filetype_filter and filetype_filter != item["filetype"].lower(): continue if tags_filter and not any(t.lower() in [tag.lower() for tag in item.get("tags", [])] for t in tags_filter): continue items.append(item) items.sort(key=lambda x: x["filename"], reverse=True) paged_items = items[offset: offset + limit] return jsonify(paged_items) # ------------------- # Related & Search # ------------------- @app.route("/related", methods=["POST"]) def related(): data = request.get_json(force=True) filename = data.get("filename", "") if not filename: return jsonify([]) md = metadata.get(filename, {}) filetype = md.get("filetype", "image") if filetype == "image": try: local = hf_hub_download(repo_id=DATASET_REPO, filename=filename, repo_type="dataset", token=HF_TOKEN) img = Image.open(local).convert("RGB") qvec = image_to_embedding(img) return jsonify(top_k_from_index(qvec, k=12, exclude_filename=filename)) except Exception: return jsonify([]) else: same_cat = [fn for fn, m in metadata.items() if fn != filename and m.get("category") == md.get("category")] random.shuffle(same_cat) res = [] for fn in same_cat[:12]: m = metadata[fn] res.append({ "url": m.get("url", dataset_url(fn)), "filename": fn, "title": m.get("title", "Untitled"), "description": m.get("description", ""), "category": m.get("category", "Uncategorized"), "uploader": m.get("uploader", "Anonymous"), "filetype": m.get("filetype", "image") }) return jsonify(res) @app.route("/search/text") def search_text(): q = request.args.get("q", "").strip() if not q: return jsonify([]) qvec = text_to_embedding(q) return jsonify(top_k_from_index(qvec, k=30)) @app.route("/search/image", methods=["POST"]) def search_image(): if "file" not in request.files: return jsonify({"error": "No file"}), 400 f = request.files["file"] if not f.filename.lower().endswith(IMAGE_EXTS): return jsonify({"error": "Unsupported file type"}), 400 img = Image.open(io.BytesIO(f.read())).convert("RGB") qvec = image_to_embedding(img) return jsonify(top_k_from_index(qvec, k=30)) # ------------------- # Init # ------------------- pull_or_init_state() if __name__ == "__main__": app.run(host="0.0.0.0", port=7860, debug=False)