init demo

app.py

@@ -7,15 +7,8 @@ import torch
 
 from app_canny import create_demo as create_demo_canny
 from app_depth import create_demo as create_demo_depth
-from app_ip2p import create_demo as create_demo_ip2p
 from app_lineart import create_demo as create_demo_lineart
-from app_mlsd import create_demo as create_demo_mlsd
-from app_normal import create_demo as create_demo_normal
-from app_openpose import create_demo as create_demo_openpose
-from app_scribble import create_demo as create_demo_scribble
-from app_scribble_interactive import create_demo as create_demo_scribble_interactive
 from app_segmentation import create_demo as create_demo_segmentation
-from app_shuffle import create_demo as create_demo_shuffle
 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

app_canny.py

@@ -16,8 +16,8 @@ def create_demo(process):
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
-                image = gr.Image()
-                prompt = gr.Textbox(label="Prompt")
+                image = gr.Image(value='images/canny_demo.jpg')
+                prompt = gr.Textbox(label="Prompt", value='BEAUTIFUL PORTRAIT PAINTINGS BY EMMA UBER')
                 run_button = gr.Button("Run")
                 with gr.Accordion("Advanced options", open=False):
                     num_samples = gr.Slider(
@@ -31,10 +31,10 @@ def create_demo(process):
                         step=256,
                     )
                     canny_low_threshold = gr.Slider(
-                        label="Canny low threshold", minimum=1, maximum=255, value=100, step=1
+                        label="Canny low threshold", minimum=0, maximum=1.0, value=0.1, step=0.05
                     )
                     canny_high_threshold = gr.Slider(
-                        label="Canny high threshold", minimum=1, maximum=255, value=200, step=1
+                        label="Canny high threshold", minimum=0, maximum=1.0, value=0.2, step=0.05
                     )
                     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=7.5, step=0.1)

app_ip2p.py

@@ -1,87 +0,0 @@
-#!/usr/bin/env python
-
-import gradio as gr
-
-from settings import (
-    DEFAULT_IMAGE_RESOLUTION,
-    DEFAULT_NUM_IMAGES,
-    MAX_IMAGE_RESOLUTION,
-    MAX_NUM_IMAGES,
-    MAX_SEED,
-)
-from utils import randomize_seed_fn
-
-
-def create_demo(process):
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image()
-                prompt = gr.Textbox(label="Prompt")
-                run_button = gr.Button("Run")
-                with gr.Accordion("Advanced options", open=False):
-                    num_samples = gr.Slider(
-                        label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
-                    )
-                    image_resolution = gr.Slider(
-                        label="Image resolution",
-                        minimum=256,
-                        maximum=MAX_IMAGE_RESOLUTION,
-                        value=DEFAULT_IMAGE_RESOLUTION,
-                        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=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="high-quality, extremely detailed, 4K")
-                    n_prompt = gr.Textbox(
-                        label="Negative prompt",
-                        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=4, object_fit="scale-down")
-        inputs = [
-            image,
-            prompt,
-            a_prompt,
-            n_prompt,
-            num_samples,
-            image_resolution,
-            num_steps,
-            guidance_scale,
-            seed,
-        ]
-        prompt.submit(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name=False,
-        )
-        run_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name="ip2p",
-        )
-    return demo
-
-
-if __name__ == "__main__":
-    from model import Model
-
-    model = Model(task_name="ip2p")
-    demo = create_demo(model.process_ip2p)
-    demo.queue().launch()

app_mlsd.py

@@ -1,99 +0,0 @@
-#!/usr/bin/env python
-
-import gradio as gr
-
-from settings import (
-    DEFAULT_IMAGE_RESOLUTION,
-    DEFAULT_NUM_IMAGES,
-    MAX_IMAGE_RESOLUTION,
-    MAX_NUM_IMAGES,
-    MAX_SEED,
-)
-from utils import randomize_seed_fn
-
-
-def create_demo(process):
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image()
-                prompt = gr.Textbox(label="Prompt")
-                run_button = gr.Button("Run")
-                with gr.Accordion("Advanced options", open=False):
-                    num_samples = gr.Slider(
-                        label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
-                    )
-                    image_resolution = gr.Slider(
-                        label="Image resolution",
-                        minimum=256,
-                        maximum=MAX_IMAGE_RESOLUTION,
-                        value=DEFAULT_IMAGE_RESOLUTION,
-                        step=256,
-                    )
-                    preprocess_resolution = gr.Slider(
-                        label="Preprocess resolution", minimum=128, maximum=512, value=512, step=1
-                    )
-                    mlsd_value_threshold = gr.Slider(
-                        label="Hough value threshold (MLSD)", minimum=0.01, maximum=2.0, value=0.1, step=0.01
-                    )
-                    mlsd_distance_threshold = gr.Slider(
-                        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=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="high-quality, extremely detailed, 4K")
-                    n_prompt = gr.Textbox(
-                        label="Negative prompt",
-                        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=4, object_fit="scale-down")
-        inputs = [
-            image,
-            prompt,
-            a_prompt,
-            n_prompt,
-            num_samples,
-            image_resolution,
-            preprocess_resolution,
-            num_steps,
-            guidance_scale,
-            seed,
-            mlsd_value_threshold,
-            mlsd_distance_threshold,
-        ]
-        prompt.submit(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name=False,
-        )
-        run_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name="mlsd",
-        )
-    return demo
-
-
-if __name__ == "__main__":
-    from model import Model
-
-    model = Model(task_name="MLSD")
-    demo = create_demo(model.process_mlsd)
-    demo.queue().launch()

