fix regex stuck

app.py

@@ -24,9 +24,9 @@ import boto3
 from datetime import datetime
 from db import Database
 
-AWS_ACCESS_KEY_ID = os.getenv('MY_AWS_ACCESS_KEY_ID')
-AWS_SECRET_KEY = os.getenv('MY_AWS_SECRET_KEY')
-AWS_S3_BUCKET_NAME = os.getenv('MY_AWS_S3_BUCKET_NAME')
+AWS_ACCESS_KEY_ID = os.getenv("MY_AWS_ACCESS_KEY_ID")
+AWS_SECRET_KEY = os.getenv("MY_AWS_SECRET_KEY")
+AWS_S3_BUCKET_NAME = os.getenv("MY_AWS_S3_BUCKET_NAME")
 
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -35,13 +35,17 @@ S3_DATA_FOLDER = Path("sd-multiplayer-data")
 
 DB_FOLDER = Path("diffusers-gallery-data")
 
-CLASSIFIER_URL = "https://radames-aesthetic-style-nsfw-classifier.hf.space/run/inference"
+CLASSIFIER_URL = (
+    "https://radames-aesthetic-style-nsfw-classifier.hf.space/run/inference"
+)
 ASSETS_URL = "https://d26smi9133w0oo.cloudfront.net/diffusers-gallery/"
 
 
-s3 = boto3.client(service_name='s3',
-                  aws_access_key_id=AWS_ACCESS_KEY_ID,
-                  aws_secret_access_key=AWS_SECRET_KEY)
+s3 = boto3.client(
+    service_name="s3",
+    aws_access_key_id=AWS_ACCESS_KEY_ID,
+    aws_secret_access_key=AWS_SECRET_KEY,
+)
 
 
 repo = Repository(
@@ -54,24 +58,31 @@ 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}")
     try:
         async with session.get(image_url) as response:
-            if response.status == 200 and (response.headers['content-type'].startswith('image') or response.headers['content-type'].startswith('application')):
-                image = Image.open(BytesIO(await response.read())).convert('RGB')
+            if response.status == 200 and (
+                response.headers["content-type"].startswith("image")
+                or response.headers["content-type"].startswith("application")
+            ):
+                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"})
+                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
     except Exception as e:
         print(f"Error uploading image {image_url}: {e}")
@@ -80,41 +91,46 @@ async def upload_resize_image_url(session, image_url):
 
 def fetch_models(page=0):
     response = requests.get(
-        f'https://huggingface.co/models-json?pipeline_tag=text-to-image&p={page}')
+        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']],
-        "numItemsPerPage": data['numItemsPerPage'],
-        "numTotalItems": data['numTotalItems'],
-        "pageIndex": data['pageIndex']
+        "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_id):
-    response = requests.get(
-        f'https://huggingface.co/{model_id}/raw/main/README.md')
+    response = requests.get(f"https://huggingface.co/{model_id}/raw/main/README.md")
     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
+REGEX = re.compile(r'---(.*?)---', re.DOTALL)
 
+def get_yaml_data(text_content):
+    matches = REGEX.findall(text_content)
+    yaml_block = matches[0].strip() if matches else None
+    if yaml_block:
+        try: 
+            data_dict = yaml.safe_load(yaml_block)
+            return data_dict
+        except yaml.YAMLError as exc:
+            print(exc)
+    return {}
 
 async def find_image_in_model_card(text):
-    image_regex = re.compile(r'https?://\S+(?:png|jpg|jpeg|webp)')
+    image_regex = re.compile(r"https?://\S+(?:png|jpg|jpeg|webp)")
     urls = re.findall(image_regex, text)
     if not urls:
         return []
 
     async with aiohttp.ClientSession() as session:
-        tasks = [asyncio.ensure_future(upload_resize_image_url(
-            session, image_url)) for image_url in urls[0:3]]
+        tasks = [
+            asyncio.ensure_future(upload_resize_image_url(session, image_url))
+            for image_url in urls[0:3]
+        ]
         return await asyncio.gather(*tasks)
 
 
@@ -123,17 +139,21 @@ def run_classifier(images):
     if len(images) > 0:
         # classifying only the first image
         images_urls = [ASSETS_URL + images[0]]
