Update app.py

app.py

@@ -1,166 +1,606 @@
-import os
-import re
+import gradio as gr
+import torch
+import spaces
+torch.jit.script = lambda f: f
+import timm
 import time
-import json
+
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
+from safetensors.torch import load_file
+from share_btn import community_icon_html, loading_icon_html, share_js
+from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
+
+import lora
 import copy
+import json
+import gc
 import random
-import requests
-import torch
-import cv2
-import numpy as np
-import gradio as gr
-import spaces
-from PIL import Image
 from urllib.parse import quote
+import gdown
+import os
+import re
+import requests
 
-# Disable Torch JIT compilation for compatibility
-torch.jit.script = lambda f: f  
-
-# Model & Utilities
-import timm
 import diffusers
 from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
 from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, UNet2DConditionModel
-from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+
 from insightface.app import FaceAnalysis
+from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
 from controlnet_aux import ZoeDetector
+
 from compel import Compel, ReturnedEmbeddingsType
+
 from gradio_imageslider import ImageSlider
 
-# Custom imports
-try:
-    from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
-    from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
-except ImportError as e:
-    print(f"Import Error: {e}. Check if modules exist or paths are correct.")
-    exit()
 
-# Device setup
-device = "cuda" if torch.cuda.is_available() else "cpu"
+#from gradio_imageslider import ImageSlider
 
-# Load LoRA configuration
 with open("sdxl_loras.json", "r") as file:
-    sdxl_loras_raw = json.load(file)
+    data = json.load(file)
+    sdxl_loras_raw = [
+        {
+            "image": item["image"],
+            "title": item["title"],
+            "repo": item["repo"],
+            "trigger_word": item["trigger_word"],
+            "weights": item["weights"],
+            "is_compatible": item["is_compatible"],
+            "is_pivotal": item.get("is_pivotal", False),
+            "text_embedding_weights": item.get("text_embedding_weights", None),
+            "likes": item.get("likes", 0),
+            "downloads": item.get("downloads", 0),
+            "is_nc": item.get("is_nc", False),
+            "new": item.get("new", False),
+        }
+        for item in data
+    ]
 
 with open("defaults_data.json", "r") as file:
     lora_defaults = json.load(file)
+    
 
-# Download required models
-CHECKPOINT_DIR = "/data/checkpoints"
-hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/config.json", local_dir=CHECKPOINT_DIR)
-hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/diffusion_pytorch_model.safetensors", local_dir=CHECKPOINT_DIR)
-hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir=CHECKPOINT_DIR)
-hf_hub_download(repo_id="latent-consistency/lcm-lora-sdxl", filename="pytorch_lora_weights.safetensors", local_dir=CHECKPOINT_DIR)
+device = "cuda" 
 
-# Download Antelopev2 Face Recognition model
-antelope_download = snapshot_download(repo_id="DIAMONIK7777/antelopev2", local_dir="/data/models/antelopev2")
-print("Antelopev2 Download Path:", antelope_download)
+state_dicts = {}
+
+for item in sdxl_loras_raw:
+    saved_name = hf_hub_download(item["repo"], item["weights"])
+    
+    if not saved_name.endswith('.safetensors'):
+        state_dict = torch.load(saved_name)
+    else:
+        state_dict = load_file(saved_name)
+    
+    state_dicts[item["repo"]] = {
+        "saved_name": saved_name,
+        "state_dict": state_dict
+    }
+
+sdxl_loras_raw = [item for item in sdxl_loras_raw if item.get("new") != True]
+    
+# download models
+hf_hub_download(
+    repo_id="InstantX/InstantID",
+    filename="ControlNetModel/config.json",
+    local_dir="/data/checkpoints",
+)
+hf_hub_download(
+    repo_id="InstantX/InstantID",
+    filename="ControlNetModel/diffusion_pytorch_model.safetensors",
+    local_dir="/data/checkpoints",
+)
+hf_hub_download(
+    repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="/data/checkpoints"
+)
+hf_hub_download(
+    repo_id="latent-consistency/lcm-lora-sdxl",
+    filename="pytorch_lora_weights.safetensors",
+    local_dir="/data/checkpoints",
+)
+# download antelopev2
+#if not os.path.exists("/data/antelopev2.zip"):
+#    gdown.download(url="https://drive.google.com/file/d/18wEUfMNohBJ4K3Ly5wpTejPfDzp-8fI8/view?usp=sharing", output="/data/", quiet=False, fuzzy=True)
+#    os.system("unzip /data/antelopev2.zip -d /data/models/")
 
