Add support for downloading loras from civitai

app.py

@@ -8,6 +8,7 @@ import logging
 import numpy as np
 import spaces
 from typing import Any, Dict, List, Optional, Union
+from civitai_utils import get_civitai_safetensors, LORA_CHECKPOINTS_CACHE
 
 import torch
 from PIL import Image
@@ -167,6 +168,42 @@ def update_selection(evt: gr.SelectData, width, height):
         height,
     )
 
+
+def load_lora_from_hub(lora: dict, lora_scale: float):
+    """Load LoRA weights from huggingface hub"""
+    with calculateDuration(f"Loading LoRA weights for {lora.get('title')}"):
+        try:
+            pipe.load_lora_weights(
+                lora.get("repo", ""), 
+                weight_name=lora.get("weights", None), 
+                adapter_name="default",
+                low_cpu_mem_usage=True
+            )
+            # Set adapter scale
+            pipe.set_adapters(["default"], adapter_weights=[lora_scale])
+        except Exception as e:
+            print(f"Error loading LoRA: {e}")
+            gr.Warning("Failed to load LoRA weights. Generating with base model.")
+
+
+def load_local_lora(lora: dict, lora_scale: float):
+    """Load LoRA weights from local cache folder"""
+    with calculateDuration(f"Loading LoRA weights for {lora.get('title')}"):
+        try:
+            pipe.load_lora_weights(
+                LORA_CHECKPOINTS_CACHE,
+                cache_dir=LORA_CHECKPOINTS_CACHE,
+                weight_name=lora.get("weights", None),
+                local_files_only=True,
+                low_cpu_mem_usage=True
+            )
+            # Set adapter scale
+            pipe.set_adapters(["default"], adapter_weights=[lora_scale])
+        except Exception as e:
+            print(f"Error loading LoRA: {e}")
+            gr.Warning("Failed to load LoRA weights. Generating with base model.")
+
+
 @spaces.GPU
 def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     # Clean up previous LoRAs in both cases
@@ -176,7 +213,6 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     # Check if a LoRA is selected
     if selected_index is not None and selected_index < len(loras):
         selected_lora = loras[selected_index]
-        lora_path = selected_lora["repo"]
         trigger_word = selected_lora["trigger_word"]
         
         # Prepare Prompt with Trigger Word
@@ -188,24 +224,12 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
                     prompt_mash = f"{prompt} {trigger_word}"
             else:
                 prompt_mash = f"{trigger_word} {prompt}"
-        else:
-            prompt_mash = prompt
 
-        # Load LoRA
-        with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
-            weight_name = selected_lora.get("weights", None)
-            try:
-                pipe.load_lora_weights(
-                    lora_path, 
-                    weight_name=weight_name, 
-                    adapter_name="default",
-                    low_cpu_mem_usage=True
-                )
-                # Set adapter scale
-                pipe.set_adapters(["default"], adapter_weights=[lora_scale])
-            except Exception as e:
-                print(f"Error loading LoRA: {e}")
-                gr.Warning("Failed to load LoRA weights. Generating with base model.")
+        # Special handling of lora loading if there's a civitai key
+        if selected_lora.get("src") == "civitai":
+            load_local_lora(selected_lora, lora_scale)
+        else:
+            load_lora_from_hub(selected_lora, lora_scale)
     else:
         # Base Model Case
         print("No LoRA selected. Running with Base Model.")
@@ -236,12 +260,12 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     yield final_image, seed, gr.update(visible=False)
 
 def get_huggingface_safetensors(link):
-  split_link = link.split("/")
-  if(len(split_link) == 2):
+    split_link = link.split("/")
+    if(len(split_link) == 2):
             model_card = ModelCard.load(link)
             base_model = model_card.data.get("base_model")
             print(base_model)
-      
+
             # Relaxed check to allow Z-Image or Flux or others, assuming user knows what they are doing
             # or specifically check for Z-Image-Turbo
             if base_model not in ["Tongyi-MAI/Z-Image-Turbo", "black-forest-labs/FLUX.1-dev"]:
