import argparse
from pathlib import Path
import os
import torch
from diffusers import StableDiffusionPipeline, AutoencoderKL
import gradio as gr
# 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):
import re
return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)
def download_thing(directory, url, civitai_api_key="", 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/")
os.system(f"aria2c --console-log-level=error --summary-interval=10 -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', '.ckpt', '.bin', '.pt', '.pth')
for file in Path(dir_path).glob("*"):
if file.suffix in valid_extensions:
file_path = str(Path(f"{dir_path}/{file.name}"))
model_list.append(file_path)
return model_list
def get_download_file(temp_dir, url, civitai_key, 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)
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"]
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=""):
if not lora_dict or not isinstance(lora_dict, dict): return
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)
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.set_adapters(a_list, adapter_weights=w_list)
pipe.fuse_lora(adapter_names=a_list, lora_scale=1.0)
pipe.unload_lora_weights()
def convert_url_to_diffusers_sd(url, civitai_key="", is_upload_sf=False, half=True, vae=None, scheduler="Euler", lora_dict={},
model_type="v1", sample_size=768, ema="ema", is_local=True, progress=gr.Progress(track_tqdm=True)):
progress(0, desc="Start converting...")
temp_dir = "."
new_file = get_download_file(temp_dir, url, civitai_key)
if not new_file:
print(f"Not found: {url}")
return ""
new_repo_name = Path(new_file).stem.replace(" ", "_").replace(",", "_").replace(".", "_") #
extract_ema = True if ema == "ema" else False
if model_type == "v1": #
config_url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"
elif model_type == "v2":
if sample_size == 512:
config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference.yaml"
else:
config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference-v.yaml"
pipe = None
if is_repo_name(url):
if half:
pipe = StableDiffusionPipeline.from_pretrained(new_file, extract_ema=extract_ema, requires_safety_checker=False, use_safetensors=True, torch_dtype=torch.float16)
else:
pipe = StableDiffusionPipeline.from_pretrained(new_file, extract_ema=extract_ema, requires_safety_checker=False, use_safetensors=True)
else:
if half:
pipe = StableDiffusionPipeline.from_single_file(new_file, extract_ema=extract_ema, requires_safety_checker=False, use_safetensors=True, torch_dtype=torch.float16)
else:
pipe = StableDiffusionPipeline.from_single_file(new_file, extract_ema=extract_ema, requires_safety_checker=False, use_safetensors=True)
new_vae_file = ""
if vae:
if is_repo_name(vae):
if half:
pipe.vae = AutoencoderKL.from_pretrained(vae, torch_dtype=torch.float16)
else:
pipe.vae = AutoencoderKL.from_pretrained(vae)
else:
new_vae_file = get_download_file(temp_dir, vae, civitai_key)
if new_vae_file and half:
pipe.vae = AutoencoderKL.from_single_file(new_vae_file, torch_dtype=torch.float16)
elif new_vae_file:
pipe.vae = AutoencoderKL.from_single_file(new_vae_file)
fuse_loras(pipe, lora_dict, temp_dir, civitai_key)
sconf = get_scheduler_config(scheduler)
pipe.scheduler = sconf[0].from_config(pipe.scheduler.config, **sconf[1])
if half:
pipe.save_pretrained(new_repo_name, safe_serialization=True, use_safetensors=True)
else:
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_repo_name(new_file) and is_upload_sf:
import shutil
shutil.move(str(Path(new_file).resolve()), str(Path(new_repo_name, Path(new_file).name).resolve()))
elif not is_local: os.remove(new_file)
progress(1, desc="Converted.")
return new_repo_name
def is_repo_exists(repo_id):
from huggingface_hub import HfApi
api = HfApi()
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}. ")
return True # for safe
def create_diffusers_repo(new_repo_id, diffusers_folder, is_private, progress=gr.Progress(track_tqdm=True)):
from huggingface_hub import HfApi
import os
hf_token = os.environ.get("HF_TOKEN")
api = HfApi()
try:
progress(0, desc="Start uploading...")
api.create_repo(repo_id=new_repo_id, token=hf_token, private=is_private)
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}. ")
print(e)
return ""
return url
def convert_url_to_diffusers_repo_sd(dl_url, hf_user, hf_repo, hf_token, civitai_key="", is_private=True, is_upload_sf=False, repo_urls=[], half=True, vae=None,
scheduler="Euler", 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,
model_type="v1", sample_size=768, ema="ema", progress=gr.Progress(track_tqdm=True)):
import shutil
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(value="")
if hf_token and not os.environ.get("HF_TOKEN"): os.environ['HF_TOKEN'] = hf_token
if not civitai_key and os.environ.get("CIVITAI_API_KEY"): civitai_key = os.environ.get("CIVITAI_API_KEY")
lora_dict = {lora1: lora1s, lora2: lora2s, lora3: lora3s, lora4: lora4s, lora5: lora5s}
new_path = convert_url_to_diffusers_sd(dl_url, civitai_key, is_upload_sf, half, vae, scheduler, lora_dict,
model_type, sample_size, ema, False)
if not new_path: return gr.update(value=repo_urls, choices=repo_urls), gr.update(value="")
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(value="")
if is_repo_exists(new_repo_id):
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(value="")
repo_url = create_diffusers_repo(new_repo_id, new_path, is_private)
shutil.rmtree(new_path)
if not repo_urls: repo_urls = []
repo_urls.append(repo_url)
md = "Your new repo:
"
for u in repo_urls:
md += f"[{str(u).split('/')[-2]}/{str(u).split('/')[-1]}]({str(u)})
"
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("--model_type", default="v1", type=str, choices=["v1", "v2"], required=False, help="Extract EMA or non-EMA?")
parser.add_argument("--sample_size", default=512, type=int, choices=[512, 768], required=False, help="Sample size (px)")
parser.add_argument("--ema", default="ema", type=str, choices=["ema", "non-ema"], required=False, help="Extract EMA or non-EMA?")
parser.add_argument("--scheduler", default="Euler", 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_sd(args.url, args.civitai_key, args.half, args.vae, args.scheduler, lora_dict,
args.model_type, args.sample_size, args.ema, True)