-# Initialize FaceAnalysis
-app = FaceAnalysis(name="antelopev2", root="/data", providers=["CPUExecutionProvider"])
+antelope_download = snapshot_download(repo_id="DIAMONIK7777/antelopev2", local_dir="/data/models/antelopev2")
+print(antelope_download)
+app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640))
 
-# Load identity & depth models
-face_adapter = os.path.join(CHECKPOINT_DIR, "ip-adapter.bin")
-controlnet_path = os.path.join(CHECKPOINT_DIR, "ControlNetModel")
+# prepare models under ./checkpoints
+face_adapter = f'/data/checkpoints/ip-adapter.bin'
+controlnet_path = f'/data/checkpoints/ControlNetModel'
 
+# load IdentityNet
+st = time.time()
 identitynet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
-zoedepthnet = ControlNetModel.from_pretrained("diffusers/controlnet-zoe-depth-sdxl-1.0", torch_dtype=torch.float16)
-
+zoedepthnet = ControlNetModel.from_pretrained("diffusers/controlnet-zoe-depth-sdxl-1.0",torch_dtype=torch.float16)
+et = time.time()
+elapsed_time = et - st
+print('Loading ControlNet took: ', elapsed_time, 'seconds')
+st = time.time()
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+et = time.time()
+elapsed_time = et - st
+print('Loading VAE took: ', elapsed_time, 'seconds')
+st = time.time()
 
-# Load main pipeline
-pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
-    "frankjoshua/albedobaseXL_v21",
-    vae=vae,
-    controlnet=[identitynet, zoedepthnet],
-    torch_dtype=torch.float16
-)
+#pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("stablediffusionapi/albedobase-xl-v21",
+pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("frankjoshua/albedobaseXL_v21",
+                                                                 vae=vae,
+                                                                 controlnet=[identitynet, zoedepthnet],
+                                                                 torch_dtype=torch.float16)
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
 pipe.load_ip_adapter_instantid(face_adapter)
 pipe.set_ip_adapter_scale(0.8)
+et = time.time()
+elapsed_time = et - st
+print('Loading pipeline took: ', elapsed_time, 'seconds')
+st = time.time()
+compel = Compel(tokenizer=[pipe.tokenizer, pipe.tokenizer_2] , text_encoder=[pipe.text_encoder, pipe.text_encoder_2], returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, requires_pooled=[False, True])
+et = time.time()
+elapsed_time = et - st
+print('Loading Compel took: ', elapsed_time, 'seconds')
 
-# Initialize Compel for text conditioning
-compel = Compel(
-    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-    requires_pooled=[False, True]
-)
-
-# Load ZoeDetector for depth estimation
+st = time.time()
 zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
+et = time.time()
+elapsed_time = et - st
+print('Loading Zoe took: ', elapsed_time, 'seconds')
 zoe.to(device)
 pipe.to(device)
 
-# LoRA Management
 last_lora = ""
 last_fused = False
+js = '''
+var button = document.getElementById('button');
+// Add a click event listener to the button
+button.addEventListener('click', function() {
+    element.classList.add('selected');
+});
+'''
+lora_archive = "/data"
 
-# --- Utility Functions ---
-def update_selection(selected_state, sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative):
-    index = selected_state.index
-    lora_repo = sdxl_loras[index]["repo"]
-    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
+def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative, is_new=False):
+    lora_repo = sdxl_loras[selected_state.index]["repo"]
+    new_placeholder = "Type a prompt to use your selected LoRA"
+    weight_name = sdxl_loras[selected_state.index]["weights"]
+    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
 
     for lora_list in lora_defaults:
