diff --git a/install_comfyui.py b/install_comfyui.py
index e1cc8a0..b943042 100644
--- a/install_comfyui.py
+++ b/install_comfyui.py
@@ -12,6 +12,22 @@ def main(install_location):
     git.Repo.clone_from(git_repo_url, install_location)
 
 
+def update(install_location):
+    print("[sd-webui-comfyui]", f"Updating comfyui at {install_location}...")
+    if not install_location.is_dir() or not any(install_location.iterdir()):
+        print("[sd-webui-comfyui]", f"Cannot update comfyui since it is not installed.", file=sys.stderr)
+        return
+
+    import git
+    repo = git.Repo(install_location)
+    current = repo.head.commit
+    repo.remotes.origin.pull()
+    if current == repo.head.commit:
+        print("[sd-webui-comfyui]", "Already up to date.")
+    else:
+        print("[sd-webui-comfyui]", "Done updating comfyui.")
+
+
 if __name__ == '__main__':
     install_location = default_install_location
     if len(sys.argv) > 1:
diff --git a/lib_comfyui/custom_extension_injector.py b/lib_comfyui/custom_extension_injector.py
index e99c4db..8bbda2f 100644
--- a/lib_comfyui/custom_extension_injector.py
+++ b/lib_comfyui/custom_extension_injector.py
@@ -54,16 +54,15 @@ async def get_extensions(request):
             continue
 
         code_patch = generate_prompt_server_init_code_patch(custom_scripts_path)
-        code_patch = textwrap.dedent(code_patch)
         extra_code = ast.parse(code_patch)
         function_to_patch.body[1:1] = extra_code.body
 
 
 def generate_prompt_server_init_code_patch(custom_scripts_path):
-    return rf'''
+    return textwrap.dedent(rf"""
         files.extend(os.path.join(self.web_root, "webui_scripts", "{os.path.basename(os.path.dirname(custom_scripts_path))}", os.path.relpath(f, r"{custom_scripts_path}")) 
         for f in glob.glob(r"{custom_scripts_path}/extensions/**/*.js", recursive=True))
-    '''
+    """)
 
 
 # patch for https://github.com/comfyanonymous/ComfyUI/blob/490771b7f495c95fb52875cf234fffc367162c7e/server.py#L487
@@ -84,7 +83,10 @@ def add_routes(self):
     for custom_scripts_path in custom_scripts_path_list:
         code_patch = generate_prompt_server_add_routes_code_patch(custom_scripts_path)
         extra_line_of_code = ast.parse(code_patch)
-        add_routes_ast_function.body[1].value.args[0].elts[0:0] = [extra_line_of_code.body[0].value]
+        try:
+            add_routes_ast_function.body[2].value.args[0].elts[0:0] = [extra_line_of_code.body[0].value]
+        except:
+            raise RuntimeError("Cannot patch comfyui as it is not up to date")
 
 
 def generate_prompt_server_add_routes_code_patch(custom_scripts_path):
diff --git a/lib_comfyui/webui/callbacks.py b/lib_comfyui/webui/callbacks.py
index 31a7b7a..8deb691 100644
--- a/lib_comfyui/webui/callbacks.py
+++ b/lib_comfyui/webui/callbacks.py
@@ -24,7 +24,8 @@ def on_ui_settings():
 
 @ipc.restrict_to_process('webui')
 def on_after_component(*args, **kwargs):
-    return patches.watch_prompts(*args, **kwargs)
+    patches.watch_prompts(*args, **kwargs)
+    settings.subscribe_update_button(*args, **kwargs)
 
 
 @ipc.restrict_to_process('webui')
diff --git a/lib_comfyui/webui/settings.py b/lib_comfyui/webui/settings.py
index 229a335..3248217 100644
--- a/lib_comfyui/webui/settings.py
+++ b/lib_comfyui/webui/settings.py
@@ -12,6 +12,10 @@ def create_section():
 
     section = ('comfyui', "ComfyUI")
     shared.opts.add_option('comfyui_enabled', shared.OptionInfo(True, 'Enable sd-webui-comfyui extension', section=section))
+
+    shared.opts.add_option("comfyui_update_button", shared.OptionInfo(
+        "Update comfyui (requires reload ui)", "Update comfyui", gr.Button, section=section))
+
     shared.opts.add_option("comfyui_install_location", shared.OptionInfo(
         install_comfyui.default_install_location, "ComfyUI install location", section=section))
     shared.opts.add_option("comfyui_additional_args", shared.OptionInfo(
@@ -66,6 +70,17 @@ def update_reverse_proxy_enabled():
     global_state.reverse_proxy_enabled = reverse_proxy_choices[reverse_proxy_enabled]() and getattr(shared.cmd_opts, "api", False)
 
 
+@ipc.restrict_to_process("webui")
+def subscribe_update_button(component, **kwargs):
+    if getattr(component, "elem_id", None) == "setting_comfyui_update_button":
+        component.click(fn=update_comfyui)
+
+
+@ipc.restrict_to_process("webui")
+def update_comfyui():
+    install_comfyui.update(get_install_location())
+
+
 ipc_strategy_choices = {
     'Default': ipc.strategies.OsFriendlyIpcStrategy,
     'Shared memory': ipc.strategies.SharedMemoryIpcStrategy,