-        response = requests.post(CLASSIFIER_URL, json={"data": [
-            {"urls": images_urls},  # json urls: list of images urls
-            False,  # enable/disable gallery image output
-            None,  # single image input
-            None,  # files input
-        ]}).json()
+        response = requests.post(
+            CLASSIFIER_URL,
+            json={
+                "data": [
+                    {"urls": images_urls},  # json urls: list of images urls
+                    False,  # enable/disable gallery image output
+                    None,  # single image input
+                    None,  # files input
+                ]
+            },
+        ).json()
 
         # data response is array data:[[{img0}, {img1}, {img2}...], Label, Gallery],
-        class_data = response['data'][0][0]
-        class_data_parsed = {row['label']: round(
-            row['score'], 3) for row in class_data}
+        class_data = response["data"][0][0]
+        class_data_parsed = {row["label"]: round(row["score"], 3) for row in class_data}
 
         # update row data with classificator data
         return class_data_parsed
@@ -143,10 +163,11 @@ def run_classifier(images):
 
 async def get_all_new_models():
     initial = fetch_models(0)
-    num_pages = ceil(initial['numTotalItems'] / initial['numItemsPerPage'])
+    num_pages = ceil(initial["numTotalItems"] / initial["numItemsPerPage"])
 
     print(
-        f"Total items: {initial['numTotalItems']} - Items per page: {initial['numItemsPerPage']}")
+        f"Total items: {initial['numTotalItems']} - Items per page: {initial['numItemsPerPage']}"
+    )
     print(f"Found {num_pages} pages")
 
     # fetch all models
@@ -154,7 +175,7 @@ async def get_all_new_models():
     for page in tqdm(range(0, num_pages)):
         print(f"Fetching page {page} of {num_pages}")
         page_models = fetch_models(page)
-        new_models += page_models['models']
+        new_models += page_models["models"]
     return new_models
 
 
@@ -169,24 +190,28 @@ async def sync_data():
     # with open(DB_FOLDER / "models.json", "r") as f:
     #     new_models = json.load(f)
 
-    new_models_ids = [model['id'] for model in all_models]
+    new_models_ids = [model["id"] for model in all_models]
 
     # get existing models
     with database.get_db() as db:
         cursor = db.cursor()
         cursor.execute("SELECT id FROM models")
-        existing_models = [row['id'] for row in cursor.fetchall()]
+        existing_models = [row["id"] for row in cursor.fetchall()]
     models_ids_to_add = list(set(new_models_ids) - set(existing_models))
     # find all models id to add from new_models
-    models = [model for model in all_models if model['id'] in models_ids_to_add]
+    models = [model for model in all_models if model["id"] in models_ids_to_add]
 
     print(f"Found {len(models)} new models")
     for model in tqdm(models):
-        model_id = model['id']
-        likes = model['likes']
-        downloads = model['downloads']
+        model_id = model["id"]
+        print(f"\n\nFetching model {model_id}")
+        likes = model["likes"]
+        downloads = model["downloads"]
+        print("Fetching model card")
         model_card = fetch_model_card(model_id)
+        print("Parsing model card")
         model_card_data = get_yaml_data(model_card)
+        print("Finding images in model card")
         images = await find_image_in_model_card(model_card)
 
         classifier = run_classifier(images)
@@ -194,58 +219,69 @@ async def sync_data():
         # update model row with image and classifier data
         with database.get_db() as db:
             cursor = db.cursor()
-            cursor.execute("INSERT INTO models(id, data, likes, downloads) VALUES (?, ?, ?, ?)",
-                           [model_id,
-                            json.dumps({
-                                **model,
-                                "meta": model_card_data,
-                                "images": images,
-                                "class": classifier
-                            }),
-                            likes,
-                            downloads
-                            ])
+            cursor.execute(
+                "INSERT INTO models(id, data, likes, downloads) VALUES (?, ?, ?, ?)",
+                [
+                    model_id,
+                    json.dumps(
+                        {
+                            **model,
+                            "meta": model_card_data,
+                            "images": images,
+                            "class": classifier,
+                        }
+                    ),
+                    likes,
+                    downloads,
+                ],
+            )
             db.commit()
-    print("Try to update images again")
+    print("\n\n\n\nTry to update images again\n\n\n")
     with database.get_db() as db:
         cursor = db.cursor()
-        cursor.execute(
-            "SELECT * from models")
+        cursor.execute("SELECT * from models")
         to_all_models = list(cursor.fetchall())
         models_no_images = []
         for model in to_all_models:
