civitai-to-hf / app.py
multimodalart's picture
Create app.py
2aeb649
raw
history blame
7.05 kB
import requests
import os
import gradio as gr
from huggingface_hub import HfApi
from slugify import slugify
import gradio as gr
import uuid
from typing import Optional
def get_json_data(url):
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):
if json_data["nsfw"]:
return False
for model_version in json_data["modelVersions"]:
for image in model_version["images"]:
if image["nsfw"] != "None":
return False
return True
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"]:
for file in model_version["files"]:
if file["primary"]:
info = {
"urls_to_download": [
{"url": file["downloadUrl"], "filename": file["name"], "type": "weightName"},
{"url": model_version["images"][0]["url"], "filename": os.path.basename(model_version["images"][0]["url"]), "type": "imageName"}
],
"id": model_version["id"],
"modelId": model_version["modelId"],
"name": json_data["name"],
"description": json_data["description"],
"trainedWords": model_version["trainedWords"],
"creator": json_data["creator"]["username"]
}
return info
return None
def download_files(info, folder="."):
downloaded_files = {
"imageName": [],
"weightName": []
}
for item in info["urls_to_download"]:
download_file(item["url"], item["filename"], folder)
downloaded_files[item["type"]].append(item["filename"])
return downloaded_files
def download_file(url, filename, folder="."):
try:
response = requests.get(url)
response.raise_for_status()
with open(f"{folder}/{filename}", 'wb') as f:
f.write(response.content)
print(f"{filename} downloaded.")
except requests.exceptions.RequestException as e:
print(f"Error downloading file: {e}")
def process_url(url, folder="."):
json_data = get_json_data(url)
if json_data:
if check_nsfw(json_data):
info = extract_info(json_data)
if info:
downloaded_files = download_files(info, folder)
return info, downloaded_files
else:
print("No model met the criteria.")
else:
print("NSFW content found.")
else:
print("Failed to get JSON data.")
def create_readme(info, downloaded_files, is_author, folder="."):
readme_content = ""
original_url = f"https://civitai.com/models/{info['id']}"
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:'
content = f"""---
license: other
tags:
- text-to-image
- stable-diffusion
- lora
- diffusers
base_model: stabilityai/stable-diffusion-xl-base-1.0
instance_prompt: {info["trainedWords"][0]}
widget:
- text: {info["trainedWords"][0]}
---
# {info["name"]}
{non_author_disclaimer if not is_author else ''}
![Image]({downloaded_files["imageName"][0]})
{info["description"]}
"""
readme_content += content + "\n"
with open(f"{folder}/README.md", "w") as file:
file.write(readme_content)
def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], url, is_author, progress=gr.Progress(track_tqdm=True)):
if not profile.name:
return gr.Error("Are you sure you are logged in?")
folder = str(uuid.uuid4())
os.makedirs(folder, exist_ok=False)
info, downloaded_files = process_url(url, folder)
create_readme(info, downloaded_files, False, folder)
try:
api = HfApi(token=hf_token)
username = api.whoami()["name"]
slug_name = slugify(info["name"])
repo_id = f"{username}/{slug_name}"
api.create_repo(repo_id=repo_id, private=True, exist_ok=True)
api.upload_folder(
folder_path=folder,
repo_id=repo_id,
repo_type="model"
)
except:
raise gr.Error("something went wrong")
return "Model uploaded!"
def swap_fill(profile: Optional[gr.OAuthProfile]):
if profile is None:
return gr.update(visible=True), gr.update(visible=False)
else:
return gr.update(visible=False), gr.update(visible=True)
css = '''
#login {
font-size: 0px;
width: 100% !important;
margin: 0 auto;
}
#login:after {
content: 'Authorize this app before uploading your model';
visibility: visible;
display: block;
font-size: var(--button-large-text-size);
}
#login:disabled{
font-size: var(--button-large-text-size);
}
#login:disabled:after{
content:''
}
#disabled_upload{
opacity: 0.5;
pointer-events:none;
}
'''
with gr.Blocks(css=css) as demo:
gr.LoginButton(elem_id="login")
with gr.Column(elem_id="disabled_upload") as disabled_area:
with gr.Row():
submit_source_civit = gr.Textbox(
label="CivitAI model URL",
info="URL of the CivitAI model, make sure it is a SDXL LoRA",
)
is_author = gr.Checkbox(label="Are you the model author?", info="If you are not the author, a disclaimer with information about the author and the CivitAI source will be added", value=False)
submit_button_civit = gr.Button("Upload model to Hugging Face and submit")
output = gr.Textbox(label="Output progress")
with gr.Column(visible=False) as enabled_area:
with gr.Row():
submit_source_civit = gr.Textbox(
label="CivitAI model URL",
info="URL of the CivitAI model, make sure it is a SDXL LoRA",
)
is_author = gr.Checkbox(label="Are you the model author?", info="If you are not the author, a disclaimer with information about the author and the CivitAI source will be added", value=False)
submit_button_civit = gr.Button("Upload model to Hugging Face")
output = gr.Textbox(label="Output progress")
demo.load(fn=swap_fill, outputs=[disabled_area, enabled_area])
submit_button_civit.click(fn=upload_civit_to_hf, inputs=[submit_source_civit, is_author], outputs=[output])
demo.queue()
demo.launch(share=True)