huggingface-projects/diffusers-gallery-bot

1import os2import re3import requests4import json5import subprocess6from io import BytesIO7import uuid8 9from math import ceil10from tqdm import tqdm11from pathlib import Path12 13from db import Database14 15DB_FOLDER = Path("diffusers-gallery-data")16 17database = Database(DB_FOLDER)18 19 20CLASSIFIER_URL = "https://radames-aesthetic-style-nsfw-classifier.hf.space/run/inference"21ASSETS_URL = "https://d26smi9133w0oo.cloudfront.net/diffusers-gallery/"22 23 24def main():25 26    with database.get_db() as db:27        cursor = db.cursor()28        cursor.execute("""29            SELECT *30            FROM models31        """)32        results = list(cursor.fetchall())33 34    for row in tqdm(results):35        row_id = row['id']36        # keep json data on row_data37        row_data = json.loads(row['data'])38        print("updating row", row_id)39        images = row_data['images']40 41        # filter nones42        images = [i for i in images if i is not None]43        if len(images) > 0:44            # classifying only the first image45            images_urls = [ASSETS_URL + images[0]]46            response = requests.post(CLASSIFIER_URL, json={"data": [47                {"urls": images_urls},  # json urls: list of images urls48                False,  # enable/disable gallery image output49                None,  # single image input50                None,  # files input51            ]}).json()52 53            # data response is array data:[[{img0}, {img1}, {img2}...], Label, Gallery],54            class_data = response['data'][0][0]55            class_data_parsed = {row['label']: round(56                row['score'], 3) for row in class_data}57 58            # update row data with classificator data59            row_data['class'] = class_data_parsed60        else:61            row_data['class'] = {}62        with database.get_db() as db:63            cursor = db.cursor()64            cursor.execute("UPDATE models SET data = ? WHERE id = ?",65                           [json.dumps(row_data), row_id])66            db.commit()67 68 69if __name__ == "__main__":70    main()71