theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-20 23:58:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #699 from modelscope/AttrCtrl

Browse Source 
Support AttriCtrl
This commit is contained in:
Zhongjie Duan
2025-07-21 19:18:18 +08:00
committed by GitHub
parent 43a8d9768c e3c5d2540b
commit 2871535f3b
9 changed files with 96 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
diffsynth/configs/model_config.py
View File

@@ -107,7 +107,7 @@ model_loader_configs = [
(None, "023f054d918a84ccf503481fd1e3379e", ["flux_dit"], [FluxDiT], "civitai"),
(None, "d02f41c13549fa5093d3521f62a5570a", ["flux_dit"], [FluxDiT], "civitai"),
(None, "605c56eab23e9e2af863ad8f0813a25d", ["flux_dit"], [FluxDiT], "diffusers"),
(None, "3ede90c44b2c161240b659f3b8393c9d", ["flux_value_controller"], [SingleValueEncoder], "civitai"),
(None, "0629116fce1472503a66992f96f3eb1a", ["flux_value_controller"], [SingleValueEncoder], "civitai"),
(None, "280189ee084bca10f70907bf6ce1649d", ["cog_vae_encoder", "cog_vae_decoder"], [CogVAEEncoder, CogVAEDecoder], "diffusers"),
(None, "9b9313d104ac4df27991352fec013fd4", ["rife"], [IFNet], "civitai"),
(None, "6b7116078c4170bfbeaedc8fe71f6649", ["esrgan"], [RRDBNet], "civitai"),

2
diffsynth/models/flux_value_control.py
View File

@@ -18,7 +18,7 @@ class MultiValueEncoder(torch.nn.Module):
class SingleValueEncoder(torch.nn.Module):
def __init__(self, dim_in=256, dim_out=3072, prefer_len=32, computation_device=None):
def __init__(self, dim_in=256, dim_out=4096, prefer_len=32, computation_device=None):
super().__init__()
self.prefer_len = prefer_len
self.prefer_proj = TemporalTimesteps(num_channels=dim_in, flip_sin_to_cos=True, downscale_freq_shift=0, computation_device=computation_device)

25
diffsynth/pipelines/flux_image_new.py
View File

@@ -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
@@ -855,18 +855,30 @@ class FluxImageUnit_InfiniteYou(PipelineUnit):
class FluxImageUnit_ValueControl(PipelineUnit):
def __init__(self):
super().__init__(
seperate_cfg=True,
input_params_posi={"prompt_emb": "prompt_emb", "text_ids": "text_ids"},
input_params_nega={"prompt_emb": "prompt_emb", "text_ids": "text_ids"},
input_params=("value_controller_inputs",),
onload_model_names=("value_controller",)
)
def add_to_text_embedding(self, prompt_emb, text_ids, value_emb):
prompt_emb = torch.concat([prompt_emb, value_emb], dim=1)
extra_text_ids = torch.zeros((value_emb.shape[0], value_emb.shape[1], 3), device=value_emb.device, dtype=value_emb.dtype)
text_ids = torch.concat([text_ids, extra_text_ids], dim=1)
return prompt_emb, text_ids
def process(self, pipe: FluxImagePipeline, value_controller_inputs):
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)
value_emb = value_emb.unsqueeze(0)
return {"value_emb": value_emb}
prompt_emb, text_ids = self.add_to_text_embedding(prompt_emb, text_ids, value_emb)
return {"prompt_emb": prompt_emb, "text_ids": text_ids}
@@ -1049,7 +1061,6 @@ def model_fn_flux_image(
flex_condition=None,
flex_uncondition=None,
flex_control_stop_timestep=None,
value_emb=None,
step1x_llm_embedding=None,
step1x_mask=None,
step1x_reference_latents=None,
@@ -1155,12 +1166,6 @@ def model_fn_flux_image(
prompt_emb, image_rotary_emb, attention_mask = dit.process_entity_masks(hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids)
else:
prompt_emb = dit.context_embedder(prompt_emb)
# Value Control
if value_emb is not None:
prompt_emb = torch.concat([prompt_emb, value_emb], dim=1)
value_text_ids = torch.zeros((value_emb.shape[0], value_emb.shape[1], 3), device=value_emb.device, dtype=value_emb.dtype)
text_ids = torch.concat([text_ids, value_text_ids], dim=1)
# Original FLUX inference
image_rotary_emb = dit.pos_embedder(torch.cat((text_ids, image_ids), dim=1))
attention_mask = None

9
examples/flux/model_inference/FLUX.1-dev-ValueControl.py → examples/flux/model_inference/FLUX.1-dev-AttriCtrl.py
View File

@@ -10,11 +10,10 @@ pipe = FluxImagePipeline.from_pretrained(
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/FLUX.1-dev-ValueController", origin_file_pattern="single/prefer_embed/value.ckpt")
ModelConfig(model_id="DiffSynth-Studio/AttriCtrl-FLUX.1-Dev", origin_file_pattern="models/brightness.safetensors")
],
)
pipe.load_lora(pipe.dit, ModelConfig(model_id="DiffSynth-Studio/FLUX.1-dev-ValueController", origin_file_pattern="single/dit_lora/dit_value.ckpt"))
for i in range(10):
image = pipe(prompt="a cat", seed=0, value_controller_inputs=[i/10])
image.save(f"value_control_{i}.jpg")
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")

9
examples/flux/model_inference_low_vram/FLUX.1-dev-ValueControl.py → examples/flux/model_inference_low_vram/FLUX.1-dev-AttriCtrl.py
View File

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

14
examples/flux/model_training/full/FLUX.1-dev-AttriCtrl.sh Normal file
View File

@@ -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

17
examples/flux/model_training/lora/FLUX.1-dev-AttriCtrl.sh Normal file
View File

@@ -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

21
examples/flux/model_training/validate_full/FLUX.1-dev-AttriCtrl.py Normal file
View File

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

19
examples/flux/model_training/validate_lora/FLUX.1-dev-AttriCtrl.py Normal file
View File

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