From e3c5d2540b5b2c81358b2ae65d7d3e4acd360cff Mon Sep 17 00:00:00 2001 From: Artiprocher Date: Mon, 21 Jul 2025 19:16:30 +0800 Subject: [PATCH] support value controller training --- diffsynth/pipelines/flux_image_new.py | 4 +++- .../model_inference/FLUX.1-dev-AttriCtrl.py | 2 +- ...alueControl.py => FLUX.1-dev-AttriCtrl.py} | 9 ++++---- .../full/FLUX.1-dev-AttriCtrl.sh | 14 +++++++++++++ .../lora/FLUX.1-dev-AttriCtrl.sh | 17 +++++++++++++++ .../validate_full/FLUX.1-dev-AttriCtrl.py | 21 +++++++++++++++++++ .../validate_lora/FLUX.1-dev-AttriCtrl.py | 19 +++++++++++++++++ 7 files changed, 79 insertions(+), 7 deletions(-) rename examples/flux/model_inference_low_vram/{FLUX.1-dev-ValueControl.py => FLUX.1-dev-AttriCtrl.py} (65%) create mode 100644 examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh create mode 100644 examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh create mode 100644 examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py create mode 100644 examples/flux/model_training/validate_lora/FLUX.1-dev-AttriCtrl.py diff --git a/diffsynth/pipelines/flux_image_new.py b/diffsynth/pipelines/flux_image_new.py index 330667c..6525dd4 100644 --- a/diffsynth/pipelines/flux_image_new.py +++ b/diffsynth/pipelines/flux_image_new.py @@ -466,7 +466,7 @@ class FluxImagePipeline(BasePipeline): flex_control_strength: float = 0.5, flex_control_stop: float = 0.5, # Value Controller - value_controller_inputs: list[float] = None, + value_controller_inputs: Union[list[float], float] = None, # Step1x step1x_reference_image: Image.Image = None, # LoRA Encoder @@ -871,6 +871,8 @@ class FluxImageUnit_ValueControl(PipelineUnit): def process(self, pipe: FluxImagePipeline, prompt_emb, text_ids, value_controller_inputs): if value_controller_inputs is None: return {} + if not isinstance(value_controller_inputs, list): + value_controller_inputs = [value_controller_inputs] value_controller_inputs = torch.tensor(value_controller_inputs).to(dtype=pipe.torch_dtype, device=pipe.device) pipe.load_models_to_device(["value_controller"]) value_emb = pipe.value_controller(value_controller_inputs, pipe.torch_dtype) diff --git a/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py b/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py index 7dd4574..6c8d870 100644 --- a/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py +++ b/examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py @@ -15,5 +15,5 @@ pipe = FluxImagePipeline.from_pretrained( ) for i in [0.1, 0.3, 0.5, 0.7, 0.9]: - image = pipe(prompt="A woman.", seed=602, value_controller_inputs=[i], rand_device="cuda") + image = pipe(prompt="a cat on the beach", seed=2, value_controller_inputs=[i]) image.save(f"value_control_{i}.jpg") diff --git a/examples/flux/model_inference_low_vram/FLUX.1-dev-ValueControl.py b/examples/flux/model_inference_low_vram/FLUX.1-dev-AttriCtrl.py similarity index 65% rename from examples/flux/model_inference_low_vram/FLUX.1-dev-ValueControl.py rename to examples/flux/model_inference_low_vram/FLUX.1-dev-AttriCtrl.py index bb6be21..52edf2c 100644 --- a/examples/flux/model_inference_low_vram/FLUX.1-dev-ValueControl.py +++ b/examples/flux/model_inference_low_vram/FLUX.1-dev-AttriCtrl.py @@ -10,12 +10,11 @@ pipe = FluxImagePipeline.from_pretrained( ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors", offload_device="cpu", offload_dtype=torch.float8_e4m3fn), ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/", offload_device="cpu", offload_dtype=torch.float8_e4m3fn), ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors", offload_device="cpu", offload_dtype=torch.float8_e4m3fn), - ModelConfig(model_id="DiffSynth-Studio/FLUX.1-dev-ValueController", origin_file_pattern="single/prefer_embed/value.ckpt", offload_device="cpu", offload_dtype=torch.float8_e4m3fn) + ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors", offload_device="cpu", offload_dtype=torch.float8_e4m3fn) ], ) -pipe.load_lora(pipe.dit, ModelConfig(model_id="DiffSynth-Studio/FLUX.1-dev-ValueController", origin_file_pattern="single/dit_lora/dit_value.ckpt")) pipe.enable_vram_management() -for i in range(10): - image = pipe(prompt="a cat", seed=0, value_controller_inputs=[i/10]) - image.save(f"value_control_{i}.jpg") \ No newline at end of file +for i in [0.1, 0.3, 0.5, 0.7, 0.9]: + image = pipe(prompt="a cat on the beach", seed=2, value_controller_inputs=[i]) + image.save(f"value_control_{i}.jpg") diff --git a/examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh b/examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh new file mode 100644 index 0000000..91dc0cf --- /dev/null +++ b/examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh @@ -0,0 +1,14 @@ +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/example_image_dataset \ + --dataset_metadata_path data/example_image_dataset/metadata_attrictrl.csv \ + --data_file_keys "image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/,black-forest-labs/FLUX.1-dev:ae.safetensors,DiffSynth-Studio/AttriCtrl-FLUX.1-Dev:models/brightness.safetensors" \ + --learning_rate 1e-5 \ + --num_epochs 1 \ + --remove_prefix_in_ckpt "pipe.value_controller.encoders.0." \ + --output_path "./models/train/FLUX.1-dev-AttriCtrl_full" \ + --trainable_models "value_controller" \ + --extra_inputs "value_controller_inputs" \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh b/examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh new file mode 100644 index 0000000..7763c5f --- /dev/null +++ b/examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh @@ -0,0 +1,17 @@ +accelerate launch examples/flux/model_training/train.py \ + --dataset_base_path data/example_image_dataset \ + --dataset_metadata_path data/example_image_dataset/metadata_attrictrl.csv \ + --data_file_keys "image" \ + --max_pixels 1048576 \ + --dataset_repeat 100 \ + --model_id_with_origin_paths "black-forest-labs/FLUX.1-dev:flux1-dev.safetensors,black-forest-labs/FLUX.1-dev:text_encoder/model.safetensors,black-forest-labs/FLUX.1-dev:text_encoder_2/,black-forest-labs/FLUX.1-dev:ae.safetensors,DiffSynth-Studio/AttriCtrl-FLUX.1-Dev:models/brightness.safetensors" \ + --learning_rate 1e-4 \ + --num_epochs 5 \ + --remove_prefix_in_ckpt "pipe.dit." \ + --output_path "./models/train/FLUX.1-dev-AttriCtrl_lora" \ + --lora_base_model "dit" \ + --lora_target_modules "a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp" \ + --lora_rank 32 \ + --extra_inputs "value_controller_inputs" \ + --align_to_opensource_format \ + --use_gradient_checkpointing diff --git a/examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py b/examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py new file mode 100644 index 0000000..17384fc --- /dev/null +++ b/examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py @@ -0,0 +1,21 @@ +import torch +from diffsynth.pipelines.flux_image_new import FluxImagePipeline, ModelConfig +from diffsynth import load_state_dict + + +pipe = FluxImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/"), + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"), + ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors") + ], +) +state_dict = load_state_dict("models/train/FLUX.1-dev-AttriCtrl_full/epoch-0.safetensors") +pipe.value_controller.encoders[0].load_state_dict(state_dict) + +image = pipe(prompt="a cat", seed=0, value_controller_inputs=0.1, rand_device="cuda") +image.save("image_FLUX.1-dev-AttriCtrl_full.jpg") diff --git a/examples/flux/model_training/validate_lora/FLUX.1-dev-AttriCtrl.py b/examples/flux/model_training/validate_lora/FLUX.1-dev-AttriCtrl.py new file mode 100644 index 0000000..f44df0d --- /dev/null +++ b/examples/flux/model_training/validate_lora/FLUX.1-dev-AttriCtrl.py @@ -0,0 +1,19 @@ +import torch +from diffsynth.pipelines.flux_image_new import FluxImagePipeline, ModelConfig + + +pipe = FluxImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="flux1-dev.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder/model.safetensors"), + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="text_encoder_2/"), + ModelConfig(model_id="black-forest-labs/FLUX.1-dev", origin_file_pattern="ae.safetensors"), + ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors") + ], +) +pipe.load_lora(pipe.dit, "models/train/FLUX.1-dev-AttriCtrl_lora/epoch-3.safetensors", alpha=1) + +image = pipe(prompt="a cat", seed=0, value_controller_inputs=0.1, rand_device="cuda") +image.save("image_FLUX.1-dev-AttriCtrl_lora.jpg")