fix(api): add pre-converted ControlNets to base sources, add missing ControlNet strings

api/onnx_web/convert/__main__.py

@@ -139,6 +139,7 @@ base_models: Models = {
     ],
     # download only
     "sources": [
+        # CodeFormer: no ONNX yet
         (
             "detection-resnet50-final",
             "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth",
@@ -163,6 +164,49 @@ base_models: Models = {
             "parsing-parsenet",
             "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth",
         ),
+        # ControlNets: already converted
+        {
+            "dest": "control",
+            "format": "onnx",
+            "name": "canny",
+            "source": "https://huggingface.co/ForserX/sd-controlnet-canny-onnx/resolve/main/model.onnx",
+        },
+        {
+            "dest": "control",
+            "format": "onnx",
+            "name": "depth",
+            "source": "https://huggingface.co/ForserX/sd-controlnet-depth-onnx/resolve/main/model.onnx",
+        },
+        {
+            "dest": "control",
+            "format": "onnx",
+            "name": "hed",
+            "source": "https://huggingface.co/ForserX/sd-controlnet-hed-onnx/resolve/main/model.onnx",
+        },
+        {
+            "dest": "control",
+            "format": "onnx",
+            "name": "mlsd",
+            "source": "https://huggingface.co/ForserX/sd-controlnet-mlsd-onnx/resolve/main/model.onnx",
+        },
+        {
+            "dest": "control",
+            "format": "onnx",
+            "name": "normal",
+            "source": "https://huggingface.co/ForserX/sd-controlnet-normal-onnx/resolve/main/model.onnx",
+        },
+        {
+            "dest": "control",
+            "format": "onnx",
+            "name": "openpose",
+            "source": "https://huggingface.co/ForserX/sd-controlnet-openpose-onnx/resolve/main/model.onnx",
+        },
+        {
+            "dest": "control",
+            "format": "onnx",
+            "name": "seg",
+            "source": "https://huggingface.co/ForserX/sd-controlnet-seg-onnx/resolve/main/model.onnx",
+        },
     ],
 }
 
@@ -242,7 +286,13 @@ def convert_models(conversion: ConversionContext, args, models: Models):
                 source = model["source"]
 
                 try:
-                    dest = fetch_model(conversion, name, source, format=model_format)
+                    dest_path = None
+                    if "dest" in model:
+                        dest_path = path.join(conversion.model_path, model["dest"])
+
+                    dest = fetch_model(
+                        conversion, name, source, format=model_format, dest=dest_path
+                    )
                     logger.info("finished downloading source: %s -> %s", source, dest)
                 except Exception:
                     logger.exception("error fetching source %s", name)

api/onnx_web/convert/diffusion/control.py

@@ -44,7 +44,7 @@ def convert_diffusion_control(
     if is_torch_2_0:
         controlnet.set_attn_processor(AttnProcessor())
 
-    cnet_path = output_path / "cnet" / ONNX_MODEL
+    cnet_path = output_path / ONNX_MODEL
     onnx_export(
         controlnet,
         model_args=(

api/onnx_web/convert/diffusion/diffusers.py

@@ -314,17 +314,20 @@ def convert_diffusion_diffusers(
 
     del pipeline.unet
 
-    convert_diffusion_diffusers_cnet(
-        conversion,
-        source,
-        device,
-        output_path,
-        dtype,
-        unet_in_channels,
-        unet_sample_size,
-        num_tokens,
-        text_hidden_size,
-    )
+    if not single_vae:
+        convert_diffusion_diffusers_cnet(
+            conversion,
+            source,
+            device,
+            output_path,
+            dtype,
+            unet_in_channels,
+            unet_sample_size,
+            num_tokens,
+            text_hidden_size,
+        )
+    else:
+        logger.debug("skipping CNet for single-VAE model")
 
     if cnet_only:
         logger.info("done converting CNet")

api/onnx_web/models/rrdb.py

@@ -77,7 +77,15 @@ class RRDB(nn.Module):
 
 
 class RRDBNet(nn.Module):
-    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4):
+    def __init__(
+        self,
+        num_in_ch=3,
+        num_out_ch=3,
+        num_feat=64,
+        num_block=23,
+        num_grow_ch=32,
+        scale=4,
+    ):
         super(RRDBNet, self).__init__()
         RRDB_block_f = functools.partial(RRDB, nf=num_feat, gc=num_grow_ch)
         self.sf = scale

api/schemas/extras.yaml

@@ -109,6 +109,8 @@ $defs:
     type: object
     required: [name, source]
     properties:
+      dest:
+        type: string
       format:
         type: string
       name:

gui/src/strings/en.ts

@@ -114,9 +114,11 @@ export const I18N_STRINGS_EN = {
           canny: 'Canny Edges',
           depth: 'Depth Map',
           hed: 'HED Edges',
+          mlsd: 'M-LSD Segmentation',
           normal: 'Normal Map',
           openpose: 'OpenPose',
           scribble: 'Scribble',
+          segment: 'Image Segmentation',
         },
       },
       modelType: {
@@ -244,17 +246,17 @@ export const I18N_STRINGS_EN = {
       },
       sourceFilter: {
         none: 'None',
-        gaussian: 'Gaussian Blur',
-        noise: 'Histogram Noise',
+        canny: 'Canny Edges',
+        depth: 'Depth Map',
         face: 'Face Detection',
-        segment: 'Image Segmentation',
+        gaussian: 'Gaussian Blur',
+        hed: 'HED Edges',
         mlsd: 'M-LSD Segmentation',
         normal: 'Normal Map',
-        hed: 'HED Edges',
-        scribble: 'Scribble',
-        depth: 'Depth Map',
-        canny: 'Canny Edges',
+        noise: 'Histogram Noise',
         openpose: 'OpenPose',
+        scribble: 'Scribble',
+        segment: 'Image Segmentation',
       },
       tab: {
         blend: 'Blend',

models/.gitignore

@@ -1 +1 @@
-*
+*