Add custom

app.py

@@ -3,7 +3,7 @@ import torch
 torch.jit.script = lambda f: f
 import timm
 import time
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard
 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
@@ -16,6 +16,8 @@ import random
 from urllib.parse import quote
 import gdown
 import os
+import re
+import requests
 
 import diffusers
 from diffusers.utils import load_image
@@ -155,6 +157,8 @@ button.addEventListener('click', function() {
     element.classList.add('selected');
 });
 '''
+lora_archive = "/data"
+
 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"
@@ -197,8 +201,8 @@ def center_crop_image_as_square(img):
     img_cropped = img.crop((left, top, right, bottom))
     return img_cropped
     
-def check_selected(selected_state):
-    if not selected_state:
+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):
@@ -224,6 +228,7 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 del lora_model
 @spaces.GPU
 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')
@@ -303,8 +308,10 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     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, progress=gr.Progress(track_tqdm=True)):
-    selected_state_index = selected_state.index
+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:
@@ -319,28 +326,35 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     print('Cropping and calculating face embeds took: ', elapsed_time, 'seconds')
 
     st = time.time()
-    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)
 
-    
+    if(custom_lora_path):
+        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("Selected State: ", selected_state_index)
-    print(sdxl_loras[selected_state_index]["repo"])
+    #print("Selected State: ", selected_state_index)
+    #print(sdxl_loras[selected_state_index]["repo"])
     if negative == "":
         negative = None
-
-    if not selected_state:
-        raise gr.Error("You must select a LoRA")
-    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("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()
@@ -368,6 +382,117 @@ def swap_gallery(order, sdxl_loras):
 def deselect():
   return gr.Gallery(selected_index=None)
 
+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
+
+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"] > 2):
+    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.")
+  elif(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):
+    try:
+        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)
+    except Exception as e:
+        print("Error: ", e)
+        return None, None, None, None
+
+def show_loading_widget():
+    return gr.update(visible=True)
+
+def load_custom_lora(link):
+    title, path, trigger_word, image = check_custom_model(link)
+    if(title):
+        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>
+                If the style isn't applied properly, modify advanced settings, specially <code>Face strength</code> and <code>Image strength</code>
+                </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}"
+    else:
+        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)
+    #except Exception as e:
+    #    gr.Info("Invalid custom LoRA")
+    #    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)
+
+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(
@@ -381,8 +506,9 @@ with gr.Blocks(css="custom.css") as demo:
         elem_id="title",
     )
     selected_state = gr.State()
+    custom_loaded_lora = gr.State()
     with gr.Row(elem_id="main_app"):
-        with gr.Column(scale=4):
+        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", )
@@ -394,14 +520,16 @@ with gr.Blocks(css="custom.css") as demo:
                 #    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="Style gallery",
+                    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="Enter a custom Hugging Face or CivitAI SDXL LoRA", interactive=False, placeholder="Coming soon...")
+                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")
@@ -431,6 +559,16 @@ with gr.Blocks(css="custom.css") as demo:
     #    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],
+        queue=False
+    )
+    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],
@@ -447,22 +585,22 @@ with gr.Blocks(css="custom.css") as demo:
     #)
     prompt.submit(
         fn=check_selected,
-        inputs=[selected_state],
+        inputs=[selected_state, custom_loaded_lora],
         queue=False,
         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],
+        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],
+        inputs=[selected_state, custom_loaded_lora],
         queue=False,
         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],
+        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)