diff --git a/api/onnx_web/convert/__main__.py b/api/onnx_web/convert/__main__.py index c9a226a8..38b1573f 100644 --- a/api/onnx_web/convert/__main__.py +++ b/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) diff --git a/api/onnx_web/convert/diffusion/control.py b/api/onnx_web/convert/diffusion/control.py index cc2e4bb5..ffc9f5c9 100644 --- a/api/onnx_web/convert/diffusion/control.py +++ b/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=( diff --git a/api/onnx_web/convert/diffusion/diffusers.py b/api/onnx_web/convert/diffusion/diffusers.py index 62e36107..5d9d59ea 100644 --- a/api/onnx_web/convert/diffusion/diffusers.py +++ b/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") diff --git a/api/onnx_web/models/rrdb.py b/api/onnx_web/models/rrdb.py index 31b2ef60..dd5a14e4 100644 --- a/api/onnx_web/models/rrdb.py +++ b/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 diff --git a/api/schemas/extras.yaml b/api/schemas/extras.yaml index c37fa1d2..6ab55910 100644 --- a/api/schemas/extras.yaml +++ b/api/schemas/extras.yaml @@ -109,6 +109,8 @@ $defs: type: object required: [name, source] properties: + dest: + type: string format: type: string name: diff --git a/gui/src/strings/en.ts b/gui/src/strings/en.ts index 8bc27b75..a7f9a644 100644 --- a/gui/src/strings/en.ts +++ b/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', diff --git a/models/.gitignore b/models/.gitignore index f59ec20a..72e8ffc0 100644 --- a/models/.gitignore +++ b/models/.gitignore @@ -1 +1 @@ -* \ No newline at end of file +*