first

app.py

@@ -1,120 +1,161 @@
-import boto3
 import os
 import re
+import asyncio
+import aiohttp
+import requests
 import json
-from pathlib import Path
-import sqlite3
-from huggingface_hub import Repository, HfFolder
-import tqdm
-import subprocess
+from tqdm import tqdm
 
-from fastapi import FastAPI
+from huggingface_hub import Repository
+
+from fastapi import FastAPI, BackgroundTasks
 from fastapi_utils.tasks import repeat_every
+from fastapi.staticfiles import StaticFiles
+
+from db import Database
 
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
-AWS_ACCESS_KEY_ID = os.getenv('AWS_ACCESS_KEY_ID')
-AWS_SECRET_KEY = os.getenv('AWS_SECRET_KEY')
-AWS_S3_BUCKET_NAME = os.getenv('AWS_S3_BUCKET_NAME')
-
-s3 = boto3.client(service_name='s3',
-                  aws_access_key_id=AWS_ACCESS_KEY_ID,
-                  aws_secret_access_key=AWS_SECRET_KEY)
-
-
-paginator = s3.get_paginator('list_objects_v2')
-
-
-S3_DATA_FOLDER = Path("sd-multiplayer-data")
-ROOMS_DATA_DB = S3_DATA_FOLDER / "rooms_data.db"
-
-
-repo = Repository(
-    local_dir=S3_DATA_FOLDER,
-    repo_type="dataset",
-    clone_from="huggingface-projects/sd-multiplayer-data",
-    use_auth_token=True,
-)
-repo.git_pull()
-
-
-if not ROOMS_DATA_DB.exists():
-    print("Creating database")
-    print("ROOMS_DATA_DB", ROOMS_DATA_DB)
-    db = sqlite3.connect(ROOMS_DATA_DB)
-    with open(Path("schema.sql"), "r") as f:
-        db.executescript(f.read())
-    db.commit()
-    db.close()
-
-
-def get_db(db_path):
-    db = sqlite3.connect(db_path, check_same_thread=False)
-    db.row_factory = sqlite3.Row
-    try:
-        yield db
-    except Exception:
-        db.rollback()
-    finally:
-        db.close()
-
-
-def sync_rooms_to_dataset():
-    for room_data_db in get_db(ROOMS_DATA_DB):
-        rooms = room_data_db.execute("SELECT * FROM rooms").fetchall()
-        cursor = room_data_db.cursor()
-        for row in tqdm.tqdm(rooms):
-            room_id = row["room_id"]
-            print("syncing room data: ", room_id)
-
-            objects = []
-            for result in paginator.paginate(Bucket=AWS_S3_BUCKET_NAME, Prefix=f'{room_id}/', Delimiter='/'):
-                results = []
-                for obj in result.get('Contents'):
-                    try:
-                        key = obj.get('Key')
-                        time = obj.get('LastModified').isoformat()
-                        split_str = re.split(r'[-/.]', key)
-                        uuid = split_str[3]
-                        x, y = [int(n)
-                                for n in re.split(r'[_]', split_str[4])]
-                        prompt = ' '.join(split_str[4:])
-                        results.append(
-                            {'x': x, 'y': y, 'prompt': prompt, 'time': time, 'key': key, 'uuid': uuid})
-                        cursor.execute(
-                            'INSERT INTO rooms_data VALUES (NULL, ?, ?, ?, ?, ?, ?, ?)', (room_id, uuid, x, y, prompt, time, key))
-                    except Exception as e:
-                        print(e)
-                        continue
-
-                objects += results
-            room_data_db.commit()
-
-            all_rows = [dict(row) for row in room_data_db.execute(
-                "SELECT * FROM rooms_data WHERE room_id = ?", (room_id,)).fetchall()]
-            data_path = S3_DATA_FOLDER / f"{room_id}.json"
-            with open(data_path, 'w') as f:
-                json.dump(all_rows, f, separators=(',', ':'))
-    print("Updating repository")
-    subprocess.Popen(
-        "git add . && git commit --amend -m 'update' && git push --force", cwd=S3_DATA_FOLDER, shell=True)
+database = Database()
 
 
-app = FastAPI()
+async def check_image_url(url):
+    async with aiohttp.ClientSession() as session:
+        async with session.head(url) as resp:
+            if resp.status == 200 and resp.content_type.startswith("image/"):
+                return url
+
+
+def fetch_models(page=0):
+    response = requests.get(
+        f'https://huggingface.co/models-json?pipeline_tag=text-to-image&p={page}&sort=likes')
+    data = response.json()
+    return {
+        "models": [model for model in data['models'] if not model['private']],
+        "numItemsPerPage": data['numItemsPerPage'],
+        "numTotalItems": data['numTotalItems'],
+        "pageIndex": data['pageIndex']
+    }
+
+
+def fetch_model_card(model):
+    response = requests.get(
+        f'https://huggingface.co/{model["id"]}/raw/main/README.md')
+    return response.text
+
+
+def find_image_in_model_card(text):
+    image_regex = re.compile(r'https?://\S+(?:png|jpg|jpeg|webp)')
+    urls = re.findall(image_regex, text)
+    # python check if arrya is not empty
+    # if urls:
+    # tasks = []
+    # for url in urls:
+    #     tasks.append(check_image_url(url))
+
+    # results = await asyncio.gather(*tasks)
+    # return [result for result in results if result]
+    return urls
+
+
+def run_inference(endpoint, img):
+    headers = {'Authorization': f'Bearer {HF_TOKEN}',
+               "X-Wait-For-Model": "true",
+               "X-Use-Cache": "true"}
 