app_normal.py

@@ -1,95 +0,0 @@
-#!/usr/bin/env python
-
-import gradio as gr
-
-from settings import (
-    DEFAULT_IMAGE_RESOLUTION,
-    DEFAULT_NUM_IMAGES,
-    MAX_IMAGE_RESOLUTION,
-    MAX_NUM_IMAGES,
-    MAX_SEED,
-)
-from utils import randomize_seed_fn
-
-
-def create_demo(process):
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image()
-                prompt = gr.Textbox(label="Prompt")
-                run_button = gr.Button("Run")
-                with gr.Accordion("Advanced options", open=False):
-                    preprocessor_name = gr.Radio(
-                        label="Preprocessor", choices=["NormalBae", "None"], type="value", value="NormalBae"
-                    )
-                    num_samples = gr.Slider(
-                        label="Images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
-                    )
-                    image_resolution = gr.Slider(
-                        label="Image resolution",
-                        minimum=256,
-                        maximum=MAX_IMAGE_RESOLUTION,
-                        value=DEFAULT_IMAGE_RESOLUTION,
-                        step=256,
-                    )
-                    preprocess_resolution = gr.Slider(
-                        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=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="high-quality, extremely detailed, 4K")
-                    n_prompt = gr.Textbox(
-                        label="Negative prompt",
-                        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=4, object_fit="scale-down")
-        inputs = [
-            image,
-            prompt,
-            a_prompt,
-            n_prompt,
-            num_samples,
-            image_resolution,
-            preprocess_resolution,
-            num_steps,
-            guidance_scale,
-            seed,
-            preprocessor_name,
-        ]
-        prompt.submit(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name=False,
-        )
-        run_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name="normal",
-        )
-    return demo
-
-
-if __name__ == "__main__":
-    from model import Model
-
-    model = Model(task_name="NormalBae")
-    demo = create_demo(model.process_normal)
-    demo.queue().launch()

app_scribble.py

@@ -1,95 +0,0 @@
-#!/usr/bin/env python
-
-import gradio as gr
-
-from settings import (
-    DEFAULT_IMAGE_RESOLUTION,
-    DEFAULT_NUM_IMAGES,
-    MAX_IMAGE_RESOLUTION,
-    MAX_NUM_IMAGES,
-    MAX_SEED,
-)
-from utils import randomize_seed_fn
-
-
-def create_demo(process):
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image()
-                prompt = gr.Textbox(label="Prompt")
-                run_button = gr.Button("Run")
-                with gr.Accordion("Advanced options", open=False):
-                    preprocessor_name = gr.Radio(
-                        label="Preprocessor", choices=["HED", "PidiNet", "None"], type="value", value="HED"
-                    )
-                    num_samples = gr.Slider(
-                        label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
-                    )
-                    image_resolution = gr.Slider(
-                        label="Image resolution",
-                        minimum=256,
-                        maximum=MAX_IMAGE_RESOLUTION,
-                        value=DEFAULT_IMAGE_RESOLUTION,
-                        step=256,
-                    )
-                    preprocess_resolution = gr.Slider(
-                        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=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="high-quality, extremely detailed, 4K")
-                    n_prompt = gr.Textbox(
-                        label="Negative prompt",
-                        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=4, object_fit="scale-down")
-        inputs = [
-            image,
-            prompt,
-            a_prompt,
-            n_prompt,
-            num_samples,
-            image_resolution,
-            preprocess_resolution,
-            num_steps,
-            guidance_scale,
-            seed,
-            preprocessor_name,
-        ]
-        prompt.submit(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name=False,
-        )
-        run_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name="scribble",
-        )
-    return demo
-
-
-if __name__ == "__main__":
-    from model import Model
-
-    model = Model(task_name="scribble")
-    demo = create_demo(model.process_scribble)
-    demo.queue().launch()

