en load refiner, track server job state

.eslintrc.json

@@ -73,7 +73,7 @@
     "register_drag_drop": "readonly",
     //extraNetworks.js
     "requestGet": "readonly",
-    "popup": "readonly",
+    "getENActiveTab": "readonly",
     // from python
     "localization": "readonly",
     // progressbar.js
@@ -82,6 +82,8 @@
     // imageviewer.js
     "modalPrevImage": "readonly",
     "modalNextImage": "readonly",
+    // logMonitor.js
+    "jobStatusEl": "readonly",
     // loader.js
     "removeSplash": "readonly"
   },

CHANGELOG.md

@@ -1,12 +1,13 @@
 # Change Log for SD.Next
 
-## Update for 2023-09-24
+## Update for 2023-09-29
 
 **TBD**: Planned before release:
 - Integrate LoRA/Lyco for *backend:original*
 - Add FreeU for *backend:diffusers*
 
 This is a big one, with some major changes and new functionality...
+And probably the biggest release since introduction of **Diffusers**  
 
 Note that for this release its recommended to perform a clean install (fresh `git clone`)  
 Upgrades are still possible and supported, but above is recommended for best experience  
@@ -17,13 +18,17 @@ Upgrades are still possible and supported, but above is recommended for best exp
   - converted submenus from checkboxes to accordion elements  
     any ui state including state of open/closed menus can be saved as default!  
     see *System -> User interface -> Set menu states*  
+  - small visual indicator bottom right of page showing internal server job state  
 - **Extra networks**:  
-  - you can scan [civitai](https://civitai.com/)  
-    for missing metadata and previews directly from extra networks  
-  - new details interface to view and save data about extra networks  
-    main ui now has a single button on each en to trigger details view  
-  - details view includes model/lora metadata parser!  
-  - details view includes civitai model metadata!  
+  - **Details**
+    - new details interface to view and save data about extra networks  
+      main ui now has a single button on each en to trigger details view  
+    - details view includes model/lora metadata parser!  
+    - details view includes civitai model metadata!  
+  - **Metadata**:
+    - you can scan [civitai](https://civitai.com/)  
+      for missing metadata and previews directly from extra networks  
+      simply click on button in top-right corner of extra networks page  
   - **Styles**
     - save/apply icons moved to extra networks  
     - can be edited in details view  
@@ -31,9 +36,13 @@ Upgrades are still possible and supported, but above is recommended for best exp
     - support for embedded previews  
   - **VAE**
     - VAEs are now also listed as part of extra networks
-  - faster search, ability to show/hide/sort networks  
-  - refactored subfolder handling  
-    *note*: this will trigger model hash recaclulation on first model use  
+  - **Refiner**
+    - you can load model from extra networks as base model or as refiner  
+      simply select button in top-right of models page  
+  - **General**
+    - faster search, ability to show/hide/sort networks  
+    - refactored subfolder handling  
+      *note*: this will trigger model hash recaclulation on first model use  
 - **Diffusers**:
   - better pipeline auto-detect when loading from safetensors  
   - **SDXL Inpaint**  
@@ -110,6 +119,7 @@ Upgrades are still possible and supported, but above is recommended for best exp
     - get browser session info in server log  
     - when running with `--debug` flag, log is force-rotated  
       so each `sdnext.log.*` represents exactly one server run  
+    - internal server job state tracking
   - **API**
     - add end-to-end example how to use API: `cli/simple-txt2img.js`  
       covers txt2img, upscale, hires, refiner  

html/locale_en.json

@@ -19,6 +19,7 @@
   {"id":"","label":"⟲","localized":"","hint":"Refresh"},
   {"id":"","label":"✕","localized":"","hint":"Close"},
   {"id":"","label":"⊜","localized":"","hint":"Fill"},
+  {"id":"","label":"⌾","localized":"","hint":"Load model as refiner model when selected, otherwise load as base model"},
   {"id":"","label":"🕸️","localized":"","hint":"Scan CivitAI for missing metadata and previews"},
   {"id":"","label":"📐","localized":"","hint":"Measure"},
   {"id":"","label":"🔍","localized":"","hint":"Search"}

javascript/extraNetworks.js

@@ -186,24 +186,31 @@ function setupExtraNetworksForTab(tabname) {
   gradioApp().querySelector(`#${tabname}_extra_tabs`).classList.add('extra-networks');
   const en = gradioApp().getElementById(`${tabname}_extra_networks`);
   const tabs = gradioApp().querySelector(`#${tabname}_extra_tabs > div`);
-  const txtSearch = gradioApp().querySelector(`#${tabname}_extra_search textarea`);
-  const txtDescription = gradioApp().getElementById(`${tabname}_description`);
+
+  // buttons
   const btnRefresh = gradioApp().getElementById(`${tabname}_extra_refresh`);
   const btnScan = gradioApp().getElementById(`${tabname}_extra_scan`);
   const btnSave = gradioApp().getElementById(`${tabname}_extra_save`);
   const btnClose = gradioApp().getElementById(`${tabname}_extra_close`);
-  txtSearch.classList.add('search');
-  txtDescription.classList.add('description');
+  const btnModel = gradioApp().getElementById(`${tabname}_extra_model`);
   const buttons = document.createElement('span');
   buttons.classList.add('buttons');
   if (btnRefresh) buttons.appendChild(btnRefresh);
+  if (btnModel) buttons.appendChild(btnModel);
   if (btnScan) buttons.appendChild(btnScan);
   if (btnSave) buttons.appendChild(btnSave);
   if (btnClose) buttons.appendChild(btnClose);
+  btnModel.onclick = () => btnModel.classList.toggle('toolbutton-selected');
   tabs.appendChild(buttons);
+
+  // search and description
   const div = document.createElement('div');
   div.classList.add('second-line');
   tabs.appendChild(div);
+  const txtSearch = gradioApp().querySelector(`#${tabname}_extra_search textarea`);
+  const txtDescription = gradioApp().getElementById(`${tabname}_description`);
+  txtSearch.classList.add('search');
+  txtDescription.classList.add('description');
   div.appendChild(txtSearch);
   div.appendChild(txtDescription);
   let searchTimer = null;
@@ -215,6 +222,7 @@ function setupExtraNetworksForTab(tabname) {
     }, 150);
   });
 
