make brushnet work

web_app/src/components/SidePanel/DiffusionOptions.tsx

@@ -109,6 +109,84 @@ const DiffusionOptions = () => {
     )
   }
 
+  const renderBrushNetSetting = () => {
+    if (!settings.model.support_brushnet) {
+      return null
+    }
+
+    return (
+      <div className="flex flex-col gap-4">
+        <div className="flex flex-col gap-4">
+          <div className="flex justify-between items-center pr-2">
+            <LabelTitle
+              text="BrushNet"
+              toolTip="BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion"
+              url="https://github.com/TencentARC/BrushNet"
+            />
+            <Switch
+              id="brushnet"
+              checked={settings.enableBrushNet}
+              onCheckedChange={(value) => {
+                updateSettings({ enableBrushNet: value })
+              }}
+            />
+          </div>
+          <div className="flex flex-col gap-1">
+            <RowContainer>
+              <Slider
+                className="w-[180px]"
+                defaultValue={[100]}
+                min={1}
+                max={100}
+                step={1}
+                disabled={!settings.enableBrushNet}
+                value={[Math.floor(settings.brushnetConditioningScale * 100)]}
+                onValueChange={(vals) =>
+                  updateSettings({ brushnetConditioningScale: vals[0] / 100 })
+                }
+              />
+              <NumberInput
+                id="controlnet-weight"
+                className="w-[60px] rounded-full"
+                disabled={!settings.enableBrushNet}
+                numberValue={settings.brushnetConditioningScale}
+                allowFloat={false}
+                onNumberValueChange={(val) => {
+                  updateSettings({ brushnetConditioningScale: val })
+                }}
+              />
+            </RowContainer>
+          </div>
+
+          <div className="pr-2">
+            <Select
+              defaultValue={settings.brushnetMethod}
+              value={settings.brushnetMethod}
+              onValueChange={(value) => {
+                updateSettings({ brushnetMethod: value })
+              }}
+              disabled={!settings.enableBrushNet}
+            >
+              <SelectTrigger>
+                <SelectValue placeholder="Select brushnet model" />
+              </SelectTrigger>
+              <SelectContent align="end">
+                <SelectGroup>
+                  {Object.values(settings.model.brushnets).map((method) => (
+                    <SelectItem key={method} value={method}>
+                      {method}
+                    </SelectItem>
+                  ))}
+                </SelectGroup>
+              </SelectContent>
+            </Select>
+          </div>
+        </div>
+        <Separator />
+      </div>
+    )
+  }
+
   const renderConterNetSetting = () => {
     if (!settings.model.support_controlnet) {
       return null
@@ -881,6 +959,7 @@ const DiffusionOptions = () => {
       {renderSeed()}
       {renderNegativePrompt()}
       <Separator />
+      {renderBrushNetSetting()}
       {renderConterNetSetting()}
       {renderLCMLora()}
       {renderMaskBlur()}

web_app/src/lib/api.ts

@@ -78,6 +78,9 @@ export default async function inpaint(
       controlnet_method: settings.controlnetMethod
         ? settings.controlnetMethod
         : "",
+      enable_brushnet: settings.enableBrushNet,
+      brushnet_method: settings.brushnetMethod ? settings.brushnetMethod : "",
+      brushnet_conditioning_scale: settings.brushnetConditioningScale,
       powerpaint_task: settings.showExtender
         ? PowerPaintTask.outpainting
         : settings.powerpaintTask,

web_app/src/lib/states.ts

@@ -99,6 +99,11 @@ export type Settings = {
   controlnetConditioningScale: number
   controlnetMethod: string
 
+  // BrushNet
+  enableBrushNet: boolean
+  brushnetMethod: string
+  brushnetConditioningScale: number
+
   enableLCMLora: boolean
   enableFreeu: boolean
   freeuConfig: FreeuConfig
@@ -306,15 +311,16 @@ const defaultValues: AppState = {
       path: "lama",
       model_type: "inpaint",
       support_controlnet: false,
+      support_brushnet: false,
       support_strength: false,
       support_outpainting: false,
       controlnets: [],
+      brushnets: [],
       support_freeu: false,
       support_lcm_lora: false,
       is_single_file_diffusers: false,
       need_prompt: false,
     },
-    enableControlnet: false,
     showCropper: false,
     showExtender: false,
     extenderDirection: ExtenderDirection.xy,
@@ -339,8 +345,12 @@ const defaultValues: AppState = {
     sdMatchHistograms: false,
     sdScale: 1.0,
     p2pImageGuidanceScale: 1.5,
-    controlnetConditioningScale: 0.4,
+    enableControlnet: false,
     controlnetMethod: "lllyasviel/control_v11p_sd15_canny",
+    controlnetConditioningScale: 0.4,
+    enableBrushNet: false,
+    brushnetMethod: "random_mask",
+    brushnetConditioningScale: 1.0,
     enableLCMLora: false,
     enableFreeu: false,
     freeuConfig: { s1: 0.9, s2: 0.2, b1: 1.2, b2: 1.4 },
@@ -1076,7 +1086,7 @@ export const useStore = createWithEqualityFn<AppState & AppAction>()(
     })),
     {
       name: "ZUSTAND_STATE", // name of the item in the storage (must be unique)
-      version: 1,
+      version: 2,
       partialize: (state) =>
         Object.fromEntries(
           Object.entries(state).filter(([key]) =>

web_app/src/lib/types.ts

@@ -48,7 +48,9 @@ export interface ModelInfo {
   support_strength: boolean
   support_outpainting: boolean
   support_controlnet: boolean
+  support_brushnet: boolean
   controlnets: string[]
+  brushnets: string[]
   support_freeu: boolean
   support_lcm_lora: boolean
   need_prompt: boolean