-            model_data = json.loads(model['data'])
-            images = model_data['images']
+            model_data = json.loads(model["data"])
+            images = model_data["images"]
             filtered_images = [x for x in images if x is not None]
             if len(filtered_images) == 0:
                 models_no_images.append(model)
 
         for model in tqdm(models_no_images):
-            model_id = model['id']
-            model_data = json.loads(model['data'])
+            model_id = model["id"]
+            model_data = json.loads(model["data"])
+            print(f"\n\nFetching model {model_id}")
             model_card = fetch_model_card(model_id)
+            print("Parsing model card")
             model_card_data = get_yaml_data(model_card)
+            print("Finding images in 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
+            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()
-                cursor.execute("UPDATE models SET data = ? WHERE id = ?",
-                               [json.dumps(model_data), model_id])
+                cursor.execute(
+                    "UPDATE models SET data = ? WHERE id = ?",
+                    [json.dumps(model_data), model_id],
+                )
                 db.commit()
 
     print("Update likes and downloads")
     for model in tqdm(all_models):
-        model_id = model['id']
-        likes = model['likes']
-        downloads = model['downloads']
+        model_id = model["id"]
+        likes = model["likes"]
+        downloads = model["downloads"]
         with database.get_db() as db:
             cursor = db.cursor()
-            cursor.execute("UPDATE models SET likes = ?, downloads = ? WHERE id = ?",
-                           [likes, downloads, model_id])
+            cursor.execute(
+                "UPDATE models SET likes = ?, downloads = ? WHERE id = ?",
+                [likes, downloads, model_id],
+            )
             db.commit()
 
     print("Updating DB repository")
@@ -288,8 +324,10 @@ class Style(str, Enum):
     nsfw = "nsfw"
 
 
-@ app.get("/api/models")
-def get_page(page: int = 1, sort: Sort = Sort.trending, style: Style = Style.all, tag: str = None):
+@app.get("/api/models")
+def get_page(
+    page: int = 1, sort: Sort = Sort.trending, style: Style = Style.all, tag: str = None
+):
     page = page if page > 0 else 1
     if sort == Sort.trending:
         sort_query = "likes / MYPOWER((JULIANDAY('now') - JULIANDAY(datetime(json_extract(data, '$.lastModified')))) + 2, 2) DESC"
@@ -311,7 +349,8 @@ def get_page(page: int = 1, sort: Sort = Sort.trending, style: Style = Style.all
 
     with database.get_db() as db:
         cursor = db.cursor()
-        cursor.execute(f"""
+        cursor.execute(
+            f"""
                 SELECT *,
                         COUNT(*) OVER() AS total,
                         isNFSW
@@ -329,36 +368,37 @@ def get_page(page: int = 1, sort: Sort = Sort.trending, style: Style = Style.all
                         ))
                     ORDER BY {sort_query}
                     LIMIT {MAX_PAGE_SIZE} OFFSET {(page - 1) * MAX_PAGE_SIZE};
-        """, (tag, tag, tag, tag))
+        """,
+            (tag, tag, tag, tag),
+        )
         results = cursor.fetchall()
-        total = results[0]['total'] if results else 0
+        total = results[0]["total"] if results else 0
         total_pages = (total + MAX_PAGE_SIZE - 1) // MAX_PAGE_SIZE
         models_data = []
         for result in results:
-            data = json.loads(result['data'])
-            images = data['images']
+            data = json.loads(result["data"])
+            images = data["images"]
             filtered_images = [x for x in images if x is not None]
             # clean nulls
-            data['images'] = filtered_images
+            data["images"] = filtered_images
             # update downloads and likes from db table
-            data['downloads'] = result['downloads']
-            data['likes'] = result['likes']
-            data['isNFSW'] = bool(result['isNFSW'])
+            data["downloads"] = result["downloads"]
+            data["likes"] = result["likes"]
+            data["isNFSW"] = bool(result["isNFSW"])
             models_data.append(data)
 
-    return {
-        "models": models_data,
-        "totalPages": total_pages
-    }
+    return {"models": models_data, "totalPages": total_pages}
 
 
 @app.get("/")
 def read_root():
     # return html page from string
-    return HTMLResponse("""
+    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" target="_blank" rel="noopener noreferrer">https://huggingface.co/spaces/huggingface-projects/diffusers-gallery</a>
-    </p>""")
+    </p>"""
+    )
 
 
 @app.on_event("startup")