+  // card hover
   let hoverTimer = null;
   let previousCard = null;
   gradioApp().getElementById(`${tabname}_extra_tabs`).onmouseover = (e) => {
@@ -233,6 +241,7 @@ function setupExtraNetworksForTab(tabname) {
     };
   };
 
+  // en style
   const intersectionObserver = new IntersectionObserver((entries) => {
     if (!en) return;
     for (const el of Array.from(gradioApp().querySelectorAll('.extra-networks-page'))) {

javascript/logMonitor.js

@@ -1,5 +1,6 @@
 let logMonitorEl = null;
 let logMonitorStatus = true;
+let jobStatusEl = null;
 
 async function logMonitor() {
   if (logMonitorStatus) setTimeout(logMonitor, opts.logmonitor_refresh_period);
@@ -51,6 +52,9 @@ async function initLogMonitor() {
     </table>
   `;
   el.style.display = 'none';
+  jobStatusEl = document.createElement('div');
+  jobStatusEl.className = 'jobStatus';
+  gradioApp().appendChild(jobStatusEl);
   fetch(`/sdapi/v1/start?agent=${encodeURI(navigator.userAgent)}`);
   logMonitor();
   log('initLogMonitor');

javascript/progressBar.js

@@ -47,7 +47,7 @@ function setProgress(res) {
   let eta = '';
   if (res?.paused) eta = 'Paused';
   else if (res?.completed || (progress > 0.99)) eta = 'Finishing';
-  else if (sec === 0) eta = 'Starting';
+  else if (sec === 0) eta = `Init${res?.job?.length > 0 ? `: ${res.job}` : ''}`;
   else {
     const min = Math.floor(sec / 60);
     sec %= 60;
@@ -106,6 +106,7 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
     debug('taskEnd:', id_task);
     localStorage.removeItem('task');
     setProgress();
+    if (jobStatusEl) jobStatusEl.style.display = 'none';
     if (parentGallery && livePreview) parentGallery.removeChild(livePreview);
     checkPaused(true);
     if (atEnd) atEnd();
@@ -113,6 +114,8 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
 
   const start = (id_task, id_live_preview) => { // eslint-disable-line no-shadow
     request('./internal/progress', { id_task, id_live_preview }, (res) => {
+      if (jobStatusEl) jobStatusEl.innerText = (res?.job || '').trim().toUpperCase();
+      if (jobStatusEl) jobStatusEl.style.display = jobStatusEl.innerText.length > 0 ? 'block' : 'none';
       lastState = res;
       const elapsedFromStart = (new Date() - dateStart) / 1000;
       hasStarted |= res.active;
@@ -124,7 +127,7 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
       if (res.live_preview && !livePreview) init();
       if (res.live_preview && galleryEl) img.src = res.live_preview;
       if (onProgress) onProgress(res);
-      setTimeout(() => start(id_task, id_live_preview), opts.live_preview_refresh_period || 250);
+      setTimeout(() => start(id_task, id_live_preview), opts.live_preview_refresh_period || 500);
     }, done);
   };
   start(id_task, 0);

javascript/style.css

@@ -120,6 +120,8 @@ div#extras_scale_to_tab div.form{ flex-direction: row; }
   z-index: var(--layer-2);
 }
 .tooltip-show { opacity: 0.9; }
+.toolbutton-selected { background: var(--background-fill-primary) !important; }
+.jobStatus { position: fixed; bottom: 1em; right: 1em; background: var(--input-background-fill); padding: 0.4em; font-size: 0.8em; color: var(--body-text-color-subdued); }
 
 /* settings */
 #si-sparkline-memo, #si-sparkline-load { background-color: #111; }

javascript/ui.js

@@ -336,13 +336,19 @@ function updateInput(target) {
 let desiredCheckpointName = null;
 function selectCheckpoint(name) {
   desiredCheckpointName = name;
-  gradioApp().getElementById('change_checkpoint').click();
+  const tabname = getENActiveTab();
+  const btnModel = gradioApp().getElementById(`${tabname}_extra_model`);
+  const isRefiner = btnModel && btnModel.classList.contains('toolbutton-selected');
+  if (isRefiner) gradioApp().getElementById('change_refiner').click();
+  else gradioApp().getElementById('change_checkpoint').click();
+  log(`Change ${isRefiner ? 'refiner' : 'model'}: ${desiredCheckpointName}`);
 }
 
 let desiredVAEName = null;
 function selectVAE(name) {
   desiredVAEName = name;
   gradioApp().getElementById('change_vae').click();
+  log(`Change VAE: ${desiredVAEName}`);
 }
 
 function currentImg2imgSourceResolution(_a, _b, scaleBy) {

modules/call_queue.py

@@ -40,7 +40,6 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
                 res[-1] = f"<div class='error'>{html.escape(str(e))}</div>"
             finally:
                 progress.finish_task(id_task)
-            shared.state.end()
         return res
     return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True, name=name)
 

modules/hashes.py

@@ -1,3 +1,4 @@
+import copy
 import hashlib
 import os.path
 from rich import progress
@@ -56,6 +57,7 @@ def sha256(filename, title, use_addnet_hash=False):
         return None
     if not os.path.isfile(filename):
         return None
+    orig_state = copy.deepcopy(shared.state)
     shared.state.begin("hashing")
     if use_addnet_hash:
         with progress.open(filename, 'rb', description=f'[cyan]Calculating hash: [yellow]{filename}', auto_refresh=True, console=shared.console) as f:
@@ -67,6 +69,7 @@ def sha256(filename, title, use_addnet_hash=False):
         "sha256": sha256_value
     }
     shared.state.end()
+    shared.state = orig_state
     dump_cache()
     return sha256_value
 

modules/processing_diffusers.py

@@ -362,6 +362,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
         desc='Base',
         **task_specific_kwargs
     )
+    # p.steps = base_args['num_inference_steps']
     p.extra_generation_params['CFG rescale'] = p.diffusers_guidance_rescale
     p.extra_generation_params["Eta"] = shared.opts.scheduler_eta if shared.opts.scheduler_eta is not None and shared.opts.scheduler_eta > 0 and shared.opts.scheduler_eta < 1 else None
     try:
@@ -414,6 +415,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                     strength=p.denoising_strength,
                     desc='Hires',
                 )
+                # p.steps += hires_args['num_inference_steps']
                 try:
                     output = shared.sd_model(**hires_args) # pylint: disable=not-callable
                 except AssertionError as e:
@@ -476,6 +478,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
                 clip_skip=p.clip_skip,
                 desc='Refiner',
             )
+            # p.steps += refiner_args['num_inference_steps']
             try:
                 refiner_output = shared.sd_refiner(**refiner_args) # pylint: disable=not-callable
             except AssertionError as e:

modules/progress.py

@@ -43,6 +43,7 @@ class ProgressRequest(BaseModel):
 
 
 class InternalProgressResponse(BaseModel):
+    job: str = Field(default=None, title="Job name", description="Internal job name")
     active: bool = Field(title="Whether the task is being worked on right now")
     queued: bool = Field(title="Whether the task is in queue")
     paused: bool = Field(title="Whether the task is paused")
@@ -64,14 +65,12 @@ def progressapi(req: ProgressRequest):
     completed = req.id_task in finished_tasks
     paused = shared.state.paused
     if not active:
-        return InternalProgressResponse(active=active, queued=queued, paused=paused, completed=completed, id_live_preview=-1, textinfo="Queued..." if queued else "Waiting...")
+        return InternalProgressResponse(job=shared.state.job, active=active, queued=queued, paused=paused, completed=completed, id_live_preview=-1, textinfo="Queued..." if queued else "Waiting...")
     progress = 0
-    job_count, job_no = shared.state.job_count, shared.state.job_no
-    sampling_steps, sampling_step = shared.state.sampling_steps, shared.state.sampling_step
-    if job_count > 0:
-        progress += job_no / job_count
-    if sampling_steps > 0 and job_count > 0:
-        progress += 1 / job_count * sampling_step / sampling_steps
+    if shared.state.job_count > 0:
+        progress += shared.state.job_no / shared.state.job_count
+    if shared.state.sampling_steps > 0 and shared.state.job_count > 0:
+        progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps
     progress = min(progress, 1)
     elapsed_since_start = time.time() - shared.state.time_start
     predicted_duration = elapsed_since_start / progress if progress > 0 else None
@@ -84,4 +83,4 @@ def progressapi(req: ProgressRequest):
         shared.state.current_image.save(buffered, format='jpeg')
         live_preview = f'data:image/jpeg;base64,{base64.b64encode(buffered.getvalue()).decode("ascii")}'
         id_live_preview = shared.state.id_live_preview
-    return InternalProgressResponse(active=active, queued=queued, paused=paused, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)
+    return InternalProgressResponse(job=shared.state.job, active=active, queued=queued, paused=paused, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo)

modules/sd_models.py

@@ -3,6 +3,7 @@ import io
 import sys
 import json
 import time
+import copy
 import logging
 import threading
 import contextlib
@@ -545,7 +546,8 @@ class ModelData:
                     if shared.backend == shared.Backend.ORIGINAL:
                         reload_model_weights(op='model')
                     elif shared.backend == shared.Backend.DIFFUSERS:
-                        load_diffuser(op='model')
+                        reload_model_weights(op='model')
+                        # load_diffuser(op='model')
                     else:
                         shared.log.error(f"Unknown Execution backend: {shared.backend}")
                     self.initial = False
@@ -1126,6 +1128,9 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
     if checkpoint_info is None:
         unload_model_weights(op=op)
         return
+    orig_state = copy.deepcopy(shared.state)
+    shared.state = shared.State()
+    shared.state.begin(f'load-{op}')
     if load_dict:
         shared.log.debug(f'Model dict: existing={sd_model is not None} target={checkpoint_info.filename} info={info}')
     else:
@@ -1165,8 +1170,11 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
             model_data.sd_dict = shared.opts.sd_model_dict
             shared.opts.data["sd_model_checkpoint"] = next_checkpoint_info.title
             reload_model_weights(reuse_dict=True) # ok we loaded dict now lets redo and load model on top of it
+        shared.state.end()
+        shared.state = orig_state
         return model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner
 
+    # fallback
     try:
         load_model_weights(sd_model, checkpoint_info, state_dict, timer)
     except Exception:
@@ -1180,7 +1188,9 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
         if sd_model is not None and not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram and not getattr(sd_model, 'has_accelerate', False):
             sd_model.to(devices.device)
             timer.record("device")
-    shared.log.info(f"Weights loaded in {timer.summary()}")
+    shared.state.end()
+    shared.state = orig_state
+    shared.log.info(f"Loaded: {op} time={timer.summary()}")
 
 
 def disable_offload(sd_model):

modules/shared.py

@@ -147,9 +147,11 @@ class State:
         self.skipped = False
         self.textinfo = None
         self.time_start = time.time()
+        log.debug(f'State begin: {self.job}')
         devices.torch_gc()
 
     def end(self):
+        log.debug(f'State end: {self.job} time={time.time() - self.time_start:.2f}s')
         self.job = ""
         self.job_count = 0
         self.job_no = 0

modules/ui.py

@@ -1168,7 +1168,14 @@ def create_ui(startup_timer = None):
             inputs=[component_dict['sd_model_checkpoint'], dummy_component],
             outputs=[component_dict['sd_model_checkpoint'], text_settings],
         )
-        button_set_vae = gr.Button('Change vae', elem_id='change_vae', visible=False)
+        button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False)
+        button_set_refiner.click(
+            fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'),
+            _js="function(v){ var res = desiredCheckpointName; desiredCheckpointName = ''; return [res || v, null]; }",
+            inputs=[component_dict['sd_model_refiner'], dummy_component],
+            outputs=[component_dict['sd_model_refiner'], text_settings],
+        )
+        button_set_vae = gr.Button('Change VAE', elem_id='change_vae', visible=False)
         button_set_vae.click(
             fn=lambda value, _: run_settings_single(value, key='sd_vae'),
             _js="function(v){ var res = desiredVAEName; desiredVAEName = ''; return [res || v, null]; }",

modules/ui_extra_networks.py

@@ -366,6 +366,7 @@ class ExtraNetworksUi:
         self.button_save: gr.Button = None
         self.button_apply: gr.Button = None
         self.button_close: gr.Button = None
+        self.button_model: gr.Checkbox = None
         self.details_components: list = []
         self.last_item: dict = None
         self.last_page: ExtraNetworksPage = None
@@ -449,12 +450,14 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         def ui_tab_change(page):
             scan_visible = page in ['Model', 'Lora', 'Hypernetwork', 'Embedding']
             save_visible = page in ['Style']
-            return [gr.update(visible=scan_visible), gr.update(visible=save_visible)]
+            model_visible = page in ['Model']
+            return [gr.update(visible=scan_visible), gr.update(visible=save_visible), gr.update(visible=model_visible)]
 
         ui.button_refresh = ToolButton(symbols.refresh, elem_id=tabname+"_extra_refresh")
         ui.button_scan = ToolButton(symbols.scan, elem_id=tabname+"_extra_scan", visible=True)
         ui.button_save = ToolButton(symbols.book, elem_id=tabname+"_extra_save", visible=False)
         ui.button_close = ToolButton(symbols.close, elem_id=tabname+"_extra_close")
+        ui.button_model = ToolButton(symbols.refine, elem_id=tabname+"_extra_model", visible=True)
         ui.search = gr.Textbox('', show_label=False, elem_id=tabname+"_extra_search", placeholder="Search...", elem_classes="textbox", lines=2)
         ui.description = gr.Textbox('', show_label=False, elem_id=tabname+"_description", elem_classes="textbox", lines=2, interactive=False)
 
@@ -466,7 +469,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
             with gr.Tab(page.title, id=page.title.lower().replace(" ", "_"), elem_classes="extra-networks-tab") as tab:
                 hmtl = gr.HTML(page.html, elem_id=f'{tabname}{page.name}_extra_page', elem_classes="extra-networks-page")
                 ui.pages.append(hmtl)
-                tab.select(ui_tab_change, _js="getENActivePage", inputs=[ui.button_details], outputs=[ui.button_scan, ui.button_save])
+                tab.select(ui_tab_change, _js="getENActivePage", inputs=[ui.button_details], outputs=[ui.button_scan, ui.button_save, ui.button_model])
 
         # ui.tabs.change(fn=ui_tab_change, inputs=[], outputs=[ui.button_scan, ui.button_save])
 

modules/ui_symbols.py

@@ -9,7 +9,13 @@ fill = '⊜'
 scan = '🕸️'
 networks = '🌐'
 paste = '⇦'
-
+refine = '⌾'
+switch = '⇅'
+detect = '📐'
+folder = '📂'
+random = '🎲️'
+reuse = '♻️'
+info = 'ℹ' # noqa
 """
 refresh = '🔄'
 close = '🛗'
@@ -21,9 +27,3 @@ fill = '⏫'
 networks = '🌐'
 paste = '📘'
 """
-switch = '⇅'
-detect = '📐'
-folder = '📂'
-random = '🎲️'
-reuse = '♻️'
-info = 'ℹ' # noqa

modules/upscaler.py

@@ -1,4 +1,5 @@
 import os
+import copy
 from abc import abstractmethod
 import PIL
 from PIL import Image
@@ -79,6 +80,8 @@ class Upscaler:
         return img
 
     def upscale(self, img: PIL.Image, scale, selected_model: str = None):
+        orig_state = copy.deepcopy(modules.shared.state)
+        modules.shared.state.begin('upscale')
         self.scale = scale
         dest_w = int(img.width * scale)
         dest_h = int(img.height * scale)
@@ -91,6 +94,8 @@ class Upscaler:
                 break
         if img.width != dest_w or img.height != dest_h:
             img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS)
+        modules.shared.state.end()
+        modules.shared.state = orig_state
         return img
 
     @abstractmethod