Merge pull request #122 from jiveabillion/lora_refresh_info_download_…

…link_reconnect

Added Link status icon and reconnect on click. Added Lora info downlo…

civitai/lib.py

@@ -9,6 +9,7 @@
 
 from tqdm import tqdm
 from modules import shared, sd_models, sd_vae, hashes
+from modules.ui_extra_networks import extra_pages
 from modules.paths import models_path
 from civitai.models import Command, ResourceRequest
 
@@ -22,6 +23,8 @@
 user_agent = 'CivitaiLink:Automatic1111'
 download_chunk_size = 8192
 cache_key = 'civitai'
+refresh_previews_function = None
+refresh_info_function = None
 #endregion
 
 #region Utils
@@ -358,15 +361,23 @@ def load_textual_inversion(resource: ResourceRequest, on_progress=None):
 
 def load_lora(resource: ResourceRequest, on_progress=None):
     isAvailable = load_if_missing(os.path.join(get_lora_dir(), resource['name']), resource['url'], on_progress)
-    # TODO: reload lora list - not sure best way to import this
-    # if isAvailable is None:
-        # lora.list_available_loras()
+    if isAvailable is None:
+        page = next(iter([x for x in extra_pages if x.name == "lora"]), None)
+        if page is not None:
+            log('Refreshing Loras')
+            page.refresh()
+            if refresh_previews_function is not None:
+                refresh_previews_function()
+            if refresh_info_function is not None:
+                refresh_info_function()
 
 def load_locon(resource: ResourceRequest, on_progress=None):
     isAvailable = load_if_missing(os.path.join(get_locon_dir(), resource['name']), resource['url'], on_progress)
-    # TODO: reload locon list - not sure best way to import this
-    # if isAvailable is None:
-        # lora.list_available_loras()
+    if isAvailable is None:
+        page = next(iter([x for x in extra_pages if x.name == "lora"]), None)
+        if page is not None:
+            log('Refreshing Locons')
+            page.refresh()
 
 def load_vae(resource: ResourceRequest, on_progress=None):
     # TODO: find by hash instead of name

civitai/link.py

@@ -215,13 +215,19 @@ def command_response(payload, history=False):
 def socketio_connect():
     if (sio.connected): return
     sio.connect(socketio_url, socketio_path='/api/socketio')
+def is_connected() -> bool:
+    return sio.connected
 
 current_key = None
 def join_room(key):
     if (current_key == key): return
     def on_join(payload):
         log(f"Joined room {key}")
     sio.emit('join', key, callback=on_join)
+
+def rejoin_room(key):
+    current_key = None
+    join_room(key)
 
 old_short_key = None
 def on_civitai_link_key_changed():

javascript/civitai.js

@@ -69,23 +69,60 @@
     }
 
     let statusElement = document.createElement('div');
+    let alphaStatusElement = document.createElement('div');
     let currentStatus = false;
+    let currentAlphaStatus = false;
     async function checkStatus() {
-        const { connected } = await fetch('/civitai/v1/link-status').then(x=>x.json());
+        const { connected,alpha_connected } = await fetch('/civitai/v1/link-status').then(x=>x.json());
         if (currentStatus != connected) {
             currentStatus = connected;
             statusElement.classList.toggle('connected', connected);
         }
+        if (currentAlphaStatus != alpha_connected) {
+            currentAlphaStatus = alpha_connected;
+            alphaStatusElement.classList.toggle('connected', alpha_connected);
+        }
+    }
+    async function checkAlphaStatus() {
+        const { connected } = await fetch('/civitai/v1/alpha-link-status').then(x=>x.json());
+        if (currentStatus != connected) {
+            currentStatus = connected;
+            alphaStatusElement.classList.toggle('alpha-connected', connected);
+        }
     }
     async function startStatusChecks() {
         statusElement.id = 'civitai-status';
+
         statusElement.classList.add('civitai-status');
+        alphaStatusElement.id = 'civitai-alpha-status';
+        alphaStatusElement.classList.add('civitai-alpha-status');
         await getElement('.gradio-container'); // wait for gradio to load
         gradioApp().appendChild(statusElement);
-
+        gradioApp().appendChild(alphaStatusElement);
+        alphaStatusElement.addEventListener('click',async function(){
+            if(true || !currentAlphaStatus)
+            {
+            const { message } = await fetch('/civitai/v1/reconnect-link').then(x=>x.json());
+                if(message == 'Civitai Link active')
+                {
+                    notify({success:true,message:message});
+                }
+                else
+                {
+                    notify({success:false,message:message});
+                }
+                console.log(message);
+            }
+            else
+            {
+                notify({success:true,message:'Not disconnected'});
+            }
+
+        });
         setInterval(checkStatus, 1000 * 10);
         checkStatus();
     }
+
 
     // Bootstrap
     const searchParams = new URLSearchParams(location.search);

scripts/api.py

@@ -1,7 +1,7 @@
 # api endpoints
 import gradio as gr
 from fastapi import FastAPI
-
+from scripts.link import reconnect_to_civitai,get_link_status
 from modules import script_callbacks as script_callbacks
 
 import civitai.lib as civitai