@@ -264,21 +288,23 @@ def get_huggingface_safetensors(link):
               print(e)
               gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
               raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            return split_link[1], link, safetensors_name, trigger_word, image_url
+            return split_link[1], link, safetensors_name, trigger_word, image_url, "", False
 
 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 "civitai" in link:
+            return get_civitai_safetensors(link)
     else: 
-        return get_huggingface_safetensors(link)
+        return ""
 
 def add_custom_lora(custom_lora):
     global loras
     if(custom_lora):
         try:
-            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
+            title, repo, path, trigger_word, image, src, nsfw = check_custom_model(custom_lora)
             print(f"Loaded custom LoRA: {repo}")
             card = f'''
             <div class="custom_lora_card">
@@ -299,7 +325,9 @@ def add_custom_lora(custom_lora):
                     "title": title,
                     "repo": repo,
                     "weights": path,
-                    "trigger_word": trigger_word
+                    "trigger_word": trigger_word,
+                    "src": src,
+                    "nsfw": nsfw
                 }
                 print(new_item)
                 existing_item_index = len(loras)

civitai_utils.py

@@ -0,0 +1,66 @@
+import os, re
+import gradio as gr
+from dotenv import load_dotenv
+
+load_dotenv() # load CIVITAI_API_TOKEN as environment variable
+from civitai_downloader import civitai_download, APIClient
+
+LORA_CHECKPOINTS_CACHE = os.path.join(os.getcwd(), "loras")
+
+def verify_url(url: str) -> Tuple[int, int]:
+    """
+    Verify if link adheres to the format below:
+    https://civitai.com/models/<number>?modelVersionId=<number>
+
+    returns (models id, model version id)
+    """
+    models_id = 0
+    model_version_id = 0
+
+    if "civitai" in url and "models" in url and "download" not in url:
+        # Extract the first sequence of numbers 
+        match = re.search(r'\d+', url)
+        if match:
+            models_id = match.group(0)
+
+        # Extract the second sequence of numbers
+        match2 = re.search(r'(?:\d+.*?)(\d+)', url)
+        if match2:
+            model_version_id = match2.group(1)
+    return int(models_id), int(model_version_id)
+
+
+def get_model_info(models_id: int, model_version_id: int) -> dict:
+    api = APIClient()
+    item = {}
+    item["title"] = api.get_model(models_id).name
+    item["src"] = "civitai"
+    get_model_version_details(item, model_version_id)
+    return item
+
+
+def get_model_version_details(item:dict, model_version_id: int) -> None:
+    api = APIClient()
+    model_version = api.get_model_version(model_version_id)
+    item["image"] = model_version.images[0].url
+    item["repo"] = LORA_CHECKPOINTS_CACHE           # local path to safetensor
+    item["weights"] =  model_version.files[0].name
+    item["trigger_word"] = model_version.trainedWords[0]
+    item["nsfw"] = model_version.model.get("nsfw", False)
+
+
+def get_civitai_safetensors(url: str):
+    """Helper function to download lora weights"""
+    models_id, model_version_id = verify_url(url)
+    if models_id != 0 and model_version_id != 0: 
+        info = get_model_info(models_id, model_version_id)
+        print(info)
+
+        # make folder and download
+        os.makedirs(LORA_CHECKPOINTS_CACHE, exist_ok=True) 
+        civitai_download(model_version_id, LORA_CHECKPOINTS_CACHE)
+        return info.get("title"), info.get("repo"), info.get("path"), info.get("trigger_word"), info.get("image"), info.get("src"), info.get("nsfw")
+    else:
+        print("Invalid URL - must be in this format https://civitai.com/models/<number>?modelVersionId=<number>") 
+        gr.Warning(f"Invalid URL - must be in this format https://civitai.com/models/<number>?modelVersionId=<number>")
+        return "", "", "", "", "", "", False

requirements.txt

@@ -10,4 +10,6 @@ kernels
 spaces
 torch
 numpy
-peft
+peft
+python-dotenv
+civitai-model-downloader==0.4.7