app_scribble_interactive.py

@@ -1,115 +0,0 @@
-#!/usr/bin/env python
-
-import gradio as gr
-import numpy as np
-
-from settings import (
-    DEFAULT_IMAGE_RESOLUTION,
-    DEFAULT_NUM_IMAGES,
-    MAX_IMAGE_RESOLUTION,
-    MAX_NUM_IMAGES,
-    MAX_SEED,
-)
-from utils import randomize_seed_fn
-
-
-def create_canvas(w, h):
-    return np.zeros(shape=(h, w, 3), dtype=np.uint8) + 255
-
-
-def create_demo(process):
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                canvas_width = gr.Slider(
-                    label="Canvas width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_RESOLUTION,
-                    value=DEFAULT_IMAGE_RESOLUTION,
-                    step=1,
-                )
-                canvas_height = gr.Slider(
-                    label="Canvas height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_RESOLUTION,
-                    value=DEFAULT_IMAGE_RESOLUTION,
-                    step=1,
-                )
-                create_button = gr.Button("Open drawing canvas!")
-                image = gr.Image(tool="sketch", brush_radius=10)
-                prompt = gr.Textbox(label="Prompt")
-                run_button = gr.Button("Run")
-                with gr.Accordion("Advanced options", open=False):
-                    num_samples = gr.Slider(
-                        label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
-                    )
-                    image_resolution = gr.Slider(
-                        label="Image resolution",
-                        minimum=256,
-                        maximum=MAX_IMAGE_RESOLUTION,
-                        value=DEFAULT_IMAGE_RESOLUTION,
-                        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=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="high-quality, extremely detailed, 4K")
-                    n_prompt = gr.Textbox(
-                        label="Negative prompt",
-                        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=4, object_fit="scale-down")
-
-        create_button.click(
-            fn=create_canvas,
-            inputs=[canvas_width, canvas_height],
-            outputs=image,
-            queue=False,
-            api_name=False,
-        )
-
-        inputs = [
-            image,
-            prompt,
-            a_prompt,
-            n_prompt,
-            num_samples,
-            image_resolution,
-            num_steps,
-            guidance_scale,
-            seed,
-        ]
-        prompt.submit(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name=False,
-        )
-        run_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-        )
-    return demo
-
-
-if __name__ == "__main__":
-    from model import Model
-
-    model = Model(task_name="scribble")
-    demo = create_demo(model.process_scribble_interactive)
-    demo.queue().launch()

app_segmentation.py

@@ -17,7 +17,7 @@ def create_demo(process):
         with gr.Row():
             with gr.Column():
                 image = gr.Image(value='images/seg_demo.png')
