Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| import torch | |
| import spaces | |
| from diffusers import DiffusionPipeline | |
| from pathlib import Path | |
| import gc | |
| import subprocess | |
| subprocess.run('pip cache purge', shell=True) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| torch.set_grad_enabled(False) | |
| models = [ | |
| "camenduru/FLUX.1-dev-diffusers", | |
| "black-forest-labs/FLUX.1-schnell", | |
| "sayakpaul/FLUX.1-merged", | |
| "John6666/blue-pencil-flux1-v001-fp8-flux", | |
| "John6666/copycat-flux-test-fp8-v11-fp8-flux", | |
| "John6666/nepotism-fuxdevschnell-v3aio-flux", | |
| "John6666/niji-style-flux-devfp8-fp8-flux", | |
| "John6666/fluxunchained-artfulnsfw-fut516xfp8e4m3fnv11-fp8-flux", | |
| "John6666/fastflux-unchained-t5f16-fp8-flux", | |
| "John6666/the-araminta-flux1a1-fp8-flux", | |
| "John6666/acorn-is-spinning-flux-v11-fp8-flux", | |
| "John6666/fluxescore-dev-v10fp16-fp8-flux", | |
| # "", | |
| ] | |
| num_loras = 3 | |
| def is_repo_name(s): | |
| import re | |
| return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s) | |
| 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}. ") | |
| print(e) | |
| return True # for safe | |
| def clear_cache(): | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| def get_repo_safetensors(repo_id: str): | |
| from huggingface_hub import HfApi | |
| api = HfApi() | |
| try: | |
| if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(value="", choices=[]) | |
| files = api.list_repo_files(repo_id=repo_id) | |
| except Exception as e: | |
| print(f"Error: Failed to get {repo_id}'s info. ") | |
| print(e) | |
| return gr.update(choices=[]) | |
| files = [f for f in files if f.endswith(".safetensors")] | |
| if len(files) == 0: return gr.update(value="", choices=[]) | |
| else: return gr.update(value=files[0], choices=files) | |
| # Initialize the base model | |
| base_model = models[0] | |
| pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16) | |
| last_model = models[0] | |
| def change_base_model(repo_id: str, progress=gr.Progress(track_tqdm=True)): | |
| global pipe | |
| global last_model | |
| try: | |
| if repo_id == last_model or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return | |
| progress(0, desc=f"Loading model: {repo_id}") | |
| clear_cache() | |
| pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16) | |
| last_model = repo_id | |
| progress(1, desc=f"Model loaded: {repo_id}") | |
| except Exception as e: | |
| print(e) | |
| return gr.update(visible=True) | |
| def compose_lora_json(lorajson: list[dict], i: int, name: str, scale: float, filename: str, trigger: str): | |
| lorajson[i]["name"] = str(name) if name != "None" else "" | |
| lorajson[i]["scale"] = float(scale) | |
| lorajson[i]["filename"] = str(filename) | |
| lorajson[i]["trigger"] = str(trigger) | |
| return lorajson | |
| def is_valid_lora(lorajson: list[dict]): | |
| valid = False | |
| for d in lorajson: | |
| if "name" in d.keys() and d["name"] and d["name"] != "None": valid = True | |
| return valid | |
| def get_trigger_word(lorajson: list[dict]): | |
| trigger = "" | |
| for d in lorajson: | |
| if "name" in d.keys() and d["name"] and d["name"] != "None" and d["trigger"]: | |
| trigger += ", " + d["trigger"] | |
| return trigger | |
| # https://huggingface.co/docs/diffusers/v0.23.1/en/api/loaders#diffusers.loaders.LoraLoaderMixin.fuse_lora | |
| # https://github.com/huggingface/diffusers/issues/4919 | |
| def fuse_loras(pipe, lorajson: list[dict]): | |
| if not lorajson or not isinstance(lorajson, list): return | |
| a_list = [] | |
| w_list = [] | |
| for d in lorajson: | |
| if not d or not isinstance(d, dict) or not d["name"] or d["name"] == "None": continue | |
| k = d["name"] | |
| if is_repo_name(k) and is_repo_exists(k): | |
| a_name = Path(k).stem | |
| pipe.load_lora_weights(k, weight_name=d["filename"], adapter_name = a_name) | |
| elif not Path(k).exists(): | |
| print(f"LoRA not found: {k}") | |
| continue | |
| else: | |
| w_name = Path(k).name | |
| a_name = Path(k).stem | |
| pipe.load_lora_weights(k, weight_name = w_name, adapter_name = a_name) | |
| a_list.append(a_name) | |
| w_list.append(d["scale"]) | |
| 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() | |
| change_base_model.zerogpu = True | |
| fuse_loras.zerogpu = True | |
| def description_ui(): | |
| gr.Markdown( | |
| """ | |
| - Mod of [multimodalart/flux-lora-the-explorer](https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer), | |
| [gokaygokay/FLUX-Prompt-Generator](https://huggingface.co/spaces/gokaygokay/FLUX-Prompt-Generator). | |
| """ | |
| ) |