resize images and use s3

app.py

@@ -1,46 +1,72 @@
 import os
 import re
-import asyncio
 import aiohttp
 import requests
 import json
+import subprocess
+import asyncio
+from io import BytesIO
+import uuid
+
+from math import ceil
 from tqdm import tqdm
 from pathlib import Path
 from huggingface_hub import Repository
-
+from PIL import Image, ImageOps
 from fastapi import FastAPI, BackgroundTasks
 from fastapi_utils.tasks import repeat_every
 from fastapi.middleware.cors import CORSMiddleware
+import boto3
 
 from db import Database
 
+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')
+
+
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
+S3_DATA_FOLDER = Path("sd-multiplayer-data")
 
 DB_FOLDER = Path("diffusers-gallery-data")
 
+s3 = boto3.client(service_name='s3',
+                  aws_access_key_id=AWS_ACCESS_KEY_ID,
+                  aws_secret_access_key=AWS_SECRET_KEY)
 
-repo = Repository(
-    local_dir=DB_FOLDER,
-    repo_type="dataset",
-    clone_from="huggingface-projects/diffusers-gallery-data",
-    use_auth_token=True,
-)
-repo.git_pull()
+
+# repo = Repository(
+#     local_dir=DB_FOLDER,
+#     repo_type="dataset",
+#     clone_from="huggingface-projects/diffusers-gallery-data",
+#     use_auth_token=True,
+# )
+# repo.git_pull()
 
 database = Database(DB_FOLDER)
 
 
-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
+async def upload_resize_image_url(session, image_url):
+    print(f"Uploading image {image_url}")
+    async with session.get(image_url) as response:
+        if response.status == 200 and response.headers['content-type'].startswith('image'):
+            image = Image.open(BytesIO(await response.read())).convert('RGB')
+            # resize image proportional
+            image = ImageOps.fit(image, (400, 400), Image.LANCZOS)
+            image_bytes = BytesIO()
+            image.save(image_bytes, format="JPEG")
+            image_bytes.seek(0)
+            fname = f'{uuid.uuid4()}.jpg'
+            s3.upload_fileobj(Fileobj=image_bytes, Bucket=AWS_S3_BUCKET_NAME, Key="diffusers-gallery/" + fname,
+                              ExtraArgs={"ContentType": "image/jpeg", "CacheControl": "max-age=31536000"})
+            return fname
+    return None
 
 
 def fetch_models(page=0):
     response = requests.get(
-        f'https://huggingface.co/models-json?pipeline_tag=text-to-image&p={page}&sort=likes')
+        f'https://huggingface.co/models-json?pipeline_tag=text-to-image&p={page}')
     data = response.json()
     return {
         "models": [model for model in data['models'] if not model['private']],
@@ -56,18 +82,16 @@ def fetch_model_card(model):
     return response.text
 
 
-def find_image_in_model_card(text):
+async 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))
+    if not urls:
+        return []
 
-    # results = await asyncio.gather(*tasks)
-    # return [result for result in results if result]
-    return urls
+    async with aiohttp.ClientSession() as session:
+        tasks = [asyncio.ensure_future(upload_resize_image_url(
+            session, image_url)) for image_url in urls[0:3]]
+        return await asyncio.gather(*tasks)
 
 
 def run_inference(endpoint, img):
@@ -79,10 +103,12 @@ def run_inference(endpoint, img):
     return response.json() if response.ok else []
 
 
-def get_all_models():
-    initial = fetch_models()
-    num_pages = initial['numTotalItems'] // initial['numItemsPerPage']
+async def get_all_models():
+    initial = fetch_models(0)
+    num_pages = ceil(initial['numTotalItems'] / initial['numItemsPerPage'])
 
+    print(
+        f"Total items: {initial['numTotalItems']} - Items per page: {initial['numItemsPerPage']}")
     print(f"Found {num_pages} pages")
 
     # fetch all models
@@ -92,13 +118,16 @@ def get_all_models():
         page_models = fetch_models(page)
         models += page_models['models']
 
+    with open(DB_FOLDER / "models_temp.json", "w") as f:
+        json.dump(models, f)
+
     # 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)
+        images = await 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])
@@ -110,20 +139,27 @@ def get_all_models():
 
 
 async def sync_data():
-    models = get_all_models()
+    print("Fetching models")
+    models = await get_all_models()
 
-    with open("data/models.json", "w") as f:
+    with open(DB_FOLDER / "models.json", "w") as f:
         json.dump(models, f)
-
+    # with open(DB_FOLDER / "models.json", "r") as f:
+    #     models = json.load(f)
+    # open temp db
+    print("Updating database")
     with database.get_db() as db:
         cursor = db.cursor()
         for model in models:
             try:
-                cursor.execute("INSERT INTO models (data) VALUES (?)",
-                               [json.dumps(model)])
+                cursor.execute("INSERT INTO models(id, data) VALUES (?, ?)",
+                               [model['id'], json.dumps(model)])
             except Exception as e:
                 print(model['id'], model)
         db.commit()
+    print("Updating repository")
+    # subprocess.Popen(
+    #     "git add . && git commit --amend -m 'update' && git push --force", cwd=DB_FOLDER, shell=True)
 
 
 app = FastAPI()
@@ -145,19 +181,14 @@ async def sync(background_tasks: BackgroundTasks):
 MAX_PAGE_SIZE = 30
 
 
-@app.get("/api/models")
+@ 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
-            )
+            SELECT *, COUNT(*) OVER() AS total
+            FROM models
             ORDER BY json_extract(data, '$.likes') DESC
             LIMIT ? OFFSET ?
         """, (MAX_PAGE_SIZE, (page - 1) * MAX_PAGE_SIZE))
@@ -175,8 +206,9 @@ def get_page(page: int = 1):
 def read_root():
     return "Just a bot to sync data from diffusers gallery"
 
+
 # @app.on_event("startup")
-# @repeat_every(seconds=1800)
-# def repeat_sync():
-#     sync_rooms_to_dataset()
+# @repeat_every(seconds=60 * 60 * 24, wait_first=False)
+# async def repeat_sync():
+#     await sync_data()
 #     return "Synced data to huggingface datasets"

requirements.txt

@@ -5,4 +5,6 @@ tqdm
 fastapi
 requests
 asyncio
-aiohttp
+aiohttp
+Pillow
+boto3

schema.sql

@@ -3,7 +3,7 @@ PRAGMA foreign_keys = OFF;
 BEGIN TRANSACTION;
 
 CREATE TABLE models (
-    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    id TEXT PRIMARY KEY NOT NULL,
     data json,
     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
 );