Upload civitai_to_hf.py

civitai_to_hf.py

@@ -251,7 +251,8 @@ CIVITAI_BASEMODEL = ["Aura Flow", "CogVideoX", "Flux.1 D", "Flux.1 S", "HiDream"
                      "SDXL 1.0 LCM", "SDXL Distilled", "SDXL Hyper", "SDXL Lightning", "SDXL Turbo", "SVD", "SVD XT", "Stable Cascade",
                      "Wan Video", "VPRED"]
 #CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Newest"]
-CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "Newest"]
+CIVITAI_SORT_EXT = ["Size", "Size (from smallest)"]
+CIVITAI_SORT = ["Highest Rated", "Most Downloaded", "Most Liked", "Most Discussed", "Most Collected", "Most Buzz", "Newest"] + CIVITAI_SORT_EXT
 CIVITAI_PERIOD = ["AllTime", "Year", "Month", "Week", "Day"]
 
 
@@ -262,7 +263,7 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
     headers = {'User-Agent': user_agent, 'content-type': 'application/json'}
     if api_key: headers['Authorization'] = f'Bearer {{{api_key}}}'
     base_url = 'https://civitai.com/api/v1/models'
-    params = {'sort': sort, 'period': period, 'limit': int(limit), 'nsfw': 'true'}
+    params = {'sort': CIVITAI_SORT[0] if sort in CIVITAI_SORT_EXT else sort, 'period': period, 'limit': int(limit), 'nsfw': 'true'}
     if len(types) != 0: params["types"] = types
     if query: params["query"] = query
     if tag: params["tag"] = tag
@@ -327,11 +328,15 @@ def search_on_civitai(query: str, types: list[str], allow_model: list[str] = [],
                     for f in model['files']:
                         i = item.copy()
                         i['dl_url'] = f['downloadUrl']
+                        i['size_kb'] = f['sizeKB']
                         if len(filetype) != 0 and f['type'] not in set(filetype): continue
                         items.append(i)
                 else:
                     item['dl_url'] = model['downloadUrl']
                     items.append(item)
+        if sort in CIVITAI_SORT_EXT:
+            if sort == "Size": items = sorted(items, key=lambda x: x.get('size_kb', 0.), reverse=True)
+            elif sort == "Size (from smallest)": items = sorted(items, key=lambda x: x.get('size_kb', 0.))
     return items if len(items) > 0 else None
 
 
@@ -349,7 +354,7 @@ def search_civitai(query, types, base_model=[], sort=CIVITAI_SORT[0], period=CIV
     gallery = []
     for item in items:
         base_model_name = "Pony🐴" if item['base_model'] == "Pony" else item['base_model']
-        name = f"{item['name']} (for {base_model_name} / By: {item['creator']})"
+        name = f"{item['name']} (for {base_model_name} / By: {item['creator']}) ({round(item['size_kb'] / 1000., 2)}MB)" if "size_kb" in item.keys() else f"{item['name']} (for {base_model_name} / By: {item['creator']})"
         value = item['dl_url']
         choices.append((name, value))
         gallery.append((item['img_url'], name))