bugfix

app.py

@@ -2,51 +2,30 @@ import os
 import gradio as gr
 import numpy as np
 import random
-import spaces
-from diffusers import DiffusionPipeline
 import torch
 import json
 import logging
-from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
-from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
+from diffusers import DiffusionPipeline
 from huggingface_hub import login
-from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
-import copy
-import random
 import time
-import boto3
-from io import BytesIO
 from datetime import datetime
-from transformers import AutoTokenizer
-
-from diffusers import UNet2DConditionModel
-
-
-
+from io import BytesIO
+from diffusers.models.attention_processor import AttentionProcessor
+import re
+import json
+# 登录 Hugging Face Hub
 HF_TOKEN = os.environ.get("HF_TOKEN")
-
 login(token=HF_TOKEN)
 
-# init
-dtype = torch.bfloat16
+# 初始化
+dtype = torch.float16  # 您可以根据需要调整数据类型
 device = "cuda" if torch.cuda.is_available() else "cpu"
-base_model = "black-forest-labs/FLUX.1-dev"
-
-# unet = UNet2DConditionModel.from_pretrained(
-#     base_model,
-#     torch_dtype=torch.float16,
-#     use_safetensors=True,
-#     variant="fp16",
-#     subfolder="unet",
-# # ).to("cuda")
-# tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-
+base_model = "black-forest-labs/FLUX.1-dev"  # 替换为您的模型
 
+# 加载管道
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
 
-
-
-MAX_SEED = 2**32-1
+MAX_SEED = 2**32 - 1
 
 class calculateDuration:
     def __init__(self, activity_name=""):
@@ -64,133 +43,285 @@ class calculateDuration:
         else:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
 
+# 定义位置、偏移和区域的映射
+valid_locations = {  # x, y in 90*90
+    'in the center': (45, 45),
+    'on the left': (15, 45),
+    'on the right': (75, 45),
+    'on the top': (45, 15),
+    'on the bottom': (45, 75),
+    'on the top-left': (15, 15),
+    'on the top-right': (75, 15),
+    'on the bottom-left': (15, 75),
+    'on the bottom-right': (75, 75)
+}
 
-def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
-    print("upload_image_to_r2", account_id, access_key, secret_key, bucket_name)
-    connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
-
-    s3 = boto3.client(
-        's3',
-        endpoint_url=connectionUrl,
-        region_name='auto',
-        aws_access_key_id=access_key,
-        aws_secret_access_key=secret_key
-    )
-
-    current_time = datetime.now().strftime("%Y/%m/%d/%H%M%S")
-    image_file = f"generated_images/{current_time}_{random.randint(0, MAX_SEED)}.png"
-    buffer = BytesIO()
-    image.save(buffer, "PNG")
-    buffer.seek(0)
-    s3.upload_fileobj(buffer, bucket_name, image_file)
-    print("upload finish", image_file)
-    return image_file
+valid_offsets = {  # x, y in 90*90
+    'no offset': (0, 0),
+    'slightly to the left': (-10, 0),
+    'slightly to the right': (10, 0),
+    'slightly to the upper': (0, -10),
+    'slightly to the lower': (0, 10),
+    'slightly to the upper-left': (-10, -10),
+    'slightly to the upper-right': (10, -10),
+    'slightly to the lower-left': (-10, 10),
+    'slightly to the lower-right': (10, 10)
+}
 
+valid_areas = {  # w, h in 90*90
+    "a small square area": (50, 50),
+    "a small vertical area": (40, 60),
+    "a small horizontal area": (60, 40),
+    "a medium-sized square area": (60, 60),
+    "a medium-sized vertical area": (50, 80),
+    "a medium-sized horizontal area": (80, 50),
+    "a large square area": (70, 70),
+    "a large vertical area": (60, 90),
+    "a large horizontal area": (90, 60)
+}
 
