app.py

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, download_files=True, folder="."):
    json_data = get_json_data(url)
    if json_data:
        if check_nsfw(json_data):
            info = extract_info(json_data)
            if info:
                if(download_files):
                    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, is_author=True, 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, 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, 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 check_civit_link(profile: Optional[gr.OAuthProfile], url):
    info, _ = process_url(url, download_files=False)
    url_creator = f"https://civitai.com/user/{info['creator']}/models" 
    # Open the target URL
    driver.get(url_creator)
    
    # Define the XPath expression
    xpath_expression = "//a[contains(@class, 'mantine-UnstyledButton-root') and contains(@class, 'mantine-ActionIcon-root') and contains(@class, 'mantine-ubxmi3') and starts-with(@href, 'https://huggingface.co/')]"
    
    # Find the element using the XPath expression
    try:
        element = WebDriverWait(driver, 10).until(
            EC.presence_of_element_located((By.XPATH, xpath_expression))
        )
        
        # Extract the href attribute
        href = element.get_attribute("href")
        
        # Extract the part after 'https://huggingface.co/'
        extracted_part = href.replace("https://huggingface.co/", "")
    
    except Exception as e:
        print("Element not found or error occurred:", e)
    
    finally:
        driver.quit()
    
    return extracted_part == profile.name
    
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)
        instructions = gr.HTML("")
        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_source_civit.change(fn=check_civit_link, inputs=[submit_source_civit], outputs=[instructions])
    submit_button_civit.click(fn=upload_civit_to_hf, inputs=[submit_source_civit], outputs=[output])
    
demo.queue()
demo.launch()