import requests
import os
import gradio as gr
from huggingface_hub import update_repo_visibility, whoami, upload_folder, create_repo, upload_file, update_repo_visibility
from slugify import slugify
import gradio as gr
import re
import uuid
from typing import Optional
import json
from bs4 import BeautifulSoup
TRUSTED_UPLOADERS = ["KappaNeuro", "CiroN2022", "multimodalart", "Norod78", "joachimsallstrom", "blink7630", "e-n-v-y", "DoctorDiffusion", "RalFinger", "artificialguybr"]
def get_json_data(url):
url_split = url.split('/')
print("url_split: ", url_split)
api_url = f"https://civitai.com/api/v1/models/{url_split[4]}"
try:
response = requests.get(api_url)
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
print(f"Error fetching JSON data: {e}")
return None
def check_nsfw(json_data, profile):
if json_data["nsfw"]:
return False
print(profile)
if(profile.username in TRUSTED_UPLOADERS):
return True
for model_version in json_data["modelVersions"]:
for image in model_version["images"]:
if image["nsfw"] > 5:
return False
return True
def get_prompts_from_image(image_id):
print("image_id: ", image_id)
url = f'https://civitai.com/images/{image_id}'
response = requests.get(url)
soup = BeautifulSoup(response.text, 'html.parser')
content = soup.find_all(class_='mantine-Code-root mantine-Code-block mantine-2v44jn')
print(content)
if(content):
if(len(content) > 1):
return content[0].text, content[1].text
else:
return content[0].text, ""
else:
return "", ""
def extract_info(json_data):
if json_data["type"] == "LORA":
for model_version in json_data["modelVersions"]:
if model_version["baseModel"] in ["SDXL 1.0", "SDXL 0.9", "SD 1.5", "SD 1.4", "SD 2.1", "SD 2.0", "SD 2.0 768", "SD 2.1 768"]:
for file in model_version["files"]:
print(file)
if "primary" in file:
# Start by adding the primary file to the list
urls_to_download = [{"url": file["downloadUrl"], "filename": file["name"], "type": "weightName"}]
# Then append all image URLs to the list
for image in model_version["images"]:
image_id = image["url"].split("/")[-1].split(".")[0]
prompt, negative_prompt = get_prompts_from_image(image_id)
if image["nsfwLevel"] > 5:
pass #ugly before checking the actual logic
else:
urls_to_download.append({
"url": image["url"],
"filename": os.path.basename(image["url"]),
"type": "imageName",
"prompt": prompt, #if "meta" in image and "prompt" in image["meta"] else ""
"negative_prompt": negative_prompt
})
model_mapping = {
"SDXL 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
"SDXL 0.9": "stabilityai/stable-diffusion-xl-base-1.0",
"SD 1.5": "runwayml/stable-diffusion-v1-5",
"SD 1.4": "CompVis/stable-diffusion-v1-4",
"SD 2.1": "stabilityai/stable-diffusion-2-1-base",
"SD 2.0": "stabilityai/stable-diffusion-2-base",
"SD 2.1 768": "stabilityai/stable-diffusion-2-1",
"SD 2.0 768": "stabilityai/stable-diffusion-2"
}
base_model = model_mapping[model_version["baseModel"]]
print(json_data)
info = {
"urls_to_download": urls_to_download,
"id": model_version["id"],
"baseModel": base_model,
"modelId": model_version.get("modelId", ""),
"name": json_data["name"],
"description": json_data["description"],
"trainedWords": model_version["trainedWords"] if "trainedWords" in model_version else [],
"creator": json_data["creator"]["username"],
"tags": json_data["tags"],
"allowNoCredit": json_data["allowNoCredit"],
"allowCommercialUse": json_data["allowCommercialUse"],
"allowDerivatives": json_data["allowDerivatives"],
"allowDifferentLicense": json_data["allowDifferentLicense"]
}
return info
return None
def download_files(info, folder="."):
downloaded_files = {
"imageName": [],
"imagePrompt": [],
"imageNegativePrompt": [],
"weightName": []
}
for item in info["urls_to_download"]:
download_file(item["url"], item["filename"], folder)
downloaded_files[item["type"]].append(item["filename"])
if(item["type"] == "imageName"):
prompt_clean = re.sub(r'<.*?>', '', item["prompt"])
negative_prompt_clean = re.sub(r'<.*?>', '', item["negative_prompt"])
downloaded_files["imagePrompt"].append(prompt_clean)
downloaded_files["imageNegativePrompt"].append(negative_prompt_clean)
return downloaded_files
def download_file(url, filename, folder="."):
headers = {}
try:
response = requests.get(url, headers=headers)
response.raise_for_status()
except requests.exceptions.HTTPError as e:
print(e)
if response.status_code == 401:
headers['Authorization'] = f'Bearer {os.environ["CIVITAI_API"]}'
try:
response = requests.get(url, headers=headers)
response.raise_for_status()
except requests.exceptions.RequestException as e:
raise gr.Error(f"Error downloading file: {e}")
else:
raise gr.Error(f"Error downloading file: {e}")
except requests.exceptions.RequestException as e:
raise gr.Error(f"Error downloading file: {e}")
with open(f"{folder}/{filename}", 'wb') as f:
f.write(response.content)
def process_url(url, profile, do_download=True, folder="."):
json_data = get_json_data(url)
if json_data:
if check_nsfw(json_data, profile):
info = extract_info(json_data)
if info:
if(do_download):
downloaded_files = download_files(info, folder)
else:
downloaded_files = []
return info, downloaded_files
else:
raise gr.Error("Only SDXL LoRAs are supported for now")
else:
raise gr.Error("This model has content tagged as unsafe by CivitAI")
else:
raise gr.Error("Something went wrong in fetching CivitAI API")
def create_readme(info, downloaded_files, user_repo_id, link_civit=False, is_author=True, folder="."):
readme_content = ""
original_url = f"https://civitai.com/models/{info['modelId']}"
link_civit_disclaimer = f'([CivitAI]({original_url}))'
non_author_disclaimer = f'This model was originally uploaded on [CivitAI]({original_url}), by [{info["creator"]}](https://civitai.com/user/{info["creator"]}/models). The information below was provided by the author on CivitAI:'
default_tags = ["text-to-image", "stable-diffusion", "lora", "diffusers", "template:sd-lora", "migrated"]
civit_tags = [t.replace(":", "") for t in info["tags"] if t not in default_tags]
tags = default_tags + civit_tags
unpacked_tags = "\n- ".join(tags)
trained_words = info['trainedWords'] if 'trainedWords' in info and info['trainedWords'] else []
formatted_words = ', '.join(f'`{word}`' for word in trained_words)
if formatted_words:
trigger_words_section = f"""## Trigger words
You should use {formatted_words} to trigger the image generation.
"""
else:
trigger_words_section = ""
widget_content = ""
for index, (prompt, negative_prompt, image) in enumerate(zip(downloaded_files["imagePrompt"], downloaded_files["imageNegativePrompt"], downloaded_files["imageName"])):
escaped_prompt = prompt.replace("'", "''")
negative_prompt_content = f"""parameters:
negative_prompt: {negative_prompt}
""" if negative_prompt else ""
widget_content += f"""- text: '{escaped_prompt if escaped_prompt else ' ' }'
{negative_prompt_content}
output:
url: >-
{image}
"""
content = f"""---
license: other
license_name: bespoke-lora-trained-license
license_link: https://multimodal.art/civitai-licenses?allowNoCredit={info["allowNoCredit"]}&allowCommercialUse={info["allowCommercialUse"][0]}&allowDerivatives={info["allowDerivatives"]}&allowDifferentLicense={info["allowDifferentLicense"]}
tags:
- {unpacked_tags}
base_model: {info["baseModel"]}
instance_prompt: {info['trainedWords'][0] if 'trainedWords' in info and len(info['trainedWords']) > 0 else ''}
widget:
{widget_content}
---
# {info["name"]}
(if you are not {info["creator"]}, you cannot submit their model at this time)'
return no_username_text, gr.update(interactive=False), gr.update(visible=True), gr.update(visible=False)
if(profile.username != hf_username):
unmatched_username_text = '