Support ControlNet

Tested with lllyasviel/ControlNet-v1-1/control_v11f1p_sd15_depth.pth
The configuration file is derived from the checkpoint name.

app.py

@@ -24,7 +24,7 @@ demo = gr.Interface(
         gr.Text(max_lines=1, label="your_hf_token"),
         gr.Text(max_lines=1, label="model_id"),
         gr.Text(max_lines=1, label="filename"),
-        gr.Radio(label="Model type", choices=["v1", "v2"]),
+        gr.Radio(label="Model type", choices=["v1", "v2", "ControlNet"]),
         gr.Radio(label="Sample size (px)", choices=[512, 768]),
         gr.Radio(label="Scheduler type", choices=["pndm", "heun", "euler", "dpm", "ddim"], value="dpm"),
         gr.Radio(label="Extract EMA or non-EMA?", choices=["ema", "non-ema"], value="ema"),

convert.py

@@ -1,19 +1,19 @@
-import argparse
+import gradio as gr
 import requests
-import json
 import os
 import shutil
-from collections import defaultdict
-from inspect import signature
+from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Dict, List, Optional, Set
+from typing import Optional
 
 import torch
 from io import BytesIO
 
 from huggingface_hub import CommitInfo, Discussion, HfApi, hf_hub_download
 from huggingface_hub.file_download import repo_folder_name
-from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt
+from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
+    download_from_original_stable_diffusion_ckpt, download_controlnet_from_original_ckpt
+)
 from transformers import CONFIG_MAPPING
 
 
@@ -33,15 +33,23 @@ def convert_single(model_id: str, filename: str, model_type: str, sample_size: i
             config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference.yaml"
         else:
             config_url = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inference-v.yaml"
+    elif model_type == "ControlNet":
+        config_url = (Path(model_id)/"resolve/main"/filename).with_suffix(".yaml")
+        config_url = "https://huggingface.co/" + str(config_url)
 
     config_file = BytesIO(requests.get(config_url).content)
 
-    pipeline = download_from_original_stable_diffusion_ckpt(ckpt_file, config_file, image_size=sample_size, scheduler_type=scheduler_type, from_safetensors=from_safetensors, extract_ema=extract_ema)
+    if model_type == "ControlNet":
+        pipeline = download_controlnet_from_original_ckpt(ckpt_file, config_file, image_size=sample_size, from_safetensors=from_safetensors, extract_ema=extract_ema)
+        to_args = {"dtype": torch.float16}
+    else:
+        pipeline = download_from_original_stable_diffusion_ckpt(ckpt_file, config_file, image_size=sample_size, scheduler_type=scheduler_type, from_safetensors=from_safetensors, extract_ema=extract_ema)
+        to_args = {"torch_dtype": torch.float16}
 
     pipeline.save_pretrained(folder)
     pipeline.save_pretrained(folder, safe_serialization=True)
 
-    pipeline = pipeline.to(torch_dtype=torch.float16)
+    pipeline = pipeline.to(**to_args)
     pipeline.save_pretrained(folder, variant="fp16")
     pipeline.save_pretrained(folder, safe_serialization=True, variant="fp16")
 
@@ -74,6 +82,8 @@ def convert(token: str, model_id: str, filename: str, model_type: str, sample_si
             new_pr  = api.upload_folder(folder_path=folder, path_in_repo="./", repo_id=model_id, repo_type="model", token=token, commit_description=COMMIT_MESSAGE.format(model_id), create_pr=True)
             pr_number = new_pr.split("%2F")[-1].split("/")[0]
             link = f"Pr created at: {'https://huggingface.co/' + os.path.join(model_id, 'discussions', pr_number)}"
+        except Exception as e:
+            raise gr.exceptions.Error(str(e))
         finally:
             shutil.rmtree(folder)