Spaces:
Running
Running
import argparse | |
from pathlib import Path | |
import os | |
import torch | |
from diffusers import StableDiffusionXLPipeline, AutoencoderKL | |
# 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=""): | |
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') | |
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): | |
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 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, civitai_key="", lora_dict={}, temp_dir="."): | |
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) | |
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="", half=True, vae=None, scheduler="Euler a", lora_dict={}): | |
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(".", "_") # | |
pipe = None | |
if is_repo_name(url): | |
if half: | |
pipe = StableDiffusionXLPipeline.from_pretrained(new_file, use_safetensors=True, torch_dtype=torch.float16) | |
else: | |
pipe = StableDiffusionXLPipeline.from_pretrained(new_file, use_safetensors=True) | |
else: | |
if half: | |
pipe = StableDiffusionXLPipeline.from_single_file(new_file, use_safetensors=True, torch_dtype=torch.float16) | |
else: | |
pipe = StableDiffusionXLPipeline.from_single_file(new_file, 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) | |
pipe = 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 __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) | |
# Usage: python convert_url_to_diffusers_sdxl.py --url https://huggingface.co/bluepen5805/anima_pencil-XL/blob/main/anima_pencil-XL-v5.0.0.safetensors | |
# python convert_url_to_diffusers_sdxl.py --url https://huggingface.co/bluepen5805/anima_pencil-XL/blob/main/anima_pencil-XL-v5.0.0.safetensors --scheduler "Euler a" | |
# python convert_url_to_diffusers_sdxl.py --url https://huggingface.co/bluepen5805/anima_pencil-XL/blob/main/anima_pencil-XL-v5.0.0.safetensors --loras ./loras |