-@spaces.GPU
-def generate_image(prompt, steps, seed, cfg_scale, width, height, progress):
-    pipe.to("cuda")
+# 解析角色位置的函数
+def parse_character_position(character_position):
+    # 定义正则表达式模式
+    location_pattern = '|'.join(re.escape(key) for key in valid_locations.keys())
+    offset_pattern = '|'.join(re.escape(key) for key in valid_offsets.keys())
+    area_pattern = '|'.join(re.escape(key) for key in valid_areas.keys())
+
+    # 提取位置
+    location_match = re.search(location_pattern, character_position, re.IGNORECASE)
+    location = location_match.group(0) if location_match else 'in the center'
+
+    # 提取偏移
+    offset_match = re.search(offset_pattern, character_position, re.IGNORECASE)
+    offset = offset_match.group(0) if offset_match else 'no offset'
+
+    # 提取区域
+    area_match = re.search(area_pattern, character_position, re.IGNORECASE)
+    area = area_match.group(0) if area_match else 'a medium-sized square area'
+
+    return {
+        'location': location,
+        'offset': offset,
+        'area': area
+}
 
-    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
+# 创建掩码的函数
+def create_attention_mask(image_width, image_height, location, offset, area):
+    # 图像在生成时通常会被缩放为 90x90，因此先定义一个基础尺寸
+    base_size = 90
+
+    # 获取位置坐标
+    loc_x, loc_y = valid_locations.get(location, (45, 45))
+    # 获取偏移量
+    offset_x, offset_y = valid_offsets.get(offset, (0, 0))
+    # 获取区域大小
+    area_width, area_height = valid_areas.get(area, (60, 60))
+
+    # 计算最终位置
+    final_x = loc_x + offset_x
+    final_y = loc_y + offset_y
+
+    # 将坐标和尺寸映射到实际图像尺寸
+    scale_x = image_width / base_size
+    scale_y = image_height / base_size
+
+    center_x = final_x * scale_x
+    center_y = final_y * scale_y
+    width = area_width * scale_x
+    height = area_height * scale_y
+
+    # 计算左上角和右下角坐标
+    x_start = int(max(center_x - width / 2, 0))
+    y_start = int(max(center_y - height / 2, 0))
+    x_end = int(min(center_x + width / 2, image_width))
+    y_end = int(min(center_y + height / 2, image_height))
+
+    # 创建掩码
+    mask = torch.zeros((image_height, image_width), dtype=torch.float32, device="cuda")
+    mask[y_start:y_end, x_start:x_end] = 1.0
+
+    # 展平成一维
+    mask_flat = mask.view(-1)  # 形状为 (image_height * image_width,)
+    return mask_flat
+
+# 自定义注意力处理器
+class CustomCrossAttentionProcessor(AttentionProcessor):
+    def __init__(self, masks, embeddings, adapter_names):
+        super().__init__()
+        self.masks = masks  # 列表，包含每个角色的掩码
+        self.embeddings = embeddings  # 列表，包含每个角色的嵌入
+        self.adapter_names = adapter_names  # 列表，包含每个角色的 LoRA 适配器名称
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, **kwargs):
+        # 获取当前的 adapter_name
+        adapter_name = getattr(attn, 'adapter_name', None)
+        if adapter_name is None or adapter_name not in self.adapter_names:
+            # 如果没有 adapter_name，直接执行默认的注意力计算
+            return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, **kwargs)
         
-    joint_attention_kwargs = {"attention_control": attention_control}
+        # 查找 adapter_name 对应的索引
+        idx = self.adapter_names.index(adapter_name)
+        mask = self.masks[idx]
         
-
+        # 标准的注意力计算
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        # 重塑以适应多头注意力
+        query = query.view(batch_size, -1, attn.heads, attn.head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, attn.head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, attn.head_dim).transpose(1, 2)
+
+        # 计算注意力得分
+        attention_scores = torch.matmul(query, key.transpose(-1, -2)) * attn.scale
+
+        # 应用掩码调整注意力得分
+        # 将 mask 调整为与 attention_scores 兼容的形状
+        # 假设 key_len 与 mask 的长度一致
+        mask_expanded = mask.unsqueeze(0).unsqueeze(0).unsqueeze(0)  # (1, 1, 1, key_len)
+        # 将掩码应用于 attention_scores
+        attention_scores += mask_expanded * 1e6  # 增强对应位置的注意力
+
+        # 计算注意力概率
+        attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1)
+
+        # 计算上下文向量
+        context = torch.matmul(attention_probs, value)
+
+        # 重塑回原始形状
+        context = context.transpose(1, 2).reshape(batch_size, -1, attn.heads * attn.head_dim)
+
+        # 输出投影
+        hidden_states = attn.to_out(context)
+        return hidden_states
+
+# 替换注意力处理器的函数
+def replace_attention_processors(pipe, masks, embeddings, adapter_names):
+    custom_processor = CustomCrossAttentionProcessor(masks, embeddings, adapter_names)
+    for name, module in pipe.unet.named_modules():
+        if hasattr(module, 'attn2'):
+            # 设置 adapter_name 为模块的属性
+            module.attn2.adapter_name = getattr(module, 'adapter_name', None)
+            module.attn2.processor = custom_processor
+
+# 生成图像的函数
+@spaces.GPU
+@torch.inference_mode()
+def generate_image_with_embeddings(prompt_embeddings, steps, seed, cfg_scale, width, height, progress):
+    pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
+
     with calculateDuration("Generating image"):
         # Generate image
-        generate_image = pipe(
-            prompt=prompt,
+        generated_image = pipe(
+            prompt_embeds=prompt_embeddings,
             num_inference_steps=steps,
             guidance_scale=cfg_scale,
             width=width,
             height=height,
             generator=generator,
-            joint_attention_kwargs=joint_attention_kwargs
         ).images[0]
-        
+    
     progress(99, "Generate success!")
-    return generate_image
-
-# 在 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
+    return generated_image
 
+# 主函数
 
-
-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)):
+def run_lora(prompt_bg, character_prompts_json, character_positions_json, lora_strings_json, prompt_details, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
+    
+    # 解析角色提示词、位置和 LoRA 字符串
+    try:
+        character_prompts = json.loads(character_prompts_json)
+        character_positions = json.loads(character_positions_json)
+        lora_strings = json.loads(lora_strings_json)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Invalid JSON input: {e}")
+    
+    # 确保提示词、位置和 LoRA 字符串的数量一致
+    if len(character_prompts) != len(character_positions) or len(character_prompts) != len(lora_strings):
+        raise ValueError("The number of character prompts, positions, and LoRA strings must be the same.")
+    
+    # 角色的数量
+    num_characters = len(character_prompts)
     
-
     # Load LoRA weights
-    if lora_strings:
-        with calculateDuration(f"Loading LoRA weights for {lora_strings}"):
-            pipe.unload_lora_weights()
-            lora_array = lora_strings.split(',')
-            adapter_names = []
-            for lora_string in lora_array:
-                parts = lora_string.split(':')
-                if len(parts) == 3:
-                    lora_repo, weights, adapter_name = parts
-                    # 调用 pipe.load_lora_weights() 方法加载权重
-                    pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
-                    adapter_names.append(adapter_name)
-                else:
-                    print(f"Invalid format for lora_string: {lora_string}")
-            
-            adapter_weights = [lora_scale] * len(adapter_names)
-            # 调用 pipeline.set_adapters 方法设置 adapter 和对应权重
-            pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
-
-
-        
+    with calculateDuration("Loading LoRA weights"):
+        pipe.unload_lora_weights()
+        adapter_names = []
+        for lora_info in lora_strings:
+            lora_repo = lora_info.get("repo")
+            weights = lora_info.get("weights")
+            adapter_name = lora_info.get("adapter_name")
+            if lora_repo and weights and adapter_name:
+                # 调用 pipe.load_lora_weights() 方法加载权重
+                pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                adapter_names.append(adapter_name)
+                # 将 adapter_name 设置为模型的属性
+                setattr(pipe.unet, 'adapter_name', adapter_name)
+            else:
+                raise ValueError("Invalid LoRA string format. Each item must have 'repo', 'weights', and 'adapter_name' keys.")
+        adapter_weights = [lora_scale] * len(adapter_names)
+        # 调用 pipeline.set_adapters 方法设置 adapter 和对应权重
+        pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+    
+    # 确保 adapter_names 的数量与角色数量匹配
+    if len(adapter_names) != num_characters:
+        raise ValueError("The number of LoRA adapters must match the number of characters.")
+    
     # Set random seed for reproducibility
     if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, progress)
+        with calculateDuration("Set random seed"):
+            seed = random.randint(0, MAX_SEED)
     
-    if upload_to_r2:
-        with calculateDuration("upload r2"):
-            url = upload_image_to_r2(final_image, account_id, access_key, secret_key, bucket)
-            result = {"status": "success", "url": url}
-    else:
-        result = {"status": "success", "message": "Image generated but not uploaded"}
+    # 编码提示词
+    with calculateDuration("Encoding prompts"):
+        # 编码背景提示词
+        bg_text_input = pipe.tokenizer(prompt_bg, return_tensors="pt").to("cuda")
+        bg_embeddings = pipe.text_encoder(bg_text_input.input_ids.to(device))[0]
     
