add reference to original, reimplement control-xs

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2023-12-26
+## Update for 2023-12-27
 
 *Note*: based on `diffusers==0.25.0.dev0`
 
@@ -59,8 +59,9 @@
       (previously via settings -> upscaler_for_img2img)  
 - **General**  
   - new **onboarding**  
-    if no models are found during startup, app will no longer ask to download default checkpoint  
-    instead, it will show message in UI with options to change model path or download any of the reference checkpoints  
+    - if no models are found during startup, app will no longer ask to download default checkpoint  
+      instead, it will show message in UI with options to change model path or download any of the reference checkpoints  
+    - *extra networks -> models -> reference* section is now enabled for both original and diffusers backend  
   - support for **Torch 2.1.2** (release) and **Torch 2.3** (dev)  
   - **Process** create videos from batch or folder processing  
       supports *GIF*, *PNG* and *MP4* with full interpolation, scene change detection, etc.  

html/reference.json

@@ -2,7 +2,8 @@
   "DreamShaper SD 1.5 v8": {
     "path": "dreamshaper_8.safetensors@https://civitai.com/api/download/models/128713",
     "desc": "Showcase finetuned model based on Stable diffusion 1.5",
-    "preview": "dreamshaper_8.jpg"
+    "preview": "dreamshaper_8.jpg",
+    "original": true
   },
   "DreamShaper SD XL Turbo": {
     "path": "dreamshaperXL_turboDpmppSDE.safetensors@https://civitai.com/api/download/models/251662",
@@ -12,7 +13,8 @@
   "Juggernaut Reborn": {
     "path": "juggernaut_reborn.safetensors@https://civitai.com/api/download/models/274039",
     "desc": "Showcase finetuned model based on Stable diffusion 1.5",
-    "preview": "juggernaut_reborn.jpg"
+    "preview": "juggernaut_reborn.jpg",
+    "original": true
   },
   "Juggernaut XL v7 RunDiffusion": {
     "path": "juggernautXL_v7Rundiffusion.safetensors@https://civitai.com/api/download/models/240840",
@@ -21,13 +23,24 @@
   },
   "RunwayML SD 1.5": {
     "path": "runwayml/stable-diffusion-v1-5",
+    "alt": "v1-5-pruned-emaonly.safetensors@https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors?download=true",
     "desc": "Stable Diffusion 1.5 is the base model all other 1.5 checkpoint were trained from. It's a latent text-to-image diffusion model capable of generating photo-realistic images given any text input. The Stable-Diffusion-v1-5 checkpoint was initialized with the weights of the Stable-Diffusion-v1-2 checkpoint and subsequently fine-tuned on 595k steps at resolution 512x512.",
-    "preview": "runwayml--stable-diffusion-v1-5.jpg"
+    "preview": "runwayml--stable-diffusion-v1-5.jpg",
+    "original": true
   },
-  "StabilityAI SD 2.1": {
+  "StabilityAI SD 2.1 EMA": {
     "path": "stabilityai/stable-diffusion-2-1-base",
-    "desc": "This stable-diffusion-2-1 model is fine-tuned from stable-diffusion-2 (768-v-ema.ckpt) with an additional 55k steps on the same dataset. Improvement over base 1.5 model, but never really took off.",
-    "preview": "stabilityai--stable-diffusion-2.1-base.jpg"
+    "alt": "v2-1_512-ema-pruned.safetensors@https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.safetensors?download=true",
+    "desc": "This stable-diffusion-2-1-base model fine-tunes stable-diffusion-2-base (512-base-ema.ckpt) with 220k extra steps taken",
+    "preview": "stabilityai--stable-diffusion-2.1-base.jpg",
+    "original": true
+  },
+  "StabilityAI SD 2.1 V": {
+    "path": "stabilityai/stable-diffusion-2-1-base",
+    "alt": "v2-1_768-ema-pruned.safetensors@https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors?download=true",
+    "desc": "This stable-diffusion-2 model is resumed from stable-diffusion-2-base (512-base-ema.ckpt) and trained for 150k steps using a v-objective on the same dataset. Resumed for another 140k steps on 768x768 images",
+    "preview": "stabilityai--stable-diffusion-2.1-base.jpg",
+    "original": true
   },
   "StabilityAI SD-XL 1.0 Base": {
     "path": "stabilityai/stable-diffusion-xl-base-1.0",
@@ -36,8 +49,10 @@
   },
   "StabilityAI SD 2.1 Turbo": {
     "path": "stabilityai/sd-turbo",
+    "alt": "sd_turbo.safetensors@https://huggingface.co/stabilityai/sd-turbo/resolve/main/sd_turbo.safetensors?download=true",
     "desc": "SD-Turbo is a distilled version of Stable Diffusion 2.1, trained for real-time synthesis. SD-Turbo is based on a novel training method called Adversarial Diffusion Distillation (ADD) (see the technical report), which allows sampling large-scale foundational image diffusion models in 1 to 4 steps at high image quality. This approach uses score distillation to leverage large-scale off-the-shelf image diffusion models as a teacher signal and combines this with an adversarial loss to ensure high image fidelity even in the low-step regime of one or two sampling steps.",
-    "preview": "stabilityai--sd-turbo.jpg"
+    "preview": "stabilityai--sd-turbo.jpg",
+    "original": true
   },
   "StabilityAI SD-XL Turbo": {
     "path": "stabilityai/sdxl-turbo",

javascript/control.js

@@ -2,6 +2,7 @@ function setupControlUI() {
   const tabs = ['input', 'output', 'preview'];
   for (const tab of tabs) {
     const btn = gradioApp().getElementById(`control-${tab}-button`);
+    if (!btn) continue; // eslint-disable-line no-continue
     btn.style.cursor = 'pointer';
     btn.onclick = () => {
       const t = gradioApp().getElementById(`control-tab-${tab}`);

javascript/sdnext.css

@@ -70,6 +70,7 @@ button.custom-button{ border-radius: var(--button-large-radius); padding: var(--
 .performance { font-size: 0.85em; color: #444; }
 .performance p { display: inline-block; color: var(--body-text-color-subdued) !important }
 .performance .time { margin-right: 0; }
+.thumbnails { background: var(--body-background-fill); }
 #control_gallery { height: 564px; }
 #control-result { padding: 0.5em; }
 #control-inputs { margin-top: 1em; }

modules/control/proc/edge.py

@@ -58,7 +58,7 @@ class EdgeDetector:
         edge_map = cv2.resize(edge_map, (W, H), interpolation=cv2.INTER_LINEAR)
 
         if output_type == "pil":
-            edge_map = edge_map.convert('L')
             edge_map = Image.fromarray(edge_map)
+            edge_map = edge_map.convert('L')
 
         return edge_map

modules/control/units/xs.py

@@ -1,16 +1,11 @@
 import os
 import time
 from typing import Union
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
 from modules.shared import log, opts
 from modules import errors
-
-ok = True
-try:
-    from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, ControlNetXSModel, StableDiffusionControlNetXSPipeline, StableDiffusionXLControlNetXSPipeline
-except Exception:
-    from diffusers import ControlNetModel
-    ControlNetXSModel = ControlNetModel # dummy
-    ok = False
+from modules.control.units.xs_model import ControlNetXSModel
+from modules.control.units.xs_pipe import StableDiffusionControlNetXSPipeline, StableDiffusionXLControlNetXSPipeline
 
 
 what = 'ControlNet-XS'
@@ -43,8 +38,6 @@ def find_models():
 
 def list_models(refresh=False):
     global models # pylint: disable=global-statement
-    if not ok:
-        return models
     import modules.shared
     if not refresh and len(models) > 0:
         return models
@@ -130,7 +123,7 @@ class ControlNetXSPipeline():
                 tokenizer_2=pipeline.tokenizer_2,
                 unet=pipeline.unet,
                 scheduler=pipeline.scheduler,
-                feature_extractor=getattr(pipeline, 'feature_extractor', None),
+                # feature_extractor=getattr(pipeline, 'feature_extractor', None),
                 controlnet=controlnet, # can be a list
             ).to(pipeline.device)
         elif isinstance(pipeline, StableDiffusionPipeline):

modules/ui.py

@@ -648,9 +648,6 @@ def create_ui(startup_timer = None):
 
             ui_extra_networks.setup_ui(extra_networks_ui, txt2img_gallery)
 
-            with FormRow():
-                gr.HTML(value="", elem_id="main_info", visible=False, elem_classes=["main-info"])
-
     timer.startup.record("ui-txt2img")
 
     import modules.img2img # pylint: disable=redefined-outer-name

modules/ui_common.py

@@ -202,6 +202,8 @@ def create_output_panel(tabname, preview=True):
 
     with gr.Column(variant='panel', elem_id=f"{tabname}_results"):
         with gr.Group(elem_id=f"{tabname}_gallery_container"):
+            if tabname == "txt2img":
+                gr.HTML(value="", elem_id="main_info", visible=False, elem_classes=["main-info"])
             # columns are for <576px, <768px, <992px, <1200px, <1400px, >1400px
             result_gallery = gr.Gallery(value=[], label='Output', show_label=False, show_download_button=True, allow_preview=True, elem_id=f"{tabname}_gallery", container=False, preview=preview, columns=5, object_fit='scale-down', height=shared.opts.gallery_height or None)
 

modules/ui_extra_networks.py

@@ -233,11 +233,10 @@ class ExtraNetworksPage:
         allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews()]
         for parentdir, dirs in {d: modelloader.directory_list(d) for d in allowed_folders}.items():
             for tgt in dirs.keys():
-                if shared.backend == shared.Backend.DIFFUSERS:
-                    if os.path.join(paths.models_path, 'Reference') in tgt:
-                        subdirs['Reference'] = 1
-                    if shared.opts.diffusers_dir in tgt:
-                        subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
+                if os.path.join(paths.models_path, 'Reference') in tgt:
+                    subdirs['Reference'] = 1
+                if shared.backend == shared.Backend.DIFFUSERS and shared.opts.diffusers_dir in tgt:
+                    subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
                 if 'models--' in tgt:
                     continue
                 subdir = tgt[len(parentdir):].replace("\\", "/")
@@ -248,7 +247,7 @@ class ExtraNetworksPage:
                     subdirs[subdir] = 1
         debug(f"Extra networks: page='{self.name}' subfolders={list(subdirs)}")
         subdirs = OrderedDict(sorted(subdirs.items()))
-        if shared.backend == shared.Backend.DIFFUSERS and self.name == 'model':
+        if self.name == 'model':
             subdirs['Reference'] = 1
             subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
             subdirs.move_to_end(os.path.basename(shared.opts.diffusers_dir))

modules/ui_extra_networks_checkpoints.py

@@ -15,21 +15,25 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
         shared.refresh_checkpoints()
 
     def list_reference(self): # pylint: disable=inconsistent-return-statements
-        if shared.backend != shared.Backend.DIFFUSERS:
-            return []
         reference_models = shared.readfile(os.path.join('html', 'reference.json'))
         for k, v in reference_models.items():
+            if shared.backend != shared.Backend.DIFFUSERS:
+                if not v.get('original', False):
+                    continue
+                url = v.get('alt', None) or v['path']
+            else:
+                url = v['path']
             name = os.path.join(reference_dir, k)
             preview = v.get('preview', v['path'])
             yield {
                 "type": 'Model',
                 "name": name,
                 "title": name,
-                "filename": v['path'],
+                "filename": url,
                 "search_term": self.search_terms_from_path(name),
                 "preview": self.find_preview(os.path.join(reference_dir, preview)),
                 "local_preview": self.find_preview_file(os.path.join(reference_dir, preview)),
-                "onclick": '"' + html.escape(f"""return selectReference({json.dumps(v['path'])})""") + '"',
+                "onclick": '"' + html.escape(f"""return selectReference({json.dumps(url)})""") + '"',
                 "hash": None,
                 "mtime": 0,
                 "size": 0,
@@ -74,4 +78,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
             yield record
 
     def allowed_directories_for_previews(self):
-        return [v for v in [shared.opts.ckpt_dir, shared.opts.diffusers_dir, reference_dir, sd_models.model_path] if v is not None]
+        if shared.backend == shared.Backend.DIFFUSERS:
+            return [v for v in [shared.opts.ckpt_dir, shared.opts.diffusers_dir, reference_dir] if v is not None]
+        else:
+            return [v for v in [shared.opts.ckpt_dir, reference_dir, sd_models.model_path] if v is not None]

pyproject.toml

@@ -40,6 +40,8 @@ exclude = [
   "repositories/blip",
   "repositories/codeformer",
   "modules/control/proc/normalbae/nets/submodules/efficientnet_repo/geffnet",
+  "modules/control/units/*_model.py",
+  "modules/control/units/*_pipe.py",
 ]
 ignore = [
 	"A003", # Class attirbute shadowing builtin

requirements.txt

@@ -24,7 +24,6 @@ lmdb
 lpips
 omegaconf
 open-clip-torch
-opencv-contrib-python-headless
 piexif
 psutil
 pyyaml
@@ -51,7 +50,7 @@ antlr4-python3-runtime==4.9.3
 requests==2.31.0
 tqdm==4.66.1
 accelerate==0.25.0
-opencv-python-headless==4.8.1.78
+opencv-contrib-python-headless==4.8.1.78
 diffusers==0.24.0
 einops==0.4.1
 gradio==3.43.2