convert_url_to_diffusers_sdxl_gr.py

import argparse
from pathlib import Path
import os
import torch
from diffusers import StableDiffusionXLPipeline, AutoencoderKL
from transformers import CLIPTokenizer, CLIPTextModel
import gradio as gr
from huggingface_hub import hf_hub_download, HfApi
import urllib.parse
import re
import shutil
import gc
# also requires aria, gdown, peft, huggingface_hub, safetensors, transformers, accelerate, pytorch_lightning


def list_sub(a, b):
    return [e for e in a if e not in b]


def is_repo_name(s):
    return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)


def split_hf_url(url: str):
    try:
        s = list(re.findall(r'^(?:https?://huggingface.co/)(?:(datasets)/)?(.+?/.+?)/\w+?/.+?/(?:(.+)/)?(.+?.\w+)(?:\?download=true)?$', url)[0])
        if len(s) < 4: return "", "", "", ""
        repo_id = s[1]
        repo_type = "dataset" if s[0] == "datasets" else "model"
        subfolder = urllib.parse.unquote(s[2]) if s[2] else None
        filename = urllib.parse.unquote(s[3])
        return repo_id, filename, subfolder, repo_type
    except Exception as e:
        print(e)


def download_hf_file(directory, url, hf_token="", progress=gr.Progress(track_tqdm=True)):
    repo_id, filename, subfolder, repo_type = split_hf_url(url)
    try:
        if subfolder is not None: hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder, repo_type=repo_type, local_dir=directory, token=hf_token)
        else: hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type, local_dir=directory, token=hf_token)
    except Exception as e:
        print(f"Failed to download: {e}")


