Spaces:
Running
Running
| 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, dtype="fp16", vae="", | |
| 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 dtype == "fp16": type_kwargs["torch_dtype"] = torch.float16 | |
| elif dtype == "fp32": type_kwargs["torch_dtype"] = torch.float32 | |
| elif dtype == "bf16": type_kwargs["torch_dtype"] = torch.bfloat16 | |
| elif dtype == "fp8": type_kwargs["torch_dtype"] = torch.float8_e4m3fn | |
| 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=[], dtype="fp16", 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, dtype, 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("--dtype", default="fp16", type=str, choices=["fp16", "fp32", "bf16", "fp8", "default"], help='Output data type. (Default: "fp16")') | |
| 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.dtype, args.vae, args.scheduler, lora_dict, True) | |