From f9abe4cddcdc6704be02633d9d5ed9640d6b9008 Mon Sep 17 00:00:00 2001
From: Sayo <ws11579@gmail.com>
Date: Mon, 8 May 2023 20:38:10 +0800
Subject: Add api method to get LoRA models with prompt

---
 extensions-builtin/Lora/lora.py        | 13 +++----------
 extensions-builtin/Lora/scripts/api.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 10 deletions(-)
 create mode 100644 extensions-builtin/Lora/scripts/api.py

(limited to 'extensions-builtin/Lora')

diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py
index 8fc1ddca..05162e41 100644
--- a/extensions-builtin/Lora/lora.py
+++ b/extensions-builtin/Lora/lora.py
@@ -2,9 +2,8 @@ import glob
 import os
 import re
 import torch
-from typing import Union, List, Optional
-from fastapi import FastAPI
-import gradio as gr
+from typing import Union
+import scripts.api as api
 
 from modules import shared, devices, sd_models, errors, scripts
 
@@ -445,12 +444,6 @@ def infotext_pasted(infotext, params):
     if added:
         params["Prompt"] += "\n" + "".join(added)
 
-def api(_: gr.Blocks, app: FastAPI):
-    @app.get("/sdapi/v1/loras")
-    async def getloras():
-        return [{"name": name, "path": available_loras[name].filename, "prompt": ""} for name in available_loras]
-
-
 available_loras = {}
 available_lora_aliases = {}
 loaded_loras = []
@@ -458,6 +451,6 @@ loaded_loras = []
 list_available_loras()
 try:
     import modules.script_callbacks as script_callbacks
-    script_callbacks.on_app_started(api)
+    script_callbacks.on_app_started(api.api)
 except:
     pass
\ No newline at end of file
diff --git a/extensions-builtin/Lora/scripts/api.py b/extensions-builtin/Lora/scripts/api.py
new file mode 100644
index 00000000..f1f2e2fc
--- /dev/null
+++ b/extensions-builtin/Lora/scripts/api.py
@@ -0,0 +1,31 @@
+from fastapi import FastAPI
+import gradio as gr
+import json
+import os
+import lora
+
+def get_lora_prompts(path):
+    directory, filename = os.path.split(path)
+    name_without_ext = os.path.splitext(filename)[0]
+    new_filename = name_without_ext + '.civitai.info'
+    try:
+        new_path = os.path.join(directory, new_filename)
+        if os.path.exists(new_path):
+            with open(new_path, 'r') as f:
+                data = json.load(f)
+                trained_words = data.get('trainedWords', [])
+                if len(trained_words) > 0:
+                    result = ','.join(trained_words)
+                    return result
+                else:
+                    return ''
+        else:
+            return ''
+    except Exception as e:
+        return ''
+
+def api(_: gr.Blocks, app: FastAPI):
+    @app.get("/sdapi/v1/loras")
+    async def get_loras():
+        return [{"name": name, "path": lora.available_loras[name].filename, "prompt": get_lora_prompts(lora.available_loras[name].filename)} for name in lora.available_loras]
+
-- 
cgit v1.2.1