Update app.py

app.py

@@ -316,74 +316,71 @@ def randomize_loras(selected_indices, loras_state):
     random_prompt = random.choice(prompt_values)
     return selected_info_1, selected_info_2, selected_info_3, selected_info_4, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, lora_image_1, lora_image_2, lora_image_3, lora_image_4, random_prompt
 
-def add_custom_lora(custom_lora, selected_indices, current_loras, gallery):
-    if custom_lora:
-        try:
-            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
-            print(f"Loaded custom LoRA: {repo}")
-            existing_item_index = next((index for (index, item) in enumerate(current_loras) if item['repo'] == repo), None)
-            if existing_item_index is None:
-                if repo.endswith(".safetensors") and repo.startswith("http"):
-                    repo = download_file(repo)
-                new_item = {
-                    "image": image if image else "/home/user/app/custom.png",
-                    "title": title,
-                    "repo": repo,
-                    "weights": path,
-                    "trigger_word": trigger_word
-                }
-                print(f"New LoRA: {new_item}")
-                existing_item_index = len(current_loras)
-                current_loras.append(new_item)
-            
-            # Update gallery
-            gallery_items = [(item["image"], item["title"]) for item in current_loras]
-            # Update selected_indices if there's room
-            if len(selected_indices) < 4:
-                selected_indices.append(existing_item_index)
-            else:
-                gr.Warning("You can select up to 4 LoRAs, remove one to select a new one.")
-
-            # Update selected_info and images
-            selected_info_1 = "Select a Celebrity as LoRA 1"
-            selected_info_2 = "Select a LoRA 2"
-            selected_info_3 = "Select a LoRA 3"
-            selected_info_4 = "Select a LoRA 4"
-            lora_scale_1 = 1.15
-            lora_scale_2 = 1.15
-            lora_scale_3 = 0.65
-            lora_scale_4 = 0.65
-            lora_image_1 = None
-            lora_image_2 = None
-            lora_image_3 = None
-            lora_image_4 = None
-            if len(selected_indices) >= 1:
-                lora1 = current_loras[selected_indices[0]]
-                selected_info_1 = f"### LoRA 1 Selected: {lora1['title']} ✨"
-                lora_image_1 = lora1['image'] if lora1['image'] else None
+def add_custom_lora(custom_lora, selected_indices, current_loras, gallery, request: gr.Request = None):
+    if not custom_lora:
+        return current_loras, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+    
+    try:
+        # Retrieve user token if running in Spaces
+        user_token = request.headers.get("Authorization", "").replace("Bearer ", "") if request else None
+
+        # Check and load custom LoRA
+        title, repo, path, trigger_word, image = check_custom_model(custom_lora, token=user_token)
+        print(f"Loaded custom LoRA: {repo}")
         
-            if len(selected_indices) >= 2:
-                lora2 = current_loras[selected_indices[1]]
-                selected_info_2 = f"### LoRA 2 Selected: {lora2['title']} ✨"
-                lora_image_2 = lora2['image'] if lora2['image'] else None
+        # Check if the LoRA already exists in the current list
+        existing_item_index = next((index for (index, item) in enumerate(current_loras) if item['repo'] == repo), None)
         
-            if len(selected_indices) >= 3:
-                lora3 = current_loras[selected_indices[2]]
-                selected_info_3 = f"### LoRA 3 Selected: {lora3['title']} ✨"
-                lora_image_3 = lora3['image'] if lora3['image'] else None
+        if existing_item_index is None:
+            # Download if a direct .safetensors URL
+            if repo.endswith(".safetensors") and repo.startswith("http"):
+                repo = download_file(repo)
+            
+            # Add the new LoRA
+            new_item = {
+                "image": image or "/home/user/app/custom.png",
+                "title": title,
+                "repo": repo,
+                "weights": path,
+                "trigger_word": trigger_word,
+            }
+            print(f"New LoRA: {new_item}")
+            existing_item_index = len(current_loras)
+            current_loras.append(new_item)
         
-            if len(selected_indices) >= 4:
-                lora4 = current_loras[selected_indices[3]]
-                selected_info_4 = f"### LoRA 4 Selected: {lora4['title']} ✨"
-                lora_image_4 = lora4['image'] if lora4['image'] else None
-            print("Finished adding custom LoRA")
-            return (current_loras, gr.update(value=gallery_items), selected_info_1, selected_info_2, selected_info_3, selected_info_4, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, lora_image_1, lora_image_2, lora_image_3, lora_image_4)
-        except Exception as e:
-            print(e)
-            gr.Warning(str(e))
-            return current_loras, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
-    else:
-        return current_loras, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+        # Update gallery items
+        gallery_items = [(item["image"], item["title"]) for item in current_loras]
+
+        # Update selected indices
+        if len(selected_indices) < 4:
+            selected_indices.append(existing_item_index)
+        else:
+            raise gr.Error("You can select up to 4 LoRAs. Please remove one to add a new one.")
+
+        # Update selection info and images
+        selected_info = [f"Select a LoRA {i + 1}" for i in range(4)]
+        lora_images = [None] * 4
+        lora_scales = [1.15, 1.15, 0.65, 0.65]
+
+        for idx, sel_idx in enumerate(selected_indices[:4]):
+            lora = current_loras[sel_idx]
+            selected_info[idx] = f"### LoRA {idx + 1} Selected: {lora['title']} ✨"
+            lora_images[idx] = lora.get("image")
+
+        print("Finished adding custom LoRA")
+        return (
+            current_loras,
+            gr.update(value=gallery_items),
+            *selected_info,
+            selected_indices,
+            *lora_scales,
+            *lora_images,
+        )
+
+    except Exception as e:
+        print(e)
+        return (current_loras, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), selected_indices, gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(),gr.update(),
+        )
 
 def remove_custom_lora(selected_indices, current_loras, gallery):
     if current_loras:
@@ -519,10 +516,10 @@ def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scal
 
 run_lora.zerogpu = True
 
-def get_huggingface_safetensors(link):
+def get_huggingface_safetensors(link, token=None):
     split_link = link.split("/")
     if len(split_link) == 2:
-        model_card = ModelCard.load(link)
+        model_card = ModelCard.load(link, use_auth_token=token)
         base_model = model_card.data.get("base_model")
         print(f"Base model: {base_model}")
         if base_model not in ["black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"]:
@@ -530,7 +527,7 @@ def get_huggingface_safetensors(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()
+        fs = HfFileSystem(token=token)
         safetensors_name = None
         try:
             list_of_files = fs.ls(link, detail=False)
@@ -549,24 +546,22 @@ def get_huggingface_safetensors(link):
     else:
         raise gr.Error("Invalid Hugging Face repository link")
 
-def check_custom_model(link):
+def check_custom_model(link, token=None):
     if link.endswith(".safetensors"):
-        # Treat as direct link to the LoRA weights
         title = os.path.basename(link)
         repo = link
-        path = None  # No specific weight name
+        path = None
         trigger_word = ""
         image_url = None
         return title, repo, path, trigger_word, image_url
     elif link.startswith("https://"):
         if "huggingface.co" in link:
             link_split = link.split("huggingface.co/")
-            return get_huggingface_safetensors(link_split[1])
+            return get_huggingface_safetensors(link_split[1], token=token)
         else:
             raise Exception("Unsupported URL")
     else:
-        # Assume it's a Hugging Face model path
-        return get_huggingface_safetensors(link)
+        return get_huggingface_safetensors(link, token=token)
 
 def update_history(new_image, history):
     """Updates the history gallery with the new image."""