-                prompt = gr.Textbox(label="Prompt", value='A large building with a pointed roof and several chimneys.')
+                prompt = gr.Textbox(label="Prompt", value='A large building with a pointed roof and several chimneys')
                 run_button = gr.Button("Run")
                 with gr.Accordion("Advanced options", open=False):
                     preprocessor_name = gr.Radio(

app_shuffle.py

@@ -1,91 +0,0 @@
-#!/usr/bin/env python
-
-import gradio as gr
-
-from settings import (
-    DEFAULT_IMAGE_RESOLUTION,
-    DEFAULT_NUM_IMAGES,
-    MAX_IMAGE_RESOLUTION,
-    MAX_NUM_IMAGES,
-    MAX_SEED,
-)
-from utils import randomize_seed_fn
-
-
-def create_demo(process):
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image()
-                prompt = gr.Textbox(label="Prompt")
-                run_button = gr.Button("Run")
-                with gr.Accordion("Advanced options", open=False):
-                    preprocessor_name = gr.Radio(
-                        label="Preprocessor", choices=["ContentShuffle", "None"], type="value", value="ContentShuffle"
-                    )
-                    num_samples = gr.Slider(
-                        label="Number of images", minimum=1, maximum=MAX_NUM_IMAGES, value=DEFAULT_NUM_IMAGES, step=1
-                    )
-                    image_resolution = gr.Slider(
-                        label="Image resolution",
-                        minimum=256,
-                        maximum=MAX_IMAGE_RESOLUTION,
-                        value=DEFAULT_IMAGE_RESOLUTION,
-                        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=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="high-quality, extremely detailed, 4K")
-                    n_prompt = gr.Textbox(
-                        label="Negative prompt",
-                        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=4, object_fit="scale-down")
-        inputs = [
-            image,
-            prompt,
-            a_prompt,
-            n_prompt,
-            num_samples,
-            image_resolution,
-            num_steps,
-            guidance_scale,
-            seed,
-            preprocessor_name,
-        ]
-        prompt.submit(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name=False,
-        )
-        run_button.click(
-            fn=randomize_seed_fn,
-            inputs=[seed, randomize_seed],
-            outputs=seed,
-            queue=False,
-            api_name=False,
-        ).then(
-            fn=process,
-            inputs=inputs,
-            outputs=result,
-            api_name="content-shuffle",
-        )
-    return demo
-
-
-if __name__ == "__main__":
-    from model import Model
-
-    model = Model(task_name="shuffle")
-    demo = create_demo(model.process_shuffle)
-    demo.queue().launch()

checkpoints/canny/controlnet/config.json

@@ -0,0 +1,52 @@
+{
+  "_class_name": "ControlNetModel",
+  "_diffusers_version": "0.26.3",
+  "_name_or_path": "work_dirs/finetune/MultiGen20M_canny/ft_controlnet_sd15_canny_res512_bs256_lr1e-5_warmup100_iter5k_fp16ft0-1000/checkpoint-5000",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "class_embed_type": null,
+  "conditioning_channels": 3,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "controlnet_conditioning_channel_order": "rgb",
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

checkpoints/canny/controlnet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3fd425077e65024addc5cf73c97195fcfd499b7a5e16868e4251b47cebb0d89
+size 1445157120

checkpoints/depth/controlnet/config.json

@@ -0,0 +1,52 @@
+{
+  "_class_name": "ControlNetModel",
+  "_diffusers_version": "0.26.3",
+  "_name_or_path": "work_dirs/finetune/MultiGen20M_depth/ft_controlnet_sd15_depth_res512_bs256_lr1e-5_warmup100_iter5k_fp16ft0-200/checkpoint-5000",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "class_embed_type": null,
+  "conditioning_channels": 3,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "controlnet_conditioning_channel_order": "rgb",
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

checkpoints/depth/controlnet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7450404d13ef888c9701433a3c17b2a86c021a6d042f9f5d2519602abd7f2f3
+size 1445157120

checkpoints/hed/controlnet/config.json

@@ -0,0 +1,51 @@
+{
+  "_class_name": "ControlNetModel",
+  "_diffusers_version": "0.19.3",
+  "_name_or_path": "work_dirs/reward_model/MultiGen20M_Hed/reward_ft5k_controlnet_sd15_hed_res512_bs256_lr1e-5_warmup100_scale-1_iter5k_fp16_train0-1k_reward0-200/checkpoint-5000",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "class_embed_type": null,
+  "conditioning_channels": 3,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "controlnet_conditioning_channel_order": "rgb",
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

checkpoints/hed/controlnet/diffusion_pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:721a7b3ec6b26bc73887f9f6d8a4fc175b01c785c0f986f3b7f15cd520cecf8e
+size 1445260234

checkpoints/lineart/controlnet/config.json

@@ -0,0 +1,51 @@
+{
+  "_class_name": "ControlNetModel",
+  "_diffusers_version": "0.19.3",
+  "_name_or_path": "work_dirs/reward_model/MultiGen20M_LineDrawing/reward_ft5k_controlnet_sd15_lineart_res512_bs256_lr1e-5_warmup100_scale-10_iter5k_fp16_train0-1k_reward0-200/checkpoint-5000",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "class_embed_type": null,
+  "conditioning_channels": 3,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "controlnet_conditioning_channel_order": "rgb",
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

checkpoints/lineart/controlnet/diffusion_pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3eab52fe2f7a3e2ad7841eeb7ce2d78271869ced226681ee83d83b8fa22a163a
+size 1445260234

checkpoints/seg/FCN_controlnet/config.json

@@ -0,0 +1,52 @@
+{
+  "_class_name": "ControlNetModel",
+  "_diffusers_version": "0.26.3",
+  "_name_or_path": "work_dirs/finetune/Captioned_ADE20K/ft_controlnet_sd15_seg_res512_bs256_lr1e-5_warmup100_iter5k_fp16/checkpoint-5000/controlnet",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "class_embed_type": null,
+  "conditioning_channels": 3,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "controlnet_conditioning_channel_order": "rgb",
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

checkpoints/seg/FCN_controlnet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c38bf06cd30bf31b4458cea39c6488a8f95c5ea7b9b5503c368aa0fef81a4e8
+size 1445157120

checkpoints/seg/controlnet/config.json

@@ -0,0 +1,51 @@
+{
+  "_class_name": "ControlNetModel",
+  "_diffusers_version": "0.19.3",
+  "_name_or_path": "work_dirs/finetune/Captioned_ADE20K/ft_controlnet_sd15_seg_res512_bs256_lr1e-5_warmup100_iter5k_fp16/checkpoint-5000/controlnet",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "class_embed_type": null,
+  "conditioning_channels": 3,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "controlnet_conditioning_channel_order": "rgb",
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

checkpoints/seg/controlnet/diffusion_pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11da8bb3cc8f400d097e49136741085af2ac87dc3508edd3fb7e15d99d963d96
+size 1445260234

model.py

@@ -18,15 +18,15 @@ from preprocessor import Preprocessor
 from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES
 
 CONTROLNET_MODEL_IDS = {
-    "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",
+    "Canny": "checkpoints/canny/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",
+    "softedge": "checkpoints/hed/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",
+    "segmentation": "checkpoints/seg/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",
+    "depth": "checkpoints/depth/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",
+    "lineart": "checkpoints/lineart/controlnet",
 }
 
 

preprocessor.py

@@ -3,6 +3,7 @@ import gc
 import numpy as np
 import PIL.Image
 import torch
+import torchvision
 from controlnet_aux import (
     CannyDetector,
     ContentShuffleDetector,
@@ -21,6 +22,32 @@ from cv_utils import resize_image
 from depth_estimator import DepthEstimator
 from image_segmentor import ImageSegmentor
 
+from kornia.core import Tensor
+from kornia.filters import canny
+
+
+class Canny:
+
+    def __call__(
+        self,
+        images: np.array,
+        low_threshold: float = 0.1,
+        high_threshold: float = 0.2,
+        kernel_size: tuple[int, int] | int = (5, 5),
+        sigma: tuple[float, float] | Tensor = (1, 1),
+        hysteresis: bool = True,
+        eps: float = 1e-6
+    ) -> torch.Tensor:
+
+        assert low_threshold is not None, "low_threshold must be provided"
+        assert high_threshold is not None, "high_threshold must be provided"
+
+        images = torch.from_numpy(images).permute(2, 0, 1).unsqueeze(0) / 255.0
+
+        images_tensor = canny(images, low_threshold, high_threshold, kernel_size, sigma, hysteresis, eps)[1]
+        images_tensor = (images_tensor[0][0].numpy() * 255).astype(np.uint8)
+        return images_tensor
+
 
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
@@ -49,7 +76,7 @@ class Preprocessor:
         elif name == "LineartAnime":
             self.model = LineartAnimeDetector.from_pretrained(self.MODEL_ID)
         elif name == "Canny":
-            self.model = CannyDetector()
+            self.model = Canny()
         elif name == "ContentShuffle":
             self.model = ContentShuffleDetector()
         elif name == "DPT":
@@ -70,7 +97,7 @@ class Preprocessor:
                 image = HWC3(image)
                 image = resize_image(image, resolution=detect_resolution)
             image = self.model(image, **kwargs)
-            return PIL.Image.fromarray(image)
+            return PIL.Image.fromarray(image).convert('RGB')
         elif self.name == "Midas":
             detect_resolution = kwargs.pop("detect_resolution", 512)
             image_resolution = kwargs.pop("image_resolution", 512)

settings.py

@@ -7,7 +7,7 @@ DEFAULT_MODEL_ID = os.getenv("DEFAULT_MODEL_ID", "runwayml/stable-diffusion-v1-5
 MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "4"))
 DEFAULT_NUM_IMAGES = min(MAX_NUM_IMAGES, int(os.getenv("DEFAULT_NUM_IMAGES", "4")))
 MAX_IMAGE_RESOLUTION = int(os.getenv("MAX_IMAGE_RESOLUTION", "768"))
-DEFAULT_IMAGE_RESOLUTION = min(MAX_IMAGE_RESOLUTION, int(os.getenv("DEFAULT_IMAGE_RESOLUTION", "768")))
+DEFAULT_IMAGE_RESOLUTION = min(MAX_IMAGE_RESOLUTION, int(os.getenv("DEFAULT_IMAGE_RESOLUTION", "512")))
 
 ALLOW_CHANGING_BASE_MODEL = os.getenv("SPACE_ID") != "hysts/ControlNet-v1-1"
 SHOW_DUPLICATE_BUTTON = os.getenv("SHOW_DUPLICATE_BUTTON") == "1"