def download_thing(directory, url, civitai_api_key="", hf_token="", progress=gr.Progress(track_tqdm=True)):
    url = url.strip()
    if "drive.google.com" in url:
        original_dir = os.getcwd()
        os.chdir(directory)
        os.system(f"gdown --fuzzy {url}")
        os.chdir(original_dir)
    elif "huggingface.co" in url:
        url = url.replace("?download=true", "")
        if "/blob/" in url:
            url = url.replace("/blob/", "/resolve/")
        user_header = f'"Authorization: Bearer {hf_token}"'
        if hf_token:
            download_hf_file(directory, url, hf_token)
            #os.system(f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
        else:
            os.system(f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
    elif "civitai.com" in url:
        if "?" in url:
            url = url.split("?")[0]
        if civitai_api_key:
            url = url + f"?token={civitai_api_key}"
            os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
        else:
            print("You need an API key to download Civitai models.")
    else:
        os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")


def get_local_model_list(dir_path):
    model_list = []
    valid_extensions = ('.safetensors')
    for file in Path(dir_path).glob("**/*.*"):
        if file.is_file() and file.suffix in valid_extensions:
            file_path = str(file)
            model_list.append(file_path)
    return model_list


def get_download_file(temp_dir, url, civitai_key, hf_token, progress=gr.Progress(track_tqdm=True)):
    if not "http" in url and is_repo_name(url) and not Path(url).exists():
        print(f"Use HF Repo: {url}")
        new_file = url
    elif not "http" in url and Path(url).exists():
        print(f"Use local file: {url}")
        new_file = url
    elif Path(f"{temp_dir}/{url.split('/')[-1]}").exists():
        print(f"File to download alreday exists: {url}")
        new_file = f"{temp_dir}/{url.split('/')[-1]}"
    else:
        print(f"Start downloading: {url}")
        before = get_local_model_list(temp_dir)
        try:
            download_thing(temp_dir, url.strip(), civitai_key, hf_token)
        except Exception:
            print(f"Download failed: {url}")
            return ""
        after = get_local_model_list(temp_dir)
        new_file = list_sub(after, before)[0] if list_sub(after, before) else ""
    if not new_file:
        print(f"Download failed: {url}")
        return ""
    print(f"Download completed: {url}")
    return new_file


from diffusers import (
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    KDPM2DiscreteScheduler,
    EulerDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
    HeunDiscreteScheduler,
    LMSDiscreteScheduler,
    DDIMScheduler,
    DEISMultistepScheduler,
    UniPCMultistepScheduler,
    LCMScheduler,
    PNDMScheduler,
    KDPM2AncestralDiscreteScheduler,
    DPMSolverSDEScheduler,
    EDMDPMSolverMultistepScheduler,
    DDPMScheduler,
    EDMEulerScheduler,
    TCDScheduler,
)


SCHEDULER_CONFIG_MAP = {
    "DPM++ 2M": (DPMSolverMultistepScheduler, {"use_karras_sigmas": False}),
    "DPM++ 2M Karras": (DPMSolverMultistepScheduler, {"use_karras_sigmas": True}),
    "DPM++ 2M SDE": (DPMSolverMultistepScheduler, {"use_karras_sigmas": False, "algorithm_type": "sde-dpmsolver++"}),
    "DPM++ 2M SDE Karras": (DPMSolverMultistepScheduler, {"use_karras_sigmas": True, "algorithm_type": "sde-dpmsolver++"}),
    "DPM++ 2S": (DPMSolverSinglestepScheduler, {"use_karras_sigmas": False}),
    "DPM++ 2S Karras": (DPMSolverSinglestepScheduler, {"use_karras_sigmas": True}),
    "DPM++ 1S": (DPMSolverMultistepScheduler, {"solver_order": 1}),
    "DPM++ 1S Karras": (DPMSolverMultistepScheduler, {"solver_order": 1, "use_karras_sigmas": True}),
    "DPM++ 3M": (DPMSolverMultistepScheduler, {"solver_order": 3}),
    "DPM++ 3M Karras": (DPMSolverMultistepScheduler, {"solver_order": 3, "use_karras_sigmas": True}),
    "DPM++ SDE": (DPMSolverSDEScheduler, {"use_karras_sigmas": False}),
    "DPM++ SDE Karras": (DPMSolverSDEScheduler, {"use_karras_sigmas": True}),
    "DPM2": (KDPM2DiscreteScheduler, {}),
    "DPM2 Karras": (KDPM2DiscreteScheduler, {"use_karras_sigmas": True}),
    "DPM2 a": (KDPM2AncestralDiscreteScheduler, {}),
    "DPM2 a Karras": (KDPM2AncestralDiscreteScheduler, {"use_karras_sigmas": True}),
    "Euler": (EulerDiscreteScheduler, {}),
    "Euler a": (EulerAncestralDiscreteScheduler, {}),
    "Euler trailing": (EulerDiscreteScheduler, {"timestep_spacing": "trailing", "prediction_type": "sample"}),
    "Euler a trailing": (EulerAncestralDiscreteScheduler, {"timestep_spacing": "trailing"}),
    "Heun": (HeunDiscreteScheduler, {}),
    "Heun Karras": (HeunDiscreteScheduler, {"use_karras_sigmas": True}),
    "LMS": (LMSDiscreteScheduler, {}),
    "LMS Karras": (LMSDiscreteScheduler, {"use_karras_sigmas": True}),
    "DDIM": (DDIMScheduler, {}),
    "DDIM trailing": (DDIMScheduler, {"timestep_spacing": "trailing"}),
    "DEIS": (DEISMultistepScheduler, {}),
    "UniPC": (UniPCMultistepScheduler, {}),
    "UniPC Karras": (UniPCMultistepScheduler, {"use_karras_sigmas": True}),
    "PNDM": (PNDMScheduler, {}),
    "Euler EDM": (EDMEulerScheduler, {}),
    "Euler EDM Karras": (EDMEulerScheduler, {"use_karras_sigmas": True}),
    "DPM++ 2M EDM": (EDMDPMSolverMultistepScheduler, {"solver_order": 2, "solver_type": "midpoint", "final_sigmas_type": "zero", "algorithm_type": "dpmsolver++"}),
    "DPM++ 2M EDM Karras": (EDMDPMSolverMultistepScheduler, {"use_karras_sigmas": True, "solver_order": 2, "solver_type": "midpoint", "final_sigmas_type": "zero", "algorithm_type": "dpmsolver++"}),
    "DDPM": (DDPMScheduler, {}),

    "DPM++ 2M Lu": (DPMSolverMultistepScheduler, {"use_lu_lambdas": True}),
    "DPM++ 2M Ef": (DPMSolverMultistepScheduler, {"euler_at_final": True}),
    "DPM++ 2M SDE Lu": (DPMSolverMultistepScheduler, {"use_lu_lambdas": True, "algorithm_type": "sde-dpmsolver++"}),
    "DPM++ 2M SDE Ef": (DPMSolverMultistepScheduler, {"algorithm_type": "sde-dpmsolver++", "euler_at_final": True}),

    "LCM": (LCMScheduler, {}),
    "TCD": (TCDScheduler, {}),
    "LCM trailing": (LCMScheduler, {"timestep_spacing": "trailing"}),
    "TCD trailing": (TCDScheduler, {"timestep_spacing": "trailing"}),
    "LCM Auto-Loader": (LCMScheduler, {}),
    "TCD Auto-Loader": (TCDScheduler, {}),
}


def get_scheduler_config(name):
    if not name in SCHEDULER_CONFIG_MAP.keys(): return SCHEDULER_CONFIG_MAP["Euler a"]
    return SCHEDULER_CONFIG_MAP[name]


def save_readme_md(dir, url):
    orig_url = ""
    orig_name = ""
    if is_repo_name(url): 
        orig_name = url
        orig_url = f"https://huggingface.co/{url}/"
    elif "http" in url:
        orig_name = url
        orig_url = url
    if orig_name and orig_url:
       md = f"""---
license: other
language:
- en
library_name: diffusers
pipeline_tag: text-to-image
tags:
- text-to-image
---
Converted from [{orig_name}]({orig_url}).
"""
    else:
        md = f"""---
license: other
language:
- en
library_name: diffusers
pipeline_tag: text-to-image
tags:
- text-to-image
---
"""
    path = str(Path(dir, "README.md"))
    with open(path, mode='w', encoding="utf-8") as f:
        f.write(md)


def fuse_loras(pipe, lora_dict={}, temp_dir=".", civitai_key="", hf_token=""):
    if not lora_dict or not isinstance(lora_dict, dict): return pipe
    a_list = []
    w_list = []
    for k, v in lora_dict.items():
        if not k: continue
        new_lora_file = get_download_file(temp_dir, k, civitai_key, hf_token)
        if not new_lora_file or not Path(new_lora_file).exists():
            print(f"LoRA not found: {k}")
            continue
        w_name = Path(new_lora_file).name
        a_name = Path(new_lora_file).stem
        pipe.load_lora_weights(new_lora_file, weight_name = w_name, adapter_name = a_name)
        a_list.append(a_name)
        w_list.append(v)
    if not a_list: return pipe
    pipe.set_adapters(a_list, adapter_weights=w_list)
    pipe.fuse_lora(adapter_names=a_list, lora_scale=1.0)
    pipe.unload_lora_weights()
    return pipe


def convert_url_to_diffusers_sdxl(url, civitai_key="", hf_token="", is_upload_sf=False, half=True, vae=None,
                                  scheduler="Euler a", lora_dict={}, is_local=True, clip="", progress=gr.Progress(track_tqdm=True)):
    progress(0, desc="Start converting...")
    temp_dir = "."
    new_file = get_download_file(temp_dir, url, civitai_key, hf_token)
    if not new_file:
        print(f"Not found: {url}")
        return ""
    new_repo_name = Path(new_file).stem.replace(" ", "_").replace(",", "_").replace(".", "_") #

    type_kwargs = {}
    kwargs = {}
    if half: type_kwargs["torch_dtype"] = torch.float16

    new_vae_file = ""
    if vae:
        if is_repo_name(vae): my_vae = AutoencoderKL.from_pretrained(vae, **type_kwargs)
        else:
            new_vae_file = get_download_file(temp_dir, vae, civitai_key, hf_token)
            if new_vae_file: my_vae = AutoencoderKL.from_single_file(new_vae_file, **type_kwargs)
        kwargs["vae"] = my_vae

    if clip:
        my_tokenizer = CLIPTokenizer.from_pretrained(clip)
        my_text_encoder = CLIPTextModel.from_pretrained(clip, **type_kwargs)
        kwargs["tokenizer"] = my_tokenizer
        kwargs["text_encoder"] = my_text_encoder

    pipe = None
    if is_repo_name(url): pipe = StableDiffusionXLPipeline.from_pretrained(new_file, use_safetensors=True, **kwargs, **type_kwargs)
    else: pipe = StableDiffusionXLPipeline.from_single_file(new_file, use_safetensors=True, **kwargs, **type_kwargs)

    pipe = fuse_loras(pipe, lora_dict, temp_dir, civitai_key, hf_token)

    sconf = get_scheduler_config(scheduler)
    pipe.scheduler = sconf[0].from_config(pipe.scheduler.config, **sconf[1])

    pipe.save_pretrained(new_repo_name, safe_serialization=True, use_safetensors=True)

    if Path(new_repo_name).exists(): save_readme_md(new_repo_name, url)

    if not is_local:
        if not is_repo_name(new_file) and is_upload_sf: shutil.move(str(Path(new_file).resolve()), str(Path(new_repo_name, Path(new_file).name).resolve()))
        else: os.remove(new_file)
    del pipe
    torch.cuda.empty_cache()
    gc.collect()

    progress(1, desc="Converted.")
    return new_repo_name


def is_repo_exists(repo_id, hf_token):
    api = HfApi(token=hf_token)
    try:
        if api.repo_exists(repo_id=repo_id): return True
        else: return False
    except Exception as e:
        print(f"Error: Failed to connect {repo_id}. {e}")
        return True # for safe


def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, hf_token, progress=gr.Progress(track_tqdm=True)):
    api = HfApi(token=hf_token)
    try:
        progress(0, desc="Start uploading...")
        api.create_repo(repo_id=new_repo_id, token=hf_token, private=is_private, exist_ok=True)
        for path in Path(diffusers_folder).glob("*"):
            if path.is_dir():
                api.upload_folder(repo_id=new_repo_id, folder_path=str(path), path_in_repo=path.name, token=hf_token)
            elif path.is_file():
                api.upload_file(repo_id=new_repo_id, path_or_fileobj=str(path), path_in_repo=path.name, token=hf_token)
        progress(1, desc="Uploaded.")
        url = f"https://huggingface.co/{new_repo_id}"
    except Exception as e:
        print(f"Error: Failed to upload to {new_repo_id}. {e}")
        return ""
    return url


