Add args.control_type to choose controlnet condition

SamSynnada · Aug 12, 2024 · a06c11d · a06c11d
1 parent 6199971
commit a06c11d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -103,23 +103,23 @@ python3 main.py \
 python3 main.py \
     --prompt="a bright blue bird in the garden, natural photo cinematic, MM full HD" \
     --repo_id "XLabs-AI/flux-controlnet-canny" \
-    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
+    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny
 ```
 ![Example Picture 1](./assets/readme/examples/picture-1-rev1.png)
 
 ```bash
 python3 main.py \
     --prompt="a dark evil mysterius house with ghosts, cinematic, MM full HD" \
     --repo_id "XLabs-AI/flux-controlnet-canny" \
-    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
+    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny
 ```
 ![Example Picture 2](./assets/readme/examples/picture-2-rev1.png)
 
 ```bash
 python3 main.py \
     --prompt="man, 4k photo" \
     --repo_id "XLabs-AI/flux-controlnet-canny" \
-    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
+    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny
 ```
 
 ![Example Picture 3](./assets/readme/examples/picture-3-rev1.png)
@@ -128,7 +128,7 @@ python3 main.py \
 python3 main.py \
     --prompt="a oil painting woman sitting at chair and smiling, cinematic, MM full HD" \
     --repo_id "XLabs-AI/flux-controlnet-canny" \
-    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4
+    --name controlnet.safetensors --device cuda --offload --use_controlnet --image "input_image.jpg" --guidance 4 --control_type canny
 ```
 
 ![Example Picture 4](./assets/readme/examples/picture-4-rev1.png)

diff --git a/main.py b/main.py
@@ -30,11 +30,11 @@ def create_argparser():
     )
     parser.add_argument(
         "--lora_repo_id", type=str, default=None,
-        help="A HuggingFace repo id to download model (Controlnet)"
+        help="A HuggingFace repo id to download model (LoRA)"
     )
     parser.add_argument(
         "--lora_name", type=str, default=None,
-        help="A filename to download from HuggingFace"
+        help="A LoRA filename to download from HuggingFace"
     )
     parser.add_argument(
         "--device", type=str, default="cuda",
@@ -55,6 +55,11 @@ def create_argparser():
     parser.add_argument(
         "--lora_weight", type=float, default=0.9, help="Lora model strength (from 0 to 1.0)"
     )
+    parser.add_argument(
+        "--control_type", type=str, default="canny",
+        choices=("canny", "openpose", "depth", "hed", "hough", "tile"),
+        help="Name of controlnet condition, example: canny"
+    )
     parser.add_argument(
         "--model_type", type=str, default="flux-dev",
         choices=("flux-dev", "flux-dev-fp8", "flux-schnell"),
@@ -98,7 +103,7 @@ def main(args):
         print('load lora:', args.lora_repo_id, args.lora_name)
         xflux_pipeline.set_lora(None, args.lora_repo_id, args.lora_name, args.lora_weight)
     if args.use_controlnet:
-        xflux_pipeline.set_controlnet("canny", args.local_path, args.repo_id, args.name)
+        xflux_pipeline.set_controlnet(args.control_type, args.local_path, args.repo_id, args.name)
 
     result = xflux_pipeline(prompt=args.prompt,
                             controlnet_image=image,