+    response = requests.post(endpoint, headers=headers, data=img)
+    return response.json() if response.ok else []
 
-@app.get("/")
-def read_root():
-    return "Just a bot to sync data to huggingface datasets and tweet tha latest data"
+
+def get_all_models():
+    initial = fetch_models()
+    num_pages = initial['numTotalItems'] // initial['numItemsPerPage']
+
+    print(f"Found {num_pages} pages")
+
+    # fetch all models
+    models = []
+    for page in tqdm(range(0, num_pages)):
+        print(f"Fetching page {page} of {num_pages}")
+        page_models = fetch_models(page)
+        models += page_models['models']
+
+    # fetch datacards and images
+    print(f"Found {len(models)} models")
+    final_models = []
+    for model in tqdm(models):
+        print(f"Fetching model {model['id']}")
+        model_card = fetch_model_card(model)
+        images = find_image_in_model_card(model_card)
+        # style = await run_inference(f"https://api-inference.huggingface.co/models/{model['id']}", images[0])
+        style = []
+        # aesthetic = await run_inference(f"https://api-inference.huggingface.co/models/{model['id']}", images[0])
+        aesthetic = []
+        final_models.append(
+            {**model, "images": images, "style": style, "aesthetic": aesthetic}
+        )
+    return final_models
+
+
+async def sync_data():
+    models = get_all_models()
+
+    with open("data/models.json", "w") as f:
+        json.dump(models, f)
+
+    with database.get_db() as db:
+        cursor = db.cursor()
+        for model in models:
+            try:
+                cursor.execute("INSERT INTO models (data) VALUES (?)",
+                               [json.dumps(model)])
+            except Exception as e:
+                print(model['id'], model)
+        db.commit()
+
+
+app = FastAPI()
 
 
-@app.get("/sync")
-def sync():
-    sync_rooms_to_dataset()
+@ app.get("/sync")
+async def sync(background_tasks: BackgroundTasks):
+    background_tasks.add_task(sync_data)
     return "Synced data to huggingface datasets"
 
 
-@app.on_event("startup")
-@repeat_every(seconds=1800)
-def repeat_sync():
-    sync_rooms_to_dataset()
-    return "Synced data to huggingface datasets"
+MAX_PAGE_SIZE = 30
+
+
+@app.get("/api/models")
+def get_page(page: int = 1):
+    page = page if page > 0 else 1
+    with database.get_db() as db:
+        cursor = db.cursor()
+        cursor.execute("""
+            SELECT *
+            FROM (
+                SELECT *, COUNT(*) OVER() AS total
+                FROM models
+                GROUP BY json_extract(data, '$.id')
+                HAVING COUNT(json_extract(data, '$.id')) = 1
+            )
+            ORDER BY json_extract(data, '$.likes') DESC
+            LIMIT ? OFFSET ?
+        """, (MAX_PAGE_SIZE, (page - 1) * MAX_PAGE_SIZE))
+        results = cursor.fetchall()
+        total = results[0][3] if results else 0
+        total_pages = (total + MAX_PAGE_SIZE - 1) // MAX_PAGE_SIZE
+
+    return {
+        "models": [json.loads(result[1]) for result in results],
+        "totalPages": total_pages
+    }
+
+
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
+
+# @app.on_event("startup")
+# @repeat_every(seconds=1800)
+# def repeat_sync():
+#     sync_rooms_to_dataset()
+#     return "Synced data to huggingface datasets"

db.py

