From a06c11deaff79e37543e899b174e635b32af9c4f Mon Sep 17 00:00:00 2001 From: Anzhella Pankratova Date: Mon, 12 Aug 2024 18:38:25 +0300 Subject: [PATCH] Add args.control_type to choose controlnet condition --- README.md | 8 ++++---- main.py | 11 ++++++++--- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index d6762d7..b4a19ae 100644 --- a/README.md +++ b/README.md @@ -103,7 +103,7 @@ 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) @@ -111,7 +111,7 @@ python3 main.py \ 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) @@ -119,7 +119,7 @@ python3 main.py \ 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 index 3e8d2e0..d40bd76 100644 --- 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,