diff --git a/civitai/lib.py b/civitai/lib.py index cc3af6c..5124a0a 100644 --- a/civitai/lib.py +++ b/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 diff --git a/civitai/link.py b/civitai/link.py index e7446a6..e4848e9 100644 --- a/civitai/link.py +++ b/civitai/link.py @@ -215,6 +215,8 @@ 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): @@ -222,6 +224,10 @@ def join_room(key): 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(): diff --git a/javascript/civitai.js b/javascript/civitai.js index 44d07bc..86e0913 100644 --- a/javascript/civitai.js +++ b/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); diff --git a/scripts/api.py b/scripts/api.py index 1c9d3c9..6986065 100644 --- a/scripts/api.py +++ b/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") diff --git a/scripts/info.py b/scripts/info.py index 2a9f890..3e98193 100644 --- a/scripts/info.py +++ b/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() diff --git a/scripts/link.py b/scripts/link.py index 923edbc..33dc2bf 100644 --- a/scripts/link.py +++ b/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) diff --git a/scripts/previews.py b/scripts/previews.py index 60bf532..1b0c83f 100644 --- a/scripts/previews.py +++ b/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() diff --git a/style.css b/style.css index c72bd55..317aaf5 100644 --- a/style.css +++ b/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; } \ No newline at end of file