@@ -0,0 +1,29 @@
+import sqlite3
+from pathlib import Path
+
+
+class Database:
+    DB_PATH = Path("data/")
+    DB_FILE = DB_PATH / "models.db"
+
+    def __init__(self):
+        if not self.DB_FILE.exists():
+            print("Creating database")
+            print("DB_FILE", self.DB_FILE)
+            db = sqlite3.connect(self.DB_FILE)
+            with open(Path("schema.sql"), "r") as f:
+                db.executescript(f.read())
+            db.commit()
+            db.close()
+
+    def get_db(self):
+        db = sqlite3.connect(self.DB_FILE, check_same_thread=False)
+        db.row_factory = sqlite3.Row
+        return db
+
+    def __enter__(self):
+        self.db = self.get_db()
+        return self.db
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.db.close()

requirements.txt

@@ -1,6 +1,9 @@
-huggingface-hub==0.10
-fastapi-utils==0.2
-uvicorn==0.19
-tqdm==4.64
-boto3==1.26
-fastapi==0.86
+huggingface-hub
+fastapi-utils
+uvicorn
+tqdm
+fastapi
+requests
+asyncio
+aiohttp
+sqlite3

schema.sql

@@ -1,59 +1,11 @@
-PRAGMA foreign_keys=OFF;
+PRAGMA foreign_keys = OFF;
+
 BEGIN TRANSACTION;
-CREATE TABLE rooms (
-    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
-    room_id TEXT NOT NULL
-);INSERT INTO rooms VALUES(1,'room-0');
-INSERT INTO rooms VALUES(2,'room-1');
-INSERT INTO rooms VALUES(3,'room-2');
-INSERT INTO rooms VALUES(4,'room-3');
-INSERT INTO rooms VALUES(5,'room-4');
-INSERT INTO rooms VALUES(6,'room-5');
-INSERT INTO rooms VALUES(7,'room-6');
-INSERT INTO rooms VALUES(8,'room-7');
-INSERT INTO rooms VALUES(9,'room-8');
-INSERT INTO rooms VALUES(10,'room-9');
-INSERT INTO rooms VALUES(11,'room-10');
-INSERT INTO rooms VALUES(12,'room-11');
-INSERT INTO rooms VALUES(13,'room-12');
-INSERT INTO rooms VALUES(14,'room-13');
-INSERT INTO rooms VALUES(15,'room-14');
-INSERT INTO rooms VALUES(16,'room-15');
-INSERT INTO rooms VALUES(17,'room-16');
-INSERT INTO rooms VALUES(18,'room-17');
-INSERT INTO rooms VALUES(19,'room-18');
-INSERT INTO rooms VALUES(20,'room-19');
-INSERT INTO rooms VALUES(21,'room-20');
-INSERT INTO rooms VALUES(22,'room-21');
-INSERT INTO rooms VALUES(23,'room-22');
-INSERT INTO rooms VALUES(24,'room-23');
-INSERT INTO rooms VALUES(25,'room-24');
-INSERT INTO rooms VALUES(26,'room-25');
-INSERT INTO rooms VALUES(27,'room-26');
-INSERT INTO rooms VALUES(28,'room-27');
-INSERT INTO rooms VALUES(29,'room-28');
-INSERT INTO rooms VALUES(30,'room-29');
-INSERT INTO rooms VALUES(31,'room-30');
-INSERT INTO rooms VALUES(32,'room-31');
-INSERT INTO rooms VALUES(33,'room-32');
-INSERT INTO rooms VALUES(34,'room-33');
-INSERT INTO rooms VALUES(35,'room-34');
-INSERT INTO rooms VALUES(36,'room-35');
-INSERT INTO rooms VALUES(37,'room-36');
-INSERT INTO rooms VALUES(38,'room-37');
-INSERT INTO rooms VALUES(39,'room-38');
-INSERT INTO rooms VALUES(40,'room-39');
-INSERT INTO rooms VALUES(41,'room-40');
-DELETE FROM sqlite_sequence;
-INSERT INTO sqlite_sequence VALUES('rooms',41);
-CREATE TABLE rooms_data (
-    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
-    room_id TEXT NOT NULL,
-    uuid TEXT NOT NULL,
-    x INTEGER NOT NULL,
-    y INTEGER NOT NULL,
-    prompt TEXT NOT NULL,
-    time DATETIME NOT NULL,
-    key TEXT NOT NULL,
-    UNIQUE (key) ON CONFLICT IGNORE
-);COMMIT;
+
+CREATE TABLE models (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    data json,
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
+);
+
+COMMIT;