attempt custom

app.py

@@ -303,7 +303,7 @@ 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)):
+def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, sdxl_loras, custom_lora_path, progress=gr.Progress(track_tqdm=True)):
     selected_state_index = selected_state.index
     st = time.time()
     face_image = center_crop_image_as_square(face_image)
@@ -334,13 +334,17 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     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"]
+    if not selected_state and not custom_loaded_lora:
+        raise gr.Error("You must select a style")
+    elif custom_loaded_lora:
+        repo_name = "Loaded custom LoRA"
+        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"]
+    
     #loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
     cross_attention_kwargs = None
     et = time.time()
@@ -368,6 +372,92 @@ def swap_gallery(order, sdxl_loras):
 def deselect():
   return gr.Gallery(selected_index=None)
 
+lora_archive = "/data"
+
+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)
+            
+            print(image_path)
+            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_url = f"https://huggingface.co/{link}/resolve/main/{image_url}"
+            except:
+              raise gr.Error("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}", 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:
+    raise gr.Error("No CivitAI model id found in your URL")
+  model_request_url = f"https://civitai.com/api/v1/models/{civitai_model_id}?token=dc9537afe3dcbbfb4696dd79ee8c23d7"
+  x = requests.get(model_request_url)
+  model_data = x.json()
+  if(model_data["nsfw"] == True):
+    raise gr.Error("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
+  elif(model_data["type"] != "LORA"):
+    raise gr.Error("The model isn't tagged at CivitAI as a LoRA")
+  model_link_download = None
+  image_url = None
+  for model in model_data["modelVersions"]:
+    if(model["baseModel"] == "SDXL 1.0"):
+      model_link_download = f"{model['downloadUrl']}/?token=dc9537afe3dcbbfb4696dd79ee8c23d7"
+      safetensors_name = model["files"][0]["name"]
+      if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
+        safetensors_file_request = requests.get(model_link_download)
+        with open(f"{lora_archive}/{safetensors_name}", 'wb') as file:
+            file.write(safetensors_file_request.content)
+      for image in model["images"]:
+        if(image["nsfwLevel"] == 1):
+          image_url = image["url"]
+          break
+      break
+  return model_data["name"], f"{lora_archive}/{safetensors_name}", image_url
+  if(not model_link_download):
+    raise gr.Error("We couldn't find a SDXL LoRA on the model you've sent")
+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):
+    title, path, image = check_custom_model(link)
+    card = f'''
+    <div class="custom_lora_card">
+      <span>Loaded custom LoRA:</span>
+      <div class="card_internal">
+        <h3>{title}</h3>
+        <img src="{image}" />
+      </div>
+    </div>
+    '''
+    return gr.update(visible=True), card, gr.update(visible=True), path, gr.Gallery(selected_index=None)
+
+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,6 +471,7 @@ 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.Group(elem_id="gallery_box"):
@@ -402,6 +493,8 @@ with gr.Blocks(css="custom.css") as demo:
                     height=550
                 )
                 custom_model = gr.Textbox(label="Enter a custom Hugging Face or CivitAI SDXL LoRA", interactive=False, placeholder="Coming soon...")
+                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 +524,19 @@ with gr.Blocks(css="custom.css") as demo:
     #    outputs=[gallery, gr_sdxl_loras],
     #    queue=False
     #)
+    custom_model.change(
+        fn=show_loading_widget,
+        outputs=[custom_model_card]
+    ).then(
+        fn=load_custom_lora,
+        inputs=[custom_model],
+        outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery],
+        queue=False
+    )
+    custom_model_button.click(
+        fn=remove_custom_lora,
+        outputs=[custom_model, custom_loaded_lora, custom_model_card, custom_loaded_lora]
+    )
     gallery.select(
         fn=update_selection,
         inputs=[gr_sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative],
@@ -452,7 +558,7 @@ with gr.Blocks(css="custom.css") as demo:
         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(
@@ -462,7 +568,7 @@ with gr.Blocks(css="custom.css") as demo:
         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)