def convert_url_to_diffusers_repo(dl_url, hf_user, hf_repo, hf_token, civitai_key="", is_private=True, is_overwrite=False, is_upload_sf=False,
                                  repo_urls=[], half=True, vae=None, clip="", scheduler="Euler a",
                                  lora1=None, lora1s=1.0, lora2=None, lora2s=1.0, lora3=None, lora3s=1.0,
                                  lora4=None, lora4s=1.0, lora5=None, lora5s=1.0, progress=gr.Progress(track_tqdm=True)):
    if not civitai_key and os.environ.get("CIVITAI_API_KEY"): civitai_key = os.environ.get("CIVITAI_API_KEY") # default Civitai API key
    if not hf_token and os.environ.get("HF_TOKEN"): hf_token = os.environ.get("HF_TOKEN") # default HF write token
    if not hf_user and os.environ.get("HF_USER"): hf_user = os.environ.get("HF_USER") # default username
    if not hf_user:
        print(f"Invalid user name: {hf_user}")
        progress(1, desc=f"Invalid user name: {hf_user}")
        return gr.update(value=repo_urls, choices=repo_urls), gr.update(visible=True)
    lora_dict = {lora1: lora1s, lora2: lora2s, lora3: lora3s, lora4: lora4s, lora5: lora5s}
    new_path = convert_url_to_diffusers_sdxl(dl_url, civitai_key, hf_token, is_upload_sf, half, vae, scheduler, lora_dict, False, clip)
    if not new_path: return ""
    new_repo_id = f"{hf_user}/{Path(new_path).stem}"
    if hf_repo != "": new_repo_id = f"{hf_user}/{hf_repo}"
    if not is_repo_name(new_repo_id):
        print(f"Invalid repo name: {new_repo_id}")
        progress(1, desc=f"Invalid repo name: {new_repo_id}")
        return gr.update(value=repo_urls, choices=repo_urls), gr.update(visible=True)
    if not is_overwrite and is_repo_exists(new_repo_id, hf_token):
        print(f"Repo already exists: {new_repo_id}")
        progress(1, desc=f"Repo already exists: {new_repo_id}")
        return gr.update(value=repo_urls, choices=repo_urls), gr.update(visible=True)
    repo_url = create_diffusers_repo(new_repo_id, new_path, is_private, hf_token)
    shutil.rmtree(new_path)
    if not repo_urls: repo_urls = []
    repo_urls.append(repo_url)
    md = "### Your new repo:\n"
    for u in repo_urls:
        md += f"[{str(u).split('/')[-2]}/{str(u).split('/')[-1]}]({str(u)})<br>"
    return gr.update(value=repo_urls, choices=repo_urls), gr.update(value=md)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--url", default=None, type=str, required=True, help="URL of the model to convert.")
    parser.add_argument("--half", default=True, help="Save weights in half precision.")
    parser.add_argument("--scheduler", default="Euler a", type=str, choices=list(SCHEDULER_CONFIG_MAP.keys()), required=False, help="Scheduler name to use.")
    parser.add_argument("--vae", default=None, type=str, required=False, help="URL of the VAE to use.")
    parser.add_argument("--civitai_key", default=None, type=str, required=False, help="Civitai API Key (If you want to download file from Civitai).")
    parser.add_argument("--lora1", default=None, type=str, required=False, help="URL of the LoRA to use.")
    parser.add_argument("--lora1s", default=1.0, type=float, required=False, help="LoRA weight scale of --lora1.")
    parser.add_argument("--lora2", default=None, type=str, required=False, help="URL of the LoRA to use.")
    parser.add_argument("--lora2s", default=1.0, type=float, required=False, help="LoRA weight scale of --lora2.")
    parser.add_argument("--lora3", default=None, type=str, required=False, help="URL of the LoRA to use.")
    parser.add_argument("--lora3s", default=1.0, type=float, required=False, help="LoRA weight scale of --lora3.")
    parser.add_argument("--lora4", default=None, type=str, required=False, help="URL of the LoRA to use.")
    parser.add_argument("--lora4s", default=1.0, type=float, required=False, help="LoRA weight scale of --lora4.")
    parser.add_argument("--lora5", default=None, type=str, required=False, help="URL of the LoRA to use.")
    parser.add_argument("--lora5s", default=1.0, type=float, required=False, help="LoRA weight scale of --lora5.")
    parser.add_argument("--loras", default=None, type=str, required=False, help="Folder of the LoRA to use.")

    args = parser.parse_args()
    assert args.url is not None, "Must provide a URL!"

    lora_dict = {args.lora1: args.lora1s, args.lora2: args.lora2s, args.lora3: args.lora3s, args.lora4: args.lora4s, args.lora5: args.lora5s}

    if args.loras and Path(args.loras).exists():
        for p in Path(args.loras).glob('**/*.safetensors'):
            lora_dict[str(p)] = 1.0

    convert_url_to_diffusers_sdxl(args.url, args.civitai_key, args.half, args.vae, args.scheduler, lora_dict, True)