Add separate VAE support (#308)

adetailer/args.py

@@ -61,6 +61,8 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
     ad_cfg_scale: NonNegativeFloat = 7.0
     ad_use_checkpoint: bool = False
     ad_checkpoint: Optional[str] = None
+    ad_use_vae: bool = False
+    ad_vae: Optional[str] = None
     ad_use_sampler: bool = False
     ad_sampler: str = "DPM++ 2M Karras"
     ad_use_noise_multiplier: bool = False
@@ -142,6 +144,10 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
             "ADetailer use separate checkpoint",
             ["ADetailer use separate checkpoint", "ADetailer checkpoint"],
         )
+        ppop(
+            "ADetailer use separate VAE",
+            ["ADetailer use separate VAE", "ADetailer VAE"],
+        )
         ppop(
             "ADetailer use separate sampler",
             ["ADetailer use separate sampler", "ADetailer sampler"],
@@ -216,6 +222,8 @@ _all_args = [
     ("ad_cfg_scale", "ADetailer CFG scale"),
     ("ad_use_checkpoint", "ADetailer use separate checkpoint"),
     ("ad_checkpoint", "ADetailer checkpoint"),
+    ("ad_use_vae", "ADetailer use separate VAE"),
+    ("ad_vae", "ADetailer VAE"),
     ("ad_use_sampler", "ADetailer use separate sampler"),
     ("ad_sampler", "ADetailer sampler"),
     ("ad_use_noise_multiplier", "ADetailer use separate noise multiplier"),

adetailer/ui.py

@@ -30,6 +30,7 @@ class WebuiInfo:
     t2i_button: gr.Button
     i2i_button: gr.Button
     checkpoints_list: Callable
+    vae_list: Callable
 
 
 def gr_interactive(value: bool = True):
@@ -441,27 +442,50 @@ def inpainting(w: Widgets, n: int, is_img2img: bool, webui_info: WebuiInfo):
                 )
 
             with gr.Column(variant="compact"):
-                w.ad_use_sampler = gr.Checkbox(
-                    label="Use separate sampler" + suffix(n),
+                w.ad_use_vae = gr.Checkbox(
+                    label="Use separate VAE (experimental)" + suffix(n),
                     value=False,
                     visible=True,
-                    elem_id=eid("ad_use_sampler"),
+                    elem_id=eid("ad_use_vae"),
                 )
 
-                w.ad_sampler = gr.Dropdown(
-                    label="ADetailer sampler" + suffix(n),
-                    choices=webui_info.sampler_names,
-                    value=webui_info.sampler_names[0],
+                vaes = ["Use same VAE"]
+                try:
+                    vaes.extend(webui_info.vae_list())
+                except Exception:
+                    vaes.extend([])
+
+                w.ad_vae = gr.Dropdown(
+                    label="ADetailer VAE" + suffix(n),
+                    choices=vaes,
+                    value=vaes[0],
                     visible=True,
-                    elem_id=eid("ad_sampler"),
+                    elem_id=eid("ad_vae"),
                 )
 
-                w.ad_use_sampler.change(
-                    gr_interactive,
-                    inputs=w.ad_use_sampler,
-                    outputs=w.ad_sampler,
-                    queue=False,
-                )
+
+        with gr.Row(), gr.Column(variant="compact"):
+            w.ad_use_sampler = gr.Checkbox(
+                label="Use separate sampler" + suffix(n),
+                value=False,
+                visible=True,
+                elem_id=eid("ad_use_sampler"),
+            )
+
+            w.ad_sampler = gr.Dropdown(
+                label="ADetailer sampler" + suffix(n),
+                choices=webui_info.sampler_names,
+                value=webui_info.sampler_names[0],
+                visible=True,
+                elem_id=eid("ad_sampler"),
+            )
+
+            w.ad_use_sampler.change(
+                gr_interactive,
+                inputs=w.ad_use_sampler,
+                outputs=w.ad_sampler,
+                queue=False,
+            )
 
         with gr.Row():
             with gr.Column(variant="compact"):

scripts/!adetailer.py

@@ -126,6 +126,7 @@ class AfterDetailerScript(scripts.Script):
             t2i_button=txt2img_submit_button,
             i2i_button=img2img_submit_button,
             checkpoints_list=modules.sd_models.checkpoint_tiles,
+            vae_list=modules.shared_items.sd_vae_items,
         )
 
         components, infotext_fields = adui(num_models, is_img2img, webui_info)
@@ -318,6 +319,13 @@ class AfterDetailerScript(scripts.Script):
             and args.ad_checkpoint not in ("None", "Use same checkpoint")
         ):
             d["sd_model_checkpoint"] = args.ad_checkpoint
+
+        if (
+            args.ad_use_vae
+            and args.ad_vae
+            and args.ad_vae not in ("None", "Use same VAE")
+        ):
+            d["sd_vae"] = args.ad_vae
         return d
 
     def get_initial_noise_multiplier(self, p, args: ADetailerArgs) -> float | None: