try to custom attention control

app.py

@@ -17,6 +17,7 @@ import time
 import boto3
 from io import BytesIO
 from datetime import datetime
+from transformers import AutoTokenizer
 
 from diffusers import UNet2DConditionModel
 
@@ -37,7 +38,8 @@ base_model = "black-forest-labs/FLUX.1-dev"
 #     use_safetensors=True,
 #     variant="fp16",
 #     subfolder="unet",
-# ).to("cuda")
+# # ).to("cuda")
+# tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
 
 
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
@@ -85,9 +87,37 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
     return image_file
 
 
-
+@spaces.GPU
 def generate_image(prompt, steps, seed, cfg_scale, width, height, progress):
     pipe.to("cuda")
+
+    text_inputs = pipe.tokenizer(prompt, return_tensors="pt").to("cuda")
+    input_ids = text_inputs.input_ids[0]
+        
+    # 获取每个主体对应的令牌 ID
+    boy_token_id = pipe.tokenizer.convert_tokens_to_ids("boy_asia_05")
+    print(boy_token_id)
+    girl_token_id = pipe.tokenizer.convert_tokens_to_ids("girl_asia_04")
+    print(girl_token_id)
+    # 找到每个主体在输入中的索引位置
+    boy_indices = (input_ids == boy_token_id).nonzero(as_tuple=True)[0]
+    girl_indices = (input_ids == girl_token_id).nonzero(as_tuple=True)[0]
+        
+    # 准备 cross_attention_kwargs
+    def attention_control(attention_probs, adapter_name):
+        # 根据 adapter_name 和令牌索引控制注意力
+        print("attention_control", adapter_name)
+        if adapter_name == "boy_asia_05":
+            # 对女孩的令牌注意力设为零
+            attention_probs[:, :, :, girl_indices] = 0
+        elif adapter_name == "girl_asia_04":
+            # 对男孩的令牌注意力设为零
+            attention_probs[:, :, :, boy_indices] = 0
+        return attention_probs
+        
+    joint_attention_kwargs = {"attention_control": attention_control}
+        
+
     generator = torch.Generator(device="cuda").manual_seed(seed)
     with calculateDuration("Generating image"):
         # Generate image
@@ -98,13 +128,26 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, progress):
             width=width,
             height=height,
             generator=generator,
-            joint_attention_kwargs={"scale": 1}
+            joint_attention_kwargs=joint_attention_kwargs
         ).images[0]
         
     progress(99, "Generate success!")
     return generate_image
 
-@spaces.GPU
+# 在 Transformer 中，自定义注意力处理器
+class CustomAttentionProcessor(torch.nn.Module):
+    def __init__(self, attention_control, adapter_name):
+        super().__init__()
+        self.attention_control = attention_control
+        self.adapter_name = adapter_name
+
+    def forward(self, attention_probs):
+        # 调用自定义的注意力控制函数
+        attention_probs = self.attention_control(attention_probs, self.adapter_name)
+        return attention_probs
+
+
+
 def run_lora(prompt, cfg_scale, steps, lora_strings, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
     
 
@@ -127,6 +170,8 @@ def run_lora(prompt, cfg_scale, steps, lora_strings, randomize_seed, seed, width
             adapter_weights = [lora_scale] * len(adapter_names)
             # 调用 pipeline.set_adapters 方法设置 adapter 和对应权重
             pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+
+
         
     # Set random seed for reproducibility
     if randomize_seed: