init work

.vscode/settings.json

@@ -3,7 +3,7 @@
         "editor.defaultFormatter": "ms-python.black-formatter",
         "editor.formatOnType": true,
         "editor.codeActionsOnSave": {
-            "source.organizeImports": true
+            "source.organizeImports": "explicit"
         }
     },
     "black-formatter.args": [

README.md

@@ -1,5 +1,5 @@
 ---
-title: ControlNet V1.1
+title: ControlNet++
 emoji: 📉
 colorFrom: yellow
 colorTo: green

app.py

@@ -20,7 +20,7 @@ from app_softedge import create_demo as create_demo_softedge
 from model import Model
 from settings import ALLOW_CHANGING_BASE_MODEL, DEFAULT_MODEL_ID, SHOW_DUPLICATE_BUTTON
 
-DESCRIPTION = "# ControlNet v1.1"
+DESCRIPTION = "# ControlNet++: Improving Conditional Controls with Efficient Consistency Feedback<p>https://arxiv.org/abs/2404.07987</p>"
 
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
@@ -38,28 +38,14 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Tabs():
         with gr.TabItem("Canny"):
             create_demo_canny(model.process_canny)
-        with gr.TabItem("MLSD"):
-            create_demo_mlsd(model.process_mlsd)
-        with gr.TabItem("Scribble"):
-            create_demo_scribble(model.process_scribble)
-        with gr.TabItem("Scribble Interactive"):
-            create_demo_scribble_interactive(model.process_scribble_interactive)
         with gr.TabItem("SoftEdge"):
             create_demo_softedge(model.process_softedge)
-        with gr.TabItem("OpenPose"):
-            create_demo_openpose(model.process_openpose)
         with gr.TabItem("Segmentation"):
             create_demo_segmentation(model.process_segmentation)
         with gr.TabItem("Depth"):
             create_demo_depth(model.process_depth)
-        with gr.TabItem("Normal map"):
-            create_demo_normal(model.process_normal)
         with gr.TabItem("Lineart"):
             create_demo_lineart(model.process_lineart)
-        with gr.TabItem("Content Shuffle"):
-            create_demo_shuffle(model.process_shuffle)
-        with gr.TabItem("Instruct Pix2Pix"):
-            create_demo_ip2p(model.process_ip2p)
 
     with gr.Accordion(label="Base model", open=False):
         with gr.Row():

app_canny.py

@@ -37,7 +37,7 @@ def create_demo(process):
                         label="Canny high threshold", minimum=1, maximum=255, value=200, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -46,7 +46,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_depth.py

@@ -37,7 +37,7 @@ def create_demo(process):
                         label="Preprocess resolution", minimum=128, maximum=512, value=384, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -46,7 +46,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_ip2p.py

@@ -31,7 +31,7 @@ def create_demo(process):
                         step=256,
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -40,7 +40,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_lineart.py

@@ -31,7 +31,7 @@ def create_demo(process):
                         ],
                         type="value",
                         value="Lineart",
-                        info='Note that "Lineart (anime)" and "None (anime)" are for anime base models like Anything-v3.',
+                        info='Please note that the only difference here is pre-processing. We use exactly the same ControlNet++ model to generate images.',
                     )
                     num_samples = gr.Slider(
                         label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
@@ -47,7 +47,7 @@ def create_demo(process):
                         label="Preprocess resolution", minimum=128, maximum=512, value=512, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -56,7 +56,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_mlsd.py

@@ -40,7 +40,7 @@ def create_demo(process):
                         label="Hough distance threshold (MLSD)", minimum=0.01, maximum=20.0, value=0.1, step=0.01
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -49,7 +49,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_normal.py

@@ -37,7 +37,7 @@ def create_demo(process):
                         label="Preprocess resolution", minimum=128, maximum=512, value=384, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -46,7 +46,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_openpose.py

@@ -37,7 +37,7 @@ def create_demo(process):
                         label="Preprocess resolution", minimum=128, maximum=512, value=512, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -46,7 +46,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_scribble.py

@@ -37,7 +37,7 @@ def create_demo(process):
                         label="Preprocess resolution", minimum=128, maximum=512, value=512, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -46,7 +46,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_scribble_interactive.py

@@ -51,7 +51,7 @@ def create_demo(process):
                         step=256,
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -60,7 +60,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
 
         create_button.click(
             fn=create_canvas,

app_segmentation.py

@@ -37,7 +37,7 @@ def create_demo(process):
                         label="Preprocess resolution", minimum=128, maximum=512, value=512, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -46,7 +46,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_shuffle.py

@@ -34,7 +34,7 @@ def create_demo(process):
                         step=256,
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -43,7 +43,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

app_softedge.py

@@ -30,7 +30,7 @@ def create_demo(process):
                             "None",
                         ],
                         type="value",
-                        value="PidiNet",
+                        value="HED",
                     )
                     num_samples = gr.Slider(
                         label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
@@ -46,7 +46,7 @@ def create_demo(process):
                         label="Preprocess resolution", minimum=128, maximum=512, value=512, step=1
                     )
                     num_steps = gr.Slider(label="Number of steps", minimum=1, maximum=100, value=20, step=1)
-                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                     a_prompt = gr.Textbox(label="Additional prompt", value="best quality, extremely detailed")
@@ -55,7 +55,7 @@ def create_demo(process):
                         value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                     )
             with gr.Column():
-                result = gr.Gallery(label="Output", show_label=False, columns=2, object_fit="scale-down")
+                result = gr.Gallery(label="Output", show_label=False, columns=4, object_fit="scale-down")
         inputs = [
             image,
             prompt,

model.py

@@ -18,19 +18,15 @@ from preprocessor import Preprocessor
 from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES
 
 CONTROLNET_MODEL_IDS = {
-    "Openpose": "lllyasviel/control_v11p_sd15_openpose",
-    "Canny": "lllyasviel/control_v11p_sd15_canny",
-    "MLSD": "lllyasviel/control_v11p_sd15_mlsd",
-    "scribble": "lllyasviel/control_v11p_sd15_scribble",
-    "softedge": "lllyasviel/control_v11p_sd15_softedge",
-    "segmentation": "lllyasviel/control_v11p_sd15_seg",
-    "depth": "lllyasviel/control_v11f1p_sd15_depth",
-    "NormalBae": "lllyasviel/control_v11p_sd15_normalbae",
-    "lineart": "lllyasviel/control_v11p_sd15_lineart",
-    "lineart_anime": "lllyasviel/control_v11p_sd15s2_lineart_anime",
-    "shuffle": "lllyasviel/control_v11e_sd15_shuffle",
-    "ip2p": "lllyasviel/control_v11e_sd15_ip2p",
-    "inpaint": "lllyasviel/control_v11e_sd15_inpaint",
+    "Canny": "../diffusers/work_dirs/reward_model/MultiGen20M_Canny/reward_ft5k_canny_res512_bs256_lr1e-5_warmup100_scale-10_iter10k_fp16_train0-1k_reward0-200_denormalized-img_gradients-with-threshold0.05-mse-loss/checkpoint-10000/controlnet",
+
+    "softedge": "../diffusers/work_dirs/reward_model/MultiGen20M_Hed/reward_ft5k_controlnet_sd15_hed_res512_bs256_lr1e-5_warmup100_scale-1_iter10k_fp16_train0-1k_reward0-200/checkpoint-10000/controlnet",
+
+    "segmentation": "../diffusers/work_dirs/reward_model/Captioned_ADE20K/reward_ft_controlnet_sd15_seg_res512_bs256_lr1e-5_warmup100_scale-0.5_iter5k_fp16_train0-1k_reward0-200_FCN-R101-d8/checkpoint-5000/controlnet",
+
+    "depth": "../diffusers/work_dirs/reward_model/MultiGen20M_Depth/reward_ft5k_controlnet_sd15_depth_res512_bs256_lr1e-5_warmup100_scale-1.0_iter10k_fp16_train0-1k_reward0-200_mse-loss/checkpoint-10000/controlnet",
+
+    "lineart": "../diffusers/work_dirs/reward_model/MultiGen20M_LineDrawing/reward_ft5k_controlnet_sd15_lineart_res512_bs256_lr1e-5_warmup100_scale-10_iter10k_fp16_train0-1k_reward0-200/checkpoint-10000/controlnet",
 }
 
 
@@ -580,6 +576,7 @@ class Model:
                 image_resolution=image_resolution,
                 detect_resolution=preprocess_resolution,
             )
+        # NOTE: We still use the general lineart model
         if "anime" in preprocessor_name:
             self.load_controlnet_weight("lineart_anime")
         else: