Support changing base model

app.py

@@ -49,7 +49,14 @@ This is an unofficial demo for [https://github.com/lllyasviel/ControlNet](https:
 
 If you are interested in trying out other base models, check out [this Space](https://huggingface.co/spaces/hysts/ControlNet-with-other-models) as well.
 '''
-if (SPACE_ID := os.getenv('SPACE_ID')) is not None:
+
+SPACE_ID = os.getenv('SPACE_ID')
+ALLOW_CHANGING_BASE_MODEL = SPACE_ID != 'hysts/ControlNet'
+
+if not ALLOW_CHANGING_BASE_MODEL:
+    DESCRIPTION += 'In this Space, the base model is not allowed to be changed so as not to slow down the demo, but it can be changed if you duplicate the Space.'
+
+if SPACE_ID is not None:
     DESCRIPTION += f'''<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.<br/>
 <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true">
 <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
@@ -59,7 +66,9 @@ if (SPACE_ID := os.getenv('SPACE_ID')) is not None:
 MAX_IMAGES = int(os.getenv('MAX_IMAGES', '3'))
 DEFAULT_NUM_IMAGES = min(MAX_IMAGES, int(os.getenv('DEFAULT_NUM_IMAGES', '1')))
 
-model = Model()
+DEFAULT_MODEL_ID = os.getenv('DEFAULT_MODEL_ID',
+                             'runwayml/stable-diffusion-v1-5')
+model = Model(base_model_id=DEFAULT_MODEL_ID, task_name='canny')
 
 with gr.Blocks(css='style.css') as demo:
     gr.Markdown(DESCRIPTION)
@@ -106,4 +115,31 @@ with gr.Blocks(css='style.css') as demo:
                                max_images=MAX_IMAGES,
                                default_num_images=DEFAULT_NUM_IMAGES)
 
+    with gr.Accordion(label='Base model', open=False):
+        with gr.Row():
+            with gr.Column():
+                current_base_model = gr.Text(label='Current base model')
+            with gr.Column(scale=0.3):
+                check_base_model_button = gr.Button('Check current base model')
+        with gr.Row():
+            with gr.Column():
+                base_model_id = gr.Text(
+                    label='Base model repo',
+                    max_lines=1,
+                    placeholder='runwayml/stable-diffusion-v1-5',
+                    info=
+                    'The base model must be compatible with Stable Diffusion v1.5.',
+                    interactive=ALLOW_CHANGING_BASE_MODEL)
+            with gr.Column(scale=0.3):
+                change_base_model_button = gr.Button('Change base model')
+
+    check_base_model_button.click(fn=lambda: model.base_model_id,
+                                  outputs=current_base_model)
+    base_model_id.submit(fn=model.set_base_model,
+                         inputs=base_model_id,
+                         outputs=current_base_model)
+    change_base_model_button.click(fn=model.set_base_model,
+                                   inputs=base_model_id,
+                                   outputs=current_base_model)
+
 demo.queue(api_open=False).launch()

model.py

@@ -39,18 +39,21 @@ CONTROLNET_MODEL_IDS = {
 
 
 class Model:
-    def __init__(self):
-        # FIXME
-        self.base_model_id = 'andite/anything-v4.0'
-        self.task_name = 'pose'
-        self.pipe = self.load_pipe()
-
-    def load_pipe(self) -> DiffusionPipeline:
-        model_id = CONTROLNET_MODEL_IDS[self.task_name]
+    def __init__(self,
+                 base_model_id: str = 'runwayml/stable-diffusion-v1-5',
+                 task_name: str = 'canny'):
+        self.base_model_id = ''
+        self.task_name = ''
+        self.pipe = self.load_pipe(base_model_id, task_name)
+
+    def load_pipe(self, base_model_id: str, task_name) -> DiffusionPipeline:
+        if base_model_id == self.base_model_id and task_name == self.task_name:
+            return self.pipe
+        model_id = CONTROLNET_MODEL_IDS[task_name]
         controlnet = ControlNetModel.from_pretrained(model_id,
                                                      torch_dtype=torch.float16)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            self.base_model_id,
+            base_model_id,
             safety_checker=None,
             controlnet=controlnet,
             torch_dtype=torch.float16)
@@ -58,8 +61,14 @@ class Model:
             pipe.scheduler.config)
         pipe.enable_xformers_memory_efficient_attention()
         pipe.enable_model_cpu_offload()
+        self.base_model_id = base_model_id
+        self.task_name = task_name
         return pipe
 
+    def set_base_model(self, base_model_id: str) -> str:
+        self.pipe = self.load_pipe(base_model_id, self.task_name)
+        return self.base_model_id
+
     def load_controlnet_weight(self, task_name: str) -> None:
         if task_name == self.task_name:
             return