Spaces:
Running
Running
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) |