updated

app.py

@@ -6,8 +6,10 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 from utils import utils, tools, preprocess
 
+BASE_MODEL_REPO_ID = "neta-art/neta-xl-2.0"
+BASE_MODEL_FILENAME = "neta-xl-v2.fp16.safetensors"
 VAE_PATH = "madebyollin/sdxl-vae-fp16-fix"
-REPO_ID = "Pbihao/ControlNeXt"
+CONTROLNEXT_REPO_ID = "Pbihao/ControlNeXt"
 UNET_FILENAME = "ControlAny-SDXL/anime_canny/unet.safetensors"
 CONTROLNET_FILENAME = "ControlAny-SDXL/anime_canny/controlnet.safetensors"
 CACHE_DIR = None
@@ -19,17 +21,17 @@ DEFAULT_NEGATIVE_PROMPT = "worst quality, abstract, clumsy pose, deformed hand,
 def ui():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model_file = hf_hub_download(
-        repo_id='neta-art/neta-xl-2.0',
-        filename='neta-xl-v2.fp16.safetensors',
+        repo_id=BASE_MODEL_REPO_ID,
+        filename=BASE_MODEL_FILENAME,
         cache_dir=CACHE_DIR,
     )
     unet_file = hf_hub_download(
-        repo_id=REPO_ID,
+        repo_id=CONTROLNEXT_REPO_ID,
         filename=UNET_FILENAME,
         cache_dir=CACHE_DIR,
     )
     controlnet_file = hf_hub_download(
-        repo_id=REPO_ID,
+        repo_id=CONTROLNEXT_REPO_ID,
         filename=CONTROLNET_FILENAME,
         cache_dir=CACHE_DIR,
     )
@@ -45,7 +47,6 @@ def ui():
         use_safetensors=True,
     )
 
-    preprocessors = ['canny']
     schedulers = ['Euler A', 'UniPC', 'Euler', 'DDIM', 'DDPM']
 
     css = """
@@ -57,7 +58,7 @@ def ui():
 
     with gr.Blocks(css=css) as demo:
         gr.Markdown(f"""
-        # [ControlNeXt-SDXL](https://github.com/dvlab-research/ControlNeXt) Demo
+        # [ControlNeXt-SDXL](https://github.com/dvlab-research/ControlNeXt) Demo (Anime Canny)
         Base model: [Neta-Art-XL-2.0](https://civitai.com/models/410737/neta-art-xl)
         """)
         with gr.Row():
@@ -80,13 +81,11 @@ def ui():
                         show_download_button=True,
                         show_share_button=True,
                     )
-                with gr.Row():
-                    processor = gr.Dropdown(
-                        label='Image Preprocessor',
-                        choices=preprocessors,
-                        value='canny',
-                    )
-                    process_button = gr.Button("Process", variant='primary', min_width=96, scale=0)
+                with gr.Accordion(label='Preprocess', open=True):
+                    with gr.Row():
+                        threshold1 = gr.Slider(minimum=-1, maximum=255, step=1, value=100, label='Threshold 1', info='-1 for auto')
+                        threshold2 = gr.Slider(minimum=-1, maximum=255, step=1, value=200, label='Threshold 2', info='-1 for auto')
+                        process_button = gr.Button("Process", variant='primary', min_width=96, scale=0)
                 with gr.Row():
                     scheduler = gr.Dropdown(
                         label='Scheduler',
@@ -142,7 +141,6 @@ def ui():
             seed,
         ):
             pipeline.scheduler = tools.get_scheduler(scheduler, pipeline.scheduler.config)
-
             generator = torch.Generator(device=device).manual_seed(max(0, min(seed, np.iinfo(np.int32).max))) if seed != -1 else None
 
             if control_image is None:
@@ -167,13 +165,12 @@ def ui():
 
         def process(
             image,
-            processor,
+            threshold1,
+            threshold2,
         ):
-            if image is None:
-                raise gr.Error('Please upload an image.')
-            processor = preprocess.get_extractor(processor)
-            image = processor(image)
-            return image
+            threshold1 = None if threshold1 == -1 else threshold1
+            threshold2 = None if threshold2 == -1 else threshold2
+            return preprocess.canny_extractor(image, threshold1, threshold2)
 
         generate_button.click(
             fn=generate,
@@ -183,7 +180,7 @@ def ui():
 
         process_button.click(
             fn=process,
-            inputs=[control_image, processor],
+            inputs=[control_image, threshold1, threshold2],
             outputs=[control_image],
         )