-        if lora_list["model"] == lora_repo:
+        if lora_list["model"] == sdxl_loras[selected_state.index]["repo"]:
             face_strength = lora_list.get("face_strength", 0.85)
             image_strength = lora_list.get("image_strength", 0.15)
             weight = lora_list.get("weight", 0.9)
             depth_control_scale = lora_list.get("depth_control_scale", 0.8)
             negative = lora_list.get("negative", "")
-
+    
+    if(is_new):
+        if(selected_state.index == 0):
+            selected_state.index = -9999
+        else:
+            selected_state.index *= -1
+    
     return (
-        updated_text, gr.update(placeholder="Type a prompt"), face_strength, 
-        image_strength, weight, depth_control_scale, negative, selected_state
+        updated_text,
+        gr.update(placeholder=new_placeholder),
+        face_strength,
+        image_strength,
+        weight,
+        depth_control_scale,
+        negative,
+        selected_state
     )
 
-def center_crop_image(img):
+def center_crop_image_as_square(img):
     square_size = min(img.size)
-    left = (img.width - square_size) // 2
-    top = (img.height - square_size) // 2
-    return img.crop((left, top, left + square_size, top + square_size))
-
-def process_face(image):
-    face_info = app.get(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR))
-    face_info = sorted(face_info, key=lambda x: (x['bbox'][2]-x['bbox'][0]) * (x['bbox'][3]-x['bbox'][1]))[-1]
-    face_emb = face_info['embedding']
-    face_kps = draw_kps(image, face_info['kps'])
-    return face_emb, face_kps
-
-def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, lora_scale):
-    global last_fused, last_lora
-    if last_lora != repo_name and last_fused:
-        pipe.unfuse_lora()
-        pipe.unload_lora_weights()
-    pipe.load_lora_weights(repo_name)
-    pipe.fuse_lora(lora_scale)
-    last_lora, last_fused = repo_name, True
+    
+    left = (img.width - square_size) / 2
+    top = (img.height - square_size) / 2
+    right = (img.width + square_size) / 2
+    bottom = (img.height + square_size) / 2
+    
+    img_cropped = img.crop((left, top, right, bottom))
+    return img_cropped
+    
+def check_selected(selected_state, custom_lora):
+    if not selected_state and not custom_lora:
+        raise gr.Error("You must select a style")
+
+def merge_incompatible_lora(full_path_lora, lora_scale):
+    for weights_file in [full_path_lora]:
+                if ";" in weights_file:
+                    weights_file, multiplier = weights_file.split(";")
+                    multiplier = float(multiplier)
+                else:
+                    multiplier = lora_scale
 
