support value controller training

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)

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")

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):
+for i in [0.1, 0.3, 0.5, 0.7, 0.9]:
-    image = pipe(prompt="a cat", seed=0, value_controller_inputs=[i/10])
+    image = pipe(prompt="a cat on the beach", seed=2, value_controller_inputs=[i])
-    image.save(f"value_control_{i}.jpg")
+    image.save(f"value_control_{i}.jpg")

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

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

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")

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")