tools / get_nfsw_civitai.py
patrickvonplaten's picture
[LCM] Better error message
c9b0b96
raw
history blame contribute delete
No virus
3.86 kB
#!/usr/bin/env python3
import requests
import json
from collections import Counter
def get_model_counts():
# Base URL of the API endpoint
url = "https://civitai.com/api/v1/models"
# Parameters for the API request
params = {
'limit': 100, # Set a high limit to reduce the number of pages to iterate through
'page': 1 # Start from the first page
}
# Initialize counts
total_models = 0
# counter
allow_commercial_use = []
creator = []
# sum
nsfw_count = []
allow_derivatives = []
allow_no_credit = []
download_count = []
favorite_count = []
comment_count = []
rating_count = []
tipped_amount_count = []
ratings = []
failures = 0
total_pages = 1000 # set to very high value, just in case it fails a lot in the beginning
while True:
# Make the GET request
try:
response = requests.get(url, params=params)
response.raise_for_status() # This will raise an error if the request fails
data = response.json()
# Process the current page of results
models = data["items"]
total_models += len(models)
for model in models:
allow_commercial_use.append(model["allowCommercialUse"])
creator.append(model["creator"]["username"])
nsfw_count.append(model["nsfw"])
allow_derivatives.append(model["allowDerivatives"])
allow_no_credit.append(model["allowNoCredit"])
download_count.append(model["stats"]["downloadCount"])
favorite_count.append(model["stats"]["favoriteCount"])
comment_count.append(model["stats"]["commentCount"])
rating_count.append(model["stats"]["ratingCount"])
tipped_amount_count.append(model["stats"]["tippedAmountCount"])
ratings.append(model["stats"]["rating"])
# Check if there are more pages to process
total_pages = int(data.get('metadata', {}).get('totalPages', 0))
# if params['page'] >= total_pages:
if params['page'] >= 2:
break
# Increment the page number for the next request
params['page'] += 1
print(f"{params['page']} / {total_pages}")
except:
failures += 1
params["page"] += 1
return {
'total_models': total_models,
'allow_commercial_use': allow_commercial_use,
'creator': creator,
'nsfw_count': nsfw_count,
'allow_derivatives': allow_derivatives,
'allow_no_credit': allow_no_credit,
'download_count': download_count,
'favorite_count': favorite_count,
'comment_count': comment_count,
'rating_count': rating_count,
'tipped_amount_count': tipped_amount_count,
'ratings': ratings,
'failures': failures,
}
outputs = get_model_counts()
def map_fn(k, v):
if k in ["total_models", 'failures']:
return v
elif k in ["creator", "allow_commercial_use"]:
return Counter(v)
elif k in ["ratings"]:
return sum(v) / len(v)
else:
return {
"sum": sum(v),
"avg": round(sum(v) / len(v), 3)
}
stats = {k: map_fn(k, v) for k,v in outputs.items()}
stats["num_creators"] = len(stats.pop("creator"))
for k in ["total_models", "download_count", "comment_count", "rating_count", "tipped_amount_count", "favorite_count"]:
if isinstance(stats[k], dict):
stats[f"{k}_per_creator"] = round(stats[k]["sum"] / stats["num_creators"], 2)
else:
stats[f"{k}_per_creator"] = round(stats[k] / stats["num_creators"], 2)
filename = "civitai_stats.json"
with open(filename, 'w') as file:
json.dump(stats, file, indent=4)