+                lora_model, weights_sd = lora.create_network_from_weights(
+                    multiplier,
+                    full_path_lora,
+                    pipe.vae,
+                    pipe.text_encoder,
+                    pipe.unet,
+                    for_inference=True,
+                )
+                lora_model.merge_to(
+                    pipe.text_encoder, pipe.unet, weights_sd, torch.float16, "cuda"
+                )
+                del weights_sd
+                del lora_model
+
+@spaces.GPU(duration=80)
+def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index, st):
+    print(loaded_state_dict)
+    et = time.time()
+    elapsed_time = et - st
+    print('Getting into the decorated function took: ', elapsed_time, 'seconds')
+    global last_fused, last_lora
+    print("Last LoRA: ", last_lora)
+    print("Current LoRA: ", repo_name)
+    print("Last fused: ", last_fused)
+    #prepare face zoe
+    st = time.time()
+    with torch.no_grad():
+        image_zoe = zoe(face_image)
+    width, height = face_kps.size
+    images = [face_kps, image_zoe.resize((height, width))]
+    et = time.time()
+    elapsed_time = et - st
+    print('Zoe Depth calculations took: ', elapsed_time, 'seconds')
+    if last_lora != repo_name:
+        if(last_fused):
+            st = time.time()
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+            pipe.unload_textual_inversion()
+            et = time.time()
+            elapsed_time = et - st
+            print('Unfuse and unload LoRA took: ', elapsed_time, 'seconds')
+        st = time.time()
+        pipe.load_lora_weights(loaded_state_dict)
+        pipe.fuse_lora(lora_scale)
+        et = time.time()
+        elapsed_time = et - st
+        print('Fuse and load LoRA took: ', elapsed_time, 'seconds')
+        last_fused = True
+        is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
+        if(is_pivotal):
+            #Add the textual inversion embeddings from pivotal tuning models
+            text_embedding_name = sdxl_loras[selected_state_index]["text_embedding_weights"]
+            embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
+            state_dict_embedding = load_file(embedding_path)
+            pipe.load_textual_inversion(state_dict_embedding["clip_l" if "clip_l" in state_dict_embedding else "text_encoders_0"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
+            pipe.load_textual_inversion(state_dict_embedding["clip_g" if "clip_g" in state_dict_embedding else "text_encoders_1"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
+            
+    print("Processing prompt...")
+    st = time.time()
     conditioning, pooled = compel(prompt)
-    negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
-    
-    images = [face_kps, zoe(face_image).resize(face_kps.size)]
-    return pipe(
-        prompt_embeds=conditioning, pooled_prompt_embeds=pooled,
-        negative_prompt_embeds=negative_conditioning, negative_pooled_prompt_embeds=negative_pooled,
-        width=1024, height=1024, image_embeds=face_emb, image=face_image,
-        strength=1-image_strength, control_image=images, num_inference_steps=20,
-        guidance_scale=guidance_scale, controlnet_conditioning_scale=[face_strength, depth_control_scale]
+    if(negative):
+        negative_conditioning, negative_pooled = compel(negative)
+    else:
+        negative_conditioning, negative_pooled = None, None
+    et = time.time()
+    elapsed_time = et - st
+    print('Prompt processing took: ', elapsed_time, 'seconds')
+    print("Processing image...")
+    st = time.time()
+    image = pipe(
+        prompt_embeds=conditioning,
+        pooled_prompt_embeds=pooled,
+        negative_prompt_embeds=negative_conditioning,
+        negative_pooled_prompt_embeds=negative_pooled,
+        width=1024,
+        height=1024,
+        image_embeds=face_emb,
+        image=face_image,
+        strength=1-image_strength,
+        control_image=images,
+        num_inference_steps=20,
+        guidance_scale = guidance_scale,
+        controlnet_conditioning_scale=[face_strength, depth_control_scale],
     ).images[0]
+    et = time.time()
+    elapsed_time = et - st
+    print('Image processing took: ', elapsed_time, 'seconds')
+    last_lora = repo_name
+    return image
+    
+def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, sdxl_loras, custom_lora, progress=gr.Progress(track_tqdm=True)):
+    print("Custom LoRA: ", custom_lora)
+    custom_lora_path = custom_lora[0] if custom_lora else None
+    selected_state_index = selected_state.index if selected_state else -1
+    st = time.time()
+    face_image = center_crop_image_as_square(face_image)
+    try:
+        face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
+        face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*x['bbox'][3]-x['bbox'][1])[-1] # only use the maximum face
+        face_emb = face_info['embedding']
+        face_kps = draw_kps(face_image, face_info['kps'])
+    except: 
+        raise gr.Error("No face found in your image. Only face images work here. Try again")
+    et = time.time()
+    elapsed_time = et - st
+    print('Cropping and calculating face embeds took: ', elapsed_time, 'seconds')
+
+    st = time.time()
+
+    if(custom_lora_path and custom_lora[1]):
+        prompt = f"{prompt} {custom_lora[1]}"
+    else:
+        for lora_list in lora_defaults:
+            if lora_list["model"] == sdxl_loras[selected_state_index]["repo"]:
+                prompt_full = lora_list.get("prompt", None)
+                if(prompt_full):
+                    prompt = prompt_full.replace("<subject>", prompt)
+
+    print("Prompt:", prompt) 
+    if(prompt == ""):
+        prompt = "a person"
+    print(f"Executing prompt: {prompt}")
+    #print("Selected State: ", selected_state_index)
+    #print(sdxl_loras[selected_state_index]["repo"])
+    if negative == "":
+        negative = None
+    print("Custom Loaded LoRA: ", custom_lora_path)
+    if not selected_state and not custom_lora_path:
+        raise gr.Error("You must select a style")
+    elif custom_lora_path:
+        repo_name = custom_lora_path
+        full_path_lora = custom_lora_path
+    else:
+        repo_name = sdxl_loras[selected_state_index]["repo"]
+        weight_name = sdxl_loras[selected_state_index]["weights"]
+        full_path_lora = state_dicts[repo_name]["saved_name"]
+    print("Full path LoRA ", full_path_lora)
+    #loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
+    cross_attention_kwargs = None
+    et = time.time()
+    elapsed_time = et - st
+    print('Small content processing took: ', elapsed_time, 'seconds')
+
+    st = time.time()
+    image = generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, full_path_lora, lora_scale, sdxl_loras, selected_state_index, st)
+    return (face_image, image), gr.update(visible=True)
+
+run_lora.zerogpu = True
+
+def shuffle_gallery(sdxl_loras):
+    random.shuffle(sdxl_loras)
+    return [(item["image"], item["title"]) for item in sdxl_loras], sdxl_loras
+
+def classify_gallery(sdxl_loras):
+    sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get("likes", 0), reverse=True)
+    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
+
+def swap_gallery(order, sdxl_loras):
+    if(order == "random"):
+        return shuffle_gallery(sdxl_loras)
+    else:
+        return classify_gallery(sdxl_loras)
+        
+def deselect():
+  return gr.Gallery(selected_index=None)
 
-# --- UI Setup ---
-with gr.Blocks() as demo:
-    photo = gr.Image(label="Upload a picture", interactive=True, type="pil", height=300)
-    gallery = gr.Gallery(label="Pick a style", allow_preview=False, columns=4, height=550)
-    prompt = gr.Textbox(label="Prompt", placeholder="Enter prompt...")
-    button = gr.Button("Run")
-    result = ImageSlider(interactive=False, label="Generated Image")
+def get_huggingface_safetensors(link):
+  split_link = link.split("/")
+  if(len(split_link) == 2):
+            model_card = ModelCard.load(link)
+            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+            trigger_word = model_card.data.get("instance_prompt", "")
+            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+            fs = HfFileSystem()
+            try:
+                list_of_files = fs.ls(link, detail=False)
+                for file in list_of_files:
+                    if(file.endswith(".safetensors")):
+                        safetensors_name = file.replace("/", "_")
+                        if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
+                          fs.get_file(file, lpath=f"{lora_archive}/{safetensors_name}")
+                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
+                      image_elements = file.split("/")
+                      image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+            except:
+              gr.Warning("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+              raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            return split_link[1], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
 
-    button.click(fn=generate_image, inputs=[prompt, gr.State(), gr.State()], outputs=result)
+def get_civitai_safetensors(link):
+  link_split = link.split("civitai.com/")
+  pattern = re.compile(r'models\/(\d+)')
+  regex_match = pattern.search(link_split[1])
+  if(regex_match):
+    civitai_model_id = regex_match.group(1)
+  else:
+    gr.Warning("No CivitAI model id found in your URL")
+    raise Exception("No CivitAI model id found in your URL")
+  model_request_url = f"https://civitai.com/api/v1/models/{civitai_model_id}?token={os.getenv('CIVITAI_TOKEN')}"
+  x = requests.get(model_request_url)
+  if(x.status_code != 200):
+      raise Exception("Invalid CivitAI URL")
+  model_data = x.json()
+  #if(model_data["nsfw"] == True or model_data["nsfwLevel"] > 20):
+   # gr.Warning("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
+   # raise Exception("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
+  if(model_data["type"] != "LORA"):
+    gr.Warning("The model isn't tagged at CivitAI as a LoRA")
+    raise Exception("The model isn't tagged at CivitAI as a LoRA")
+  model_link_download = None
+  image_url = None
+  trigger_word = ""
+  for model in model_data["modelVersions"]:
+    if(model["baseModel"] == "SDXL 1.0"):
+      model_link_download = f"{model['downloadUrl']}/?token={os.getenv('CIVITAI_TOKEN')}"
+      safetensors_name = model["files"][0]["name"]
+      if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
+        safetensors_file_request = requests.get(model_link_download)
+        if(safetensors_file_request.status_code != 200):
+            raise Exception("Invalid CivitAI download link")
+        with open(f"{lora_archive}/{safetensors_name}", 'wb') as file:
+            file.write(safetensors_file_request.content)
+      trigger_word = model.get("trainedWords", [""])[0]
+      for image in model["images"]:
+        if(image["nsfwLevel"] == 1):
+          image_url = image["url"]
+          break
+      break
+  if(not model_link_download):
+    gr.Warning("We couldn't find a SDXL LoRA on the model you've sent")
+    raise Exception("We couldn't find a SDXL LoRA on the model you've sent")
+  return model_data["name"], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
+  
+def check_custom_model(link):
+    if(link.startswith("https://")):
+        if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
+            link_split = link.split("huggingface.co/")
+            return get_huggingface_safetensors(link_split[1])
+        elif(link.startswith("https://civitai.com") or link.startswith("https://www.civitai.com")):
+            return get_civitai_safetensors(link)
+    else: 
+        return get_huggingface_safetensors(link)
+
+def show_loading_widget():
+    return gr.update(visible=True)
+
+def load_custom_lora(link):
+    if(link):
+        try:
+            title, path, trigger_word, image = check_custom_model(link)
+            card = f'''
+            <div class="custom_lora_card">
+              <span>Loaded custom LoRA:</span>
+              <div class="card_internal">
+                <img src="{image}" />
+                <div>
+                    <h3>{title}</h3>
+                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
+                </div>
+              </div>
+            </div>
+            '''
+            return gr.update(visible=True), card, gr.update(visible=True), [path, trigger_word], gr.Gallery(selected_index=None), f"Custom: {path}"
+        except Exception as e:
+            gr.Warning("Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content")
+            return gr.update(visible=True), "Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), "", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
+
+def remove_custom_lora():
+    return "", gr.update(visible=False), gr.update(visible=False), None
+with gr.Blocks(css="custom.css") as demo:
+    gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
+    title = gr.HTML(
+        """<h1><img src="https://i.imgur.com/DVoGw04.png">
+<span>Face to All<br><small style="
+    font-size: 13px;
+    display: block;
+    font-weight: normal;
+    opacity: 0.75;
+">🧨 diffusers InstantID + ControlNet<br> inspired by fofr's <a href="https://github.com/fofr/cog-face-to-many" target="_blank">face-to-many</a></small></span></h1>""",
+        elem_id="title",
+    )
+    selected_state = gr.State()
+    custom_loaded_lora = gr.State()
+    with gr.Row(elem_id="main_app"):
+        with gr.Column(scale=4, elem_id="box_column"):
+            with gr.Group(elem_id="gallery_box"):
+                photo = gr.Image(label="Upload a picture of yourself", interactive=True, type="pil", height=300)
+                selected_loras = gr.Gallery(label="Selected LoRAs", height=80, show_share_button=False, visible=False, elem_id="gallery_selected", )
+                #order_gallery = gr.Radio(choices=["random", "likes"], value="random", label="Order by", elem_id="order_radio")
+                #new_gallery = gr.Gallery(
+                #    label="New LoRAs",
+                #    elem_id="gallery_new",
+                #    columns=3,
+                #    value=[(item["image"], item["title"]) for item in sdxl_loras_raw_new], allow_preview=False, show_share_button=False)
+                gallery = gr.Gallery(
+                    #value=[(item["image"], item["title"]) for item in sdxl_loras],
+                    label="Pick a style from the gallery",
+                    allow_preview=False,
+                    columns=4,
+                    elem_id="gallery",
+                    show_share_button=False,
+                    height=550
+                )
+                custom_model = gr.Textbox(label="or enter a custom Hugging Face or CivitAI SDXL LoRA", placeholder="Paste Hugging Face or CivitAI model path...")
+                custom_model_card = gr.HTML(visible=False)
+                custom_model_button = gr.Button("Remove custom LoRA", visible=False)
+        with gr.Column(scale=5):
+            with gr.Row():
+                prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, info="Describe your subject (optional)", value="a person", elem_id="prompt")
+                button = gr.Button("Run", elem_id="run_button")
+            result = ImageSlider(
+                interactive=False, label="Generated Image", elem_id="result-image", position=0.1
+            )
+            with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
+                community_icon = gr.HTML(community_icon_html)
+                loading_icon = gr.HTML(loading_icon_html)
+                share_button = gr.Button("Share to community", elem_id="share-btn")
+            with gr.Accordion("Advanced options", open=False):
+                negative = gr.Textbox(label="Negative Prompt")
+                weight = gr.Slider(0, 10, value=0.9, step=0.1, label="LoRA weight")
+                face_strength = gr.Slider(0, 2, value=0.85, step=0.01, label="Face strength", info="Higher values increase the face likeness but reduce the creative liberty of the models")
+                image_strength = gr.Slider(0, 1, value=0.15, step=0.01, label="Image strength", info="Higher values increase the similarity with the structure/colors of the original photo")
+                guidance_scale = gr.Slider(0, 50, value=7, step=0.1, label="Guidance Scale")
+                depth_control_scale = gr.Slider(0, 1, value=0.8, step=0.01, label="Zoe Depth ControlNet strenght")
+            prompt_title = gr.Markdown(
+                value="### Click on a LoRA in the gallery to select it",
+                visible=True,
+                elem_id="selected_lora",
+            )
+    #order_gallery.change(
+    #    fn=swap_gallery,
+    #    inputs=[order_gallery, gr_sdxl_loras],
+    #    outputs=[gallery, gr_sdxl_loras],
+    #    queue=False
+    #)
+    custom_model.input(
+        fn=load_custom_lora,
+        inputs=[custom_model],
+        outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery, prompt_title],
+    )
+    custom_model_button.click(
+        fn=remove_custom_lora,
+        outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora]
+    )
+    gallery.select(
+        fn=update_selection,
+        inputs=[gr_sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative],
+        outputs=[prompt_title, prompt, face_strength, image_strength, weight, depth_control_scale, negative, selected_state],
+        show_progress=False
+    )
+    #new_gallery.select(
+    #    fn=update_selection,
+    #    inputs=[gr_sdxl_loras_new, gr.State(True)],
+    #    outputs=[prompt_title, prompt, prompt, selected_state, gallery],
+    #    queue=False,
+    #    show_progress=False
+    #)
+    prompt.submit(
+        fn=check_selected,
+        inputs=[selected_state, custom_loaded_lora],
+        show_progress=False
+    ).success(
+        fn=run_lora,
+        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora],
+        outputs=[result, share_group],
+    )
+    button.click(
+        fn=check_selected,
+        inputs=[selected_state, custom_loaded_lora],
+        show_progress=False
+    ).success(
+        fn=run_lora,
+        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora],
+        outputs=[result, share_group],
+    )
+    share_button.click(None, [], [], js=share_js)
+    demo.load(fn=classify_gallery, inputs=[gr_sdxl_loras], outputs=[gallery, gr_sdxl_loras], js=js)
 
-demo.queue()
-demo.launch(share=True)
+demo.queue(default_concurrency_limit=None, api_open=True)
+demo.launch(share=True)