tags data

app.py

@@ -8,6 +8,7 @@ import subprocess
 import asyncio
 from io import BytesIO
 import uuid
+import yaml
 
 from math import ceil
 from tqdm import tqdm
@@ -53,6 +54,9 @@ repo.git_pull()
 
 database = Database(DB_FOLDER)
 
+REGEX_YAML_BLOCK = re.compile(
+    r"^(\s*---[\r\n]+)([\S\s]*?)([\r\n]+---(\r\n|\n|$))")
+
 
 async def upload_resize_image_url(session, image_url):
     print(f"Uploading image {image_url}")
@@ -92,6 +96,16 @@ def fetch_model_card(model_id):
     return response.text
 
 
+def get_yaml_data(text_content):
+    match = REGEX_YAML_BLOCK.search(text_content)
+    if match:
+        yaml_block = match.group(2)
+        data_dict = yaml.safe_load(yaml_block)
+    else:
+        data_dict = {}
+    return data_dict
+
+
 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)
@@ -207,11 +221,13 @@ async def sync_data():
         for model in tqdm(models_no_images):
             model_id = model['id']
             model_data = json.loads(model['data'])
-            print("Updating model", model_id)
             model_card = fetch_model_card(model_id)
+            model_card_data = get_yaml_data(model_card)
             images = await find_image_in_model_card(model_card)
             classifier = run_classifier(images)
-
+            model_data['images'] = images
+            model_data['class'] = classifier
+            model_data['meta'] = model_card_data
             # update model row with image and classifier data
             with database.get_db() as db:
                 cursor = db.cursor()
@@ -234,7 +250,7 @@ async def sync_data():
     time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     cmd = f"git add . && git commit --amend -m 'update at {time}' && git push --force"
     print(cmd)
-    subprocess.Popen(cmd, cwd=DB_FOLDER, shell=True)
+    # subprocess.Popen(cmd, cwd=DB_FOLDER, shell=True)
 
 
 app = FastAPI()
@@ -327,7 +343,7 @@ def read_root():
     # return html page from string
     return HTMLResponse("""
     <p>Just a bot to sync data from diffusers gallery please go to
-    <a href="https://huggingface.co/spaces/huggingface-projects/diffusers-gallery">https://huggingface.co/spaces/huggingface-projects/diffusers-gallery</a>
+    <a href="https://huggingface.co/spaces/huggingface-projects/diffusers-gallery" target="_blank" rel="noopener noreferrer">https://huggingface.co/spaces/huggingface-projects/diffusers-gallery</a>
     </p>""")
 
 

requirements.txt

@@ -7,4 +7,5 @@ requests
 asyncio
 aiohttp
 Pillow
-boto3
+boto3
+yaml