-    progress(100, "Completed!")
+        # 编码角色提示词
+        character_embeddings = []
+        for prompt in character_prompts:
+            char_text_input = pipe.tokenizer(prompt, return_tensors="pt").to("cuda")
+            char_embeddings = pipe.text_encoder(char_text_input.input_ids.to(device))[0]
+            character_embeddings.append(char_embeddings)
+    
+        # 编码互动细节提示词
+        details_text_input = pipe.tokenizer(prompt_details, return_tensors="pt").to("cuda")
+        details_embeddings = pipe.text_encoder(details_text_input.input_ids.to(device))[0]
+    
+        # 合并背景和互动细节的嵌入
+        prompt_embeddings = torch.cat([bg_embeddings, details_embeddings], dim=1)
+    
+    # 解析角色位置
+    character_infos = []
+    for position_str in character_positions:
+        info = parse_character_position(position_str)
+        character_infos.append(info)
+    
+    # 创建角色的掩码
+    masks = []
+    for info in character_infos:
+        mask = create_attention_mask(width, height, info['location'], info['offset'], info['area'])
+        masks.append(mask)
+    
+    # 替换注意力处理器
+    replace_attention_processors(pipe, masks, character_embeddings, adapter_names)
+    
+    # Generate image
+    final_image = generate_image_with_embeddings(prompt_embeddings, steps, seed, cfg_scale, width, height, progress)
+    
+    # 您可以在此处添加上传图片的代码
+    result = {"status": "success", "message": "Image generated"}
 
-    yield final_image, seed, json.dumps(result)
+    progress(100, "Completed!")
 
+    return final_image, seed, json.dumps(result)
 
+# Gradio 界面
 css="""
 #col-container {
     margin: 0 auto;
@@ -199,12 +330,16 @@ css="""
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("Flux with lora")
+    gr.Markdown("Flux with LoRA")
     with gr.Row():
         
         with gr.Column():
-            prompt = gr.Text(label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False)
-            lora_strings = gr.Text( label="lora_strings", max_lines=1, placeholder="Enter a lora strings", visible=True)        
+            
+            prompt_bg = gr.Text(label="Background Prompt", placeholder="Enter background/scene prompt", lines=2)
+            character_prompts = gr.Text(label="Character Prompts (JSON List)", placeholder='["Character 1 prompt", "Character 2 prompt"]', lines=5)
+            character_positions = gr.Text(label="Character Positions (JSON List)", placeholder='["Character 1 position", "Character 2 position"]', lines=5)
+            lora_strings_json = gr.Text(label="LoRA Strings (JSON List)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1"}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2"}]', lines=5)
+            prompt_details = gr.Text(label="Interaction Details", placeholder="Enter interaction details between characters", lines=2)
             run_button = gr.Button("Run", scale=0)
 
             with gr.Accordion("Advanced Settings", open=False):
@@ -215,11 +350,11 @@ with gr.Blocks(css=css) as demo:
                     lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.5)
 
                 with gr.Row():
-                    width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
-                    height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
+                    width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=512)
+                    height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=512)
 
                 with gr.Row():
-                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
+                    cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=7.5)
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28) 
 
                 upload_to_r2 = gr.Checkbox(label="Upload to R2", value=False)
@@ -231,13 +366,35 @@ with gr.Blocks(css=css) as demo:
 
         with gr.Column():
             result = gr.Image(label="Result", show_label=False)
-            json_text = gr.Text()
-
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn = run_lora,
-        inputs = [prompt, cfg_scale, steps, lora_strings, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket],
-        outputs=[result, seed, json_text]
+            seed_output = gr.Text(label="Seed")
+            json_text = gr.Text(label="Result JSON")
+
+    inputs = [
+        prompt_bg,
+        character_prompts,
+        character_positions,
+        lora_strings_json,
+        prompt_details,
+        cfg_scale,
+        steps,
+        randomize_seed,
+        seed,
+        width,
+        height,
+        lora_scale,
+        upload_to_r2,
+        account_id,
+        access_key,
+        secret_key,
+        bucket
+    ]
+
+    outputs = [result, seed_output, json_text]
+
+    run_button.click(
+        fn=run_lora,
+        inputs=inputs,
+        outputs=outputs
     )
 
-demo.queue().launch()
+demo.queue().launch()