@@ -11,6 +11,12 @@ def civitaiAPI(_: gr.Blocks, app: FastAPI):
     @app.get('/civitai/v1/link-status')
     def link_status():
         return { "connected": civitai.connected }
-
+    @app.get('/civitai/v1/reconnect-link')
+    def reconnect_link():
+        msg = reconnect_to_civitai()
+        return { "message": msg }
+    @app.get('/civitai/v1/alpha-link-status')
+    def alpha_link_status():
+        return { "connected": get_link_status() }
 script_callbacks.on_app_started(civitaiAPI)
 civitai.log("API loaded")

scripts/info.py

@@ -83,6 +83,7 @@ def load_info():
 
 # Automatically pull model with corresponding hash from Civitai
 def start_load_info(demo: gr.Blocks, app):
+    civitai.refresh_info_function = load_info
     thread = threading.Thread(target=load_info)
     thread.start()
 

scripts/link.py

@@ -12,6 +12,26 @@ def connect_to_civitai(demo: gr.Blocks, app):
     link.log('Connecting to Civitai Link Server')
     link.socketio_connect()
     link.join_room(key)
+def get_link_status()-> bool:
+    return link.is_connected()
+
+def reconnect_to_civitai() -> str:
+    key = shared.opts.data.get("civitai_link_key", None)
+    # If key is empty or not set, don't connect to Civitai Link
+    if not key: 
+        msg = 'Civitai Link Key is empty'
+        link.log(msg)
+        return msg
+
+    link.log('Reconnecting to Civitai Link Server')
+    link.socketio_connect()
+    link.rejoin_room(key)
+    if link.is_connected():
+        msg = 'Civitai Link active'
+        link.log(msg)
+        return msg
+    else:
+        return 'Civitai Link not connected'
 
 script_callbacks.on_app_started(connect_to_civitai)
 

scripts/previews.py

@@ -57,6 +57,7 @@ def load_previews():
 
 # Automatically pull model with corresponding hash from Civitai
 def start_load_previews(demo: gr.Blocks, app):
+    civitai.refresh_previews_function = load_previews
     thread = threading.Thread(target=load_previews)
     thread.start()
 

style.css

@@ -19,7 +19,29 @@ div.civitai-status:before {
     right:1px;
     border: 1px solid rgba(255,255,255,0.3);
 }
-
+div.civitai-alpha-status{
+    position: absolute;
+    top: 36px;
+    right: 5px;
+    width:24px;
+    height:24px;
+    background-repeat: no-repeat;
+    background-size: cover;
+	background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' height='1em' viewBox='0 0 512 512'%3E%3C!--! Font Awesome Free 6.4.2 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license (Commercial License) Copyright 2023 Fonticons, Inc. --%3E%3Cstyle%3Esvg%7Bfill:%23ff0000%7D%3C/style%3E%3Cpath d='M228.3 469.1L47.6 300.4c-4.2-3.9-8.2-8.1-11.9-12.4h87c22.6 0 43-13.6 51.7-34.5l10.5-25.2 49.3 109.5c3.8 8.5 12.1 14 21.4 14.1s17.8-5 22-13.3L320 253.7l1.7 3.4c9.5 19 28.9 31 50.1 31H476.3c-3.7 4.3-7.7 8.5-11.9 12.4L283.7 469.1c-7.5 7-17.4 10.9-27.7 10.9s-20.2-3.9-27.7-10.9zM503.7 240h-132c-3 0-5.8-1.7-7.2-4.4l-23.2-46.3c-4.1-8.1-12.4-13.3-21.5-13.3s-17.4 5.1-21.5 13.3l-41.4 82.8L205.9 158.2c-3.9-8.7-12.7-14.3-22.2-14.1s-18.1 5.9-21.8 14.8l-31.8 76.3c-1.2 3-4.2 4.9-7.4 4.9H16c-2.6 0-5 .4-7.3 1.1C3 225.2 0 208.2 0 190.9v-5.8c0-69.9 50.5-129.5 119.4-141C165 36.5 211.4 51.4 244 84l12 12 12-12c32.6-32.6 79-47.5 124.6-39.9C461.5 55.6 512 115.2 512 185.1v5.8c0 16.9-2.8 33.5-8.3 49.1z'/%3E%3C/svg%3E");
+    z-index: 1000;
+}
+div.civitai-alpha-status:before {
+    pointer-events:all;
+    width:6px;
+    height:6px;
+    background: red;
+    border-radius: 50%;
+    content: '';
+    position:absolute;
+    top:-4px;
+    right:1px;
+    border: 1px solid rgba(255,255,255,0.3);
+}
 /* blinking animation */
 @keyframes blink {
     0% { opacity: 0.2; }
@@ -30,4 +52,8 @@ div.civitai-status:before {
 div.civitai-status.connected:before {
     background:green;
     animation: blink 1s ease-in-out infinite;
+}
+div.civitai-alpha-status.connected:before {
+    background:green;
+    animation: blink 1s ease-in-out infinite;
 }